* 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>
* 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.
