Module Ir.Schedule

For_loop i in [0, N)] becomes [i_o in [0, ceil(N/factor)) { i_i in [0, factor) }] with [i := factor*i_o + i_i] substituted throughout the body (index vectors and [Embed_index]), and — when [factor] does not divide [N] — the body wrapped in an [If (factor*i_o + i_i < N)] remainder guard (construct-then-fold: {!apply}'s trailing simplify erases it whenever the loop extents prove it). The split loop must start at 0, which lowering guarantees. Splitting a serial loop preserves iteration order; retyping the results to hardware axes carries the iteration-independence obligation exactly as for [Low_level.validate_parallel].

Interchange two perfectly nested loops (the outer loop's body must be exactly the inner loop; fails loudly otherwise). Reorders iterations — legal for the associative-commutative accumulation patterns lowering emits; bitwise reproducibility is the caller's concern.

Change a loop's axis type in place. Retyping to a hardware kind requires from_ = 0 and iteration independence (the caller's obligation; structure is checked downstream by Low_level.validate_parallel). Retyping to Vectorized likewise asserts iteration independence: the C backends render the loop under vectorization pragmas (gh-ocannl-164), backends without pragmas render it as a plain serial loop.

materialize = false: set the axis type to Unrolled — codegen repeats the body with the index bound as a per-block constant (after simplify/CSE have run, so the copies are opaque to the optimizer). materialize = true: unroll in the IR by substituting index constants, so that apply's trailing simplify + CSE see the copies — constant-folding Affine indices and deduplicating repeated loads. Register blocktiling is Split + materializing Unroll + the existing CSE (schedule-ir-optops §4).

Index-set splitting (gh-ocannl-508): replace the Serial loop For_loop axis in [from_, to_] by consecutive segment loops s_0 in [from_, b_1) ; s_1 in [b_1, b_2) ; ... ; s_m in [b_m, to_]], each binding a fresh symbol substituted for [axis] in a copy of the body (scalar-local scope ids refreshed per copy, as for materializing [Unroll]). Segment ranges stay absolute — no index arithmetic changes, iteration order is preserved exactly — and each segment's narrowed range lets {!apply}'s trailing simplify interval-fold the guards it decides (statement [If]s and scalar [Where] range guards alike). This is the segmented-rendering replacement for in-loop range guards: an inlined concatenation's per-component guards (partition the consumer at the component boundaries, e.g. from {!partition_breakpoints}) and [Split]'s construct-then-fold remainder guard (partition at the last tile-multiple first, then [Split] the dividing main segment — clean main nest plus epilogue) both specialize into guard-free segment nests, and the fresh symbols make each segment individually addressable by subsequent ops (per-segment scheduling).

Pad-to-tile (gh-ocannl-485, PADTO): extend the Serial loop For_loop axis in [0, N)] to [[0, M)] where [M] is the least multiple of [to_multiple_of] with [M >= N], and guard every effectful leaf statement of the body with [If (axis < N)] (leaf statements, not the whole body, so barriers inserted by a later shared {!Stage} stay under uniform control flow). The pad iterations are no-ops, so the op is unconditionally semantics-preserving ([op_legality]: [Op_legal]); when [to_multiple_of] already divides the extent it is the identity. Purpose: downstream [Split]s by factors of [M] divide cleanly (no remainder guard), {!Stage} tiles minted over the padded tile loops zero-fill their fringe (see {!constructor-Stage}), and {!constructor-Tensorize} recognizes the guards as pad masks — moving row/column masks to the accumulator transfers and discharging reduction masks against zero-filled staged operands — so tensorized paths cover arbitrary extents. The surviving guards are exactly the flip points {!partition_breakpoints} detects: a later {!constructor-Partition} of an enclosing block loop at the last fully valid block specializes them away in the interior segments.

