• Lei Wang's avatar
    [Language] Support tile operator `T.cumsum` (#423) · 88747fcd
    Lei Wang authored
    * [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
    88747fcd
test_tilelang_language_cumsum.py 3.35 KB