`hexp` is #define'd to `cutlass::fast_exp`, so the hand-written
`cutlass::fast_exp(bfloat16_t)` overload's `return ::hexp(x)` expanded into a
call to itself: bf16 `T.exp` recursed and returned garbage.
Route through float, mirroring cutlass's own `half_t fast_exp`
(cutlass/fast_math.h), so the call resolves to `cutlass::fast_exp(float)`
(= ::expf). Verified against torch.exp using the #2383 repro.
Fixes#2383.
* Fix TMA 1D load test lowering
* Refactor test execution in tilelang TMA 1D test script
Updated the main execution flow to directly call the test function instead of running an elementwise add operation, improving clarity and focus on testing functionality.
The SM75 GEMM emitter derives MMA shapes from input dtype bits alone, so it
already emits the (16,8,8) shape for fp16->fp16 accumulation (identical to the
registered fp16->fp32 case) and the (8,8,16) shape for uint8 inputs (identical
to signed int8). But the C++ dispatcher table only covered the three configs
tested by #2198 (fp16->fp32, int8, int4), so requesting fp16->fp16-accum or
uint8 GEMM on Turing failed at nvcc with an opaque
`tl::mma_sync: unsupported configuration` static_assert.
Add the two missing tl::detail SM75 atoms and register their dispatchers,
following the precedent set by #2198's hand-written int4/uint4 structs (CuTe
ships only fp16->fp32 and signed-int8 SM75 atoms). The uint8 atom uses the
non-saturating PTX form to match the existing signed-int8 path.
Verified on a TITAN RTX (SM75): emitted PTX is
mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 and
mma.sync.aligned.m8n8k16.row.col.s32.u8.u8.s32, and end-to-end GEMM numerics
match torch.
* fix: name
* [Transform] Enhance role management and scalar dependency propagation
- Introduced a new Role enumeration value `kNeutral` to better represent variable states.
- Added `CombineRoles` function to manage role combinations effectively.
- Implemented `ScalarAccess` structure and `ScalarAccessCollector` class for tracking variable definitions and usages.
- Enhanced `PropagateScalarDependencyRoles` function to propagate roles based on scalar dependencies across statements.
- Updated role assignments in various statement visitors to utilize the new role management logic, improving clarity and correctness in role handling.
* Update main execution flow in test_tilelang_language_let.py to enable tilelang testing framework
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* Refactor kernel launch materialization
* [Refactor] Improve git commit ID retrieval and version labeling logic
* Enhance the get_git_commit_id function to check for the existence of the .git directory before attempting to retrieve the commit hash, ensuring more robust behavior.
* Update the dynamic_metadata function to conditionally set the NO_VERSION_LABEL environment variable based on the presence of the .git directory.
* Modify the GitHub Actions workflow to unset the NO_VERSION_LABEL environment variable during the installation and version check of the package, allowing for accurate version labeling in SDist builds.
* [Refactor] Update WebGPU backend registration
* Import the WebGPU package to automatically register its pass pipeline.
* Remove the explicit registration of the WebGPU pass pipeline, streamlining the backend setup process.
* [BugFix] CuTeDSL backend support for TileKernels benchmarks
TileKernels benchmark coverage with TILELANG_TARGET=cutedsl exercises backend paths that were either missing or not legal for CuTeDSL-generated Python/MLIR code. The failures were not limited to one kernel: packed FP4/FP8 quantization flows need sub-byte vector materialization, MoE mapping uses warp vote/shuffle and grid synchronization, and benchmark launch wrappers need dtype-aware host plumbing for generated CuTeDSL modules.
Before this change the H100 benchmark run exposed several classes of failures: missing shuffle/vote intrinsics such as shfl_xor/match_any/popc in CuTeDSL lowering, invalid scalar/vector casts for FP4 E2M1 and packed UE8M0 paths, dynamic if expressions whose branches narrowed or changed integer types, unsupported sub-byte load/store/copy patterns, and a fragile one-counter software grid barrier that could reset before all waiting blocks observed release.
Lower the missing TileLang operations to CuTeDSL runtime helpers, add Float4E2M1FN software pack/unpack and zero-fill support, make sub-byte vector loads/stores use packed Uint8 storage where required, cast conditional branches to the expected result dtype, and teach the CuTeDSL wrapper about float4/float8 dtype plumbing and C++ min/max expression rendering. Also replace the sync_grid single-counter barrier with a counter plus generation sense barrier so a waiting block cannot miss the release when the arrival counter is reset.
This keeps the fixes in the CuTeDSL backend rather than adding TileKernels-specific workarounds. For example, TileKernels expand_to_fused_with_sf relies on TileLang safe-memory legalization for aligned SF loops; the CuTeDSL generated code now preserves the hidden_sf guard instead of requiring the benchmark source to special-case packed UE8M0 layouts.
Validation on H100 (sm_90, 114 SMs) with CUDA visible device 0, Python 3.12, TILELANG_TARGET=cutedsl, TK_FULL_TEST=1, TILELANG_DISABLE_CACHE=1, and PYTHONPATH=/home/jayce.su/new/tilelang:
- cmake --build build -j 16
- ruff check tilelang/contrib/cutedsl tilelang/jit/adapter/cutedsl tilelang/jit/adapter/utils.py
- git diff --check
- Original failing node: tests/moe/test_expand_to_fused.py::test_expand_to_fused_with_sf_benchmark[num_send_tokens=4001-num_topk=2-num_experts=4-num_ep_ranks=64-hidden=576-num_per_channels=32-use_tma_aligned_col_major_sf=True-round_sf=True-use_packed_ue8m0=True], 1 passed with -n 2
- Full-run local pressure slice around the failing MoE benchmark region, 80 passed with -n 2 in 379.53s
Author-email: jayce.su@enflame-tech.com
Co-authored-by: dingsg <shengge.ding@enflame-tech.com>
* [BugFix] Complete CuTeDSL support for TileKernels benchmarks
TileKernels benchmark coverage with TILELANG_TARGET=cutedsl exposes several general CuTeDSL backend gaps. These are not TileKernels-specific workarounds, H100-only specializations, or changes required only because of a particular CUDA/CUTLASS DSL version pairing. The H100 CUDA 13.2 + nvidia-cutlass-dsl-libs-base 4.5.0 environment is the reproducer that exercises the missing backend coverage; the underlying issues are generic to the emitted CuTeDSL code shapes.
Root causes:
- Predicated T.copy and safe memory legalization lower BufferLoad-or-zero values into conditional expressions. Emitting these values through tl.where evaluates both branches eagerly, which can issue unsafe packed loads and can force CuTeDSL/MLIR through unsupported vector select and unrealized_conversion_cast paths for FP8/FP4 storage.
- Safe memory legalization can bind the zero branch through a local zero-like variable, let binding, cast, or broadcast-zero tensor. The CuTeDSL codegen must recognize those zero-like values when deciding whether a conditional load can be lowered lazily and when materializing Float4E2M1FN broadcasts.
- CuTeDSL generated code has to distinguish rmem Tensor objects from TensorSSA values. Appending .load() blindly is invalid once a branch or cast already produced TensorSSA, and this showed up in the TileLang CI cast example as AttributeError: 'TensorSSA' object has no attribute 'load'.
- Scalar FP8 widening and narrowing require CuTeDSL tensor casts on aligned vector values. Plain scalar construction falls back to unsupported scalar FP8 conversion paths, so the backend has to bridge scalar casts through a packed vector form.
- Sub-byte vectors such as Float4E2M1FN cannot be copied lane by lane because CuTeDSL does not support sub-byte scalar dereference. The backend must preserve TensorSSA/packed-vector store form for these values, and scalarized local/shared sub-byte vector buffers must use the sub-byte packing factor as the make_tensor_at_offset div_by assumption.
- Generated CuTeDSL kernels can call CUDA-template math helpers such as __habs and __float2half_rz, and math wrappers such as log1p/isfinite. Those helper names must be provided by tilelang.contrib.cutedsl.math when the Python CuTeDSL backend is selected.
- sync_grid publishes block arrival through thread 0, but the memory ordering has to cover all threads in the block. A CTA sync alone does not provide device-scope ordering for global writes performed by non-zero threads.
Fixes:
- Recognize zero-like vars in conditional store lowering, including let bindings and bind values that are casted or broadcast-zero tensors. This keeps predicated loads lazy and avoids eager tl.where evaluation for packed and vectorized T.copy paths.
- Add CuTeDSL helper bridges for Tensor/TensorSSA normalization. tl.as_tensor_ssa avoids calling .load() on TensorSSA values, and tl.as_rmem_tensor spills TensorSSA values into rmem tensors before byte recasts for sub-byte stores.
- Add scalar FP8 bridge lowering that packs Float16/Float32/BFloat16 or FP8 scalar values into a four-lane rmem tensor, uses tl.cast_tensor, and reads lane zero. This matches CuTeDSL/NVGPU alignment requirements without forcing front-end kernels to perturb tile shapes just to avoid scalar casts.
- Keep sub-byte vector stores in TensorSSA form instead of copying lane by lane, keep non-load vector store RHS materialized as rmem tensors so CuTeDSL .store() receives a TensorSSA value, and use 8 / element_bits as the div_by packing factor for scalarized local/shared sub-byte vector buffer views.
- Gate the Float16 nvgpu.cvt_fpext fast path to floating-point source TensorSSA types only, so integer TensorSSA casts do not enter an FP-only conversion op.
- Strengthen sync_grid by fencing all block threads at device scope before thread 0 publishes arrival, then waiting at block scope and fencing again before user code resumes.
- Expose the missing CuTeDSL math helpers. __habs clears the sign bit for Float16/BFloat16 scalars and returns the original dtype. __float2half_rz uses inline PTX cvt.rz.f16.f32 so generated E5M6 cast-back code can preserve CUDA round-toward-zero semantics. log1p and isfinite are wired through the CuTeDSL math module and fastmath canonicalization.
- Keep the defensive dynamic-symbolic None fallback visible by logging a warning before substituting 0.
Version/specialization assessment:
- The patch is backend-generic: it fixes how TileLang emits CuTeDSL for predicated loads, FP8 scalar casts, Tensor/TensorSSA boundaries, sub-byte vector stores, grid synchronization, and generated math helpers. It does not specialize TileKernels cases, tensor shapes, H100 SM counts, or benchmark parameters.
- A future CuTeDSL version might accept some currently rejected IR forms, but relying on eager branch evaluation, missing helper symbols, incorrect TensorSSA .load() calls, or thread-0-only grid-sync ordering would still be fragile. The corrected lowering is semantically closer to TIR: predicated loads are lazy, sub-byte storage stays packed, TensorSSA values are not treated as tensors, and scalar FP8 conversion follows the vector cast path required by CuTeDSL/NVGPU.
- CUDA 13.2 and CUTLASS DSL libs 4.5.0 are recorded as the verified reproducer, not as the reason for the patch. The same TileLang backend gaps can be triggered by any workload that emits these CuTeDSL patterns.
Verified on H100 PCIe (sm_90, CUDA_VISIBLE_DEVICES=0 maps to NVIDIA H100 PCIe), CUDA 13.2, Python 3.12.3, PyTorch 2.12.0+cu132, nvidia-cutlass-dsl-libs-base 4.5.0, and local TileLang built from this branch:
- cmake --build build -j$(nproc)
- python3.12 -m py_compile tilelang/contrib/cutedsl/math.py tilelang/contrib/cutedsl/utils.py tilelang/contrib/cutedsl/grid_sync.py tilelang/jit/adapter/cutedsl/adapter.py
- pre-commit run --files $(git diff --name-only origin/main...HEAD)
- examples/cast/test_example_cast.py::test_example_group_per_split_token_cast_to_fp8: 1 passed
- representative per_token_cast, cast_back, cast_back_e5m6, per_block_cast_lossless, per_block_cast, and per_channel_cast benchmark failures: 6 passed
- TileKernels full benchmark with TILELANG_TARGET=cutedsl, TK_BENCHMARK_BACKEND=event, and pytest -n 2: 10165 passed, 160 warnings in 2:44:58
Jayce: JayceSu98 <jayce.su@enflame-tech.com>
Co-authored-by: dingsg <shengge.ding@enflame-tech.com>
* [BugFix] Avoid zero-sized fake tensors in CuTeDSL cubin generation
CuTeDSL cubin generation builds fake compact tensors to compile wrapper code before the real runtime tensors are launched. TileKernels exercises benchmark and correctness cases with empty dynamic dimensions, such as zero-token inputs. Passing those zero extents directly into make_fake_compact_tensor can make the CuTeDSL fake layout invalid even though the shape is only a compilation placeholder and the real runtime tensor metadata is passed separately.
Clamp only the fake tensor shape to positive extents before constructing the CuTeDSL placeholder. This keeps the cubin generation path valid for empty-input cases without changing runtime tensor shapes, strides, launch parameters, or kernel behavior.
This is not a TileKernels-specific workaround or a CUDA toolkit version workaround. It fixes the CuTeDSL adapter boundary so compile-time fake tensors satisfy CuTeDSL layout construction requirements while preserving the actual runtime contract.
Verified with:
- python -m py_compile tilelang/jit/adapter/cutedsl/wrapper.py
- git diff --check -- tilelang/jit/adapter/cutedsl/wrapper.py
Jayce: JayceSu98 <jayce.su@enflame-tech.com>
Co-authored-by: dingsg <shengge.ding@enflame-tech.com>
* [BugFix] Stabilize batched BF16 reduce test input
CUDA CI exposed a flaky failure in python/language/test_tilelang_language_reduce.py::test_reduce[sum-bfloat16-64x128-f2f-t256-b4]. The failing run had only one mismatched output element, with an absolute difference slightly above the 0.1 BF16 tolerance and a large relative error because the reference value was small.
The runtime code on this branch only changes CuTeDSL backend paths, and the same test passed locally on the ordinary CUDA backend when run in isolation. The failure is caused by the test input depending on the worker-local torch RNG state: the module seeds once at import time, but xdist load scheduling lets each worker consume a different number of random tensors before this parameterized case runs. That makes the BF16 sum case occasionally sample an input whose rounded reduction error crosses the tolerance boundary.
Make the parameterized reduce input schedule-independent by deriving a deterministic seed from the case parameters and using a local CPU torch.Generator for both integer and floating inputs. This preserves randomized coverage per case while decoupling correctness from xdist worker assignment, prior cases in the same worker, and full-suite ordering.
Verified on H100 PCIe with CUDA visible device 1, Python 3.12.3, PyTorch 2.12.0+cu132, and the local TileLang dev build:
- ruff check testing/python/language/test_tilelang_language_reduce.py
- python -m pytest -q -s --tb=short testing/python/language/test_tilelang_language_reduce.py::test_reduce[sum-bfloat16-64x128-f2f-t256-b4]
- python -m pytest --verbose --color=yes --durations=0 --showlocals --cache-clear --maxfail=3 --numprocesses=8 testing/python/language/test_tilelang_language_reduce.py
- 10 consecutive reruns of the failed BF16 reduce case
Jayce: JayceSu98 <jayce.su@enflame-tech.com>
Co-authored-by: dingsg <shengge.ding@enflame-tech.com>
---------
Co-authored-by: dingsg <shengge.ding@enflame-tech.com>
LowerSharedBarrier lowers T.alloc_barrier() to hardware mbarrier code (the
cutlass Barrier type, tl::tl_shuffle_elect, tl::fence_barrier_init), all of
which only exist on sm_90+. On pre-Hopper targets (e.g. sm_86) the frontend
accepted the call silently and nvcc later failed with cryptic
"identifier 'Barrier' is undefined" / "namespace 'tl' has no member ..."
errors.
Add a static compute-capability guard in CUDAPassPipelineBody, before
LowerSharedBarrier: if the target is < sm_90 and the module uses a shared
barrier, raise a clear ValueError naming the offending primitives and pointing
to __syncthreads()/named barriers as the pre-Hopper alternative.
Fixes the sm_86 alloc_barrier compile crash.
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* Fix eager JIT packed FP4 shape binding
* Add logical shape and stride calculations for buffers in eager builder
This update introduces helper functions to compute logical shape and stride values for buffers, enhancing the handling of packed storage factors. The changes streamline the process of determining dimensions and strides in the context of tensor operations, improving overall efficiency and clarity in the eager execution flow.
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* Add TMA copy for packed global FP4 to unpacked SMEM.
Wire IsValidTMACopyDtypePair and ALIGN16B tensor map descriptors for
float4_e2m1fn global to float4_e2m1_unpacked shared transfers, and
extend language tests for TMA and SIMT copy restrictions.
* Extend tcgen05 MMA metadata for FP4 f8f6f4 and dtype helpers.
* Refactor tcgen05_gemm_blockscaled API and lowering.
* Use separate A/B dtypes in GEMM tileop emitters across backends.
* Add maint/gemm correctness evaluation scripts and update gemm_sm100 examples.
* Fix profiler assert_allclose to compare outputs only.
* lint
* canonicalize data types
* add dpgm fp8-fp4 kernel
* resolve comments
* Refine FP4 unpacked SMEM dtype handling
Treat float4_e2m1_unpacked as a SMEM/TMA-only storage tag instead of a general vector/register dtype. Remove vectorized Python aliases and keep codegen uses constrained to shared-memory storage/reference paths that lower to uint8_t.
Allow packed FP4 global tensors to TMA-load into unpacked shared memory with the ALIGN16B TensorMap format, while keeping unpacked storage out of TMA store paths. Use payload-bit accounting for global strides and mbarrier transactions.
Add focused dtype and TMA descriptor tests, and validate the real tcgen05 fp4-unpacked GEMM paths.
* Restrict FP4 unpack TMA to loads
* Hide vectorized unpacked FP4 language aliases
* Update TVM for unpacked FP4 symbol exports
* Rebind TVM submodule to merged unpacked FP4 export fix
* Enable cache for persistent tcgen5 GEMM example
* Place auto WS producers in first warp group
* [Transform] Update layout inference to bind thread range for fully replicated fragments; add tests for TMA fragment mask copy and producer-consumer warp specialization
* [Refactor] Update thread variable calculation to use zero-based indexing across GEMM operations
* [Transform] Enhance simplification process with ContextSingletonFloorDivFolder and add tests for floor division simplification
* [Refactor] Simplify tensor conversion in dequantization functions and improve performance in GEMM examples
* [AMD][RDNA4] Fix RDNA4 (gfx1201 / Wave32) CI failures
- reduce.h: Replace hardcoded wave64 assumptions with runtime
__builtin_amdgcn_wavefrontsize() in SharedReduceWarp and AllReduce
- scan.h: Add SEG=32 support in InclusiveScan/CumSum/CumMax for RDNA wave32
- finalize_reducer.cc: Use TargetGetWarpSize() instead of hardcoded 64
- carver/arch/__init__.py + rdna.py: Add gfx12 (RDNA4) support alongside gfx11
- wmma_macro_generator.py: Fix k_pack indexing in ldmatrix_a/b for k_pack > 1
These changes correspond to PR #2210 (Fix RDNA4 gfx1201/Wave32 CI Failures),
rebased cleanly onto upstream main (a8d93798).
* [AMD][HIP] Fix HIPMath dispatch for vector dtype in T.Parallel
When BLOCK_M > threads, TileLang generates vector types (e.g. float32x4)
for T.Parallel loops. HIPMath::operator() only handled scalar types, so
vector dtypes fell through to return causing Unresolved call errors
in codegen for math intrinsics like tirx.exp2.
Fix: add a vector-type early-return that recurses on element_of(), so
float32xN correctly resolves to exp2f the same as scalar float32.
This complements the tirx.* registration fixes in #2282 and enables
Online Softmax on gfx1151 with BM > threads (e.g. BM=1024, TH=256).
* [Review] Fix 3 issues found in code review
1. tilelang/carver/arch/__init__.py: error message was not updated
when gfx12 support was added — still said gfx11 only. Now says
gfx11/gfx12 targets only, consistent with rdna.py.
2. src/tl_templates/hip/reduce.h: the static_assert(Threads % kWarpSize == 0)
was removed when kWarpSize became a runtime value, leaving no safety
check. Added equivalent runtime assert().
3. src/tl_templates/hip/scan.h: run_auto() else branch silently fell
through to SEG=64 for any non-32 wavefront size. Added assert to
catch unsupported wavefront sizes early.
* update codes for format checking
* [Review] Address review feedback and rebase on upstream main
- intrin_rule_hip.cc: remove redundant HIPMath scalar recursion; upstream
PR #2286 already handles vector dtypes via DispatchPureExternScalarized
at the IR level (correct for HIP packed carrier types)
- wmma_macro_generator.py: rename l/r to row_off/col_off to fix Ruff E741
ambiguous variable name warnings
- scan.h: guard run_auto() branches with threads>=32 / threads>=64 so a
threads==32 kernel on Wave64 doesn't compute TILE=0 and divide by zero
- reduce.h: replace assert (stripped in Release) with unconditional early
return when Threads % kWarpSize != 0, preventing silent num_warps==0
* update codes for format checking
* fix
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* Scalarize vector math intrinsics for HIP codegen
* Clarify HIP fp32 vector-math comment in scalarizer
* fix comment
* Scalarize HIP vector math via PrintCallExtern, not generic Shuffle
The generic Shuffle re-pack emits one ctor arg per logical lane, which is
wrong for HIP packed carrier types (float32x8->ulonglong4, float16x4->uint2).
Emit a vector call_pure_extern with the element-derived extern name instead
and let CodeGenTileLangHIP::PrintCallExtern store each lane (PrintVecElemStore),
which is carrier-aware.
* Add HIP vector math scalarizer test
* JIT-decorate scalarizer test helper
[BugFix][Transform] Freshen buffer defs after loop unroll
UnrollLoop clones each loop body with Substitute, which preserves Buffer object identity for definitions inside the body. This can define the same DeclBuffer handle in multiple unrolled copies after flattening and violate TIR SSA. Freshen buffer definitions per cloned body while preserving DeclBuffer backing data aliases.
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
* [Backend] Refactor Pipeline to support different backends
Introduce Pipeline abstraction in backend/pipeline.py with per-backend
registration. Each backend (cuda, hip, c, llvm) now registers its own
compilation pass pipeline. engine/lower.py resolves the pipeline from
the target instead of hardcoding phase imports.
* enhance the structure
* fuse two stage
* fix after rebase
* update template pass pipeline
* pipeline
* fix test
* fix test
* fix cuda pipeline prologue
* Add float4_e2m1_unpacked dtype
* Fix TIR cast for float4_e2m1_unpacked dtype.
Route custom[float4_e2m1_unpacked] through Float4E2M1Unpacked FFI
so T.float4_e2m1_unpacked(0.0) works in prim_func builders.
* Limit float4_e2m1_unpacked vector lanes and fix TMA transaction bytes.
Cap unpacked FP4 vector dtypes at x32 and count TMA/mbarrier
transactions using 4-bit payload size (0.5 byte per element).
* lint
* [Refactor] Extract shared-access analysis into a reusable header
- Move AccessEntry/AccessType definitions from thread_storage_sync.cc to
src/transform/common/shared_access_analysis.h so other passes (notably
MergeSharedMemoryAllocations) can reuse the same precise access
representation.
- Add tl.debug_thread_storage_sync_summary pass-config option for optional
ThreadSync planner summary logging.
- Add regression tests covering explicit shared sync deduplication and
shared.dyn alias-aware dependency planning.
* [Feature] Rewrite MergeSharedMemoryAllocations with precise access analysis
- Build per-access live ranges using shared_access_analysis utilities so
branch-alternative dataflow and barrier-driven epoch boundaries are
modeled directly instead of through whole-buffer approximations.
- Plan arena layout starting from arena_size_const so allocations whose
starting offsets are non-zero no longer collide.
- Preserve single-slot semantics for mutually exclusive branch writes
followed by a post-dominating read (fixes the K_norm_pairs_shared style
miscompile where stage-1 alternatives were split across epochs).
- Cover BufferRegion block reads/writes, tvm_access_ptr, tl.access_ptr,
cp_async branches, loop-carried buffers and partial-subregion liveness
cases.
- Keep aggressive-merge disabled under warp specialization and retain
whole-buffer liveness as the conservative default.
- Add 18 transform regressions in
test_tilelang_transform_merge_shared_memory_allocations.py covering
basic merge, sync, branch, cp_async, BufferRegion and opaque-pointer
paths.
* [Cleanup] Remove dead FreeList/FreeBlock/ActiveInterval helpers
These helpers were superseded by the segment-aware LinearScanPack
implementation that performs interval-overlap walking directly. They
have no remaining callers after the rewrite; remove them to reduce
maintenance burden.
* [Test] Cover tail-allocation arena growth scenario
Add a regression case mirroring upstream b4c913be (#2106): when the only
free slot reusable by an incoming buffer sits at the arena tail and is
smaller than the request, the arena must grow by the deficit only, not
push the new buffer past the freed slot.
The legacy free-list-based packer would emit a 384 B arena for this
pattern (X=128, Y=64-then-freed, Z=192), bumping Z to offset 192. The
interval-based LinearScanPack reuses Y's offset and grows the tail to
fit Z, producing a 320 B arena -- a 16.7%% reduction that matches the
peak-live lower bound.
* [Feature] Add EpochGraph and use per-epoch liveness to relax aliasing
Introduce src/transform/common/epoch_graph.h, a control-flow-aware
builder that partitions a flattened StmtEntry sequence into hard-sync
and cp_async-wait epochs across nested For/IfThenElse scopes, exposing
loop-carry, branch-join and parent-edge structure required for sound
liveness analysis.
Wire MergeSharedMemoryAllocations to compute per-buffer live epochs as
the intersection of:
* forward def-reach: epochs reachable from any write of the buffer
along epoch-graph successors;
* backward use-reach: epochs that can reach any read of the buffer
along predecessors.
A buffer is live in epoch e iff e is in both sets, modeling 'data
written upstream is consumed downstream' rather than the legacy
[start_stmt, end_stmt) bounding box. The result feeds two relaxers:
* LinearScanPack: two intervals whose live-epoch sets are disjoint
may share an offset even if their stmt ranges overlap.
* PlanMemory: the symmetric overlap probe consults the same set
before forcing a non-aliased layout.
Hard syncs (kSync) and cp_async waits are treated as epoch boundaries;
ptx_cp_async_wait is conservatively classified as a hard barrier.
IfThenElse joins take the union of branch live sets so a buffer
written in only one branch remains live at the join.
The whole-buffer EpochGraph + liveness pipeline always runs; verbose
[MSMA-EPOCH-GRAPH] / [MSMA-EPOCH-LIVE] / [MSMA-EPOCH-ACCESS] /
[MSMA-EPOCHSEG] / [MSMA-STAGE0-RAW] dumps remain gated behind the
existing tl.debug_merge_shared_memory_allocations pass-config flag.
Also remove SplitLocalSegmentsAtEpochSyncs, a Phase-1 helper that was
defined but never wired into the planner; the new liveness pipeline
subsumes its role.
Validated on a real lowbit MHA decode kernel
(seq_lens=[(3,96),(2,64)], heads=4, dim=64, page_size=32, bits=4):
arena drops from 35,232 B (upstream baseline) to 31,136 B, an 11.6%%
reduction driven by K_shared and V_shared sharing one offset thanks
to disjoint live-epoch sets.
Add testing/python/transform/test_tilelang_transform_epoch_graph.py
covering the five S1.4 epoch-graph topology cases (sequential
__syncthreads, ptx_wait_group, For body/back/exit edges, IfThenElse
branch_join, nested For). Existing
test_tilelang_transform_merge_shared_memory_allocations.py and
test_tilelang_transform_thread_sync.py suites continue to pass.
* [Feature] Gate per-epoch alias relaxer on annotated layout equality
Two shared-memory buffers carrying distinct layouts under
`block.annotations[attr::kLayoutMap]` (e.g. via `T.annotate_layout(...)`)
must not share storage even when per-epoch dataflow proves their
liveness sets disjoint: downstream layout-driven access lowering can
emit indexing arithmetic that is only sound under each buffer's
original layout, so naively aliasing them risks silent miscompiles.
Add a `LayoutSignatureCollector` that harvests the per-Var `Layout`
entries from `Block` annotations and a `LayoutAliasIncompatible`
helper that returns true iff both candidates have non-trivial layouts
and the layouts are not structurally equal. The two relaxer call
sites in `SharedMemoryRewriter` (`LinearScanPack` and the post-pack
overlap audit) consult the helper before dropping a conflict, so the
legacy decision is preserved exactly when the gate fires.
Buffers without a layout entry remain layout-agnostic and continue to
participate in arena packing, keeping the existing reuse for the
common un-annotated case. Verified that the lowbit reference kernel
arena footprints stay at 31136/61856/61808/61792 B and the merge /
thread-sync / epoch-graph trio passes 49/49.
* [Test] Add transform regression that fails on upstream MergeSharedMemoryAllocations
Add test_general_cond_branch_alternatives_share_offset_with_postdom_buffer
to lock the per-epoch / branch-alternative analysis behavior on a
mutually-exclusive if/else with a runtime cond plus a post-dominator
buffer after a sync. Upstream's linear-scan plan keeps the two branch
buffers at distinct offsets (128-byte arena); the per-epoch relaxer
collapses them onto the same offset and reuses it for the
post-dominator buffer (64-byte arena, 50% smaller).
* [Example] Demo MergeSharedMemoryAllocations branch-alternative arena reuse (64 vs 128 B)
Add runnable demo under examples/merge_shared_memory_allocations/
that constructs a three-buffer PrimFunc with mutually-exclusive if/else
branches, lowers through the merge pass, and reads the merged arena
size from the dyn_shared_memory_buf attr. The demo also compiles the
kernel end-to-end and validates output correctness with torch.allclose.
On this branch: arena = 64 B (saved 128 B / 66.7% vs naive 192).
On upstream main (1fbd994c): arena = 128 B (-50% worse).
* [Test] Add end-to-end torch correctness regression for merge pass
Add two tests (test_merged_branch_kernel_generates_correct_results,
test_merged_sync_phased_buffers_generate_correct_results) that compile
the merged kernel and verify bit-exact output via torch.allclose,
ensuring the arena packing does not alter numerical results.
* [Test] Add torch correctness and dynamic-buffer fallback regression
- test_merged_branch_kernel_generates_correct_results: compiles a
merged branch-alternative kernel with torch.allclose validation.
- test_merged_sync_phased_buffers_generate_correct_results: sync-phased
two-buffer kernel end-to-end check.
- test_dynamic_size_buffer_fallback_does_not_incorrectly_alias:
verifies that a symbolic-extent shared.dyn buffer correctly aliases
with a sync-separated constant-size buffer without unsafe overlap.
* [Fix] Disable per-epoch relaxer for loop-live buffers to avoid pipelined iteration overlap
* [Fix] Fallback to structural_segments_by_var for buffers missing planner segments to avoid buffer_byte_offsets crash and shared memory bloat
* Merge branch 'main' into merge_pass and fix merge comflicts
* [Fix] Resolve post-merge artifacts in CMakeLists, hipcc, and deepseek example
- Remove duplicate TILELANG_OUTPUT_TARGETS in CMakeLists.txt
- Fix duplicate _resolve_artifact_paths function in hipcc.py
- Add missing G=kv_group and fix assert in sparse_mla_fwd_pipelined.py
- Add ruff noqa pragma to suppress lint warnings
* [Style] Apply clang-format to modified source files
* [Pipeline] Refactor software pipeline transforms
* enhance
* Refine pipeline handling for guarded binds
* Refactor IfStmtBindingRewriter and related classes for improved access analysis
- Removed unused Target parameter from IfStmtBindingRewriter.
- Introduced IfStmtAccessCollector to handle buffer access collection.
- Updated access analysis methods to utilize the new collector.
- Cleaned up includes and removed obsolete functions in access_analysis.
- Enhanced pipeline utility headers with necessary includes for consistency.
* [Fix] Suppress spurious __syncthreads() between T.ws blocks (#2192) and add T.named_barrier_arrive (#2188)
Issue #2192 - BugFix: Unexpected CTA-wide __syncthreads() between T.ws producer/consumer
- Root cause: FindConflict() in thread_storage_sync.cc did not check for mutually
exclusive thread constraints. When T.ws(0) (tx<128) reads and T.ws(1) (tx>=128)
writes the same shared buffer, no thread satisfies both constraints, so no
__syncthreads() is needed — mbarriers already handle the synchronization.
- Fix: add mutual-exclusivity check (prove NOT(P AND C)) early in FindConflict().
If the access constraints are mutually exclusive (disjoint thread sets), return
false (no conflict) immediately, before any index-range analysis.
- Measured improvement: barrierpipe_stage2 gemm latency 44ms → 25ms on H200.
Issue #2188 - Feature: expose native T.named_barrier_arrive(barrier_id, thread_count)
- Adds tl.named_barrier_arrive intrinsic op (bar.arrive PTX, one-sided arrive).
- Registers op in builtin.cc, declares in builtin.h.
- Adds tl::__named_barrier_arrive<id, cnt>() template in common.h.
- Adds codegen case in codegen_cuda.cc.
- Exposes T.named_barrier_arrive() in tilelang/language/builtin.py.
- Adds regression test testing/python/issue/test_tilelang_issue_2188_2192.py.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* lint: clang-format + ruff fixes after #2188/#2192 changes
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* test: simplify issue 2188/2192 test, keep only #2192, add requires_cuda
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
* [Fix] Make named_barrier_arrive accept variable barrier_id/thread_count
* Update documentation for named barriers and remove obsolete test file
- Clarified documentation for `named_barrier_arrive` in `builtin.h` and `builtin.py` to specify the barrier index range.
- Removed obsolete regression test for issues #2188 and #2192 as it is no longer needed.
* lint fix
---------
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Tong WU