Repository navigation
[Fix] Keep the source register in the int8/int4b to float16 cast - #171
Merged
Merged
Conversation
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 <dezhen.lu@student.uni-tuebingen.de>
- 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 <noreply@anthropic.com>
Member
|
/ok to test c7147f8 |
Member
|
thank you for the recent 4 PRs! @0z5a |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Casting an
int8orint4bregister tofloat16changes the source register. Anything that reads the register after the cast gets wrong values. This PR leaves the source intact.Why
cast_i8_to_f16andcast_i4_to_f16bias the input to unsigned before the conversion. They do it in place:src_uint32points at the source register, so the emitted kernel containssrc_uint32[0] = (src_uint32[0] ^ ...). The cast result is correct, but a later read of the same register sees every element with its sign bit flipped. Forint4bthe mask is0x88888888. No error is raised. A second.to(...)or a store of the register silently produces wrong data.Only these two specialized paths are affected. The other paths in
cast.pywrite only to the destination and their own temporaries.How
Put the biased value in a local variable and read that in the following
prmtand shift steps. The arithmetic is otherwise the same.Testing
Host tests ran on the Grace CPU of a GH200: GCC 13.3, CUDA 13.2, PyTorch 2.14.1+cu130, tilus
mainat f23018a.New tests in
tests/kernels/test_cast.py. Each script casts the register, then stores both the result and the register itself:maintest_cast_to_float16_does_not_write_source_register_codegen(int8, int4b): no store intosrc_uint32in the generated sourcetest_cast_to_float16_keeps_source_register(int8, int4b): checks the cast values and the stored registerGPU tests ran on GH200 (
sm_90, CUDA 13.2). Not run on B300.The rest of the CPU-runnable suite (
tests/lang,tests/ir,tests/drivers,tests/transformsand the top-level tests) gives the same result onmainand here: 192 passed, and the same 5 failed because they need a GPU.ruff checkandruff format --check(0.11.0 frompyproject.toml, 0.12.3 from pre-commit),mypy(1.15.0) and the copyright check pass on the changed files.