Skip to content

fix(rocke): decide every tiling bound check with one rule, so ragged tiles never read or write outside a tensor - #15026

Draft
cgmillette wants to merge 11 commits into
developfrom
users/cgmillette/rocke/ragged-clipping
Draft

cgmillette wants to merge 11 commits into
developfrom
users/cgmillette/rocke/ragged-clipping

Conversation

@cgmillette

@cgmillette cgmillette commented Nov 4, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

The tiling verbs load_fragment / store_fragment could read and write outside a tensor without any error. Which sides of a window access got a run-time check was decided in three places, each in its own way, and each trusted what it could not see: a runtime origin was taken as grid-aligned, a tile-aligned length as a proof, and the upper end of a leading axis as inside. A tile could overhang its tensor unmasked. Alignment claims, 32-bit element indexes and stores through overlapping views had wrong-answer defects of their own.

This PR fixes every one of those wrong answers. One rule decides every bound check: the shared bound-check rule of tensor_transforms. A kernel author writes the same calls for a ragged tile as for an aligned one:

M = b.param("M", I32, ge=1, le=1 << 20)                  # facts KernelLauncher.from_artifact checks before each launch
K = b.param("K", I32, ge=1, le=1 << 20)
a_td = make_tensor_desc((M, K), (K, 1), F16)             # lengths and strides may be parameter Values
m0 = exit_if_tile_outside(b, b.mul(b.block_id_y(), b.const_i32(TM)), M)  # guard the block origin
a_win = make_window(a_td, (m0, k))                       # a window positions a tile; no length of its own
print(explain_bounds_checks(b, a_ptr, a_win, a_desc))   # which sides are checked, and how
a_frag = load_fragment(b, a_ptr, a_win, a_desc, lane)    # outside the tensor -> 0, never touched

A window that is proven inside its tensor emits no check. A checked window keeps its vector width and masks each access by its validity. No instruction touches an element outside the tensor. The fast path for tile-multiple shapes is a second kernel whose lengths are stated as multiples of the tile, picked at launch with first_supported.

It no longer depends on #14448 (parameter facts and the checking launch door): #14448 has merged, and this branch is rebased onto develop with only its own eleven commits. One test (test_runtime_extents::test_a_window_over_runtime_extents) needs #14846 (unique value names); it passes with #14846 merged in.

ISSUE ID : AICK-2466
ISSUE ID : AICK-2579

Technical Details

