* [Feature] Implement Hoist Loop-Invariant If Pass
- Added a new pass to hoist loop-invariant if statements out of loops, improving optimization opportunities.
- Introduced classes for collecting and checking conditions, as well as for rewriting statements.
- Integrated the new pass into the optimization pipeline and provided a corresponding API for usage.
* Add tilelang semantics guide to programming guides (#1740)
* Add CallNodeChecker to Loop Unswitching Pass and integrate tilelang cache disabling in benchmark script
- Introduced CallNodeChecker class to identify CallNode expressions in loop conditions, enhancing loop-invariant checks.
- Updated IsLoopInvariant function to reject conditions containing CallNodes, preventing potential side effects.
- Added tilelang.disable_cache() in benchmark_mha_sink_fwd.py to optimize performance during benchmarking.
* lint fix
* lint fix
* Enhance Loop Unswitching Pass with Let Binding Support
- Added support for Let-bound variables in the WrittenBufferReadChecker to improve buffer read checks.
- Introduced UsesLoopVarThroughLetBindings function to check if conditions depend on loop variables through Let bindings.
- Updated IsLoopInvariant function to account for Let bindings when determining loop invariance.
- Enhanced HoistableIfFinder to track Let bindings for variables bound to BufferLoad expressions.
- Added debug print statements in the OptimizeForTarget function to visualize the module state before and after loop unswitching.
* fix
* lint fix
* Enhance Loop Unswitching with CallNode Exclusion Check
- Added CallCheckerExcludingIf class to ensure function calls outside of hoisted if statements are identified, preventing potential synchronization issues during loop unswitching.
- Updated loop unswitching logic to incorporate the new call checker, enhancing safety and correctness.
- Integrated debug print statements in OptimizeForTarget to visualize module state before and after loop unswitching.
- Disabled tilelang cache in the benchmark script for improved performance.
* Update TVM subproject to latest commit and remove debug print statements in OptimizeForTarget function
* Add support for disabling loop unswitching optimization
- Introduced a new configuration option `tl.disable_loop_unswitching` to allow users to disable the loop unswitching optimization.
- Updated the Loop Unswitching pass to check this configuration and return the original function if the option is enabled.
- Added relevant documentation in the PassConfigKey enumeration for clarity.
* update tvm
* Refactor: Update TileLang to use tirx instead of tir
- Removed references to `tvm.tir` and replaced them with `tvm.tirx` across various files, including examples and backend operations.
- Updated target configuration in documentation to reflect the new usage of target config dictionaries instead of CLI-style strings.
- Cleaned up `pyproject.toml` and `load_tvm.cmake` by removing obsolete paths.
- Enhanced examples for dequantization and GEMM operations to utilize the new `tirx` constructs.
- Adjusted various backend operations to ensure compatibility with the new `tirx` namespace.
* pass fix
* Update submodule 'tvm' to latest commit f24c7b19
* Apply formatting fixes
* Update CMakeLists.txt to set HIDE_PRIVATE_SYMBOLS for TileLang and modify target_host initialization in lower.py for compatibility with TVM's target handling.
* Refactor CMakeLists.txt and kernel_cache.py to update target names and library dependencies for TVM integration. Adjusted output targets and library names for better compatibility across platforms.
* Update repair-wheel-command in pyproject.toml to exclude libz3.so with wildcard for better compatibility with dependency packages.
* Fix macOS Z3 dylib lookup
* Move macOS Z3 install-name fix to install step
* Update TVM submodule for runtime import fix
* Update ROCm backend for TVM module factory changes
* Fix build-tree Z3 runtime lookup
* Fix macOS build-tree Z3 lookup
* Support legacy target option strings
* Normalize CUDA target in transform test
* Qualify Metal stream FFI lookup
* Fix PyPI Z3 packaging links
* Update TVM for Windows DLL lookup
* Register Windows dependency DLL directories
* Bump tvm-ffi to match tirx
* tvm update
* Apply lint fixes
* Enhance reduction operations with support for nan propagation and vectorization. Introduce new functions for preferred vectorized size and update existing reduction logic to handle packed operations for bfloat16 and float16 types. Add nan-aware min and max operations in CUDA and ROCm backends, and update related tests to validate functionality.
* Refactor ReduceLowerer: Remove unnecessary buffer allocation logic for packed operations in reduce.h
* Normalize PyPI Z3 dylib before linking
* Drop legacy target option string support
* Fix type reference for barrier annotations in CUDA copy operations. Updated `BufferLoadNode` to `tirx::BufferLoadNode` in `GetBarrier` and `LowerCluster` methods to ensure compatibility with recent changes in the TIR API.
---------
Co-authored-by: Yu Cheng <yu.cheng@pku.edu.cn>
Co-authored-by: Yichen Yan <oraluben@outlook.com>
* [Refactor] Enhance Producer-Consumer WS Transformation
- Introduced a new utility function `ExprUsesAnyVar` to check variable usage in expressions.
- Updated `BuildWSBlock` to collect enclosing loops for better context in transformations.
- Modified global linear index calculations to account for nested loops.
- Enhanced phase counter logic to handle non-rectangular loop nests.
- Added functionality to hoist workspace branches out of loops for improved performance.
- Implemented `CollectEnclosingForLoops` to facilitate loop ancestry collection.
- Refactored related methods to accommodate new logic and ensure correctness in transformations.
* Fix nested warp-specialized pipeline sync
* lint
* Reject autotune default scalar inputs
* [Autotune] Defer scalar input validation and add do_not_specialize option
- Move _validate_input_supply_requirements from __call__ into run(),
after the "tunable parameters already provided" check. This prevents
spurious errors when all tunable params are supplied by the caller
and the autotuner skips profiling entirely (fixes test_example_mha_fwd_varlen).
- Add `do_not_specialize` parameter to @autotune decorator. Parameters
listed here are excluded from the autotune cache key, so changing
their values reuses the previously tuned configuration instead of
triggering re-autotuning.
Usage:
@autotune(configs=get_configs(), do_not_specialize=["batch_size", "UQ"])
@tilelang.jit
def kernel(batch_size, UQ, ...):
...
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* Revert unrelated WS changes from autotune PR
* Remove CPU count logging and related calculations from AutoTuner class in tuner.py
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* [Feature] Extend t.copy to support TMA multicast and SM-to-SM cluster copy
Add multicast copy (tma_load_multicast) and shared-memory-to-shared-memory
cluster copy (tma_store_cluster / ptx_cluster_store) to the t.copy operator.
Multicast copy:
- A single CTA issues tma_load_multicast to broadcast a tile to multiple
CTAs in the cluster simultaneously; other masked CTAs receive passively.
- New `cluster_mask` parameter on T.copy() controls which CTAs participate.
SM-to-SM cluster copy:
- Fast path via cp.async.bulk.shared::cluster with mbarrier synchronisation
(new `remote_barrier` parameter on T.copy()).
- Slow path via cooperative_groups::map_shared_rank for element-by-element
fallback (new `dst_block` parameter on T.copy()).
Supporting changes:
- New builtins: tma_load_multicast, tma_store_cluster, ptx_cluster_store,
cluster_sync, get_cluster_id, get_cluster_block_rank, get_cluster_block_nums,
mbarrier_arrive (cluster-scope).
- Codegen (codegen_cuda.cc) and device templates (barrier.h, copy_sm90.h)
for all new intrinsics; unified mbarrier API via tl::mbarrier_* free functions.
- inject_tma_barrier: handle tma_load_multicast; distinguish thread guards
from cluster-rank conditions for correct expect_tx injection.
- lower_hopper_intrin: migrate barrier allocation to builtin::create_barriers;
hoist user ptx_init_barrier_thread_count calls alongside compiler barriers.
- warp_specialized_rewriter: support tma_load_multicast as producer op;
UserBarrierInitExtractor to separate user barrier inits from body;
ragged (dynamic-extent) pipeline prefix counter; fix nested IfThenElse
state corruption in shuffle-elect optimisation.
- multi_version_buffer_rewriter: ragged pipeline support with runtime
prefix counter for correct ping-pong buffer versioning.
- thread_storage_sync: reserve fixed barrier IDs (kProducer/kConsumer)
for 256-thread CTA warp-specialised split to prevent deadlocks.
- pipeline_planning: defensive handling of non-BufferLoad mbarrier args.
- Shared memory alignment bumped to 128 bytes for TMA.
Co-authored-by: Guangda Sun <2012661711@qq.com>
* [Docs] Add programming guide for Cluster TMA features
Add docs/programming_guides/cluster_tma.md covering the two new
T.copy extensions introduced in the t_copy_extend feature branch:
- TMA multicast (cluster_mask): explains how a single TMA transaction
broadcasts one global tile to multiple CTAs in a cluster simultaneously,
with API usage, a per-rank behaviour table, and a complete code example.
- SM-to-SM cluster copy (dst_block / remote_barrier): documents the fast
path (cp.async.bulk.shared::cluster + mbarrier, single-thread async DMA)
and the slow path (map_shared_rank element-by-element SIMT fallback),
including the synchronisation contract for source and destination CTAs.
Also covers cluster helper builtins (T.get_cluster_block_rank,
T.cluster_sync, T.mbarrier_init/arrive/wait_parity, etc.) and a
split-K sketch combining both features end-to-end.
* change cluster_size -> cluster_dims
* fix merge conflict
* fix TILELANG CHECK bug
* unify cluster function
* fix pre-commit errors
* docs(cluster_tma): clarify multicast non-issuer behavior in lowering note
* fix(copy): gate dst_block cluster lowering on target capability
* fix(cuda/cluster): hard-trap when cluster intrinsics are unavailable instead of assert(false)+dummy return.
* fix inject_tma & lower_hopper & lower_tile bugs
* [warp specialize]Fix mbarrier retarget for tma_load_multicast in warp-specialized producer path
* fix(tma-barrier): correct arrive thread count under equality-guarded threadIdx
* fix: overlapped TMA barrier IDs
* fix: multicast checks
* minor fix
* minor fix
* fix: tma_load_multicast() in the stage-local mbarrier rewrite
* minor fix
* format
* minor fix
* fix: Track allocation from every hoisted barrier init.
* fix: Scope the 128-byte alignment to kernels that actually use TMA.
* fix: Assert the tl::tma_store_cluster lowering, not just the output.
* format fix
* fix: Don't assume both if arms transfer the same number of bytes.
* fix: Don't wrap the kept pipelined loop in a SeqStmt before the For post-processing runs.
* rm mbarrier_init and use alloc_cluster_barrier
* fix: Don't make cluster copy a global exemption for unrelated TMA loads.
* fix: Don't drop dependency modeling for handle-based mbarrier_wait_parity().
* fix: Don't hoist extracted barrier-init if statements to the block root.
* fix(copy): honor remapped dst layout in cluster-copy slow path
* fix(copy): gate cluster TMA fast path on provable contiguity
* fix(copy): add barrier completion for SIMT cluster-copy fallback
* fix: remove dup cluster_sync
* fix: remove not used _get_mbarrier
* fix: remove mbarrier related code in codegen_cuda.cc and barrier.h
* fix: codegen_cuda.cc: reuse functions in cluster.h and remove dependency on cg
* fix testting
* fix(transform): prevent ragged_prefix free var from breaking MakePackedAPI
* minor fix
* format
* fix inject tma barrier: fix crash and correct mbarrier thread-count injection
* revert inject tma
* fix(inject_tma_barrier): avoid SIGSEGV and correctly infer barrier arrive thread counts
* fix: add tma_load_multicast same as tma_load
* T.copy_cluster
* fix: renaming cluster mbarrier_arrive to ptx_arrive_cluster_barerier
* fix: remove mbarrier_arrive in builtin.cc
* test_tma_dsmem: fallback test
* [Feature] Multi-TMA fallback for non-contiguous T.copy_cluster regions
Previously, a non-contiguous copy region (e.g. a 2-D slice where the
inner dim does not span the full buffer width) with a remote_barrier
fell back to element-wise SIMT stores, losing TMA throughput.
This commit introduces a "multi-TMA" middle path in LowerClusterCopy:
the ND copy region is recursively decomposed into individual contiguous
rows, each emitting one tma_store_cluster call. The mbarrier
arrive_count is updated to N (the number of rows) so the destination
CTA's wait() still works correctly without any API change.
Key changes:
- src/op/copy.cc: add static MakeTMARows() that recursively splits an
ND region into contiguous rows; static extents are unrolled at
compile time, symbolic extents produce TIR For loops.
- src/op/operator.h: add UpdateBarrierArriveCallback (Var → PrimExpr)
to LowerArgs so LowerClusterCopy can propagate the new arrive_count.
- src/transform/lower_tile_op.cc: collect arrive_count overrides in
barrier_arrive_updates_ and apply them to the barrier_init block
annotation after visiting the block body, before LowerSharedBarrier
consumes it.
- testing/python/cuda/test_tma_dsmem.py: replace the old SIMT-barrier
test with test_store_cluster_multi_tma_barrier (2-D, M rows) and add
test_store_cluster_3d_multi_tma (3-D, D×M rows) to cover the
two-level recursive decomposition.
* bug fix
* fix: unified continuous check
* fix doc
* fix: remote redundant annotation converting
* fix: remove unused code
* fix: remove unused code
* rm unused comments
* minor fix
* fix: rm unused helper func
* minor
* minor
* Clean up copy cluster merge formatting
* Reduce copy cluster PR churn
---------
Co-authored-by: Guangda Sun <2012661711@qq.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* bugfix: legalize mixed-sign ramp negative indices
* fix: update dtype handling in legalize_negative_index and add test for mixed-sign ramp legalization
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* [CuTeDSL] Support CUTLASS DSL 4.5 setup
Update the CUDA test requirements to the 4.5.0 CuTeDSL wheel stack and keep the runtime compatibility check accepting older supported CuTeDSL versions. Also accept the relocated Constexpr export while retaining the fallback for older wheels.
The 4.5.0 pin is a CI/test environment choice. CuTeDSL PDL APIs were introduced upstream in 4.3.4, so the pin should not be read as the minimum PDL version.
* [CuTeDSL] Wire PDL lowering and launch support
TileLang PDL intrinsics were already available for the CUDA backend, but the CuTeDSL backend did not lower them to cute.arch.griddepcontrol_* or enable PDL on generated kernel launches.
Lower T.pdl_trigger() and T.pdl_sync() to the CuTeDSL griddepcontrol APIs, re-export those APIs from tilelang.contrib.cutedsl when available, and pass use_pdl / programmatic stream serialization in Python and C++ launchers. Add CuTeDSL tests that cover device lowering and launcher generation.
---------
Co-authored-by: Jayce <jayce.su@enflame-tech.com>
* [Refactor] refactor gemm_sp following 5d729eee
* [Doc] update doc
* [Refactor] remove gemm_sp CUTLASS templates
* [cute] add mma sp tempaltes for sm80
* [templates] use templates in codegen
* Add WGMMA_SP templates
* Pass layout_map to py lowering
* Fix type
* Add wgmma.sp checking
* Implement wgmma_sp_ss
* Add wgmma_sp_rs
* Add sparse selector
* Fix layout and param pass
* Fix transpose metadata
* Add mma.sp fp8
* Fix integer rounding
* Fix metadata layout
* Update compress and test cases
* Remove debug print statements from sparse compress and wgmma emitter
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Clean up gemm_sp test file
- Remove unused imports (tvm, _TORCH_DTYPE_TO_STR)
- Remove print("pass") calls
- Normalize all parametrize tuples to compact aligned format
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Remove unused tvm import from gemm_sp_wgmma
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Update examples, tests, and docs for new compress() API
- Replace removed compress_sm90 with compress() in sparse_tensorcore
example and test_compress_utils
- Fix example_gemm_sp.py: remove ARCH_INFO and make_cutlass_metadata_layout;
compress() now produces natural-layout int16 metadata on all architectures
- Update matmul_sparse.md to reflect the simplified compress() API and
remove the CUTLASS layout annotation requirement
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Refactor sparse constants into sparse_config.py and clean up examples
- Add tilelang/utils/sparse_config.py as single source of truth for
SPARSE_PARAMS, E_FACTOR_MAP, E_REPLICATE_FACTOR, and get_e_factor()
- Remove inline E_FACTOR_MAP/E_REPLICATE_FACTOR from SparseTensorCoreIntrinEmitter;
import from sparse_config.py with backward-compatible class aliases
- Update compress() API: remove transposed/block_k/arch params, use T.int16 metadata
- Delete example_custom_compress.py; move torch_compress reference to
testing/python/utils/test_compress_utils.py with float8 .view() support
- Update examples (example_gemm_sp, tilelang_example_sparse_tensorcore) and
test files to use new compress() API and get_e_factor()
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Compute e_factor and e_replicate_factor instead of hardcoding tables
- Remove E_FACTOR_MAP and E_REPLICATE_FACTOR dicts from sparse_config.py
- get_e_factor: derived from SPARSE_PARAMS group size and meta_dtype.bits
(hardware always packs 4 bits per sparsity group)
- get_e_replicate_factor: derived from a_dtype.bits (<=8 bit → 1, else → 2)
- SparseTensorCoreIntrinEmitter uses the two functions directly instead of
class-level dict aliases
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Use DataType keys in SPARSE_PARAMS
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Move sparse_config.py from utils/ to intrinsics/
The constants (SPARSE_PARAMS, get_e_factor, get_e_replicate_factor) are
hardware-level parameters primarily consumed by mma_sp_macro_generator;
intrinsics/ is the natural home. utils/sparse.py re-exports them for
backward compatibility.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Rename sparse_config.py to sparse_params.py
"params" better reflects the content: hardware sparsity format parameters
(SPARSE_PARAMS, get_e_factor, get_e_replicate_factor), not runtime config.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Fix sparse.py breaking TIR type hints by removing future annotations
from __future__ import annotations makes all annotations lazy strings,
which breaks TIR prim_func's get_type_hints() on the inner kernel.
Use Optional[torch.dtype] instead of the union syntax to avoid needing it.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* Fix compress utils
* Update example
* Add a compress benchmark
* Remove print
* Remove unused MetaType
* Polish
* Refactor gemm_sp op
* Update doc && remove unused file
* Fix bug spotted by coderabbit
* Fix
* Add note
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
* fix: T.Parallel + call_extern + address_of correctness in warp-specialized pipelines
Two bugs caused incorrect results when T.Parallel loops use
T.call_extern with T.address_of(shared_buffer) inside pipelined,
warp-specialized kernels:
Bug 1 (lower_tile_op.cc): has_non_local_store did not recognize
address_of(shared_buffer) inside call_extern as a non-local access,
so parallel_loop was set to false and thread partitioning was skipped
entirely — every thread ran the full serial loop.
Bug 3 (producer_consumer_ws.cc): wait_insert_pos was only computed
for kTmaProducer statements. SIMT/cp.async producers (e.g. T.Parallel
global→shared copies) were skipped, so the consumer read from
SIMT-produced shared buffers before the barrier wait — a data race.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* lin
* fix ws async producer wait placement
---------
Co-authored-by: Rachmanino <wutong.1109@bytedance.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* Enhance intrinsic modules by adding WGMMATensorCoreIntrinEmitter and TCGEN05TensorCoreIntrinEmitter. Introduce descriptor parameter classes for WGMMA and TCGEN05, improving the initialization and atom offset computation processes. Update the TensorCoreIntrinEmitter to support atom-level interfaces and macro definitions for efficient matrix multiplication operations.
* lint
* add test for mma/wgmma/tcgen5mma atom
* lint
* typo
* fix import path err
---------
Co-authored-by: Rachmanino <wutong.1109@bytedance.com>
* [Refactor] Move backend-specific GEMM implementations and transforms into backend directories
Restructure the codebase so that each backend (cpu, cuda, rocm) owns its
GEMM implementations, sparse GEMM implementations, and transform passes
under a consistent op/ subdirectory layout. The shared GEMM registry and
base classes remain in tileop/ as the platform-agnostic dispatch layer.
- Move gemm/gemm_sp registries from backend/ into tileop/ as registry.py
- Move CUDA GEMM impls (mma, mma_sm70, wgmma, tcgen05) into backend/cuda/op/gemm/
- Move CUDA sparse GEMM impl into backend/cuda/op/gemm_sp/
- Move CPU GEMM impl (scalar) into backend/cpu/op/gemm/
- Move ROCm GEMM impls (mfma, wmma) into backend/rocm/op/gemm/
- Move CUDA-specific transform passes from src/transform/ into src/backend/cuda/transform/
- Move CUDA runtime sources from src/runtime/ into src/backend/cuda/
- Remove dead backend-importing wrappers from transform/__init__.py
- Update phase.py and tests to import CUDA transforms from their canonical location
Each backend now has a symmetric op/ directory structure. Adding a new
backend no longer requires modifying shared transform or tileop modules.
* Refactor Python backend package layout
* Remove deprecated intrinsic implementations and related tests for matrix multiplication and element-wise addition. This includes the deletion of files for `benchmark_matmul_intrinsic.py`, `example_tilelang_gemm_fp8_intrinsic.py`, and associated test files, streamlining the codebase by eliminating unused components.
* Move CUDA transform passes back to common transform
* [Feature] Add Windows support with cross-platform minimal diff
* [Feature] Add ccache support for Windows
* [Feature] Add parallel autotune support for Windows
* [Feature] Add nvrtc support for Windows
* [Bugfix] Fix tilelibrary support for Windows
* [Bugfix] Fix alignment for Windows
* [Feature] Add clang-cl support for Windows
* [Bugfix] Fix numeric limit for Windows
* [Bugfix] Fix rand for Windows
* [Misc] Use dynamic test workload to avoid hang
* [Bugfix] Fix TMA alignment for Windows and sm120
* [Bugfix] Add unix-like hex formatter for Windows
* [CI] Add Windows wheel job
* [Feature] Flexible CUDA loader
* fix rebase
---------
Co-authored-by: SiriusNEO <chaofan@deepseek.com>
* [Refactor] Enhance SetMaxNRegCollector and SetMaxNRegInjector for improved register management
* Refactor SetMaxNRegCollector to return a structured result, including a flag for preserving explicit register settings.
* Update SetMaxNRegInjector to utilize the new result structure, ensuring proper handling of register hints.
* Modify logic to conditionally inject register allocation based on the presence of SIMT copies and explicit settings.
* Improve test cases to validate the injection of set_max_nreg calls and ensure correct behavior in various scenarios.
* lint
* [Refactor] Update register allocation logic to honor explicit hints
* Modify the register allocation logic to ensure that explicit user hints are respected even when SIMT copy-like statements are present.
* Adjust conditions for injecting register allocation to improve handling of register hints, enhancing overall register management.
* Ensure default behavior remains conservative while allowing for user-defined configurations.
* Fix warp-specialized reg hint annotation
---------
Co-authored-by: wutong.1109 <wutong.1109@bytedance.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* fix: TMA alignment to 1024 bytes on Blackwell
Previously 1024 byte alignemnt was guarded only by `IsHopper`.
This commit changes it to be guarded by `TargetHasBulkCopy` with the
same threshold as Hopper is using.
The kernel added in the test was failing on my RTX5090 GPU.
* Fix lint
* add gfx950 guard to test_gfx950_cp_async_gs_16_in_codegen
* add gfx950 checking and update other cases for CI issues
* update
* [AMD][gfx950] Fix async-copy swizzle conflict, re-enable pipelining, fix 32x32 MFMA store layout
* update for format checking
* upgrade CI rocm sdk version to be rocm7.2
* update the changes according to the review info
* update the codes for format checking
* onlu run copy async on gfx950
* move T.alloc_shared and T.alloc_fragment outside the T.Pipelined loop
* update the test case
* revert the change to pyproject.toml
* update test_tilelang_hip_codegen
* find a possible issue about failed CI test
---------
Co-authored-by: root <root@dell300x-ccs-aus-f06-01.cs-aus.dcgpu>
* [Feature] Batched AllReduce for better T.reduce performance (#1928)
When each thread holds multiple values to reduce, the previous code
emitted one AllReduce call per element inside a loop, causing redundant
synchronization barriers (e.g. 8 values × 6 syncs = 48 syncs).
This change adds batch support to AllReduce so all per-thread values are
reduced in a single butterfly pass sharing barriers (6 syncs total).
- Add `batch_size` and `workspace_stride` template parameters to
`tl::AllReduce` (defaults preserve existing scalar behaviour).
- In ReduceOpNode::Lower, detect when batching is beneficial
(batch_size > 1 and cross-warp reduction) and restructure into three
phases: loop(init + local reduce), batched AllReduce, loop(copy-back).
- Apply same optimisation to FinalizeReducerOpNode::Lower.
- Mirror template changes in HIP AllReduce.
Closes#1928
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Test] Add codegen tests for batched AllReduce
Verify that the batched AllReduce path (batch_size > 1, workspace_stride)
is emitted in the generated CUDA source for reduce_max, reduce_sum, and
reduce_min with cross-warp reductions.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Bugfix] Fix batched AllReduce codegen for HIP/ROCm targets
HIP AllReduce template has no Barrier type parameter, so the batched
path must omit tl::SyncThreadsBarrier for ROCm targets. Adds
TargetIsRocm() checks in both reduce.cc and finalize_reducer.cc.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Fix] Use target-aware warp size for cross-warp reduction threshold
Address CodeRabbit review: ROCm wavefronts are 64-wide, so the cross-warp
threshold should be 64 (not 32) on HIP targets. This avoids unnecessary
workspace allocation and batched codegen for single-wavefront reductions.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Style] Fix ruff format for test assertion
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Bugfix] Support batched AllReduce in CuteDSL Python codegen
The CuteDSL codegen translates C++ AllReduce template calls to Python.
With the new batch_size and workspace_stride template params, the
generated Python calls AllReduce() with 7 args instead of 5. Add
batch_size and workspace_stride params to the Python AllReduce function
and implement batched reduction logic for both standard and Hopper paths.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Feature] Add explicit batch parameter to T.reduce for user-controlled batched AllReduce
Expose a `batch` parameter on T.reduce (and all T.reduce_* wrappers,
default 1 = scalar / current behaviour). When batch > 1 the compiler:
- validates batch divides the per-thread output element count N;
- emits ceil(N/batch) calls to the new AllReduce::run_batch interface,
each sharing a single pair of barriers across all batch elements;
- allocates workspace of size reducing_threads * batch (only for cross-
warp reduce, i.e. reducing_threads > 32).
Key differences from the existing auto-detection approach:
* The user controls when batching is enabled — no silent smem increase.
* workspace_stride = reducing_threads (not block_size), so workspace is
minimal and does not blow past the GPU shared-memory limit.
* run_batch (not run) avoids C++ overload-resolution ambiguity when a
pointer is passed as the first argument.
Files changed:
- src/op/reduce.h: add `int batch{1}` field to ReduceOpNode
- src/op/reduce.cc: read "batch" annotation; validate; emit N/batch
batched AllReduce calls with address_of per-chunk pointers
- src/tl_templates/cuda/reduce.h: rename batch run() -> run_batch(),
fix recursive call
- src/tl_templates/hip/reduce.h: same
- tilelang/language/reduce_op.py: add batch param to reduce() and all
reduce_max/min/sum/absmax/abssum/bitand/bitor/bitxor wrappers
- testing/python/language/test_tilelang_language_reduce.py: update
test_batched_allreduce_codegen to use explicit batch=2 and match
run_batch in generated source
* [Fix] Use run_batch in finalize_reducer batched AllReduce codegen
finalize_reducer.cc also generates batched AllReduce calls (used by
T.alloc_reducer / T.finalize_reducer), but was still emitting ::run
instead of ::run_batch, causing a compiler error because the batched
pointer-based interface is now named run_batch to avoid overload
ambiguity.
* rocm fix
* [Feature] Add explicit batch parameter to T.finalize_reducer
Mirrors the batch= API already available on T.reduce:
- batch=1 (default) keeps the existing scalar AllReduce::run path
- batch>1 emits AllReduce::run_batch, sharing barriers across batch
output elements and reducing barrier count by batch×
C++ changes:
- FinalizeReducerOpNode gains an int batch{1} field (registered via
reflection) and reads it from the 'batch' annotation in the constructor
- Lower() validates batch against the layout output-element count and
uses effective_batch to drive use_batch / codegen
Python changes:
- finalize_reducer(reducer, batch=1) – new optional argument
- batch < 1 raises ValueError (consistent with T.reduce)
Tests: add test_batched_finalize_reducer_codegen covering sum/max/min
with explicit batch values, verifying run_batch is emitted with the
correct template argument.
* [Test] Add comprehensive tests for T.finalize_reducer batch parameter
Covers all cases introduced by the explicit batch= API:
Codegen tests:
- batch>1 emits AllReduce::run_batch with correct template argument
- batch=1 (default) does NOT emit run_batch (scalar path)
Correctness tests:
- batch=1 scalar path: sum/max/min across float32/float16 (PASS)
- batch==block_M full-batch: correct by coincidence (PASS)
- batch>1 partial-batch: xfail(strict=True) – known bug where
run_batch receives buffer->data directly but fragment layout is
non-contiguous per thread; root cause documented in
work/bug-finalize-reducer-batch.md
Error-input tests:
- batch=0 / batch<0 raise ValueError at Python layer
- batch > total output elements raises at compile/lower time
- batch that does not evenly divide output elements raises
* [Test] Rewrite finalize_reducer tests to follow clean outline
Replace the previous over-engineered test suite (4 case lists, 2 split
correctness functions, per-case xfail) with a minimal set organised by
the three categories that matter:
1. Codegen – single parametrised test covers both batch=1 (no run_batch)
and batch>1 (run_batch with correct template arg), reusing one case
list for both codegen and correctness.
2. Correctness – only batch=1 (scalar path), which is currently correct.
batch>1 correctness tests are omitted until the underlying fragment
layout bug is fixed (tracked in work/bug-finalize-reducer-batch.md).
3. Error inputs – batch<1 collapsed into one parametrised test; compile-
time errors (batch>output, batch not divisible) each get one test
using a shared _make_error_kernel helper.
* [Test] Rewrite reduce tests following clean outline
Replace fragmented per-function tests with four parametrized suites:
- test_reduce: op × dtype × src_scope × dst_scope × threads, covers
fragment→fragment, shared→shared, fragment→shared for sum/max/min/
abssum/absmax with float32/float16/int32/int64
- test_reduce_clear: op × src_scope × dst_scope with clear=False,
verifies pre-filled dst is accumulated correctly
- test_reduce_batch_codegen: verify batch>1 emits AllReduce::run_batch
- test_finalize_reducer_codegen/correctness/invalid_batch: batch
parameter feature tests, sharing one FINALIZE_REDUCER_CASES list
Use tilelang.jit lazy style and torch.testing.assert_close directly,
no profiler wrapper.
* [Test] Merge batch codegen+correctness into test_reduce
Add batch column to REDUCE_CASES; batch>1 cases verify both that
run_batch appears in generated source and that numerical output matches
reference, within the same test_reduce parametrization.
Remove the now-redundant test_reduce_batch_codegen.
* [Test] Add more batch>1 reduce cases covering all ops and dtypes
Add batch=2/4/8 cases for sum/max/min/abssum/absmax across float32/
float16/bfloat16, with both fragment and shared src/dst scopes.
Relax tolerance to 1e-1 for float16/bfloat16 to account for larger
rounding error in multi-element reductions.
* [Test] Limit int64 input range to avoid ROCm shfl_xor 64-bit bug
* Revert "[Test] Limit int64 input range to avoid ROCm shfl_xor 64-bit bug"
This reverts commit 9d60d6b197.
* fix(HIP): skip __launch_bounds__ for sub-wavefront threads to avoid miscompile of 64-bit shuffle ops
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: kurisu6912 <kurisu6912@users.noreply.github.com>
* Fix HIP codegen for sync_warp, sync_grid, and local.var initialisation
* [AMD/HIP] Fix warp_reduce VGPR bug, ShuffleNode packing, and Pipelined LDS overflow
Extends PR #2096 with three additional fixes for CDNA (MI350) targets:
Fix 1 — src/tl_templates/hip/reduce.h: warp_reduce width=32
The old 6-step butterfly called __shfl_xor(value, 32) without a width
argument. On CDNA (wave64) with 32 active threads, lanes 32-63 are
inactive and hold uninitialised VGPRs, producing NaN in reduce_max /
reduce_sum / AllReduce. Fix: remove the step-32 shuffle; pass width=32
to all remaining 5 steps so every shuffle stays within the [0,31] group.
Fix 2 — src/target/codegen_hip.cc + src/tl_templates/hip/common.h:
ShuffleNode bfloat16x2 / float16x2 packing
CodeGenC emitted `uint1(a, b)` for bfloat16x2 construction, which is an
invalid HIP constructor call. Fix: override VisitExpr_(ShuffleNode) in
CodeGenTileLangHIP to emit `uint1{__pack_bfloat162(a, b)}` / `uint1{
__pack_half2(a, b)}` using aggregate initialisation. Also add five
bfloat16x2 math overloads for uint1 carrier (abs2/max2/min2/add2/mul2).
Fix 3 — src/transform/pipeline_planning.cc: skip T.Pipelined(num_stages>1)
Double-buffering doubled LDS per loop-body buffer. On CDNA (≤128 KB LDS
per workgroup), this caused hipModuleLaunchKernel EINVAL. Fix: when
TargetIsRocm() && num_stages > 1, skip pipeline planning and fall back
to a plain sequential loop with synchronous T.copy.
Also: fix __habs(hip_bfloat16) and __habs(float16_t) in common.h to use
__builtin_memcpy instead of reinterpret_cast to avoid strict-aliasing UB
(as flagged by CodeRabbit on PR #2096).
Tests: 19 new cases added to testing/python/amd/test_tilelang_hip_codegen.py
covering all three fixes. All 42 tests pass on MI350 (gfx950).
* [AMD/HIP] Merge test_tilelang_hip_bugfixes.py into test_tilelang_hip_codegen.py
Consolidate all HIP regression tests into a single file. The merged file
covers all six fixes with 32 tests total (previously split across two files
with duplicated test cases for warp_reduce, pipelined GEMM, and ShuffleNode).
Changes versus the two individual files:
- Deduplicated test_warp_reduce_no_nan (identical in both files)
- Deduplicated test_pipelined_no_lds_overflow / test_pipelined_shared_mem_no_launch_error
- Deduplicated test_pipelined_multi_stage_fp16_gemm
- Merged bfloat16 shuffle tests: source check + runtime correctness in one function
- Kept PR #2096 source-inspection tests (alloc_var, sync_warp, sync_grid)
- Added runtime tests from bugfixes: inf init, serial loop accumulation,
float scalar readback, two-group wave64 reduce, float16 shuffle
* fixup: correct LDS size comment — gfx950 has 160 KB, not 128 KB
gfx942 (CDNA3 / MI300X) has 64 KB LDS per workgroup.
gfx950 (CDNA4 / MI350) has 160 KB LDS per workgroup (see PR #2058).
The old comment said '128 KB' which is wrong for both generations.
Updated pipeline_planning.cc and the test docstrings to reflect the
correct per-architecture limits.
* update for format checking
Updated the async pipeline logic to progressively relax loop waits and modify the trailing drain suffix. The test case was also adjusted to reflect changes in the expected behavior, ensuring that the pipeline maintains the correct number of groups in flight and descends through the drain suffix as intended.
Co-authored-by: wutong.1109 <wutong.1109@bytedance.com>
* add gfx950 new feature, 160K LDS and compy async
* update accoding to format check result
* fix a typo error
* add gfx950 copy async test codes
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* [Cleanup] Remove redundant int4 MMA code paths
* [CI] Apply format.sh updates
* lint fix
* fix
* Refactor FP8 type determination to utilize `T` module constants
Updated the `determine_fp8_type` function to return constants from the `T` module instead of string literals for FP8 formats. This change enhances consistency and aligns with recent updates in type handling. Adjusted return statements for both HIP and CUDA scenarios to use `T.float8_*` constants.
* Refactor FP8 type usage in tests to directly call `determine_fp8_type`
Updated test files to replace `T.dtype(determine_fp8_type())` with direct calls to `determine_fp8_type()` for improved clarity and consistency in FP8 type handling. This change aligns with recent refactoring efforts in type determination.
* Remove print statements from `assert_tl_matmul_correctness` and `test_assert_tl_matmul` functions to clean up test output. This change enhances the clarity of test results by eliminating unnecessary console output.
* Enhance CI workflow for CuTeDSL examples with Python 3.12 and CUDA-12.8
- Added a step to list generated files and clean up Python cache directories.
- Introduced a new job for running CuTeDSL examples on self-hosted NVIDIA runners.
- Set up environment variables for CUDA and Python, including caching configurations.
- Implemented error handling to clear cache on setup failure.
- Updated steps for installing dependencies and running examples, ensuring compatibility with the latest Python and CUDA versions.
* Remove redundant `map_torch_type` in favor of `T.dtype.as_torch()`
`map_torch_type` was a standalone function that converted dtype strings or
TVM DataType objects to `torch.dtype`. This was redundant since
`T.dtype.as_torch()` (and `T.dtype(x).as_torch()`) already provides the
same functionality with better handling (HIP/CUDA awareness, proper dtype
lookup tables, etc.).
Changes:
- Remove `map_torch_type` from `tilelang/utils/tensor.py` and its re-export
- Make `T.dtype()` idempotent for existing dtype objects (pass-through)
- Add backward-compat for `"e4m3fnuz_float8"` TVM internal name in `as_torch()`
- Replace all usages across production code, tests, and examples
- Rebuild Cython extension to reflect the import change
* Refactor output and accumulation data types in GEMM FP8 example
Updated the `out_dtype` and `accum_dtype` variables in the `example_tilelang_gemm_fp8_intrinsic.py` file to use `T.float32` instead of string literals. This change enhances type consistency and aligns with recent updates in type handling.
* ugrade tvm
* Remove upper version bound on apache-tvm-ffi
The <0.1.10 pin was introduced to avoid a derived_object regression,
which has since been resolved. Removing the cap allows compatibility
with newer versions of apache-tvm-ffi.
* Enhance ptxas usage level handling in CUDA compilation
Added type conversion for ptxas_usage_level in both tilelang_callback_cuda_compile and LibraryGenerator classes to ensure it is treated as an integer when specified. This change improves the robustness of the configuration handling for CUDA compilation.
* Fix IR string representation issues from tvm-ffi upgrade
The new tvm-ffi version changes str() output for IR nodes from compact
script format to verbose repr format. This caused two issues:
1. T.call_packed("name") is now tir.tvm_call_packed with value="name"
2. tir.Var str() now produces "tir.Var(span=None, ...)" instead of name
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* Simplify TMA descriptor globalAddress handling in cutedsl wrapper
Since global_address is now converted to string by pythonic_expr_func
in parse_tma_descriptor_args, remove the unnecessary tir.Var type check
and .name extraction.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* Refactor CUDA test for vectorization and update IR string representation
Updated the test for vectorization in CUDA by replacing the direct function definition with a JIT-compiled kernel. Additionally, modified the string representation of IR modules in multiple tests to use the script format instead of the default string format, ensuring consistency and clarity in assertions.
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
* Support int4 T.gemm and direct packed cp.async lowering
* Keep cp.async vectorization as a late fallback
* Change tl.ptx_cp_async to use num_elems semantics
* fix test
* Refactor cache namespace layout
* Update atomic cache tests for namespaced paths
* Ensure namespaced cache dirs exist before save
* Use auto backend in grouped gemm ptr example
* Keep grouped gemm ptr example off cutedsl
* Always use auto backend for grouped gemm ptr
* Skip grouped gemm examples in cutedsl CI
* Ignore grouped gemm examples in cutedsl workflow
* [BugFix] Keep shared-prelude local vars in producer-consumer WS
Preserve prelude definitions needed by shared LetStmt values during warp-specialized pipeline rewriting, and avoid cloning local.var buffers into the producer branch. This keeps grouped GEMM batch index calculations shared and fixes incorrect TMA batch selection.
* add test
* lint
* [Feature] Add Blackwell Cluster Launch Control (CLC) primitives
Expose cluster launch control query operations (try_cancel,
try_cancel_multicast, is_canceled, get_first_ctaid_{x,y,z}) in both
the C++ codegen path and the Python tilelang.language API, plus a new
mbarrier_arrive_expect_tx builtin used with CLC multicast completion.
Enables writing persistent SM100 kernels that use the hardware scheduler
to fetch the next tile coordinate.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* [Example] Add SM100 persistent 2-CTA GEMM example using CLC scheduling
Adds examples/gemm_sm100/gemm_tcgen5mma_ws_clc.py, a warp-specialized
persistent GEMM for Blackwell (sm_100) that uses tcgen05 2-CTA MMA
with cluster launch control tile scheduling. The kernel is split into
producer/MMA/scheduler/epilogue warp groups, with the scheduler warp
running on both CTAs and feeding next-tile coordinates via the CLC
hardware multicast into schedule_arrived directly (no intermediate
barrier). Storage sync is disabled and all syncs are placed manually.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* lint
* typo
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
* remove unnecessary test
* Add CLC query intrinsics and async fence
* Refactor CLC response handling by introducing CLCResponseDecode struct. Update clc_get_first_ctaid_x/y/z and clc_is_canceled functions to utilize the new decode response method for improved clarity and maintainability.
* Enhance CLC scheduling in persistent GEMM by introducing shared variables for schedule validity and tile ID. Update cancellation checks to utilize the new shared variables, improving clarity and maintainability of the scheduling logic.
* test fix
---------
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* [Enhancement] Use atomic directory rename for cache writes
Replace per-file os.replace with write-to-staging + os.rename to
ensure other processes never see an incomplete cache directory.
Also add cleanup of stale staging dirs on init.
* [Bugfix] Fix atomic cache save: abort on critical file failure and fix staging cleanup
- Let kernel lib and params save failures propagate instead of swallowing
them, so the staging directory is cleaned up rather than renamed as an
incomplete cache entry.
- Move autotuner staging directory to TILELANG_CACHE_DIR so that
_cleanup_stale_staging_dirs() can find and remove stale entries.
* Fix atomic cache save recovery and error handling
* Remove unnecessary blank lines in _FakeAdapter and _FakeKernel classes in autotune and cache test files.
* Tighten atomic cache completeness checks
* Simplify atomic cache save error handling
* submodule update
* remove legacy test
---------
Co-authored-by: Freebase6912 <amid-gauze-racing@duck.com>
* [Transform] Add InjectTcgen05Fence pass for Blackwell (SM100+)
On Blackwell GPUs, the tcgen05 accumulator (TMEM) resides in a separate
address space that is not synchronized by regular thread barriers like
__syncthreads() or mbarrier. Two PTX fence instructions are required to
ensure cross-thread visibility of TMEM state:
tcgen05.fence::before_thread_sync -- flush TMEM before barrier
tcgen05.fence::after_thread_sync -- pull TMEM after barrier
This commit introduces the `InjectTcgen05Fence` TIR pass that
automatically wraps every `tvm_storage_sync("shared")` call with the
fence pair when the target is SM100+ and the function uses tcgen05/TMEM
operations.
Changes:
- Define two new TIR intrinsic Ops: `tcgen05_before_thread_sync` and
`tcgen05_after_thread_sync` (builtin.h/cc)
- Add codegen support to emit `tl::tcgen05_before_thread_sync()` and
`tl::tcgen05_after_thread_sync()` (codegen_cuda.cc)
- Implement the `InjectTcgen05Fence` pass (inject_tcgen05_fence.cc)
- Register the pass in the Python transform module (__init__.py)
- Insert the pass in OptimizeForTarget after ThreadSync (phase.py)
* fix
* Add conservative tcgen05 fence injection
Extend InjectTcgen05Fence to cover shared storage sync and linear wait/use, use/arrive handoffs on SM100+ kernels, update the surrounding docs, and add transform coverage. Also lower tcgen05 ld/st copies through dedicated intrinsics instead of call_extern strings, with matching CUDA and CuTeDSL codegen support.
* Use tcgen05 fence injection in SM100 examples
Remove the hand-written tcgen05 before/after thread-sync fences from the ws and persistent GEMM examples now that InjectTcgen05Fence covers these linear wait/use and use/arrive handoffs.
* lint fix
* Drop tcgen05 call_extern fence compatibility
Simplify InjectTcgen05Fence to recognize only the lowered tcgen05 intrinsics seen after LowerTileOp, and update the transform tests to construct tcgen05_ld directly instead of relying on legacy call_extern strings.
---------
Co-authored-by: Freebase6912 <amid-gauze-racing@duck.com>
Lei Wang