Stage reads of source through a tile: a fresh Local-mode node registered in the traced store, its dims derived per source axis from the range of the index terms over tile_loops (schedule-ir-optops §5). All reads of source must use one index vector (v1) whose per-axis terms split cleanly into tile-loop terms (positive coefficients) and outer terms. With shared = true the tile is added to workgroup_shared and a cooperative-load nest plus barriers are inserted at the deepest loop carrying an outer-part symbol or a reused Workgroup tile axis: Workgroup-typed tile loops are reused as the cooperating thread indices, Serial tile loops are iterated under fresh symbols, per-axis edge guards are constructed then folded — in Where form storing 0 (the add-reduce accumulation identity) to out-of-range slots, so edge tiles of a non-dividing or padded staging are safe to read over their whole index space; every minted tile is recorded in Low_level.optimized.zero_fringe accordingly (gh-ocannl-485) — and redundant loading along non-participating workgroup axes is restricted to one representative thread. A shared stage with no anchor (no outer-part symbol, no reused workgroup axis — e.g. staging a broadcast vector) wraps the outermost tile loop instead of the routine root, so enclosing workgroup axes can guard the loads; every workgroup slot active in the kernel's launch must be reused or bound by a loop enclosing the staging point, otherwise threads differing in the uncovered slot would race on the shared tile and the op raises. Note that Split's whole-body remainder guards would place the inserted barriers under divergent control flow (rejected by validate_parallel), so v1 shared staging requires tile sizes dividing the extents. With shared = false (CPU operand packing) all tile loops must be Serial and a plain serial copy nest is inserted, no barriers. The source must not be written in the routine.

