* [Backend] Refactor Pipeline to support different backends
Introduce Pipeline abstraction in backend/pipeline.py with per-backend
registration. Each backend (cuda, hip, c, llvm) now registers its own
compilation pass pipeline. engine/lower.py resolves the pipeline from
the target instead of hardcoding phase imports.
* enhance the structure
* fuse two stage
* fix after rebase
* update template pass pipeline
* pipeline
* fix test
* fix test
* fix cuda pipeline prologue
* Add float4_e2m1_unpacked dtype
* Fix TIR cast for float4_e2m1_unpacked dtype.
Route custom[float4_e2m1_unpacked] through Float4E2M1Unpacked FFI
so T.float4_e2m1_unpacked(0.0) works in prim_func builders.
* Limit float4_e2m1_unpacked vector lanes and fix TMA transaction bytes.
Cap unpacked FP4 vector dtypes at x32 and count TMA/mbarrier
transactions using 4-bit payload size (0.5 byte per element).
* lint
* [Refactor] Extract shared-access analysis into a reusable header
- Move AccessEntry/AccessType definitions from thread_storage_sync.cc to
src/transform/common/shared_access_analysis.h so other passes (notably
MergeSharedMemoryAllocations) can reuse the same precise access
representation.
- Add tl.debug_thread_storage_sync_summary pass-config option for optional
ThreadSync planner summary logging.
- Add regression tests covering explicit shared sync deduplication and
shared.dyn alias-aware dependency planning.
* [Feature] Rewrite MergeSharedMemoryAllocations with precise access analysis
- Build per-access live ranges using shared_access_analysis utilities so
branch-alternative dataflow and barrier-driven epoch boundaries are
modeled directly instead of through whole-buffer approximations.
- Plan arena layout starting from arena_size_const so allocations whose
starting offsets are non-zero no longer collide.
- Preserve single-slot semantics for mutually exclusive branch writes
followed by a post-dominating read (fixes the K_norm_pairs_shared style
miscompile where stage-1 alternatives were split across epochs).
- Cover BufferRegion block reads/writes, tvm_access_ptr, tl.access_ptr,
cp_async branches, loop-carried buffers and partial-subregion liveness
cases.
- Keep aggressive-merge disabled under warp specialization and retain
whole-buffer liveness as the conservative default.
- Add 18 transform regressions in
test_tilelang_transform_merge_shared_memory_allocations.py covering
basic merge, sync, branch, cp_async, BufferRegion and opaque-pointer
paths.
* [Cleanup] Remove dead FreeList/FreeBlock/ActiveInterval helpers
These helpers were superseded by the segment-aware LinearScanPack
implementation that performs interval-overlap walking directly. They
have no remaining callers after the rewrite; remove them to reduce
maintenance burden.
* [Test] Cover tail-allocation arena growth scenario
Add a regression case mirroring upstream b4c913be (#2106): when the only
free slot reusable by an incoming buffer sits at the arena tail and is
smaller than the request, the arena must grow by the deficit only, not
push the new buffer past the freed slot.
The legacy free-list-based packer would emit a 384 B arena for this
pattern (X=128, Y=64-then-freed, Z=192), bumping Z to offset 192. The
interval-based LinearScanPack reuses Y's offset and grows the tail to
fit Z, producing a 320 B arena -- a 16.7%% reduction that matches the
peak-live lower bound.
* [Feature] Add EpochGraph and use per-epoch liveness to relax aliasing
Introduce src/transform/common/epoch_graph.h, a control-flow-aware
builder that partitions a flattened StmtEntry sequence into hard-sync
and cp_async-wait epochs across nested For/IfThenElse scopes, exposing
loop-carry, branch-join and parent-edge structure required for sound
liveness analysis.
Wire MergeSharedMemoryAllocations to compute per-buffer live epochs as
the intersection of:
* forward def-reach: epochs reachable from any write of the buffer
along epoch-graph successors;
* backward use-reach: epochs that can reach any read of the buffer
along predecessors.
A buffer is live in epoch e iff e is in both sets, modeling 'data
written upstream is consumed downstream' rather than the legacy
[start_stmt, end_stmt) bounding box. The result feeds two relaxers:
* LinearScanPack: two intervals whose live-epoch sets are disjoint
may share an offset even if their stmt ranges overlap.
* PlanMemory: the symmetric overlap probe consults the same set
before forcing a non-aliased layout.
Hard syncs (kSync) and cp_async waits are treated as epoch boundaries;
ptx_cp_async_wait is conservatively classified as a hard barrier.
IfThenElse joins take the union of branch live sets so a buffer
written in only one branch remains live at the join.
The whole-buffer EpochGraph + liveness pipeline always runs; verbose
[MSMA-EPOCH-GRAPH] / [MSMA-EPOCH-LIVE] / [MSMA-EPOCH-ACCESS] /
[MSMA-EPOCHSEG] / [MSMA-STAGE0-RAW] dumps remain gated behind the
existing tl.debug_merge_shared_memory_allocations pass-config flag.
Also remove SplitLocalSegmentsAtEpochSyncs, a Phase-1 helper that was
defined but never wired into the planner; the new liveness pipeline
subsumes its role.
Validated on a real lowbit MHA decode kernel
(seq_lens=[(3,96),(2,64)], heads=4, dim=64, page_size=32, bits=4):
arena drops from 35,232 B (upstream baseline) to 31,136 B, an 11.6%%
reduction driven by K_shared and V_shared sharing one offset thanks
to disjoint live-epoch sets.
Add testing/python/transform/test_tilelang_transform_epoch_graph.py
covering the five S1.4 epoch-graph topology cases (sequential
__syncthreads, ptx_wait_group, For body/back/exit edges, IfThenElse
branch_join, nested For). Existing
test_tilelang_transform_merge_shared_memory_allocations.py and
test_tilelang_transform_thread_sync.py suites continue to pass.
* [Feature] Gate per-epoch alias relaxer on annotated layout equality
Two shared-memory buffers carrying distinct layouts under
`block.annotations[attr::kLayoutMap]` (e.g. via `T.annotate_layout(...)`)
must not share storage even when per-epoch dataflow proves their
liveness sets disjoint: downstream layout-driven access lowering can
emit indexing arithmetic that is only sound under each buffer's
original layout, so naively aliasing them risks silent miscompiles.
Add a `LayoutSignatureCollector` that harvests the per-Var `Layout`
entries from `Block` annotations and a `LayoutAliasIncompatible`
helper that returns true iff both candidates have non-trivial layouts
and the layouts are not structurally equal. The two relaxer call
sites in `SharedMemoryRewriter` (`LinearScanPack` and the post-pack
overlap audit) consult the helper before dropping a conflict, so the
legacy decision is preserved exactly when the gate fires.
Buffers without a layout entry remain layout-agnostic and continue to
participate in arena packing, keeping the existing reuse for the
common un-annotated case. Verified that the lowbit reference kernel
arena footprints stay at 31136/61856/61808/61792 B and the merge /
thread-sync / epoch-graph trio passes 49/49.
* [Test] Add transform regression that fails on upstream MergeSharedMemoryAllocations
Add test_general_cond_branch_alternatives_share_offset_with_postdom_buffer
to lock the per-epoch / branch-alternative analysis behavior on a
mutually-exclusive if/else with a runtime cond plus a post-dominator
buffer after a sync. Upstream's linear-scan plan keeps the two branch
buffers at distinct offsets (128-byte arena); the per-epoch relaxer
collapses them onto the same offset and reuses it for the
post-dominator buffer (64-byte arena, 50% smaller).
* [Example] Demo MergeSharedMemoryAllocations branch-alternative arena reuse (64 vs 128 B)
Add runnable demo under examples/merge_shared_memory_allocations/
that constructs a three-buffer PrimFunc with mutually-exclusive if/else
branches, lowers through the merge pass, and reads the merged arena
size from the dyn_shared_memory_buf attr. The demo also compiles the
kernel end-to-end and validates output correctness with torch.allclose.
On this branch: arena = 64 B (saved 128 B / 66.7% vs naive 192).
On upstream main (1fbd994c): arena = 128 B (-50% worse).
* [Test] Add end-to-end torch correctness regression for merge pass
Add two tests (test_merged_branch_kernel_generates_correct_results,
test_merged_sync_phased_buffers_generate_correct_results) that compile
the merged kernel and verify bit-exact output via torch.allclose,
ensuring the arena packing does not alter numerical results.
* [Test] Add torch correctness and dynamic-buffer fallback regression
- test_merged_branch_kernel_generates_correct_results: compiles a
merged branch-alternative kernel with torch.allclose validation.
- test_merged_sync_phased_buffers_generate_correct_results: sync-phased
two-buffer kernel end-to-end check.
- test_dynamic_size_buffer_fallback_does_not_incorrectly_alias:
verifies that a symbolic-extent shared.dyn buffer correctly aliases
with a sync-separated constant-size buffer without unsafe overlap.
* [Fix] Disable per-epoch relaxer for loop-live buffers to avoid pipelined iteration overlap
* [Fix] Fallback to structural_segments_by_var for buffers missing planner segments to avoid buffer_byte_offsets crash and shared memory bloat
* Merge branch 'main' into merge_pass and fix merge comflicts
* [Fix] Resolve post-merge artifacts in CMakeLists, hipcc, and deepseek example
- Remove duplicate TILELANG_OUTPUT_TARGETS in CMakeLists.txt
- Fix duplicate _resolve_artifact_paths function in hipcc.py
- Add missing G=kv_group and fix assert in sparse_mla_fwd_pipelined.py
- Add ruff noqa pragma to suppress lint warnings
* [Style] Apply clang-format to modified source files
* [Pipeline] Refactor software pipeline transforms
* enhance
* Refine pipeline handling for guarded binds
* Refactor IfStmtBindingRewriter and related classes for improved access analysis
- Removed unused Target parameter from IfStmtBindingRewriter.
- Introduced IfStmtAccessCollector to handle buffer access collection.
- Updated access analysis methods to utilize the new collector.
- Cleaned up includes and removed obsolete functions in access_analysis.
- Enhanced pipeline utility headers with necessary includes for consistency.
* [Fix] Suppress spurious __syncthreads() between T.ws blocks (#2192) and add T.named_barrier_arrive (#2188)
Issue #2192 - BugFix: Unexpected CTA-wide __syncthreads() between T.ws producer/consumer
- Root cause: FindConflict() in thread_storage_sync.cc did not check for mutually
exclusive thread constraints. When T.ws(0) (tx<128) reads and T.ws(1) (tx>=128)
writes the same shared buffer, no thread satisfies both constraints, so no
__syncthreads() is needed — mbarriers already handle the synchronization.
- Fix: add mutual-exclusivity check (prove NOT(P AND C)) early in FindConflict().
If the access constraints are mutually exclusive (disjoint thread sets), return
false (no conflict) immediately, before any index-range analysis.
- Measured improvement: barrierpipe_stage2 gemm latency 44ms → 25ms on H200.
Issue #2188 - Feature: expose native T.named_barrier_arrive(barrier_id, thread_count)
- Adds tl.named_barrier_arrive intrinsic op (bar.arrive PTX, one-sided arrive).
- Registers op in builtin.cc, declares in builtin.h.
- Adds tl::__named_barrier_arrive<id, cnt>() template in common.h.
- Adds codegen case in codegen_cuda.cc.
- Exposes T.named_barrier_arrive() in tilelang/language/builtin.py.
- Adds regression test testing/python/issue/test_tilelang_issue_2188_2192.py.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* lint: clang-format + ruff fixes after #2188/#2192 changes
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* test: simplify issue 2188/2192 test, keep only #2192, add requires_cuda
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* [Fix] Make named_barrier_arrive accept variable barrier_id/thread_count
* Update documentation for named barriers and remove obsolete test file
- Clarified documentation for `named_barrier_arrive` in `builtin.h` and `builtin.py` to specify the barrier index range.
- Removed obsolete regression test for issues #2188 and #2192 as it is no longer needed.
* lint fix
---------
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* Optimize disk cache source loading
* Add tests for kernel cache frontend hit and improve memory cache handling
- Introduced a new test to validate the round-trip functionality of real primitive functions in the kernel cache.
- Enhanced the memory cache logic to prevent overwriting existing entries.
- Improved error handling when tagging kernel cache entries with logging for better debugging.
* [Cache] Refine lazy source restore
Guard memory-cache writes after compile, make cache tagging failures diagnosable, and document the lazy frontend-cache gate.
Wrap cached source text/path in CachedTextSource and add a real PrimFunc frontend-cache round-trip regression test.
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* feat(tma_copy): add thread_extent parameter for per-warp TMA copy
Add a `thread_extent` parameter to `T.tma_copy()` that controls the
elect size for TMA leader thread selection. Defaults to full block
size, can be set to warp_size (32) for per-warp TMA operations.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* refactor(tma_copy): rename leader election scope parameter
* remove: delete test_per_warp_tma.py and its associated tests; add per_warp_tma_copy_kernel and block_tma_copy_kernel tests in test_tilelang_language_tma_copy.py
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* [Bugfix] Fix issue 2123 access_ptr lowering
Add the issue 2123 regression test and keep the access_ptr base BufferLoad intact through safe-memory legalization and lowering.
Closes#2123
* Deduplicate access_ptr base load handling
Move shared tl.access_ptr BufferLoad base visitation into a common transform helper so LowerAccessPtr and safe-memory legalization stay aligned.
Extend the issue 2123 regression with a direct LowerAccessPtr pass check while keeping the LowerAndLegalize pipeline coverage.
* Format access_ptr helper header
* Narrow access_ptr helper scope
* Share access_ptr base helper
* Address access_ptr helper review comments
* [Transform] Preserve access_ptr buffer loads in safe memory legalization
* Refactor AccessPtrInfo return structure in legalize_safe_memory_access.cc
Updated the return structure of AccessPtrInfo to improve readability and maintainability. The logic for handling the tvm_access_ptr operation has been streamlined, ensuring clearer checks and consistent formatting of returned values.
* fix
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* add amd gfx950 mxfp4 support
* update the codes for format checking
* [AMD][CDNA4] Address review comments for MXFP4 gfx950 support
Four issues from CodeRabbit review fixed:
1. hip_fp4.h: Fix FP4 E2M1 denormal decoder round-trip bug.
The decoder returned 0.25f for mant==1 (exp==0) but the encoder
maps nibble 1 to 0.5f, breaking float→fp4→float round-trips.
Fix: return 0.5f to match the encoder and LUT tables.
2. codegen_hip.cc: Fix UB address-of-prvalue in FP4→float16 cast.
The generated code took &(__tl_cvt_fp4x2_to_half2(...)), which
takes the address of a temporary (C++ UB). Fix: materialize the
uint1 return value into a local variable first, matching the
pattern already used in the float32/bfloat16 cast paths.
3. codegen_hip.cc: Extend fp4_pair_cast to cover FP4x16 and FP4x32.
Previously only 2/4/8-lane FP4 casts were handled; x16/x32 fell
through silently to generic vector code that cannot index fp4_e2_*
aggregates. The pairwise uint8_t byte logic is correct for any
even lane count, so extend to include 16 and 32 lanes.
4. mxfp.py: Narrow bare except to expected import-time errors.
The catch-all except Exception silently swallowed all errors from
target_is_gfx950(), causing HIP targets to fall back to CUDA/PTX.
Fix: catch only ImportError/ModuleNotFoundError/AttributeError.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* update for format checking
---------
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* [Metal] Add Metal GEMM support with simdgroup_matrix MMA
Add T.gemm support for Apple Metal using simdgroup_matrix 8x8 operations
(simdgroup_load/store/multiply_accumulate). Works on all Apple Silicon
(M1-M5) without requiring a TVM fork.
Key changes:
- codegen_metal.cc/h: Fork TVM Metal codegen to tilelang with
simdgroup intrinsic emission and 128-bit vectorized copy
- gemm_metal.py: GemmMetal tile operator for sharedxshared GEMM
- metal_macro_generator.py: MPSIntrinEmitter for simdgroup MMA macros
- metal_fragment_to_simdgroup.py: Pass rewrites local.fragment GEMM
accumulators to metal.simdgroup scope before layout inference
- LowerSIMDGroupCopy in copy.cc for fragment->device simdgroup_store
24 Metal tests (codegen cross-platform + correctness on device).
* Move Metal buffer helpers to Metal backend
* fix(metal): remove layout_inference bypass, fix build & clarify docs
- Remove duplicate BuildTileLangMetal in rt_mod_metal.cc (fixes linker error)
- Remove TargetIsMetal bypass in layout_inference.cc. The bypass is
unnecessary because MetalFragmentToSimdgroup already converts all
GEMM accumulator fragment buffers to metal.simdgroup scope before
LayoutInference runs, so IsFragmentBuffer never matches them.
- Clarify docstrings for MetalFragmentToSimdgroup pass explaining why
it runs before LayoutInference (simdgroup matrices are opaque and
have no explicit thread-level layout).
- Restore is_metal_simdgroup check in GemmMetal.lower() so the GEMM
lowering correctly handles buffers already converted to simdgroup
scope by the pass.
* fix: tir->tirx migration build fixes for metal backend
- Replace tir::Buffer with Buffer, tir:: with tirx:: in metal backend
- Update codegen_metal to use AllocBufferNode instead of AllocateNode
- Create runtime/metal/metal_module.h with MetalModuleCreate factory
- Fix ICHECK -> TVM_FFI_ICHECK macros
- Fix 'from tvm import tir' -> 'from tvm import tirx as tir' in Python
- Restore tilelang/backend/__init__.py for metal gemm registration
- Add metal backend import to tilelang/__init__.py
- Fix metal_fragment_to_simdgroup.py: Block->SBlock, Allocate->AllocBuffer
* fix: improve metal_fragment_to_simdgroup buffer remapping
- Add AllocBuffer node handling for standalone allocations
- Remap buffer references in SBlock reads/writes/match_buffers
- Add _remap_buffer_region and _remap_match_buffer helpers
* lint
* fix: add full-body var subst in MetalFragmentToSimdgroup to reach gemm call_intrin args
The per-SBlock substitute() in _rewrite_scope missed Var references
buried inside BufferRegion objects nested in call_intrin arguments (e.g.,
the GEMM op's C buffer). This leaked the old local.fragment-typed C_local
var into the subsequent gemm lowering, producing a duplicate variable that
MakePackedAPI flagged as undefined.
Fix: apply tir.stmt_functor.substitute() to the entire body before the
ir_transform pass, so all Var references — including those inside opaque
Call args — are remapped to the metal.simdgroup-typed replacements.
Also remove the redundant per-SBlock substitute() and the duplicate
_remap_buffer function.
* fix: add missing <cmath> include for std::isinf/isnan in codegen_metal
* [Metal] Fix review findings: dedup warp partition, kernel_only, barrier comment, immutable gemm ops
- Extract shared ComputeSquareWarpPartition to src/backend/metal/op/utils.h,
used by both gemm.cc (Square policy) and copy.cc (simdgroup store lowering).
- Implement kernel_only parameter in MetalKernelAdapter.get_kernel_source.
- Add comment documenting that Metal's simdgroup_barrier synchronizes the
full threadgroup (no per-simdgroup barrier exists).
- Replace module-level mutable global in metal_fragment_to_simdgroup.py
with @lru_cache and remove CUDA-only gemm ops from the set.
* lint
* Refactor MetalFragmentToSimdgroup import path in LowerAndLegalize function
* Remove metal_fragment_to_simdgroup.py file and update import in metal/__init__.py
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
[CuTeDSL] Integrate host codegen call sites
Drive CuTeDSL launcher generation from lowered host tvm_call_packed call sites so multi-kernel wrappers preserve host IR order and repeated launches can be represented.
Expose separate device, host, and generated module sources through the CuTeDSL adapter and persist them separately in the kernel/autotune caches.
Normalize bare cutedsl targets through TileLang CUDA arch detection so Hopper runners use sm_90a for CuTeDSL compilation instead of the generic sm_90 Target("cuda") result.
Add CuTeDSL host-codegen integration tests covering multi-call launcher ordering, source accessors, cache restore behavior, and bare target architecture normalization.
Co-authored-by: Jayce <jayce.su@enflame-tech.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* [Example] Migrate eligible examples to eager style
* fix wrong shape annotation in gemm_split_k examples
* fix incorrect dtype in gemm_streamk example
* fix
* revert some changes
* Refactor attention functions to use new parameter structures and improve readability. Updated `example_amd_flash_attn_bwd.py` and `example_amd_flash_attn_fwd.py` to define input parameters more clearly. Enhanced shared memory usage in `flashattn_fwd` and `flashattn_sink` functions across various examples, ensuring compatibility with new configurations. Adjusted kernel definitions and added necessary checks for GPU capabilities.
* lint fix
* FIX
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* Preserve shared memory aliases after dynamic merge
* Enhance SharedMemoryRewriter to support alias preservation and buffer remapping during shared memory allocation merging
* lint fix
* [Feature] Introduce T.tfloat32
* fix: change tfloat32 type code to 130
* minor fix
* update tvm to latest
* minor fix
* minor fix
* update tvm
* add conversion from __nv_bfloat16 to tfloat32_t, update mhc example
* minor fix
* enable wgmma support for tfloat32, update gdn example
* update tf32 gemm test
* Add TL_DISABLE_MERGE_SHARED_MEMORY_ALLOCATIONS pass config
* Rename to TL_DISABLE_SHARED_MEMORY_REUSE and implement disable_reuse in C++
Instead of skipping MergeSharedMemoryAllocations entirely, add a disable_reuse
flag that still merges allocations into one but lays out each buffer sequentially
without lifetime-based sharing.
* Fix build: use AllocBufferNode instead of tir::AllocateNode in PlanSequentialLayout
* Add test for TL_DISABLE_SHARED_MEMORY_REUSE pass config
* [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>
Lei Wang