[Language] Support tile operator T.cumsum
#423
Merged
Conversation
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
* Added CumSumOp class for cumulative sum operations, including argument validation and lowering logic. * Introduced CumSum2D template for CUDA, supporting both forward and reverse cumulative sums. * Created tests for CumSum functionality in shared memory and fragment contexts. * Updated language interface to include cumsum operation, enhancing the reduction capabilities of TileLang. * Refactored reduce.py to support cumsum functionality with appropriate memory allocation and copying mechanisms.
LeiWang1999
added a commit
to LeiWang1999/tilelang
that referenced
this pull request
Jul 18, 2025
* [Feature] Implement CumSum operation in TileLang * Added CumSumOp class for cumulative sum operations, including argument validation and lowering logic. * Introduced CumSum2D template for CUDA, supporting both forward and reverse cumulative sums. * Created tests for CumSum functionality in shared memory and fragment contexts. * Updated language interface to include cumsum operation, enhancing the reduction capabilities of TileLang. * Refactored reduce.py to support cumsum functionality with appropriate memory allocation and copying mechanisms. * lint fix
LeiWang1999
added a commit
to LeiWang1999/tilelang
that referenced
this pull request
Jul 20, 2025
* [Feature] Implement CumSum operation in TileLang * Added CumSumOp class for cumulative sum operations, including argument validation and lowering logic. * Introduced CumSum2D template for CUDA, supporting both forward and reverse cumulative sums. * Created tests for CumSum functionality in shared memory and fragment contexts. * Updated language interface to include cumsum operation, enhancing the reduction capabilities of TileLang. * Refactored reduce.py to support cumsum functionality with appropriate memory allocation and copying mechanisms. * lint fix
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
This pull request introduces a new cumulative sum (
cumsum) operation in the TileLang framework, along with its implementation, integration, and testing. The changes include adding theCumSumOpoperator, implementing the CUDA kernel forcumsum, updating the TileLang language interface, and adding comprehensive tests for both shared memory and fragment scopes.New Feature:
cumsumOperationOperator Implementation:
CumSumOpclass insrc/op/reduce.hand its corresponding registration insrc/op/reduce.cc. This operator supports cumulative summation along a specified dimension, with an option for reverse mode. It includes methods for lowering to target-specific code and layout inference. [1] [2]CUDA Kernel:
CumSum2Dtemplate insrc/tl_templates/cuda/reduce.hfor efficient cumulative summation on 2D data. It supports both forward and reverse modes and handles different thread configurations.Integration into TileLang
cumsummacro and function to the TileLang language interface intilelang/language/reduce.py, allowing users to invoke thecumsumoperation on buffers. The implementation supports both shared memory and fragment scopes.cumsumin the TileLang language module (tilelang/language/__init__.py).Testing
New Tests for
cumsum:test_tilelang_language_cumsum.pyto validate thecumsumoperation for both shared memory and fragment scopes. It includes tests for different data types, dimensions, and reverse modes.Improved Debugging:
torch_assert_closeutility intilelang/utils/tensor.pyto include detailed mismatch information, including the left-hand side (LHS) and right-hand side (RHS) tensors.