Repository navigation
fix(rocke): decide every tiling bound check with one rule, so ragged tiles never read or write outside a tensor - #15026
cgmillette wants to merge 11 commits into
Conversation
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
|
GPU check at head 7b60d53 on gfx90a (MI210): the whole tiling test tree ( |
7b60d53 to
0882e32
Compare
|
Applied four decisions from today, each folded into the commit that owns the code (still eleven commits on #14448):
|
0882e32 to
e077c8d
Compare
|
Rebased onto #14448's new head ( Rebase: this PR now reads stated facts and range markers only through #14448's homes ( Quality fixes:
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. |
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>
e077c8d to
10f92e1
Compare
|
Rebased onto develop ( 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. |
Motivation
The tiling verbs
load_fragment/store_fragmentcould 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: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:
fix(rocke): require TileDesc shape to equal its layout extentTileDescwhose 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.feat(rocke): TensorDesc lengths and strides take parameter Valuesb.paramdeclared with them (checked on the host before each launch throughKernelLauncher.from_artifact) or ab.gpu_checked_valueresult (checked on the GPU).TensorDescreads 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.check_gemm_shapesagrees 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.feat(rocke): fold runtime strides and slice offsets in tensor_transforms addressesfeat(rocke): window positions with no length of their ownmake_window(tensor, origin)keeps one origin per tensor axis and no region.bounds=is gone, andwith_originmoves one axis.feat(rocke): guard block origins in the kernel prologueexit_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.EXIT_IF_TILE_OUTSIDE_WRITERin Python,ROCKE_EXIT_IF_TILE_OUTSIDE_WRITERin 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 anif. 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.fix(rocke): decide clipping with the shared bound-check ruleexplain_bounds_checks(b, ptr, window, tile_desc)returns aBoundsCheckPlanand 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.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, theb.gpu_checked_valuecall 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.safe_accessmodule 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.addr_space="constant"is refused at build, by the verbs andexplain_bounds_checksalike, on every form except the one-index form. Those forms do arithmetic on the 64-bit pointer, and core'sglobal_ptr_*ops take only a global pointer. The refusal is aTypeError, as core's refusal of a pointer in the wrong space is.KernelLauncher.from_artifact. Every global window in them is proven and keeps the one-index form.plan.form(GlobalForm.INDEX32/PLACED/MASKED, exported besideAccessPath); its value is the textprint(plan)shows.load_fragmentrefuses a non-zeropad, and both verbs refuse acoherencyother 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 aTypeError.safe_accessshare their helpers through public module names (not the package door), never another module's private names:clipping'splan_window,contiguous_run,element_bytes,global_store_widthand others;encoding'slane_terms/register_coord_offsets;tensor_transforms'required_checks_within(renamed from_required_checks) andInterval.numeric_ends.dtype_info(element_bytes, which refuses a dtype that is not whole bytes); the width table this PR used to move out ofemit.pyis gone, and the dtypes the verbs vectorize have one home.global_store_widthfor stores): a checked or placed store is one wide store per run, so the coverage witness accounts for every emitted store.fix(rocke): claim only the alignment a wide access really hasalignan access claims is derived from facts only: the pointer'salign=, 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 throughtensor_transforms'end_residue, now public.fix(rocke): refuse runtime extents on LDS tensorsTensorDescrefuses 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_fragmentandstore_fragmenton an LDS buffer, the cooperative LDS copies, the built-in swizzle andTensorDesc.element_spanall use it.fix(rocke): never take the 32-bit-index path for an offset that does not fit i32fix(rocke): refuse stores through aliasing tensor viewsstore_fragmentrefuses, at build, a store through a view whose overlap is proven from the description, naming two positions that meet.feat(rocke): tiling proofs read relations stated between kernel parametersb.param(ldc ge N,sb ge product_of(M, ldc)) prove padded and batched C views free of aliasing.BoundsCheckPlan.address_factsnames 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 threetensor_transformsemit goldens and the demo IR snapshot, and updates one SSA-name pattern intest_exit_if_tile_outside.py.Test Plan
explain_bounds_checks,load_fragmentandstore_fragmenton the masked and placed forms. That refusal emits no IR, and the same pointer is not refused on the one-index form.IRBuilder.at_kernel_top_levelat the top level, inside loops,ifs and both branches of anif/else, nested, and after leaving them.padandcoherencyrefusals on the proven, checked and LDS paths, with exact type and text and nothing emitted;element_bytesover byte and sub-byte dtypes and an unknown name; the store widthglobal_store_widthgives on each form against the stores the verb emits.check_gemm_shapesgives its "check it once" advice only for twogpu_checked_valuecalls on one value, not for twoexit_if_tile_outsideguards.element_span.Test Result
On develop
21e3ca4350d:rocke_tensor_transforms_pytestCTest entry passes.Submission Checklist
🤖 Generated with Claude Code