[CuTeDSL] Fix export_to_c shape slot width for 64-bit dynamic dims - #3448
Open
wilyan09007 wants to merge 1 commit into
Open
[CuTeDSL] Fix export_to_c shape slot width for 64-bit dynamic dims#3448wilyan09007 wants to merge 1 commit into
wilyan09007 wants to merge 1 commit into
Conversation
The C header generator derives a tensor descriptor's stride type from the tensor but passes shape_type="int32_t" unconditionally, so a tensor built with cute.sym_int64() dimensions gets int32_t shape slots while the compiled function reads them as i64. A C caller that fills in the generated struct passes a descriptor half the expected size, and every field from the second shape slot on is misread. The widths are already known: _tensor_binding keeps a SymInt dimension at its declared width and gives every other dynamic dimension 32 bits. Derive the shape type the same way at both call sites. One array cannot hold dimensions of two widths, so a tensor mixing them is refused, and the list/tuple descriptor check now compares width alongside the masks. Fixes NVIDIA#3447 Signed-off-by: William <wilyan090@gmail.com>
Contributor
|
LGTM |
Contributor
|
Just curious, in what real cases do we need i64 shape representation? |
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
3 participants
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.
Summary
_generate_argumentsderives a tensor descriptor's stride type from the tensor but passesshape_type="int32_t"unconditionally, socute.sym_int64()dimensions getint32_tshape slots while the compiled function reads them asi64. A C caller filling in that struct hands over a descriptor half the expected size.The widths are already known one layer down:
_tensor_bindingkeeps aSymIntdimension at its declared width and gives every other dynamic dimension 32 bits. This derives the shape type the same way at both call sites. One array cannot hold dimensions of two widths, so a tensor mixing them is refused, and the list/tuple descriptor check now compares width alongside the masks it already had.Verification
Added
test/python/CuTeDSL/test_export_c_header_shape_type.py, 8 tests. Four fail on main: thesym_int64slot width, the mixed-width rejection, and both tensor-list cases. Thesym_int32, static-shape, and DLPack-dynamic cases pass before and after.Fixes #3447