mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
256fbc6e60cbcfce1b72cbc9aa3aa7363ef2ab53
433
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Lei Wang
|
256fbc6e60 |
[Pipeline] Fix guarded TMA pipeline handling (#2224)
* Fix guarded TMA pipeline handling * Apply ruff fixes for latest main |
||
|
KurisandLeiWang1999
|
719d23dce5 |
feat: auto-vectorize bf16/fp16 reduce with packed add2 intrinsics (#2112)
* fix: add missing TvmLogDebugSettings::ParseSpec and VerboseEnabledImpl for TVM_LOG_CUSTOMIZE builds When TVM_LOG_CUSTOMIZE=1, TVM's logging.cc skips compiling TvmLogDebugSettings::ParseSpec and VerboseEnabledImpl (guarded by #if TVM_LOG_CUSTOMIZE == 0). However libtilelang.so calls these functions via the inline FromFlag(), causing a runtime symbol lookup error. Add the missing implementations to tilelang's own logging.cc. * feat: auto-vectorize bf16/fp16 reduce with packed add2 intrinsics * style: clang-format * fix: fp16 batch pack — fix Shuffle codegen taking address of temporary - Add tl::pack_half2 helper in common.h that returns uint1, avoiding taking address of a temporary __pack_half2 return value - Update ShuffleNode codegen for fp16x2 to use pack_half2 instead of &(__pack_half2(...)) - Re-enable fp16 batch pack (remove bf16-only restriction in can_batch_pack) - Add SumOp_fp16x2/MaxOp_fp16x2/MinOp_fp16x2 packed reducer structs * refactor: simplify reduce packing API - Replace ReducePackConfig with simple GetPreferedVectorizedSize(DataType, Target) -> int - Merge getBatchReducerSuffix into MakeCodegenReducer(vsize) - MakeReduce(vsize, acc, b) -> optional<PrimExpr> - MakeInitValue(vsize) handles Broadcast internally - Guard nan_propagate in MakeCodegenReducer vsize>1 path (no packed nan ops exist) * fix: guard nan_propagate in MakeReduce packed path, mirror scalar structure * feat: support nan_propagate in packed reduce (max2_nan/min2_nan) * fix: add bitwise reduce ops to MakeReduce vsize=1 path * fix lint error --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
|
f11954cb36 |
[Backend] Refactor gemm_sp (#2048)
* [Refactor] refactor gemm_sp following
|
||
|
Leonard
|
c2a9d9872d |
[Fix][AMD] Fix Roller autotuner for RDNA3 WMMA targets (#2208)
* Tune RDNA WMMA Roller defaults * Fix matmul Roller benchmark callback * lint fix |
||
|
Chenhao Xu
|
fa76eaf014 |
Allow variable barrier id in T.sync_threads() (#2197)
* fix: allow variable barrier id in T.sync_threads() * ruff format |
||
|
ZihaoMu
|
7bf8de13ca |
[ROCm] Try to fix ROCm CI error (#2179)
* Fix ROCm tvm-ffi DLPack handling
* Fix CuTeDSL FP4 torch dtype map
* Fix CuTeDSL FP4 packed stores
* Mark DeepSeek V4 FP4 CuTeDSL xfail
* Keep ROCm DLPack PR focused
* Revert "Keep ROCm DLPack PR focused"
This reverts commit
|
||
|
|
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> |
||
|
Leonard
|
d863f1d89d |
Add RDNA gfx1151 ROCm target support (#2127)
* Add RDNA gfx1151 ROCm target support * Fix optional annotation handling for gfx1100 * Address CodeRabbit ROCm cleanup feedback * Scope MFMA tests to CDNA targets * Scope gfx950 copy async GEMM test * Reject non-divisible warp layouts * Skip CUDA/CDNA GEMM dtype tests on RDNA * Scope RDNA device model to gfx11 * Use RDNA arch helper in block scoring * Make RDNA tensor intrinsics generation-aware * Exclude gfx1010 from SDOT4 guard * Store RDNA tensor instructions as list * Reject unsupported RDNA carver targets |
||
|
ZihaoMu
|
beef5cf471 | Fix float4 storage dtype torch mapping (#2174) | ||
|
Tong WUandRachmanino
|
1fbd994cbb |
[Refactor] Refactor multiple TensorCoreIntrinEmitter to provide atom-level mma control interface (#2161)
* Enhance intrinsic modules by adding WGMMATensorCoreIntrinEmitter and TCGEN05TensorCoreIntrinEmitter. Introduce descriptor parameter classes for WGMMA and TCGEN05, improving the initialization and atom offset computation processes. Update the TensorCoreIntrinEmitter to support atom-level interfaces and macro definitions for efficient matrix multiplication operations. * lint * add test for mma/wgmma/tcgen5mma atom * lint * typo * fix import path err --------- Co-authored-by: Rachmanino <wutong.1109@bytedance.com> |
||
|
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 |
||
|
Lei Wang
|
0fdd0f8a60 |
[Refactor] Move backend stubs out of codegen (#2164)
* Move backend stubs out of codegen * lint fix |
||
|
SepcntandSiriusNEO
|
c2cabd9b3f |
[Feature] Add full Windows support and fix related cross-platform issues (#2093)
* [Feature] Add Windows support with cross-platform minimal diff * [Feature] Add ccache support for Windows * [Feature] Add parallel autotune support for Windows * [Feature] Add nvrtc support for Windows * [Bugfix] Fix tilelibrary support for Windows * [Bugfix] Fix alignment for Windows * [Feature] Add clang-cl support for Windows * [Bugfix] Fix numeric limit for Windows * [Bugfix] Fix rand for Windows * [Misc] Use dynamic test workload to avoid hang * [Bugfix] Fix TMA alignment for Windows and sm120 * [Bugfix] Add unix-like hex formatter for Windows * [CI] Add Windows wheel job * [Feature] Flexible CUDA loader * fix rebase --------- Co-authored-by: SiriusNEO <chaofan@deepseek.com> |
||
|
|
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 |
||
|
Kasper Nielsen
|
d135bd1cd2 |
fix: TMA alignment to 1024 bytes on Blackwell (#2134)
* fix: TMA alignment to 1024 bytes on Blackwell Previously 1024 byte alignemnt was guarded only by `IsHopper`. This commit changes it to be guarded by `TargetHasBulkCopy` with the same threshold as Hopper is using. The kernel added in the test was failing on my RTX5090 GPU. * Fix lint |
||
|
Lei Wang
|
9452312af2 |
[codex] Fix 1D TMA store layout inference (#2137)
fix tma store 1d layout |
||
|
Zhang Jasonandroot
|
4639c279ea |
[AMD][CI issue] add gfx950 guard to fix the CI issues (#2105)
* add gfx950 guard to test_gfx950_cp_async_gs_16_in_codegen * add gfx950 checking and update other cases for CI issues * update * [AMD][gfx950] Fix async-copy swizzle conflict, re-enable pipelining, fix 32x32 MFMA store layout * update for format checking * upgrade CI rocm sdk version to be rocm7.2 * update the changes according to the review info * update the codes for format checking * onlu run copy async on gfx950 * move T.alloc_shared and T.alloc_fragment outside the T.Pipelined loop * update the test case * revert the change to pyproject.toml * update test_tilelang_hip_codegen * find a possible issue about failed CI test --------- Co-authored-by: root <root@dell300x-ccs-aus-f06-01.cs-aus.dcgpu> |
||
|
|
6548c05f0c |
[Feature] Batched AllReduce for better T.reduce performance (#1976)
* [Feature] Batched AllReduce for better T.reduce performance (#1928)
When each thread holds multiple values to reduce, the previous code
emitted one AllReduce call per element inside a loop, causing redundant
synchronization barriers (e.g. 8 values × 6 syncs = 48 syncs).
This change adds batch support to AllReduce so all per-thread values are
reduced in a single butterfly pass sharing barriers (6 syncs total).
- Add `batch_size` and `workspace_stride` template parameters to
`tl::AllReduce` (defaults preserve existing scalar behaviour).
- In ReduceOpNode::Lower, detect when batching is beneficial
(batch_size > 1 and cross-warp reduction) and restructure into three
phases: loop(init + local reduce), batched AllReduce, loop(copy-back).
- Apply same optimisation to FinalizeReducerOpNode::Lower.
- Mirror template changes in HIP AllReduce.
Closes #1928
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Test] Add codegen tests for batched AllReduce
Verify that the batched AllReduce path (batch_size > 1, workspace_stride)
is emitted in the generated CUDA source for reduce_max, reduce_sum, and
reduce_min with cross-warp reductions.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Bugfix] Fix batched AllReduce codegen for HIP/ROCm targets
HIP AllReduce template has no Barrier type parameter, so the batched
path must omit tl::SyncThreadsBarrier for ROCm targets. Adds
TargetIsRocm() checks in both reduce.cc and finalize_reducer.cc.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Fix] Use target-aware warp size for cross-warp reduction threshold
Address CodeRabbit review: ROCm wavefronts are 64-wide, so the cross-warp
threshold should be 64 (not 32) on HIP targets. This avoids unnecessary
workspace allocation and batched codegen for single-wavefront reductions.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Style] Fix ruff format for test assertion
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Bugfix] Support batched AllReduce in CuteDSL Python codegen
The CuteDSL codegen translates C++ AllReduce template calls to Python.
With the new batch_size and workspace_stride template params, the
generated Python calls AllReduce() with 7 args instead of 5. Add
batch_size and workspace_stride params to the Python AllReduce function
and implement batched reduction logic for both standard and Hopper paths.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
* [Feature] Add explicit batch parameter to T.reduce for user-controlled batched AllReduce
Expose a `batch` parameter on T.reduce (and all T.reduce_* wrappers,
default 1 = scalar / current behaviour). When batch > 1 the compiler:
- validates batch divides the per-thread output element count N;
- emits ceil(N/batch) calls to the new AllReduce::run_batch interface,
each sharing a single pair of barriers across all batch elements;
- allocates workspace of size reducing_threads * batch (only for cross-
warp reduce, i.e. reducing_threads > 32).
Key differences from the existing auto-detection approach:
* The user controls when batching is enabled — no silent smem increase.
* workspace_stride = reducing_threads (not block_size), so workspace is
minimal and does not blow past the GPU shared-memory limit.
* run_batch (not run) avoids C++ overload-resolution ambiguity when a
pointer is passed as the first argument.
Files changed:
- src/op/reduce.h: add `int batch{1}` field to ReduceOpNode
- src/op/reduce.cc: read "batch" annotation; validate; emit N/batch
batched AllReduce calls with address_of per-chunk pointers
- src/tl_templates/cuda/reduce.h: rename batch run() -> run_batch(),
fix recursive call
- src/tl_templates/hip/reduce.h: same
- tilelang/language/reduce_op.py: add batch param to reduce() and all
reduce_max/min/sum/absmax/abssum/bitand/bitor/bitxor wrappers
- testing/python/language/test_tilelang_language_reduce.py: update
test_batched_allreduce_codegen to use explicit batch=2 and match
run_batch in generated source
* [Fix] Use run_batch in finalize_reducer batched AllReduce codegen
finalize_reducer.cc also generates batched AllReduce calls (used by
T.alloc_reducer / T.finalize_reducer), but was still emitting ::run
instead of ::run_batch, causing a compiler error because the batched
pointer-based interface is now named run_batch to avoid overload
ambiguity.
* rocm fix
* [Feature] Add explicit batch parameter to T.finalize_reducer
Mirrors the batch= API already available on T.reduce:
- batch=1 (default) keeps the existing scalar AllReduce::run path
- batch>1 emits AllReduce::run_batch, sharing barriers across batch
output elements and reducing barrier count by batch×
C++ changes:
- FinalizeReducerOpNode gains an int batch{1} field (registered via
reflection) and reads it from the 'batch' annotation in the constructor
- Lower() validates batch against the layout output-element count and
uses effective_batch to drive use_batch / codegen
Python changes:
- finalize_reducer(reducer, batch=1) – new optional argument
- batch < 1 raises ValueError (consistent with T.reduce)
Tests: add test_batched_finalize_reducer_codegen covering sum/max/min
with explicit batch values, verifying run_batch is emitted with the
correct template argument.
* [Test] Add comprehensive tests for T.finalize_reducer batch parameter
Covers all cases introduced by the explicit batch= API:
Codegen tests:
- batch>1 emits AllReduce::run_batch with correct template argument
- batch=1 (default) does NOT emit run_batch (scalar path)
Correctness tests:
- batch=1 scalar path: sum/max/min across float32/float16 (PASS)
- batch==block_M full-batch: correct by coincidence (PASS)
- batch>1 partial-batch: xfail(strict=True) – known bug where
run_batch receives buffer->data directly but fragment layout is
non-contiguous per thread; root cause documented in
work/bug-finalize-reducer-batch.md
Error-input tests:
- batch=0 / batch<0 raise ValueError at Python layer
- batch > total output elements raises at compile/lower time
- batch that does not evenly divide output elements raises
* [Test] Rewrite finalize_reducer tests to follow clean outline
Replace the previous over-engineered test suite (4 case lists, 2 split
correctness functions, per-case xfail) with a minimal set organised by
the three categories that matter:
1. Codegen – single parametrised test covers both batch=1 (no run_batch)
and batch>1 (run_batch with correct template arg), reusing one case
list for both codegen and correctness.
2. Correctness – only batch=1 (scalar path), which is currently correct.
batch>1 correctness tests are omitted until the underlying fragment
layout bug is fixed (tracked in work/bug-finalize-reducer-batch.md).
3. Error inputs – batch<1 collapsed into one parametrised test; compile-
time errors (batch>output, batch not divisible) each get one test
using a shared _make_error_kernel helper.
* [Test] Rewrite reduce tests following clean outline
Replace fragmented per-function tests with four parametrized suites:
- test_reduce: op × dtype × src_scope × dst_scope × threads, covers
fragment→fragment, shared→shared, fragment→shared for sum/max/min/
abssum/absmax with float32/float16/int32/int64
- test_reduce_clear: op × src_scope × dst_scope with clear=False,
verifies pre-filled dst is accumulated correctly
- test_reduce_batch_codegen: verify batch>1 emits AllReduce::run_batch
- test_finalize_reducer_codegen/correctness/invalid_batch: batch
parameter feature tests, sharing one FINALIZE_REDUCER_CASES list
Use tilelang.jit lazy style and torch.testing.assert_close directly,
no profiler wrapper.
* [Test] Merge batch codegen+correctness into test_reduce
Add batch column to REDUCE_CASES; batch>1 cases verify both that
run_batch appears in generated source and that numerical output matches
reference, within the same test_reduce parametrization.
Remove the now-redundant test_reduce_batch_codegen.
* [Test] Add more batch>1 reduce cases covering all ops and dtypes
Add batch=2/4/8 cases for sum/max/min/abssum/absmax across float32/
float16/bfloat16, with both fragment and shared src/dst scopes.
Relax tolerance to 1e-1 for float16/bfloat16 to account for larger
rounding error in multi-element reductions.
* [Test] Limit int64 input range to avoid ROCm shfl_xor 64-bit bug
* Revert "[Test] Limit int64 input range to avoid ROCm shfl_xor 64-bit bug"
This reverts commit
|
||
|
Lei Wang
|
5680e39e45 |
[TMA] Support FP4 TensorMap TMA copies (#2107)
Support FP4 TensorMap TMA copies |
||
|
Zhang Jason
|
53a4c9866b |
[AMD] [gfx950]Fix multiple HIP codegen bugs to support TileKernel (#2099)
* Fix HIP codegen for sync_warp, sync_grid, and local.var initialisation * [AMD/HIP] Fix warp_reduce VGPR bug, ShuffleNode packing, and Pipelined LDS overflow Extends PR #2096 with three additional fixes for CDNA (MI350) targets: Fix 1 — src/tl_templates/hip/reduce.h: warp_reduce width=32 The old 6-step butterfly called __shfl_xor(value, 32) without a width argument. On CDNA (wave64) with 32 active threads, lanes 32-63 are inactive and hold uninitialised VGPRs, producing NaN in reduce_max / reduce_sum / AllReduce. Fix: remove the step-32 shuffle; pass width=32 to all remaining 5 steps so every shuffle stays within the [0,31] group. Fix 2 — src/target/codegen_hip.cc + src/tl_templates/hip/common.h: ShuffleNode bfloat16x2 / float16x2 packing CodeGenC emitted `uint1(a, b)` for bfloat16x2 construction, which is an invalid HIP constructor call. Fix: override VisitExpr_(ShuffleNode) in CodeGenTileLangHIP to emit `uint1{__pack_bfloat162(a, b)}` / `uint1{ __pack_half2(a, b)}` using aggregate initialisation. Also add five bfloat16x2 math overloads for uint1 carrier (abs2/max2/min2/add2/mul2). Fix 3 — src/transform/pipeline_planning.cc: skip T.Pipelined(num_stages>1) Double-buffering doubled LDS per loop-body buffer. On CDNA (≤128 KB LDS per workgroup), this caused hipModuleLaunchKernel EINVAL. Fix: when TargetIsRocm() && num_stages > 1, skip pipeline planning and fall back to a plain sequential loop with synchronous T.copy. Also: fix __habs(hip_bfloat16) and __habs(float16_t) in common.h to use __builtin_memcpy instead of reinterpret_cast to avoid strict-aliasing UB (as flagged by CodeRabbit on PR #2096). Tests: 19 new cases added to testing/python/amd/test_tilelang_hip_codegen.py covering all three fixes. All 42 tests pass on MI350 (gfx950). * [AMD/HIP] Merge test_tilelang_hip_bugfixes.py into test_tilelang_hip_codegen.py Consolidate all HIP regression tests into a single file. The merged file covers all six fixes with 32 tests total (previously split across two files with duplicated test cases for warp_reduce, pipelined GEMM, and ShuffleNode). Changes versus the two individual files: - Deduplicated test_warp_reduce_no_nan (identical in both files) - Deduplicated test_pipelined_no_lds_overflow / test_pipelined_shared_mem_no_launch_error - Deduplicated test_pipelined_multi_stage_fp16_gemm - Merged bfloat16 shuffle tests: source check + runtime correctness in one function - Kept PR #2096 source-inspection tests (alloc_var, sync_warp, sync_grid) - Added runtime tests from bugfixes: inf init, serial loop accumulation, float scalar readback, two-group wave64 reduce, float16 shuffle * fixup: correct LDS size comment — gfx950 has 160 KB, not 128 KB gfx942 (CDNA3 / MI300X) has 64 KB LDS per workgroup. gfx950 (CDNA4 / MI350) has 160 KB LDS per workgroup (see PR #2058). The old comment said '128 KB' which is wrong for both generations. Updated pipeline_planning.cc and the test docstrings to reflect the correct per-architecture limits. * update for format checking |
||
|
Jiaxing DingandJiaxing Ding
|
6a29c76aee |
feat: support cdna4 v_mfma_i32_16x16x64_i8 & v_mfma_i32_32x32x32_i8 (#2097)
Co-authored-by: Jiaxing Ding <csding99@outlook.com> |
||
|
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> |
||
|
Zhang JasonandLeiWang1999
|
6d0bffb9e7 |
[AMD][Gfx950] Add the support of 160K LDS and copy.async (#2058)
* add gfx950 new feature, 160K LDS and compy async * update accoding to format check result * fix a typo error * add gfx950 copy async test codes --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Zhang Jason
|
b88bbcb5fa |
[AMD][gfx950] Add ds_read_tr16_b64 / ds_read_tr8_b64 support for gfx950 LDS transpose reads (#2085)
* add gfx950 ds_read_tr support * update with format checking |
||
|
Lei Wang
|
6fee85002e |
[Release] Bump version into 0.1.9 (#2060)
* bump version into 0.1.9 * remove legacy test |
||
|
Lei Wang
|
15309f5cba |
[CUDA] Improve int4 GEMM lowering and packed codegen support (#2073)
* [Cleanup] Remove redundant int4 MMA code paths * [CI] Apply format.sh updates * lint fix * fix * Refactor FP8 type determination to utilize `T` module constants Updated the `determine_fp8_type` function to return constants from the `T` module instead of string literals for FP8 formats. This change enhances consistency and aligns with recent updates in type handling. Adjusted return statements for both HIP and CUDA scenarios to use `T.float8_*` constants. * Refactor FP8 type usage in tests to directly call `determine_fp8_type` Updated test files to replace `T.dtype(determine_fp8_type())` with direct calls to `determine_fp8_type()` for improved clarity and consistency in FP8 type handling. This change aligns with recent refactoring efforts in type determination. * Remove print statements from `assert_tl_matmul_correctness` and `test_assert_tl_matmul` functions to clean up test output. This change enhances the clarity of test results by eliminating unnecessary console output. * Enhance CI workflow for CuTeDSL examples with Python 3.12 and CUDA-12.8 - Added a step to list generated files and clean up Python cache directories. - Introduced a new job for running CuTeDSL examples on self-hosted NVIDIA runners. - Set up environment variables for CUDA and Python, including caching configurations. - Implemented error handling to clear cache on setup failure. - Updated steps for installing dependencies and running examples, ensuring compatibility with the latest Python and CUDA versions. |
||
|
Chenhao Xu
|
64bd742071 |
[Bugfix] Fix reduce layout (#2074)
* fix: fix reduce layout * test: add targeting test |
||
|
Lei Wang
|
9c95a42dfb |
[Refactor] Phaseout legacy util map_torch_type with T.dtype.as_torch (#2075)
* Remove redundant `map_torch_type` in favor of `T.dtype.as_torch()` `map_torch_type` was a standalone function that converted dtype strings or TVM DataType objects to `torch.dtype`. This was redundant since `T.dtype.as_torch()` (and `T.dtype(x).as_torch()`) already provides the same functionality with better handling (HIP/CUDA awareness, proper dtype lookup tables, etc.). Changes: - Remove `map_torch_type` from `tilelang/utils/tensor.py` and its re-export - Make `T.dtype()` idempotent for existing dtype objects (pass-through) - Add backward-compat for `"e4m3fnuz_float8"` TVM internal name in `as_torch()` - Replace all usages across production code, tests, and examples - Rebuild Cython extension to reflect the import change * Refactor output and accumulation data types in GEMM FP8 example Updated the `out_dtype` and `accum_dtype` variables in the `example_tilelang_gemm_fp8_intrinsic.py` file to use `T.float32` instead of string literals. This change enhances type consistency and aligns with recent updates in type handling. |
||
|
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 |
||
|
Lei Wang
|
e1d63880f6 |
[Cache] Refactor cache namespace layout (#2057)
* Refactor cache namespace layout * Update atomic cache tests for namespaced paths * Ensure namespaced cache dirs exist before save * Use auto backend in grouped gemm ptr example * Keep grouped gemm ptr example off cutedsl * Always use auto backend for grouped gemm ptr * Skip grouped gemm examples in cutedsl CI * Ignore grouped gemm examples in cutedsl workflow |
||
|
TerminusAkiviliandLeiWang1999
|
04468a3425 |
[Bugfix] Fix stage-expanded annotated-layout aliases in LayoutInference (#2031)
* Fix annotated layout inference for dtype-changing views * lint fix --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
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 |
||
|
Chaofan Lin
|
77cbe6da87 |
[Feature] Introduce T.CUDASourceCodeKernel (#1970)
* [WIP][Feature] Introduce CUDARunner * refactor * refactor to T.SourceCodeKernel api * remove streamk change * remove runner * remove prim func params * enhance API * add entry_name * add IsCodeBlockKey * remove misc file * cleanup code * move parse logic to python through callback * fix * comments |
||
|
|
27f1f81c61 |
[Feature][Example] Introduce CLC tile schedule and add example for sm100 GEMM (#2029)
* [Feature] Add Blackwell Cluster Launch Control (CLC) primitives
Expose cluster launch control query operations (try_cancel,
try_cancel_multicast, is_canceled, get_first_ctaid_{x,y,z}) in both
the C++ codegen path and the Python tilelang.language API, plus a new
mbarrier_arrive_expect_tx builtin used with CLC multicast completion.
Enables writing persistent SM100 kernels that use the hardware scheduler
to fetch the next tile coordinate.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* [Example] Add SM100 persistent 2-CTA GEMM example using CLC scheduling
Adds examples/gemm_sm100/gemm_tcgen5mma_ws_clc.py, a warp-specialized
persistent GEMM for Blackwell (sm_100) that uses tcgen05 2-CTA MMA
with cluster launch control tile scheduling. The kernel is split into
producer/MMA/scheduler/epilogue warp groups, with the scheduler warp
running on both CTAs and feeding next-tile coordinates via the CLC
hardware multicast into schedule_arrived directly (no intermediate
barrier). Storage sync is disabled and all syncs are placed manually.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* lint
* typo
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
* remove unnecessary test
* Add CLC query intrinsics and async fence
* Refactor CLC response handling by introducing CLCResponseDecode struct. Update clc_get_first_ctaid_x/y/z and clc_is_canceled functions to utilize the new decode response method for improved clarity and maintainability.
* Enhance CLC scheduling in persistent GEMM by introducing shared variables for schedule validity and tile ID. Update cancellation checks to utilize the new shared variables, improving clarity and maintainability of the scheduling logic.
* test fix
---------
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
|
||
|
Lei WangandFreebase6912
|
d098f51b8f |
[Enhancement] Use atomic directory rename for cache writes (#1982)
* [Enhancement] Use atomic directory rename for cache writes Replace per-file os.replace with write-to-staging + os.rename to ensure other processes never see an incomplete cache directory. Also add cleanup of stale staging dirs on init. * [Bugfix] Fix atomic cache save: abort on critical file failure and fix staging cleanup - Let kernel lib and params save failures propagate instead of swallowing them, so the staging directory is cleaned up rather than renamed as an incomplete cache entry. - Move autotuner staging directory to TILELANG_CACHE_DIR so that _cleanup_stale_staging_dirs() can find and remove stale entries. * Fix atomic cache save recovery and error handling * Remove unnecessary blank lines in _FakeAdapter and _FakeKernel classes in autotune and cache test files. * Tighten atomic cache completeness checks * Simplify atomic cache save error handling * submodule update * remove legacy test --------- Co-authored-by: Freebase6912 <amid-gauze-racing@duck.com> |
||
|
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>
|
||
|
Nguyen Huy HoangandLeiWang1999
|
235ad7eda4 |
Add regression test for 1D TMA load compilation and execution (#1989)
* refactor: add regression test for 1d tma load compilation and execution Signed-off-by: Nguyen Huy Hoang <181364121+huyhoang171106@users.noreply.github.com> * lint fix --------- Signed-off-by: Nguyen Huy Hoang <181364121+huyhoang171106@users.noreply.github.com> Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
KurisandClaude Opus 4.6
|
f309d8147e |
[Bugfix][Subtype] Fix scalar fp4 store/load codegen for non-packed buffers (#2037)
* [Bugfix][Subtype] Fix scalar fp4 load/store codegen and dynamic stride inference for sub-byte types The scalar fp4 buffer access path in codegen_cuda.cc used GetBufferRef which applies a /2 index division (mapping two fp4 elements to one byte), then performed a plain byte assignment — destroying the neighboring nibble. This caused incorrect results whenever fp4 stores couldn't be vectorized (e.g. with dynamic/non-constant strides in StridedTensor). Fix: use tl_fp4_packed_load/tl_fp4_packed_store with (fp4_e2_2_t*) cast for scalar fp4 accesses to non-packed buffers, which correctly handles nibble-level read-modify-write via set_x/set_y. Also add stride_scale factor to _process_dynamic_symbolic in the cython and tvm_ffi adapters, so torch strides (in storage units) are correctly converted to logical element strides for sub-byte dtypes. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * [Test][Subtype] Add correctness tests for fp4 dynamic-stride store codegen Add regression test that scatters fp4 rows into a non-contiguous StridedTensor with dynamic strides and verifies byte-for-byte match against a static-stride reference. Parametrized over block_size. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * [Enhancement][Subtype] Inject stride divisibility assume for sub-byte dtypes and add scalar fp4 store codegen test - Extend InjectAssumes pass to emit `assume(stride % pack_factor == 0)` for non-last-dimension strides of sub-byte buffers (e.g. fp4 with pack_factor=2). This provides the TIR analyzer with divisibility information before LetStmt inlining. - Add `test_subtype_fp4_scalar_store_codegen`: scatter fp4 elements via indirection to force the scalar (non-vectorized) store codegen path and verify nibble-level correctness. - Convert new tests to eager-style @tilelang.jit. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * style: apply clang-format and ruff formatting Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * fix lint error * fix: unpack 4-tuple in cython_wrapper dynamic_symbolic_map lookup --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
||
|
|
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 Wang
|
b3d5981b00 |
[API] Default warp-lane mask to 0xFFFFFFFF for warp-sync builtins (#2039)
Reorder the `mask` parameter to the end with a default of `_FULL_WARP_MASK` for shfl_xor/down/up/sync, any_sync, all_sync, ballot_sync, and match_any/all_sync. Full-warp is the overwhelmingly common case, so making mask an optional trailing kwarg removes boilerplate at call sites and lets old `(value, delta)`-style callers keep working. Update the test suite and docs/instructions.md to match the new signatures. |
||
|
|
19236b44c0 |
[Feature] Add TIR builtins for warp-level vote and block-level predicate sync (#1858)
* [Feature] Add TIR builtins for warp-level vote and block-level predicate sync * [Refactor] Lower warp-level intrinsics through tl.* ops Convert the warp-vote, warp-shuffle, predicated block-sync, and warp-match builtins from raw `tir.call_extern` wrappers to proper TIR ops registered under `tl.*` and lowered in `codegen_cuda` / `codegen_hip`. This removes Python-side `_IS_HIP_AVAILABLE` branching and pushes the CUDA/HIP split into codegen, where it belongs. * Register new tl ops: any_sync, all_sync, ballot_sync, ballot, activemask, syncthreads_count/and/or, shfl_sync, shfl_xor_sync, shfl_down_sync, shfl_up_sync, match_any_sync, match_all_sync. * Codegen lowering on both CUDA and HIP. uint32->uint64 zero-extension for ballot/activemask now happens in codegen. HIP drops the mask argument for shfl/vote and emits LOG(FATAL) for match_*_sync. * `__match_all_sync`'s int* pred output is hidden behind an immediately-invoked lambda so the wrapper stays expression-form. * Python wrappers normalize Python int masks to uint32 TIR consts via `_as_uint32_mask`, so the emitted source prints `(uint)0xFFFFFFFF` instead of `(int64_t)4294967295`. * Unify `shfl_xor/down/up` signatures to `(mask, value, delta, width=32)`, matching CUDA convention and `shfl_sync`. No in-tree callers existed. * Add tests for match_any_sync and match_all_sync. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
||
|
|
8243f7e089 |
Nan propagation option for bf16 and half16 (#1958)
* register op, and add to build file * update semantics of NaN for Max and Min ops in Reduce * lower Max op to cutlass fastmax in general case, _hmax for bf16 and _hmax_nan for hf16 * update semantics of NaN for Max and Min ops in Reduce in reduce.h * develope codes transforming ReduceOpNode in TIR and handling CallNode in CUDA code generator's runtime template * test file for reduce maxmin nan * refine test case message and explanation“ ” ‘ “ * invert the senamtics of global config TL_REDUCE_MAXMIN_NAN_PROPAGATE * lint fix * refactor reduce maxmin NaN propagate into per-call annotation Replace the global tl.reduce_maxmin_nan_propagate pass-config with a per-op nan_propagate flag carried on the ReduceOp annotation map. Frontend exposes it as a kwarg on T.reduce_max/min/absmax (default False, preserving prior behavior). The CUDA codegen MinNode/MaxNode visitors are reverted to plain __hmin/__hmax; only the new tl::max_nan /tl::min_nan CallNode handler emits __hmax_nan/__hmin_nan, so non-reduce Min/Max ops are no longer silently rewritten. ReduceOp::Lower now errors early with a clear message if nan_propagate=True is requested for fp16 /bf16 on a non-CUDA target, instead of emitting an undefined symbol on HIP/CPU. Tests rewritten to parametrize on the kwarg and assert runtime NaN-propagation behavior, not just generated source substrings. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Haoran <chaoran@umich.edu> Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
||
|
Lei Wang
|
5d729eeebc |
[Refactor] Remove GEMM v1 and promote gemm_py to be the canonical gemm op (#2033)
* [Refactor] Remove GEMM v1 and merge gemm_py into the canonical gemm op
GEMM v2 has been the default for a while and no tests or examples still
depend on v1. This removes the legacy C++ GemmNode path entirely and
promotes the former GemmPyNode (Python-delegated lowering) to be the one
and only GEMM tile op.
C++:
- Merge src/op/gemm_py.{h,cc} into src/op/gemm.{h,cc}. GemmPyNode becomes
GemmNode, GemmPy becomes Gemm, FFI type key "tl.GemmPy" becomes "tl.Gemm",
and the op names "tl.tileop.{gemm_py,wgmma_gemm_py,tcgen05_gemm_py}"
collapse to "tl.tileop.{gemm,wgmma_gemm,tcgen05_gemm}". The shared
infrastructure (GemmWarpPolicyNode, GemmInst enum,
GemmWarpPolicyNode::computeWarpPartition) is preserved in place so
GemmSP/GemmSPPy keep compiling unchanged.
- Delete the tl_gemm builtin, its codegen dispatch in
codegen_{cuda,hip,cutedsl}.cc, and its matching branches in
inject_fence_proxy.cc and merge_shared_memory_allocations.cc. tl_gemm_sp
is kept as the remaining async-proxy gemm builtin.
- Drop #include "gemm_py.h" and GemmPyNode/GemmPy references from
inject_pipeline.cc, producer_consumer_ws.cc, and lower_blackwell_2sm.cc.
Python:
- tileop/gemm: rename GemmPy -> Gemm, register as "tl.Gemm", and rename
the FFI callbacks tl.gemm_py.{lower,infer_layout} -> tl.gemm.{lower,
infer_layout}. _ffi_api.GemmPyGemmInst becomes GemmGetGemmInst.
- language/gemm_op.py: delete gemm_v1/gemm_v2 wrappers. gemm() now emits
tl.tileop.gemm directly; wgmma_gemm/tcgen05_gemm drop the _py suffix.
- env.py: remove TILELANG_USE_GEMM_V1 and the use_gemm_v1() helper.
- ir.py: drop the empty v1-era `class Gemm(Node, Scriptable): ...`
placeholder that was colliding with the real tileop.gemm.Gemm wrapper
under the "tl.Gemm" FFI type key.
- contrib/cutedsl: delete cutedsl/gemm_v1.py (runtime helper for the v1
cutedsl lowering path) and its re-export in cutedsl/__init__.py.
Tests / maint:
- testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py switches
from T.gemm_v2 to T.gemm.
- testing/python/transform/test_tilelang_transform_inject_fence_proxy.py::
test_lower_fence_proxy is rewritten to exercise tl.tl_gemm_sp instead
of the removed tl.tl_gemm, preserving the fence-proxy assertion.
- Delete maint/gemm_v2/ (v1-vs-v2 benchmarking scripts) and drop the
v1-specific comment from maint/scripts/run_local_ci_test.sh.
C++ headers under src/tl_templates/{cuda,hip,cpp}/gemm*.h are left in
place as reference implementations; they are no longer reachable from
the generated code path since the only consumer (tl_gemm builtin) is
gone, but keeping them avoids a large, hard-to-reverse deletion.
* [Fix] Restore cutedsl/gemm_v1.py (Gemm_SM80 / Gemm_SM90 helpers)
The previous commit deleted cutedsl/gemm_v1.py on the assumption that it
was the CuteDSL lowering helper for the removed T.gemm_v1 path. It is
actually a runtime helper module that defines SM80/SM90 CUTE DSL
primitives (make_smem_layout_AB, get_tma_atom, etc.) which are emitted
into generated cutedsl kernel code via
`import tilelang.contrib.cutedsl as tl; tl.Gemm_SM90.get_tma_atom(...)`
(see jit/adapter/cutedsl/wrapper.py CUBIN_TMA_ATOM_INIT_TEMPLATE).
The "v1" in the filename refers to the CuteDSL helper version, not to
the deleted T.gemm_v1 op, so removing it broke the cutedsl code path for
attention_sink, blocksparse_attention, blocksparse_gemm, and deepgemm
examples with `module 'tilelang.contrib.cutedsl' has no attribute
'Gemm_SM90'`. Restore the file and its `from .gemm_v1 import *` wildcard
re-export.
* [Refactor] Restore tl_gemm builtin as a call_intrin path to tl:: templates
The previous commit in this PR deleted the `tl_gemm` builtin entirely on
the argument that no code emitted it anymore. That was too aggressive:
keeping the builtin preserves a deliberate extension point for any
Python-side gemm lowering backend (for example a future
GemmLegacyExtern) to delegate into the existing C++ templates in
src/tl_templates/<target>/gemm*.h via
T.call_intrin("handle", tir.op.Op.get("tl.tl_gemm"),
op_instance_str, A_ptr, B_ptr, C_ptr)
without any extra C++ plumbing. It is also the only thing that keeps
those tl_templates headers reachable rather than dead code.
Restore:
- tl_gemm builtin declaration in src/op/builtin.h (with an updated
docstring describing the call_intrin entry point) and its
TIR_DEFINE_TL_BUILTIN registration in src/op/builtin.cc.
- Codegen dispatch branches in codegen_cuda.cc, codegen_hip.cc, and
codegen_cutedsl.cc that turn a tl_gemm Call into a PrintCallExtern
emission of `tl::gemm_xx<...>(A, B, C)`.
- SMEM alignment and async-proxy fence handling for tl_gemm in
merge_shared_memory_allocations.cc and inject_fence_proxy.cc. These
passes are tied to the builtin's semantics (not to the removed v1
GemmNode) and are still correct the moment anything emits tl_gemm.
- test_lower_fence_proxy in test_tilelang_transform_inject_fence_proxy
reverts to exercising `tl.tl_gemm` instead of the temporary
`tl.tl_gemm_sp` substitution introduced when the builtin was gone.
Note that the current sole gemm lowering path (the merged GemmNode that
delegates to Python via tl.gemm.lower) still emits inline mma/wgmma/
tcgen5mma via the Python macro generators, not tl_gemm. So restoring
the builtin is a no-op for the default path - it only re-enables the
"call into v1 templates from Python" option for anyone who wants it.
* [Fix] Default HIP execution backend to tvm_ffi, not cython
resolve_execution_backend() only promoted cuda and metal to tvm_ffi for
the "auto" path, so hip fell into the else branch and picked cython. But
allowed_backends_for_target() already lists tvm_ffi as an allowed
backend for hip, and the cuda default is tvm_ffi for good reasons:
- The cython SourceWrapper path in src/tilelang/jit/adapter/wrapper.py
has known gaps (for example T.alloc_global is not handled in
create_dispatch_func, which breaks any test that uses an internally
allocated global scratch buffer).
- libgen.py appends user-supplied compile_flags verbatim to the
hipcc/nvcc command. nvcc accepts --use_fast_math, clang++/hipcc does
not, so a CUDA-style flag leaks into HIP compilation.
Neither of these leak through the tvm_ffi backend. Aligning hip's
default with cuda makes the two paths symmetric and unbreaks
test_alloc_global and test_all_attrs_together_lazy on ROCm without
having to mark the tests skip or rewrite the cython wrapper.
cython is still available on hip via execution_backend='cython' for
anyone who explicitly wants it.
* [Test] Gate CUDA-only language tests with requires_cuda
Three language tests hardcode target='cuda' (or rely on CUDA-only
codegen behavior) but lacked a CUDA availability marker, so they failed
on the ROCm CI runner with "Cannot find global function
target.build.tilelang_cuda" once the earlier pre-existing failures got
past --maxfail=3. Add @tilelang.testing.requires_cuda to:
- test_tilelang_if_range
- test_annotate_min_blocks_per_sm_launch_bounds
- test_tilelang_transpose and test_tilelang_transpose_square
Also replace the hand-rolled __main__ block in test_tilelang_transpose
with tilelang.testing.main() so the skip markers are honored when the
file is run directly.
* [Test] Delete tile library gemm test
Remove testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py.
|
||
|
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
|
||
|
Lei Wang
|
37c1c0c688 |
[codex] Reduce slow pytest runtime in testing/python (#2018)
Optimize bitwise reduce test runtime |
||
|
Lei WangandZhiwen Mo
|
cabdb86f25 |
[codex] Fuse packed x2 mul-add into fma2 in CUDA codegen (#2017)
* fuse packed x2 mul-add into fma2 * document packed x2 fma2 fusion rationale --------- Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk> |
||
|
Lei Wang
|
4f75940612 |
[BugFix] Enhance CUDA vectorization for binary operations (#2015)
Enhance CUDA vectorization for binary operations by supporting wider tensor dimensions. Update kernel generation to accommodate variable width, improving auto-vectorization for float32 types. Add tests for width-8 scenarios to ensure correct emission of packed operations. |