mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
c31ff58a5fbaa6f42f8a30d930e7e778ee93264c
374
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Chenhao Xu
|
c31ff58a5f |
Unified packed x2 intrinsics with multi-dtype support and bug fixes (#1978)
* fix: bits product to have lanes included * fix: fix T.view to view non vectorized type to vectorized type with same total bytes * feat: bfloat16 vec type along with all vectorized ops * fix: fix tests |
||
|
Tong WU
|
88049820fe |
[Feature] Introduce annotation for minBlocksPerMultiprocessor in __launch_bounds__ (#1979)
* Add annotate_min_blocks_per_sm for __launch_bounds__(..., minBlocksPerSM) - Python: T.annotate_min_blocks_per_sm(n) sets tl.min_blocks_per_sm on the prim_func - Codegen: emit second launch_bounds argument from attr; default remains 1 when unset - Export from tilelang.language; document attr key in builtin.h Made-with: Cursor * test * lint |
||
|
Lei Wang
|
e2c88335c6 |
Fix wrapped pre-loop TMA prefixes in producer-consumer WS (#1975)
* Support wrapped pre-loop TMA prefixes in producer-consumer WS * Format producer-consumer WS changes * Keep pure-TMA barrier grouping unchanged |
||
|
|
a2d6e01866 |
[Bugfix] Fix double buffer versioning when TMA is used without warp specialization (#1962)
* [Bugfix] Fix double buffer versioning when TMA is used without warp specialization When TMA is enabled but warp specialization is disabled, both MultiVersionBuffer and InjectSoftwarePipeline independently added version dimensions to shared data buffers, resulting in 4x buffer allocation instead of the intended 2x. Add a `barrier_only` mode to MultiVersionBuffer so that in the no-WS path it only versions barrier buffers (for mbarrier parity rewriting), while data buffer versioning is left to InjectSoftwarePipeline. Also flatten the if/else branches in OptimizeForTarget into a single linear flow with a three-way branch (WS / TMA-noWS / noTMA), removing the duplicated LowerOpaqueBlock call. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * [Enhancement] Introduce TMA attribute handling in LowerTileOp and optimization pipeline - Added a new PrimFunc attribute `tl.has_tma` to indicate whether TMA operations were generated during the lowering process. - Updated `OptimizeForTarget` to utilize the `tl.has_tma` attribute for determining the appropriate optimization path, enhancing clarity and efficiency in the pipeline. - Refactored the `module_has_tma` function to check for TMA operations across all functions in the module, streamlining the decision-making process for TMA-related optimizations. This change improves the maintainability of the code and ensures that TMA operations are correctly accounted for in the optimization pipeline. * fix * remove barrier related test for cutedsl. --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
||
|
Chaofan Lin
|
d0090b428e |
[Feature] Support alloc global workspace (#1940)
* add docstring * add hoist global allocation pass * fix rebase * fix * fix tvm * fix |
||
|
Lei WangandZhiwen Mo
|
715ca303da |
[Bugfix] Handle int64 offsets in ThreadSync for tvm_access_ptr (#1952)
Updated the ThreadSync implementation to support int64 offsets in tvm_access_ptr, preventing crashes during index reconstruction. Added a regression test to ensure correct behavior with shared/shared.dyn pointers. Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei Wang
|
b5bd642479 |
[Refactor] Separate gemm into explicit wgmma_gemm and tcgen05_gemm functions (#1949)
* Refactor GEMM operations to use explicit WGMMA and TCGEN5MMA functions - Replaced instances of T.gemm with T.wgmma_gemm and T.tcgen05_gemm in various example scripts to enhance performance and control over asynchronous execution. - Updated related documentation to reflect the changes in GEMM function usage, ensuring clarity on the new explicit interfaces. - Removed unnecessary wg_wait parameters from GEMM calls, aligning with the new function signatures. This refactor improves the efficiency of matrix multiplication operations in the TileLang framework, particularly for Hopper and Blackwell architectures. * lint fix * Refactor GEMM operations to unify function signatures - Replaced deprecated GEMM function variants (gemm_v1, gemm_v2, tcgen05_gemm_v1, tcgen05_gemm_v2) with a single unified function (gemm, tcgen05_gemm) across examples and tests. - Updated test cases to reflect the new function signatures, ensuring consistency and clarity in GEMM operation usage. - This change simplifies the API and enhances maintainability within the TileLang framework. * Refactor GEMM initialization to improve annotation checks - Simplified the handling of "is_wgmma" and "is_tcgen05" annotations by removing redundant checks and adding explicit error messages for type validation. - Enhanced the validation for the "mbar" argument to ensure it is a BufferLoadNode, improving error handling and code clarity. - These changes streamline the initialization process for GEMM operations, enhancing maintainability and robustness in the TileLang framework. |
||
|
Lei WangandZhiwen Mo
|
5a538f3b37 |
Support packed subtype views during layout reshape (#1947)
* Support packed subtype views in layout reshape * lint fix --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei WangandClaude Opus 4.6
|
9246b82b24 |
[Refactor] Replace create_list_of_mbarrier with buffer-based T.alloc_barrier (#1944)
* [Refactor] Replace create_list_of_mbarrier/get_mbarrier with buffer-based T.alloc_barrier Unify the two mbarrier representation paths into a single buffer-based approach using shared.barrier scope buffers with barrier_init block annotations. This eliminates the create_list_of_mbarrier and get_mbarrier intermediate intrinsics, making the IR cleaner and the barrier handling consistent across user-explicit (T.alloc_barrier) and compiler-generated (T.copy/T.tma_copy) paths. Key changes: - LowerTileOp now injects a shared.barrier buffer + barrier_init annotation into the root block instead of emitting create_list_of_mbarrier statements - copy.cc uses BufferLoad on the barrier buffer instead of get_mbarrier - WarpSpecializedRewriter and ProducerConsumerWS create barrier buffers with annotations instead of create_list_of_mbarrier calls - LowerSharedBarrier runs a second time after WS passes to process newly created barriers - LowerHopperIntrin no longer handles mbarrier allocation/init (fully delegated to LowerSharedBarrier) - Removed create_list_of_mbarrier and get_mbarrier op definitions from builtin.h/cc - Removed get_mbarrier codegen handling from CUDA, CuTeDSL, HIP backends - Updated WarpSpecializedDetector, inject_tma_barrier, and eliminate_storage_sync passes to detect barriers via shared.barrier scope buffers instead of the removed intrinsics Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * style: apply clang-format and ruff-format Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Refactor mbarrier handling in copy operations Updated the handling of mbarriers in the `LowerBulkCopy` and `LowerBulkCopy1D` methods to allocate a single barrier slot instead of multiple for each pipeline stage. This change simplifies the barrier management by allowing the `MultiVersionBuffer` to expand the barrier for pipelining stages. Additionally, modified the `LowerArgs` structure to use a pointer for the `mbarrier_buffer`, ensuring compatibility with lazy buffer creation. Key changes: - Simplified mbarrier allocation logic in `copy.cc`. - Updated `LowerArgs` to use a pointer for `mbarrier_buffer`. - Adjusted related comments for clarity. This refactor enhances the consistency and readability of the code while maintaining functionality. * Remove deprecated TMA barrier and warp specialization files This commit deletes the `inject_tma_barrier.cc`, `warp_specialized_rewriter.cc`, and associated test files as part of a cleanup effort. These files are no longer necessary following recent refactoring and updates to the TMA pipeline and warp specialization logic. The removal simplifies the codebase and eliminates outdated implementations that have been superseded by newer approaches. Key changes: - Deleted `inject_tma_barrier.cc` and `warp_specialized_rewriter.cc`. - Removed related test file `test_tilelang_transform_warp_specialized.py`. - Updated `phase.py` and `__init__.py` to reflect the removal of the WarpSpecialized and InjectTmaBarrier functionalities. This cleanup enhances maintainability and clarity in the transformation pipeline. --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
||
|
Lei WangandZhiwen Mo
|
f8dc61c02c |
test: reduce CI runtime for slow Python suites (#1932)
* test: trim CI runtime for slow Python suites * remove useless test * Update TVM submodule to latest commit and remove manual device assert test from CI. Replace with a call to tilelang.testing.main(). * lint fix * fix * revert test for eager jit --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei WangandClaude Opus 4.6
|
ded6a99221 |
[Feature] Add Producer-Consumer Warp Specialization and T.tma_copy() API (#1909)
* [Feature] Add Producer-Consumer Warp Specialization and T.tma_copy() API
This PR introduces Producer-Consumer Warp Specialization for sm90+ TMA
pipelines and a new T.tma_copy() API for explicit mbarrier management.
Key changes:
1. **ProducerConsumerWarpSpecialized pass** (new):
Splits pipelined TMA loops into producer (TMA loads) and consumer
(compute) warp groups with back-pressure barriers for buffer reuse.
Works with num_stages >= 1. Producer warp (128 threads) handles
arrive_expect_tx + tma_load; consumer warps handle compute + arrive
on back-pressure barriers.
2. **T.tma_copy() API** (new):
Fire-and-forget TMA copy with a required `barrier` parameter. Unlike
T.copy() which emits producer+wait pairs, T.tma_copy() emits only
arrive_expect_tx + tma_load. User manages synchronization via
T.mbarrier_wait_parity().
3. **MultiVersionBuffer barrier expansion**:
Extended to handle `shared.barrier` scope buffers, expanding them for
pipelining (1D size multiplication instead of prepending a dimension).
Auto-computes mbarrier parity as `(k // num_stages) % 2`.
4. **Pass ordering** (phase.py):
LowerSharedBarrier now runs after MultiVersionBuffer in the TMA path
so barrier buffers retain their `shared.barrier` scope during expansion.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* fix for annotate reg alloc
* Enhance debugging and performance measurement in examples
- Added print statements to output kernel source for both forward and backward sparse MLA examples, aiding in debugging.
- Updated `sparse_mla_fwd_pipelined.py` to run performance regression tests and print average latency, improving performance tracking.
- Disabled auto cache in `example_mha_sink_fwd_bhsd.py` for better control during testing.
* fix
* Enhance debugging and performance tracking in examples
- Added print statements to output kernel source in various example scripts, including `example_gqa_sink_bwd_bhsd.py`, `example_blocksparse_gemm.py`, and `example_dequant_groupedgemm_bf16_mxfp4_hopper.py`, aiding in debugging.
- Updated default `window_size` in `example_gqa_sink_bwd_bhsd.py` for improved configuration.
- Disabled auto cache in `example_gemm_schedule.py` to enhance control during testing.
- Introduced a new test for preserving protected auto-injected wait groups in the async copy optimization process.
* clean
* Remove unnecessary print statements for kernel source in example scripts
- Eliminated print statements from `example_gqa_sink_bwd_bhsd.py`, `example_blocksparse_gemm.py`, and `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` to clean up output and improve readability.
- These changes streamline the examples while maintaining functionality.
* clean
* lint fix
* Refactor T.tma_copy() and barrier synchronization in examples and tests
- Updated `example_warp_specialize_gemm_copy_1_gemm_0.py` to use `T.tma_copy()` for optimized memory copying with explicit barrier management.
- Enhanced test cases in `test_tilelang_issue_tma_no_ws.py` and `test_tilelang_language_tma_copy.py` to reflect changes in barrier synchronization, ensuring proper functionality with the new T.tma_copy() API.
- Removed unnecessary print statements and adjusted barrier allocations for clarity and performance.
- Cleaned up the `test_tilelang_transform_inject_tma_barrier.py` file by deleting it, as it was no longer needed.
* Add buffer data mapping for remapped buffers in multi_version_buffer_rewriter
- Introduced a check to ensure that the data variable of a buffer is discoverable, allowing the barrier_init annotation update to find the remapped buffer correctly.
- This change enhances the functionality of the buffer remapping process, ensuring proper synchronization in the transformation pipeline.
* Implement TMA im2col load synchronization with mbarrier allocation
- Added support for allocating mbarriers to synchronize TMA im2col loads, enhancing the pipeline's barrier management.
- Introduced logic to handle mbarrier creation and integration into the TMA copy process, ensuring proper synchronization and data transfer.
- Updated the handling of the mbarrier in the TMA copy statement to improve performance and maintain functionality within the thread-gated block.
* Implement TMA im2col load synchronization with mbarrier allocation
- Added support for allocating mbarriers to synchronize TMA im2col loads, enhancing the pipeline's barrier management.
- Introduced logic to handle mbarrier initialization and integration within the TMA copy process, ensuring proper synchronization during memory operations.
- Updated the handling of the mbarrier in the im2col operation to improve performance and maintain consistency in multi-threaded environments.
* disable tma lower for test_cutedsl_barrier
* Refactor code formatting in test_tilelang_jit_cutedsl.py
- Improved readability by adjusting the formatting of the `tilelang.compile` function call.
- Enabled the main testing function in the script to run without being commented out.
* Enhance sparse_mla_fwd.py and thread_storage_sync.cc with code cleanup and new functionality
- Added a call to disable cache in the main execution block of `sparse_mla_fwd.py` for improved performance.
- Commented out unused benchmarking code in `test_sparse_mla_fwd` to streamline testing.
- Introduced new utility functions in `thread_storage_sync.cc` to analyze and optimize shared memory access patterns, enhancing synchronization management.
- Improved readability and organization of the code by restructuring and commenting on key sections.
* Add thread storage synchronization utilities and optimizations
- Introduced functions to identify shared-like scopes, buffers, and pointer variables, enhancing memory management.
- Implemented a `BarrierTransparentStmtAnalyzer` to analyze statement transparency concerning shared memory.
- Developed a `ThreadStorageSyncOptimizer` to optimize sequences of statements by removing unnecessary synchronization calls, improving performance in multi-threaded environments.
- Enhanced handling of shared memory access in various call types, ensuring proper synchronization during memory operations.
* remove optimize thread_storage_sync
* Refactor pipeline planning to enhance cp.async group scheduling
- Introduced a last-use index for cp.async groups to improve scheduling accuracy.
- Removed unnecessary tracking of first consumers, simplifying the scheduling logic.
- Updated the handling of cp.async and commit statements to ensure proper group ordering based on last-use.
- Added a new test to verify that cp.async groups are ordered correctly by their last-use index.
* lint fix
* Add CUDA requirements to TMA copy tests and update main execution call
- Added decorators to require CUDA and a minimum compute version for the `test_tma_copy_pipeline_2_stages` and `test_tma_copy_pipeline_3_stages` functions.
- Updated the main execution block to call `tilelang.testing.main()` instead of directly invoking the test functions, improving test execution management.
* Refactor attention sink examples by removing pipelined implementations
- Updated imports in `benchmark_gqa_sink_fwd.py` and `benchmark_mha_sink_fwd.py` to use the non-pipelined versions of the attention sink functions.
- Deleted the pipelined implementations of `example_gqa_sink_fwd_bhsd_wgmma_pipelined.py` and `example_mha_sink_fwd_bhsd_wgmma_pipelined.py`.
- Cleaned up regression and test files to remove references to the deleted pipelined examples, ensuring consistency across the codebase.
* performance fix
* Enhance example_blocksparse_gemm.py and inject_pipeline.cc with debugging and utility improvements
- Added a print statement to output the kernel source in `example_blocksparse_gemm.py` for debugging purposes.
- Updated `inject_pipeline.cc` to include a new utility function for stripping TMA copy write buffer attributes, improving code clarity and functionality.
- Modified `producer_consumer_ws.cc` to utilize the new utility function, enhancing the handling of TMA barriers.
- Expanded the `_compile_tvm_ffi` function in the test file to accept additional keyword arguments, improving flexibility for test configurations.
- Introduced a new test to verify that pure TMA warp specialization correctly emits barriers, ensuring compliance with expected behavior.
* lint fix
* Refactor guard variable declarations in producer_consumer_ws.cc for clarity
- Changed guard variable declarations from non-const to const references to improve performance and readability.
- This minor adjustment enhances the handling of producer guards in the context of TMA barriers.
* lint fix
* fix
* lint fix
* perf fix
* Add print statement for kernel source and update main execution in example_dequant_gemm_w4a8.py
- Introduced a print statement to output the kernel source for debugging purposes.
- Updated the main execution block to call `run_regression_perf()` and print the latency, enhancing performance measurement capabilities.
* fix
* Add debugging output and update execution flow in example_tilelang_block_sparse_attn.py
- Added a print statement to display the kernel source for better debugging.
- Updated the main execution block to disable cache and call `run_regression_perf()`, printing the resulting latency for performance measurement.
- Enhanced the producer-consumer logic in producer_consumer_ws.cc with new utility functions to manage local variable access and improve code clarity.
- Introduced a new test case in test_tilelang_issue_tma_no_ws.py to ensure consumer-only local initializations do not leak into the producer, verifying correct behavior in TMA contexts.
* Revert "Add debugging output and update execution flow in example_tilelang_block_sparse_attn.py"
This reverts commit
|
||
|
Lei WangandZhiwen Mo
|
fee041c3cd |
Fix predicated cp.async pipeline scheduling (#1937)
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei WangandZhiwen Mo
|
8fe367c81f |
[Runtime] Improve TMA descriptor diagnostics (#1931)
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei WangandZhiwen Mo
|
d338c4ff95 |
[Bugfix] Fix concurrent TempDirectory creation during CUDA compilation (#1926)
* [Refactor] Update TVM subproject and enhance CUDA artifact path resolution - Updated the TVM subproject to commit 470430cf. - Introduced a new helper function `_resolve_artifact_paths` in `nvcc.py` to streamline the resolution of CUDA artifact paths, improving code clarity and maintainability. - Refactored the `compile_cuda` function to utilize the new helper for managing temporary file paths based on the specified output directory. * [Enhancement] Add CUDA requirement decorator to matmul test - Introduced the `@tilelang.testing.requires_cuda` decorator to the `test_matmul` function, ensuring that the test is only executed in environments with CUDA support. This change improves test reliability and clarity regarding dependencies. * lint fix --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Liu Yunuo
|
7b6b64f93e |
[Enhancement] Add eager-mode support for tilelang.autotune (#1906)
* [Enhancement] Add eager-mode support for tilelang.autotune * Fix type annotation and improve code structure in tuner.py * Fix misleading docstring in autotune tests |
||
|
Lei WangandZhiwen Mo
|
ccfc127d91 |
[Feature] Add CUDA intrinsic for isfinite operation (#1925)
* [Feature] Add CUDA intrinsic for isfinite operation - Introduced DispatchCUDAIsFinite function to handle the isfinite intrinsic for CUDA, enabling support for checking finite values in floating-point computations. - Registered the new intrinsic with TVM under the operation name "tir.isfinite", enhancing the CUDA intrinsic capabilities. Updated submodule reference for TVM to the latest commit. * merge cluster mod --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei WangandZhiwen Mo
|
bdbb32c99b |
Support ptr-table grouped GEMM kernels (#1923)
Add ptr-table grouped GEMM support Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
xyyy1420
|
ff1b8cc56c |
[Feature] Add T.gemm support for CPU target (#1904)
* [Feature] Add T.gemm support for CPU target - Add GemmInst::kCPU enum and TargetIsCPU() helper for CPU target detection - Add CPU branch in getGemmInst() for both GemmNode and GemmPyNode - Add CPU early return in GemmNode::Lower() and InferLayout() (gemm_v1 no-op) - Add GemmScalar: triple-loop scalar gemm implementation (gemm_v2 path) - Add CPU dispatch in GemmPy._get_implementation_class() - Add pytest tests for codegen and correctness (NN/TN/NT/TT, multiple sizes) * rename kCPU to kScalar for consistency with existing GemmInst naming |
||
|
Lei Wang
|
4929ad89cc |
[Refactor] Improve cp.async lowering and add async_copy op (#1887)
* Add CPAsync copy selection/lowering and vectorized ptx_cp_async folding
* [transform] Rename SimplifyCPAsyncSync to OptimizeCPAsyncSync
* ThreadSync: treat ptx_cp_async as async-copy
* Implement async_copy operator and enhance cp.async handling
- Introduced the async_copy operator for explicit asynchronous global-to-shared memory copies, ensuring no automatic waits are inserted.
- Updated the copy operator to normalize copy regions and handle both synchronous and asynchronous copy semantics.
- Enhanced pipeline planning to recognize explicit cp.async calls and bind commit and wait operations to the appropriate stages.
- Added tests to validate the behavior of async_copy and its integration with pipeline planning.
This commit improves the handling of asynchronous memory operations and clarifies the semantics of copy operations in the context of the new async_copy functionality.
* docs: clarify T.copy vs T.async_copy
* style: apply format.sh fixes
* async_copy: improve errors and defer cp.async check
* test: check T.copy vs T.async_copy codegen equivalence on SM80
* copy: allow predicated cp.async for OOB global->shared
* pipeline: schedule cp.async commits with cp.async groups
* lower: preserve manual pipeline statement arity
* pipeline: skip commit-order check when fused with cp.async
* pipeline: enforce cp.async ordering in PipelinePlanning
PipelinePlanning may float pure control intrinsics like ptx_commit_group earlier because they have no buffer regions, which can generate commit-before-cp_async.\n\nAdd an order-repair step that models cp.async group semantics (cp_async* -> commit, commit(group_i) -> cp_async(group_{i+1})) and wait->consumer constraints, then applies a stable topological sort to produce a valid software_pipeline_order.\n\nAlso adds a regression test to assert cp.async is ordered before commit when split into separate blocks.
* Enhance cp.async synchronization handling in pipeline planning
- Introduced logic to relax loop wait conditions and split epilogue waits for cp.async operations, improving overlap between consumer phases.
- Added a new method to identify and manage consumer phases based on async memory reads, ensuring correct synchronization.
- Updated tests to validate the new behavior, confirming that wait groups are correctly relaxed and inserted as needed in the optimized pipeline.
This commit refines the handling of cp.async synchronization, enhancing performance and correctness in the generated code.
* lint fix
* Refactor cp.async copy handling and improve vectorization logic
- Removed the unused IsCPAsyncElementBytes function to streamline the code.
- Enhanced the CPAsyncStoreRewriter class with a new method to track vectorized lanes during statement rewriting.
- Updated CheckCPAsyncCopy to remove the buffer_oob parameter, simplifying its signature.
- Added logging for vectorized loop iterations to aid in debugging and performance analysis.
- Improved error handling in LowerCPAsyncCopy to provide clearer warnings when cp.async rewrites fail.
These changes aim to optimize the cp.async copy process and improve code maintainability.
* lint fix
* Update cp.async copy handling to improve rewrite success tracking
- Set the successfully_rewritten_ flag to true when a copy operation is successfully rewritten.
- Adjust the flag to false in cases where the load is not global-like or the buffer is not shared.
- Removed logging for vectorized loop iterations to streamline the code.
These changes enhance the tracking of rewrite success in the cp.async copy process, contributing to better error handling and maintainability.
* fix
* test: cover commit-merge + wait relax
* fix: skip wait relaxation under conditional async copy
* test fix
* Implement async copy support in TileLang
- Introduced a new configuration option `tl.enable_async_copy` to control the lowering of global to shared memory copies to PTX `cp.async`.
- Updated the `CopyNode::LowerCPAsyncCopy` method to handle async copy semantics based on the new configuration.
- Enhanced the `InjectPTXAsyncCopy` pass to rewrite eligible global->shared copies into `cp.async`, including necessary commit and wait statements.
- Added tests to verify the correct behavior of async copy lowering and the impact of the configuration option.
These changes improve the flexibility and performance of memory copy operations in TileLang.
* lint fix
* Refactor async copy handling in TileLang
- Renamed `InjectPTXAsyncCopy` to `LowerPTXAsyncCopy` to better reflect its functionality.
- Updated the `OptimizeForTarget` function to use `LowerPTXAsyncCopy` instead of the deprecated `InjectPTXAsyncCopy`.
- Added a deprecated alias for `InjectPTXAsyncCopy` to maintain backward compatibility.
- Adjusted documentation to clarify the purpose of the new function name.
These changes enhance clarity and maintainability in the async copy transformation process.
* fix
* lint fix
* test fix
* Enhance PTX async copy integration in TileLang
- Added the `InjectPTXAsyncCopy` function to facilitate the injection of PTX `cp.async` operations into the lowering process.
- Updated the `LowerTileOp` function to conditionally apply `InjectPTXAsyncCopy` during parallel loop lowering when async copy is enabled.
- Removed deprecated calls to `LowerPTXAsyncCopy` in the `OptimizeForTarget` function to streamline the optimization process.
These changes improve the handling of async copy operations, ensuring better performance and compatibility with CUDA targets.
* lint fix
* fix
* fix
* fix
* set parallisim to 1 to test the ci
* fix
* Refactor async copy handling in PTX lowering
- Improved the `PrepareCopyIndexInfo` function to clarify the calculation of transfer bytes and ensure valid async copy operations.
- Updated synchronization state management for injected `cp.async` transfers, enhancing clarity and consistency in handling open copy regions and uncommitted transfers.
- Added helper functions for zero value checks and linear offset flattening to streamline index processing.
These changes enhance the robustness and readability of the async copy implementation.
* Refactor nested pipeline functions and update test cases
- Removed unused parameters `order` and `stage` from several functions, simplifying their signatures.
- Updated nested pipeline implementations to use `num_stages=2` for better clarity and performance.
- Adjusted test cases to reflect the changes in function signatures, ensuring they run correctly without the removed parameters.
These modifications enhance code readability and maintainability while ensuring the tests remain functional.
* Refactor async copy statement generation in PTX lowering
- Updated the `MakeCPAsyncStmt` function to `MakeCPAsyncStmtFromLoads`, improving the handling of base indices for vectorized accesses.
- Introduced `ExtractVectorBaseIndices` to streamline the extraction of base indices from vectorized loads and stores.
- Enhanced error handling for predicated async copy operations when base indices cannot be extracted.
These changes improve the clarity and robustness of async copy statement generation in the PTX lowering process.
* lint fix
* Enhance async copy operation handling in CodeGenTileLangCuTeDSL
- Added support for conditional async copy operations in the `VisitExpr_` method, allowing for an optional predicate argument.
- Improved error checking to ensure the correct number of arguments for `tl::ptx_cp_async`.
- Streamlined the generation of async copy statements for both standard and conditional cases.
These updates improve the flexibility and robustness of async copy operations in the code generation process.ix
* lint fix
* disable auto cp async without pipeline.
* lint fix
* Enhance cp.async synchronization handling in TileLang
- Introduced `ThreadSyncAfterWaitGroupInserter` to ensure correct insertion of storage synchronization after `ptx_wait_group` statements, addressing potential visibility issues in shared memory.
- Updated `WarpSpecialized` to utilize `tl.ptx_cp_async_barrier_noinc` for releasing producer-consumer dependencies, preventing incomplete memory observations.
- Added tests to verify the correct insertion of synchronization and the handling of cp.async barriers under warp specialization.
These changes improve the correctness and reliability of async copy operations in the TileLang transformation pipeline.
* fix
* fix
* test fix
* Refactor sparse_mla_fwd_pipelined.py and enhance phase.py synchronization
- Removed unused post-processing function and related code from `sparse_mla_fwd_pipelined.py` to streamline the implementation.
- Added debug print statements in `phase.py` to log the state of the module before and after thread synchronization, aiding in troubleshooting and verification of synchronization behavior.
* Enhance async copy handling and loop synchronization in TileLang
- Introduced new loop annotations for async copy operations, including `parallel_prefer_async` and `parallel_async_without_async_commit_wait`, to improve control over async behavior in parallel loops.
- Updated the `InjectPTXAsyncCopy` function to accommodate the new annotations, ensuring correct async copy injection based on user-defined preferences.
- Enhanced the pipeline planning logic to prioritize cp.async groups by their earliest consumers, improving scheduling efficiency and reducing wait times.
- Added tests to validate the new async copy behaviors and ensure correct synchronization in pipelined loops.
These changes enhance the flexibility and performance of async copy operations in the TileLang transformation pipeline.
* lint fix
* Add Regression Perf for mhc
* clean
* clean
* lint fix
* lint fix
|
||
|
xyyy1420
|
26ac2fe45e |
[BugFix] Add vector type definitions to common.h for CPU codegen (#1901)
This is the follow-up fix mentioned in #1899 for the remaining float4 compilation issue. The CPU C codegen emits vector types (e.g. float4, int32_t4) when VectorizeLoop vectorizes T.copy operations, but common.h had no definitions for these types, causing g++ compilation failures. Add a generic vec_type<T,N> template and type aliases for all combinations that PrintType() can emit (float/double/half/intN_t, lanes 2/4/8/16). Update test_matmul_with_copy_cython to verify end-to-end correctness now that both the wrapper and vector type issues are resolved. |
||
|
xyyy1420
|
00267e43aa |
[BugFix] add target context and avoid redundant re-lowering in TLCPUSourceWrapper (#1899)
fix(cpu): skip redundant re-lowering in TLCPUSourceWrapper when mods already exist TLCPUSourceWrapper.parse_source_information() unconditionally called get_annotated_mod() to re-lower the IR, even when device_mod and host_mod were already passed in from tilelang.lower(). This caused a crash when the kernel uses T.copy (which generates T.Parallel loops), because VectorizePlanner requires Target::Current() during lowering. Fix: check if device_mod/host_mod already exist before re-lowering, matching the pattern used by TLCUDASourceWrapper. Also add self.target to the context manager for the fallback path. Added test_matmul_with_copy_cython to reproduce the issue. |
||
|
Chaofan Lin
|
c214f0ae0e |
[BugFix] Fix type mismatch when lowering to AtomicAddx2 template (#1898)
* [BugFix] Fix type mismatch when lowering to AtomicAddx2 template * remove debug output * add tests |
||
|
Lei WangandZhiwen Mo
|
df5ade9992 |
Refactor CUDA version checks for compute 9.0 (#1893)
* Refactor CUDA version checks in test files to use equality comparison for compute version 9.0 - Updated multiple test files to change `requires_cuda_compute_version_ge(9, 0)` to `requires_cuda_compute_version_eq(9, 0)`. - Removed deprecated test functions in `test_example_blocksparse_attention.py`. - Simplified configuration generation in `test_tilelang_autotune.py` by removing the `with_roller` parameter. - Ensured consistency across various examples and tests regarding CUDA version requirements. * revert unexpected changes --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
HaleandTong WU
|
eb2755750b |
[CUDA] Support tcgen5mma gemm ts (#1866)
* feat: add mma_ts init * feat: mma_ts valid * fix: tcgen05.st & .ld mismatch on 16-bit pack/unpack, cause function mistake in gemm * fix: diable unpack for bf16's tcgen05.st, enforce each tmem cell pack with 2 val, to be consistent with mma_ts's behaviour * test: add test example in testing/python/kernel/test_tilelang_kernel_bf16_gemm_tcgen5_ts.py * docs: add docstrings to gemm_base,gemm_tcgen05 and test files * style: apply clang-format and ruff format * refactor: merge ld/st meta, remove example, rename getTcgen05Meta to getTcgen05MetaLd_* * Restore TODO for tcgen05.cp support in copy.cc Added a TODO comment for future support of tcgen05.cp in conjunction with LowerTmemCopy. * fix: swap Ramp/PrimExpr isinstance order in access_ptr_from --------- Co-authored-by: Tong WU <109033598+Rachmanino@users.noreply.github.com> |
||
|
Tong WU
|
2e854b068e |
[Feature] Support cluster launch, query, synchronization and barrier operations (#1874)
* support cluster launch for tvm_ffi and cython backend * lint * add test * lint * update tvm * Add cluster synchronization operations and related functions - Introduced new built-in operations for cluster synchronization: `cluster_arrive_relaxed`, `cluster_arrive`, `cluster_wait`, `cluster_sync`, and `block_rank_in_cluster`. - Updated the CUDA code generator to handle these new operations. - Added corresponding Python bindings and documentation for the new cluster functions. - Included necessary header files for cluster operations in the CUDA code generation process. * add test for intrinsics * lint * remove legacy code * Add cluster barrier support and related functionality - Introduced `ptx_arrive_cluster_barrier` built-in for arriving at cluster barriers. - Added `alloc_cluster_barrier` function for allocating cluster barrier buffers. - Updated CUDA code generation to handle the new cluster barrier operations. - Enhanced existing functions to support shared and cluster barrier scopes. - Added tests for cluster barrier functionality in the TileLang testing suite. * lint * fix typo * add storage sync for cluster |
||
|
Chaofan Lin
|
7fcd62b0ba |
[Enhancement] Clarify the semantic rule of copy operator and add shape mismatched tests (#1883)
* [Enhancement] Clarify the semantic rule of copy operator and add shape mismatched tests * remove debug msg * fix after rebase * fix * fix |
||
|
Taoyu Zhu
|
0fa78e235f |
[AMD] Fix gfx950 ci and add 16x16x32_bf16/fp16 instructions support (#1878)
* add * add * add * zz * zz * tmp not right * right * right * right * right * right * clear * pre-commit lint * pre-commit lint * pre-commit lint * pre-commit lint * pre-commit lint * fix commit * rerun ci? * add tvm ffi target support for cdna4 * add tvm ffi target support for cdna4 * add tvm ffi target support for cdna4 * add tvm ffi target support for cdna4 * tmp 1 process to run rocm cis * restore * fix |
||
|
Lei Wang
|
808b0cefcc |
[Feature] Introduce higher-dimensional gemm layout support (#1798)
* [Refactor] Update swizzle layout functions to accept Buffer objects directly * Refactor swizzle layout functions to take Buffer parameters instead of stride and continuous values, improving usability and consistency. * Introduce helper functions to extract stride and continuous dimensions from Buffer objects. * Update related functions and documentation to reflect the new parameter types. * Enhance layout inference logic to merge swizzle layouts based on Buffer objects. * Add new reshape functionality to the Layout class for better shape management. * lint fix * use layout.expand instead of reshape to lift a swizzled layout * lint fix |
||
|
Henry
|
911385fdd1 |
[BugFix] Fix Hopper TMA lowering without warp specialization (#1840)
* [BugFix] Fix Hopper TMA lowering without warp specialization Decouple the Hopper pipeline entry condition from warp specialization so TMA lowering works correctly when `tl.disable_warp_specialized=True`. Also fix barrier id promotion for multi-dim TMA loads in inject_tma_barrier.cc. Fixes #737, #1648 * Unify is_1d_tma_load logic and improve test diagnostics Address CodeRabbit review feedback: - Extract shared Is1DTmaLoad() helper to ensure consistent barrier argument index selection across TmaExpectTxRewriter and TmaBarrierRewriter - Add failure message to regex assertion in regression test * Fix ruff-format: merge f-string into single line |
||
|
Lei Wang
|
9bd62f5b5f |
[Layout] Add Layout.expand to lift a layout into higher dimensions and improve repeat errors (#1876)
* layout: add Layout.expand and improve repeat errors * revert unexpected changes |
||
|
Lei Wang
|
a22fa61dfe |
layout: add Layout.repeat for tiling atom layouts (#1875)
layout: add Layout.repeat via FFI |
||
|
Lei Wang
|
7ab2b2be92 |
[BugFix] ArgBinder: relax shared-shape binding for unused nullable buffers (#1870)
* [BugFix] tvm_ffi: handle unused nullable buffers with shared symbolic shapes * [BugFix] ArgBinder: relax shared-shape binding for unused nullable buffers |
||
|
Lei Wang
|
9f259544d4 |
[Refactor] Refactor Pass InjectFenceProxy (#1863)
* tir: add T.cdiv alias for T.ceildiv * docs: Update InjectFenceProxy documentation and enhance code comments - Clarified the description of `tl.InjectFenceProxy` to specify the transition between generic and async proxy operations. - Improved explanations of the pass's functionality, including state tracking and the handling of TMA store synchronization. - Added details on the new `ProxyStateSet` class and its role in managing proxy states. - Updated usage instructions for proxy hints to include new options for custom operations. - Enhanced test coverage for handling unknown external calls and proxy hint overrides. * fix * enhance * fix * remove tl.proxy_hint * refactor tma_store_arrive and tma_store_wait * refactor * fix * InjectFenceProxy: hoist fence out of pure-async loops * InjectFenceProxy: hoist fence for if/while pure-async regions * Remove LowerTileOp transform test |
||
|
Gabriel Wuandyuxic
|
391cf5b6b9 |
[Enhancement] GEMM V2 on SM90/SM100 CuTeDSL backend (#1855)
* fix(cutedsl): add rsqrt/rsqrtf support in codegen
- Add rsqrt/rsqrtf to fastmath function mapping
- Implement ieee_frsqrt using tl.rsqrt
- Export rsqrt from contrib/cutedsl/math.py
Fixes: norm/test_rms_norm.py
* fix(cutedsl): add atomic_add support and fix shared.dyn access
- Add atomic_add_elem_op and atomic_add_ret_elem_op handling
- Exclude shared.dyn from direct indexing to avoid out-of-bounds access
(shared.dyn shape is set to (1,) but accessed with arbitrary indices)
Fixes: gemv/test_example_gemv.py
* fix(cutedsl): add builtin::ptx_ldmatrix support
Implement builtin::ptx_ldmatrix by reusing tl.ptx_ldmatrix_x* functions.
* feat(cutedsl): add ptx_mma support with all MMA types
Add complete PTX MMA intrinsic support for CuTeDSL backend:
- Implement ptx_mma codegen in codegen_cutedsl.cc
- Add GetVarPtr_ helper for local buffer pointer extraction
- Create ptx_mma.py with inline PTX assembly wrappers for:
- FP16: m16n8k16 -> f16/f32
- BF16: m16n8k16 -> f32
- INT8/UINT8: m16n8k32 -> i32
- INT4/UINT4: m16n8k32 -> i32
- TF32: m16n8k4, m16n8k8 -> f32
- FP64: m8n8k4 -> f64
- FP8 e4m3/e5m2: m16n8k32 -> f16/f32
- Fix BroadcastNode to emit typed values for make_filled_tensor
- Update ptx_ldmatrix to use GetVarPtr_ for local buffers
* feat(cutedsl): comprehensive atomic operations support
- Add atomic.py module with full atomic operations:
- AtomicAdd: float32 (nvvm.atomicrmw), float16 (inline PTX with i16 bitcast),
int32/int64 (nvvm.atomicrmw)
- AtomicAddx2/x4: vectorized float32 via PTX atom.global.add.v2/v4.f32
- AtomicMax/Min: int32/int64 (nvvm.atomicrmw), float32 (CAS loop via inline PTX)
- AtomicLoad/Store: all memory orderings via inline PTX (relaxed/acquire/release/seq_cst)
- Add codegen support for atomic_load/store/max/min/addx2/addx4 in codegen_cutedsl.cc
- Move AtomicAdd from __init__.py to atomic.py, extend with float16/int64 support
* feat(cutedsl): add fabsf/fabs support in codegen and math module
* feat(cutedsl): add wait_wgmma support via nvvm.wgmma_wait_group_sync_aligned
* feat(cutedsl): support i32 min/max
* fix(cutedsl): handle FP8 narrow precision alignment in codegen
CuTeDSL requires all FP8 vector operations to be 32-bit aligned, but
TIR generates float8x2 (16-bit) vectors. Additionally, MLIR vector
extractelement on FP8 types fails with unrealized_conversion_cast.
Fix CastNode, BufferLoadNode, and BufferStoreNode to:
- Pad unaligned FP8 casts to aligned width (e.g. f32x4 -> fp8x4)
- Return rmem tensor names (not .load()) so downstream uses rmem
element access instead of MLIR vector extractelement
- Use padded aligned stores for local rmem, scalar element stores
for global/shared memory (to handle non-aligned offsets)
* fix(cutedsl): use pointer arithmetic for shared.barrier buffer access
_Pointer type from tl.alloc_smem() does not support subscript [i],
but supports pointer arithmetic (ptr + i). This fixes warp_specialize
tests that use mbarrier.
- GetBufferRef_: use (vid + index) for shared.barrier instead of vid[index]
- BufferLoadNode: skip .load() suffix for shared.barrier (returns pointer, not value)
Fixes 3/4 warp_specialize tests.
* fix(cutedsl): handle single coordinate in tma_load
When tma_load receives a single Int32 coordinate (for 1D tensors) instead
of a tuple, wrap it in a tuple before iterating. This fixes the error:
'Int32' object is not iterable
Fixes flash_decoding/test_example_flash_decoding.py
* feat(cutedsl): add decode_i4u/s_to_f16 and warp primitives
- Add quantize.py with decode_i4u_to_f16 and decode_i4s_to_f16
using inline PTX (lop3.b32 + sub.f16x2)
- Add warp.py with __activemask, __shfl_down_sync, __shfl_sync
using CuTeDSL's shuffle_sync functions
- Update codegen to map decode_* and __* functions to tl.* namespace
Fixes dequant_gemv_fp16xint4 test.
* feat(cutedsl): add decode_fp4_to_bf16_twiddling
Implement FP4 to BF16 conversion using bit twiddling technique.
- Use prmt.b32 for endianness handling
- Use mul.bf16x2 for scaling
- Pack output with correct high/low BF16 layout
Fixes test_example_dequant_gemm_bf16_mxfp4_hopper.
* fix(cutedsl): support reinterpret for arbitrary expressions and fix shift type promotion
- reinterpret: handle non-BufferLoad args with tl.bitcast()
- shift_right/shift_left: wrap result in dtype() to preserve type
- CuTeDSL Int8>>4 returns Int32 (Python promotion), CUDA keeps int8_t
- Add bitcast() function using llvm.bitcast MLIR op
* feat: disable v2->v1 fallback
* fix(cutedsl): TensorSSA for store (Select/scalar) and TF32 MMA dispatch
- ArithValue: emit tl.where() for Select and wrap scalar store RHS in
make_filled_tensor().load() so CuTeDSL Tensor.store() gets TensorSSA.
- TF32: accept fp32/f32/float32 for m16n8k4 and m16n8k8 in ptx_mma dispatch
(e.g. deepseek_mhc).
* feat(cutedsl): implement CumSum1D and CumSum2D for cumulative sum operations
- Add CumSum1D and CumSum2D classes in reduce.py based on CUDA templates
- Add __shfl_up_sync wrapper in warp.py using CuTeDSL's shuffle_sync_up
- Implement warp-level prefix sum (forward/reverse) using shuffle operations
- Use @cute.jit decorator and nested if for DSL-compatible control flow
- Fixes gdn/test_example_cumsum_compilation test (7/8 gdn tests now pass)
* feat(cutedsl): implement tma_reduce for TMA atomic add reduction
* fix(cutedsl): handle if_then_else as tl.where and unify Select/where arm dtypes
- Override builtin::if_then_else in CallNode to emit tl.where() instead of
Python ternary (which yields ArithValue, not TensorSSA)
- Use make_rmem_tensor + element assignment instead of make_filled_tensor
(which creates unsupported vector<1xT> types for int)
- Unify arm dtypes in both SelectNode and if_then_else: when true/false arms
differ (e.g. BFloat16 vs Float32), cast both to common type via .to()
- Fix BufferStore: skip make_filled_tensor wrapping for Select/if_then_else
values (already TensorSSA); extract [0] for scalar element assignment
- Fix GetBufferRef_: extract Ramp base for make_tensor_at_offset offset
parameter to avoid emitting tuple offsets
Fixes: cast/group_per_split_token, topk, deepseek_mhc, flash varlen
* fix(cutedsl): convert non-integer offsets to Int64 in make_tensor_at_offset
Complex arithmetic on cutlass.Int64 values (involving tl.min, //, %) can
produce Float64 ArithValue that CuTeDSL Pointer.__add__ does not accept.
Detect non-integer offsets via is_integer() and wrap with cutlass.Int64().
Fixes blocksparse_attention varlen_indice and varlen_mask (6/6 passed).
* fix(cutedsl): fix narrow unaligned load/store for sub-32-bit types
Three issues in the narrow unaligned (e.g. uint8x2, fp8x2) code paths:
1. Load from shared/global: using aligned_lanes (4) elements caused OOB
near buffer boundaries and MisalignedAddress for non-aligned offsets.
Fix: load exactly value_lanes (2) elements with no cute.assume.
2. Load from rmem: cute.assume(offset, divby=aligned_lanes) silently
truncates offsets not divisible by aligned_lanes (e.g. offset=6 with
divby=4 → 4). Fix: drop cute.assume for rmem (no hardware constraint).
3. Store to rmem: padded (aligned_lanes,) .store() writes beyond
value_lanes elements, causing overlapping writes and OOB at tensor
boundaries. Also triggers CuTeDSL's 32-bit alignment check for FP8.
Fix: use element-by-element assignment (vid[i]=val).
Fixes dequantize_gemm w4a8 (cudaErrorMisalignedAddress → correct results).
* fix(cutedsl): handle signless int types in tl.min/tl.max fallback
* fix(cutedsl): handle _Pointer src_values in AtomicAddx2/x4
* fix(cutedsl): remove double PrintIndent in cp_async_gs_conditional codegen
* fix(example): support SM100 in attention_sink bwd configs
Change get_bwd_configs() to use >= 90 instead of == 90 so that SM100
(Blackwell) uses the same config as SM90 (Hopper): block_M=128,
block_N=32, num_stages=2, threads=256.
* feat: implement WGMMA V2 (ss + rs) for CuTeDSL backend
Implement explicit WGMMA descriptor management ops in the CuTeDSL backend,
replacing LOG(FATAL) stubs with working codegen for all 8 WGMMA-related ops.
New Python files:
- typing.py: WGMMA type map (string dtype -> nvvm.WGMMATypes)
- gemm_V2.py: GmmaDescriptor class, initialize/increase descriptor,
warpgroup sync ops, wgmma_ss (via nvvm.wgmma_mma_async),
wgmma_rs (via PTX inline asm)
Key findings and fixes:
- nvvm.mma_smem_desc does not correctly pack the layout_type/swizzle
field in the WGMMA descriptor; replaced with manual bit packing
- wgmma_rs implemented using llvm.inline_asm with dynamically generated
PTX and constraint strings based on M/N/K Constexpr parameters
C++ codegen changes (codegen_cutedsl.cc):
- warpgroup_arrive/commit_batch/wait/fence_operand
- ptx_wgmma_ss (15 args) and ptx_wgmma_rs (14 args)
- initialize_wgmma_descriptor (arg reorder: TIR->Python)
- increase_descriptor_offset
All GEMM tests (4/4) and flash attention tests (14/14) pass.
* fix(cutedsl): replace break with guard variable pattern for CuTeDSL compatibility
* feat(cutedsl): implement tcgen05 (SM100/Blackwell) MMA support
Add CuTeDSL codegen and Python runtime support for tcgen05 MMA operations:
- New gemm_tcgen05.py: Tcgen05SmemDescriptor, descriptor initialization,
MMA variants (ss/ws_ss/ts), mbarrier arrive, TMEM alloc/dealloc/load
- All MMA and commit instructions guarded by elect_one_sync via inline PTX
(@q predicate on individual instructions, not on block scopes)
- TMEM load uses recursive power-of-2 splitting with direct MLIR emission
- C++ codegen: replace 6 LOG(FATAL) stubs with proper code generation for
tcgen05 MMA/descriptor/TMEM ops, handle descriptor scope allocation
- Fix shared.barrier pointer access (skip .iterator for alloc_smem results)
- Fix mbarrier_wait: use blocking PTX loop instead of single try_wait call
* fix(cutedsl): correct tcgen05_ld 64b/128b/256b PTX variants and cleanup
- Fix tcgen05_ld_32dp64bNx/128bNx/256bNx: use correct PTX instructions
(16x64b/16x128b/16x256b) with two-half decomposition instead of
hardcoded 32x32b
- Generalize _emit_tmem_ld_segment to accept ptx_type and regs_per_x
- Remove dead expression in ptx_mma.py (c_ptr + a_offset)
- Remove unused GemmCuTeDSL and is_cutedsl_target imports
* fix(cutedsl): review cleanup — warp shuffle, mbar timeout, atomic addr space, dead code
- Fix warp.py: correct mask_and_clamp formulas to match CUDA semantics
for __shfl_down_sync, __shfl_up_sync, __shfl_sync (width param was ignored)
- Fix mbar.py: use timeout_ns parameter instead of hardcoded 0x989680
- Fix atomic.py: use generic address space (ld.f32, atom.cas.b32) instead of
.global for AtomicMax/AtomicMin CAS loops, supporting shared memory pointers
- Add ICHECK for non-unit Ramp stride in codegen_cutedsl.cc GetBufferRef_
- Remove unreachable duplicate elif gemm_inst.is_tcgen5mma() branch
- Remove unused _to_ir_value helper from ptx_mma.py
* fix(cutedsl): skip BulkLoad1D for CuTeDSL + fix uint broadcast type mismatch
Two fixes for dequant/groupedgemm_bf16_mxfp4_hopper on CuTeDSL:
1. Skip BulkLoad1D/BulkStore1D for CuTeDSL target in copy.cc:
cp_async_bulk_shared_cluster_global (raw 1D TMA) combined with WGMMA
in the same kernel triggers a ptxas ICE in the NVPTX backend.
Falls through to descriptor-based BulkLoad/BulkStore instead.
2. Use signed int type for BroadcastNode fill values in codegen_cutedsl.cc:
CuTeDSL/MLIR normalizes unsigned integer tensor loads to signed types
(Uint8 pointer -> i8 elements). BroadcastNode fill values now use
matching signed types to avoid Int8/Uint8 mismatch in tl.where().
* fix(cutedsl): support fp16/bf16 in AtomicAddx2 via atom.add.noftz.v2.f16/bf16
* ci(cutedsl): add CuTeDSL examples to CI pipeline
Add a new GitHub Actions step that runs all examples with
TILELANG_TARGET=cutedsl. Known failures (sparse ops) are marked as
xfail and unsupported tests (convolution, stream-K, flash_decoding FA3)
are auto-skipped via pytest_collection_modifyitems in conftest.py.
The step uses continue-on-error: true initially until stabilized.
* chore: fix pre-commit
* fix(adapter): use name-based fallback for tir.Var dynamic symbolic lookup
tir.Var uses object identity for hash/eq, so when param_shapes (from
the original PrimFunc) and dynamic_symbolic_map (from the transformed
PrimFunc in ir_module) contain different tir.Var objects with the same
name, direct dict lookup fails with KeyError.
Add a secondary name-based index (_dynamic_symbolic_name_map) and a
_lookup_dynamic_symbolic() helper that falls back to name matching.
This fixes KeyError: batch under pytest-xdist parallel execution.
Affected adapters: cutedsl, nvrtc. (cython and tvm_ffi already had
a str()-based workaround.)
* chore: remove dead GemmCuTeDSL class and update stale V1 comments
- Delete `gemm_cutedsl.py` (no references anywhere in codebase)
- Remove outdated comment in `gemm_op.py`
- Update `run_local_ci_test.sh` comment: CuTeDSL no longer requires V1
- Keep `TILELANG_USE_GEMM_V1` env var as a debugging escape hatch
* refactor(cutedsl): reorganize contrib/cutedsl module structure
- Rename gemm_V1.py/gemm_V2.py to snake_case (gemm_v1.py/gemm_v2.py)
- Merge typing.py (19 lines) and mbar.py (53 lines) into utils.py and cpasync.py
- Extract utility functions from __init__.py into new utils.py module
- Slim __init__.py to pure re-exports
- Add __all__ to all submodules for explicit public API documentation
* fix(ci): use .empty() instead of .size() >= 1 in pipeline_planning.cc
* fix: coderabbit comments
* fix(cutedsl): correct WGMMA result register count for f16/bf16 output
wgmma_ss/wgmma_rs used M*N//128 for result registers, which is only
correct for 32-bit output types. For f16/bf16, each i32 register holds
2 packed elements, halving the count. Derive from cutlass type .width
instead of hardcoding.
Also update barrier test assertion to match codegen's pointer-arithmetic
syntax: (barriers + i) instead of barriers[i].
* refactor(cutedsl): templatize ptx_mma.py with factory function
Replace 13 near-identical @dsl_user_op PTX MMA functions (~840 lines)
with a single _make_ptx_mma factory and a config table (~120 lines).
Also correctly interpolates a_layout/b_layout into the PTX mnemonic
instead of hardcoding row.col.
* fix(cutedsl): fix NaN CAS loop, wgmma_rs int accum, and stale pipeline state
- atomic.py: Use integer-domain comparison (setp.ne.b32) in AtomicMax/Min
CAS loops to prevent infinite spin when values are NaN
- gemm_v2.py: Use "=r"/i32 constraints for wgmma_rs when C_dtype is s32
(int8 GEMM) instead of hardcoded "=f"/f32
- reduce.py: Add assert H <= 32 guard for CumSum2D dim=0 column-wise path
- pipeline_planning.cc: Clear pending tcgen05 state unconditionally after
arrive to prevent stale entries from leaking to future arrives
* fix(cutedsl): fix ptx_mma brace escaping and wgmma_rs DSL compatibility
- ptx_mma.py: Double-brace register lists in PTX template so
.format(a_layout=, b_layout=) doesn't interpret {$0, ...} as
format fields (KeyError regression from templatization)
- gemm_v2.py: Replace if/else branches inside DSL-traced loops with
ternary expressions to avoid CuTeDSL "operands type change inside
dynamic if" error
- ci.yml: Remove continue-on-error from cutedsl CI, reduce maxfail
to 3, increase parallelism to 4
* feat(cutedsl): implement ptx_mma_sp (sparse MMA) codegen and runtime
* fix(cutedsl): guard saturate param, move CuTeDSL CI step after core tests
- ptx_mma/ptx_mma_sp: raise NotImplementedError if saturate=True
- warpgroup_fence_operand: add comment explaining no-op rationale
- ci.yml: move CuTeDSL examples after CUDA/ROCm/Metal test steps
so a CuTeDSL failure no longer blocks core tests
* feat(cutedsl): implement 16 previously unsupported ops
Add support for thread index queries (get_lane_idx, get_warp_idx,
get_warp_idx_sync, get_warp_group_idx), exp10 math function,
warp-level reductions (sum, max, min, bitand, bitor), and
IEEE-754 compliant arithmetic with explicit rounding modes
(add, sub, mul, fma, rcp, sqrt, div).
Reduces unsupported op count from 26 to 10. The remaining ops
are either high-complexity (tma_load_im2col, tl_gemm_sp),
require wrapper changes (sync_grid), or are legacy/unused.
* fix(cutedsl): wrap non-int32 integer literals with CuTeDSL type constructors
CuTeDSL's as_numeric() converts bare Python ints to Int32, causing
"Type mismatch" errors when storing to narrower tensors (e.g. Int16).
Override IntImmNode to emit cutlass.Int16(0) instead of bare 0 for
non-int32 literals. Also add a BufferStore narrowing guard as safety net.
Fixes compress_kernel Int32→Int16 error in gemm_sp/custom_compress.
* feat(cutedsl): implement sync_grid with cooperative kernel launch
Add grid-level synchronization via a software barrier using
llvm.mlir.global counter + inline PTX atomic spin-wait. The CuTeDSL
wrapper detects sync_grid usage and automatically switches to
cuLaunchCooperativeKernel to guarantee all blocks are resident.
- grid_sync.py: module-global counter + PTX barrier protocol
- wrapper.py: CPP_COOPERATIVE_KERNEL_LAUNCH_TEMPLATE
- codegen_cutedsl.cc: emit tl.sync_grid() instead of LOG(FATAL)
---------
Co-authored-by: yuxic <yuxic@nvidia.com>
|
||
|
Yichen Yan
|
fbbda0f22b | [Codegen] Metal codegen on Linux (#1857) | ||
|
Tong WUandLeiWang1999
|
ffa0cdd1e1 |
[BugFix] Fix eager mode where there is no tensor args (#1851)
* [BugFix] Fix eager mode where there is no tensor args * fix --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Lei Wang
|
4956b5835f | ThreadSync: avoid barriers between atomic ops (#1852) | ||
|
|
48a8f4a6ce |
[Feature] Hierarchical reduction and warp reduction intrinsics support (#1762)
* [Feature] Add new hierarchical reduce with less workspace size and new ptx * Lint * Add intrinsic functions and more types * Fix sm100a redux instruction * Fix data type * [Test] Add tests for reduce operations * [Refactor] Unify AllReduce barrier mechanism via policy template Replace the duplicated run()/run_hopper() methods in AllReduce with a single run() that accepts a Barrier policy template parameter (SyncThreadsBarrier or NamedBarrier<N>). Extract the shared inter-warp reduction logic into a warp_inter_reduce helper, and split the dispatch into private hierarchical_reduce/butterfly_reduce methods. Update codegen in reduce.cc and finalize_reducer.cc to emit NamedBarrier<all_threads> for SM >= 90 targets instead of the old all_threads + run_hopper pattern. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Add handling for non-float types in warp_reduce function This update introduces a new case in the warp_reduce function to cast non-float types to float using static_cast. This enhancement improves type flexibility and ensures compatibility with a broader range of input types during reduction operations. * [Refactor] Remove hierarchical reduce because of performance regression * [BugFix] Buffer for reduction --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
||
|
Lei WangandZhiwen Mo
|
6d10d1602c |
[CUDA][Feature] Add packed FP32x2 math intrinsics and auto vectorized support (#1839)
* [Feature] Add packed FP32x2 math intrinsics and support for CUDA/HIP This commit introduces new packed FP32x2 math operations: fadd2, fmul2, and fma2, which leverage PTX instructions on supported architectures. The changes include: - Definitions and implementations of fadd2, fmul2, and fma2 in both CUDA and HIP codegen files. - New Python API functions for these operations, ensuring they validate input types and handle fallbacks for unsupported architectures. - Documentation updates to reflect the new intrinsics in the math language module. These enhancements improve performance for vectorized floating-point operations in TileLang. * revert hip changes * annotate cuda only --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei Wang
|
4ff81c7d40 |
[Refactor] Introduce T.access_of to combine T.address_of and access_ptr (#1827)
* [Refactor] Update atomic operations to use access_ptr for buffer access * Replaced instances of address_of with access_ptr in atomic operations to improve clarity and maintainability. * Enhanced atomic operation implementations in various files, including atomic_add, atomic_max, and atomic_min. * Updated related logic in loop vectorization and tile language to support new access patterns. * Added a utility function to extract destination scalar dtype for atomic operations, supporting both BufferLoad and tvm_access_ptr. * Improved kernel cache logic to include a build-stamp for the TileLang runtime library, ensuring cache validity during development. * lint fix * Fix access_ptr layout rewrite and atomic vectorization * lint fix * Refactor atomic operations to remove const qualifiers from pointer parameters in atomicAddx2 and ToFloat functions. This change enhances type flexibility and consistency across atomic operations. * Refactor buffer access logic in lower_tile_op and vectorize_loop to enhance clarity and maintainability. Removed unused code and improved offset calculations for tile access. Updated atomic operation handling to streamline dtype extraction and ensure correct vectorization behavior. * Add logging for loop layout and lowered loop in atomic_add; implement tvm_access_ptr handling in loop_vectorize and lower_tile_op This commit introduces logging statements to track the loop layout and the result of the lowered loop in the AtomicAddNode. Additionally, it implements a new method to handle tvm_access_ptr calls in the loop_vectorize and lower_tile_op files, enhancing the handling of memory access patterns and ensuring correct vectorization behavior. * fix * lint fix * fix * Update ToFloat2 and ToFloat4 functions to use const qualifiers for pointer parameters in atomic.h; modify test script to call test_atomic_addx2_float instead of main. * Refactor tvm_access_ptr handling in loop_vectorize to improve argument validation and buffer resolution logic. Ensure at least three arguments are provided and streamline the search for matching buffers in layout_map_. Update vector size calculations to enhance clarity and maintainability. * Add handling for tvm_access_ptr in vectorize_loop to support offset substitution and buffer creation This commit introduces logic to handle tvm_access_ptr calls, allowing for the substitution of loop variables in offsets and the creation of dummy buffers. The new method ensures correct base address calculation and enhances the overall handling of memory access patterns in vectorization. * Add handling for tvm_access_ptr offset validation in loop_vectorize This commit enhances the loop_vectorize functionality by adding checks for the element offset in tvm_access_ptr calls. It ensures that offsets are invariant within the vector boundary, improving the robustness of vectorization logic and memory access patterns. * fix * Enhance tvm_access_ptr handling in loop_vectorize by adding offset validation checks. This update improves the robustness of vectorization logic and ensures offsets remain invariant within vector boundaries. * Enhance tvm_access_ptr handling in loop_vectorize by adding offset validation checks. This update ensures offsets remain invariant within vector boundaries, improving memory access robustness and vectorization logic. * fix * Add tl.access_ptr support in TileLang This commit introduces the `tl.access_ptr` operation, which carries pointer access metadata in the frontend. It allows for better handling of buffer loads by retaining necessary information for downstream analysis. The new operation is integrated into various components, including the loop vectorization and memory access legalizations, ensuring compatibility with existing functionality while enhancing the robustness of memory access patterns. * lint fix * fix * Add support for tl.access_ptr in atomic operations This commit updates the AtomicAddNode and AtomicOpBaseNode implementations to utilize the new tl.access_ptr operation for building access pointers to destination elements. It enhances the handling of buffer loads by incorporating type checks for index data types, improving the robustness of memory access patterns. Additionally, it removes unnecessary logging statements in the lowering process, streamlining the codebase. * lint fix * Refactor buffer load extraction in vectorize_loop to check for non-empty arguments in access_ptr calls. This change improves the robustness of buffer load handling by ensuring valid argument presence before processing. * Remove unused fallback logic in KernelCache and clean up import statements in test_tilelang_language_access_ptr.py for improved code clarity. * Refactor memory access in BitNet and dequantization examples to use tl.access_ptr This update replaces T.address_of calls with T.access_ptr in multiple files, enhancing memory access handling by specifying read and write permissions. This change improves code clarity and aligns with recent enhancements in pointer access management. |
||
|
Lei Wang
|
e666d2d3cc |
[Refactor] Treat local.var as local buffers when deciding vectorization for stable actions (#1835)
lint fix |
||
|
Lei WangandCopilot
|
f8b4d9e677 |
[Bugfix] Fix threadIdx variable lookup by thread_tag instead of position in ThreadSync (#1829)
* Refactor OptimizeForTarget function by removing debug print statements and updating module state visualization * Add test for sync hoisting in non-uniform if within loop using shared memory This commit introduces a new test case to verify the correct behavior of sync hoisting when a non-uniform if statement is present inside a loop that utilizes shared memory. The test ensures that the synchronization occurs before the if statement, confirming the expected transformation in the module's intermediate representation. * fix * Refactor thread variable handling in thread_storage_sync.cc This commit removes the optional thread variable lookup by tag and replaces it with direct indexing based on the expected position of threadIdx variables. This change addresses a bug where the last three elements of env_threads_ were incorrectly assumed to always correspond to threadIdx.x/y/z, improving the accuracy of thread variable access in the synchronization logic. * Update testing/python/transform/test_tilelang_transform_thread_sync.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> |
||
|
Lei Wang
|
c99879ad12 |
[Bugfix] Fix ast builder error for value -= 1 (#1825)
* Enhance augmented assignment methods in BaseBuilder and Builder classes - Updated `aug_assign` method to include an optional `name` parameter for better handling of variable names during augmented assignments. - Improved handling of immutable variables in the `Builder` class to treat augmented assignments as re-bindings, providing warnings for potential issues. - Added comments to clarify the changes and ensure user expectations are met regarding variable assignments. * lint fix |
||
|
Lei Wang
|
6722690697 |
[BugFix] Fix fp16 annotate_l2_hit_ratio host stub compilation (issue #1810) (#1818)
* Fix L2 persistent lowering for fp16 host stubs Avoid typed access_ptr casts in stream access policy window setup and add issue #1810 regression test. * test and lint fix * Add regression test for issue #1810 Add a regression test for issue #1810 related to CUDA stream access policy. |
||
|
Lei Wang
|
e9d056954c |
[BugFix] LoopUnswitching: gate non-trivial else behind PassConfig (#1816)
* [BugFix] LoopUnswitching: gate non-trivial else behind PassConfig * lint fix |
||
|
Kuris
|
c47d6df379 |
[BugFix] Missing Recursive Loop Var Checking in Loop Unswitching (#1801)
* [Fix] Update loop unswitching logic to handle multiple let bindings and add corresponding test case * remove debug print |
||
|
Jiawei Lin
|
af30ac2152 |
[Bugfix] Support loop-dependent conditions in IfThenElse within T.Pipelined (#1799)
* [BugFix] Fix loop-dependent conditions in IfThenElse within T.Pipelined This commit applies the same strategy used for LetStmt to IfThenElse conditions: 1.Introduced IfWrapper struct to track if conditions that depend on the loop variable 2.Added dependency detection that checks if an if condition uses: - The pipeline loop variable directly, OR - Any variable transitively dependent on the loop variable 3.Loop-dependent conditions are pushed inside each pipeline stage with the loop variable properly substituted for that iteration * Add test for loop-dependent condtions within T.Pipelined * Fix code format |
||
|
Chaofan Lin
|
4349b2c50b |
[Feature] Support passing PrimExpr value in tile-level atomic operation (#1796)
* [Feature] Support passing PrimExpr value in tile-level atomic operation * fix after rebase * address comments * fix tvm ver * fix |
||
|
Lei Wang
|
b01cb93586 |
[Refactor] Reorganize ParallelOp code structure and move ProveFragmentContains to layout utils (#1779)
* [Feature] Implement ProveFragmentContains Function for Fragment Thread Validation - Added the ProveFragmentContains function to check if the threads accessing elements of a smaller fragment are a subset of those accessing a larger fragment. - This function ensures valid access when transitioning from a smaller to a larger fragment layout. - Updated layout.cc and utils.cc to incorporate this new functionality, enhancing the layout validation process. - Removed the previous implementation of ProveFragmentContains from parallel.cc to streamline the codebase. * fix * Refactor ParallelOpNode Layout Handling - Removed the initial DeReplicate attempt from InferLayout to streamline layout inference. - Added DeReplicate logic to ComputeLoopLayoutFromBuffer to reduce replication when validating layout candidates. - Updated test cases to disable caching and ensure proper functionality of loop layout kernels. * fix * Refactor Test Cases for Loop Layout - Removed caching disablement and print statements from the loop layout identity test for cleaner output. - Updated the main execution block to directly call the testing framework, enhancing test execution flow. |
||
|
Lei Wang
|
191d879ddd |
[Bugfix] Fix thread storage sync conflict detection for loop carry write-after-read (#1781)
* Fix thread storage synchronization logic in `thread_storage_sync.cc` to correctly identify conflicts between read and write operations based on loop carry conditions. * lint fix * Refactor `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` to use shared memory for `sorted_token_ids` instead of local memory, improving thread synchronization. Adjust default argument values for M, N, and K in the main function for better testing scenarios. * Add UniformExprChecker to enforce thread synchronization rules Introduce the UniformExprChecker class to determine if expressions are uniform across threads, crucial for safe synchronization in conditional statements. Update the TileLangThreadSyncPlanner to hoist synchronization points out of non-uniform if-statements to prevent potential deadlocks. Enhance tests to validate sync hoisting behavior for various non-uniform conditions involving thread indices and shared memory access. * lint fix * Enhance `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` with cache disabling and kernel source printing for debugging. Update thread synchronization logic in `thread_storage_sync.cc` to check for runtime-dependent conditions, preventing potential deadlocks by hoisting sync points as necessary. * Update submodule `tvm` to latest commit and remove deprecated `example_gqa_decode_varlen_logits_paged.py` file. Refactor `example_gqa_decode_varlen_logits.py` to enhance performance and maintainability by removing unused imports and optimizing shared memory usage. Adjust test cases to reflect the removal of the paged example. * fix * Enhance thread synchronization logic in `thread_storage_sync.cc` by adding a configurable warp size parameter to `RuntimeDependentConditionChecker` and `TileLangThreadSyncPlanner`. This allows for better adaptability to different target architectures. Update the logic to ensure thread extent is a constant and improve handling of runtime-dependent conditions. * lint fix * Refactor thread extent validation in `thread_storage_sync.cc` to use pointer checks instead of optional values. This change improves clarity and ensures that the thread extent is correctly validated as a constant. * Adjust loop variable constraints in `thread_storage_sync.cc` for loop-carry analysis by modifying the extent calculation. This change ensures valid iteration comparisons by reducing the extent by one, allowing for accurate analysis of loop iterations. * lint fix |