Eleven commits:

  1. fix(rocke): require TileDesc shape to equal its layout extent
    • A TileDesc whose shape is smaller than the extent its layout addresses let registers fall outside the region the bound checks cover. It is now refused at construction.
  2. feat(rocke): TensorDesc lengths and strides take parameter Values
    • A length or stride may be an i32 Value with stated facts: a b.param declared with them (checked on the host before each launch through KernelLauncher.from_artifact) or a b.gpu_checked_value result (checked on the GPU). TensorDesc reads them through feat(rocke): parameter facts checked on the host, GPU checks for device-only values, and a checked launcher #14448's homes (has_stated_facts, stated_facts_of, marker_facts, Value.param) and keeps no copy.
    • Each value that cannot be used is refused, with a message that names the fix.
    • check_gemm_shapes agrees runtime extents only when they are the same Value. The static readers (element_span, the built-in LDS swizzle, the cooperative fit check, the recorded-pipeline address analyses) refuse a runtime extent.
  3. feat(rocke): fold runtime strides and slice offsets in tensor_transforms addresses
    • The address emitters multiply by a runtime stride at run time and fold each axis's slice begin.
    • A slice can start past 2^31 elements.
    • A tensor whose begins are all 0 emits the same IR as before.
  4. feat(rocke): window positions with no length of their own
    • make_window(tensor, origin) keeps one origin per tensor axis and no region. bounds= is gone, and with_origin moves one axis.
    • A negative origin (a convolution halo) is accepted and masked.
  5. feat(rocke): guard block origins in the kernel prologue
    • exit_if_tile_outside(b, origin, length) emits one unsigned compare and an early return in the prologue, then a range marker (by="exit_if_tile_outside") whose result the IR range walk reads.
    • With a tile-multiple length, the whole tile is proven inside.
    • The marker's writer list gains the name in both engines, as one constant each (EXIT_IF_TILE_OUTSIDE_WRITER in Python, ROCKE_EXIT_IF_TILE_OUTSIDE_WRITER in C++).
    • IRBuilder.at_kernel_top_level, a new read-only property, says the builder is emitting into the kernel body itself, not inside a loop or an if. The guard uses it to keep its early exit in the prologue, instead of reading the builder's private region stack. The C++ builder needs no twin: it already exposes the current region publicly, and it has no guard.
  6. fix(rocke): decide clipping with the shared bound-check rule
    • explain_bounds_checks(b, ptr, window, tile_desc) returns a BoundsCheckPlan and emits nothing. For each side of each tensor axis it says whether the side is proven, or checked by one scalar compare or one compare per lane, and how a vector's validity is decided.
    • The verbs take their decision from the same function.
    • When a side is checked only because a loop stops at a runtime length that is not stated a multiple of the loop step, print(plan) ends that side's line with the fix: -- state multiple_of=16 on K's b.param declaration, or, for a length checked on the GPU, the b.gpu_checked_value call to state it in. The IR range walk records such a stop on the loop counter's part of the split. The plan keeps it only when the bound-check rule proves the side with the range that fact would give.
    • A new internal safe_access module delivers every access that is not a proven window with int strides. The per-wave address goes on the 64-bit pointer, the per-lane part is a 32-bit offset, and a checked access runs under its validity. No address an invalid lane computes carries an in-bounds promise, and no load is clamped to element 0.
    • A pointer parameter declared addr_space="constant" is refused at build, by the verbs and explain_bounds_checks alike, on every form except the one-index form. Those forms do arithmetic on the 64-bit pointer, and core's global_ptr_* ops take only a global pointer. The refusal is a TypeError, as core's refusal of a pointer in the wrong space is.
    • The one-index form stays exempt on purpose. It passes the pointer on to core's global loads and stores unchanged, and core gives it valid IR: a load reads through a constant-space pointer, and a store through one is refused at lowering (fix(rocke): refuse buffer_rsrc operands its lowerings cannot take #14842, on develop).
    • The manual, interleaved (LDS-staged included) and CRC demos guard their block origins and launch through KernelLauncher.from_artifact. Every global window in them is proven and keeps the one-index form.
    • The plan names the form an access takes, plan.form (GlobalForm.INDEX32 / PLACED / MASKED, exported beside AccessPath); its value is the text print(plan) shows.
    • load_fragment refuses a non-zero pad, and both verbs refuse a coherency other than 0, before emitting anything and on every path: neither is implemented, so the value used to be ignored. An IR value given for either is a TypeError.
    • The verbs, the plan and safe_access share their helpers through public module names (not the package door), never another module's private names: clipping's plan_window, contiguous_run, element_bytes, global_store_width and others; encoding's lane_terms / register_coord_offsets; tensor_transforms' required_checks_within (renamed from _required_checks) and Interval.numeric_ends.
    • An element's byte width comes from core's dtype_info (element_bytes, which refuses a dtype that is not whole bytes); the width table this PR used to move out of emit.py is gone, and the dtypes the verbs vectorize have one home.
    • The pipeline recorder records each global access at the width the verb emits it, from the access's plan (global_store_width for stores): a checked or placed store is one wide store per run, so the coverage witness accounts for every emitted store.
  7. fix(rocke): claim only the alignment a wide access really has
    • The align an access claims is derived from facts only: the pointer's align=, the begin offset, the origin's multiple-of fact and every other stride, capped at the access width and 16 bytes. A runtime stride's multiple-of fact is read through tensor_transforms' end_residue, now public.
    • Develop's width rule is unchanged.
  8. fix(rocke): refuse runtime extents on LDS tensors
    • rocKE allocates LDS statically, with a fixed size, so an LDS tensor needs int lengths and strides. One check on TensorDesc refuses a runtime length or stride before anything reads the layout, with one message opened by the caller's name: load_fragment: lengths[0] is a runtime value (S) -- rocKE allocates LDS statically, with a fixed size, so an LDS tensor needs int lengths and strides.
    • load_fragment and store_fragment on an LDS buffer, the cooperative LDS copies, the built-in swizzle and TensorDesc.element_span all use it.
  9. fix(rocke): never take the 32-bit-index path for an offset that does not fit i32
    • A constant element offset outside i32 now goes on the 64-bit pointer. Before, it was passed as an index that LLVM wrapped by 2^32 elements.
  10. fix(rocke): refuse stores through aliasing tensor views
    • store_fragment refuses, at build, a store through a view whose overlap is proven from the description, naming two positions that meet.
    • A view the build cannot decide either way is trusted. Loads are unchanged.
  11. feat(rocke): tiling proofs read relations stated between kernel parameters
    • Relations stated on b.param (ldc ge N, sb ge product_of(M, ldc)) prove padded and batched C views free of aliasing.
    • A relation can also bound a proven window's offsets inside one 32-bit index. BoundsCheckPlan.address_facts names the relations a proof used.

No core emission changes: byte identity is green at every LLVM flavor, and no representative-IR or differential golden changes. The tiling demo IR snapshot is re-blessed in commit 6, with the reason stated there: the block origins are guarded in the prologue, and the size parameters state the sizes the kernel bakes. The K-loop ISA is unchanged up to register numbers.

One test is red until #14846 (never hand out the same SSA name twice) lands: test_runtime_extents.py::test_a_window_over_runtime_extents_builds_and_both_engines_lower_it[True], where the C++ engine reads back an SSA name the Python builder issued twice. With #14846 merged in, it passes in both engines. Whichever of this PR and #14846 lands second re-records the three tensor_transforms emit goldens and the demo IR snapshot, and updates one SSA-name pattern in test_exit_if_tile_outside.py.

