Merge upstream main (85fd8fc2) into the Ascend comparison snapshot. This brings
in #3177, #3181, #3182, #3185 (AutoSchedule -> AutoWarpSpecialization) and
#3186 (target-neutral T.Kernel with dialect-owned launch annotations).
Conflict resolution:
- AutoSchedule -> AutoWarpSpecialization (#3185) supersedes the snapshot's own
disambiguation rename to `tl.cuda_auto_schedule`; upstream's name is already
distinct from Ascend's boolean `tl.enable_auto_schedule`, so the CUDA rename is
taken wholesale and `TL_ENABLE_AUTO_SCHEDULE` (Ascend-only) is kept.
`testing/python/backend/test_tilelang_backend_auto_schedule.py` follows.
- `tilelang/language/kernel.py` is now upstream's target-neutral implementation,
unchanged. The Ascend additions it used to carry (`check_ascend_availability`,
`_current_target_is_ascend`, the `is_cpu`/`is_ascend` branches in `Kernel`, the
annotation-sniffing grid/thread split in `KernelLaunchFrame.__enter__`, and the
SimtVF hooks in the thread accessors) are relocated, not adapted: see below.
- `tilelang/language/__init__.py` keeps the fork's Ascend facade and gains
upstream's typed `Kernel` re-export.
Ascend now owns its launch surface, so the shared launch path carries no backend
branches and no runtime hardware probe:
- `tilelang/ascend/language/kernel.py` defines the dialect's `Kernel(*blocks,
prelude=)`: the NPU 1-D core grid, with no `threads=`/`cluster_dims=` to be
rejected at runtime because Python rejects the keywords. It also owns the
SimtVF thread-scope state and the `get_thread_binding(s)`/`get_thread_extent(s)`
accessors that resolve against an active `T.SimtVF` scope.
- `MaterializeKernelLaunch` splits "how a launch dimension is lowered" from
"whether SIMT threads exist", which the old single `lower_thread_binding` flag
conflated: Ascend needs the grid materialized as a real block-level
thread_extent while having no threadIdx at kernel scope. `lower_grid_binding`
defaults to None and follows `lower_thread_binding`, so every existing caller
and test keeps its meaning. `launch_dim_tags` lets a backend name extra launch
dimensions; Ascend passes `cthread` for `T.MixedKernel`'s sub-block-id.
- `MixedKernelLaunch` populates the frame's new grid_vars/thread_vars fields and
emits the thread placeholders like `KernelLaunch` does, so a body referencing a
thread index gets the actionable diagnostic instead of an out-of-range lookup.
- `tilelang.is_cpu_kernel_frame` / `tilelang.is_npu_kernel_frame` are dropped from
`src/transform/common/attr.h`: the explicit grid_vars/thread_vars fields replaced
every reader the snapshot had added for them.
Combining them also removes the cross-backend landmine where, on any host with
torch_npu available, the shared `T.Kernel(N, threads=...)` raised even for CUDA
targets.
* [Refactor][Language] Make T.Kernel target-neutral with dialect-owned launch annotations
T.Kernel is traced before the Target is known, yet it used to emit
threadIdx.* thread_binding loops with a hard-coded default of 128 threads.
That is a SIMT launch baked into the frontend: CPU lowering silently pinned
the thread index to 0 (so per-thread code computed 1/128 of the work), and
every backend inherited CUDA's launch vocabulary.
Split the launch into the part every target shares and the parts a backend
interprets:
- The frontend records the grid as blockIdx.* thread_binding loops (the
program-index space all targets share) and reserves thread identities as
`tx/ty/tz = tl.launch_thread_idx(axis)` Bind placeholders. `threads=`
becomes the `tl.launch_threads` annotation on the tilelang_root block;
the frontend no longer chooses a default.
- MaterializeKernelLaunch(lower_thread_binding, default_threads,
unsupported_annotations) is where the target-dependent meaning is decided.
SIMT backends rebind each placeholder as a threadIdx.* thread_extent over
the same Var (extent from the annotation, else the backend default; the
CUDA/ROCm/Metal/WebGPU pipelines pass 128). Backends without SIMT drop the
placeholders and reject a body that references a thread index instead of
miscomputing. Launch annotations a backend cannot honor (cluster_dims on
cpu/rocm/metal/webgpu) are rejected here rather than ignored downstream.
- KernelLaunchFrame exposes grid_vars/grid_extents/thread_vars/
thread_extents explicitly; get_thread_extent() raises at trace time when
threads= was omitted, since only the backend knows the extent.
- Each language dialect owns its `Kernel` with explicit keyword parameters:
cuda offers threads/prelude/cluster_dims, rocm/metal threads/prelude,
webgpu threads, cpu prelude, and the common `Kernel(*grid)` none. All of
them funnel into `tilelang.language.kernel.launch_kernel`. This is what
editors show on hover and what type checkers verify; unknown keywords are
rejected by Python itself. ClusterKernel keeps its interface.
- Cluster identity: `T.get_cluster_id/ids/dims/size/extent(s)` are
program-space arithmetic (`bx // cluster_dims`) computed at trace time and
therefore target-neutral, while `T.block_rank_in_cluster()` stays the
CUDA hardware intrinsic. The two agree under the cluster-aware threadblock
swizzle, which permutes the grid at cluster granularity.
Tests that applied passes directly to a traced function now materialize the
launch first (thread indices otherwise carry no extent), fixtures that wrote
per-thread code and compiled it for the `c` target are rewritten at the tile
level, and the span-preservation check asserts that no statement loses its
span rather than that the spanned count never decreases.
* [Fix][Language] Unwind kernel frames after tracing errors
Pass the active exception to the pending frame exits when tracing fails, so
the Python launch state (KernelLaunchFrame stack) is cleared without
finalizing incomplete TIR. Previously the next trace observed the failed
kernel's launch frame as KernelLaunchFrame.Current(). Covered for both the
prim_func and eager JIT tracing paths.
* [Fix][Language] Recognize every dialect's Kernel as a launch in the eager rewriter
Eager JIT phase 1 infers the signature with symbolic T.const values and must
not execute the kernel body; the AST rewriter guards `with T.Kernel(...)`
with `if __tb.skip_kernel_ctx(): return`. It found the launch by object
identity against the common `Kernel`, which no longer holds now that each
dialect (and the `tilelang.language` facade) exports its own `Kernel`: the
guard was not injected, phase 1 ran the body, and e.g. `T.gemm` on a tile
whose extent comes from `T.const` failed its static-shape check.
Mark every launch factory (`Kernel` of each dialect, `ClusterKernel`) with
`kernel_launch_factory` and let the rewriter test that mark instead. It now
resolves plain names and attribute chains only, so aliases such as
`K = T.Kernel` are recognized and no factory expression is executed at
rewrite time. Covered by a rewriter-level test over all dialects and an
end-to-end eager gemm whose tile extent depends on a const.
Comparison-only snapshot of Ascend at 925bfbb9aff0fb57a4ed9fecc3591b8dcac21c8f.
Use upstream #3174 (62bba8d20d) as the logical synchronization point.
Preserve the complete tree of the previous comparison head d0d2f40be354cc0a4d6ed037b34278c8746fa031,
including submodule references. No source changes or additional integration
resolutions are introduced. This snapshot is for upstream comparison only.
* [CUDA] Unify TMA copy lowering on CuTe algebra and support general swizzled layouts
Every descriptor-based TMA path (bulk copy, atomic add, gather4/scatter4,
im2col) now shares one CuTe-algebra box decomposition. Boxes follow the
driver contract (inner box bytes at most the swizzle span, per-dim caps),
and the recovered swizzle is widened on demand under an equivalence proof,
so general annotated layouts -- including phase-shifted swizzles -- lower
correctly.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
* Address review concerns
---------
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
[Layout] Preserve vectorized reducer update plans
Treat reducer_update as a transparent expression during layout vector-shape analysis. This lets contribution loads determine the planned width instead of forcing update nests onto scalar-shaped layouts before reducer materialization.\n\nAdd a CUDA regression that checks the shared-to-fragment copy remains packed and verifies the reducer result.
The probe inverts a candidate loop layout into symbolic per-thread
coordinates plus a symbolic thread variable, then checks whether any
fragment physical index still contains the thread variable. Proving a
strided distribution thread-free is range-dependent (e.g.
`(o0*512 + t*4 + o1) // 512 -> o0` only holds for t*4 + o1 < 512), so
with unbound ranges Simplify keeps the thread term and even an identity
access (loop layout == fragment layout) is falsely charged as a register
spill. The spill pricing then steers layout inference away from
conflict-free strided shared-memory loads toward chunk-contiguous
layouts whose 16-word stride bank-conflicts on every shared load,
observed as a 1.23-1.29x latency regression on a production
KL-divergence kernel (8.8ms vs 6.8ms at n=8192, d=262144 on H800).
Bind every per-thread output coordinate to its output-shape range and
the thread variable to the loop's thread range/extent so the legitimate
cancellations go through. Genuine thread dependence is unaffected: range
information only enables valid simplification, it cannot remove a real
thread term. Regression test distilled from the affected kernel.
[Bugfix] Make TMA linear layouts region-aware to keep versioned slices contiguous
The inferred "linear" shared layout hoisted every 256-split quotient above
all dims, so a pipeline-versioned slice became two gapped chunks: LowerBulk's
bijection check fired on warp-specialized wide reductions, and the TMA
atomic-add path (which has no such check) silently read the other version.
MakeTmaLinearLayout now takes the copied region and orders modes as
[region-fixed dims, box quotients, box contents], so a fixed slice owns one
contiguous run of whole TMA boxes and the existing rest-loop machinery emits
one instruction per box over the true (unaligned) global shape. The
atomic-add path shares the helper and computes the region base offset through
the layout instead of hand-rolled row-major strides over the physical shape.
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
A synchronous T.gemm(mbar=...) inside T.Pipelined waited on the wrong
mbarrier parity. InjectSoftwarePipeline annotated every copy/im2col/gemm
call with the ring parity (k / num_stages) % 2, but that law belongs to
the compiler-generated pipeline_mbar ring: an explicit mbarrier whose
arrive happens inside a tile op is deliberately excluded from
ExpandPipelineBarriers, so it remains a single slot that flips once per
iteration. The auto-emitted wait asked for a parity that had already
completed, returned immediately, and the next iteration's copies
overwrote a shared-memory stage the still-running MMA was reading
(wrong results observed at num_stages >= 3).
Annotate only the ops that consume the phase -- those that auto-emit a
paired arrive+wait on an explicit mbarrier: sync gemm with an mbar, and
copies carrying an explicit "barrier" that are not is_tma_copy -- and
annotate them with the logical iteration parity (k - min) % 2, written
against the original loop var. EmitImpl's existing per-block
Substitute(loop_var -> normalized_access_index) specializes the
expression for every prologue/body/epilogue clone, since the tirx
mutators already rewrite PrimExprs inside call annotations.
Pipeline-managed TMA copies keep their ring parity in the explicit
waits the pass inserts, and the explicit-async T.tcgen05_gemm never
auto-waits, so neither is annotated.
Drop the matching ring branch from LowerTileOp's fallback: it keyed on
the "num_stages" annotation, which InjectSoftwarePipeline strips, so it
was dead on pipelined loops and wrong on unpipelined ones -- a
compiler-generated per-site barrier flips once per loop iteration
regardless of pipeline depth. The fallback is now the zero-based epoch
of the nearest enclosing serial loop.
Add transform-level tests asserting the annotated phase and the emitted
wait parities on sm_90a, so the fix is covered without sm_100 hardware;
the loops start at a nonzero min to pin the epoch normalization. Add a
tcgen05 correctness test running the failing kernel at num_stages 2
and 3.
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Remove an unconsumed cluster-planning pass, duplicate sparse GEMM and fast-math implementations, the unused JIT environment module, and the superseded autotuner timeout machinery. Document and test the surviving portable timeout behavior.
* [Op] Add read-before-write semantics to tile operators
GetAccessRegions() answers a may-read question: could this operator read this
region? It is deliberately conservative, which is correct for the dependency
analysis it was built for. There is no way to ask the stricter question of
whether an operator consumes a region.s preexisting contents before
establishing a value of its own.
GetReadBeforeWriteRegions() asks it. The default returns GetAccessRegions()
.reads, so operators whose read set does not depend on an argument value need
no change, and existing consumers of GetAccessRegions() are untouched.
Gemm and GemmSP override it. GetAccessRegions() lists the accumulator as a
read whenever the clear cannot be proven (!is_one(clearAccum_)), which is
right for pipelining: a conditional clear genuinely is a read dependency. It
is wrong for a read-before-write question, because the common pipelined idiom
clear_accum=(k == 0) is neither literal true nor literal false. The overrides
use is_zero() instead.
Note that operators such as ReduceOp need no override: its clear is a plain
bool rather than a PrimExpr, so the default is already exact.
* [Transform] Add VerifyBufferInit, a general buffer-initialization check
Reading a non-global buffer that nothing has written yields whatever those
locations last held. That is undefined behaviour: it can look correct on one
architecture and produce NaN on another (issue #2936).
VerifyBufferInit walks a PrimFunc before LowerTileOp, where tile operators
still declare their access regions, so writes performed by T.copy, T.clear and
friends are visible rather than opaque. It warns and returns the function
unchanged; it never fails a build.
The analysis is conservative by construction, in three ways:
- A call the pass cannot interpret may write through its arguments, so it is
assumed to. TCallEffectKind distinguishes these from pure calls; an
operator with no registered effect kind is assumed to write. This also
covers tl.access_ptr, which keeps a BufferLoad rather than a data Var.
- Shared memory ignores source order. A warp-specialized kernel fills a shared
buffer from a producer branch that appears after the consumer branch that
reads it; the two run concurrently and coordinate through barriers, so
source order is not an approximation of happens-before. For shared scopes
the question is whether any *other* operation writes the buffer at all.
Per-thread storage keeps order-sensitive tracking, where source order is
that thread.s execution order.
- Global memory and mbarrier scopes are never reported.
Excluding the reading operator.s own write is what keeps #2936 reported: a
gemm.s only writer of its accumulator is itself.
* [Transform] Enable the buffer-initialization check in every pipeline
Wired into all five backends immediately after LayoutReducer, and on by
default with tl.disable_buffer_init_check to opt out.
The slot matters in four ways. LegalizeNegativeIndex has run, so buffer
indices are canonical. LowerTileOp has not, so tile operators still carry
their access regions. InjectSoftwarePipeline has not, so loop bodies are in
source order rather than rotated across stages. And LayoutReducer has run:
before it, a tl.tileop.finalize_reducer call carries a single argument while
its builder reads args[1], so ParseOperator throws on it. Every pre-existing
ParseOperator caller sits after that point, which is why the incomplete form
had not been parsed before.
22 tests cover the gemm cases, the general cases, the two scope-ordering
rules, report aggregation, and the config plumbing end to end.
* [Transform] Route both buffer-init walks through one opaque-write rule
The whole-function collector and the ordered walk each decided for
themselves what an uninterpretable call may write, and they disagreed.
The collector recorded any BufferLoad argument, so a read-only
tl.access_ptr counted as a writer and could suppress a real finding on
shared memory, while a write through tvm_access_ptr passed a data Var
rather than a BufferLoad and was missed entirely. It also applied the
rule to pure calls, whose arguments only read.
Both now share ForEachOpaqueWrite and the same escape predicate, so the
access-mask and pointer-form handling cannot drift apart again.
Also adds a regression test for the unguarded loop-carried read. That
shape is reported and the report is correct: the first iteration reads
the fragment before the write at the bottom of the body reaches it. The
case the analysis does not decide is a read the loop guards out of its
own first iteration, which is path sensitivity rather than ordering.
Verified: 23 verifier tests pass; testing/python/language is 997 passed
with the same 5 pre-existing failures and zero warnings; all 284 example
scripts emit one warning, the known uninitialized lse_max_local in
examples/flash_decoding/example_mha_inference.py.
* [Transform] Record a store's destination after evaluating it
A BufferStore reads its own right-hand side before it establishes
anything, so marking the destination written first hid every
self-referential read: `x[i] = x[i] + 1` with nothing written earlier
went unreported.
I had defended that ordering on the grounds that flagging
`idx = T.if_then_else(cond, i, idx)` turns the check into
definite-assignment analysis over conventional idioms. Looking at the
two examples that motivated it, that defence does not survive.
examples/grouped_gemm/example_grouped_gemm_fwd_ptr.py writes
`cur_batch_idx = 0` before exactly this loop; example_grouped_gemm_fwd.py
and example_grouped_gemm_bwd.py do not. The initializer is already there
in the sibling, so this is an oversight in two of three copies rather
than an idiom the check misreads.
Cross-thread scopes are unaffected: they ignore source order and ask
only whether another node writes the buffer.
Verified: 24 verifier tests pass; testing/python/language is 997 passed
with the same 5 pre-existing failures and zero warnings; all 284 example
scripts yield 3 warnings in 2 places, the lse_max_local reduce
destination and the two cur_batch_idx sites, with no false positives.
Note the two grouped_gemm examples also fail their own torch.allclose
check, but that is unrelated: initializing cur_batch_idx leaves the
mismatch at an unchanged 0.25 max absolute error over ~3.4k of 1572864
elements.
* [Transform] Discount a store's own write when checking its reads
CheckRead already took the reading node, so a cross-thread scope could
ask whether some other node writes the buffer. That is the rule keeping
a gemm accumulating into an untouched accumulator reported. Loads passed
no reader, so a store that read its own destination counted itself as
its own initializer: `s[i] = s[i] + A[i]` on shared memory was silent
while the identical code on a fragment was reported.
The enclosing store is now published while its subexpressions are
visited. A store is recorded as a writer of nothing but its own
destination -- an opaque call inside the value is recorded under the
call node -- so discounting it cannot suppress a report for any other
buffer.
Also covers the reducer v2 idiom introduced by #2940. reducer_init,
reducer_update and finalize_reducer inherit the default
GetReadBeforeWriteRegions and need no override, which is the same
result the pass already gets for T.reduce.
* [Transform] Exempt a parameter buffer from the check at any scope
SeedParams treats every parameter as written by the caller, but the
cross-thread branch of CheckRead asks only whether another node in the
body writes the buffer, and the caller is not one. A shared-scope
parameter therefore reported while an identical local-scope parameter
stayed silent.
Record the parameters in their own set and admit them ahead of the
scope-dependent question, so WrittenByAnotherNode is left untouched for
every buffer the body has to establish itself.
* [Transform] Print the buffer name and widen the parse fallback
Streaming the Buffer sent a one-word diagnostic through the reflection
repr printer. The name is a String with a direct stream overload and is
what the message wanted either way, so print it directly.
TryParseOperator caught tvm::ffi::Error, which every builder failure
seen so far raises. It derives from std::exception, so catching the base
extends the documented degrade-rather-than-abort rule to the failures
that would not.
* [Example] Initialize the buffers the new check reports
Three sites, all read before anything writes them. None is new and none
is known to misbehave; they are benign on the inputs these examples use,
which is why nothing has surfaced them.
grouped_gemm fwd and bwd allocate `cur_batch_idx` and immediately carry
it forward through `T.if_then_else(cond, i, cur_batch_idx)`, so a false
predicate on the first iteration returns whatever the variable held.
`cur_batch_idx = 0` is what example_grouped_gemm_fwd_ptr.py already does,
and has done since its first commit (#1923) — a second author writing the
same loop initialized it unprompted.
flash_decoding passes an uninitialized `lse_max_local` to
`T.reduce_max(..., clear=False)`, which makes the destination a read by
the operator's own semantics. It is the only write to that buffer and it
sits two lines below `T.clear(lse_logsum_local)`, so `clear=True` is what
was meant.
Verified: all 284 example scripts now compile with zero warnings.
example_mha_inference.py still reports "All checks passed!" and its
output is bit-identical to the current form on an RTX 3060 Ti.
The two grouped_gemm examples still fail their own torch.allclose check.
That is a separate, pre-existing issue untouched here: the mismatch is
unchanged at 0.25 max absolute error over ~3.4k of 1572864 elements
whether or not cur_batch_idx is initialized.
* [Refactor] Read the buffer-init config inside the pass
The pass function already receives the PassContext, so it can gate
itself on tl.disable_buffer_init_check rather than have every backend
pipeline consult a Python helper first. This matches ThreadSync, which
reads tl.disable_thread_storage_sync the same way and is scheduled
unconditionally.
Drops should_enable_buffer_init_check and the five gated call sites.
The two tests that asserted on the helper now call the pass directly,
which also covers the case the Python gate never did: invoking
VerifyBufferInit outside the pipeline under a disabling PassContext.
This PR adds a new set of annotations for describing arbitrary warp specialization scheme, with the abstraction of pipelines. A pipeline supports 4 kinds of operations: producer_acquire, producer_commit, consumer_wait, and consumer_release. Behind the scenes, mbarriers are created to implement the synchronization between warps.
This PR introduces a new pass, MaterializeWSSchedule, that transforms an annotated kernel to a warp-specialized kernel.
The set of infrastructure enables further development of fully automatic warp specialization in TileLang. I am working on this.
* [BugFix] Skip DecoupleTypeCast on Evaluate roots to keep cp.async operands in their address spaces
* [Test] Assert cp_async_gs lines exist before checking their operands
The codegen regression test filtered kernel source with
, so a source with no global->shared calls
left the operand assertion with nothing to inspect and the test passed
vacuously. The preceding did not catch
that: cp_async_commit/cp_async_wait come from the pipeline's sync
skeleton, a separate emit site, and match that substring on their own.
Collect the cp_async_gs lines, assert the collection is non-empty, then
check each one for _local_cast operands.
A large number of tests import CUDA-only internals directly (for example
`tilelang.cuda.pipeline.CUDAPassPipelineBodyPrologue`) or exercise
NVIDIA-specific features such as TMA, LDGSTG, mbarrier, warp specialization,
tcgen05 and tfloat32. `src/cuda/CMakeLists.txt` only compiles a small
"ALWAYS" subset when `USE_CUDA=OFF`, so those FFI symbols are never
registered and the tests fail with:
AttributeError: module 'tilelang.cuda._ffi_api' has no attribute
'AnnotateDeviceBoundTmaCopies'
This is not ROCm-specific -- the same tests fail on any build without the
CUDA backend, including Metal and CPU-only.
Add `@tilelang.testing.requires_cuda` to those tests, plus three
portability fixes that keep coverage on non-CUDA backends instead of
skipping it:
* `test_tilelang_transform_layout_inference.py`: CUDA emits
`__launch_bounds__(N, 1)` while HIP emits `__launch_bounds__(N)`.
Match both.
* `test_tilelang_transform_pipeline_planning.py`: PipelinePlanning
deliberately strips pipeline annotations on ROCm targets other than
gfx950 (see the guard in src/transform/pipeline_planning.cc), so assert
the documented sequential fallback there.
* `test_tilelang_kernel_gemm.py`: the two f32f32f32 cases use `T.tfloat32`,
which is NVIDIA-only, so narrow `requires_cuda_or_cdna` to
`requires_cuda`.
Also refresh three stale RDNA tests: gfx12 is now a supported generation
(`get_arch` accepts 11 and 12), so they can no longer assert that gfx1200
is rejected. Force an unsupported generation instead, which is the only
way to reach the guard now that `target_is_rdna` reports only gfx11/gfx12.
Verified on 8x MI300X (gfx942, ROCm 7.2 container, torch nightly rocm7.2):
the ROCm suite goes from 105 failures to 0 with no loss of passing tests.
The VerifyParallelLoop data race check can report false positives on
shared buffer stores whose per-thread addresses cannot be proven
distinct (e.g. indices loaded from fragments). Disable it by default;
users can opt in via the TILELANG_ENABLE_DATA_RACE_CHECK environment
variable or by explicitly setting the tl.disable_data_race_check pass
config to False.
* [CUDA] Pack logical TMEM buffers into shared tcgen05.alloc arenas
LowerSharedTmem gave every T.alloc_tmem buffer its own tcgen05.alloc,
rounded up to a power of two. Small buffers therefore each burned a
32-column allocation: a 384-column accumulator next to three 4-column
block-scale buffers needs 396 columns but asked the hardware for
512 + 32 + 32 + 32 = 608 of the 512 a CTA has, so that kernel could not
be expressed at all. Kernels worked around it by declaring one wide
buffer and slicing scale factors out of it by hand.
Plan the allocations instead. PlanTmemArenas packs logical buffers into
shared allocations, and each buffer's column offset inside the one it
shares is folded into every address it forms -- both the encoded
coordinate of a BufferLoad, which tcgen05.ld/st consume, and the
(base, offset) argument pairs of the tl.ptx_tcgen05_* intrinsics.
Buffers at least 32 columns wide keep the 32-column alignment
tcgen05.alloc would have given them alone; narrower ones use the
4-column granularity of tcgen05.cp.32x128b.warpx4, the same boundary
CUTLASS and DeepGEMM place these operands on.
A buffer joins an existing allocation only when that strictly reduces the
total column count, so packing can never grow a kernel's TMEM footprint
and a kernel whose buffers already fit lowers exactly as before. The
mxfp8 block-scaled GEMM drops from 320 to 288 columns; flash attention
and single-accumulator GEMMs are untouched.
Buffers with an explicit T.deallocate_tmem keep their own allocation,
since releasing one early would release whatever shares it. A base
address that reaches an operation this pass cannot offset is rejected
rather than silently addressing the start of the allocation. A block
that allocates more than 512 columns is now reported at compile time
instead of failing inside tcgen05.alloc on device.
* [Test] Cover TMEM arena planning
Pins the plan the packer produces, not just that it produces one:
- 396 columns across four buffers become one 512-column allocation at
offsets 0 / 384 / 388 / 392, matching DeepGEMM's hand-written layout
for this kernel column for column
- block-scale factors share an allocation while the accumulator keeps
its own, because packing all three would cost more than it saves
- equal-cost cases keep separate allocations, so a kernel that already
fits is untouched
- allocations are issued widest first, as PTX requires
- packing never allocates more than separate allocations would, over six
shape combinations
- an explicit T.deallocate_tmem keeps its own allocation
- a dynamic pipeline-stage coordinate still carries the column offset
- a block over the 512-column budget is reported
* [Test] Roundtrip two TMEM buffers that share one allocation
Covers the tcgen05.st/ld address path for a buffer packed at a nonzero
column offset, which no existing kernel exercises: the block-scale
buffers that share an allocation today are never read back into
registers.
The wide buffer's roundtrip is the assertion that matters. The narrow
one would still round-trip if its store and its load dropped the column
offset together -- but its store would land on the wide buffer's first
columns, so only the wide readback catches it. Verified by mutation:
dropping the offset in LowerSharedTmem fails this test.
* [Test] Guard the TMEM transform tests on Blackwell
Gate every TMEM test the same way, matching the decorators the tcgen05
language and kernel tests already carry.
These three lower IR for a hardcoded sm_100 target and need no device, so
this does trade running them on any machine for consistency with the rest
of the TMEM suite.
* [BugFix] Skip descriptor TMA for device-bound copy bases
* [BugFix] Run TMA base annotation independently of WS
* [Refactor] Expose TMA base annotation as a standalone pass
* [Docs] Restore warp-specialization pipeline comment
* [Transform] Never classify side-effecting binds as replayable
IsReplayableScalarBind (introduced in #2245) decides whether a scalar
bind may be dissolved and its value re-evaluated at every use site
(IfStmtBinding, default-on) or replayed across pipeline stages
(PipelinePlanning, InjectSoftwarePipeline). The predicate checked the
dtype and read/write disjointness, but never the value's effect kind,
so a bind like
pos = T.atomic_add(counter[0], 1, return_prev=True)
used at two sites (bound check + store index) was inlined into both use
sites, executing the atomic once per use site instead of exactly once
and silently corrupting compaction/histogram/top-k kernels since
v0.1.11. The atomic intrinsic is correctly registered kOpaque; its
pointer argument enters the expression through tvm_access_ptr, which
the access collector records as a read, so the write-disjointness check
passed and the bind was declared replayable.
Reject any bind whose value has SideEffect > kReadState, mirroring the
effect guard StmtSimplifier already applies to let inlining. Read-only
values remain replayable; read consistency is enforced by the existing
write-set disjointness check.
The shipped examples/deepseek_v32/topk_selector.py uses exactly this
pattern and is miscompiled on v0.1.11/v0.1.12; its test only printed
recall and could not fail, so it also gains a recall assertion.
Regression tests:
- transform: IfStmtBinding must keep a bind whose value is an atomic
RMW returning the previous value.
- language (CUDA): compaction kernel using atomic_add(return_prev=True)
must contain exactly one AtomicAddRet and produce a permutation with
counter == N.
* [Transform] Validate scalar bind replay safety
Track access-pointer write masks during bind and pipeline analysis, and reject unsafe precomputed replay annotations during pipeline injection.
Add regressions for atomic target reads and side-effecting replay masks.
* Remove recall calculation from top-k selector
Removed recall calculation and assertion for top-k selection.
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
* [BugFix] Track GEMM accumulator writes in warp-specialization liveness collector
* [BugFix] Track scale-factor regions and region extents in gemm liveness
* [BugFix] Return previous value for atomic min/max/add (scalar + vectorized)
`atomic_min`/`atomic_max` (scalar) and `atomic_addx2`/`atomic_addx4`
(vectorized) typed their call with a non-void return under
`return_prev=True` but always selected the void `*_elem_op`, so codegen
emitted an empty binding (`prev = ;`) and nvcc failed.
Select the `*_ret_elem_op` variant under `return_prev`, as `atomic_add`
already does:
- min/max: the ret ops, codegen handlers, and device templates
(AtomicMaxRet/AtomicMinRet) already existed.
- addx2/addx4 (float32): register the ret ops (builtin.cc/.h) and add
codegen handlers emitting AtomicAddx2Ret/AtomicAddx4Ret (already in
the device headers) as expressions; the return is typed as the packed
2-/4-lane vector so the store matches.
- addx2 (fp16/bf16): the ret store LHS is `uint1`; codegen wraps the
half2/bf162 result with `tl::to_uint1`.
- addx4 (fp16/bf16): add AtomicAddx4Ret(half_t*)/(bfloat16_t*) device
overloads that run two AtomicAddx2Ret halves and return them packed as
uint2 (a fp16x4 is stored as uint2).
The 16-bit x4 path is per-pair atomic (two half2 atomics), not
whole-quad single-atomic -- hardware has no vector atomic wider than
`.f16x2` for 16-bit types. This matches the existing fp32-x4 scalar
fallback, and the docstring only promises "quad-width atomic addition".
Adds tests across scalar/x2/x4 and float32/fp16/bf16.
Fixes#2667
* minor fix for missing conditiions
---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
LeiWang1999