* [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.
* [CUDA][Cache] Publish immutable binary cache directories
* [Example] Drop stale torch JIT build locks in minference example
The vertical_slash_index extension builds into a cache directory shared
across CI runs. torch's FileBaton lock file is not tied to the owning
process, so a run cancelled mid-compile leaves it behind and every later
load() polls on it forever (hung CI at test_vs_sparse_attn). Guard the
build with an OS-level flock, which dies with its holder, and remove any
leftover FileBaton lock before invoking torch.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---------
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
The kernel cache and the autotuner result cache loaded params.pkl with
cloudpickle.load(), so a writable cache directory allowed arbitrary code
execution on the next cache hit. Kernel params are plain data, so store
them as JSON via TVM's reflection (tvm.ir.save_json / load_json), which
rebuilds IR nodes without executing code. Entries without params.json are
treated as incomplete and rebuilt.
Fixes#2817
PrimExpr overloads unary minus and invert but not unary plus, so
`Out[i] = +A[i]` raised TypeError at trace time. Route unary plus through
the builder, the same way `not` already is, and return PrimExpr operands
unchanged; Python values keep their native unary plus.
Fixes#2991
CUDA lowering emits hpow and hfmod for 16-bit pow and fmod, but no CUDA
toolkit declares either name, so fp16/bf16 fmod and variable-exponent
fp16 pow fail to compile. Integer-exponent pow is unaffected: TileLang
expands it into repeated multiplication before lowering.
Add half_t and bfloat16_t overloads that evaluate with powf or fmodf in
float32 and convert back, matching the existing bridge pattern. Remove
math.h's hpow-to-powf macro: it shadows the new overload at the call
site, and compiles only where the result is assigned directly.
Add codegen and numerical coverage for fp16 pow and both fmod dtypes.
The bfloat16 hpow bridge completes the backend layer, while T.pow on
bfloat16 remains blocked by an upstream TVM dtype check.
* [BugFix][Blackwell] Fix 256-bit vectorized codegen for fp4/int4/fp16
256-bit vectorization (sm_100+) broke nvcc compilation for several
packed dtypes because every (dtype x lanes) combination relied on
hand-written per-type glue:
- fp4 x 64 lanes: predicated copy materializes a typed vector temporary,
which needs `fp4_e2_64_t::operator=(ulonglong4)` (never defined; the
operator was copy-pasted onto the 128-bit fp4_e2_32_t with a broken
body) and `make_fp4_e2_64_t` (constructors stop at 32 lanes).
- int4/uint4 x 64 lanes: PrintType emitted the nonexistent CUDA type
`(u)int8`, and the packed-nibble broadcast only handled 16/32 lanes.
- fp16 x 16 lanes: broadcast calls tl::pack_float16x4 whose parameters
were `__half`, but codegen prints fp16 scalars as cutlass::half_t
(the bf16 variant already took bfloat16_t).
Instead of patching each cell, make the contracts generic:
- tl::load_global_256 / _conditional now return the pointee type
(static_assert sizeof == 32) and bit-cast internally, so assignment
type-checks for any 32-byte vector struct. The hand-written
operator=(ulonglong4) on fp8_e{4,5,8}_32_t become dead and are
removed, along with the misplaced fp4_e2_32_t one. The T.ldg256
builtin keeps its uint32x8 result via an explicit ulonglong4 cast.
- New tl::broadcast<V>(unit) replicates a packed unit across any vector
width; fp4 broadcast now emits one packed fp4x2 byte + tl::broadcast
for every lane count instead of a 64-arg make_ call.
- int4x64 maps to (u)longlong4 (like int8x32) and its broadcast packs
two 8-nibble fields per 64-bit component.
Adds a compile matrix test (12 dtypes x {copy, predicated copy, fill}
x {128b, 256b}) plus numeric checks for the previously broken cells, so
the next dtype/width gap fails in CI instead of on a user's machine.
On < sm_100 targets the 256b params compile at 128 bits and stay green.
* [Refactor][CUDA] Replace per-lane-count vector constructors with variadic packers
Follow-up to #3153: the make_<type>_<N>_t constructor ladders (5 for fp4,
15 for fp8) and several BroadcastNode special cases enumerated the same
(dtype x lanes) matrix by hand, so every new dtype or vector width needed
more glue. Replace the enumeration with two generic templates:
- tl::make_vec<V>(e0, e1, ...): builds a vector struct from one scalar
element per lane (byte-sized or larger elements).
- tl::make_fp4_vec<V>(e0, e1, ...): same for nibble-packed fp4.
and route the codegen through them:
- PrintVecConstructor emits tl::make_fp4_vec<V> / tl::make_vec<V> for
fp4/fp8 vectors, covering the Shuffle and element-wise assembly paths
(whose previous output for these types did not even compile).
- fp8 and int8 broadcasts of side-effect-free values emit
tl::broadcast<carrier>(byte) instead of per-lane make_ calls; the int4
16/32/64-lane packed broadcasts collapse into a single
tl::broadcast branch.
- Side-effecting broadcast values (e.g. casts of rng calls) keep their
evaluate-once-per-lane semantics: they fall through to the generic
per-lane path, guarded via tirx::SideEffect. The int32 rng broadcast
path and the fp16/bf16/f32x8 paths are intentionally untouched.
The make_fp4_e2_{2..32}_t and make_fp8_{e4,e5,e8}_{2..32}_t ladders are
deleted; the char-packing make_int* overloads in common.h stay because
the int8 Shuffle path still assembles vectors element-wise.
Verified on sm_103a: the 12-dtype x 3-op x 2-width compile matrix stays
green, new numeric tests cover non-constant fp4/fp8/int8 scalar fills,
and the ldg/ldst/vectorize/vectorized_cast/quantize suites pass
(228 tests) with TILELANG_DISABLE_CACHE=1.
* [Test] Update issue #2480 assertions for the tl::broadcast emission
The int8x32 broadcast is now emitted as tl::broadcast<(u)longlong4>(byte)
instead of make_(u)longlong4 with 32 scalar args; the exact-value checks
that guard the original issue (upper 32 bits of each field zeroed) are
unchanged.
plan_rasterization compared the runtime compute-capability string lexicographically. That classified sm_100 and sm_120 as older than sm_80 and disabled rasterization on Blackwell targets.
Use the target-derived integer sm_version for CUDA targets while preserving existing non-CUDA behavior. Cover pre-Ampere, Ampere-and-newer, L2-fit, and non-CUDA cases without requiring a GPU.
`for ... else` and `while ... else` inside a kernel traced without error
while the `else` body was silently dropped from the IR. Raise a clear
NotImplementedError from the AST rewriter instead; `if ... else` is
unaffected.
Fixes#2946
* [BugFix][Carver] Gate _legalize_info fallback on sm_version
TensorCorePolicy._legalize_info compared the bare compute-capability values reported by CUDA against sm_-prefixed literals. The fallback therefore selected one pipeline stage and disabled asynchronous copies even for SM 80 and SM 90.
Gate the fallback on the normalized numeric sm_version while preserving non-CUDA behavior. Keep the membership at {80, 90}, matching the existing matmul_analysis tag path; do not enable targets above SM 90 without real-machine validation of the pipeline and asynchronous-copy effects.
Constraint: No GPU above SM 90 was available for functional or performance validation, so the supported fallback set remains {80, 90}.
Rejected: Broaden the fallback to additional SM versions | It would exceed the existing tag-path policy, and above-SM-90 behavior has not been validated on hardware.
Confidence: high
Scope-risk: narrow
Directive: Keep fallback and tag-path SM membership aligned until additional architectures have real-machine evidence.
Tested: Ruff check and format; Python syntax compilation; 10-case _legalize_info source behavior harness.
* Apply suggestion from @coderabbitai[bot]
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
---------
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
* [BugFix][Hopper][Blackwell] Read the WGMMA/UMMA K-panel stride from the layout
A shared operand sliced along MN silently computed wrong results whenever
block_K spanned more than one swizzle atom:
T.gemm(a_s, b_s[0:64, :], c0, transpose_B=True) # wrong past the first K atom
When block_K exceeds one swizzle atom the shared layout splits K into panels
stored panel-major -- a (128, 128) bf16 tile decodes as
(128,(64,2)):(64,(1,8192)) -- so stepping to the next K panel moves by the
*buffer's* MN extent. The atom offset formulas instead reconstructed that step
as `mn_extent * swizzle_atom_elems` using the extent of the operand, which is
smaller for a slice. The slice origin was already layout-derived and correct, so
only K atoms past the first were affected: the two halves of a split buffer read
each other's rows, in bounds and without any diagnostic.
block_K == 64 was correct only by accident -- at 16-bit dtypes that is exactly
one 128B atom, so `ki // k_atom_size` is always zero and the bad multiplier was
multiplied away. MN-major operands never reach the term at all.
Take the panel stride from the decoded layout instead of reconstructing it:
cute.restrict already returns the sublayout carrying it, and slicing does not
change a stride. Expose it as `k_panel_elems` on WGMMADescriptorParams /
TCGEN05DescriptorParams, None when the layout has a single K panel or the
operand is MN-major -- exactly the cases where the term is multiplied by zero,
so those keep the previous expression and behaviour.
The same defect was duplicated across all seven atom emitters: dense Hopper
(ss/rs), sparse Hopper (ss/rs) and Blackwell (ss/ts/block-scaled).
Verified on H200 (sm_90a):
- Whole-buffer path is unchanged: across 120 K-major layouts (fp16/bf16/fp8_e4m3/
fp32 x 5 MN extents x 6 K extents), k_panel_elems is either None or exactly
equal to the old reconstruction, for both descriptor builders. The fix is a
strict generalization of the path that already worked.
- New test_tilelang_kernel_gemm_sliced_operand.py cross-checks a sliced operand
against one-buffer-per-slice over operand x majorness x block_K; it fails on
the four K-major block_K>64 combinations before this change and passes after.
- A layout-level assertion covers the Blackwell descriptor, which has no numeric
coverage here for lack of the hardware.
Closes#2964
* [BugFix][Hopper][Blackwell] Refuse a multi-panel K mode with no constant stride
Review follow-up on the K-panel stride decode.
The decode returned None whenever the K mode was not a canonical 2-tuple, and
every consumer then fell back to `mn_extent * swizzle_atom_elems` -- the exact
reconstruction this series removes. The fallback is only sound for a single-panel
layout, where `ki // k_atom_size` is always zero, and nothing checked that. So a
swizzled multi-panel layout that failed to decode would silently re-enter the
wrong-code path: the same shape as the original defect, where a formula's
precondition was assumed but never enforced.
Move the decode into a shared `decode_k_panel_elems` used by both descriptor
builders, and make it assert instead of returning None when a swizzled operand
spans more than one panel without exposing a constant panel stride. Panel count
is rounded up, so a K extent that is not a whole number of atoms cannot slip
through as single-panel.
Also tighten the accept condition: reading `stride[-1]` as the panel step is only
valid once the leading sub-mode is confirmed to be the contiguous atom. On a
panels-first `(2, atom):(panel_stride, 1)` ordering the old condition would have
returned 1 -- a silently wrong step rather than a refusal.
Consumers now read the step through `k_panel_stride(mn_extent)` on the params
dataclass, so the contract and its single-panel default live in one place
instead of being restated at eleven call sites.
New test_tilelang_cuda_wgmma_operand_layout.py covers the decode directly
(hardware-independent): canonical multi-panel strides, the single-panel and
MN-major None cases, and each refusal -- panels-first, flat multi-atom, a
three-sub-mode K, and a partial-atom extent.
Re-verified on H200: the whole-buffer invariant is unchanged by the tightening
(120 K-major layouts x 2 builders, 60 single-panel, 0 disagreements), every
multi-panel K slice is still accepted with the buffer-derived stride, and the
full suite set passes (174 passed, 40 skipped).
16-bit lowering emits half-style names for sinh, cosh, tanh, atan, erf,
and nearbyint, but no CUDA overload of those names accepts CUTLASS
half_t or bfloat16_t. Most toolkits do not declare the names at all, and
the overloads that do exist take native __half or __nv_bfloat16, which
the wrappers have no implicit conversion to, so these operations fail
everywhere.
Add hsinh, hcosh, htanh, hatan, herf, and hnearbyint overloads for both
wrapper types. Each evaluates in float32 and converts back, following
the existing htan bridge. This completes 16-bit bridge coverage for the
unary floating-point operations that CUDA lowering already handles.
Extend the 16-bit intrinsics test kernel to cover all six operations.
Introduce two CUDA-specific math intrinsics with guaranteed instruction
selection:
- T.fma(x, y, z): fused multiply-add whose fusion (single rounding) is
guaranteed regardless of NVCC contraction settings.
- T.fmul(x, y): multiply with an explicit boundary that is never
contracted into a neighbouring FMA.
Unlike the ieee_* family (which carries an explicit rounding mode and
stays scalar), both ops fix round-to-nearest and are marked
TVectorizable, so T.Parallel loops containing them keep their vectorized
loads/stores:
- scalar: __fmul_rn/__fmaf_rn (fp32), __dmul_rn/__fma_rn (fp64),
__hmul_rn/__hfma (fp16/bf16)
- vectorized on SM100+: packed x2 helpers tl::mul2/tl::fma2
- vectorized elsewhere: per-lane scalar calls
- CuTeDSL: tl.ieee_fmul/tl.ieee_fmaf with rounding="rn", expanded per
lane through a register tensor for vector dtypes
The packed x2 pair-decomposition logic is extracted from
PrintVecBinaryOp into a reusable EmitPackedX2Call helper (no behavior
change), alongside a new EmitPerLaneScalarCall fallback.
* [BugFix] Always vectorize `T.Parallel` loops
* Fix CI failures
* Restore always-vectorize accidentally reverted by the CI-fix commit
f447919b captured a transient working-tree state from a concurrent
baseline-vs-patched benchmark flip, reverting the gate removal in src/
while keeping the new-policy tests: CI then ran the old gated compiler
against them (scalar var-loop -> no float2; positional if/for check
broken by the copy loops the larger tile introduces).
Bring back the parameter-less LowerParallelLoop and rewrite the issue
1734 assertion as containment: the guarded loop must sit inside the
hoisted if block, which holds on both sm90 and sm100 regardless of how
many copy loops precede it.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---------
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Some ldmatrix lane maps step beyond the final SM75 MMA micro-tile, and
column overflow must carry into the next row. Layout remapping previously
linearized the source indices and optional offset before converting them
back to multidimensional coordinates, which normalized these addresses as
a side effect. Since #2836 preserves the multidimensional indices, the
out-of-range coordinates reach the swizzle layout unchanged and the native
SM75 MMA GEMM paths fault on illegal shared-memory accesses.
Wrap each complete row-major source offset within its operand region in
ldmatrix_a and ldmatrix_b before rebuilding the source coordinates.
Leading dimensions remain untouched, so the wrap cannot cross into another
stage of a multi-version buffer.
When the coordinates are provably in range, the simplifier removes the
wrap. SM70 and SM80 generated sources remain byte-identical, while the
SM75 addresses match the pre-#2836 address mapping.
Add multi-stage SM75 GEMM coverage and a lowered-offset bounds regression
that cross-compiles explicitly for sm_75 and therefore runs on any CUDA
runner.
_tir_packed_to_unsigned_convert_with_zeros performed the zero-point subtraction in the unsigned packed storage dtype before casting to the requested output dtype. Values below the zero-point therefore wrapped modulo the storage width.
Cast the unpacked field and zero-point to the output dtype before subtraction so asymmetric dequantization preserves negative values.
Add CUDA regression coverage for uint32-packed 4-bit values decoded with a non-zero zero-point.
Fixes#2947
* [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>
* [Compiler][Z3] Bump 3rdparty/tvm to materialize Z3 solvers lazily
Z3Prover::Impl created a z3::solver eagerly in the Analyzer constructor
and translated every Bind/EnterConstraint into z3 constraints up front.
Analyzers are scratch objects created in huge numbers during lowering (a
reducer-heavy kernel compile creates ~67k of them, ~100k binds), while
only ~0.3% ever reach the Z3 fallback of CanProve, so the eager setup
dominated layout-inference-heavy compiles.
With the bump, binds and constraints are journaled and the solver is
created and replayed on first use. End-to-end compile of the worst
regressing reducer kernels drops 40-60%, and tilelang.lower roughly
halves across the board. Language suite 1168 passed / 137 skipped,
reducer_v2 75 passed, arith 89 passed; generated CUDA byte-identical
under disabled ASLR.
tvm changeset: https://github.com/tile-ai/tvm/pull/68
* [Cleanup] Drop Z3ContextScope wrappers
With the bumped tvm, every materialized Z3 solver owns a private context
(tile-ai/tvm#69), which subsumes the per-compilation isolation these
scopes provided; the scope functions are now documented no-ops. Remove
the wrappers around lowering and grouped compilation, and rewrite the
scope-lifetime arith test to assert what it actually guarded: a cloned
Analyzer's Z3 handles stay valid after the source is destroyed.
Validation (cold caches): layout suite 298 passed (including the
z3-context isolation tests, which cover the grouped-compile path),
language suite 1172 passed / 137 skipped, arith 89 passed, reducer_v2
75 passed.
* Point 3rdparty/tvm at the merged lazy-Z3 changes
tile-ai/tvm#68 and tile-ai/tvm#69 are both merged; replace the PR-branch
pin with the tilelang_main commit.
[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.
math.h aliases several half-style names to cutlass::fast_* when a kernel
uses selected fast-math intrinsics. CUTLASS provides only float and double
overloads for fast_log, fast_sqrt, fast_sin, and fast_cos, and adds half_t
overloads only for fast_exp and fast_tanh. Overload resolution can therefore
leave a float result that cannot copy-initialize the CUTLASS wrapper.
Once a kernel includes math.h, for example through T.__exp, these aliases
also affect matching plain math calls in the same kernel. Their compilation
could therefore depend on whether another intrinsic pulled math.h in.
Add half_t and bfloat16_t overloads for fast_log, fast_sqrt, fast_sin, and
fast_cos, plus bfloat16_t for fast_tanh. Each evaluates in float32 and
converts back, following the existing fast_exp(bfloat16_t) bridge. Calls
that already compiled retain the same numerical path.
Add compile-and-run coverage for both wrapper types.
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.
* Steer reducer destinations during layout inference
* Remove extra reducer steering attempt filtering
* handle reducer problem3
* bug fix
* [Refactor] Give reducers a first-class PartialFragment layout solved by layout inference
The reducer's physical plan (narrow vs wide partials) used to be decided
after inference by ReducerPlanAndMaterialize, with dst-steering emulating
the missing propagation through side channels (wake tables, hint plumbing,
ownership gates); the same induced layout was computed twice and kept
consistent by convention only, and the reducer split the use-def graph
into disconnected components.
This makes the reducer's layout a first-class solver value that fully
determines the physical plan:
* New layout kind `tl.PartialFragment` (PartialFragmentNode : FragmentNode):
Fragment algebra whose replication coordinate carries an explicit
combine decomposition — under a low-bits convention enforced at
construction, `_rep % combine_size` enumerates addend lanes the finalize
collective must reduce and `_rep / combine_size` enumerates equal-value
copy groups. (storage algebra, combine width) is the complete plan: the
collective steps are derivable from the node alone (CombineSteps), the
FullParticipant object (combine == replicate == participants, thread ==
_rep) is the wide fallback, and combine width is part of layout identity
(same-storage different-decomposition proposals widen at the commit
point instead of unifying silently). A boundary-preserving condensation
(CondenseReplicateVarKeepingBoundary) keeps the convention true through
projection; kind-aware IsEqual in both dispatch directions prevents
silent unification with plain Fragments.
* Update nests (ParallelOp) record their reducer_update sites and, once
their loop layout solves, propose the induced partial (per-site
AnalyzeReducerUpdateSite proofs) or the participant-wide plan. The
commit point resolves multi-site disagreement with a monotone
widen-on-conflict rule; zero-update epochs are seeded wide before any
inference level runs; tl.reducer_force_baseline pins the proposal.
* finalize_reducer reads the solved partial and steers an unconstrained
destination to its post-collective reading. The wake table, the
LayoutInferArgs.reducer_update_sites hint plumbing and the
reducer_use_list_ component glue are deleted: reducers travel ordinary
use_list_ buffer edges.
* ReducerPlanAndMaterialize reads the plan off the node: storage =
AsPostCollective(), steps = CombineSteps(), wide lowering iff the node
is the FullParticipant object. TryNarrowPlan is retired; what remains
(DeriveNarrowLowering) derives only the lowering details that live
outside the layout — the packed-lane heuristic and the
destination-containment gate (which still demotes gracefully to the
wide plan, preserving the unsatisfiable-steering fallback contract).
The destination override chains are kept as a verified-no-op fallback
with retirement telemetry.
* T.annotate_layout accepts a PartialFragment on reducer buffers, pinning
plans the automatic solution cannot reach (e.g. a middle plan between
narrow and wide). Annotated partials are strict: a disagreeing update
site raises with the narrow-rejection reason, and
ComputeReducerPinnedCandidate back-propagates the annotation through
the update loop into its input fragments.
* Layout cost models price register-array spills: a fragment access whose
physical slot index depends on the thread variable (detected by the
shared FragmentThreadIndexProbe, extracted from
HasThreadDependentFragmentIndex) demotes the whole array to local
memory, charged as byte-denominated traffic in both models' `mem`
channel so it competes with — rather than vetoes — global-bandwidth
estimates. ChooseBestCandidate's anti-spill preference now applies
regardless of where the loop stores (a fragment->fragment publication
copy spills exactly like a fragment->global one). With generation and
pricing combined, inference produces the wide publication chain's
replicated layouts on its own: the reducer_v2 suite and the
mhc_post_bwd production kernels pass with the destination override
fully disabled.
Testing (all with the kernel cache disabled): reducer_v2 suite 70 passed
(new coverage: pinned middle plans, annotation steering input fragments,
conflict diagnostics, combine-step roundtrip, same-storage
different-combine widening at commit); language suite 1261 passed;
layout-inference cost-model tests and maint/layout_inference goldens
unchanged; generated CUDA for existing reducer kernels is byte-identical;
mhc_post_bwd correctness passes and its six benchmark sizes show
no latency change with zero STL/LDL in SASS.
---------
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
[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.
Lei Wang