Test Plan

  • Every refusal, with its exact text, including the constant-space pointer from explain_bounds_checks, load_fragment and store_fragment on the masked and placed forms. That refusal emits no IR, and the same pointer is not refused on the one-index form.
  • The plan equals the shared rule on the walked positions: guarded origins, loop counters, the tail-tile origin, leading axes, scalar checks and the vector axis of a rank-3 window. The verb applies the plan it explains.
  • The fix the print names: for a parameter and for a GPU-checked length, keeping a multiple the length already states (24 with step 16 gives 48). No fix is printed when K is stated a tile multiple, when the side is checked for another reason, or when the length's facts are not stated by the author.
  • IRBuilder.at_kernel_top_level at the top level, inside loops, ifs and both branches of an if/else, nested, and after leaving them.
  • The pad and coherency refusals on the proven, checked and LDS paths, with exact type and text and nothing emitted; element_bytes over byte and sub-byte dtypes and an unknown name; the store width global_store_width gives on each form against the stores the verb emits.
  • A recorded ragged copy: its wide checked load and store are recorded at width 4, the coverage witness accounts for every emitted op, and the recorded kernel is byte-identical to the plain build.
  • check_gemm_shapes gives its "check it once" advice only for two gpu_checked_value calls on one value, not for two exit_if_tile_outside guards.
  • The LDS refusal text from both verbs (with and without the swizzle), the cooperative copies and element_span.
  • The IR and ISA of each form:
    • no in-bounds promise and no clamp on a checked access;
    • the pointer-plus-offset form with no compare for a runtime stride;
    • the row compare of an M-ragged GEMM outside its K loop on gfx90a and gfx942.
  • On a gfx90a GPU, ragged and misplaced windows read and write exactly the valid elements, with every guard zone intact: f16/f32, M- and K-ragged, runtime length 0, negative origin, odd stride, a batch index outside the batch, a convolution halo against im2col, offsets past 2^31, and the CRC demo.
  • The demos' global accesses stay unmasked, and each demo's parameter contract accepts exactly the launch of its shape.
  • Every commit builds and passes ctest and its own tests on its own, with both engines in strict mode.

Test Result

On develop 21e3ca4350d:

  • Byte identity green at llvm20, llvm22, llvm23 and llvm24.
  • ctest: all pass.
  • Tiling tests, both engines strict: all pass except the one SSA-name test above, which passes with fix(rocke): never give two values the same name #14846 merged in.
  • The installed-layout rocke_tensor_transforms_pytest CTest entry passes.
  • Full platform suite: the same failures as develop, none new.
  • GPU cases (gfx90a): not yet run on this head; the host runs above had no GPU visible, so those cases were skipped. The tiling demo IR snapshot is unchanged by the rebase.
  • pre-commit passes, and ruff reports nothing new.

Submission Checklist

🤖 Generated with Claude Code

@therock-pr-bot

therock-pr-bot Bot commented Nov 4, 2026 •

Copy link
Copy Markdown

✅ All Checks Passed — Ready for Review

Check Status Details
📝 PR Description ✅ Pass —
⛔ Forbidden Files ✅ Pass —
🧪 Unit Test ✅ Pass —
🔎 pre-commit ✅ Pass —
🚫 Draft PR 🔜 To Be Enabled —
🚩 Feature Flag 🔜 To Be Enabled —
📊 Code Coverage 🔜 To Be Enabled —
🤖 therock-pr-bot ✅ Pass —

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

@therock-pr-bot

therock-pr-bot Bot commented Nov 4, 2026

Copy link
Copy Markdown

🎉 All checks passed! This PR is ready for review.

@cgmillette

Copy link
Copy Markdown
Contributor Author

GPU check at head 7b60d53 on gfx90a (MI210): the whole tiling test tree (python/rocke/tests/helpers/tiling, default engine) -- 1854 passed, 104 skipped, 0 failed.

@cgmillette
cgmillette force-pushed the users/cgmillette/rocke/ragged-clipping branch from 7b60d53 to 0882e32 Compare November 4, 2026 04:05
@cgmillette

Copy link
Copy Markdown
Contributor Author

