Ir.ScheduleHalide-style schedules over the lowered, optimized IR: a list of optops applied as a pure Low_level.optimized -> Low_level.optimized pass at the ?lowered_transform seam of backend compile. See docs/schedules_and_autotuning.md for the system view (composition recipes, default presets, fission, autotuning) and docs/proposals/schedule-ir-optops.md for the design rationale, including the pass-ordering contract (§2): schedules run after the whole optimize_proc pipeline (so they see fused code), apply folds freshly constructed guards by re-running simplify_llc (plus CSE and hoisting when a transform duplicated code), and there is no re-virtualization.
module Tn = Tnodetype optop = | Split of {axis : Indexing.symbol;The loop to split, identified by its index symbol.
*)factor : Base.int;Extent of the new inner loop.
*)outer : Low_level.axis_type;Axis type of the new outer loop (Serial = no retype).
inner : Low_level.axis_type;Axis type of the new inner loop.
*)outer_index : Indexing.symbol;inner_index : Indexing.symbol;}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.
| Swap of {outer : Indexing.symbol;inner : Indexing.symbol;}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.
*)| Retype of {axis : Indexing.symbol;ty : Low_level.axis_type;}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.
| Unroll of {axis : Indexing.symbol;materialize : Base.bool;}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).
| Partition of {axis : Indexing.symbol;The loop to partition, identified by its index symbol.
*)breakpoints : Base.int Base.list;Segment starts: strictly increasing, each strictly inside the loop range (from_ < b <= to_, so every segment is non-empty).
segment_indices : Indexing.symbol Base.list;}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 Ifs 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 of {axis : Indexing.symbol;The loop to pad, identified by its index symbol.
*)to_multiple_of : Base.int;The padded extent is the least multiple >= the loop extent.
}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 Splits by factors of M divide cleanly (no remainder guard), Stage tiles minted over the padded tile loops zero-fill their fringe (see Stage), and 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. A SCALAR pipeline keeps the leaf guards instead, and Privatize classifies them rather than rejecting them (gh-ocannl-730), which is what lets the register-blocktiled GPU family pad too. The surviving guards are exactly the flip points partition_breakpoints detects: a later Partition of an enclosing block loop at the last fully valid block specializes them away in the interior segments.
| Coalesce of {outer : Indexing.symbol;The outer loop of the pair, by its index symbol.
*)inner : Indexing.symbol;The inner loop, which must be the outer loop's whole body.
*)merged : Indexing.symbol;}Loop coalescing (gh-ocannl-1165): the perfectly nested Serial pair outer in [0, P) { inner in [0, Q) { body } } becomes one Serial loop merged in [0, P*Q) { body' }, where every access of body mentioning either symbol reads ... Iterator outer; Iterator inner ... on two adjacent axes whose dims are exactly P and Q, and is rewritten to ... Sub_axis; Iterator merged ...: the composed row-major index Q*outer + inner, which Indexing.Sub_axis expresses affinely (the run adds no term, the merged component ranges over the run's extent product) and Affine.view reads as one flattened coordinate. Anything else declines with Invalid_argument: a mention elsewhere (an Embed_index, a guard, an Affine term, one symbol alone, a dynamic or vector access, a Tile_mma, a Scan_loop), a pair right after an existing Sub_axis run (whose flattening the merged index would extend), or an axis whose dim differs from its loop extent — a padded inner axis in particular, whose stride is not Q. Iteration order and every address are unchanged, so the op is unconditionally semantics-preserving (op_legality: Op_legal once it applies). Purpose: a tile spanning both loops — the q/k/v projections' head axis merged into the column role, so a 64-wide column tile covers two 32-wide heads of useful columns.
| Stage of {source : Tn.t;tile_loops : Indexing.symbol Base.list;cooperative : Base.int Base.option;hoisted : Base.bool;swizzle : Low_level.swizzle_kind Base.option;pad_stride : Base.int Base.option;pipeline_depth : Base.int;tile_prec : Ops.prec Base.option;}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). The reads of source whose index vector mentions a tile loop are staged, and must use one index vector (v1) whose per-axis terms split cleanly into tile-loop terms (positive coefficients) and outer terms; every other read of source (a second consumer, gh-ocannl-1147) keeps reading the unwritten source. 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 each in-scope workgroup axis the tile does not use is folded into a serial copy loop of its own (gh-ocannl-1661): over e / w steps of w * s + t where its extent w divides the loop's e, else replacing a loop with e <= w under t < e; only an axis no remaining loop admits restricts the copy to one representative thread (t == 0). Shared stages sharing an anchor compose into ONE phase per iteration of it: a later stage's load nest is grouped in front of the phase barrier the earlier one left, back to back with its loads (gh-ocannl-567 — the loads write distinct tiles and depend on nothing of each other's, so a barrier between them would only serialize two independent copies), and the phase closer is shared rather than duplicated. Grouping requires that the phase already in place read nothing this stage's remap turns into a read of its tile. 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. A Local_scope body (the accumulator scope a materializing Unroll or Partition mints) is hoisted out of its position and may hold neither loads nor barriers (gh-ocannl-1467), so the staging point never lies inside one (gh-ocannl-1534): an anchor bound inside a scope body is refused, naming the anchor and the scope -- the tile holds the cells of one anchor value, so ahead of the scope its loads could not see the anchor, and the tile is sized by tile_loops alone, never grown to the scope's footprint behind the caller's back (adding the scope-bound loops of the staged index to tile_loops is that larger stage); an unanchored shared stage whose outermost tile loop is bound inside a scope body wraps the statement holding the scope instead of that loop. 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.
pipeline_depth = d with d > 1 software-pipelines a cooperative staging (gh-ocannl-487; an int rather than a bool so a search has a dimension, not a switch — 1 is the identity, taking exactly the unpipelined code path; d = 2 only: both the portable form (phase 1) and CUDA's cp.async arm (phase 2) have single-step lookahead — one prefetch per iteration, completed behind one barrier resp. one wait-all — so deeper depths would only consume shared memory and are rejected until the async arms grow commit-group/wait-group bookkeeping). Requires cooperative = Some _ and an anchor loop L* that is Serial, starts at 0 and has at least one iteration (the rotation needs a well-defined iteration order and the prologue fills buffer copy 0 — a dead anchor must not execute it; the no-anchor broadcast case has no rotor). The staging composition then becomes: a prologue copy of iteration from_ before L*, per iteration k a single barrier followed by the copy for k + 1 under an If (k < to_) guard (folded away on 1-trip loops) followed by the compute, and a trailing barrier after L* — 2N barriers become N + 1, and the prefetch is issued before the compute whose latency hides it. Pipelined stages sharing the anchor compose into the same phase — the later stage's prefetch is grouped behind the existing barrier, back to back with the earlier ones, keeping one barrier per iteration for all of them (a barrier between prefetches would force the earlier copy to complete before the compute, forfeiting its overlap). When a later Tensorize leaves the rotor body ending in a Tile_mma — a barrier by the emission contract, bracketed on every barrier-capable backend — apply's finalization elides the pipeline's own barriers (the loop-leading one, and the after-loop one where still adjacent): the intrinsic's bracket is the phase, so the tensorized pipelined kernel pays NO explicit barrier per iteration, one fewer than the depth-1 form (whose k-block loads feed the SAME iteration's compute, so its phase barrier is load-bearing — only the intrinsic's trailing bracket is emitted per iteration on every rendering form; see apply). A scalar pipelined compute keeps its explicit barriers. The tile is recorded in Low_level.optimized.pipelined: codegen allocates d rotating copies and selects the buffer by the loop counter (reads k mod d, in-loop writes (k+1) mod d, the prologue write copy 0), so the compute reads exactly the values the unpipelined form reads, in the same order — the pipelined rendering is bitwise identical to pipeline_depth = 1, a pure prefetch-timing transform (check_hardware_limits accounts the tile's shared-memory bytes times d). Tile_mma in the compute is compatible: its bracketing barriers are uniformly reached and separate same-buffer phases even more strongly than required. Backends that cannot render workgroup-shared staging (the serial C backends) reject the composition exactly as they reject unpipelined shared staging.
tile_prec = Some p mints the tile at storage precision p instead of the source's, folding the widening conversion into the staging copy (gh-ocannl-575, the gh-ocannl-517 packing seam): a packed panel of a narrow-storage operand becomes e.g. f32 scratch, converted once per element at pack time instead of once per read at the memory boundary, and precision-uniformity consumers (the register-tiled Tile_mma micro-kernel) see the compute precision directly. Only widenings are accepted — p must equal the source's storage precision or be single/double over a narrow float — because a widening is exact and its round-trip is the identity, making the transform unconditionally numerics-preserving; a narrowing tile would change values and is rejected. Composes with hoisted (the host-side pack converts) and with the in-kernel copy nest; the edge-guard zero fill is at the tile's precision either way.
| Privatize of {target : Tn.t;over : Indexing.symbol;acc_prec : Ops.prec;}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, whose conditions are classified (gh-ocannl-730). An iteration-INVARIANT condition (no memory reads, no symbols bound inside over's subtree) also gates the init-load and store-back, so lane-restricted accumulations (e.g. w == 0) privatize correctly. A condition that varies is kept on the update alone, and only in the two shapes that provably cannot select threads: one mentioning no hardware-typed loop symbol at all (a mask over one thread's own iterations — what a reduction-axis Pad leaves behind), and one that is literally target's own index on some axis compared against a bound no larger than that axis's dimension (a row/column pad mask; the transfers already carry that same edge guard). Any other varying condition — data dependent, or mixing a hardware symbol into a comparison that is not target's index — is rejected, since it could restrict which threads accumulate while the transfers write back an accumulator that never received the update. A Zero_out of target elsewhere is left in place unless the shared zero-seed proof licenses forwarding it directly into the tile: unconditional whole-node coverage and every repeated-cell loop inside over. Compose as: Splits → Stages → Privatize → materializing Unrolls (the unrolls then turn the tile accesses into constant-indexed, register-allocatable form).
The tile is minted at acc_prec, which must be target's storage precision or an exact widening of it (a narrow float to single/double), else the op is rejected. It is the accumulator residency: construct the op with privatize, which resolves it from the rendering backend's codegen_capabilities.accum_prec, so the privatized reduction widens and narrows exactly where the serial rendering's scope local does (gh-ocannl-1116). Like the backend's own residency, it applies only to what code generation would widen, and the decision is code generation's own, Low_level.decide_accum_width, asked of the program the schedule was APPLIED TO — never of the nest as earlier ops in the same schedule left it, nor rebuilt from it (gh-ocannl-1151). over's subtree must hold exactly ONE update of target, itself a base Low_level.accum_base_widens accepts (sibling accumulations narrow between each other in the serial rendering); the write of entry it descends from — the same position among target's Low_level.accumulating_writes, which schedule ops never reorder, or every such write where a materialization changed their count — must be one the serial rendering widens (Low_level.serial_accum_widths, asked once vouching for no hardware level and once for every one, since which hardware levels a backend serializes is code generation's knowledge — where the two answers differ the op is rejected); and the tile narrows once per execution of over, so where one execution of over covers a loop of that program — over is that loop, the outer half of a Split of a loop that covers it (a Split_reduce's block loop is one), or the loop a Coalesce merged from an outer loop that covers it, by a record apply keeps of its own Splits, Split_reduces, Coalesces and Partitions — the serial rendering must localize the write at that loop. The inner half of such a Split (a Split_reduce's in-block loop), and a segment of a Partition of such a loop, cover only part of it per execution, so a tile over either is at storage: wide, it would narrow once per execution of that part, where the serial rendering narrows once per execution of the whole loop. A loop any other op minted, or one minted from a loop with no record, carries no record and is judged by the decision alone (a Split_reduce's combine loops, over which no Privatize is admissible). A Split_reduce's own target accumulates into the minted partials node, which that program lacks, so a Privatize of it is at storage; one of a node accumulated beside it is judged as above. A widened scope-form update is re-keyed to the tile. So a Stage applied by an EARLIER apply is part of the program this one judges: its loads are sibling statements of the reduction, which code generation does not localize, and the tile stays at storage. Likewise sibling statements of any kind, a reduction level the cell varies under, a non-reduction or mixed-operator recurrence, an RNG-bearing contribution and routine logging all get a storage-precision tile, which code generation renders exactly as it would target, so the narrowing points do not move. After the finalization passes apply re-checks each wide tile (privatized_tile_widens) and refuses one whose accumulation the decision no longer widens — a backstop no known op sequence reaches: index-rewriting ops map the update and its self-read alike, the materializing Unroll and Partition mint exactly the scope grammar the decision reads back, no op rewrites a contribution to read the tile, and the sequences that would break it (a hardware-typed loop or a tensor-node write put into such a minted scope) are refused at their op by the scope grammar (gh-ocannl-1467, apply). A target whose storage is already its residency (f32, f64) is unaffected by all of this. A hardware-typed loop between over and the update is rejected — one target's cell does not mention would split the accumulation across threads, and the inner index loops of the tile part must be Serial — and so is any LATER op of the same apply that puts a hardware-typed loop inside the tile's per-thread lifetime (between its init-load and its store-back, Privatize's own over included, gh-ocannl-1467). Any other admissible value is an explicit numerics choice of the schedule, e.g. the storage precision reproduces per-step narrowing. Part of the schedule's identity (the cache saves it).
| Expand_zero of {tn : Tn.t;indices : Indexing.symbol Base.list;}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.
| Tensorize of {i : Indexing.symbol;j : Indexing.symbol;k : Indexing.symbol;lane : Indexing.symbol;simd_width : Base.int;Extent of the lane loop (the backend's Backend_intf.mma_capability.mma_simd_width; 32 on Metal and CUDA).
tile : Register_tile.t Base.option;The C-tile geometry of the register-tiled CPU rendering (gh-ocannl-619): rm rows by rn vectors of lanes, held across the k-loop. None (the default, and what every pre-existing schedule means) lets the renderer choose by its ranking model (Register_tile.default); Some is honoured exactly, or declined to the scalar fallback with a diagnostic when it does not fit (Register_tile.check: a width the file renders and the extents fill, the live-register budget) — never silently replaced, so a candidate timed under a geometry label ran that geometry or ran none. Part of the schedule's identity (the cache saves it, the canonical render digests it). The intrinsic (tensor-core) emissions ignore it.
}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.
| Fold_mma of {query : Indexing.symbol;The fold's query row loop, identified by its index symbol.
*)lane : Indexing.symbol;Fresh: the lane loop, one lane per query row of a block.
*)block : Indexing.symbol;Fresh: the query-block loop.
*)width : Base.int;The SIMD width: the query block, and the lane loop's extent.
*)}The online-softmax block fold on matrix units (gh-ocannl-1003): rewrites the whole fold nest Online_softmax emits under online_softmax_block -- the query loop split by width into block (Grid) and lane (Workgroup), the other row loops Grid, the per-row tiles re-minted with the block's rows in workgroup-shared memory, the score and value contractions in the scan body each one Tile_mma over the block between barriers, everything else per lane. GPU-only (a Workgroup loop enclosing barriers has no serial rendering); the query count must be a multiple of width, and the fold must have no key tail. Emitted by default_gpu; see apply_fold_mma.
| Fuse_epilogue of {target : Tn.t;}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. At the fragment and whole-K sites the store-back must complete target: one under an enclosing loop of extent > 1 that does not index it, or followed by another write of target (a later segment of a Partitioned block loop, gh-ocannl-1670), is refused. A later segment that is itself the last writer — the top-level edge segment of an unbatched k-split — is judged: its store-back completes target and takes the tail (gh-ocannl-1733), unless it runs more than once or its own index covers only part of the output. 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.
| Split_reduce of {axis : Indexing.symbol;The Serial reduction loop to split, by its index symbol.
*)target : Tn.t;The accumulated node whose reduction over axis is split.
num_blocks : Base.int;Number of per-block partials (the new block loop's extent).
*)block_index : Indexing.symbol;inner_index : Indexing.symbol;combine_indices : Indexing.symbol Base.list;}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 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 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):
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).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 Swap of accumulations). The result is deterministic 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).
val sexp_of_optop : optop -> Sexplib0.Sexp.tval sexp_of_schedule : schedule -> Sexplib0.Sexp.tval split :
axis:Indexing.symbol ->
factor:Base.int ->
outer:Low_level.axis_type ->
inner:Low_level.axis_type ->
optop * Indexing.symbol * Indexing.symbolBuilds a Split with fresh outer and inner index symbols (via Indexing.get_symbol) and returns them, so subsequent ops in a programmatically built schedule can reference the new loops.
val partition :
axis:Indexing.symbol ->
breakpoints:Base.int Base.list ->
optop * Indexing.symbol Base.listBuilds a Partition with one fresh index symbol per segment (via Indexing.get_symbol) and returns them in segment order, so subsequent ops in a programmatically built schedule can address individual segments.
val pad_edge_partition :
block_loop:Indexing.symbol ->
extent:Base.int ->
block:Base.int ->
optop Base.listThe Partition a Pad of an axis of extent to a multiple of block prescribes, once that axis is Split by block into block_loop and an inner loop: block_loop partitioned at its last full block, so every guard the pad left on the inner loop (block*block_loop + inner < extent, e.g. a zero-fringe Stage's edge guard on each packed element) is interval-decided in the full blocks and folds there, and survives only in the edge block (gh-ocannl-1631). Empty when block divides extent (no pad) or the axis has no full block. The one spelling of this split for the packed GEBP builders: the autotuner's packed sketch (Sketch_families) and bin/bench_gebp.ml (gh-ocannl-1670).
val partition_breakpoints :
axis:Indexing.symbol ->
Low_level.t ->
Base.int Base.listDerives candidate Partition breakpoints of the axis loop from the guards already present in its body: statement If conditions and scalar Where conditions (e.g. the virtualizer's per-component range guards of an inlined concatenation, Split's remainder guard, or a clamped window's range guards, gh-ocannl-504) are scanned for comparisons whose two sides differ by k*axis + off. With everything besides axis constant, such a comparison flips truth value at one point of the axis range. Non-axis symbols bound by other loops (inside the axis loop — e.g. the window symbol of a clamped-window guard — or enclosing it) are bounded by their loop ranges, putting off in an interval: the comparison is then always-true / mixed / always-false over the axis range, and both transition points are recorded — the mixed (boundary) segments are exactly delimited while the decided segments fold their guards. Returns the collected flip points that fall strictly inside the loop range, sorted and deduplicated (possibly empty — e.g. when every guard is already interval-decided); partitioning at them makes every such guard interval-decided within each segment (for interval-ranged comparisons: outside the mixed segments), so apply's trailing simplify erases them. Loops are located the way apply rewrites them (gh-ocannl-668): inside a Local_scope body, where the accumulation mints of a materializing Unroll or a Partition put the segment and inner loops; and over EVERY loop binding axis, not the first — a materializing Unroll leaves one copy per unrolled step with a different constant substituted for the unrolled index, so one source guard's copies flip at different points and the result is their union (each copy filtered by its own range). Raises Invalid_argument when no loop binds axis.
val fold_mma :
query:Indexing.symbol ->
width:Base.int ->
optop * Indexing.symbol * Indexing.symbolfold_mma ~query ~width mints the lane and block symbols of a Fold_mma.
val coalesce :
outer:Indexing.symbol ->
inner:Indexing.symbol ->
optop * Indexing.symbolcoalesce ~outer ~inner mints the merged loop's symbol of a Coalesce.
val expand_zero : tn:Tn.t -> optop * Indexing.symbol Base.listBuilds an Expand_zero with one fresh symbol per axis of tn (forcing tn's dims) and returns the symbols for subsequent Split/Retype ops.
val privatize :
accum_prec:(Ops.prec -> Ops.prec) ->
target:Tn.t ->
over:Indexing.symbol ->
optopBuilds a Privatize whose tile resides at accum_prec of target's storage precision (forcing it). Pass the rendering backend's resolution, (Context.codegen_capabilities ctx).accum_prec — the same function its code generation widens the serial rendering's accumulator by — so the privatized reduction's width is the policy's, not the schedule's (gh-ocannl-1116).
val privatize_index_vectors_refusal : Tn.t -> Base.stringThe detail of Privatize's v1 refusal of target: under the accumulation loop it is accessed at more than one index vector. Unlike Stage's (gh-ocannl-1147), this refusal is not relaxed: the target is written under the loop, so a second accessor at another vector would read a stale target between the tile's init-load and its store-back unless proved disjoint from every cell the loop writes. No tuned workload reaches it (gh-ocannl-1395's census, pinned by test/operations/privatize_census), which classifies Autotune.privatize_extension's drops by this text.
val split_reduce :
axis:Indexing.symbol ->
target:Tn.t ->
num_blocks:Base.int ->
optop * Indexing.symbol * Indexing.symbol * Indexing.symbol Base.listBuilds a Split_reduce with fresh block and inner index symbols plus one fresh combine symbol per axis of target (forcing target's dims), and returns them (op, block, inner, combine) so subsequent ops can annotate the pass-1 block loop and the combine nest's loops.
val fused_accumulation :
d:Tnode.t ->
di:Indexing.axis_index Base.array ->
Low_level.scalar_t ->
Low_level.scalar_t Base.optionLow_level.fused_accumulation: the fused form of a Tile_mma fallback's accumulation into d[di], None for an integer accumulation (gh-ocannl-824) or no product accumulated into d[di]. Tensorize no longer writes its fallback through it (gh-ocannl-1457): the renderer fuses the fallback unless the backend's compile pins contraction off. The sketch seeding asks it whether the register tile can render a site.
val tensorize :
?tile:Register_tile.t ->
i:Indexing.symbol ->
j:Indexing.symbol ->
k:Indexing.symbol ->
simd_width:Base.int ->
Base.unit ->
optop * Indexing.symbolBuilds a Tensorize with a fresh lane symbol (via Indexing.get_symbol) and returns it. ?tile is the register-tile geometry request (gh-ocannl-619), absent by default.
val add_symbol :
Indexing.axis_index ->
Indexing.symbol ->
Indexing.axis_index Base.optionadd_symbol idx s is idx plus s with coefficient 1: the per-tile coordinate Tensorize and Fuse_epilogue add to an accumulator's base index. None -- which each caller refuses in its own words -- for a Concat and for a Sub_axis: the latter has no per-axis affine form, and extended as the empty sum it would turn into an ordinary Iterator s coordinate where the flattened component after it carries the address (gh-ocannl-1165).
val split_reduce_hoist :
Low_level.optimized ->
optop ->
Indexing.symbol Base.listThe loops that would have to enclose the reduction for this Split_reduce to be recognized: the accumulation cell's symbols that are currently bound inside the reduction loop (gh-ocannl-537). This is the one rejection cause a loop interchange can remove — OCANNL lowers conv bias/weight gradients with the accumulated channel loops innermost and the reduction loops (batch, y, x) outside them, the exact inverse of the static form's pinning discipline — so autotune seeding hoists these symbols with a Swap chain and re-probes, rather than dropping the site. The answer comes from the recognizer itself (a hermetic probe of its own apply), not from a re-implementation of the discipline or a parse of its message.
[] for every other outcome: a legal op, a non-Split_reduce op, and any rejection an interchange cannot remove (the cell mentioning the reduction loop itself, an enclosing loop that pins no component, an unrecognized accumulation form).
val fuse_epilogue_witness :
target:Tn.t ->
Low_level.optimized ->
Base.string Base.optionWhy Fuse_epilogue on target would fail on this code — the recognizer's own refusal message, probed hermetically like split_reduce_hoist — or None when an eligible elementwise tail immediately follows the reduction over target and the fusion applies. The autotune family tree judges its fused-flavor branch with it (the check runs on the pre-schedule code, where the plain accumulation-nest site applies), so a site without a fusable tail carries the recognizer's reason as a construction-time refutation rather than silently minting no twins.
val can_fuse_epilogue : target:Tn.t -> Low_level.optimized -> Base.boolOption.is_none (fuse_epilogue_witness ~target opt): whether the fusion applies.
val hoistable_constant : Tn.t -> Base.boolWhether the node is eligible as a hoisted Stage source: declared value-constant (Tnode.known_host_constant) with registered host-init data to pack from. Shared by the autotune sketch and by Schedule_cache.canonicalize, which renders it per tensor node so that same-shape programs differing in operand constancy do not share cached schedules.
val hoist_known_constant : Tn.Placements.t -> Tn.t -> Base.boolThe constancy half of a hoisted Stage's source preconditions: Tnode.known_host_constant, or a known_constant placement in the lineage.
val may_replace_by_packed_tile : Tn.Placements.t -> Tn.t -> Base.boolgh-ocannl-1554: whether some schedule can take every read of the node off it — a hoisted Stage reads a host-packed tile instead and emits no load nest, so a candidate whose store is filtered to its code footprint holds the tile, sized by the staged accesses, and not the node. Exactly the hoisted-source preconditions apply checks on the node itself (registered host-init data and hoist_known_constant).
val settled_on_device : Tn.Placements.t -> Tn.t -> Base.boolWhether the node's placement is settled On_device in this lineage: the one placement the shared epilogue placement of Fuse_epilogue leaves alone.
val apply :
?static_indices:Indexing.static_symbol Base.list ->
schedule ->
Low_level.optimized ->
Low_level.optimizedApplies the ops left to right to the optimized code, then elides the staging barriers a later Tensorize made redundant, then re-runs Low_level.simplify_llc (which folds remainder and edge guards the loop extents prove) and, when a materializing Unroll duplicated code, CSE + cross-statement hoisting. Stage registers its tile in the traced store (and workgroup_shared when shared); the optimization context and merge node are never changed. Raises Invalid_argument when an op references a loop that does not exist at its point in the schedule, or violates an op precondition (see optop) -- the rendering of the typed refusal apply_classified passes through. An empty schedule is the identity. Every op is applied with the program the schedule started from in hand, which Privatize judges its tile width on, and its wide tiles are re-checked on the finalized program (privatized_tile_widens).
Each op is refused, naming it, when the program it leaves breaks the scope grammar (gh-ocannl-1467): a Local_scope body is hoisted out of its position, so it may hold neither a hardware-annotated loop (Low_level.scope_hardware_loop) nor an effect beyond its own locals (Low_level.scope_purity_violation) — e.g. a hardware Retype of a loop inside the scope a materializing Unroll minted, or a Stage whose loads land there. Only a violation the op added is refused; an entry program already breaking the grammar (hand-built IR) is left to the compile gates. Likewise when a Privatize tile minted earlier in the same application gets a hardware-typed loop inside its per-thread lifetime (see Privatize in optop); a tile minted by an earlier apply is not seen.
The barrier elision (gh-ocannl-487, widened to every staged anchor by gh-ocannl-567) is a synchronization-only rewrite — bitwise identical output, on every backend path. It rests on the Tile_mma emission contract, whose two halves are NOT symmetric: every rendering form ENDS the intrinsic block with a workgroup barrier, so a staging barrier that follows one (inside a staged anchor's body, or right after the loop) is dropped at any pipeline depth; but only some forms OPEN one per iteration — the fragment-scope form opens it once, on the scope wrapping the anchor loop — so a barrier is never elided against a following Tile_mma. Hence the loop-carried arm (a rotor body's leading barrier, against the previous iteration's bracket) stays pipelined-only, and a depth-1 k-block keeps the barrier between its loads and its compute.
val apply_classified :
?static_indices:Indexing.static_symbol Base.list ->
schedule ->
Low_level.optimized ->
Low_level.optimizedInternal candidate-facing variant of apply. An op precondition refusal (an op referencing a loop absent at its point in the schedule, or violating a precondition of optop) raises Schedule_outcome.Cause_at (Transform, Illegal_schedule { check = "Schedule.apply"; _ }), which a search counts as a declined candidate. Every other exception escaping the transforms or the finalization passes -- a bare Invalid_argument, an array bound, a List.zip_exn -- is a transform bug and propagates untyped, hence fatal under strict_failure_classification (gh-ocannl-1323).
val privatized_tile_widens : tile:Tn.t -> Low_level.t -> Base.boolThe re-check of a Privatize tile minted at the accumulator residency (gh-ocannl-1151): every write of tile in the program is either a transfer (it does not read the tile) or a self-accumulating update whose base Low_level.accum_base_widens — a Tile_mma's scalar fallback judged like any statement. apply asks it of the finalized program and refuses a tile failing it. Exposed for tests.
val with_transform_fault : (optop -> Base.unit) -> (Base.unit -> 'a) -> 'aTest seam (gh-ocannl-1323): with_transform_fault fault body runs body with fault called on each op just before apply / apply_classified applies it, so a test can make a transform raise what a transform bug raises — a genuine bug has no stable trigger. The calling domain's previous fault is restored on return or exception; the default is a no-op and no configuration selects it.
val sexp_of_op_verdict : op_verdict -> Sexplib0.Sexp.tval equal_op_verdict : op_verdict -> op_verdict -> Base.booltype pipeline_violation = | Read_outside_rotor of {tile : Tn.t;rotor : Indexing.symbol;}A read of tile has no loop binding rotor around it: there is no rotor value to select the copy with (writes outside the rotor loop are the prologue load, which fills copy 0). A Split or a materializing Unroll of the rotor loop removes its binder.
| Rotor_not_serial of {tile : Tn.t;rotor : Indexing.symbol;axis : Low_level.axis_type;}A loop binding rotor is no longer Serial (an Unroll in place, a retype): the rotation renders off the serial counter.
gh-ocannl-1231: how an IR breaks the staging scope's invariant for one of its pipelined tiles (Low_level.pipelined_tile) — the two positional facts the renderer's buffer rotation needs.
val sexp_of_pipeline_violation : pipeline_violation -> Sexplib0.Sexp.tval read_outside_rotor :
tile:Tn.t ->
rotor:Indexing.symbol ->
pipeline_violationval rotor_not_serial :
tile:Tn.t ->
rotor:Indexing.symbol ->
axis:Low_level.axis_type ->
pipeline_violationval is_read_outside_rotor : pipeline_violation -> Base.boolval is_rotor_not_serial : pipeline_violation -> Base.boolval read_outside_rotor_val :
pipeline_violation ->
([ `tile of Tn.t ] * [ `rotor of Indexing.symbol ]) Base.optionval rotor_not_serial_val :
pipeline_violation ->
([ `tile of Tn.t ]
* [ `rotor of Indexing.symbol ]
* [ `axis of Low_level.axis_type ])
Base.optionmodule Variants_of_pipeline_violation : sig ... endval pipeline_violation : Low_level.optimized -> pipeline_violation Base.optionThe first violation of the staging scope's invariant in opt.llc, over every tile of opt.pipelined: every loop binding a rotor is Serial, and every read of a pipelined tile sits inside a loop binding its rotor. None when opt has no pipelined tiles. These are the facts C_syntax checks while rendering (its typed Backend_codegen declines remain as the backstop); op_legality refuses the moves that would break them.
val pipeline_violation_message : pipeline_violation -> Base.stringval pipeline_violation_feature : pipeline_violation -> Base.stringThe decline feature naming the violation's kind — the key a search report aggregates Unsupported declines under.
val pipeline_violation_cause :
site:Base.string ->
pipeline_violation ->
Schedule_outcome.causeThe renderer's typed decline for a violation it meets while rendering (an IR that bypassed op_legality): Unsupported with pipeline_violation_feature and site ^ ": " followed by pipeline_violation_message. C_syntax builds both of its pipelined-tile declines through it, so the oracle and the renderer state each fact in one wording; pipelined_move_legality pins that, on the same IR, they reach the same violation.
type fragment_scope_violation = | Tile_replicated of {fragment : Tn.t;tiles : Base.int;}The fragment's load and store-back bracket tiles > 1 statements into it: a materializing Unroll or a Partition of a loop inside the scope.
| Scope_not_leading of {fragment : Tn.t;}The bracket holds the one statement, but the scope does not lead the statement sequence the renderer tries it in: after another recognized scope (a materializing Unroll of a loop around the scope).
| Tile_unbracketed of {fragment : Tn.t;}No load/store-back pair of the fragment encloses the statement among its siblings.
*)How a Tile_mma into a marked fragment falls outside every scope the renderer recognizes — rendered, then, as the lane-0 scalar fallback.
val sexp_of_fragment_scope_violation :
fragment_scope_violation ->
Sexplib0.Sexp.tval tile_replicated :
fragment:Tn.t ->
tiles:Base.int ->
fragment_scope_violationval scope_not_leading : fragment:Tn.t -> fragment_scope_violationval tile_unbracketed : fragment:Tn.t -> fragment_scope_violationval is_tile_replicated : fragment_scope_violation -> Base.boolval is_scope_not_leading : fragment_scope_violation -> Base.boolval is_tile_unbracketed : fragment_scope_violation -> Base.boolval tile_replicated_val :
fragment_scope_violation ->
([ `fragment of Tn.t ] * [ `tiles of Base.int ]) Base.optionval scope_not_leading_val :
fragment_scope_violation ->
[ `fragment of Tn.t ] Base.optionval tile_unbracketed_val :
fragment_scope_violation ->
[ `fragment of Tn.t ] Base.optionmodule Variants_of_fragment_scope_violation : sig ... endval fragment_scope_violations :
Low_level.optimized ->
fragment_scope_violation Base.listOne violation per Tile_mma into a fragment of opt.simdgroup_fragments that no recognized scope covers, in program order. A scope is Low_level.fragment_region over a marked fragment — the renderer's recognition — and coverage follows C_syntax.pp_ll's own route to it (each Seq tried whole, a recognized scope's trailing statements rendered one by one). [] when opt marks no fragment.
val fragment_scope_violation_message : fragment_scope_violation -> Base.stringval written_nodes : Low_level.t -> Base.Set.M(Ir.Tnode).tThe nodes the code writes: Set-family and Zero_out statements, Tile_mma accumulators, and the statements inside Local_scope bodies and scan-loop inits.
val op_legality : Low_level.optimized -> optop -> op_verdictgh-494 waypoint 3: the op-legality oracle. A schedule op is a thread-pairing transform over the routine's access relations; its obligation is decided by the Affine queries — pairing the op's loop symbol(s) as thread identity must leave every write-involving access pair of every written node Disjoint or Same_thread (with the reduction reassociation license for Vectorized retypes and Swaps of accumulations). Three-valued and sound in both proven directions: Op_legal proves the annotation race-free; Op_illegal proves a violation (e.g. a materialized node's unguarded write provably independent of the retyped axis); Op_unknown means "compile and see" — the op may be valid under semantics the queries do not model (per-thread scratch copies, renderer fallbacks) — and must never be treated as a rejection. Hardware annotations interleave sibling statements' threads with no grid-wide synchronization, so a hardware retype/split of a loop sharing an accessed node across its statement boundary (with a write on either side) reports Op_unknown: such pairs need the default annotator's aligned-mapping analysis, which the single-op oracle does not model. Vectorized retypes and Swaps only reorder within the nest and keep the intra-nest scope.
Tensorize's role assignment is decided: the micro-kernel recognition of apply (a pure function of the code) is probed, so an invalid role permutation — e.g. the reduction loop assigned to i/j, violating the discipline that the k role is the only rmw-carrying loop of the accumulation — is a proven Op_illegal; on structural success the affine queries decide i/j iteration independence under the reduction-reassociation license (the intrinsic reassociates k), with the hardware-lane cross-statement downgrade to Op_unknown as above. Stage is likewise probed hermetically (a raising apply is Op_illegal), then its implicit contract — the staged tile covers the reads it replaces within the staging scope — is the containment query Affine.read_covered_before over the fresh tile's accesses: a covered non-shared (packing) stage is Op_legal; shared staging (barrier placement and launch geometry are validated downstream) and hoisted staging (link-time host-side packing) report Op_unknown. Privatize/Expand_zero/Fuse_epilogue report Op_unknown; their own apply-time preconditions remain in force.
Before any of that, on code carrying pipelined tiles (gh-ocannl-1231): an op whose application leaves a pipeline_violation is a proven Op_illegal carrying pipeline_violation_message — the staging scope's invariant, stated once for every op shape, so that the beam's menu, seeding's schedule_legality and the renderer agree. Likewise, on code marking simdgroup fragments (gh-ocannl-1342): an op whose application leaves more fragment_scope_violations than it found is Op_illegal with fragment_scope_violation_message of the first it added. First of all, on code where some Local_scope body holds a loop: an op whose application adds a scope-grammar violation (gh-ocannl-1467, see apply) is Op_illegal with the refusal apply would raise — the race queries alone call a Grid retype of a scope-nested loop legal. The Privatize lifetime check needs the application's record of its tiles, so a single op judged here does not get it; schedule_legality does, as it applies the chain.
val schedule_legality :
Low_level.optimized ->
schedule ->
(optop * op_verdict) Base.listPer-op verdicts, each against the code with the preceding ops applied (on a hermetic copy — checking never mutates the argument). Stops after a proven-illegal op or a refused application; an op whose application is refused reports Op_illegal with the refusal. Any other exception propagates: it is a transform bug, not a verdict (gh-ocannl-1323).
val aligned_chains :
?max_chain:Base.int ->
?expanded_zeros:Tn.t Base.list ->
Low_level.optimized ->
(Low_level.t * (Indexing.symbol * Base.int) Base.list) Base.list Base.optionThe cross-nest analysis behind default_gpu (the CPU preset proves nests separately), as data: per top-level loop nest (paired with the nest statement itself), the outermost loops that may carry hardware geometry — already trimmed to the common aligned prefix each dependency component admits, so chain position k of two linked nests denotes the same hardware thread coordinate and the two extents at position k are equal. None when the analysis bails, i.e. when the kernel is not safely parallelizable at all.
max_chain caps each nest's chain length; the default 2 mirrors the presets' one-Grid + one-Workgroup shape. A caller supplying its own geometry per chain position passes the arity it needs — a batched matmul site's chain is its batch loops plus row plus column, and capping at 2 made every rank-3+ site's companion coverage decline (gh-ocannl-569). The alignment rule itself is arity-independent.
The presets consume this to emit their own Grid/Workgroup-per-nest geometry. A sketch pipeline (autotune) consumes it to cover the nests it does not build — a tensorized accumulation nest binds two Grid slots and a Workgroup lane, geometry no preset emits, and its companion nests must be annotated with the matching slot structure or Low_level.validate_parallel rejects their writes (gh-ocannl-521). Callers supplying their own geometry must keep the positional pairing: the alignment argument is exactly that thread (c0, c1, ...) covers the same index slice in every linked nest.
expanded_zeros names nodes whose whole-node Zero_out the caller's schedule will turn into a per-element nest (Expand_zero) carrying the caller's geometry. Such a statement is skipped for the query instead of bailing it as a bare materialized write — the write it stands for is not bare by the time the code is validated.
How the default GPU schedule's lane geometry treats a preamble reduction — a serial loop accumulating into the fused attention backward's dp scope local ahead of the lane loop, dp = Σ_e dO·v (gh-ocannl-1124); only that local qualifies (Online_softmax.reassociable_local), since the cooperative form reassociates under the fused backward's approximate-tier license. Preamble_refused: such a nest keeps its plain plan (the gh-ocannl-1003 stage-1 rule). Preamble_duplicated: the nest takes lanes and every lane recomputes the reduction serially, in its summation order. Preamble_cooperative: a reduction whose extent is the lane's whole one-loop workgroup and k of the device's simdgroup_width, k at most lane_all_reduce_simdgroups_for, is retyped Workgroup_reduce, sharing the lane's slot, and renders as an all-reduce leaving the sum in every lane (a butterfly within each simdgroup, then, past one, the partials through workgroup-shared memory between two barriers, gh-ocannl-1168); a kernel holding any other admitted reduction keeps its plain plan, never duplicated lanes (a renderer that still cannot shuffle runs the loop in every lane, the retype's serial meaning). Config gpu_lane_preamble_reduction, whose auto default resolves per device (lane_preamble_reduction_for).
val sexp_of_lane_preamble_reduction :
lane_preamble_reduction ->
Sexplib0.Sexp.tval equal_lane_preamble_reduction :
lane_preamble_reduction ->
lane_preamble_reduction ->
Base.boolval gpu_lane_preamble_reduction :
Base.unit ->
lane_preamble_reduction Base.optionConfig gpu_lane_preamble_reduction (auto, the default, is None | refused | duplicated | cooperative).
val lane_preamble_reduction_for :
Backend_intf.hardware_limits ->
lane_preamble_reductionThe treatment default_gpu applies at limits absent an explicit ?preamble_reduction: the configured mode, or under auto the device's economics — Preamble_cooperative where Backend_intf.hardware_limits' lane_scalar_recompute_cheap holds (measured: Metal, CUDA), Preamble_refused otherwise (HIP, the C backends, anything unmeasured), since the lanes recompute each pair's scalar preamble once per lane (gh-ocannl-1124).
Config gpu_lane_all_reduce_simdgroups (auto, the default, is None | a positive count).
val lane_all_reduce_simdgroups_for : Backend_intf.hardware_limits -> Base.intThe most simdgroups a Preamble_cooperative reduction may span at limits absent an explicit ?all_reduce_simdgroups: the configured count, or under auto the device's measured Backend_intf.hardware_limits' lane_all_reduce_simdgroups (1 wherever unmeasured, so a head wider than one simdgroup keeps its plain plan; gh-ocannl-1168).
How default_gpu weighs its lane geometry (gh-ocannl-1167). Every lane recomputes its nest's per-pair scalar preamble, which the plain plan pays once per thread. Lanes_cut: taken wherever a kernel admits it, and fission_keep_mapping cuts a kernel rather than merge a lane nest out of its lanes. Lanes_admitted: taken wherever a kernel admits it, but merges between nests with a lane path are judged on the plain plans, so such a merge costing a nest nothing but its lanes is taken; a neighbour with no lane path is still cut off rather than cost a nest its lanes (gh-ocannl-1239, mapping_probe). Lanes_off: never taken, so the preamble-reduction treatment has nothing to apply to. Config gpu_serial_lanes, whose auto default resolves per device (serial_lanes_for).
val sexp_of_serial_lanes : serial_lanes -> Sexplib0.Sexp.tval equal_serial_lanes : serial_lanes -> serial_lanes -> Base.boolval gpu_serial_lanes : Base.unit -> serial_lanes Base.optionConfig gpu_serial_lanes (auto, the default, is None | cut | admitted | off).
val serial_lanes_for : Backend_intf.hardware_limits -> serial_lanesThe configured mode, or under auto the device's economics: Lanes_cut where Backend_intf.hardware_limits' lane_scalar_recompute_cheap holds (measured: Metal, CUDA), Lanes_admitted otherwise (HIP: the forward value pass wins on lanes in its own kernel, the fused backward's dV loses when cut from dK for them).
val resolve_search_shaping_value :
Backend_intf.hardware_limits ->
key:Base.string ->
Base.string ->
Base.stringWhat the canonical value of the configuration key key (Utils.canonical_config_value) resolves to on the device limits describes, for Utils.config_class_fingerprint ~resolve (gh-ocannl-1272): an auto lane mode becomes the mode lane_preamble_reduction_for, lane_all_reduce_simdgroups_for or serial_lanes_for picks at limits, and an explicit one stays itself; a gpu_schedule_block_size becomes its clamp to the device's caps. Any other key, and a value its parser rejects, is returned unchanged.
val default_gpu :
?block_size:Base.int ->
?min_parallel:Base.int ->
?workgroup_fill:Base.int ->
?preamble_reduction:lane_preamble_reduction ->
?all_reduce_simdgroups:Base.int ->
?lanes:Base.bool ->
?limits:Backend_intf.hardware_limits ->
Low_level.optimized ->
scheduleThe default GPU annotator preset (schedule-ir-optops §6): for each top-level loop nest whose parallelism is provable from the lowered code alone, annotate every loop of its proved parallel chain (gh-ocannl-1133) — the leading loops Grid (Grid slots >= 2 fold onto .z), the innermost the Workgroup lane, split Grid blocks x Workgroup by block_size when it is longer (default from config gpu_schedule_block_size = 256, clamped to limits' Backend_intf.hardware_limits.max_threads_per_workgroup and .x workgroup cap when given — the configured block size is a target, the device's workgroup capacity a hard cap). While the workgroup holds fewer than workgroup_fill threads (default from config gpu_schedule_workgroup_fill = 256), the loops just above the lane join it as further Workgroup dimensions (at most three, the product within the block size, each extent within its dimension's cap). The plan is taken only when every chain-carrying nest of the kernel fits one common hardware topology (a nest one Grid slot short may split its lane into a one-block Grid loop), the launch fits the device's per-dimension caps, and no nest does less useful parallel work than under the two-loop presets below, with more work somewhere; otherwise the presets apply. The presets annotate exactly one Grid and one Workgroup loop per nest (splitting a single parallel loop by block_size; a second loop longer than the block size keeps its outer part Serial). A loop is parallelizable when its index occurs as a plain Iterator component in every materialized write vector beneath it — the same coverage property Low_level.validate_parallel enforces, used generatively — and the kernel passes a conservative race analysis (all accesses to written nodes agree on parallel-index components, no Zero_out of materialized nodes, no barriers or opaque statements; reduction loops stay serial). Both plans map proved chains: the lane plans propose each nest's whole chain to that analysis, the presets a pair — when the leading parallel extent is below min_parallel, a later pair with no fewer grid groups and a larger grid-times-clamped-workgroup product. Skipped leading loops remain serial. This choice precedes the race analysis; if it fails, or alignment loses groups, active lanes, or the launch threshold, the original outermost pair is used instead. Expanded whole-node zeros (zero_expansion) use the same plans. Unless lanes is false (default: the serial_lanes_for mode at limits is not Lanes_off), a lane geometry takes priority over both (gh-ocannl-1003): when every nest carrying a chain has a parallel loop under a serial loop past loop-free lane-uniform scalar work (declarations and assignments of scope locals -- the online-softmax hoist's value pass, for (b, s, h) { for t { p := P[s, t]; for e { O[s, e] += p * V[t, e] } } }; a preamble holding an inlined reduction is excluded, since every lane would recompute it; a preamble reduction loop into a scope local is admitted per ?preamble_reduction, see lane_preamble_reduction, spanning at most ?all_reduce_simdgroups simdgroups when cooperative), the chain extends through that preamble uncapped: the loops above the serial loop become Grid loops (slots >= 2 fold onto .z) and the loop past it a Workgroup lane, Grid (b, s, h) -> Serial t -> Workgroup e — taken when the planner finds such nests one common topology, the launch fits the device's caps, and it has more threads than the plain plan's geometry. The race analysis is the same: thread identity is the tuple of chain symbols wherever they sit, the serial loop runs whole in every thread, and Low_level.validate_parallel already accepts a hardware loop under a serial one. Cross-nest producer/consumer (or WAW/WAR) pairs over a written node are allowed only when aligned: the linked nests' chains are trimmed to a common equal-extent prefix — identical annotation geometry, so each hardware thread covers the same index slice in every linked nest — and per axis position the paired accesses either both use plain Iterators of same-chain-position parallel symbols or neither mentions one; otherwise the analysis bails. For non-materialized (per-thread copy) scratch the edge additionally requires value thread-invariance at the chosen trim: a chain symbol feeding a scratch write's value without pinning the written cell would leave each consumer thread's copy holding its own chunk's last value where the serial reference holds the last chunk's, so the trim search serializes that loop (gh-494; direct syntactic dependence only). Returns the empty schedule when any check fails or when the largest parallelizable nest has fewer than min_parallel iterations (default from config gpu_schedule_min_parallel = 64: a kernel launches either way, so any real parallelism beats the serial 1x1 fallback — a single GPU thread is 1-2 orders of magnitude slower than a CPU core; the remaining small threshold keeps sub-simdgroup-scale programs fully serial so their segments coalesce and placements stay unchanged).
val default_cpu : ?min_work:Base.int -> Low_level.optimized -> scheduleThe default CPU annotator preset (gh-ocannl-1458): per top-level nest, at most one loop is retyped to Grid -- pool-backed Grid rendering in the C backend (docs/proposals/gh-ocannl-164.md) partitions it into contiguous chunks on a process-global native thread pool, and a Workgroup split would only add loop structure that runs serially inside a chunk. The candidate is the nest's first path loop whose index is a plain Iterator in every materialized write -- in every scratch write for a nest writing only scratch (Local, one function-scope array on the C backends; gh-ocannl-1343) -- with at least two iterations and the per-nest hazard analysis passing over that ONE symbol: the loop actually made parallel is the one proven independent.
No cross-nest alignment, unlike default_gpu. The C renderer runs each pool-parallel Grid loop as one synchronous dispatch (dispatch_apply / #pragma omp parallel for) that joins before the next statement, and everything outside those loops once, on the calling thread, in program order. So two accesses in different top-level statements are ordered by a join for every placement -- materialized buffers and shared function-scope locals alike; a Local is privatized per chunk only when all its accesses in the kernel sit inside that loop's body, so no privatized value crosses statements; within one dispatch the hazard proof, C_syntax.parallel_grid_safe (locals) and the renderer's store-separation check decide; a Grid loop under serial loops forks and joins per instance. Low_level.validate_parallel's ~fork_join mode is the validator's side of the same model: writes outside every Grid loop are not subject to coverage, so a serial nest may sit beside a parallel one. A statement the access scan cannot model, or a nest with no candidate or a failed proof, stays serial without affecting its siblings.
Profitability is per dispatch: a candidate is retyped only when one run of it does at least min_work scalar updates (Cost_model.fork_join_updates, the renderer's launch-floor units; enclosing serial iterations and sibling nests do not count). Default from config cpu_schedule_min_work = 16384; a non-positive value disables the gate (the tuner's ungated presets pass ~min_work:0).
val channel_vectors : base:schedule -> Low_level.optimized -> schedulePer-channel vector accumulators for direct conv nests (gh-ocannl-1630): per top-level statement that is a perfect Serial accumulation nest of write loops then reduction loops, whose destination's minor axis is the innermost write loop oc, with one operand reading oc (and otherwise only reduction loops, not written in the routine) and some operand reading a sliding window, at single or double storage throughout: Swaps sinking oc below every reduction loop (each cell keeps its terms in order), a packing Stage of the oc operand channel-minor at the routine root (skipped when it is channel-minor already), a Privatize of the destination over the outermost reduction loop at its storage precision, and oc retyped Vectorized. Bitwise identical to the scalar-chain rendering. A site whose ops do not apply after base (on a scratch copy of the store) is left out, and so is everything when the sites do not apply together.
val cpu_preset : ?min_work:Base.int -> Low_level.optimized -> scheduleThe default CPU preset: default_cpu, followed, under config cpu_schedule_channel_vectors (default true), by channel_vectors over it. What maybe_default_schedule and maybe_default_schedules apply on CPU backends, and the tuner's segment presets.
Config schedule_log_launches: backend compile logs one stderr line per compiled segment with its launch grid/block dims and statement count — for diffing what two compiles of nominally identical code actually emit.
val code_footprint : Low_level.t -> Base.Set.M(Tn).t * Base.boolAll tensor nodes the code references (scope ids' backing nodes included) — the set a fission segment's traced store is filtered to — and whether it reads the merge buffer.
val may_keep_off_device :
facts:Backend_intf.backend_facts ->
limits:Backend_intf.hardware_limits ->
Tn.Placements.t ->
Tn.t ->
Low_level.traced_array ->
Base.boolgh-ocannl-1554: whether some schedule can keep this existing node off the device, the premise on which the device-memory floor (Backends.link_floor) leaves it out. The one schedule op that turns an existing node Local is Fuse_epilogue's shared placement, of a fusion target (written, read by the fused tail) not settled_on_device; it exists only where the backend states workgroup_shared_placement, and charges the node's whole size to the workgroup memory that check_hardware_limits_classified holds to max_workgroup_memory_bytes for every kernel a compile emits, so a node larger than that limit cannot leave the device. None as the limit bounds nothing.
val search_shaping_key_read :
facts:Backend_intf.backend_facts ->
Base.string ->
Base.boolWhether a search on a backend with facts can read the default pipeline's configuration key key (gh-ocannl-1331): the GPU-only keys (automatic_gpu_schedule, gpu_schedule_*, the lane modes, gpu_fission_keep_mapping) exactly where Backend_intf.binds_hardware_grid holds, the CPU-only ones (automatic_cpu_schedule, cpu_schedule_min_work) where Backend_intf.renders_on_host_pool does, and any other key everywhere. Autotune.search_shaping_key_read adds the seeded plan's keys.
val automatic_schedule_active : facts:Backend_intf.backend_facts -> Base.boolWhether maybe_default_schedule (and maybe_default_schedules) would apply a default preset on this backend rather than the identity: the backend binds Grid loops, the corresponding config gate (automatic_gpu_schedule / automatic_cpu_schedule) is on, and runtime kernel logging (debug_log_from_routines) is off. Callers that substitute their own default-schedule choice (e.g. the cost-model default selection of gh-ocannl-491) must respect this gate.
The config schedule_fission gate consulted by maybe_default_schedules: whether the untuned default pipeline fissions into per-segment kernels (subject to automatic_schedule_active), or applies the whole-routine maybe_default_schedule instead.
val default_schedule_fingerprint :
facts:Backend_intf.backend_facts ->
Base.stringA stable summary of the policy and configuration that shape the untuned default pipeline on this backend: the automatic_schedule_active gate ("inactive" when it is off — the untuned default is then the unscheduled serial form), the policy's version tags (bumped when a change redefines the untuned default), the schedule_fission gate, and the preset thresholds (gpu_schedule_block_size / gpu_schedule_min_parallel / cpu_schedule_min_work). Diagnostics recorded against one default pipeline (e.g. the autotuner's default_ms, gh-ocannl-552) compare fingerprints to detect that a config change redefined the default (Codex P2 on PR #279). Does not cover per-device hardware limits — cache consumers already key per backend.
val maybe_default_schedule :
facts:Backend_intf.backend_facts ->
?limits:Backend_intf.hardware_limits ->
static_indices:Indexing.static_symbol Base.list ->
Low_level.optimized ->
Low_level.optimizedThe implicit transform applied by backend compile when the caller passes no ?lowered_transform: apply of default_gpu on GPU backends, of default_cpu on CPU backends, the identity otherwise. limits (default Backend_intf.no_hardware_limits) should be the compiling backend's Backend_intf.Backend_device_common.hardware_limits. Disabled by config automatic_gpu_schedule=false / automatic_cpu_schedule=false respectively, and skipped when runtime kernel logging (debug_log_from_routines) is active, to keep logs serial and deterministic.
val zero_expansion :
?block_size:Base.int ->
?min_parallel:Base.int ->
?workgroup_fill:Base.int ->
limits:Backend_intf.hardware_limits ->
Tnode.t Base.list ->
scheduleThe expand-and-annotate schedule maybe_default_schedules applies to a fission segment of materialized whole-node Zero_outs on GPU backends: expand_zero plus the same geometry policy as default_gpu — the lane plans over each node's non-singleton axes (its whole loop set is a proved chain: each cell is written once), under one common topology across the segment's nodes, else the two-loop presets. Below min_parallel (largest node) the zeros stay whole-node (a serial kernel renders them as memset). Exposed for callers (e.g. the autotuner) that replicate the default fission pipeline with custom per-segment schedules.
val statement_mappings :
Low_level.t ->
schedule ->
(Base.int * Base.int) Base.listThe hardware mapping schedule gives each top-level statement of the code, as (groups, active threads) of the statement's OWN loops, read off the ops: what schedule-aware fission (fission_keep_mapping) compares between a merged kernel and a statement alone. A Workgroup_reduce retype shares the slot of the statement's widest Workgroup loop (the lane geometry's cooperative preamble reduction, gh-ocannl-1124), adding threads only past its width.
type mapping_probe = {mp_kernel : Low_level.optimized -> schedule;mp_lane_trade : (Low_level.optimized -> schedule) Base.option Base.Lazy.t;}The probes of schedule-aware fission's merge rule (gh-ocannl-1126): mp_kernel, the schedule a merged kernel will receive, against which a merge is refused when it gives some statement fewer groups or active threads of its own loops than it gets alone; and, where lanes are only Lanes_admitted (gh-ocannl-1167), mp_lane_trade, the same schedule without the lane geometry, on which a merge is judged instead -- unless the merged kernel would hold a loop nest with no lane path beside one with a lane path, which must pass on both (gh-ocannl-1239: such a kernel declines the lanes for a neighbour that could never take them).
val mapping_probe :
?lane_trade:(Low_level.optimized -> schedule) ->
(Low_level.optimized -> schedule) ->
mapping_probemapping_probe ?lane_trade kernel: the probes, no lane trade unless given.
val fission_keep_mapping :
is_gpu:Base.bool ->
limits:Backend_intf.hardware_limits ->
mapping_probe Base.optionThe keep_mapping probes maybe_default_schedules passes to fission_scheduled: default_gpu at limits on a GPU backend while config gpu_fission_keep_mapping is on (the default), None otherwise -- with its lane geometry unless serial_lanes_for is Lanes_off, and under Lanes_admitted with the lane-free default_gpu as the lane trade, so a merge costing a nest only its lanes is taken where every nest of the kernel has a lane path. A caller replicating the default segmentation (the autotuner's fissioned candidates) passes the same, so its segments are the untuned pipeline's.
val fission_scheduled :
?promote_locals:Base.bool ->
?arity_cuts:Base.bool ->
?fold_zeros:Base.bool ->
?keep_mapping:mapping_probe ->
preset:(Low_level.optimized -> schedule) ->
zero_sched:(Tnode.t Base.list -> schedule) ->
static_indices:Indexing.static_symbol Base.list ->
Low_level.optimized ->
([ `Normal | `Zeros | `Solo ]
* Low_level.optimized
* schedule
* Low_level.optimized)
Base.listThe kernel-fission pipeline underlying maybe_default_schedules, with caller-supplied per-segment schedules and a per-segment result: the routine's top-level statements are partitioned at cross-workgroup dependency edges exactly as described there (edges the aligned cross-nest rule of default_gpu proves race-free without losing any nest's standalone parallelism do not cut), preset is called on each `Normal segment's (pre-schedule) optimized slice and zero_sched on each `Zeros segment's nodes, and each result tuple carries the segment kind, the pre-schedule segment, the schedule chosen for it, and the scheduled segment (apply of the schedule). `Solo segments (opaque to the analysis, or coalesced runs of unannotated segments) get the empty schedule.
fold_zeros (default false, gh-ocannl-1175): a `Zeros segment holding a single zero that initializes the reduction opening the next segment folds into that segment, expanded by zero_sched's Expand_zero, when the merge rules below re-admit the whole segment with the per-cell zero at its head; no other boundary moves, so the fold never merges two reductions the zero kept apart. Only the autotuner's and model selector's per-segment sketch candidates pass it, so the fold ships only where a timed (or modeled) sketch pipeline carries it: the untuned default keeps the separate zero kernel, because an untiled accumulation with its zero folded measured 7% slower on HIP (unified memory) although its emitted source only drops the opening read. A replayed segmentation refolds where its record says.
When fission does not apply (single segment, unfissionable crossings, or everything coalesces back) the result is a single `Normal tuple over the whole routine with preset's schedule. Callers compile each scheduled segment as its own kernel in order (the plural transform seam of backend compile); see maybe_default_schedules for the synchronization contract.
arity_cuts (default false, gh-ocannl-574): segment for the full-arity sketch pipelines instead of the default presets. The no-parallelism-loss guard normally compares chains under the presets' Grid+Workgroup cap, so a merge that trims a rank-3+ site's minor axis reads as lossless — correct for the default annotators, which never bind more than that prefix, but the GPU matmul sketches annotate the site's whole chain and cover every materialized-writing nest of the kernel with the same geometry (companion_geometry, gh-ocannl-521/569). Under arity_cuts chains are analyzed uncapped and a merge additionally requires every materialized-writing nest to keep an identical extent list, so a companion that cannot follow the site's arity — a reduction over the site's minor axis and the reduction target's initialization nest, the lm_head's max-logits row — is cut into its own downstream kernel (stream order supplies the synchronization) and the site's kernel seeds at full arity. The finer segmentation costs launches the default pipeline does not want to pay unconditionally, so this is a candidate-generation mode (the autotuner times it), never the default.
keep_mapping (gh-ocannl-1126): the default GPU schedule, as fission_keep_mapping probes it -- the schedule each candidate kernel would actually receive, except that where lanes are Lanes_admitted a merge between nests with a lane path is judged without the lane geometry, so no boundary is kept for lanes between them (mapping_probe, gh-ocannl-1239). A merge the rules above admit -- an aligned dependent merge, and a conflict-free one, which they admit unconditionally -- is still refused when some statement of the merged kernel gets fewer Grid groups or fewer active threads of its OWN loops under keep_mapping than it gets in a kernel of its own: the kernel boundary is kept rather than a nest's mapping lost. Only the refused merges add cuts, so a merge that keeps every mapping (an elementwise tail over the same chain) still saves its launch. The same rule holds a `Zeros segment together (gh-ocannl-1169): a zero joins one only when no node of the merged segment gets fewer groups or active threads under zero_sched than alone -- a union of zeros with no common lane-plan topology sends every node to the two-loop presets. Each maximal run of whole-node zeros is stably sorted by the zeroed node's rank and dimensions in every mode (the zeros commute), so the cuts fall at shape changes rather than at every alternation of a backward pass's gradient zeros. Ignored under arity_cuts. None (the default): the legality rules alone decide.
promote_locals (default false): promote statement-crossing Local scratch to On_device before segmentation. A nest whose only writes land in Local scratch gets no parallel chain (the annotator's coverage property quantifies over materialized writes), and Local producer/consumer edges do not cut — so such producers either drag their segment down to a serial 1x1 launch or are redundantly re-executed by every hardware thread; small reduction intermediates (layer-norm statistics, softmax max/denominator) land exactly here via the Local stack threshold. Promotion lets the ordinary materialized machinery apply; promotions fission does not end up needing (single-kernel fallbacks, or all accesses within one segment after coalescing) are restored. maybe_default_schedules passes true on GPU backends, where a serial nest costs orders of magnitude more than on CPU.
A recorded segment's kind: a derived segment's (`Normal, `Zeros, `Solo), or `Unfolded (gh-ocannl-1480) -- the unfissioned routine's one segment whose head is a whole-node zero left whole although its fold is provable. A `Normal segment opening with such a zero is a folded one and replays with the zero expanded; an `Unfolded one replays as a `Normal segment with the zero as it stands, so a replay rebuilds exactly the code its recording compile ran.
val sexp_of_segment_kind : segment_kind -> Sexplib0.Sexp.tval segment_kind_of_sexp : Sexplib0.Sexp.t -> segment_kindval __segment_kind_of_sexp__ : Sexplib0.Sexp.t -> segment_kindval equal_segment_kind : segment_kind -> segment_kind -> Base.booltype segmentation = (segment_kind * Base.int) Base.listA fission segmentation as data (gh-ocannl-1164): the kind and the length, in units, of each segment in order — a unit being one top-level statement with the comments before it, with each maximal run of whole-node zeros stably sorted by the zeroed node's shape, so the units are a function of the code alone. One segment over every unit is the unfissioned routine. The schedule cache persists it with a fissioned winner, so a replay cuts where the winner was cut instead of re-deriving the cuts under the current policy.
val sexp_of_segmentation : segmentation -> Sexplib0.Sexp.tval segmentation_of_sexp : Sexplib0.Sexp.t -> segmentationval equal_segmentation : segmentation -> segmentation -> Base.boolval segment_class : segment_kind -> [ `Normal | `Zeros | `Solo ]The kind a recorded segment is replayed as: `Unfolded is a `Normal segment.
val fission_segmented :
?promote_locals:Base.bool ->
?arity_cuts:Base.bool ->
?fold_zeros:Base.bool ->
?keep_mapping:mapping_probe ->
?replay:
(segmentation * (Base.int -> Low_level.optimized -> schedule Base.option)) ->
preset:(Low_level.optimized -> schedule) ->
zero_sched:(Tnode.t Base.list -> schedule) ->
static_indices:Indexing.static_symbol Base.list ->
Low_level.optimized ->
segmentation
* ([ `Normal | `Zeros | `Solo ]
* Low_level.optimized
* schedule
* Low_level.optimized)
Base.listfission_scheduled, also returning the segmentation of its result (one element per tuple), or — given replay = (segmentation, schedule) — applying that recorded segmentation instead of deriving one. A recorded segmentation is neither re-grouped nor coalesced, so arity_cuts and keep_mapping do not reach it; the promotions and scope-local replicas are recomputed for its boundaries exactly as for derived ones, and the ith segment's (pre-schedule) slice gets schedule i whatever its kind — by position, because segments of one structure can differ in placements, hence in kind and schedule. Where schedule i is None the segment gets its kind's schedule as in a derived segmentation (preset, zero_sched, or none for a `Solo): a per-segment substitution over a recorded segmentation (the autotuner's sketch candidates, gh-ocannl-1243) names only the segments it schedules. A segmentation that does not fit the routine — its lengths do not sum to the routine's units, a segment is empty, or a boundary separates a scope-local definition from a use it cannot be replicated for — raises Schedule_outcome.Cause_at with an Illegal_schedule cause, a candidate's ordinary decline; it is never repaired into another segmentation. Whatever raises, the promotions made so far are restored first.
val maybe_default_schedules :
facts:Backend_intf.backend_facts ->
?limits:Backend_intf.hardware_limits ->
static_indices:Indexing.static_symbol Base.list ->
Low_level.optimized ->
Low_level.optimized Base.listLike maybe_default_schedule, but with kernel fission: the routine's top-level statements are partitioned into segments at the cross-workgroup dependency edges the default annotator's interference analysis would otherwise reject wholesale — materialized producer/consumer (and WAW/WAR) pairs of sibling statements, bare materialized writes, materialized whole-node Zero_outs, statements opaque to the analysis — and each segment receives the default schedule independently, on its own launch geometry. The caller compiles each returned optimized as its own kernel and runs them in order on the routine's stream, chaining a device-side event at each boundary — that event supplies the grid-wide synchronization at each cut (queue FIFO alone does not order overlapping command buffers over Metal's untracked resources). Scalar scope-locals hoisted across statements are replicated into consuming segments when provably value-preserving (segments merge back and run serially otherwise); Local-placed scratch whose live range crosses a cut is promoted to On_device in the compile's placement fork. On GPU, segments consisting of materialized Zero_outs are expanded (expand_zero) and annotated like ordinary nests. Adjacent segments that end up unannotated are coalesced, so an all-serial routine stays a single kernel. Returns a singleton equal to maybe_default_schedule's result when fission does not apply (non-GPU/CPU backend, automatic scheduling disabled, config schedule_fission=false, kernel logging active, or nothing to split).
type launch_geometry = {lg_grid_y : Base.int Base.option;lg_grid_z : Base.int Base.option;lg_block_x : Base.int Base.option;lg_block_y : Base.int Base.option;lg_block_z : Base.int Base.option;}As much of a candidate's launch geometry as its caller knows (gh-ocannl-709): the two grid dimensions backends cap and the workgroup's three, None meaning "not predicted here" — which exempts that dimension rather than refusing on it. Five fields, not six: grid.(0) is 2^31-scale on every backend that binds hardware axes, so no backend has a cap to report for it. check_hardware_limits_classified fills all five from the lowered code (launch_geometry_of_dims); autotune's seeding pre-filter predicts what it can from the parameters it is about to commit to. A prediction is a lower bound — it describes the site's own nest, while Low_level.launch_dims maxes over the kernel's zeroing and companion nests too — so an under-prediction costs a compile the gate then declines, and only an over-prediction could withhold a legal candidate.
val unknown_launch_geometry : launch_geometryAll-None: predicts nothing, refuses nothing. The base a family builds its prediction on.
val launch_geometry_of_dims : Low_level.launch_dims -> launch_geometryThe geometry a lowered kernel will actually launch with — every field Some.
val launch_geometry_of_nests :
(Base.int Base.list * Base.int Base.list) Base.list ->
launch_geometry Base.optionThe geometry a kernel launches with when its nests' hardware loops have these extents — per nest, its Grid and its Workgroup loop extents in nest order, outermost first (gh-ocannl-1133). The one encoding of the positional slot rule outside Low_level.launch_dims, which reads the same rule off lowered code: among loops of one kind the innermost binds .x, the next .y, the next .z; the maximum is taken per slot across nests, and Grid slots >= 2 fold the product of their per-slot maxima onto .z. Every field is Some; None when that product overflows. The default annotators' lane plans and the sketch families' predictions (Autotune.matmul_launch_geometry / Autotune.conv_launch_geometry) both read it.
type launch_excess = {lx_resource : Schedule_outcome.resource;lx_requested : Base.int;lx_limit : Base.int;lx_phrase : Base.string;}The first dimension of a geometry the device refuses. lx_phrase is the verb phrase both callers render — e.g. "requests a .z workgroup extent of 128, exceeding the device limit of 64" — so a gate detail and a seeding refutation witness say the same thing about the same candidate.
val launch_geometry_excess :
limits:Backend_intf.hardware_limits ->
launch_geometry ->
launch_excess Base.optionThe one static reading of a device's per-dimension launch caps (gh-ocannl-709), consulted both by check_hardware_limits_classified (on the lowered code's geometry) and by autotune's seeding pre-filter (on a candidate's predicted geometry), so the two cannot disagree about what a device permits and a cap is never encoded twice. None = every predicted dimension fits. Rows are tested in the order .x/.y/.z workgroup, .y grid, .z grid fold; a None cap (Backend_intf.hardware_limits.max_workgroup_dims on the C backends, Backend_intf.hardware_limits.max_grid_yz on Metal) exempts its dimensions. The workgroup thread PRODUCT is deliberately not a row: it is not a per-dimension geometry question, and only the gate checks it.
val kernel_workgroup_bytes :
capabilities:Backend_intf.codegen_capabilities ->
arrays:Base.int Base.list ->
mma_triples:(Ops.prec * Ops.prec * Ops.prec) Base.list ->
Base.intThe workgroup-shared charge of a kernel declaring arrays (each one array's declared bytes) and tile-MMA emission scopes of the (d, a, b) storage triples mma_triples: every array rounded up to the backend's Backend_intf.codegen_capabilities.workgroup_allocation_alignment (gh-ocannl-1330), plus the backend's per-kernel Backend_intf.codegen_capabilities.mma_kernel_workgroup_bytes over the triples, once (gh-ocannl-1255). What workgroup_memory_bytes charges the code it reads, and what a seeding prediction of it charges the tiles and scopes every completion of a sketch declares (Sketch_families), so the two round and count alike.
val workgroup_memory_bytes :
capabilities:Backend_intf.codegen_capabilities ->
Low_level.optimized ->
Base.intThe static workgroup-shared bytes the kernel will declare, charged by kernel_workgroup_bytes (Metal rounds every threadgroup array up to 4 bytes, gh-ocannl-1330): its workgroup_shared tiles (a pipelined tile, one array of rotating copies, counted once per copy) plus, once for the kernel, the rendering backend's Backend_intf.codegen_capabilities.mma_kernel_workgroup_bytes over the storage triples of its tile-MMA emission scopes — Metal declares one coordinate table per kernel when any of them converts its destination boundary (gh-ocannl-1205, gh-ocannl-1255). A scope is one Tile_mma statement, each copy listed where a loop transform duplicated it; a simdgroup-fragment scope is rendered only around exactly one Tile_mma, so it is the same list. Plus, per lane all-reduce site that stages per-simdgroup partials (Low_level.lane_all_reduce_sites' lar_cross_simdgroup, gh-ocannl-1168), the backend's Backend_intf.codegen_capabilities.lane_all_reduce_workgroup_bytes for its extent at the larger of its local's compute and accumulator precisions. Plus, per warp-shuffle reduction into a cell (Low_level.warp_reduce_cell_sites, gh-ocannl-1452), the backend's Backend_intf.codegen_capabilities.warp_reduce_workgroup_bytes for its extent at the accumulator residency of the cell's storage precision: the per-warp partials its multi-warp form declares. Pass the compiling context's Context.codegen_capabilities; the capability-free Backend_intf.no_codegen_capabilities counts the tiles alone. Forces no node's size_in_bytes (Tnode.size_in_bytes_peek).
val thread_private_bytes : Low_level.optimized -> Base.intThe per-work-item scratch bytes the kernel will declare (gh-ocannl-1313): one function-scope array per node of its traced store that settles Local (Tnode.Placements.is_local_peek), sized at its storage precision, workgroup-shared tiles excepted. Settles no placement and forces no node's size_in_bytes (Tnode.size_in_bytes_peek). An upper estimate of what codegen declares — the CPU backends' per-chunk privatization and Metal's simdgroup fragments are rendering decisions that can drop arrays from the function scope — and of the device compiler's frame, which may drop an array it proves dead; checked by check_hardware_limits_classified against Backend_intf.hardware_limits' max_thread_private_bytes.
val check_hardware_limits :
name:Base.string ->
limits:Backend_intf.hardware_limits ->
capabilities:Backend_intf.codegen_capabilities ->
Low_level.optimized ->
Base.unitValidates the scheduled kernel name against the device limits, raising Utils.User_error on violation: the launch's workgroup size (the product of Low_level.launch_dims' block dimensions) against Backend_intf.hardware_limits.max_threads_per_workgroup, the workgroup_memory_bytes estimate against Backend_intf.hardware_limits.max_workgroup_memory_bytes, the thread_private_bytes estimate against Backend_intf.hardware_limits.max_thread_private_bytes, and then the launch geometry dimension by dimension through the shared launch_geometry_excess — the workgroup's .x/.y/.z extents against Backend_intf.hardware_limits.max_workgroup_dims (a separate hardware fact from the thread product: CUDA caps maxThreadsDim.z at 64 while the product cap is 1024 — gh-ocannl-679), and both 16-bit-capped grid dimensions, the .y extent (grid.(1), a blocktiled matmul's row-block count) and the folded .z extent (grid.(2), the product of the Grid slots >= 2 — gh-ocannl-643), against Backend_intf.hardware_limits.max_grid_yz. grid.(0) is deliberately ungated: it is 2^31-scale wherever hardware axes bind. Backend compile calls this after any ?lowered_transform (or the default annotator, which already respects the limits), turning driver-level launch failures into early, named errors. A no-op for all-None limits.
val with_gate_capabilities :
(Backend_intf.codegen_capabilities -> Backend_intf.codegen_capabilities) ->
f:(Base.unit -> 'a) ->
'aTest seam (gh-ocannl-1330): runs f with check_hardware_limits_classified judging by adjust capabilities instead of the capabilities it is handed, in the calling domain only, restoring the previous adjustment on return or exception. The default is the identity and no configuration selects another. What it is for: reaching a refusal downstream of the gate (the driver's own) with a kernel the exact gate would refuse first — test/operations/metal_threadgroup_overflow removes the per-allocation rounding. Not a production seam: a schedule the gate misjudges is a defect of the estimate.
val check_hardware_limits_classified :
name:Base.string ->
limits:Backend_intf.hardware_limits ->
capabilities:Backend_intf.codegen_capabilities ->
Low_level.optimized ->
Base.unitInternal candidate-facing variant of check_hardware_limits; transports excess thread, workgroup-memory, per-work-item scratch (Schedule_outcome.resource.Thread_scratch at phase Hardware_limits, before any codegen — the post-link HIP budget check reports the same resource at Backend_link), per-workgroup-dimension, .y-grid or folded-.z-grid requests as typed Schedule_outcome.cause.Resource_exceeded causes, one variant per launch dimension (Schedule_outcome.resource.Workgroup_x_extent .. Schedule_outcome.resource.Grid_z_extent) so an autotune search's declines group by the dimension that asked too much.