* [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>
* [Bugfix][Subtype] Fix scalar fp4 load/store codegen and dynamic stride inference for sub-byte types
The scalar fp4 buffer access path in codegen_cuda.cc used GetBufferRef which
applies a /2 index division (mapping two fp4 elements to one byte), then
performed a plain byte assignment — destroying the neighboring nibble. This
caused incorrect results whenever fp4 stores couldn't be vectorized (e.g.
with dynamic/non-constant strides in StridedTensor).
Fix: use tl_fp4_packed_load/tl_fp4_packed_store with (fp4_e2_2_t*) cast for
scalar fp4 accesses to non-packed buffers, which correctly handles nibble-level
read-modify-write via set_x/set_y.
Also add stride_scale factor to _process_dynamic_symbolic in the cython and
tvm_ffi adapters, so torch strides (in storage units) are correctly converted
to logical element strides for sub-byte dtypes.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* [Test][Subtype] Add correctness tests for fp4 dynamic-stride store codegen
Add regression test that scatters fp4 rows into a non-contiguous
StridedTensor with dynamic strides and verifies byte-for-byte match
against a static-stride reference. Parametrized over block_size.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* [Enhancement][Subtype] Inject stride divisibility assume for sub-byte dtypes and add scalar fp4 store codegen test
- Extend InjectAssumes pass to emit `assume(stride % pack_factor == 0)` for
non-last-dimension strides of sub-byte buffers (e.g. fp4 with pack_factor=2).
This provides the TIR analyzer with divisibility information before LetStmt
inlining.
- Add `test_subtype_fp4_scalar_store_codegen`: scatter fp4 elements via
indirection to force the scalar (non-vectorized) store codegen path and
verify nibble-level correctness.
- Convert new tests to eager-style @tilelang.jit.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* style: apply clang-format and ruff formatting
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* fix lint error
* fix: unpack 4-tuple in cython_wrapper dynamic_symbolic_map lookup
---------
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* draft
* fix lint
* fix lint
* Improve pass comments
* treat loop-invariant memory access as scalar in vectorization
* Add unit test
* fix lint
* fix b[i] = a[i] + a[i+32] case
* Minor fix
* Tighten DecoupleTypeCast cast detection and add RMW regression test
Skip BufferLoad/BufferStore indices when searching for Cast nodes so an
index-type conversion does not spuriously trigger the decoupling
transformation. Clarify the load-replacement table in visit_for_ — store
entries must feed into it so RMW loads map to the store-side cast buffer —
and cover the a[i] = a[i] + a[i+32] case with a regression test.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Keep loop-invariant stores in memory bucket during vectorize planning
Commit 03bb0706 diverted all loop-invariant global/shared accesses into
the local_fragment bucket, but that bucket's constraint is dropped by the
has_global_or_shared_buffer strategy. ComputeBufferVectorSize already
returns 1 for a reduction-like store such as shared[tx] += a[...+j], yet
that 1 was silently lost, so vectorization proceeded with vector_size=2
and emitted two scalar writes to the same shared[tx] — the second
clobbered the first and dropped a lane of the accumulation.
Only loop-invariant loads (genuine broadcast reads) are safe to divert;
stores must stay in the memory bucket so their vector_size=1 constraint
is honored and the loop is left scalar.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Remove issue 1106 test
Each thread only touches its own shared[tx] slot, so no __syncthreads
is actually required and the test asserts an over-conservative behavior.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Reorder the `mask` parameter to the end with a default of `_FULL_WARP_MASK`
for shfl_xor/down/up/sync, any_sync, all_sync, ballot_sync, and
match_any/all_sync. Full-warp is the overwhelmingly common case, so making
mask an optional trailing kwarg removes boilerplate at call sites and lets
old `(value, delta)`-style callers keep working. Update the test suite and
docs/instructions.md to match the new signatures.
* [Feature] Add TIR builtins for warp-level vote and block-level predicate sync
* [Refactor] Lower warp-level intrinsics through tl.* ops
Convert the warp-vote, warp-shuffle, predicated block-sync, and warp-match
builtins from raw `tir.call_extern` wrappers to proper TIR ops registered
under `tl.*` and lowered in `codegen_cuda` / `codegen_hip`. This removes
Python-side `_IS_HIP_AVAILABLE` branching and pushes the CUDA/HIP split
into codegen, where it belongs.
* Register new tl ops: any_sync, all_sync, ballot_sync, ballot,
activemask, syncthreads_count/and/or, shfl_sync, shfl_xor_sync,
shfl_down_sync, shfl_up_sync, match_any_sync, match_all_sync.
* Codegen lowering on both CUDA and HIP. uint32->uint64 zero-extension
for ballot/activemask now happens in codegen. HIP drops the mask
argument for shfl/vote and emits LOG(FATAL) for match_*_sync.
* `__match_all_sync`'s int* pred output is hidden behind an
immediately-invoked lambda so the wrapper stays expression-form.
* Python wrappers normalize Python int masks to uint32 TIR consts via
`_as_uint32_mask`, so the emitted source prints `(uint)0xFFFFFFFF`
instead of `(int64_t)4294967295`.
* Unify `shfl_xor/down/up` signatures to `(mask, value, delta, width=32)`,
matching CUDA convention and `shfl_sync`. No in-tree callers existed.
* Add tests for match_any_sync and match_all_sync.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* register op, and add to build file
* update semantics of NaN for Max and Min ops in Reduce
* lower Max op to cutlass fastmax in general case, _hmax for bf16 and _hmax_nan for hf16
* update semantics of NaN for Max and Min ops in Reduce in reduce.h
* develope codes transforming ReduceOpNode in TIR and handling CallNode in CUDA code generator's runtime template
* test file for reduce maxmin nan
* refine test case message and explanation“
”
‘
“
* invert the senamtics of global config TL_REDUCE_MAXMIN_NAN_PROPAGATE
* lint fix
* refactor reduce maxmin NaN propagate into per-call annotation
Replace the global tl.reduce_maxmin_nan_propagate pass-config with a
per-op nan_propagate flag carried on the ReduceOp annotation map.
Frontend exposes it as a kwarg on T.reduce_max/min/absmax (default
False, preserving prior behavior). The CUDA codegen MinNode/MaxNode
visitors are reverted to plain __hmin/__hmax; only the new tl::max_nan
/tl::min_nan CallNode handler emits __hmax_nan/__hmin_nan, so non-reduce
Min/Max ops are no longer silently rewritten. ReduceOp::Lower now errors
early with a clear message if nan_propagate=True is requested for fp16
/bf16 on a non-CUDA target, instead of emitting an undefined symbol on
HIP/CPU. Tests rewritten to parametrize on the kwarg and assert runtime
NaN-propagation behavior, not just generated source substrings.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: Haoran <chaoran@umich.edu>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* [Refactor] Remove GEMM v1 and merge gemm_py into the canonical gemm op
GEMM v2 has been the default for a while and no tests or examples still
depend on v1. This removes the legacy C++ GemmNode path entirely and
promotes the former GemmPyNode (Python-delegated lowering) to be the one
and only GEMM tile op.
C++:
- Merge src/op/gemm_py.{h,cc} into src/op/gemm.{h,cc}. GemmPyNode becomes
GemmNode, GemmPy becomes Gemm, FFI type key "tl.GemmPy" becomes "tl.Gemm",
and the op names "tl.tileop.{gemm_py,wgmma_gemm_py,tcgen05_gemm_py}"
collapse to "tl.tileop.{gemm,wgmma_gemm,tcgen05_gemm}". The shared
infrastructure (GemmWarpPolicyNode, GemmInst enum,
GemmWarpPolicyNode::computeWarpPartition) is preserved in place so
GemmSP/GemmSPPy keep compiling unchanged.
- Delete the tl_gemm builtin, its codegen dispatch in
codegen_{cuda,hip,cutedsl}.cc, and its matching branches in
inject_fence_proxy.cc and merge_shared_memory_allocations.cc. tl_gemm_sp
is kept as the remaining async-proxy gemm builtin.
- Drop #include "gemm_py.h" and GemmPyNode/GemmPy references from
inject_pipeline.cc, producer_consumer_ws.cc, and lower_blackwell_2sm.cc.
Python:
- tileop/gemm: rename GemmPy -> Gemm, register as "tl.Gemm", and rename
the FFI callbacks tl.gemm_py.{lower,infer_layout} -> tl.gemm.{lower,
infer_layout}. _ffi_api.GemmPyGemmInst becomes GemmGetGemmInst.
- language/gemm_op.py: delete gemm_v1/gemm_v2 wrappers. gemm() now emits
tl.tileop.gemm directly; wgmma_gemm/tcgen05_gemm drop the _py suffix.
- env.py: remove TILELANG_USE_GEMM_V1 and the use_gemm_v1() helper.
- ir.py: drop the empty v1-era `class Gemm(Node, Scriptable): ...`
placeholder that was colliding with the real tileop.gemm.Gemm wrapper
under the "tl.Gemm" FFI type key.
- contrib/cutedsl: delete cutedsl/gemm_v1.py (runtime helper for the v1
cutedsl lowering path) and its re-export in cutedsl/__init__.py.
Tests / maint:
- testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py switches
from T.gemm_v2 to T.gemm.
- testing/python/transform/test_tilelang_transform_inject_fence_proxy.py::
test_lower_fence_proxy is rewritten to exercise tl.tl_gemm_sp instead
of the removed tl.tl_gemm, preserving the fence-proxy assertion.
- Delete maint/gemm_v2/ (v1-vs-v2 benchmarking scripts) and drop the
v1-specific comment from maint/scripts/run_local_ci_test.sh.
C++ headers under src/tl_templates/{cuda,hip,cpp}/gemm*.h are left in
place as reference implementations; they are no longer reachable from
the generated code path since the only consumer (tl_gemm builtin) is
gone, but keeping them avoids a large, hard-to-reverse deletion.
* [Fix] Restore cutedsl/gemm_v1.py (Gemm_SM80 / Gemm_SM90 helpers)
The previous commit deleted cutedsl/gemm_v1.py on the assumption that it
was the CuteDSL lowering helper for the removed T.gemm_v1 path. It is
actually a runtime helper module that defines SM80/SM90 CUTE DSL
primitives (make_smem_layout_AB, get_tma_atom, etc.) which are emitted
into generated cutedsl kernel code via
`import tilelang.contrib.cutedsl as tl; tl.Gemm_SM90.get_tma_atom(...)`
(see jit/adapter/cutedsl/wrapper.py CUBIN_TMA_ATOM_INIT_TEMPLATE).
The "v1" in the filename refers to the CuteDSL helper version, not to
the deleted T.gemm_v1 op, so removing it broke the cutedsl code path for
attention_sink, blocksparse_attention, blocksparse_gemm, and deepgemm
examples with `module 'tilelang.contrib.cutedsl' has no attribute
'Gemm_SM90'`. Restore the file and its `from .gemm_v1 import *` wildcard
re-export.
* [Refactor] Restore tl_gemm builtin as a call_intrin path to tl:: templates
The previous commit in this PR deleted the `tl_gemm` builtin entirely on
the argument that no code emitted it anymore. That was too aggressive:
keeping the builtin preserves a deliberate extension point for any
Python-side gemm lowering backend (for example a future
GemmLegacyExtern) to delegate into the existing C++ templates in
src/tl_templates/<target>/gemm*.h via
T.call_intrin("handle", tir.op.Op.get("tl.tl_gemm"),
op_instance_str, A_ptr, B_ptr, C_ptr)
without any extra C++ plumbing. It is also the only thing that keeps
those tl_templates headers reachable rather than dead code.
Restore:
- tl_gemm builtin declaration in src/op/builtin.h (with an updated
docstring describing the call_intrin entry point) and its
TIR_DEFINE_TL_BUILTIN registration in src/op/builtin.cc.
- Codegen dispatch branches in codegen_cuda.cc, codegen_hip.cc, and
codegen_cutedsl.cc that turn a tl_gemm Call into a PrintCallExtern
emission of `tl::gemm_xx<...>(A, B, C)`.
- SMEM alignment and async-proxy fence handling for tl_gemm in
merge_shared_memory_allocations.cc and inject_fence_proxy.cc. These
passes are tied to the builtin's semantics (not to the removed v1
GemmNode) and are still correct the moment anything emits tl_gemm.
- test_lower_fence_proxy in test_tilelang_transform_inject_fence_proxy
reverts to exercising `tl.tl_gemm` instead of the temporary
`tl.tl_gemm_sp` substitution introduced when the builtin was gone.
Note that the current sole gemm lowering path (the merged GemmNode that
delegates to Python via tl.gemm.lower) still emits inline mma/wgmma/
tcgen5mma via the Python macro generators, not tl_gemm. So restoring
the builtin is a no-op for the default path - it only re-enables the
"call into v1 templates from Python" option for anyone who wants it.
* [Fix] Default HIP execution backend to tvm_ffi, not cython
resolve_execution_backend() only promoted cuda and metal to tvm_ffi for
the "auto" path, so hip fell into the else branch and picked cython. But
allowed_backends_for_target() already lists tvm_ffi as an allowed
backend for hip, and the cuda default is tvm_ffi for good reasons:
- The cython SourceWrapper path in src/tilelang/jit/adapter/wrapper.py
has known gaps (for example T.alloc_global is not handled in
create_dispatch_func, which breaks any test that uses an internally
allocated global scratch buffer).
- libgen.py appends user-supplied compile_flags verbatim to the
hipcc/nvcc command. nvcc accepts --use_fast_math, clang++/hipcc does
not, so a CUDA-style flag leaks into HIP compilation.
Neither of these leak through the tvm_ffi backend. Aligning hip's
default with cuda makes the two paths symmetric and unbreaks
test_alloc_global and test_all_attrs_together_lazy on ROCm without
having to mark the tests skip or rewrite the cython wrapper.
cython is still available on hip via execution_backend='cython' for
anyone who explicitly wants it.
* [Test] Gate CUDA-only language tests with requires_cuda
Three language tests hardcode target='cuda' (or rely on CUDA-only
codegen behavior) but lacked a CUDA availability marker, so they failed
on the ROCm CI runner with "Cannot find global function
target.build.tilelang_cuda" once the earlier pre-existing failures got
past --maxfail=3. Add @tilelang.testing.requires_cuda to:
- test_tilelang_if_range
- test_annotate_min_blocks_per_sm_launch_bounds
- test_tilelang_transpose and test_tilelang_transpose_square
Also replace the hand-rolled __main__ block in test_tilelang_transpose
with tilelang.testing.main() so the skip markers are honored when the
file is run directly.
* [Test] Delete tile library gemm test
Remove testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py.