Applied four decisions from today, each folded into the commit that owns the code (still eleven commits on #14448):

  • The printed bound-check report now names the fix when a side is checked only because a loop stops at a length that is not stated a multiple of the step, for example -- state multiple_of=16 on K's b.param declaration (commit 6). For a length checked on the GPU it names the b.gpu_checked_value call instead. It prints nothing when the side is checked for any other reason.
  • exit_if_tile_outside reads a new read-only builder property, IRBuilder.at_kernel_top_level, instead of the builder's private region stack (commit 5). The guard is still allowed only at the start of the kernel.
  • A runtime length or stride on an LDS tensor now gets one message from one place: the verbs, the cooperative LDS copies, the swizzle and element_span all say that rocKE allocates LDS statically, with a fixed size, so an LDS tensor needs int lengths and strides (commit 8).
  • The one-index global form stays exempt from the constant-space pointer refusal. The comment and the test now say that it relies on fix(rocke): refuse buffer_rsrc operands its lowerings cannot take #14842 for valid IR (commit 6).

@cgmillette
cgmillette force-pushed the users/cgmillette/rocke/ragged-clipping branch from 0882e32 to e077c8d Compare November 4, 2026 07:51
@cgmillette

Copy link
Copy Markdown
Contributor Author

Rebased onto #14448's new head (edc7c2da6ce) and folded in the quality fixes from the latest review, each into the commit that owns the code.

Rebase: this PR now reads stated facts and range markers only through #14448's homes (has_stated_facts, marker_facts, Value.param, ASSUME_RANGE_OP, GPU_CHECKED_VALUE_WRITER). The exit_if_tile_outside writer is one constant in each engine. Every text that said parameter facts are checked "before every launch" now says which launch path checks them: KernelLauncher.from_artifact.

Quality fixes:

  • load_fragment refuses a non-zero pad, and both verbs refuse a coherency other than 0, at the top of the verb and on every path (both used to be ignored). Tests pin the type, the text, and that nothing is emitted.
  • The byte-width table moved out of emit.py is gone. clipping.element_bytes reads the width from core's dtype_info and refuses a dtype that is not whole bytes. The set of dtypes the verbs vectorize has one home. fits_i32 and the descriptor's i32 limit read core's int_type_limits. Core's own element_bytes (in another open PR) is not on this base; when it lands, clipping.element_bytes should be replaced by it.
  • The pipeline recorder now records each global access at the width the verb emits it, read from the access's plan. A new test records a ragged copy, checks that the coverage witness accounts for every emitted op, and checks that the recorded kernel is byte-identical to the plain build.
  • No new private cross-module imports in library code. The helpers the verbs share are public module names. GlobalForm is exported beside AccessPath as the type of the new plan.form, whose value is the text print(plan) shows. Test modules keep the white-box imports the repo's tests already use.

Gates on the new head: every commit builds and passes ctest and its own tests with both engines in strict mode. Byte identity is green at llvm20, llvm22, llvm23 and llvm24, and the demo IR snapshot is unchanged. The full platform suite has no new failure against the base except the known SSA-name test, which waits for #14846. pre-commit passes, and ruff reports nothing new.

cgmillette and others added 11 commits November 7, 2026 08:28
The tiling verbs bound-check a tile by its TileDesc shape but place each
register by its layout. A shape smaller than the extent the layout
addresses therefore let registers be read and written outside the region
the bound checks cover, with no error.

TileDesc now checks on construction that its shape equals, on every
axis, the product of that axis's hierarchical lengths, and refuses any
other shape (smaller, larger, or a different number of axes) with a
message naming both.

The new test pins the message, covers each kind of mismatch, checks that
swap_dims and reorder_registers still build matching descriptors, and
builds every tracked in-tree demo configuration so a demo relying on a
mismatched shape would fail.

Co-Authored-By: Claude <noreply@anthropic.com>
A tensor description could only take compile-time lengths and strides,
so a kernel whose sizes are known only at launch had no honest way to
describe its tensors.

A length or a stride may now be an i32 Value whose facts are stated: a
kernel parameter declared with them (b.param("M", I32, ge=1, le=4096)),
which KernelLauncher.from_artifact checks on the host before each
launch, or the Value b.gpu_checked_value returned, which the GPU checks. TensorDesc reads
the facts with stated_facts_of and keeps no copy. It refuses, each with
a message that names the fix: a Value with no stated facts (the raw
value a gpu_checked_value call checked included), a length that could
be negative, a length with no upper bound, a stride that could be below
1, a non-i32 Value, a Value that can differ between the lanes of a wave,
and a runtime stride where the contiguous axis must keep the int stride
1. A relation carries a lower bound forward, so ldc declared ge=N with
N ge=1 needs no word of its own. TensorDesc also gains an internal
per-axis begin field, all 0 for a tensor that is not a slice.

The wave-uniform rule of the IR range walk becomes one function,
wave_uniform, that needs no builder; the walk's splitter and TensorDesc
both use it.

check_gemm_shapes agrees runtime extents only when they are one and the
same Value; two Values that two gpu_checked_value calls returned for one
loaded value are refused with a message to check it once.

The readers of lengths and strides handle a runtime Value or refuse it
with a message: the verbs clip against a runtime length, multiply by a
runtime stride and stay scalar beside one; element_span, the built-in
LDS swizzle and the cooperative LDS fit check refuse runtime extents
(LDS is sized at build); the recorded-pipeline analyses that compute
concrete addresses refuse a runtime stride. The surface doc gains a
section on runtime extents.

Tests pin every message; that a window over a runtime-length,
runtime-stride tensor builds and both engines lower it at every LLVM
flavor; the check_gemm_shapes identity rule; that the parameter contract
accepts an aligned-shape kernel exactly for tile-multiple shapes and a
ragged one for every in-range shape; that the removed wrapper names are
gone from the package and its docs; and, on the GPU, that
first_supported picks between the two kernels, each copies exactly its
rows, and a launch outside the contract is refused before it runs.

Co-Authored-By: Claude <noreply@anthropic.com>
…rms addresses

A tensor description can carry a runtime stride and a per-axis begin
(a slice's start in its parent), but the tensor_transforms address
emitters refused a runtime stride and ignored the begin.

The stride fold now multiplies by a runtime stride at run time: a
uniform part times it is one i64 multiply in its input set's term, a
per-lane part times it one i32 multiply shared by the call, with a lane
range that follows the stride's stated range. A constant times a runtime
stride cannot be the int constant: the call's least constant on that
axis joins the uniform term once, and what each access has above it is
its uniform_delta, a wave-uniform value >= 0 (i32 when its range fits,
else i64). A stride whose stated range can reach 0 is refused at bind,
for a description built without TensorDesc's own check.

Each axis's begin is split by the IR range walk and folded like a
per-wave origin part: its constant times an int stride joins the
constant, and every other product joins one uniform term keyed by no
graph input, summed in tensor-axis order, i32 or i64 by its range, so a
slice can start past 2^31 elements; a begin of unknown range gives an
i64 term. A begin with a per-lane part is refused, naming its axis. The
checks stay those of the tensor's own lengths, and a tensor whose begins
are all 0 emits the same IR as before.

address_of and transform_address_of add the begins in Python integers
and take launch_values, (Value, int) pairs matched by identity, for the
int each runtime stride and begin takes at launch; a missing one is
refused by name. Interval.times gives the range of a product of two
ranges. The API doc describes runtime strides, slice begins, the new
signatures and the new errors.

Tests pin the multiply counts, the delta and lane ranges, the least
constant joining the uniform term, the i32 / i64 choice of the delta and
of a begin term, a vector along a runtime-stride axis refused as not
contiguous, every new error before any IR is emitted, the IR of a plain
tensor, and every address against the golden on the host. On the GPU,
every lane's parts sum to the golden for packed rows with a runtime K
(f32 and f16, K 37, 64 and 65), a BSHD batch through its runtime batch
stride, and a THD slice whose begin is loaded through readfirstlane,
including a begin term past 2^31.

Co-Authored-By: Claude <noreply@anthropic.com>
A TensorWindow now does one job: place a tile on a tensor. It keeps one
origin per tensor axis and no region; the valid region is the tensor's
own lengths, so a bound check reaches the very length Value the author
declared. A narrower region is a narrower tensor description.

- make_window(tensor, origin) takes no bounds argument (passing one
  raises Python's own keyword error). It checks the origin count and
  each origin's kind (an int or an i32 IR value), with actionable
  messages.
- TensorWindow.with_origin(axis, origin) replaces one axis's origin,
  refusing an axis the window does not have.
- A negative origin is accepted (a convolution halo). The verbs mask
  the elements before the tensor start: an axis whose origin is a
  negative int gets the lower compare and never takes the unmasked
  wide path, so those elements read 0 and are not written.
- The leading-axis bounds refusal and the LDS-window clip guard go
  with the field they checked.

The tests cover the new errors, with_origin, negative origins (IR and a
guard-zone GPU copy), and the wide-load gate rewritten on tensor
lengths, including runtime lengths. In-tree kernel IR is unchanged.

Co-Authored-By: Claude <noreply@anthropic.com>
A block origin such as block_id_y * TM is a plain i32 value: nothing in
the IR says the tile starts inside the tensor, and an over-launched grid
or an in-kernel grid remap can place it outside.

exit_if_tile_outside(b, origin, length) makes 0 <= origin < length true.
At its call it emits the tensor_transforms bound check of the RAW origin
(one unsigned compare, so a negative origin fails it too), a uniform
branch to an early return when it fails, then a range marker stating the
fact, written by="exit_if_tile_outside", and returns the marker's
result. The IR range walk reads the fact from every use of that value,
snapped to the origin's multiple-of fact: with a length that is a
multiple of the tile the origin walks to the last tile start, so the
whole tile is inside. A runtime length is the marker's upper operand,
so the walk's end is the very Value the tensor's length is.

The range marker's writer list gains exit_if_tile_outside in both
engines (EXIT_IF_TILE_OUTSIDE_WRITER beside GPU_CHECKED_VALUE_WRITER),
so the verifiers accept the marker and the parameter contract keeps
listing only gpu_checked_value markers as GPU checks.

IRBuilder gains a read-only at_kernel_top_level property, true while the
builder emits into the kernel body itself and not inside a loop or an
if; the guard reads it to keep its early exit in the prologue.

It refuses, each with an actionable message, a call inside a loop or an
if, an origin with a per-lane part, an origin or a length that can
differ between the waves of a workgroup (a guard on one would end some
waves and not others), an int origin, and a length that is not a
positive int or an i32 Value whose stated facts make it a length (the
raw value a gpu_checked_value call checked is refused with the Value to
pass instead). The walk's own errors (a non-i32 origin, a missing
workgroup size) pass through. The length rules are the ones a tensor
description applies, shared by both.

The tests pin the emitted guard and its op order, identical IR from both
engines at every LLVM flavor, a scalar compare and a branch to the end
of the program in the gfx90a ISA, the facts the walk reads through the
returned value for the origins kernels write (raw products and shifts, a
grouped grid remap, parameters, read-only loads, loop counters, the
tail-tile origin, wave offsets added to the guarded value, a batch
index), every refusal, the builder property at the top level, inside
loops, ifs and both if/else branches and after leaving them, that two
guards of one value used as lengths get check_gemm_shapes' layout advice
and not its gpu_checked_value advice, and on the GPU that an
over-launched grid touches nothing outside the tensor and a negative
remapped origin does nothing.
The surface doc gains the guard and its contract.

Co-Authored-By: Claude <noreply@anthropic.com>
Which sides of a global window access needed a run-time check was
decided in three places -- the clip mask, the wide-load gate and the
origin-alignment test -- each differently, and each trusting what it
could not see: a runtime origin was taken as grid-aligned, a
tile-aligned length as a proof, the upper end of a leading axis as
inside. A tile could overhang its tensor unmasked.

Now one decision, made by the tensor_transforms bound-check rule:

- clipping.py (new): explain_bounds_checks(b, ptr, window, tile_desc)
  returns a BoundsCheckPlan (SideCheck, VectorValidity, AccessPath,
  AxisChecks, DomainChecks). Per tensor axis the positions are the
  origin's range as the IR range walk reads it plus the tile's extent;
  required_checks over the tensor's identity graph against [0, length)
  says which sides can fail. A failing side is one scalar compare when
  the position is the same for every lane, else a per-lane compare. The
  plan also gives the vector width (the unchanged width rule, moved
  here), the vector axis as a window axis and how a vector's validity is
  decided. It emits nothing; print(plan) gives one line per side. The
  verbs take the same decision from the same function. When a side is
  checked only because a loop stops at a runtime length its facts do
  not make a multiple of the loop step, its printed line ends with the
  fix, e.g. "-- state multiple_of=16 on K's b.param declaration" (or
  the b.gpu_checked_value call to state it in): the IR range walk
  records such a stop on the loop counter's part of the split, and the
  plan keeps it when the bound-check rule proves the side with the
  range that fact would give. An LDS buffer is
  refused: LDS accesses are not bound-checked. A pointer parameter
  declared addr_space="constant" is refused at build, by the plan and
  the verbs alike, for every form but the one-index form: those forms
  do arithmetic on the 64-bit pointer, which core's global_ptr ops take
  only as a global pointer. The one-index form stays exempt; it relies
  on core's global load and store lowering to handle such a pointer.
- safe_access.py (new, internal): delivers every global access that is
  not a proven window with int strides. Its address and validity come
  from emit_accesses (with the vector axis and width); the per-wave part
  goes on the 64-bit pointer, the per-lane part is a zero-extended
  32-bit byte offset, the constant a trailing index. A checked access
  runs under its validity: a wholly valid vector is one wide access, a
  vector that may cross a bound has each valid element accessed alone.
  No instruction touches an element outside the tensor, no address an
  invalid lane computes carries an in-bounds promise and no load is
  clamped to element 0. These forms claim the element's alignment.
  The plan names the form the verbs deliver the access in (GlobalForm,
  plan.form): one 32-bit element index, a placed address, or a placed
  address under the access's validity.
- emit.py: a proven window with int strides and offsets inside i32
  keeps the one-index form, without the per-register positions only the
  deleted clip read. The clip mask, the in-bounds gate of the width rule
  and the trusted origins are deleted. load_fragment refuses a non-zero
  pad and both verbs refuse a coherency other than 0, before emitting
  anything and on every path: neither is implemented, so the value would
  be ignored (an IR value given for either is a TypeError).
- The verbs and the plan share their helpers through public module
  names, never another module's private ones: clipping's plan_window,
  contiguous_run, leading_rank, check_tile_rank, is_lds,
  global_store_width and element_bytes; encoding's lane_terms and
  register_coord_offsets; tensor_transforms' required_checks_within (the
  shared rule over Interval ranges) and Interval.numeric_ends. An
  element's width comes from core's dtype_info (element_bytes, which
  refuses a dtype that is not whole bytes; the layout renderers call it
  too, in place of their own wrapper), and the dtypes the verbs access
  as vectors have one home.
- The pipeline recorder records each global access at the width the
  verb emits it, read from the access's plan: a checked or placed store
  is one wide store per run, a store on the one-index form one store
  per register, so its coverage witness accounts for every emitted
  store.
- The manual, interleaved (LDS-staged included) and CRC demos guard
  their block origins with exit_if_tile_outside, so every global window
  is proven and keeps the one-index form. Their size parameters state
  the sizes the kernel bakes (eq=), and their run functions launch
  through KernelLauncher.from_artifact, which reads the argument
  signature from the kernel and checks those facts before every launch;
  the hand-written signatures are gone. A demo whose C store adds a wave
  offset records the workgroup size the walk read.
- The IR range walk snaps a loop counter whose stop is a runtime length
  plus constants to the counter's multiple-of fact, as it already did
  for an exact stop: the loop [0, K) step TK over a K stated a multiple
  of TK ends at its last trip, K - TK, so the aligned kernel's natural
  K loop proves its K tiles whole; Interval.snapped takes the supposed
  residue of an upper end for the range a fact not yet stated would
  give.
- explain_bounds_checks and its values (GlobalForm included) are
  exported; the surface and
  contract docs state the rule, and the surface doc shows the aligned
  and the ragged K loop side by side.

The demo IR snapshot is re-blessed: the demos with the reason "block
origins guarded in the prologue, and the size parameters state the
sizes the kernel bakes" (and the dead position adds gone; the K-loop
ISA is unchanged up to register numbers), the bank-conflict
probes because their band window, checked before and after, is now
masked at full width instead of element by element.

