* Fix physical wait counts for async pipeline groups
* Add commit group tracking for async pipeline in inject_pipeline.cc
- Introduced buffer_to_commit_group mapping and commit_group_count in AsyncStateGlobal.
- Enhanced wait handling for tail consumers in async pipeline to ensure correct physical wait counts.
- Updated tests to validate fine-grained physical waits and tail drain waits for async pipeline.
* [Refactor] Remove unsupported tl_gemm and tl_gemm_sp operations
* Removed references to tl_gemm and tl_gemm_sp from various files, including code generation and built-in definitions, as these operations are currently unsupported.
* Updated related comments and diagnostics to reflect the removal of these operations, ensuring clarity in the codebase.
* This change simplifies the handling of unsupported operations and improves overall code maintainability.
* Reduce CUDA template include overhead
* Fix CUDA reduce fast min max helpers
* lint fix
* Fix CUDA swizzle ceil div helper
* Restore CUDA wait_wgmma intrinsic helper
* Fix CUDA cython TMA host adapter include
* Remove unnecessary check for Cutlass fast math header in CUDA code generation
* Update LibraryGenerator to correctly handle CUDA gencode flags for shared-library compilation, ensuring compatibility with target architecture and avoiding Hopper-only instruction issues.
* Refactor smem_ptr_to_uint function to use __cvta_generic_to_shared for improved clarity and performance.
* Refactor example scripts and CUDA math intrinsics for improved performance and clarity
- Updated `example_dequant_gemm_fp4_hopper.py` to directly call `run_regression_perf()` and print latency, removing unused argument parsing code.
- Modified `example_topk.py` to specify the backend as "cupti" in the benchmarking function.
- Enhanced CUDA code generation by introducing `RequiresTileLangMathHeader` to manage math header dependencies for specific functions.
- Updated barrier operations in `barrier.h` to use the correct syntax for shared memory barriers.
- Improved testing for bfloat16 math intrinsics by replacing pytest skips with decorators for CUDA availability checks.
* Refactor CUDA code for improved readability and consistency
- Reformatted the `RequiresTileLangMathHeader` function for better readability.
- Updated assembly syntax in `barrier.h` for clarity and consistency.
- Cleaned up whitespace in the bfloat16 math intrinsics test file to enhance code cleanliness.
* Remove unnecessary whitespace in CUDA code for improved readability
* Add TileLang intrinsic for prefetching TMA descriptor
- Introduced `prefetch_tma_descriptor` intrinsic in C++ and updated relevant files to use `call_intrin` instead of `call_extern`.
- Updated CUDA code generation to handle the new intrinsic correctly.
- Modified tests to ensure the new intrinsic is called appropriately and that no `call_extern` is used for internal TileLang operations.
- Enhanced documentation in SKILL.md regarding the addition of new TileLang intrinsics.
* Add bfloat16 hexp overload and update common.h for compatibility
- Introduced `hexp` overload for `bfloat16_t` in `common.h` to bridge TileLang and CUDA's handling of exponential functions.
- Updated tests to verify the presence of the new `hexp` definition in `common.h`.
- Enhanced CUDA code generation to support the new intrinsic for better performance and compatibility.
* Update common.h for CUDA compatibility with aligned double4 types
- Adjusted preprocessor condition in `common.h` to maintain compatibility with CUDA versions prior to 13 by ensuring proper definition of aligned double4 types.
- This change enhances the compatibility of the code with older NVRTC built-in vector headers.
* Refactor kernel launch materialization
* [Refactor] Improve git commit ID retrieval and version labeling logic
* Enhance the get_git_commit_id function to check for the existence of the .git directory before attempting to retrieve the commit hash, ensuring more robust behavior.
* Update the dynamic_metadata function to conditionally set the NO_VERSION_LABEL environment variable based on the presence of the .git directory.
* Modify the GitHub Actions workflow to unset the NO_VERSION_LABEL environment variable during the installation and version check of the package, allowing for accurate version labeling in SDist builds.
* [Refactor] Update WebGPU backend registration
* Import the WebGPU package to automatically register its pass pipeline.
* Remove the explicit registration of the WebGPU pass pipeline, streamlining the backend setup process.
* Place auto WS producers in first warp group
* [Transform] Update layout inference to bind thread range for fully replicated fragments; add tests for TMA fragment mask copy and producer-consumer warp specialization
* [Refactor] Update thread variable calculation to use zero-based indexing across GEMM operations
* [Transform] Enhance simplification process with ContextSingletonFloorDivFolder and add tests for floor division simplification
* [Refactor] Simplify tensor conversion in dequantization functions and improve performance in GEMM examples
[BugFix][Transform] Freshen buffer defs after loop unroll
UnrollLoop clones each loop body with Substitute, which preserves Buffer object identity for definitions inside the body. This can define the same DeclBuffer handle in multiple unrolled copies after flattening and violate TIR SSA. Freshen buffer definitions per cloned body while preserving DeclBuffer backing data aliases.
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* [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
* [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.
* [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>
* [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
* 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>
* 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>
* 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>
* [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
* [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>
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>
* 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
* [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
* [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>
* 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>
* [Codegen] Add lexical_alloc_scope for scoped local variable lifetime
Introduce a `lexical_alloc_scope` AttrStmt that generates `{ ... }` in
C/CUDA codegen, giving the underlying compiler accurate variable lifetime
information for better register allocation.
- Define `tl::attr::kLexicalAllocScope` constant
- LowerOpaqueBlock wraps block-local allocations in the new AttrStmt
- StorageRewrite treats it as a scope boundary with proper thread_scope_
save/restore so allocations are not hoisted past the boundary
- CUDA and HIP codegen emit scoped `{ }` blocks
- Add tests for IR insertion, StorageRewrite preservation, and codegen output
* lint fix
* [Codegen] Refine lexical_alloc_scope: skip top-level blocks and decouple from thread_scope
Two improvements to the lexical_alloc_scope mechanism:
1. LowerOpaqueBlock: only insert lexical_alloc_scope for blocks inside
loops (inside_loop_ > 0). Top-level blocks already have function-body
lifetime, so the extra `{ }` in codegen is pointless.
2. StorageRewrite: introduce a separate `lexical_scope_` / `effective_scope()`
instead of overriding `thread_scope_`. This avoids breaking
PlanNewScope's toggle protocol when lexical_alloc_scope is nested
inside thread_extent, fixing a `ICHECK(thread_scope_ == op)` crash.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* [Transform] Fix alloc-scope placement regressions
* [Infra] Remove clang-tidy integration
* Remove unused local descriptor allocation pass and related tests; update example to disable main execution and print kernel source for debugging.
* Preserve lexical alloc scopes for nested register buffers
* Limit lexical alloc scopes to local storage
* refactor
* Unify lexical alloc scope annotations
* Clean up dead code and fix unsafe cast
- Remove unused block_nesting_ member from OpaqueBlockLower
- Use static_cast instead of reinterpret_cast in ResolveAllocationSite
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Lei Wang