[Fix] Keep the source register in the int8/int4b to float16 cast

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)

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)