Tensor load padding support #408
Conversation
Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com>
Hardcode84
requested review from
Megan0704-1,
Copilot,
harsh-nod and
raikonenfnu
There was a problem hiding this comment.
Pull Request Overview
This PR adds support for hardware padding in tensor load operations to shared memory, allowing the compiler to generate padded memory layouts when constraints are met. Previously, all padding was cleared unconditionally.
Key changes:
- Introduced
is_pow2()utility function to check if a value is a power of 2 - Added validation logic to determine when padding is supported by the tensor operation hardware constraints
- Modified padding clearing to be conditional based on hardware support rather than unconditional
- Renamed property
get_unpadded_dimstounpadded_dimsand added complementaryunpadded_shapeproperty
Reviewed Changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| wave_lang/kernel/wave/utils/general_utils.py | Added is_pow2() utility function for power-of-2 checks |
| wave_lang/kernel/wave/tensor_load_to_shared.py | Added padding validation and conditional clearing logic; imported is_literal and is_pow2 |
| wave_lang/kernel/wave/generate_bounds_exprs.py | Updated to use renamed unpadded_dims property |
| wave_lang/kernel/wave/codegen/read_write.py | Added padding configuration to tensor descriptor; imported subs_idxc |
| wave_lang/kernel/ops/wave_ops.py | Renamed get_unpadded_dims property to unpadded_dims and added unpadded_shape property |
| lit_tests/kernel/wave/mma.py | Updated test expectations to verify padded memory allocation |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Megan0704-1
left a comment
Megan0704-1
left a comment
There was a problem hiding this comment.
Thank you @Hardcode84!
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com>
Megan0704-1
left a comment
Megan0704-1
left a comment
There was a problem hiding this comment.
LGTM! Thanks @Hardcode84
Tensor load op support shared memory padding with some restrictions. The the padded allocation shape and only drop padding if it's not supported. --------- Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com>
2 participants
Tensor load op support shared memory padding with some restrictions.
The the padded allocation shape and only drop padding if it's not supported.