Tests: the plan equals the shared rule on walked positions, the walk's
facts through the plan (guarded products and shifts, full-tile loop
counters, the loop up to a K stated a tile multiple, the tail-tile
origin, int lengths), leading axes on both sides, scalar checks, the
vector axis of a rank-3 window, the plan the verb applies, the print,
the fix it names (by parameter or GPU-checked length, keeping a stated
multiple) and none when K is stated a tile multiple, the side is checked
for another reason or the length's facts are not the author's, the
walk's record of such a loop stop, and the refusals (a constant-space
pointer refused off the one-index form, and passed on unchanged on it); the IR and ISA of the forms (no in-bounds promise or clamp
on a checked access, the pointer-plus-offset form with no compare for a
runtime stride, the per-lane byte limit, the row compare of an M-ragged
GEMM outside its K loop on gfx90a and gfx942); on the GPU, ragged and
misplaced windows (f16/f32, M- and K-ragged, runtime length 0, negative
origin, odd stride, a tile past the tensor, runtime origins before and
past it, a batch index outside the batch, a convolution halo against
im2col, a wide batched load, the CRC demo at 256x128x64) read and write
exactly the valid elements with every guard zone intact; the demos'
global accesses unmasked, their workgroup-size record, and a parameter
contract that accepts exactly the launch of each demo's shape; the pad
and coherency refusals on the proven, checked and LDS paths, with
nothing emitted; element_bytes over byte and sub-byte dtypes; the store
width of every form; a recorded ragged copy whose wide checked load and
store the coverage witness accounts for, built byte-identical. The tests
that pinned the trusted behaviour now pin the rule.

