From 3ade880e19d05b1113e808f2fbdd97e5c379cc31 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sat, 7 Nov 2026 12:43:50 +0100 Subject: [PATCH 1/2] [Fix] Keep the source register in the int8/int4b to float16 cast cast_i8_to_f16 and cast_i4_to_f16 biased the source in place with buffer_store(src_uint32, ...). Any instruction that read the register after the cast saw the values with their sign bits flipped. Put the biased value in a temporary instead. Signed-off-by: 0z5a --- python/tilus/backends/emitters/cast.py | 17 +++--- tests/kernels/test_cast.py | 72 +++++++++++++++++++++++++- 2 files changed, 80 insertions(+), 9 deletions(-) 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..2b87c065 100644 --- a/tests/kernels/test_cast.py +++ b/tests/kernels/test_cast.py @@ -12,14 +12,22 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import re import struct +from pathlib import Path import pytest import tilus import torch -from tilus.hidet.ir.dtypes import bfloat16, float16, float32, int8, int16, int32 +from tilus.backends.codegen import generate_ir_module +from tilus.drivers import BuildOptions, optimize_ir_module, optimize_program +from tilus.hidet.backend.codegen import codegen +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.target import nvgpu_sm80, scope +from tilus.testing import requires from tilus.utils import dtype_to_torch from tqdm import tqdm @@ -251,3 +259,65 @@ 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)] + + +@pytest.mark.parametrize("dtype, vector", CAST_REUSE_CASES) +def test_cast_to_float16_does_not_write_source_register_codegen(dtype, vector, tmp_path: Path): + with scope(nvgpu_sm80): + script = CastThenReuseSource(dtype, vector) + program = script._jit_instance_for(None, None, None).transpiled_programs[0] # type: ignore + ir_module = generate_ir_module(optimize_program(program, BuildOptions(), tmp_path)) + ir_module = optimize_ir_module(ir_module, tmp_path) + codegen(ir_module, src_out_path=str(tmp_path / "source.cu"), target="cuda") + + source = (tmp_path / "source.cu").read_text() + assert "src_uint32" in source + assert re.search(r"src_uint32\[[^\]]*\] = ", source) is None + + +@requires.nvgpu_sm80 +@pytest.mark.parametrize("dtype, vector", CAST_REUSE_CASES) +def test_cast_to_float16_keeps_source_register(dtype, vector): + device = "cuda" if torch.cuda.is_available() else "cpu" + 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) From c7147f8a3932e84f967ad50492baa7b06ee2950e Mon Sep 17 00:00:00 2001 From: Yaoyao Ding Date: Mon, 9 Nov 2026 12:23:21 -0500 Subject: [PATCH 2/2] Address review: simplify the cast source-register tests - Remove the hand-rolled codegen test. It drove optimize_program, generate_ir_module, optimize_ir_module and codegen directly, which is brittle against internal APIs. The normal build path cannot replace it without losing its no-GPU value, because compile_source requires CUDA. The GPU test covers the behaviour. - Use "cuda" directly in the GPU test. The "cpu" fallback was dead: compile-only mode also needs CUDA to build. - Add larger vectors to CAST_REUSE_CASES so the non-unrolled loop body of the int8/int4b to float16 cast is compiled and run. Co-Authored-By: Claude Opus 5.5 --- tests/kernels/test_cast.py | 24 ++---------------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/tests/kernels/test_cast.py b/tests/kernels/test_cast.py index 2b87c065..e9bad3b6 100644 --- a/tests/kernels/test_cast.py +++ b/tests/kernels/test_cast.py @@ -12,21 +12,15 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import re import struct -from pathlib import Path import pytest import tilus import torch -from tilus.backends.codegen import generate_ir_module -from tilus.drivers import BuildOptions, optimize_ir_module, optimize_program -from tilus.hidet.backend.codegen import codegen 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.target import nvgpu_sm80, scope from tilus.testing import requires from tilus.utils import dtype_to_torch from tqdm import tqdm @@ -286,27 +280,13 @@ def __call__(self, src_ptr: tilus.void_p, cast_ptr: tilus.void_p, copy_ptr: tilu self.store_global(g_copy, r_src, offsets=[0]) -CAST_REUSE_CASES = [(int8, 4), (int4b, 8)] - - -@pytest.mark.parametrize("dtype, vector", CAST_REUSE_CASES) -def test_cast_to_float16_does_not_write_source_register_codegen(dtype, vector, tmp_path: Path): - with scope(nvgpu_sm80): - script = CastThenReuseSource(dtype, vector) - program = script._jit_instance_for(None, None, None).transpiled_programs[0] # type: ignore - ir_module = generate_ir_module(optimize_program(program, BuildOptions(), tmp_path)) - ir_module = optimize_ir_module(ir_module, tmp_path) - codegen(ir_module, src_out_path=str(tmp_path / "source.cu"), target="cuda") - - source = (tmp_path / "source.cu").read_text() - assert "src_uint32" in source - assert re.search(r"src_uint32\[[^\]]*\] = ", source) is None +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" if torch.cuda.is_available() else "cpu" + 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)