* 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
* 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.
* 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>
* [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 4995504098.
* Enhance producer-consumer logic and add tests for TMA behavior
- Updated the logic in producer_consumer_ws.cc to avoid reserving unused pre-loop barriers, improving efficiency.
- Introduced a new test case in test_tilelang_transform_producer_consumer_ws.py to verify that consumer-only local initializations do not leak into the producer.
- Added assertions in test_tilelang_issue_tma_no_ws.py to ensure correct behavior of TMA contexts.
* Add local access tracking and buffer management in producer_consumer_ws.cc
- Introduced LocalAccessSummary and LocalLiveSet structures to track read/write buffers and variable definitions.
- Implemented BufferDataToBufferCollector and LocalAccessCollector classes for efficient buffer data collection and access summary generation.
- Enhanced the logic for handling TMA copy write buffer data extraction, improving the producer-consumer interaction in TMA contexts.
* fix
* fix
* fix
* Refactor producer-consumer logic and enhance kernel execution in tilelang
- Updated the SparseFlashAttn class to improve kernel invocation clarity by separating kernel definition and execution.
- Enhanced producer-consumer workflow in producer_consumer_ws.cc by introducing optional guards for wait statements, ensuring correct behavior in concurrent contexts.
- Added utility functions to manage loop body lets, improving code readability and maintainability.
- Expanded test coverage in test_tilelang_transform_producer_consumer_ws.py to verify the preservation of guarded waits and backpressure logic.
* Enhance test coverage and execution flow in attention sink examples
- Added pytest timeout decorators to tests in test_example_attention_sink.py and test_example_blocksparse_attention.py to prevent long-running tests.
- Updated the main execution block in test_example_attention_sink.py to disable cache and directly call the sliding window test function, improving test execution clarity.
* update
* lint fix
* lint fix
* Refactor and enhance barrier handling in tilelang transformations
- Removed unused VisitStmt_ method from fuse_mbarrier_arrive_expect_tx.cc to streamline code.
- Introduced ProducerSimtCopyDetector in producer_consumer_ws.cc to identify SIMT copy patterns, improving buffer access tracking.
- Enhanced test coverage in test_tilelang_transform_fuse_mbarrier_arrive_expect_tx.py and test_tilelang_transform_producer_consumer_ws.py to validate new functionality and ensure correct barrier behavior.
* lint fix
* lint fix
* Refactor example scripts and clean up utility functions
- Removed commented-out print statements in example_warp_specialize_gemm_copy_1_gemm_0.py and example_warp_specialize_gemm_softpipe_stage2.py for cleaner code.
- Simplified buffer scope checks in utils.h by removing inline functions and directly using scope comparisons in relevant functions.
- Enhanced buffer handling in lower_tile_op.cc and producer_consumer_ws.cc by streamlining condition checks for buffer types.
* Remove unnecessary cache disabling in tilelang tests and examples
- Eliminated calls to tilelang.disable_cache() in example_gqa_decode_varlen_logits.py, test_tilelang_issue_tma_no_ws.py, and test_tilelang_language_tma_copy.py for cleaner code and improved performance.
- Updated test execution flow in test_tilelang_issue_tma_no_ws.py to directly call the main testing function.
* fix
* Add LocalAccessSummary collection for loop bounds in producer_consumer_ws.cc
- Introduced CollectExpr method in LocalAccessCollector to gather access summaries for expressions.
- Updated pre-loop liveness assignment to include variables used in pipeline loop bounds, ensuring accurate classification of scalar setups.
* enhance
* fix
* enhance
* fix
* fix
---------
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* [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>
* [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>
* 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
* 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
* [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>
* [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.
* 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
* Fix type annotations for T.reshape and T.view
* Fix issue #1666: reduce_sum with clear=False not accumulating correctly
* address comments and add testcases
* add more tests
---------
Co-authored-by: SiriusNEO <chaofan@deepseek.com>
* [Enhancement] Add global memory load/store intrinsics for various vector widths (#1710)
* Introduced new intrinsics: ldg32, ldg64, ldg128, ldg256 for loading data from global memory.
* Added corresponding store intrinsics: stg32, stg64, stg128, stg256 for storing data to global memory.
* Updated built-in functions and templates to support these new operations.
* Enhanced documentation for usage examples and type checks in the language module.
* Updated CUDA code generation to utilize new intrinsics for improved performance.
* lint fix
* test fix
* Refactor buffer type annotations in global memory load/store intrinsics
Updated the type annotations for the `ldg` and `stg` functions to use `tir.Buffer` instead of the generic `Buffer` type. This change enhances type specificity and aligns with the TVM type system for better clarity and correctness in memory operations.
* lint fix
* [Enhancement] Introduce LDG/STG lowering options and transform pass
* Added configuration options to enable/disable lowering of non-predicated and predicated global memory loads/stores to ldg/stg intrinsics.
* Implemented the LowerLDGSTG transformation pass to optimize global memory access for better performance.
* Updated tests to require a minimum CUDA compute version for relevant functionality.
* Enhanced documentation for new configuration options and transformation behavior.
* lint fix
* Refactor LDG/STG intrinsic selection logic for improved readability
* Replaced switch-case statements with if-else conditions for determining the appropriate LDG operation based on byte width.
* Enhanced error handling for unsupported byte widths with a clear fatal log message.
* lint fix
* fix
* testfix
* Enhance LDGSTG pass to skip processing for CuTeDSL backend
* Added a check to skip the transformation if the target has the "cutedsl" key.
* Updated tests to require a minimum CUDA compute version of 10.0 for LDG and STG code generation tests.
* Refactor global memory load/store functions to include conditional variants
* Updated the CUDA code generation to support conditional versions of LDG and STG functions, enhancing flexibility in memory operations.
* Modified existing functions to include new conditional variants (e.g., `ldg32_conditional`, `stg32_conditional`) for better control over memory access based on predicates.
* Adjusted comments to reflect the changes in function signatures and usage.
* Update comments and input handling for LDG intrinsics to include address and optional predicate
* Modified comments for `ldg32`, `ldg64`, `ldg128`, and `ldg256` to clarify the function signatures, indicating the use of 'address' and an optional 'predicate'.
* Changed the number of inputs for these LDG functions from 2 to -1 to accommodate the new signature format.
* Add correctness verification to LDG code generation tests
* Introduced assertions to verify the correctness of generated CUDA source for `ldg32`, `ldg64`, `ldg128`, and `ldg256` functions.
* Added checks to ensure the output matches the expected reference values within specified tolerances.
* Updated test cases to enhance validation of LDG intrinsic functionality.
* test fix
* Refactor LDG and STG intrinsics to use unified naming convention
* Updated all instances of `ldg` and `stg` functions to their new names: `load_global` and `store_global`, respectively, across the codebase.
* Modified related test cases to reflect the new function names, ensuring consistency in the testing framework.
* Enhanced comments and documentation to clarify the changes in function signatures and usage.
* Update reference values in LDG256 predicated code generation test to improve accuracy
* lint fix
* Rename test function for clarity and add new test for predicated load/store global intrinsics. Update assertions to reflect new naming convention in generated kernel source.
* lint fix
* subtype cast fix
[Refactor] Update fp4 data structures and CUDA code generation
- Refactored `fp4_e2_2_t` structure to include methods for accessing and setting low and high bits.
- Updated `PrintVecElemLoad` and `PrintVecElemStore` methods in `codegen_cuda.cc` to reflect changes in member access for `fp4_e2_4_t` and `fp4_e2_2_t`.
- Adjusted test for `tilelang_copy_fp4` to ensure correct input dimensions and output validation based on data type consistency.
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
* Remove RewriteUnsafeSelect Pass
* Add test for
* Add test for T.Select
* Add test for T.Select
* Add test for T.Select
* Update test_tilelang_language_select.py
---------
Co-authored-by: wangxiangwen <xiangwen.wang@deepseek.com>
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
* Enhance loop vectorization logic and buffer handling
- Refactored vectorization logic to utilize configuration settings for enabling/disabling vectorization.
- Introduced new checks for buffer access invariance in call nodes to improve vectorization accuracy.
- Updated GCD calculations for vector sizes to include call node constraints, enhancing overall performance.
- Added verbose logging for better insights during vectorization planning.
- Introduced tests for vectorization of infinity and bitwise logical operations to validate new functionality.
* lint fix
* fix
* test fix
* move to vectorize_loop
* fix
* lint fix
* Refactor atomic add functions to improve type handling and reduce redundancy
- Introduced utility functions ToFloat2 and ToFloat4 for type conversion, enhancing the flexibility of atomic add operations.
- Updated AtomicAddx2 and AtomicAddx4 functions to accept generic value types, streamlining the code and improving readability.
- Removed outdated function overloads and consolidated logic for atomic addition, ensuring consistency across implementations.
- Enhanced memory order handling in atomic operations for better performance and clarity.
* add related test
* Remove atomic add vectorization implementation
- Deleted the atomicadd_vectorize.cc and atomicadd_vectorize.h files, which contained the automatic vectorization pass for atomic add operations.
- Removed references to the deleted files in atomic_add.cc and loop_partition.cc to clean up the codebase and eliminate unused dependencies.
* [Feature] Implement atomic reduction operations and enhance atomic add functionality
- Added support for atomic max and min operations with corresponding classes and methods.
- Introduced vectorization capabilities for atomic add operations based on data type and target architecture.
- Refactored atomic add intrinsic calls to improve clarity and consistency.
- Enhanced layout inference for atomic operations and integrated new utility functions for target architecture checks.
- Updated tests to validate the new atomic operations and their vectorization behavior.
* Refactor atomic operation code for improved readability and consistency
- Cleaned up whitespace and formatting in atomic add and reduce implementations for better code clarity.
- Enhanced comments in CUDA code generation for atomic operations to improve understanding.
- Updated function signatures and layout definitions for consistency across atomic operation files.
- Ensured proper alignment of code for better maintainability and readability.
* lint fix
* Enhance loop vectorization and layout handling across multiple operations
- Updated function signatures in atomic add, atomic reduce, copy, and fill operations to include layout map parameters for improved layout handling.
- Refactored vectorization logic to utilize layout maps, ensuring better performance and compatibility with various buffer layouts.
- Enhanced the LowerParallelLoop function to accept layout maps, facilitating more efficient loop transformations.
- Added checks for buffer contiguity in vectorization processes to ensure correctness when using layout maps.
- Updated tests to validate the new layout handling and vectorization behavior.
* Add atomic operation checks in legalize_safe_memory_access
- Introduced a new method to identify atomic operations within the legalize_safe_memory_access transformation.
- Updated the VisitStmt_ function to handle both CallExtern and atomic operations, ensuring recursive condition collection for these cases.
- Enhanced comments for clarity on the handling of atomic operations in the context of memory access legality.
* Add atomic operations for HIP code generation
- Implemented support for various atomic operations including atomic add, atomic load, atomic store, atomic max, and atomic min in the HIP code generation.
- Enhanced the handling of atomic operations to include optional memory order parameters.
- Improved code readability with added comments explaining the purpose of each atomic operation.
- Ensured consistency with existing atomic operation implementations in the codebase.
* Refactor AtomicOpBaseNode and related classes for improved clarity
- Changed the GetElemOpStatic method to a virtual GetElemOp method in AtomicOpBaseNode, enhancing polymorphism.
- Updated AtomicAddNode, AtomicMaxNode, and AtomicMinNode to override the new GetElemOp method.
- Removed unnecessary template parameters from AtomicOpBaseNode, simplifying the class structure.
- Cleaned up includes in atomic_reduce.cc to remove unused dependencies, improving code organization.
* Enhance argument validation for atomic operations
- Added checks in AtomicAdd, AtomicMax, and AtomicMin constructors to ensure at least two arguments are provided, improving error handling and user feedback.
- Removed the unused LowerTMA method declaration from the AtomicAdd class, streamlining the codebase.
* lint fix
* fix
* test fix
* performance fix
* remove useless comments
* Enhance CUDA code generation for BroadcastNode by implementing compile-time constant folding and runtime broadcasting for various lane configurations. Improved handling for 4-bit and 8-bit integer types, ensuring correct replication and type casting in output expressions. This update increases performance and correctness in CUDA kernel generation.
* add test
* lint fix
* fix
* support cp.reduce.async.bulk.tensor and add test
* Refactor flash attention example by removing unnecessary layout annotations
* support swizzle layout for tma reduce
* auto swizzle for non-1d tma atomic add
* upd example and test
* lint
* typo
* add constraint for test
* Refactor CUDA data type mapping by moving the to_CUtensorMapDataType function to utils.cc and utils.h, while removing redundant definitions from atomic_add.cc and copy.cc.
* lint
* [Fix] Refactor type hint extraction logic in DSLMutator for better clarity and handling of annotations
* [Refactor] Remove redundant tensor creation in loop layout tests and update kernel compilation parameters
* Add unroll loop functionality and corresponding test case
- Introduced a new `UnrollLoop` function in the transform module to unroll loops based on various configuration options.
- Added a test case in `test_tilelang_language_unroll.py` to validate the behavior of `T.unroll` with only the extent parameter, ensuring correct kernel generation with unroll pragmas.
* Refactor unroll kernel implementation and update test case
- Changed the kernel function in `test_tilelang_language_unroll.py` to use a new `unroll_kernel` function that compiles and returns the output tensor, improving clarity and structure.
- Updated the `OptimizeForTarget` function in `phase.py` to ensure the `UnrollLoop` transformation is applied correctly, maintaining consistency in optimization phases.
* lint fix
* lint fix
- Move `tilelang/language/eager/dtypes.py` to `tilelang/language/dtypes.py`
since dtype definitions are general-purpose, not eager-specific
- Add `bytes` property to dtype (alias for `itemsize`) for convenience
- Add type hints for `bits` and `bytes` properties in TYPE_CHECKING block
- Create `tilelang/dtypes.py` for direct import via `from tilelang.dtypes import ...`
- Update all import paths accordingly
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
* Refactor: Replace @tilelang.lazy_jit with @tilelang.jit in examples and tests
Updated all instances of the @tilelang.lazy_jit decorator to @tilelang.jit in the lazyjit example notebooks and related test files. This change aligns the code with the new JIT compilation approach, enhancing consistency across the codebase. Additionally, removed the lazy_jit import from the module initialization to streamline the API.
* Enhancement: Improve layout annotation handling in T.copy and parallel loops
Updated the T.copy function to accept a new keyword-only parameter, loop_layout, allowing users to specify layout hints for the outermost parallel loop. Enhanced layout annotation handling in CopyNode and AtomicAddNode classes to ensure compatibility with SIMT operations. Added tests to validate the functionality of loop layout annotations, improving robustness and clarity in layout management for parallel loops.
* Add out_idx parameter to @jit and validate it's only used in lazy mode
* lintfix
* Refactor: centralize mode inference in JITImpl.__call__ before parse_args
* test fix
* Enhancement: Introduce JITNoBuilderError for eager mode checks in T.Kernel and T.const functions. Ensure proper error handling when Builder is unavailable, improving robustness in JIT context management.
* lint fix
* lint fix
* Enhancement: Update example_triton_sparse_gqa_decode_varlen_indice.py to improve performance metrics output by adding average FLOPS calculations. Modify layout_inference.cc to refine buffer layout inference checks by ensuring fragment buffers are considered.
* lint fix
* Enhancement: Refactor JIT implementation to centralize mode initialization and improve error handling for eager mode. Update matmul tests to utilize both Cython and TVM FFI JIT kernels, ensuring consistency in output validation.
* Refactor: Remove unused simplify_prim_func decorator from example_gemm_intrinsics.py and related test files. Enhance layout.cc to simplify integer set analysis for multi-threaded compilation, ensuring constant integer outputs. Update lower_tile_op.cc to enforce constant integer checks in layout output shapes. Improve error handling in builder.py for JIT context management.
* remove claude.md
* Refactor: Clean up code formatting in layout.cc, lower_tile_op.cc, and proxy.py for improved readability. Add missing import in test_tilelang_kernel_int4_gemm_mma.py and ensure consistent spacing in builder.py.
* Refactor: Replace logical conditions with tir_all in complex expressions within test_arith_hard.py for improved clarity and consistency. Update main execution flow to call test_hard_prove instead of tilelang.testing.main.
* Enhancement: Add @tilelang.testing.requires_cuda decorator to test functions in test_tilelang_issue_1549.py and test_tilelang_issue_1601.py to ensure CUDA requirements are explicitly stated for relevant tests.
* lint fix
* Refactor: Simplify output shape calculation in layout.cc by directly setting max value from ist.max(). Update test_valid_loop in test_tilelang_fragment_loop_checker.py to disable certain tests and ensure proper execution flow. Enhance error handling in builder.py to correctly manage PrimFunc return types.
* Refactor: Update boolean condition in test_boolop to use Or and Not for improved clarity and consistency in logical expressions.
* Refactor: Update test execution flow in test_tilelang_fragment_loop_checker.py to call tilelang.testing.main. Modify JIT implementation to handle callable functions more generically. Enhance TirTemplate class in builder.py to return an empty tuple if matcher is None, improving robustness.
* Refactor JIT function handling: Replace LazyJITFunc with JITFunc across the codebase, updating related function signatures and logic to support eager JIT compilation mode. This change enhances the clarity and functionality of JIT operations.
* Update JITImpl documentation: Clarify output tensor index attribute description and remove outdated comment regarding implicit TIR construction.
* Refactor TileLang to transition from v2 to eager language module. Update imports across multiple files to replace references to the deprecated v2 module with the new eager module, ensuring compatibility with the latest JIT compilation features. Remove obsolete v2 module files to streamline the codebase.
* Enhance JIT compilation support: Update JIT function handling to improve compatibility with eager execution. Refactor related function signatures and logic, ensuring seamless integration with the latest compilation features. This change aims to streamline JIT operations and enhance overall code clarity.
* Fix some bugs on ci and open rocm ci test
* fix pre-commit
* fix build issue on hip
* refactor on hip
* pre-commit fix
* fix test with pdl
* fix pdl test
* [Enhancement] Improve shape and stride binding for subbyte types in ArgBinder
* Updated the shape binding logic in ArgBinder to handle subbyte types (bits < 8) by introducing a packing factor for logical shapes.
* Enhanced the handling of strides for subbyte types, ensuring correct mapping between logical and runtime strides.
* Refactored the binding process to accommodate both symbolic and constant dimensions, improving the robustness of the shape binding mechanism.
* Removed obsolete test file `test_ir_kernel_frame.py` as part of cleanup efforts.
* lint fix
* [Feat] PDL Support
* remove unused comments/codes
* fix mismatch funcname
* Resolve duplicate config issue for AI review
* fix typo error && clang-tidy error
* chore: retrigger CI
* change file comment
* rename test_tilelang_jit_ctypes.py to test_tilelang_jit_cython.py
* remove ctype backend
* Move pdl attributes to header and remove redundant VisitStmt_
* chore: trigger CI
* modify the cuda kernel with pdl due to nvcc's bug
* throw an error when invoking __ldg with pdl due to nvcc's bug
* Throw an error when pdl is not supported
* fix checking about pdl
* add comments about MarkCudaSyncCalls
* remove pdl support with compute_capability<90
* ruff check
* fix nvrtc about pdl
* update tvm & extend pdl support for tvm_ffi & add test for tvm_ffi
* ensure that kUseDynamicSharedMemoryTag is the last tag in launch_param_tags
* fix bug introduced by branch merging
* [Feature] Add get_host_source method to CythonKernelAdapter
* Introduced a new method, get_host_source, in the CythonKernelAdapter class to return the source code of the host function. This addition enhances the functionality of the adapter by providing direct access to the host kernel source.
---------
Co-authored-by: silentCoder-dev <silentcoder@foxmail.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* fp4 related update, require_cu13
* Enhance CUDA type conversion handling and optimize dtype management
- Updated CUDA vectorized cast functions to ensure proper handling of float16, float32, bfloat16, and float8 conversions, adding checks for bit sizes.
- Refactored dtype conversion logic in `cuda_fp4.h` to utilize `cudaRoundZero` for improved accuracy in floating-point conversions.
- Introduced a new method in `KernelParam` to convert TVM DataType to TileLang dtype.
- Adjusted argument binding logic in `arg_binder.cc` to allow for better subtype matching based on total bit counts.
- Enhanced dtype handling in `dtypes.py` to accommodate new float4_e2m1fn types and ensure compatibility with PyTorch.
This update aims to improve type safety and conversion accuracy across the codebase.
* lint fix
* lint fix
* typo fix
---------
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
* Refactor CUDA vectorized cast generation and remove unsupported FP8 type
* test fix
* lint fix
* Refactor CUDA vectorized cast function naming for clarity
* Add support for float4_e2m1fn type conversions in CUDA vectorized casts
- Implemented conversions between float4_e2m1fn and float32, half2, and float2 in utils.cc and cuda_fp4.h.
- Updated test_tilelang_language_vectorized_cast.py to validate new conversions and ensure correctness.
- Enhanced dtype conversion in dtypes.py to handle float4_e2m1fn appropriately, logging a warning for unsupported types in PyTorch.
* Enhance vectorized cast tests for new data types
- Added tests for vectorized casting of float8 and float4 data types, ensuring compatibility with CUDA compute versions.
- Refactored existing test functions to improve clarity and organization, separating tests for different data types.
- Updated parameterization to include additional test cases for new conversions.
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
Tong WU