mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
44690f67769c91ae4942968ca5b87dc25ae31ebf
108
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Lei Wang
|
44690f6776 |
[Transform] Preserve bind scope when splitting if statements (#2232)
Fix guarded bind scope splitting |
||
|
|
b939fa0137 |
[TIR][IR] Update to use tirx (#2216)
* [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> |
||
|
AkeliaandLeiWang1999
|
a7d9380e5a |
[Enhancement]Support mixed-sign ramp indices in LegalizeNegativeIndex (#2225)
* 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> |
||
|
|
431c85f100 |
[BugFix] Consider non-local store in external call and SIMT producer for warp specialize (#2166)
* 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
|
a797e5120f |
[Refactor] Move backend-specific GEMM implementations and transforms into backend directories (#2165)
* [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 |
||
|
|
2905fde91d |
[Refactor] Refactor register annotation lowering (#2088)
* [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> |
||
|
Lei Wang
|
3e87f1e872 |
[TMA] Fix TMA descriptor init placement (#2151)
Fix TMA descriptor init placement |
||
|
Tong WUandwutong.1109
|
0f29f9c986 |
[BugFix] Relax loop wait and adjust trailing drain behavior in async pipeline tests (#2092)
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> |
||
|
Lei WangandClaude Opus 4.6
|
046b1bdc69 |
[FFI] Remove upper version bound on apache-tvm-ffi (#2071)
* 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>
|
||
|
Lei Wang
|
99ee74b343 |
[CUDA] Support int4 T.gemm (#2063)
* 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 |
||
|
Tong WU
|
9c5fe44c7a |
[BugFix] Keep shared-prelude local vars in producer-consumer WS (#2055)
* [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 |
||
|
Lei WangandFreebase6912
|
891109ea87 |
[Transform] Add InjectTcgen05Fence pass (#2003)
* [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>
|
||
|
|
74fc9806fb |
[Refactor] Refactor DecoupleTypeCast Pass (#2026)
* 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
|
||
|
Lei WangandClaude Opus 4.6
|
853e805b2a |
[Codegen] Add lexical_alloc_scope for scoped local variable lifetime (#2023)
* [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 WangandClaude Opus 4.6
|
3ee0988c86 |
[Refactor][Pipeline] Run pipeline rewriting before layout inference and stabilize tiled WS (#2002)
* Refactor optimization pipeline in phase.py to streamline barrier handling
- Consolidated the handling of shared barriers and pipeline planning by removing redundant conditional checks.
- Ensured that `LowerSharedBarrier`, `PipelinePlanning`, and `InjectSoftwarePipeline` are consistently applied, enhancing the clarity and efficiency of the optimization process.
This change improves the maintainability of the code while preserving existing functionality.
* Enhance tile operation handling with instruction annotation and refactor gemm.h
- Added a new `InstructionAnnotation` pass to annotate tile operations with their instruction kind before layout inference, improving the optimization pipeline's ability to reason about instruction mixes.
- Refactored `gemm.h` to move the `allowTcgen5Mma` and `allowWgmma` methods under the private section, enhancing code organization and encapsulation.
These changes improve the clarity and maintainability of the code while preserving existing functionality.
* Enhance multi-version buffer handling and introduce tile-level warp specialization
- Updated `multi_version_buffer_rewriter.cc` to improve read/write access detection for tile operations by analyzing `tl.tileop.region` calls, ensuring accurate buffer access tracking.
- Modified `phase.py` to integrate `ProducerConsumerWarpSpecializedTiled` before layout inference, allowing for high-level tile-op IR transformations that enhance producer/consumer splits.
- Added a new `ProducerConsumerWarpSpecializedTiled` function in `__init__.py` to facilitate tile-level warp specialization, improving the optimization pipeline's efficiency.
These changes enhance the handling of multi-version buffers and optimize the transformation process for tiled operations.
* Enable tiled WS for stage-1 pipelines
* Move Hopper pipeline planning before layout inference
* Run pipeline rewriting before layout inference
* Add pipeline refactor WIP: subtree_modified_ guard, debug prints, and plan
- inject_pipeline.cc: guard reads/writes recalculation with subtree_modified_
flag to prevent local.var buffer promotion to kernel parameters
- phase.py: add temporary debug prints after PipelinePlanning and
InjectSoftwarePipeline (to be removed during CI fix work)
- example_group_per_split_token_cast_to_fp8.py: add disable_cache for debugging
- docs/plan.md: implementation plan for fixing CI test failures
- draft.md: original design draft
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Remove debug artifacts: prints in phase.py, disable_cache in example
- Remove print("After PipelinePlanning"), print(mod),
print("After InjectSoftwarePipeline"), print(mod) from
LowerAndLegalize in tilelang/engine/phase.py
- Remove tilelang.disable_cache() from
examples/cast/example_group_per_split_token_cast_to_fp8.py
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix LayoutInference crash on multi-versioned pipeline buffers
After InjectSoftwarePipeline, multi-versioned buffers share the same
data Var as the original but have an extra leading dimension (num_stages).
LayoutInference's alias propagation and annotation handling tried to
Reshape layouts between these buffers, which failed because the total
element counts differ.
Guard three Reshape call sites in layout_inference.cc to skip sibling
buffers whose total storage size is incompatible with the source layout.
This lets multi-versioned buffers get their own layout inference instead
of inheriting an incompatible layout from the original buffer.
Fixes compilation failures in dequantize_gemm, GDN, and other kernels
that use software pipelining with shared memory buffers.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Strip pipeline annotations on WS fallback for TMA kernels
When ProducerConsumerWarpSpecializedTiled identifies a TMA kernel as a
warp-specialization candidate but the tiled rewriter cannot handle it
(e.g., conditional loop bodies like sparse block masks), the fallback
previously returned the original function with num_stages annotations
intact. PipelinePlanning and InjectSoftwarePipeline would then generate
non-WS TMA pipeline code with broken barrier phase tracking for
conditional pipeline bodies (barrier waits outside conditionals cause
deadlocks when the condition is false).
Fix: on WS fallback, strip num_stages annotations from pipeline loops
so that the pipeline passes skip the function. The kernel runs
unpipelined but correctly.
Fixes CUDA_ERROR_LAUNCH_FAILED in blocksparse_gemm and related TMA
kernels with conditional loop bodies.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Add guarded phase-counter support to tiled WS pass
Port PhaseCounter and StageExprReplacer from the legacy
ProducerConsumerWarpSpecialized pass into the tiled WS pass to handle
conditional loop bodies (e.g., sparse block masks).
When the pipeline loop body is wrapped in an IfThenElse without else:
1. Unwrap the condition before classifying statements
2. Create separate producer/consumer PhaseCounters (local int32 buffers)
3. Use counter-based stage/parity expressions instead of loop-variable
4. Wrap producer and consumer bodies in the original condition
5. Increment counters at end of each guarded iteration
6. Rewrite shared-buffer stage indices via StageExprReplacer
This ensures barrier parity stays correct when iterations are
conditionally skipped, fixing CUDA_ERROR_LAUNCH_FAILED in
blocksparse_gemm and related TMA kernels with conditional execution.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Add GemmSPNode support to pipeline planning and injection
Add GemmSPNode handling to:
- inject_pipeline.cc AddReadsWritesForTileOp: model A, E, B as reads
and C as write (E is the sparse metadata buffer)
- pipeline_planning.cc: same access model for dependency analysis
This makes sparse GEMM visible to the pipeline machinery for correct
stage assignment and buffer multi-versioning. However, the tile-op's
consumer-side buffer accesses still don't get stage-indexed because
the pipeline body rewriter can't rewrite high-level tile-op Call
arguments (an architectural limitation of running InjectSoftwarePipeline
before LowerTileOp).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix GemmSPNode::Lower to use MakeAccessPtrFromRegion
Replace whole-buffer access_ptr(1) calls with MakeAccessPtrFromRegion
for A, B, C, and E buffers in sparse GEMM lowering. This preserves
stage-specific region offsets from pipeline multi-versioning, matching
the dense GemmNode::Lower pattern.
CUDA output now shows correct stage-indexed consumer accesses:
gemm_sp_ss(..., (k%3)*8192, (k%3)*8192+27648, C_local, (k%3)*2048+49152)
instead of always using stage-0 offsets.
Note: gemm_sp still produces incorrect results because the kernel needs
warp specialization but TiledWSCandidate::Check doesn't recognize it
as a TMA candidate. The non-WS pipeline path generates structurally
different code from the reference WS path.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix gemm_sp and seer_attention regressions
Three changes in producer_consumer_ws_tiled.cc:
1. Require num_stages >= 2 for WS candidacy (was >= 1). Single-stage
kernels like seer_attention don't need WS and the transformation
produces incorrect results for them.
2. Add HasTmaPipeline() check to detect TMA kernels with pipeline
annotations that are rejected by the full WS candidate check (e.g.,
kernels with manual layout annotations like gemm_sp).
3. Strip num_stages annotations for rejected TMA pipeline kernels to
prevent InjectSoftwarePipeline from generating broken non-WS TMA
pipeline code.
One change in gemm_sp.cc:
- Use MakeAccessPtrFromRegion for A, B, C, E buffer access pointers
instead of whole-buffer access_ptr. This preserves stage-specific
region offsets from pipeline multi-versioning.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix WS conditional body unwrapping for complex loop bodies
Only unwrap IfThenElse wrapper when the then-branch is a simple flat
sequence of tile-op Evaluate calls. Skip unwrapping for complex bodies
with LetStmt, For, or other control flow that could break variable
scoping when split into producer/consumer for WS.
Fixes variable-used-before-definition error in
blocksparse_attention sparse_gqa_decode_varlen_indice, which has a
conditional loop body containing LetStmt bindings.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Revert num_stages >= 2 to >= 1 for WS candidacy
The >= 2 threshold broke test_num_stages_one_pure_tma_keeps_auto_warp_specialize.
Pure TMA kernels with num_stages=1 should still be WS candidates.
The seer_attention issue (num_stages=1 with manual layout) is handled
by the has_manual_layout_ check, not the num_stages threshold.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Propagate manual layouts onto MVB-expanded buffers via Layout::Expand
When LayoutInference encounters an MVB-expanded buffer (with leading
stage dimensions) whose trailing dimensions match the original layout,
use Layout::Expand to propagate the manual layout instead of rejecting
or skipping the buffer. Applied to all 3 layout propagation paths:
annotated layout map, alias propagation, and finalization.
Also remove the blanket !has_manual_layout_ WS candidate rejection
since manual layouts now survive onto versioned shared buffers via
Layout::Expand.
Fixes test_sparse_ws_regular_metadata_copy_stays_in_producer.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Extend conditional body unwrapping for LetStmt chains and restore
has_manual_layout_ guard
1. LetStmt chain peeling: when IfThenElse then_case starts with
LetStmt bindings, peel them and append to let_bindings before
checking the simple-body guard. This allows WS for conditional
bodies with variable definitions (e.g., sparse attention patterns).
2. Restore !has_manual_layout_ in WS candidacy check: removing it
caused dequant_groupedgemm_bf16_mxfp4_hopper to fail because
MXFP4 layouts don't survive MVB expansion. The Layout::Expand fix
handles sparse metadata layouts but not all manual layout types.
3. Layout::Expand propagation (from Round 6) remains in place for
future use when MVB learns to handle all manual layout types.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Remove is_simple_body guard from conditional WS unwrapping
The simple-body guard only accepted flat Evaluate sequences inside
IfThenElse, blocking legitimate WS for complex conditional bodies
like sparse flash attention (T.clear, T.reduce_max, T.gemm inside
the guard). The LetStmt peeling already handles variable scoping.
Fixes test_pure_tma_consumer_local_init_does_not_leak_into_producer
and test_sparse_ws_regular_metadata_copy_stays_in_producer.
The remaining test_mixed_tma_cp_async_shared_stage_barriers failure
is a pre-existing issue on the original branch: the tiled WS pass
produces SIMT copies instead of cp.async because LowerPTXAsyncCopy
is not in the pass pipeline (the comment at phase.py:308 says it
runs earlier but no actual call exists on either branch or reference).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Preserve pipeline context for WS producer cp.async injection
Wrap the WS producer loop body in a kPipelineContextNumStages AttrStmt
so that LowerTileOp's pipelined_depth_ is > 0 when processing SIMT
producer copies. This enables InjectPTXAsyncCopy to generate cp.async
for global-to-shared copies in the WS producer branch.
Without this, the WS rewriter strips all pipeline annotations from
the rewritten loops, causing LowerTileOp to skip cp.async injection
for SIMT producers. The consumer loop stays annotation-free since it
doesn't need cp.async.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix mixed barrier protocol, consumer-init sinking, restore manual-layout guard
1. Mixed TMA+cp.async barrier: use ptx_cp_async_barrier_noinc for
forward barrier arrival in mixed producer groups, matching the
reference producer_consumer_ws.cc protocol.
2. Consumer-only pre-loop init sinking: in ReplacePipelineLoopInStmt,
guard pre-loop siblings as consumer-only when they're not classified
as producer (TMA/SIMT/cp.async). Fragment init (T.fill, T.clear)
and local buffer init are placed in the consumer branch instead of
the shared prelude.
3. Restored blanket has_manual_layout_ guard: the dtype-based heuristic
to distinguish sparse metadata from MXFP4 layouts doesn't work
because both use uint8. Dequant_groupedgemm_bf16_mxfp4_hopper
requires the guard to prevent broken WS.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Narrow consumer-init sinking, targeted manual-layout check, per-group barrier
1. Consumer-init sinking: only sink pre-loop stmts that are FillNode
writing to fragment/local buffers. Keep block_mask setup and shared
state in the shared prelude.
2. Manual-layout: attempt targeted check that only rejects when TMA
copy destinations match layout_map entries. Collect layout_map vars
and compare against TMA copy destinations.
3. Per-group cp.async barrier: use group-level cp.async flag (single
group for now) instead of function-wide boolean.
The 3 WS issue tests still fail because the layout_map annotation
parsing falls through to the conservative rejection path.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix layout_map parsing with Var→Buffer mapping for manual-layout check
The layout_map annotation uses Map<Var, Layout> before LayoutInference
(not Map<Buffer, Layout>). Parse as Map<ObjectRef, ObjectRef> and
handle both key types: Buffer (post-inference) and Var (pre-inference).
For Var keys, look up the corresponding alloc_buffer by data Var match.
Compare collected manual-layout buffers against pipeline copy
destinations: reject only when a manually-laid-out buffer is also
a producer copy target (TMA/SIMT/cp.async) inside the pipeline.
This allows sparse metadata (E_shared, SIMT-copied) onto the WS path
while rejecting dequant MXFP4 (B_shared, SIMT-copied with swizzle).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix SEGFAULT in manual-layout check, use DetectSwizzleMode
1. Fixed SEGFAULT: removed dangling layout_map_layouts_ vector that
was never populated, causing OOB access. Now stores Buffer+Layout
pairs in layout_map_entries_.
2. Use DetectSwizzleMode to distinguish swizzled layouts (MXFP4,
incompatible with MVB) from non-swizzled (sparse metadata, safe).
Swizzled layouts reject WS candidacy; non-swizzled layouts allow it.
3. Removed debug LOG(WARNING) from hot path.
4. Parse layout_map annotation keys as both Buffer and Var (via
Map<ObjectRef, ObjectRef>), resolving Var keys to alloc_buffers.
Results: sparse metadata WS test PASSES, dequant PASSES (protected).
Down to 2 failures: mixed barrier pattern + consumer init sinking.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Simplify consumer-init sinking to Evaluate-only classification
Sink pre-loop Evaluate nodes classified as kConsumer into consumer
branch. Keep For loops (block_mask setup), producer copies, and other
control flow in the shared prelude. This is simpler and safer than
the FillNode scope check approach.
The test_pure_tma_consumer_local_init test still fails because the
T.fill statements are at a different structural level than the
SeqStmt where the pipeline loop lives. Fixing this requires deeper
IR structure analysis.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Sink consumer-only pre-loop init into WS consumer branch
Extract consumer-only pre-loop Evaluate statements (T.fill on
fragments) from the shared prelude and prepend them to the consumer
branch inside the WS if/else structure. This ensures fragment init
like acc_o, logsum, scores_max appears only in the consumer branch,
not in the shared prelude or producer branch.
Uses a two-pass approach: first ReplacePipelineLoopInStmt extracts
consumer-only stmts into extracted_consumer_init_, then the WS body
is rebuilt with the extracted stmts prepended to the consumer branch.
Fixes test_pure_tma_consumer_local_init_does_not_leak_into_producer.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Finalize pipeline refactor fixes
* Refactor CopyNode stride checks for TMA bulk load/store
- Extracted the global stride validation logic into a new static method `CheckGlobalStrides` for better readability and reusability.
- Updated `CheckBulkLoad` and `CheckBulkStore` methods to utilize the new stride checking function, improving code clarity and maintainability.
- Enhanced documentation for the new method to clarify its purpose and requirements.
* Remove redundant copy direction notes from documentation in `copy.h` for stride checks. This simplifies the comments while maintaining clarity on TMA requirements.
* refactor
* layout related fix
* refactor mbarrier with software pipeline.
* remove legacy pass
* Enhance TMA barrier handling and merging logic
- Updated `CopyToTmaCopyRewriter` to conditionally emit arrive barriers based on the last TMA copy.
- Refactored barrier creation to allow merging of TMA barriers when conditions are met, reducing the number of barriers needed.
- Adjusted consumer and producer logic to accommodate merged barriers, ensuring correct synchronization across TMA copies.
- Improved documentation and comments for clarity on the new barrier merging behavior.
* Update example_mhc_post.py to disable main execution and add test function call
- Commented out the main function call in `example_mhc_post.py`.
- Added a call to `tilelang.disable_cache()` and a new `test(n=4096, h=2560)` function for testing purposes.
* Enhance pipeline barrier handling and TMA copy detection
- Updated `RewritePipelineTmaBarriers` to accept additional parameters for loop variable and stage count, improving barrier synchronization.
- Modified `HandleTileOp` to ensure only user-defined T.copy patterns are marked for pipeline barriers, preventing interference with T.tma_copy.
- Adjusted test cases to enforce specific CUDA compute version requirements for better compatibility and consistency.
* lint fix
* Enhance convolution example and TMA handling
- Added kernel source printing in `example_convolution.py` for debugging.
- Disabled the main function call and invoked `run_regression_perf()` to streamline execution.
- Introduced `MakeTmaLeaderCondition` function in `copy.cc` to improve TMA leader-thread condition handling.
- Updated `LowerBulkCopy` and `LowerBulkCopy1D` methods to utilize the new TMA leader condition for better thread management.
- Enhanced `Conv2DIm2ColOpNode` to support barrier annotations, improving synchronization in TMA operations.
* enhance
* Remove example_dequant_groupedgemm_bf16_mxfp4_hopper.py and clean up phase.py by removing unnecessary whitespace. Enhance TMA handling in producer_consumer_ws_tiled.cc to improve barrier synchronization and streamline TMA copy operations.
* fix
* fix
* Eliminate MVB(barrier_only=true) late fixup from OptimizeForTarget
Move pipeline barrier ownership into InjectSoftwarePipeline: create
pipeline_mbar[num_stages] at final expanded size instead of
pipeline_mbar[1] that required late MVB expansion.
Key changes:
- RewritePipelineTmaBarriers creates barriers at num_stages size
- Barrier indices use FloorMod(loop_var - loop_min, num_stages)
- barrier_init has num_stages entries (one per slot, arrive_count=1)
- CopyToTmaCopyRewriter accepts PrimExpr barrier_id (was int)
- phase.py: remove MVB(barrier_only=True), unify both paths under
PlanAndUpdateBufferAllocationLocation
- Fix pre-existing test expectation for tma_copies annotation
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix Codex review issues: num_stages=1 regression, im2col pipeline barriers, test fixes
1. Fix num_stages=1 regression: use tl_pipelined_num_stages annotation
for barrier sizing instead of max_stage+1. Gate barrier creation on
pipeline_depth > 1 so num_stages=1 kernels don't get multi-versioned
pipeline barriers.
2. Extend RewritePipelineTmaBarriers to handle c2d_im2col: annotate
im2col calls with pipeline barrier in CopyToTmaCopyRewriter. Add
im2col TMA recognition to PipelinePlanning. Fix im2col Lower() to
respect emit_arrive annotation from pipeline barriers.
3. Fix test_simple_pipeline: use annotation-checking approach instead of
structural equality with hardcoded tma_copies annotation that fails
on non-Hopper targets.
4. Remove dead kPipelineMVBStageExpr/kPipelineMVBParityExpr/
kPipelineMVBContextNumStages emission from inject_pipeline.cc since
pipeline barriers are now created at final size.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix depth-1 barrier ownership, add regression tests, remove dead kPipelineMVB*
1. Fix depth-1 ownership gap: always create shared pipeline barrier for
TMA copies even when pipeline_depth=1. This prevents LowerTileOp from
allocating separate per-copy internal barriers, keeping num_stages=1
kernels at pipeline_mbar[1] (single slot).
2. Add checked-in lowering regression tests:
- non-WS num_stages=3 TMA GEMM → asserts pipeline_mbar[3], no fallback
- non-WS num_stages=1 TMA GEMM → asserts pipeline_mbar[1], no multi-slot
- non-WS num_stages=3 im2col → asserts pipeline_mbar[3] feeds tma_load_im2col
3. Remove dead kPipelineMVB* definitions from pipeline_utils.h and
consumption code from multi_version_buffer_rewriter.cc (stacks,
explicit parity/version index logic, attr stripping).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix GemmWMMA.lower() signature and autotuner cache backwards compatibility
[P1] Add missing mbar_phase_expr parameter to GemmWMMA.lower() to match
the interface expected by GemmPy.lower() dispatch. Without this, any
RDNA kernel using the WMMA path would fail with TypeError.
[P2] Make out_idx.json loading optional in autotuner load_from_disk().
Older cache directories don't have this file; fallback to
compile_args.out_idx when the file is absent.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Add .humanize to gitignore
* Fix HIP swizzle codegen and CPU scalar GEMM region handling
[P2] Update codegen_hip.cc to parse tvm_tuple(device_func, panel_size)
format for threadblock_swizzle_pattern, matching the annotation format
now emitted by T.use_swizzle(). Previously expected StringImmNode which
would ICHECK-fail.
[P2] Fix gemm_scalar.py: clear only the output tile region (not the
whole buffer), and use the last two dimensions from regions to handle
rank>2 buffers with leading singleton dimensions.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix HIP codegen to parse tvm_tuple swizzle annotation format
Update codegen_hip.cc threadblock_swizzle_pattern handler to parse
tvm_tuple(device_func, panel_size) format, matching CUDA and cutedsl
codegens. Previously expected StringImmNode which ICHECK-fails with
the current annotation format.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Restore kPipelineMVB* annotations — still needed by WS tiled full MVB
The kPipelineMVBStageExpr/ParityExpr/ContextNumStages annotations and
their MVB consumption code cannot be removed: they are emitted by
EmitImpl for ALL pipelines and consumed by the full MVB call inside
ProducerConsumerWarpSpecializedTiled. Removing them broke non-TMA
pipeline kernels like mhc_post that go through the WS path.
Restores: annotation emission in inject_pipeline.cc, constant
definitions in pipeline_utils.h, and consumption code in
multi_version_buffer_rewriter.cc.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Unify barrier multi-versioning in InjectSoftwarePipeline (Direction A)
Move all pipeline barrier multi-versioning into InjectSoftwarePipeline
via ExpandPipelineBarriers, eliminating the need for late
MVB(barrier_only=true) in OptimizeForTarget.
ExpandPipelineBarriers runs before BuildPipeline and handles:
- ISP-created pipeline_mbar (for non-WS TMA pipelines)
- User-written T.alloc_barrier (for manual WS pipelines like softpipe)
Key design: only barriers with explicit ptx_arrive_barrier calls OR
ISP-created local barriers are expanded. Barriers whose arrival is
managed internally by tile-ops (e.g., tcgen05 MMA arrive) are left
unchanged. This distinguishes pipeline sync barriers from hardware-
managed barriers.
Buffer expansion: barrier[N] -> barrier[N * num_stages]
Index rewriting: barrier[idx] -> barrier[stage_expr * N + idx]
Parity rewriting: user_parity -> (iteration_block + offset) % 2
barrier_init replication: [c0,c1] -> [c0,c1,c0,c1] for num_stages=2
Expanded buffers propagate to outer blocks via pending_buffer_remap_.
barrier_init annotations are replicated in VisitStmt_(BlockNode).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix Python 3.9 compat: replace star expression in index with tuple concat
* fix
* fix
* revert changes
* fix
* enhance pipeline
* enhance pipeline
* implement descriptor reuse pass
* Refactor pipeline management and enhance async copy handling
- Updated `example_mhc_pre.py` to replace the main function call with a test function for better testing flexibility.
- Modified `example_gqa_decode.py` to disable argument parsing and added latency measurement for regression performance.
- Enhanced `copy.cc` and `copy.h` with new checks for pipeline-managed cp.async synchronization and improved async copy handling.
- Updated `inject_pipeline.cc` and `pipeline_planning.cc` to refine the handling of tile operations and global/shared buffer checks.
- Added tests to ensure correct behavior of async pipeline and descriptor allocation reuse.
This commit improves the overall pipeline management and async copy capabilities, ensuring better performance and flexibility in the codebase.
* enhance pipeline
* enhance pipeline
* enhance pipeline
* enhance pipeline
* enhance pipeline
* Revert "enhance pipeline"
This reverts commit
|
||
|
WilliamandLeiWang1999
|
a82fa719b8 |
fix: fix copy+cast vectorize loop to use wider vector load/store instrcution (#2004)
* fix copy+cast vectorize loop to use wider vector load/store instrcution * clean test * fix test * fix format * test fix --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Lei WangandClaude Opus 4.6
|
1ede0bd748 |
Introduce T.deallocate_tmem and T.transpose (#1971)
* Fix tcgen05 barrier allocation planning regression * Add explicit TMEM deallocation and shared transpose copy * Add transpose operation to documentation and update implementation in copy_op.py * LINT FIX * Refactor: extract transpose from CopyNode into standalone TransposeNode op Remove the transpose annotation logic from CopyNode (GetTranspose, MakeIndices transpose branch, MakePredicate transpose branch, and GetCopyInst early return). Transpose is now handled by the independent TransposeNode registered as tl.tileop.transpose. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Refactor CopyNode to remove transpose handling from index generation and predicate creation. Simplify MakeIndices and MakePredicate methods by eliminating unnecessary transpose checks and related logic. Update associated checks to ensure index consistency. Clean up unused GetTranspose method in copy.h. * lint fix --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
||
|
Lei Wang
|
e2c88335c6 |
Fix wrapped pre-loop TMA prefixes in producer-consumer WS (#1975)
* Support wrapped pre-loop TMA prefixes in producer-consumer WS * Format producer-consumer WS changes * Keep pure-TMA barrier grouping unchanged |
||
|
|
a2d6e01866 |
[Bugfix] Fix double buffer versioning when TMA is used without warp specialization (#1962)
* [Bugfix] Fix double buffer versioning when TMA is used without warp specialization When TMA is enabled but warp specialization is disabled, both MultiVersionBuffer and InjectSoftwarePipeline independently added version dimensions to shared data buffers, resulting in 4x buffer allocation instead of the intended 2x. Add a `barrier_only` mode to MultiVersionBuffer so that in the no-WS path it only versions barrier buffers (for mbarrier parity rewriting), while data buffer versioning is left to InjectSoftwarePipeline. Also flatten the if/else branches in OptimizeForTarget into a single linear flow with a three-way branch (WS / TMA-noWS / noTMA), removing the duplicated LowerOpaqueBlock call. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * [Enhancement] Introduce TMA attribute handling in LowerTileOp and optimization pipeline - Added a new PrimFunc attribute `tl.has_tma` to indicate whether TMA operations were generated during the lowering process. - Updated `OptimizeForTarget` to utilize the `tl.has_tma` attribute for determining the appropriate optimization path, enhancing clarity and efficiency in the pipeline. - Refactored the `module_has_tma` function to check for TMA operations across all functions in the module, streamlining the decision-making process for TMA-related optimizations. This change improves the maintainability of the code and ensures that TMA operations are correctly accounted for in the optimization pipeline. * fix * remove barrier related test for cutedsl. --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
||
|
Lei WangandZhiwen Mo
|
715ca303da |
[Bugfix] Handle int64 offsets in ThreadSync for tvm_access_ptr (#1952)
Updated the ThreadSync implementation to support int64 offsets in tvm_access_ptr, preventing crashes during index reconstruction. Added a regression test to ensure correct behavior with shared/shared.dyn pointers. Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei WangandClaude Opus 4.6
|
9246b82b24 |
[Refactor] Replace create_list_of_mbarrier with buffer-based T.alloc_barrier (#1944)
* [Refactor] Replace create_list_of_mbarrier/get_mbarrier with buffer-based T.alloc_barrier Unify the two mbarrier representation paths into a single buffer-based approach using shared.barrier scope buffers with barrier_init block annotations. This eliminates the create_list_of_mbarrier and get_mbarrier intermediate intrinsics, making the IR cleaner and the barrier handling consistent across user-explicit (T.alloc_barrier) and compiler-generated (T.copy/T.tma_copy) paths. Key changes: - LowerTileOp now injects a shared.barrier buffer + barrier_init annotation into the root block instead of emitting create_list_of_mbarrier statements - copy.cc uses BufferLoad on the barrier buffer instead of get_mbarrier - WarpSpecializedRewriter and ProducerConsumerWS create barrier buffers with annotations instead of create_list_of_mbarrier calls - LowerSharedBarrier runs a second time after WS passes to process newly created barriers - LowerHopperIntrin no longer handles mbarrier allocation/init (fully delegated to LowerSharedBarrier) - Removed create_list_of_mbarrier and get_mbarrier op definitions from builtin.h/cc - Removed get_mbarrier codegen handling from CUDA, CuTeDSL, HIP backends - Updated WarpSpecializedDetector, inject_tma_barrier, and eliminate_storage_sync passes to detect barriers via shared.barrier scope buffers instead of the removed intrinsics Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * style: apply clang-format and ruff-format Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Refactor mbarrier handling in copy operations Updated the handling of mbarriers in the `LowerBulkCopy` and `LowerBulkCopy1D` methods to allocate a single barrier slot instead of multiple for each pipeline stage. This change simplifies the barrier management by allowing the `MultiVersionBuffer` to expand the barrier for pipelining stages. Additionally, modified the `LowerArgs` structure to use a pointer for the `mbarrier_buffer`, ensuring compatibility with lazy buffer creation. Key changes: - Simplified mbarrier allocation logic in `copy.cc`. - Updated `LowerArgs` to use a pointer for `mbarrier_buffer`. - Adjusted related comments for clarity. This refactor enhances the consistency and readability of the code while maintaining functionality. * Remove deprecated TMA barrier and warp specialization files This commit deletes the `inject_tma_barrier.cc`, `warp_specialized_rewriter.cc`, and associated test files as part of a cleanup effort. These files are no longer necessary following recent refactoring and updates to the TMA pipeline and warp specialization logic. The removal simplifies the codebase and eliminates outdated implementations that have been superseded by newer approaches. Key changes: - Deleted `inject_tma_barrier.cc` and `warp_specialized_rewriter.cc`. - Removed related test file `test_tilelang_transform_warp_specialized.py`. - Updated `phase.py` and `__init__.py` to reflect the removal of the WarpSpecialized and InjectTmaBarrier functionalities. This cleanup enhances maintainability and clarity in the transformation pipeline. --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
||
|
Lei WangandClaude Opus 4.6
|
ded6a99221 |
[Feature] Add Producer-Consumer Warp Specialization and T.tma_copy() API (#1909)
* [Feature] Add Producer-Consumer Warp Specialization and T.tma_copy() API
This PR introduces Producer-Consumer Warp Specialization for sm90+ TMA
pipelines and a new T.tma_copy() API for explicit mbarrier management.
Key changes:
1. **ProducerConsumerWarpSpecialized pass** (new):
Splits pipelined TMA loops into producer (TMA loads) and consumer
(compute) warp groups with back-pressure barriers for buffer reuse.
Works with num_stages >= 1. Producer warp (128 threads) handles
arrive_expect_tx + tma_load; consumer warps handle compute + arrive
on back-pressure barriers.
2. **T.tma_copy() API** (new):
Fire-and-forget TMA copy with a required `barrier` parameter. Unlike
T.copy() which emits producer+wait pairs, T.tma_copy() emits only
arrive_expect_tx + tma_load. User manages synchronization via
T.mbarrier_wait_parity().
3. **MultiVersionBuffer barrier expansion**:
Extended to handle `shared.barrier` scope buffers, expanding them for
pipelining (1D size multiplication instead of prepending a dimension).
Auto-computes mbarrier parity as `(k // num_stages) % 2`.
4. **Pass ordering** (phase.py):
LowerSharedBarrier now runs after MultiVersionBuffer in the TMA path
so barrier buffers retain their `shared.barrier` scope during expansion.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* fix for annotate reg alloc
* Enhance debugging and performance measurement in examples
- Added print statements to output kernel source for both forward and backward sparse MLA examples, aiding in debugging.
- Updated `sparse_mla_fwd_pipelined.py` to run performance regression tests and print average latency, improving performance tracking.
- Disabled auto cache in `example_mha_sink_fwd_bhsd.py` for better control during testing.
* fix
* Enhance debugging and performance tracking in examples
- Added print statements to output kernel source in various example scripts, including `example_gqa_sink_bwd_bhsd.py`, `example_blocksparse_gemm.py`, and `example_dequant_groupedgemm_bf16_mxfp4_hopper.py`, aiding in debugging.
- Updated default `window_size` in `example_gqa_sink_bwd_bhsd.py` for improved configuration.
- Disabled auto cache in `example_gemm_schedule.py` to enhance control during testing.
- Introduced a new test for preserving protected auto-injected wait groups in the async copy optimization process.
* clean
* Remove unnecessary print statements for kernel source in example scripts
- Eliminated print statements from `example_gqa_sink_bwd_bhsd.py`, `example_blocksparse_gemm.py`, and `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` to clean up output and improve readability.
- These changes streamline the examples while maintaining functionality.
* clean
* lint fix
* Refactor T.tma_copy() and barrier synchronization in examples and tests
- Updated `example_warp_specialize_gemm_copy_1_gemm_0.py` to use `T.tma_copy()` for optimized memory copying with explicit barrier management.
- Enhanced test cases in `test_tilelang_issue_tma_no_ws.py` and `test_tilelang_language_tma_copy.py` to reflect changes in barrier synchronization, ensuring proper functionality with the new T.tma_copy() API.
- Removed unnecessary print statements and adjusted barrier allocations for clarity and performance.
- Cleaned up the `test_tilelang_transform_inject_tma_barrier.py` file by deleting it, as it was no longer needed.
* Add buffer data mapping for remapped buffers in multi_version_buffer_rewriter
- Introduced a check to ensure that the data variable of a buffer is discoverable, allowing the barrier_init annotation update to find the remapped buffer correctly.
- This change enhances the functionality of the buffer remapping process, ensuring proper synchronization in the transformation pipeline.
* Implement TMA im2col load synchronization with mbarrier allocation
- Added support for allocating mbarriers to synchronize TMA im2col loads, enhancing the pipeline's barrier management.
- Introduced logic to handle mbarrier creation and integration into the TMA copy process, ensuring proper synchronization and data transfer.
- Updated the handling of the mbarrier in the TMA copy statement to improve performance and maintain functionality within the thread-gated block.
* Implement TMA im2col load synchronization with mbarrier allocation
- Added support for allocating mbarriers to synchronize TMA im2col loads, enhancing the pipeline's barrier management.
- Introduced logic to handle mbarrier initialization and integration within the TMA copy process, ensuring proper synchronization during memory operations.
- Updated the handling of the mbarrier in the im2col operation to improve performance and maintain consistency in multi-threaded environments.
* disable tma lower for test_cutedsl_barrier
* Refactor code formatting in test_tilelang_jit_cutedsl.py
- Improved readability by adjusting the formatting of the `tilelang.compile` function call.
- Enabled the main testing function in the script to run without being commented out.
* Enhance sparse_mla_fwd.py and thread_storage_sync.cc with code cleanup and new functionality
- Added a call to disable cache in the main execution block of `sparse_mla_fwd.py` for improved performance.
- Commented out unused benchmarking code in `test_sparse_mla_fwd` to streamline testing.
- Introduced new utility functions in `thread_storage_sync.cc` to analyze and optimize shared memory access patterns, enhancing synchronization management.
- Improved readability and organization of the code by restructuring and commenting on key sections.
* Add thread storage synchronization utilities and optimizations
- Introduced functions to identify shared-like scopes, buffers, and pointer variables, enhancing memory management.
- Implemented a `BarrierTransparentStmtAnalyzer` to analyze statement transparency concerning shared memory.
- Developed a `ThreadStorageSyncOptimizer` to optimize sequences of statements by removing unnecessary synchronization calls, improving performance in multi-threaded environments.
- Enhanced handling of shared memory access in various call types, ensuring proper synchronization during memory operations.
* remove optimize thread_storage_sync
* Refactor pipeline planning to enhance cp.async group scheduling
- Introduced a last-use index for cp.async groups to improve scheduling accuracy.
- Removed unnecessary tracking of first consumers, simplifying the scheduling logic.
- Updated the handling of cp.async and commit statements to ensure proper group ordering based on last-use.
- Added a new test to verify that cp.async groups are ordered correctly by their last-use index.
* lint fix
* Add CUDA requirements to TMA copy tests and update main execution call
- Added decorators to require CUDA and a minimum compute version for the `test_tma_copy_pipeline_2_stages` and `test_tma_copy_pipeline_3_stages` functions.
- Updated the main execution block to call `tilelang.testing.main()` instead of directly invoking the test functions, improving test execution management.
* Refactor attention sink examples by removing pipelined implementations
- Updated imports in `benchmark_gqa_sink_fwd.py` and `benchmark_mha_sink_fwd.py` to use the non-pipelined versions of the attention sink functions.
- Deleted the pipelined implementations of `example_gqa_sink_fwd_bhsd_wgmma_pipelined.py` and `example_mha_sink_fwd_bhsd_wgmma_pipelined.py`.
- Cleaned up regression and test files to remove references to the deleted pipelined examples, ensuring consistency across the codebase.
* performance fix
* Enhance example_blocksparse_gemm.py and inject_pipeline.cc with debugging and utility improvements
- Added a print statement to output the kernel source in `example_blocksparse_gemm.py` for debugging purposes.
- Updated `inject_pipeline.cc` to include a new utility function for stripping TMA copy write buffer attributes, improving code clarity and functionality.
- Modified `producer_consumer_ws.cc` to utilize the new utility function, enhancing the handling of TMA barriers.
- Expanded the `_compile_tvm_ffi` function in the test file to accept additional keyword arguments, improving flexibility for test configurations.
- Introduced a new test to verify that pure TMA warp specialization correctly emits barriers, ensuring compliance with expected behavior.
* lint fix
* Refactor guard variable declarations in producer_consumer_ws.cc for clarity
- Changed guard variable declarations from non-const to const references to improve performance and readability.
- This minor adjustment enhances the handling of producer guards in the context of TMA barriers.
* lint fix
* fix
* lint fix
* perf fix
* Add print statement for kernel source and update main execution in example_dequant_gemm_w4a8.py
- Introduced a print statement to output the kernel source for debugging purposes.
- Updated the main execution block to call `run_regression_perf()` and print the latency, enhancing performance measurement capabilities.
* fix
* Add debugging output and update execution flow in example_tilelang_block_sparse_attn.py
- Added a print statement to display the kernel source for better debugging.
- Updated the main execution block to disable cache and call `run_regression_perf()`, printing the resulting latency for performance measurement.
- Enhanced the producer-consumer logic in producer_consumer_ws.cc with new utility functions to manage local variable access and improve code clarity.
- Introduced a new test case in test_tilelang_issue_tma_no_ws.py to ensure consumer-only local initializations do not leak into the producer, verifying correct behavior in TMA contexts.
* Revert "Add debugging output and update execution flow in example_tilelang_block_sparse_attn.py"
This reverts commit
|
||
|
Lei WangandZhiwen Mo
|
fee041c3cd |
Fix predicated cp.async pipeline scheduling (#1937)
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei WangandZhiwen Mo
|
bdbb32c99b |
Support ptr-table grouped GEMM kernels (#1923)
Add ptr-table grouped GEMM support Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei Wang
|
4929ad89cc |
[Refactor] Improve cp.async lowering and add async_copy op (#1887)
* Add CPAsync copy selection/lowering and vectorized ptx_cp_async folding
* [transform] Rename SimplifyCPAsyncSync to OptimizeCPAsyncSync
* ThreadSync: treat ptx_cp_async as async-copy
* Implement async_copy operator and enhance cp.async handling
- Introduced the async_copy operator for explicit asynchronous global-to-shared memory copies, ensuring no automatic waits are inserted.
- Updated the copy operator to normalize copy regions and handle both synchronous and asynchronous copy semantics.
- Enhanced pipeline planning to recognize explicit cp.async calls and bind commit and wait operations to the appropriate stages.
- Added tests to validate the behavior of async_copy and its integration with pipeline planning.
This commit improves the handling of asynchronous memory operations and clarifies the semantics of copy operations in the context of the new async_copy functionality.
* docs: clarify T.copy vs T.async_copy
* style: apply format.sh fixes
* async_copy: improve errors and defer cp.async check
* test: check T.copy vs T.async_copy codegen equivalence on SM80
* copy: allow predicated cp.async for OOB global->shared
* pipeline: schedule cp.async commits with cp.async groups
* lower: preserve manual pipeline statement arity
* pipeline: skip commit-order check when fused with cp.async
* pipeline: enforce cp.async ordering in PipelinePlanning
PipelinePlanning may float pure control intrinsics like ptx_commit_group earlier because they have no buffer regions, which can generate commit-before-cp_async.\n\nAdd an order-repair step that models cp.async group semantics (cp_async* -> commit, commit(group_i) -> cp_async(group_{i+1})) and wait->consumer constraints, then applies a stable topological sort to produce a valid software_pipeline_order.\n\nAlso adds a regression test to assert cp.async is ordered before commit when split into separate blocks.
* Enhance cp.async synchronization handling in pipeline planning
- Introduced logic to relax loop wait conditions and split epilogue waits for cp.async operations, improving overlap between consumer phases.
- Added a new method to identify and manage consumer phases based on async memory reads, ensuring correct synchronization.
- Updated tests to validate the new behavior, confirming that wait groups are correctly relaxed and inserted as needed in the optimized pipeline.
This commit refines the handling of cp.async synchronization, enhancing performance and correctness in the generated code.
* lint fix
* Refactor cp.async copy handling and improve vectorization logic
- Removed the unused IsCPAsyncElementBytes function to streamline the code.
- Enhanced the CPAsyncStoreRewriter class with a new method to track vectorized lanes during statement rewriting.
- Updated CheckCPAsyncCopy to remove the buffer_oob parameter, simplifying its signature.
- Added logging for vectorized loop iterations to aid in debugging and performance analysis.
- Improved error handling in LowerCPAsyncCopy to provide clearer warnings when cp.async rewrites fail.
These changes aim to optimize the cp.async copy process and improve code maintainability.
* lint fix
* Update cp.async copy handling to improve rewrite success tracking
- Set the successfully_rewritten_ flag to true when a copy operation is successfully rewritten.
- Adjust the flag to false in cases where the load is not global-like or the buffer is not shared.
- Removed logging for vectorized loop iterations to streamline the code.
These changes enhance the tracking of rewrite success in the cp.async copy process, contributing to better error handling and maintainability.
* fix
* test: cover commit-merge + wait relax
* fix: skip wait relaxation under conditional async copy
* test fix
* Implement async copy support in TileLang
- Introduced a new configuration option `tl.enable_async_copy` to control the lowering of global to shared memory copies to PTX `cp.async`.
- Updated the `CopyNode::LowerCPAsyncCopy` method to handle async copy semantics based on the new configuration.
- Enhanced the `InjectPTXAsyncCopy` pass to rewrite eligible global->shared copies into `cp.async`, including necessary commit and wait statements.
- Added tests to verify the correct behavior of async copy lowering and the impact of the configuration option.
These changes improve the flexibility and performance of memory copy operations in TileLang.
* lint fix
* Refactor async copy handling in TileLang
- Renamed `InjectPTXAsyncCopy` to `LowerPTXAsyncCopy` to better reflect its functionality.
- Updated the `OptimizeForTarget` function to use `LowerPTXAsyncCopy` instead of the deprecated `InjectPTXAsyncCopy`.
- Added a deprecated alias for `InjectPTXAsyncCopy` to maintain backward compatibility.
- Adjusted documentation to clarify the purpose of the new function name.
These changes enhance clarity and maintainability in the async copy transformation process.
* fix
* lint fix
* test fix
* Enhance PTX async copy integration in TileLang
- Added the `InjectPTXAsyncCopy` function to facilitate the injection of PTX `cp.async` operations into the lowering process.
- Updated the `LowerTileOp` function to conditionally apply `InjectPTXAsyncCopy` during parallel loop lowering when async copy is enabled.
- Removed deprecated calls to `LowerPTXAsyncCopy` in the `OptimizeForTarget` function to streamline the optimization process.
These changes improve the handling of async copy operations, ensuring better performance and compatibility with CUDA targets.
* lint fix
* fix
* fix
* fix
* set parallisim to 1 to test the ci
* fix
* Refactor async copy handling in PTX lowering
- Improved the `PrepareCopyIndexInfo` function to clarify the calculation of transfer bytes and ensure valid async copy operations.
- Updated synchronization state management for injected `cp.async` transfers, enhancing clarity and consistency in handling open copy regions and uncommitted transfers.
- Added helper functions for zero value checks and linear offset flattening to streamline index processing.
These changes enhance the robustness and readability of the async copy implementation.
* Refactor nested pipeline functions and update test cases
- Removed unused parameters `order` and `stage` from several functions, simplifying their signatures.
- Updated nested pipeline implementations to use `num_stages=2` for better clarity and performance.
- Adjusted test cases to reflect the changes in function signatures, ensuring they run correctly without the removed parameters.
These modifications enhance code readability and maintainability while ensuring the tests remain functional.
* Refactor async copy statement generation in PTX lowering
- Updated the `MakeCPAsyncStmt` function to `MakeCPAsyncStmtFromLoads`, improving the handling of base indices for vectorized accesses.
- Introduced `ExtractVectorBaseIndices` to streamline the extraction of base indices from vectorized loads and stores.
- Enhanced error handling for predicated async copy operations when base indices cannot be extracted.
These changes improve the clarity and robustness of async copy statement generation in the PTX lowering process.
* lint fix
* Enhance async copy operation handling in CodeGenTileLangCuTeDSL
- Added support for conditional async copy operations in the `VisitExpr_` method, allowing for an optional predicate argument.
- Improved error checking to ensure the correct number of arguments for `tl::ptx_cp_async`.
- Streamlined the generation of async copy statements for both standard and conditional cases.
These updates improve the flexibility and robustness of async copy operations in the code generation process.ix
* lint fix
* disable auto cp async without pipeline.
* lint fix
* Enhance cp.async synchronization handling in TileLang
- Introduced `ThreadSyncAfterWaitGroupInserter` to ensure correct insertion of storage synchronization after `ptx_wait_group` statements, addressing potential visibility issues in shared memory.
- Updated `WarpSpecialized` to utilize `tl.ptx_cp_async_barrier_noinc` for releasing producer-consumer dependencies, preventing incomplete memory observations.
- Added tests to verify the correct insertion of synchronization and the handling of cp.async barriers under warp specialization.
These changes improve the correctness and reliability of async copy operations in the TileLang transformation pipeline.
* fix
* fix
* test fix
* Refactor sparse_mla_fwd_pipelined.py and enhance phase.py synchronization
- Removed unused post-processing function and related code from `sparse_mla_fwd_pipelined.py` to streamline the implementation.
- Added debug print statements in `phase.py` to log the state of the module before and after thread synchronization, aiding in troubleshooting and verification of synchronization behavior.
* Enhance async copy handling and loop synchronization in TileLang
- Introduced new loop annotations for async copy operations, including `parallel_prefer_async` and `parallel_async_without_async_commit_wait`, to improve control over async behavior in parallel loops.
- Updated the `InjectPTXAsyncCopy` function to accommodate the new annotations, ensuring correct async copy injection based on user-defined preferences.
- Enhanced the pipeline planning logic to prioritize cp.async groups by their earliest consumers, improving scheduling efficiency and reducing wait times.
- Added tests to validate the new async copy behaviors and ensure correct synchronization in pipelined loops.
These changes enhance the flexibility and performance of async copy operations in the TileLang transformation pipeline.
* lint fix
* Add Regression Perf for mhc
* clean
* clean
* lint fix
* lint fix
|
||
|
Lei Wang
|
7ab2b2be92 |
[BugFix] ArgBinder: relax shared-shape binding for unused nullable buffers (#1870)
* [BugFix] tvm_ffi: handle unused nullable buffers with shared symbolic shapes * [BugFix] ArgBinder: relax shared-shape binding for unused nullable buffers |
||
|
Lei Wang
|
9f259544d4 |
[Refactor] Refactor Pass InjectFenceProxy (#1863)
* tir: add T.cdiv alias for T.ceildiv * docs: Update InjectFenceProxy documentation and enhance code comments - Clarified the description of `tl.InjectFenceProxy` to specify the transition between generic and async proxy operations. - Improved explanations of the pass's functionality, including state tracking and the handling of TMA store synchronization. - Added details on the new `ProxyStateSet` class and its role in managing proxy states. - Updated usage instructions for proxy hints to include new options for custom operations. - Enhanced test coverage for handling unknown external calls and proxy hint overrides. * fix * enhance * fix * remove tl.proxy_hint * refactor tma_store_arrive and tma_store_wait * refactor * fix * InjectFenceProxy: hoist fence out of pure-async loops * InjectFenceProxy: hoist fence for if/while pure-async regions * Remove LowerTileOp transform test |
||
|
Lei Wang
|
4956b5835f | ThreadSync: avoid barriers between atomic ops (#1852) | ||
|
Lei Wang
|
4ff81c7d40 |
[Refactor] Introduce T.access_of to combine T.address_of and access_ptr (#1827)
* [Refactor] Update atomic operations to use access_ptr for buffer access * Replaced instances of address_of with access_ptr in atomic operations to improve clarity and maintainability. * Enhanced atomic operation implementations in various files, including atomic_add, atomic_max, and atomic_min. * Updated related logic in loop vectorization and tile language to support new access patterns. * Added a utility function to extract destination scalar dtype for atomic operations, supporting both BufferLoad and tvm_access_ptr. * Improved kernel cache logic to include a build-stamp for the TileLang runtime library, ensuring cache validity during development. * lint fix * Fix access_ptr layout rewrite and atomic vectorization * lint fix * Refactor atomic operations to remove const qualifiers from pointer parameters in atomicAddx2 and ToFloat functions. This change enhances type flexibility and consistency across atomic operations. * Refactor buffer access logic in lower_tile_op and vectorize_loop to enhance clarity and maintainability. Removed unused code and improved offset calculations for tile access. Updated atomic operation handling to streamline dtype extraction and ensure correct vectorization behavior. * Add logging for loop layout and lowered loop in atomic_add; implement tvm_access_ptr handling in loop_vectorize and lower_tile_op This commit introduces logging statements to track the loop layout and the result of the lowered loop in the AtomicAddNode. Additionally, it implements a new method to handle tvm_access_ptr calls in the loop_vectorize and lower_tile_op files, enhancing the handling of memory access patterns and ensuring correct vectorization behavior. * fix * lint fix * fix * Update ToFloat2 and ToFloat4 functions to use const qualifiers for pointer parameters in atomic.h; modify test script to call test_atomic_addx2_float instead of main. * Refactor tvm_access_ptr handling in loop_vectorize to improve argument validation and buffer resolution logic. Ensure at least three arguments are provided and streamline the search for matching buffers in layout_map_. Update vector size calculations to enhance clarity and maintainability. * Add handling for tvm_access_ptr in vectorize_loop to support offset substitution and buffer creation This commit introduces logic to handle tvm_access_ptr calls, allowing for the substitution of loop variables in offsets and the creation of dummy buffers. The new method ensures correct base address calculation and enhances the overall handling of memory access patterns in vectorization. * Add handling for tvm_access_ptr offset validation in loop_vectorize This commit enhances the loop_vectorize functionality by adding checks for the element offset in tvm_access_ptr calls. It ensures that offsets are invariant within the vector boundary, improving the robustness of vectorization logic and memory access patterns. * fix * Enhance tvm_access_ptr handling in loop_vectorize by adding offset validation checks. This update improves the robustness of vectorization logic and ensures offsets remain invariant within vector boundaries. * Enhance tvm_access_ptr handling in loop_vectorize by adding offset validation checks. This update ensures offsets remain invariant within vector boundaries, improving memory access robustness and vectorization logic. * fix * Add tl.access_ptr support in TileLang This commit introduces the `tl.access_ptr` operation, which carries pointer access metadata in the frontend. It allows for better handling of buffer loads by retaining necessary information for downstream analysis. The new operation is integrated into various components, including the loop vectorization and memory access legalizations, ensuring compatibility with existing functionality while enhancing the robustness of memory access patterns. * lint fix * fix * Add support for tl.access_ptr in atomic operations This commit updates the AtomicAddNode and AtomicOpBaseNode implementations to utilize the new tl.access_ptr operation for building access pointers to destination elements. It enhances the handling of buffer loads by incorporating type checks for index data types, improving the robustness of memory access patterns. Additionally, it removes unnecessary logging statements in the lowering process, streamlining the codebase. * lint fix * Refactor buffer load extraction in vectorize_loop to check for non-empty arguments in access_ptr calls. This change improves the robustness of buffer load handling by ensuring valid argument presence before processing. * Remove unused fallback logic in KernelCache and clean up import statements in test_tilelang_language_access_ptr.py for improved code clarity. * Refactor memory access in BitNet and dequantization examples to use tl.access_ptr This update replaces T.address_of calls with T.access_ptr in multiple files, enhancing memory access handling by specifying read and write permissions. This change improves code clarity and aligns with recent enhancements in pointer access management. |
||
|
Lei WangandCopilot
|
f8b4d9e677 |
[Bugfix] Fix threadIdx variable lookup by thread_tag instead of position in ThreadSync (#1829)
* Refactor OptimizeForTarget function by removing debug print statements and updating module state visualization * Add test for sync hoisting in non-uniform if within loop using shared memory This commit introduces a new test case to verify the correct behavior of sync hoisting when a non-uniform if statement is present inside a loop that utilizes shared memory. The test ensures that the synchronization occurs before the if statement, confirming the expected transformation in the module's intermediate representation. * fix * Refactor thread variable handling in thread_storage_sync.cc This commit removes the optional thread variable lookup by tag and replaces it with direct indexing based on the expected position of threadIdx variables. This change addresses a bug where the last three elements of env_threads_ were incorrectly assumed to always correspond to threadIdx.x/y/z, improving the accuracy of thread variable access in the synchronization logic. * Update testing/python/transform/test_tilelang_transform_thread_sync.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> |
||
|
Lei Wang
|
e9d056954c |
[BugFix] LoopUnswitching: gate non-trivial else behind PassConfig (#1816)
* [BugFix] LoopUnswitching: gate non-trivial else behind PassConfig * lint fix |
||
|
Kuris
|
c47d6df379 |
[BugFix] Missing Recursive Loop Var Checking in Loop Unswitching (#1801)
* [Fix] Update loop unswitching logic to handle multiple let bindings and add corresponding test case * remove debug print |
||
|
Lei Wang
|
191d879ddd |
[Bugfix] Fix thread storage sync conflict detection for loop carry write-after-read (#1781)
* Fix thread storage synchronization logic in `thread_storage_sync.cc` to correctly identify conflicts between read and write operations based on loop carry conditions. * lint fix * Refactor `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` to use shared memory for `sorted_token_ids` instead of local memory, improving thread synchronization. Adjust default argument values for M, N, and K in the main function for better testing scenarios. * Add UniformExprChecker to enforce thread synchronization rules Introduce the UniformExprChecker class to determine if expressions are uniform across threads, crucial for safe synchronization in conditional statements. Update the TileLangThreadSyncPlanner to hoist synchronization points out of non-uniform if-statements to prevent potential deadlocks. Enhance tests to validate sync hoisting behavior for various non-uniform conditions involving thread indices and shared memory access. * lint fix * Enhance `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` with cache disabling and kernel source printing for debugging. Update thread synchronization logic in `thread_storage_sync.cc` to check for runtime-dependent conditions, preventing potential deadlocks by hoisting sync points as necessary. * Update submodule `tvm` to latest commit and remove deprecated `example_gqa_decode_varlen_logits_paged.py` file. Refactor `example_gqa_decode_varlen_logits.py` to enhance performance and maintainability by removing unused imports and optimizing shared memory usage. Adjust test cases to reflect the removal of the paged example. * fix * Enhance thread synchronization logic in `thread_storage_sync.cc` by adding a configurable warp size parameter to `RuntimeDependentConditionChecker` and `TileLangThreadSyncPlanner`. This allows for better adaptability to different target architectures. Update the logic to ensure thread extent is a constant and improve handling of runtime-dependent conditions. * lint fix * Refactor thread extent validation in `thread_storage_sync.cc` to use pointer checks instead of optional values. This change improves clarity and ensures that the thread extent is correctly validated as a constant. * Adjust loop variable constraints in `thread_storage_sync.cc` for loop-carry analysis by modifying the extent calculation. This change ensures valid iteration comparisons by reducing the extent by one, allowing for accurate analysis of loop iterations. * lint fix |
||
|
Chaofan Lin
|
378a8f2573 |
[Refactor] Unify the usage of cast-related operators (#1757)
* [Refactor] Unify the usage of cast-related operators * reinterpret auto detect |
||
|
Lei Wang
|
7bf43447d5 |
[Feature] Add PassConfig for Controlling Let Statement Inlining in Simplify Pass (#1769)
[Feature] Enhance Simplify Pass with Let Inline and Common Subexpression Elimination - Added a new configuration option `tilelang_enable_simplify_let_inline` to enable inlining of let statements during simplification. - Updated the Simplify pass to conditionally inline let statements based on the new configuration. - Integrated common subexpression elimination into the device code generation process to optimize performance. - Enhanced test coverage for the Simplify pass, including various scenarios to validate simplification behavior. |
||
|
Lei Wang
|
c4748da9d8 |
[Bugfix] Add predicate to loads inside predicated stores in LowerLDGSTG pass (#1767)
* [Feature] Add support for predicated loads and stores in LDGSTG lowering - Introduced a new configuration option to disable loop unswitching. - Enhanced the lowering logic to ensure that loads within predicated stores are also treated as predicated, preventing out-of-bounds memory access. - Added a test case to verify the correct behavior of predicated loads and stores in the transformation pipeline. - Included debug print statements to visualize the module state before and after the lowering process. * lint fix |
||
|
Lei Wang
|
5f60c09060 |
[Feature] Enhance Loop Unswitching with Let Binding and Condition Handling (#1766)
- Updated the loop unswitching pass to support hoisting Let-bound variables alongside their conditions, improving optimization for loops with variable dependencies. - Modified the CallCheckerExcludingIf class to exclude matching If nodes when checking for CallNode presence, ensuring safety during unswitching. - Enhanced IfBranchReplacer to replace all instances of If nodes with matching conditions, streamlining the transformation process. - Added comprehensive tests to validate the new behavior, including scenarios with multiple Let-bound variables and identical conditions across If statements. |
||
|
Yu ChengandLeiWang1999
|
8c45417584 |
[Feature] Support Pass LoopUnswitching (#1747)
* [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 --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Lei Wang
|
fbf28a143f |
[ThreadSync] Use Z3 for constraint equivalence checking (#1760)
* [ThreadSync] Use Z3 for constraint equivalence checking Replace const_int_bound based range comparison with Z3 SMT solver for checking thread constraint equivalence in FindConflict. * [ThreadSync] Clarify conflict detection for constraint equivalence Added comments to improve understanding of the logic determining whether constraints are equivalent or in conflict in the thread storage synchronization process. * add test * Update test_issue_1106 to check syncthreads is outside for loop * test fix |
||
|
|
fe5bf3569d |
[Refactor] Enhance T.alloc_barrier with new features and deprecate legacy mbarrier related intrinsics (#1733)
* enhance `T.alloc_barrier` and let `T.wait/arrive` only accept buffers * lint * Refactor mbarrier handling and update related documentation - Updated the `T.alloc_barrier` function to accept an array of arrive counts instead of individual thread counts. - Modified the `get_mbarrier` function documentation to clarify its purpose and usage. - Enhanced the `CreateMBarrierBuffer` function to create mbarrier buffers with shared storage scope. - Updated tests to reflect changes in mbarrier allocation and usage. - Adjusted comments in various files to improve clarity regarding mbarrier functionality. * lint * fix test * tcgen5 check input * lint * fix name conflict * add name check * apply ai suggestions * fix shared mem disalignment and missing fence for mbarrier init * fix cutedsl codegen * bugfix * add test for cutedsl codegen * Update testing/python/jit/test_tilelang_jit_cutedsl.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> * Remove tilelang.disable_cache() call Remove cache disabling from the test file. --------- Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com> |
||
|
Lei Wang
|
350f987a5b |
[Refactor] Disable Predicated LDG PTX Lowering by default (#1739)
[Refactor] Update predicated global load/store lowering configuration - Changed the configuration option from `kDisableLowerLDGSTGPredicated` to `kEnableLowerLDGSTGPredicated` to enable predicated lowering of global load/store operations. - Updated relevant comments and documentation to reflect the new configuration behavior. - Adjusted tests to utilize the new configuration option for enabling predicated lowering, ensuring consistency across the codebase. |
||
|
Lei Wang
|
8bd0293740 |
[Enhancement] Add explicit global memory load/store intrinsics (ldg/stg 32/64/128) (#1717)
* [Enhancement] Add global memory load/store intrinsics for various vector widths (#1710) * Introduced new intrinsics: ldg32, ldg64, ldg128, ldg256 for loading data from global memory. * Added corresponding store intrinsics: stg32, stg64, stg128, stg256 for storing data to global memory. * Updated built-in functions and templates to support these new operations. * Enhanced documentation for usage examples and type checks in the language module. * Updated CUDA code generation to utilize new intrinsics for improved performance. * lint fix * test fix * Refactor buffer type annotations in global memory load/store intrinsics Updated the type annotations for the `ldg` and `stg` functions to use `tir.Buffer` instead of the generic `Buffer` type. This change enhances type specificity and aligns with the TVM type system for better clarity and correctness in memory operations. * lint fix * [Enhancement] Introduce LDG/STG lowering options and transform pass * Added configuration options to enable/disable lowering of non-predicated and predicated global memory loads/stores to ldg/stg intrinsics. * Implemented the LowerLDGSTG transformation pass to optimize global memory access for better performance. * Updated tests to require a minimum CUDA compute version for relevant functionality. * Enhanced documentation for new configuration options and transformation behavior. * lint fix * Refactor LDG/STG intrinsic selection logic for improved readability * Replaced switch-case statements with if-else conditions for determining the appropriate LDG operation based on byte width. * Enhanced error handling for unsupported byte widths with a clear fatal log message. * lint fix * fix * testfix * Enhance LDGSTG pass to skip processing for CuTeDSL backend * Added a check to skip the transformation if the target has the "cutedsl" key. * Updated tests to require a minimum CUDA compute version of 10.0 for LDG and STG code generation tests. * Refactor global memory load/store functions to include conditional variants * Updated the CUDA code generation to support conditional versions of LDG and STG functions, enhancing flexibility in memory operations. * Modified existing functions to include new conditional variants (e.g., `ldg32_conditional`, `stg32_conditional`) for better control over memory access based on predicates. * Adjusted comments to reflect the changes in function signatures and usage. * Update comments and input handling for LDG intrinsics to include address and optional predicate * Modified comments for `ldg32`, `ldg64`, `ldg128`, and `ldg256` to clarify the function signatures, indicating the use of 'address' and an optional 'predicate'. * Changed the number of inputs for these LDG functions from 2 to -1 to accommodate the new signature format. * Add correctness verification to LDG code generation tests * Introduced assertions to verify the correctness of generated CUDA source for `ldg32`, `ldg64`, `ldg128`, and `ldg256` functions. * Added checks to ensure the output matches the expected reference values within specified tolerances. * Updated test cases to enhance validation of LDG intrinsic functionality. * test fix * Refactor LDG and STG intrinsics to use unified naming convention * Updated all instances of `ldg` and `stg` functions to their new names: `load_global` and `store_global`, respectively, across the codebase. * Modified related test cases to reflect the new function names, ensuring consistency in the testing framework. * Enhanced comments and documentation to clarify the changes in function signatures and usage. * Update reference values in LDG256 predicated code generation test to improve accuracy * lint fix * Rename test function for clarity and add new test for predicated load/store global intrinsics. Update assertions to reflect new naming convention in generated kernel source. * lint fix * subtype cast fix |
||
|
Lei WangandZhiwen Mo
|
56dbc9f8f0 |
[Refactor] Enhance variable substitution in device function generation (#1723)
* [Refactor] Enhance variable substitution in device function generation - Updated `wrapBodyWithHostSideAssumes` to accept a mapping of host-side variable names to device-side variables, ensuring consistent variable references during SSA conversion. - Introduced a substitution mechanism for old parameter variables to new ones in the device function, preventing variable sharing between host and device. - Adjusted buffer declarations to utilize new variable mappings, enhancing the robustness of the device function generation process. * lint fix * test fix * lint fix --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei Wang
|
2ff1dd50dc |
[Feature] Introduce DecoupleTypeCast pass for mixed-precision vectorization (#1644)
* [Feature] Introduce DecoupleTypeCast pass to optimize vectorization constraints - Added a new pass to decouple type cast vectorization constraints, allowing for optimal vectorization in mixed-precision scenarios. - Implemented transformations for local-to-memory and memory-to-local cases, inserting intermediate local buffers as needed. - Updated pass configuration to enable verbose output for vectorization planning. - Enhanced loop vectorization logic to handle complex cases with multiple statements. - Added tests to validate the new functionality and ensure correctness in various scenarios. * Refactor verbose output in loop vectorization for improved readability - Updated the output formatting in the VectorizePlanner to use newline characters instead of std::endl for better performance and clarity. - Ensured consistent output style across various verbose logging statements related to buffer sizes and vectorization strategies. * Refactor verbose output in loop vectorization for consistency - Simplified the output formatting in the verbose logging of vector size to enhance readability by removing unnecessary line breaks. * use eager style. * lint fix * Enhance loop vectorization by introducing GetBufferStrides function - Added GetBufferStrides function to compute buffer strides when not defined, improving clarity and reusability. - Updated VectorizeFindGlobalAccess to utilize the new function for stride calculations, simplifying the logic. - Refactored related code to remove redundant stride calculations and improve overall readability. - Introduced a test case to ensure no transformation occurs when different dtypes are only in if-then-else conditions, validating the new behavior. * Refactor loop vectorization logic and enhance target support checks - Updated the logical condition in TargetIsSm100 to use the correct logical AND operator. - Introduced TargetSupportVectorize256 function to encapsulate vectorization support checks for specific architectures. - Refactored VectorizeFindGlobalAccess to VectorizeFindMemoryAccess for improved clarity and added support for shared memory access checks. - Updated vectorization logic to utilize the new support checks, enhancing overall performance and maintainability. - Cleaned up redundant logging statements and improved code formatting for better readability. * Update TargetSupportVectorize256 to simplify architecture check - Modified the architecture check in TargetSupportVectorize256 to return true for architectures 100 and above, removing the upper limit of 110. This change enhances compatibility with newer architectures while maintaining existing functionality. * lint fix * Refactor tilelang JIT configurations and optimize GEMM calls in fused MoE example - Updated JIT pass configurations to use enum keys for better clarity. - Removed unused parameters `k_pack` and `coalesced_width` from the `moe_forward_tilelang_routed` function. - Simplified GEMM calls by removing the `k_pack` argument, enhancing readability. - Adjusted swizzle usage to remove the `enable` flag for consistency. - Added a call to `tilelang.disable_cache()` in the main function to improve performance. * lint fix * Update TVM submodule and enhance vectorization logging - Updated the TVM submodule to the latest commit for improved functionality. - Added logging in the VectorizePlanner to provide insights during the planning phase of for loops. - Ensured GCD calculation for vector size aligns with loop extent for better vectorization accuracy. - Enabled verbose logging for the vectorize planner in the test configuration to aid debugging. * lint fix * performance fix * fix * fix * fix * fix * bench fix |
||
|
Lei Wang
|
b27fb9286a |
[Feature] Atomic Reduction Operations and Vectorization Enhancement (#1676)
* [Feature] Implement atomic reduction operations and enhance atomic add functionality - Added support for atomic max and min operations with corresponding classes and methods. - Introduced vectorization capabilities for atomic add operations based on data type and target architecture. - Refactored atomic add intrinsic calls to improve clarity and consistency. - Enhanced layout inference for atomic operations and integrated new utility functions for target architecture checks. - Updated tests to validate the new atomic operations and their vectorization behavior. * Refactor atomic operation code for improved readability and consistency - Cleaned up whitespace and formatting in atomic add and reduce implementations for better code clarity. - Enhanced comments in CUDA code generation for atomic operations to improve understanding. - Updated function signatures and layout definitions for consistency across atomic operation files. - Ensured proper alignment of code for better maintainability and readability. * lint fix * Enhance loop vectorization and layout handling across multiple operations - Updated function signatures in atomic add, atomic reduce, copy, and fill operations to include layout map parameters for improved layout handling. - Refactored vectorization logic to utilize layout maps, ensuring better performance and compatibility with various buffer layouts. - Enhanced the LowerParallelLoop function to accept layout maps, facilitating more efficient loop transformations. - Added checks for buffer contiguity in vectorization processes to ensure correctness when using layout maps. - Updated tests to validate the new layout handling and vectorization behavior. * Add atomic operation checks in legalize_safe_memory_access - Introduced a new method to identify atomic operations within the legalize_safe_memory_access transformation. - Updated the VisitStmt_ function to handle both CallExtern and atomic operations, ensuring recursive condition collection for these cases. - Enhanced comments for clarity on the handling of atomic operations in the context of memory access legality. * Add atomic operations for HIP code generation - Implemented support for various atomic operations including atomic add, atomic load, atomic store, atomic max, and atomic min in the HIP code generation. - Enhanced the handling of atomic operations to include optional memory order parameters. - Improved code readability with added comments explaining the purpose of each atomic operation. - Ensured consistency with existing atomic operation implementations in the codebase. * Refactor AtomicOpBaseNode and related classes for improved clarity - Changed the GetElemOpStatic method to a virtual GetElemOp method in AtomicOpBaseNode, enhancing polymorphism. - Updated AtomicAddNode, AtomicMaxNode, and AtomicMinNode to override the new GetElemOp method. - Removed unnecessary template parameters from AtomicOpBaseNode, simplifying the class structure. - Cleaned up includes in atomic_reduce.cc to remove unused dependencies, improving code organization. * Enhance argument validation for atomic operations - Added checks in AtomicAdd, AtomicMax, and AtomicMin constructors to ensure at least two arguments are provided, improving error handling and user feedback. - Removed the unused LowerTMA method declaration from the AtomicAdd class, streamlining the codebase. * lint fix * fix * test fix * performance fix * remove useless comments |
||
|
silentCoder-devandLeiWang1999
|
1d71bed2db |
[Feature] Reimplement Threadsync with ConstrVisitor (#1631)
* reimplement ThreadSync pass * fix bugs * remove old storage_access * remove deadcode for debugging * format * add WhileOp & expose MakeGuard in constr_visitor * bugfix * bugfix * bugfix * bugfix * trigger ci * add try for checking * bugfix * update tvm with better analyzer * update tvm * Log for debugging in thread_storage_sync * format * consider let node * typo * remove debugging log * consider while op in thread_sync & fix bug sbout letnode * add testing for thread_sync issues * update testing for thread_sync * format * consider T.Ramp expr * remove useless header * format --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Lei Wang
|
a56212de10 |
[Bugfix] Introduce a flag to avoid unnecessary broadcast hoist and enable for let stmt (#1638)
* [Enhancement] Implement hoisting of broadcast values in Let statements and update tests for validation * lint fix * [Enhancement] Update hoisting logic for broadcast values in Let statements and add a new test case for nested BufferStore broadcasts * lint fix |
||
|
silentCoder-devandkurisu6912
|
b5be1a16ab |
[Feature] Add hoist_broadcast_values pass (#1606)
* add hoist_broadcast_values pass * only hoist intImm or floatImm * refactor * add test for hoist_broadcast_values * add a test for issue 1601 * add transform test for hoist_broadcast_value --------- Co-authored-by: kurisu6912 <227995639+kurisu6912@users.noreply.github.com> |
||
|
Kuris
|
5d806298de |
[Feat] Add strong checker to detect data racing in T.Parallel (#1615)
* add verifier for parallel loop data racing * fix lint error * fix typo * fix lint error * fix bug in metal test |
||
|
Lei Wang
|
6fd0bd31fe |
[Refactor] Use access_ptr instead of buffer and offsets for cp async params (#1590)
* refactor params of ptx_cp_aync * Enhancement: Introduce ptx_cp_async intrinsic and improve async copy handling - Added the `ptx_cp_async` intrinsic for asynchronous memory copy operations in TileLang, supporting both predicated and non-predicated versions. - Updated the code generation for `ptx_cp_async` to handle offsets correctly and added inline functions for simplified usage. - Enhanced the `inject_ptx_async_copy` transformation to support predicated vectorized copies. - Updated example scripts to disable caching and print kernel sources for better debugging and performance insights. * Remove tilelang.disable_cache() calls from example scripts for cleaner execution * Refactor comments and formatting in CUDA code generation and PTX async copy transformation - Updated comments for clarity in `codegen_cuda.cc` regarding `ptx_cp_async` arguments. - Improved formatting in `inject_ptx_async_copy.cc` for better readability of access pointer creation and logging. - Added a blank line in `op.py` for consistency in the `ptx_cp_async` function definition. * Remove unused ptx_cp_async handling from CUDA code generation - Eliminated the implementation for `ptx_cp_async` and `tl::ptx_cp_async` from `codegen_cuda.cc`, streamlining the code by removing unnecessary checks and assembly generation. - This change focuses on cleaning up the codebase and improving maintainability by removing deprecated functionality. |