cooperative = Some w is the lane-aware mode for composing with Tensorize (docs/proposals/tensorize-mma.md, "Lane-aware Stage"): shared staging with all-Serial tile loops whose load nest is wrapped in a fresh extent-w Workgroup lane loop — positionally the same slot 0 the tensorized micro-kernel's lane loop binds, with extent agreement enforced downstream by Low_level.validate_parallel's barrier-strength uniformity (pass the backend's Backend_intf.mma_capability.mma_simd_width to both ops). The lane is folded linearly into the innermost fresh copy loop: extent <= w replaces the loop by the lane index under a lane < extent guard (folds when equal); an extent divisible by w iterates extent / w chunks at w*step + lane; otherwise the whole nest is restricted to lane 0 (division is not expressible in affine indices). The lane loop covers workgroup slot 0 for the staging-point coverage rule by construction.

swizzle = Some kind stores the tile in an XOR-swizzled layout (the bank-conflict-avoidance follow-up of docs/proposals/tensorize-mma.md): the tile node is recorded in the optimized record's Low_level.optimized.swizzled map under kind and codegen remaps every element access P*C + col to a per-row bijection of the minor axis, so semantics are unchanged while same-column accesses (the classic strided read of a staged tile) spread across shared-memory banks. Requires shared = true and a tile with at least two axes.

The two flavors differ in the unit the XOR permutes — hence in their minor-extent requirement, neither of which implies the other — and in who can consume them (gh-ocannl-481 item 3):

• constructor-Low_level.Swizzle_elem remaps P*C + col to P*C + (col lxor (P land (C-1))), requiring a power-of-two minor tile dim > 1. Renderings that assume row-major storage decline it: Tile_mma intrinsic/register-tiled paths fall back to the scalar micro-kernel (which reads elementwise and stays correct), so do not combine it with Tensorize when the intrinsics are the goal — this flavor is for scalar/register-blocktiled staged kernels.
• constructor-Low_level.Swizzle_b128 permutes whole 16-byte units, requiring the minor tile dim to span a power-of-two count > 1 of 16-byte units. This is the layout the CUDA inline-PTX mma.sync arms' ldmatrix loads consume, so it DOES combine with Tensorize: the tile is both bank-de-conflicted and fragment-loadable in one instruction. Backends without such loads still decline it to the scalar micro-kernel, which stays correct.

pad_stride = Some p rounds the tile's MINOR dim up to a multiple of p (gh-ocannl-481 item 4). The tile's leading-dimension stride is that dim — every consumer reads it off the node — so this changes the stride while the iterated index space stays the unpadded extents. Two payoffs, both about the stride rather than the data: shared memory bank conflicts on a strided read of the tile (proposal §5's "correct, possibly bank-conflicted" v1), and layout rules stated on the stride — a fragment load's ld-multiple constraint, and Swizzle_b128's 16-byte-unit count, which is why the two compose as "pad first, then check". Requires a tile with at least two axes and p > 1.

The padded slots hold nothing under a row-major layout: no loop reaches them, so they are neither written nor read, and the Low_level.optimized.zero_fringe contract — about the fringe of the staged source region within the iterated space — is unaffected. Under a swizzle they do carry data, the XOR being a bijection of the whole padded row, and reads go through the same map.

hoisted = true packs a compile-time-constant operand once, out of the routine (gh-ocannl-470, the compiler-native analog of ggml's CPU_REPACK set_tensor hook): instead of a per-invocation scratch tile refilled by an in-kernel load nest, the packed layout covers the whole source — one packed-buffer axis per outer coordinate (outer part / tile dim, requiring outer-part coefficients and offsets divisible by the tile dim on tiled axes) followed by the tile axes — and the reads are remapped to it directly; no load nest, no barriers. The packed node is minted as a host-initialized constant: its buffer is computed on the host when first forced (a Host_inits lazy driven by the same affine index maps; pad slots of edge tiles are zero-filled) and uploaded once per device into the constant pool (constant_buffer_cache), so re-linking into sibling contexts reuses the same packed buffer. Requires shared = false (and hence all-Serial tile loops), a source with registered host-init data that is known constant (declared Effectively_constant intent or a constant placement), no padding on the source, and every outer-part symbol bound by an enclosing loop (static/dynamic indices are rejected — packing runs at link time with no bindings). Note: later host-side writes to the source (e.g. set_values) do NOT refresh the packed copy.

Accumulator privatization: contract the read-modify-write accumulation of the materialized target across the (Serial) over loop's whole subtree into a per-thread Local accumulator tile — initialized from target before the loop, accumulated in place, stored back after, one final write per element. This recovers, for materialized nodes, the scope-local form virtualization gives virtual accumulators (Local_scope) — and because a routine-local tile cannot alias the kernel's device pointers, downstream compilers register-allocate it without waiting for restrict (gh-ocannl-164). Tile shape: per target axis, the index terms over loops nested inside over (required Serial with from_ = 0); no such terms yields a scalar accumulator. All accesses of target under over must use one index vector, which must not mention over itself, and must sit under one If guard chain, which must be iteration-invariant (no memory reads, no symbols bound inside over's subtree) — the same predicate then gates the init-load and store-back, so lane-restricted accumulations (e.g. w == 0) privatize correctly; per-iteration guards are rejected. A Zero_out of target elsewhere is left in place — the init-load observes it, so semantics are preserved without a surjectivity analysis. Compose as: Splits → Stages → Privatize → materializing Unrolls (the unrolls then turn the tile accesses into constant-indexed, register-allocatable form).

Expand the unique Zero_out tn statement into an ordinary loop nest over the supplied symbols (one per axis of tn's padded dims; see expand_zero). Whole-node Zero_out of a materialized node is rejected by Low_level.validate_parallel in multi-threaded kernels — expanding it first lets the schedule split and annotate the zeroing with the same hardware geometry as the computation that follows.

Tensor-core emission (docs/proposals/tensorize-mma.md §3): replace the perfectly nested serial i × j × k matmul micro-kernel — whose body is the single accumulation d[..., i, j] += a[..., i, k] * b[..., k, j] (plain-add or FMA form; each operand's tile spans its last two axes with unit coefficients, transposed layouts rejected in v1) — with a Low_level.Tile_mma block statement covering the full m×n×k extents, wrapped in a fresh Workgroup-typed lane loop of extent simd_width. The original nest is kept as the statement's scalar fallback: backends without an MMA hook (or declining a particular precision/shape) render it once per simdgroup under an if (lane == 0) guard, so the op is always semantics-preserving. Apply after Splits and Stages; Stage/Privatize must come before it. Divisibility by the intrinsic tile (8 on Metal) is a per-call emission concern, not checked here.

Epilogue fusion (gh-ocannl-486): fold the sole-consumer, index-space-compatible elementwise tail that re-reads target — the typical bias add / activation / residual after a reduction — into target's store-back site, eliminating the tail's separate memory pass and keeping the fused routine a single kernel/segment. Recognized sites: the lane-0 fragment store-back synthesized by Tensorize's accumulator contraction (the tail becomes a fourth, lane-0-guarded statement of the marked region, rendered after the backend's intrinsic block; a pad-masked store-back — gh-ocannl-485's range guards on a non-dividing site — is recognized through its guards, which are re-imposed on the relocated tail, gh-ocannl-521), the whole-K Tile_mma writing target directly (the unstaged tensorized pipelines: the intrinsic block completes the accumulator's m x n tile, so the tail becomes a sibling lane-0 nest over the tile at the accumulator's base indices, gh-ocannl-521), the Privatize tile store-back (per-element), and the plain accumulation nest (the tail slides inside the output loops after the serial reduction loop). The tail must be the first real statement after the last statement writing target: a perfect Serial nest over exactly target's dims, assigning a different node at the identity index tuple, elementwise, with every read of target at that tuple, and no later statement mentioning target. The store-back of target itself is kept. Elementwise tails never reorder the reduction, so on the C backends the fused values are bitwise equal to the two-kernel form. Apply after Tensorize/Privatize.

shared (GPU only, requires the fragment site): the fused tail is often target's last consumer, so placement makes target routine-local — a per-thread array the fragment hooks cannot simdgroup_load from. With shared, target is placed in workgroup-shared memory (like Stage's shared tiles) unless already settled on-device, so the intrinsic fragment path fires against threadgroup memory. CPU backends reject shared placement.

Deterministic two-pass split reduction (gh-ocannl-484): parallel reductions without atomics. The Serial loop axis in [0, N)] — carrying the single accumulation of [target] in its subtree — becomes [block in [0, num_blocks) { inner in [0, ceil(N/num_blocks)) }] with [axis := ceil(N/num_blocks)*block + inner] substituted and a construct-then-fold remainder guard as for {!constructor-Split}; the accumulation is redirected into a fresh scratch node [partials] of dims [num_blocks x target-dims] (tile namespace, placed [On_device], registered in the traced store) at the original cell prefixed by [block]; and a synthesized combine statement, inserted right after the enclosing top-level statement, folds the partials into [target] in a {e fixed balanced-tree order} over [combine_indices]-bound loops: [target[c..] := target[c..] ⊕ tree(partials[0..B-1][c..])]. The block loop is freely annotatable afterwards ([Retype]/the default presets): its index pins the partials row, so parallelizing it is race-free by construction, and the [partials] producer/consumer pair is exactly the materialized cross-nest edge kernel fission cuts at — under [fission_scheduled]/[maybe_default_schedules] the two passes (and the scatter form's zeroing) compile as separate kernels with the event chain supplying the grid-wide synchronization the combine needs. Never annotate the block loop to a hardware axis while compiling both passes into one kernel. Recognized accumulation forms (the [axis] subtree must contain no other access of [target], and the rest of the statement must not touch [target] — its combined value exists only after the combine statement): - Static: a single rmw [Set], [target[idcs] := target[idcs] ⊕ e] with [⊕ ∈ {Add, Max, Min, Mul}] (FMA counts as [Add]), [idcs] free of [axis]. Every loop enclosing [axis] within the statement must pin exactly one component of [idcs] (injectively: at most one symbol per component), so distinct enclosing iterations use distinct partial cells and the combine re-iterates exactly the written cells. Each block initializes its partial cell to the accumulation identity in-nest (no separate zeroing pass). - Dynamic (the gh-466 embedding-backward scatter): a single [Set_dynamic] add-accumulation of its own row (the [Get_dynamic] rmw form built by [rewrite_one_hot_reductions]), possibly under guards. The scatter is redirected to [partials] with the block index prepended (dynamic axis shifted by one): within a block colliding rows stay serial, across blocks rows land in disjoint partials slices — the block loop parallelizes what the scatter alone cannot. Rows are data-dependent, so [partials] is zeroed by a preceding whole-node [Zero_out] statement (its own fission segment) and the combine covers all of [target]. Within each chunk the original serial order and rounding are preserved; across chunks the reduction is reassociated (the same license as {!constructor-Swap} of accumulations). The result is deterministic {e per schedule} — bitwise-reproducible run to run and across serial/parallel renderings of the same schedule — but not bitwise-equal to the unsplit serial reduction: schedule identity pins numerics (retuning may change [num_blocks] and hence the combine tree).