diff --git a/python/tilus/backends/emitters/cast.py b/python/tilus/backends/emitters/cast.py index c1a14cfb..42d44994 100644 --- a/python/tilus/backends/emitters/cast.py +++ b/python/tilus/backends/emitters/cast.py @@ -159,12 +159,12 @@ def cast_i8_to_f16(self, src: Var, dst: Var) -> None: # cast 4 int8 to 4 float16 at each time # range for int8: -128 ~ 127, uint8: 0 ~ 255 - # transform int8[a, b, c, d] to uint8[a, b, c, d] + 128 - self.buffer_store(buf=src_uint32, indices=[i], value=src_uint32[i] ^ uint32(0x80808080)) + # transform int8[a, b, c, d] to uint8[a, b, c, d] + 128, without changing the source register + biased = self.declare_var("biased", uint32, init=src_uint32[i] ^ uint32(0x80808080)) # transform int8[a, b, c, d] + 128 to float16[a, b, c, d] + (1024 + 128) - self.append(prmt(d=~dst_uint32[i << 1], a=src_uint32[i], b=uint32(0x00000064), c=uint32(0x4140))) - self.append(prmt(d=~dst_uint32[i << 1 | 1], a=src_uint32[i], b=uint32(0x00000064), c=uint32(0x4342))) + self.append(prmt(d=~dst_uint32[i << 1], a=biased, b=uint32(0x00000064), c=uint32(0x4140))) + self.append(prmt(d=~dst_uint32[i << 1 | 1], a=biased, b=uint32(0x00000064), c=uint32(0x4342))) # transform float16[a, b, c, d] + (1024 + 128) to float16[a, b, c, d] # where 0x6480 = 0b0110010010000000 = 1152 @@ -186,13 +186,14 @@ def cast_i4_to_f16(self, src: Var, dst: Var) -> None: # cast 8 int4b to float16 at each time # range for int4b: -8 ~ 7 - # transform int4b[7, 6, 5, 4, 3, 2, 1, 0] to uint4b[7, 6, 5, 4, 3, 2, 1, 0] + 8 - self.buffer_store(buf=src_uint32, indices=[i], value=src_uint32[i] ^ uint32(0x88888888)) + # transform int4b[7, 6, 5, 4, 3, 2, 1, 0] to uint4b[7, 6, 5, 4, 3, 2, 1, 0] + 8, without changing the source + # register + biased = self.declare_var("biased", uint32, init=src_uint32[i] ^ uint32(0x88888888)) # transform int4b[7, 6, 5, 4, 3, 2, 1, 0] + 8 to float16 [4, 0], [5, 1], [6, 2], [7, 3] imm_lut = (0xF0 & 0xCC) | 0xAA - a = self.declare(v=Var("a", type=uint32), init=src_uint32[i]) - b = self.declare(v=Var("b", type=uint32), init=src_uint32[i] >> 8) + a = self.declare(v=Var("a", type=uint32), init=biased) + b = self.declare(v=Var("b", type=uint32), init=biased >> 8) h = self.declare(v=tensor_var("h", shape=[4], dtype=uint32)) self.append( lop3(~h[0], a, uint32(0x000F000F), uint32(0x64006400), imm_lut=imm_lut) diff --git a/tests/kernels/test_cast.py b/tests/kernels/test_cast.py index 2bfda5f7..e9bad3b6 100644 --- a/tests/kernels/test_cast.py +++ b/tests/kernels/test_cast.py @@ -17,9 +17,11 @@ import pytest import tilus import torch -from tilus.hidet.ir.dtypes import bfloat16, float16, float32, int8, int16, int32 +from tilus.hidet.ir.dtypes import bfloat16, float16, float32, int4b, int8, int16, int32 from tilus.hidet.ir.dtypes.floats_subbyte import FloatSubbyteType from tilus.hidet.ir.dtypes.integer_subbyte import IntegerSubbyteType +from tilus.ir.layout.ops import spatial +from tilus.testing import requires from tilus.utils import dtype_to_torch from tqdm import tqdm @@ -251,3 +253,51 @@ def test_cast_int_subbyte_to_int32(dtype: IntegerSubbyteType) -> None: y = tilus.kernels.cast(x, dtype=tilus.int32) y_expected = tilus.full(shape=[128], fill_value=cast_ix_to_i32_ref(n, dtype=dtype), dtype=tilus.int32) torch.testing.assert_close(actual=y.torch(), expected=y_expected.torch(), equal_nan=True) + + +class CastThenReuseSource(tilus.Script): + """Cast a register to float16, then store the register itself, which the cast must leave intact.""" + + def __init__(self, dtype, vector): + super().__init__() + self.dtype = dtype + self.layout = spatial(32).local(vector) + + def __call__(self, src_ptr: tilus.void_p, cast_ptr: tilus.void_p, copy_ptr: tilus.void_p): + self.attrs.blocks = 1 + self.attrs.warps = 1 + + shape = self.layout.shape + g_src = self.global_view(src_ptr, dtype=self.dtype, shape=shape) + g_cast = self.global_view(cast_ptr, dtype=float16, shape=shape) + g_copy = self.global_view(copy_ptr, dtype=self.dtype, shape=shape) + + r_src = self.load_global(g_src, offsets=[0], shape=shape) + self.annotate_layout(r_src, self.layout) + r_cast = r_src.to(float16) + self.annotate_layout(r_cast, self.layout) + self.store_global(g_cast, r_cast, offsets=[0]) + self.store_global(g_copy, r_src, offsets=[0]) + + +CAST_REUSE_CASES = [(int8, 4), (int4b, 8), (int8, 16), (int4b, 16), (int4b, 32)] + + +@requires.nvgpu_sm80 +@pytest.mark.parametrize("dtype, vector", CAST_REUSE_CASES) +def test_cast_to_float16_keeps_source_register(dtype, vector): + device = "cuda" + n = 32 * vector + src = torch.randint(0, 256, (n * dtype.nbits // 8,), dtype=torch.uint8, device=device) + casted = torch.zeros(n, dtype=torch.float16, device=device) + copy = torch.zeros_like(src) + + CastThenReuseSource(dtype, vector)(src, casted, copy) + + if dtype == int8: + expected = src.to(torch.int8) + else: + nibbles = torch.stack([src & 0xF, src >> 4], dim=1).flatten().to(torch.int8) + expected = torch.where(nibbles >= 8, nibbles - 16, nibbles) + torch.testing.assert_close(casted, expected.to(torch.float16)) + torch.testing.assert_close(copy, src)