Co-Authored-By: Claude <noreply@anthropic.com>
A wide global access told the compiler it was aligned to its own
width (align = width * element size) whatever the facts said. An origin
the run time picks, a slice's begin offset or a pointer declared less
aligned can make that claim false, and a false claim is a wrong answer:
a uniform address claimed 4-byte aligned or more may be read by a
scalar load, which ignores the low address bits.

- safe_access.py: access_alignment decides the claim from facts only:
  the largest power of two dividing the pointer's alignment (the align=
  of its kernel parameter, else the element size) and every access's
  element offset sum((position + begin) * stride), at most the access's
  own byte width and 16 bytes. The residues of the parts are combined
  exactly: an int by its value, a value by the multiple-of fact the IR
  range walk proves, a runtime stride by its stated multiple_of; a
  position is the origin plus any lane's and any run's offset. The
  global loads and stores of every form take this claim; an element of
  a vector crossing a bound claims at most its own size.
- emit.py: the verbs compute the claim once per access call and pass it
  to the one-index form and to safe_access.
- clipping.py: the width rule is unchanged (the width divides every
  other stride, else 1), and now reads a runtime stride's stated
  multiple_of instead of keeping the access scalar next to it, through
  tensor_transforms' end_residue, which becomes public for it.
- docs: tiling_api_surface.md gains "Pointer alignment" (declare align=
  on pointer parameters; what a claim is built from) and its width and
  defaults lines; tiling_api_contract.md states that pointer alignment
  is read from the pointer, never assumed.

Tests (test_access_alignment.py): the claim equals the true alignment
over stride classes K + {0, 1, 2, 4, 8}, int origins and begin offsets
(brute force), and runtime strides, origins and begin offsets over
every value their facts allow, for three pointer alignments; the stated
examples (odd stride scalar at 2 bytes, runtime lda multiple_of 1 / 8
with align 16 / 4, a pointer with no align= or computed from a
parameter, an odd begin, a checked window's edge elements); a uniform
row at an odd runtime column is read by a vector-memory load while the
same row at a multiple of 8 is a scalar load; on the GPU present, the
manual and interleaved GEMMs are bit-exact at every padded stride with
NaN operand pads and the C pad untouched, and an 8-wide load at an odd
runtime column copies exactly. The demo IR snapshot is unchanged.

Co-Authored-By: Claude <noreply@anthropic.com>
rocKE allocates LDS statically, with a fixed size, so the tensor
descriptor that is an LDS buffer's layout needs int lengths and strides.
A runtime length or stride on an LDS tensor is now refused by one check
on TensorDesc, with one message naming the first runtime extent, before
anything reads the layout:

  load_fragment: lengths[0] is a runtime value (S) -- rocKE allocates
  LDS statically, with a fixed size, so an LDS tensor needs int lengths
  and strides

Every reader that lays a tensor out in LDS calls it and opens the
message with its own name: load_fragment and store_fragment on an LDS
buffer, the cooperative LDS copies, the built-in swizzle and
TensorDesc.element_span. Before, the verbs failed later, inside the span
or swizzle computation, and the copies and element_span used older
wording. A runtime window origin stays accepted.

Tests: both verbs, with and without the built-in swizzle, the
cooperative copies and element_span refuse a runtime length and a
runtime stride with that message; a window at a runtime stage origin of
a static buffer builds; the runtime-extent tests pin the new text. The
tiling API surface doc states the rule.

Co-Authored-By: Claude <noreply@anthropic.com>
…not fit i32

A proven global access placed on the 64-bit pointer handed its
constant element offset to the instruction as an i32 index. A window
whose int origin lies 2^31 or more elements into its tensor (a late row
of a 65536 x 65536 tensor) emitted an index outside i32, which LLVM
reads wrapped by 2^32 elements: the access read and wrote the wrong
address with no error.

Now a constant outside i32 joins the per-wave part on the 64-bit
pointer, as one 64-bit byte offset added before the per-lane offset,
and the instruction's index is 0. A constant inside i32 is placed as
before, so no in-tree kernel's IR changes. Whether an element offset
fits i32 has one home in the tiling verbs, clipping.fits_i32 over core's
int_type_limits, shared by the one-index form's choice and the
placement; the instruction index asserts it. The surface doc says so.

Tests: every kind of window of a 65536 x 65536 tensor (guarded block
origins, int origins near and at the end, past and before the tensor,
unguarded origins) builds for load and store on the global path and
emits no i32 literal outside i32; the one-index form is taken for the
tile ending at offset 2^31 - 1 and not one row further; a tile starting
at offset 2^31 is placed with a 2^32-byte pointer offset and index 0;
on the GPU, a row 2^32 elements into an i8 tensor is read and written
where it lies and the tensor's first row is untouched.

Co-Authored-By: Claude <noreply@anthropic.com>
A tensor description whose strides map two positions to one element (a
stride 0, rows longer than their stride, a sliding-window view) is fine
to read, but a store through it has two lanes writing one element and
the result depends on which lands last. store_fragment wrote through
such views without a word.

store_fragment now refuses, at build and before it emits anything, a
global or LDS store through a view whose overlap is proven from the
description, and names two positions that reach one element:

  store_fragment: positions (1, 0) and (0, 1) of the tensor reach the
  same element -- a store through an aliasing view races; store through
  a view whose positions are all different elements

A view the build can decide neither way -- a runtime row stride over a
runtime row length with nothing stated between them -- is trusted, as
the caller describes it. Loads are unchanged. On an LDS buffer the
static-extent rule runs first, as it does on a load.

Two internal questions on TensorDesc decide it. _strides_injective
proves that no two positions share an element: some order of the axes
must exist in which each stride is larger than the farthest element the
axes before it reach together, each step proven from ints and stated
facts, from packed rows (the stride is the very Value of the previous
axis's length over stride 1), or from a stride gpu_checked_value
returned for the previous length times its stride, accepted only when
that product cannot wrap in i32. _aliasing_positions finds two positions
that meet at every launch the facts allow: an int stride 0, two int
strides whose smallest common step lies inside both axes' guaranteed
lengths (a runtime length holds its stated lower bound), or two axes on
one runtime stride. A store is refused only when the first fails and the
second finds a pair. The surface doc states the rule.

Tests: packed, padded, permuted, reversed and length-1 int views, packed
runtime rows, a BSHD tensor with a runtime sequence length (a batch
stride declared large enough, and one checked as the product of the
sequence length), a batch stride over runtime rows and one THD sequence
of a packed token buffer are proven not to alias; a stride 0, a sliding
window, rows longer than their stride, reversed rows, a runtime row
length stated long enough and two axes on one runtime stride are proven
to; a three-axis overlap, a runtime row stride over a runtime row
length, a product that could wrap and a product over another axis's
length are decided neither way; a global and an LDS store through a
proven overlap are refused with the message and emit nothing; a store
through a runtime-ldc padded C view builds and lowers in both engines; a
load through an aliasing view lowers identically in both engines; on the
GPU a load through a sliding-window view reads exactly numpy's view of
the same bytes. One LDS test that stored through rows longer than their
padded stride now describes the padded rows it meant.

Co-Authored-By: Claude <noreply@anthropic.com>
…eters

A kernel parameter can state a relation to earlier ones (ldc ge N, sb
ge product_of(M, ldc)), and KernelLauncher.from_artifact checks it
before each launch, but the tiling layer's build-time proofs read only
numbers: a padded C stated ldc >= N was still decided neither way, and
the 32-bit index proof ignored a relation that bounds a window's
largest offset.

The new internal module stated_relations reads every relation stated on
the parameters a tensor description uses, and those declared before
them, in one shape: big >= small, both integer polynomials over
parameter names. x ge t is x >= t; x le t is t >= x; gt and lt add one
to the smaller side and eq gives both directions. A quotient_of on the
bigger side moves to a product (x le quotient_of(a, y) reads a >= x *
y, since every divisor is stated at least 1) and a difference_of to a
sum (x le difference_of(a, y) reads a >= x + y), so one fact reads the
same in whichever shape the declaration order allows. A relation with a
quotient left on its smaller side gives no exact product and is left
out. proving_relation finds the relation that shows big >= small when
its smaller side exceeds small only by terms that cannot be negative,
and keeps the text the author stated it with.

The aliasing proof clears an axis whose stride is a parameter stated at
least the previous axis's length times its stride: ldc ge N clears
padded rows, sb ge product_of(M, ldc) (or ldc le quotient_of(sb, M))
clears a batch of them. TensorDesc._injective_proof returns the
relations the proof used.

The 32-bit index proof of a proven window with int strides, when the
numbers alone do not keep its offsets inside i32, bounds the largest
offset by a stated relation with a constant bigger side (N le
difference_of(2147483647, product_of(M, 4096)) over rows of 4096), so
the window takes the one-index form. BoundsCheckPlan gains
address_facts, the relations that proof used -- only when it used them
-- and str(plan) ends with how an element is addressed and those
relations. The surface doc states both uses.

Tests: the shape of every way a relation can be stated and the shape
left out; padded and batched C views proven not to alias in every ABI
order, naming exactly the relations used, and decided neither way
without them (the store still builds); a window over runtime lengths
that takes the one-index form only through its relation, names it, and
names nothing when the numbers suffice; both engines lower it at every
LLVM flavor, and on the GPU a copy through it, launched through
KernelLauncher.from_artifact, writes exactly its elements.

Co-Authored-By: Claude <noreply@anthropic.com>
@cgmillette
cgmillette force-pushed the users/cgmillette/rocke/ragged-clipping branch from e077c8d to 10f92e1 Compare November 7, 2026 13:51
@cgmillette

Copy link
Copy Markdown
Contributor Author

Rebased onto develop (21e3ca4350d) now that #14448 has merged. The branch carries only this PR's own eleven commits; the old copy of #14448 underneath is gone, and the description no longer says it depends on #14448. The commits are unchanged apart from context lines in one test file's imports.

Checks on the new head: byte identity green at llvm20, llvm22, llvm23 and llvm24; ctest passes; the tiling tests pass with both engines in strict mode except the one test that waits for #14846 (as before); the installed-layout tensor_transforms CTest entry passes; the full platform suite shows the same failures as develop and none new; pre-commit passes. The GPU cases are still to run on this head. It stays a draft.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant