Module Ir.Low_level

A for-loop-based array language and backend-agnostic optimization

Global references

module Scope_id : sig ... end
type scope_id = Scope_id.t = {
  1. tn : Tnode.t;
  2. scope_id : Base.int;
}
val sexp_of_scope_id : scope_id -> Sexplib0.Sexp.t
val equal_scope_id : scope_id -> scope_id -> Base.bool
val hash_fold_scope_id : Ppx_hash_lib.Std.Hash.state -> scope_id -> Ppx_hash_lib.Std.Hash.state
val hash_scope_id : scope_id -> Ppx_hash_lib.Std.Hash.hash_value
val compare_scope_id : scope_id -> scope_id -> Base.int
val get_scope : Tnode.t -> scope_id

Low-level representation

type axis_type =
  1. | Serial
  2. | Grid
  3. | Workgroup
  4. | Workgroup_reduce
  5. | Unrolled
  6. | Vectorized

How a loop's iterations map to hardware; see docs/proposals/axis-types-for-loops.md. Serial is an ordinary for-loop. Grid / Workgroup bind the loop index to a GPU grid / workgroup (block, threadgroup) hardware index instead of looping; Workgroup_reduce is a Workgroup axis participating in a workgroup-cooperative reduction (see the contract below). Unrolled is emitted as the repeated body with substituted constants. Vectorized renders eligible bodies as explicit SIMD code — elementwise statements via vector extensions / packed loads (gh-ocannl-164 / gh-ocannl-463), a single recognized accumulation as independent accumulator chains with a horizontal reduce at exit on CPU backends (gh-ocannl-468) — and everything else as a serial loop annotated with the backend's vectorization pragmas when it provides them (a plain, un-annotated serial loop for accumulating bodies, whose loop-carried dependency the pragmas would deny) — like the hardware kinds, the annotating pass asserts iteration independence or, for a recognized accumulation, licenses reassociating it. Hardware slots are positional: among a kernel's loops of one kind, the innermost binds .x, then .y, .z. Annotated loops must have from_ = 0 and iterations with no cross-iteration dependencies (Vectorized accumulations again excepted). Workgroup_reduce is the labelled exception; its body must either stage its communication explicitly through workgroup-shared nodes and barriers (rendered by binding the index like Workgroup), or be a single accumulation statement acc = op(acc, contrib) over an associative-commutative op with the accumulator's indices free of the loop index — the renderer then owns the communication: warp/simdgroup shuffles on GPU backends (gh-ocannl-462), the plain serial loop on CPU backends. Like Vectorized, the annotation licenses reassociating the (floating-point) reduction.

val sexp_of_axis_type : axis_type -> Sexplib0.Sexp.t
val axis_type_of_sexp : Sexplib0.Sexp.t -> axis_type
val compare_axis_type : axis_type -> axis_type -> Base.int
val equal_axis_type : axis_type -> axis_type -> Base.bool
val axis_type_label : axis_type -> Base.string

Loop keyword used by the human-readable printers: plain "for" for Serial, "for@<axis>" otherwise.

type scope_mint =
  1. | Inlined_computation
  2. | Schedule_minted

Which pass minted a Local_scope (gh-ocannl-687). The construct has two producers, and a consumer that walks the IR looking for schedulable structure means only one of them:

  • Inlined_computation -- virtualization's inline of a virtual node's computation at a read site, and the rewrites that carry those scopes along. The loops inside such a body are the inlined node's own iteration space, re-instantiated per use site; no Schedule op has ever targeted them.
  • Schedule_minted -- the accumulator localization built by Schedule's materializing Unroll and by Partition (gh-ocannl-639), and by C_syntax.try_localize_serial_reduce: a running value for a MATERIALIZED cell, whose body holds the very per-step / per-segment loops Schedule.rewrite_loop retargets.

This is the fact Autotune.collect_loops needs: it enumerates loops inside Schedule_minted scopes only, so the action menu does not spend its per-unit budget proposing splits, swaps, unrolls and vectorize retypes for inlined interpolation and reduction loops that no schedule op has ever been able to reach.

Deliberately NOT the mechanism behind the scope-target contract input_scope_ids serves: that one asks whether a scope was in the program a given optimize call was HANDED, which is a per-call fact -- a virtualizer-minted scope handed back into a second optimize still carries Inlined_computation and still is not that call's to retract -- and hand-built IR has no honest way to spell "not mine".

val sexp_of_scope_mint : scope_mint -> Sexplib0.Sexp.t
val scope_mint_of_sexp : Sexplib0.Sexp.t -> scope_mint
val compare_scope_mint : scope_mint -> scope_mint -> Base.int
val equal_scope_mint : scope_mint -> scope_mint -> Base.bool
type scan_direction =
  1. | Forward
  2. | Backward

The iteration order of a t.Scan_loop (gh-ocannl-696): Forward runs the index from from_ up to to_, Backward from to_ down to from_. Part of the construct from day one because every adjoint of a forward scan is a backward scan over the same range.

val sexp_of_scan_direction : scan_direction -> Sexplib0.Sexp.t
val scan_direction_of_sexp : Sexplib0.Sexp.t -> scan_direction
val compare_scan_direction : scan_direction -> scan_direction -> Base.int
val equal_scan_direction : scan_direction -> scan_direction -> Base.bool
type t =
  1. | Noop
  2. | Comment of Base.string
  3. | Staged_compilation of Base.unit -> PPrint.document
    (*

    Source text the C-family renderers emit verbatim. Its accesses are invisible to every analysis, and nothing renames inside it: a node it references must be spelled as the dialect emits the node's identifier, its code name behind the backend's Backend_intf.backend_facts.node_ident_prefix (t_<name> on Metal, CUDA and HIP, gh-ocannl-1338).

    *)
  4. | Seq of t * t
  5. | For_loop of {
    1. index : Indexing.symbol;
    2. from_ : Base.int;
    3. to_ : Base.int;
    4. body : t;
    5. axis : axis_type;
    }
  6. | Scan_loop of {
    1. index : Indexing.symbol;
    2. from_ : Base.int;
    3. to_ : Base.int;
    4. direction : scan_direction;
    5. carried : carried Base.list;
    6. body : t;
    }
    (*

    A loop with declared loop-carried scalar state (gh-ocannl-696): the minimal recurrence construct behind cumulative ops, online softmax and top-k. Semantics, for carried = c_1 .. c_n: each c.prev is set to c.init once, before the first iteration; then for every value of index in from_ .. to_ taken in direction, body runs reading the previous iteration's state through Get_local c.prev and producing the next through Set_local c.next; after the body, every c.prev takes its c.next simultaneously (phi-style rotation, so a body may read any old value after any new one is written). A dead range (to_ < from_) is malformed and refused by the validator: a producer with no iterations emits Noop, so no walker needs a dead-scan convention. The final state is not readable after the loop: a body that wants a trajectory or a final value writes it to a tensor node itself.

    Contract, enforced by validate_scan_loops at both ends of the pipeline: c.prev and c.next are ids, pairwise distinct across carried, over one node DECLARED virtual (c.prev.tn == c.next.tn), which names and types the state and is never a buffer; c.next is written exactly once, as a top-level statement of body (not under a guard or a nested loop), and read only by statements after that write; nothing writes c.prev; c.init reads no carried local and does not mention index. The state is scalars only, of unbounded arity -- small fixed extents unroll into it; there is no dynamic indexing into state.

    Placement contract: a virtualization candidate whose captured computation contains a scan -- one written inside the body, one fed by a scan through a scope local, or one merely enclosing a sibling scan -- is refused (Non_virtual 148): a cell's value depends on the whole prefix through state the index does not parameterize, so per-cell recomputation at a read site is unbounded and wrong. Reads inside the body inline as usual (an independent producer replays soundly at any position of the scan). Schedule transforms neither target the scan's own index nor reach the loops inside its body (the loop is opaque to Schedule.find_loops_env / rewrite_loop, so an op naming either symbol declines with the usual no-such-loop Invalid_argument); enclosing loops keep their full menu, since the state is per-iteration-of-the-enclosing-loop local scratch. The index is an ordinary affine loop symbol for footprint purposes (loop_ranges, affine_accesses, interval analysis) -- a scan's accesses stay affine and dense even though its values are serial.

    *)
  7. | Zero_out of Tnode.t
  8. | Set of {
    1. tn : Tnode.t;
    2. idcs : Indexing.axis_index Base.array;
    3. llsc : scalar_t;
    4. mutable debug : Base.string;
    }
  9. | Set_dynamic of {
    1. tn : Tnode.t;
    2. idcs : Indexing.axis_index Base.array;
      (*

      Static everywhere except dyn_axis (a Fixed_idx 0 placeholder there).

      *)
    3. dyn_axis : Base.int;
      (*

      Which idcs slot is replaced by dyn_value at codegen time.

      *)
    4. dyn_value : scalar_arg;
      (*

      Integer-valued index spliced into the row-major offset at dyn_axis.

      *)
    5. llsc : scalar_t;
    6. mutable debug : Base.string;
    }
    (*

    A scatter: like Set but the write lands at a runtime row of axis dyn_axis — the write counterpart of scalar_t.Get_dynamic. gh-466: produced only by rewrite_one_hot_reductions (transposed one-hot pattern, the embedding-table gradient); never constructed by Assignments lowering. Schedule analyses must treat this write as statically unknown: loops whose index reaches dyn_value carry a cross-iteration write dependency and must stay serial (the deterministic no-atomics invariant).

    *)
  10. | Set_from_vec of {
    1. tn : Tnode.t;
    2. idcs : Indexing.axis_index Base.array;
    3. length : Base.int;
    4. vec_unop : Ops.vec_unop;
    5. arg : scalar_arg;
    6. mutable debug : Base.string;
    }
  11. | Set_local of scope_id * scalar_t
  12. | Declare_local of {
    1. id : scope_id;
    2. needs_init : Base.bool;
    }
  13. | Workgroup_barrier
    (*

    Workgroup-scoped synchronization (__syncthreads() / threadgroup_barrier). An opaque effectful statement: no CSE, hoisting, or code motion across it. Grid-scoped synchronization is deliberately not representable.

    *)
  14. | If of {
    1. cond : scalar_arg;
    2. body : t;
    }
    (*

    Guarded statement: body executes iff cond is nonzero (renders as if (cond != 0) { body }). Introduced by launch-extent guards on hardware-annotated loops (docs/proposals/axis-types-for-loops.md §2); simplify_llc erases a guard whose condition an interval proves, and simplifies a surviving guard's body under the bounds the condition implies. A conditional write is never a definite write; virtualization treats guarded computations as non-inlineable in v1.

    *)
  15. | Tile_mma of {
    1. d : Tnode.t * Indexing.axis_index Base.array;
      (*

      Accumulator block base.

      *)
    2. a : Tnode.t * Indexing.axis_index Base.array;
    3. b : Tnode.t * Indexing.axis_index Base.array;
    4. ta : Base.bool;
      (*

      a is stored transposed: its tile axes are k, i-major rather than i, k.

      *)
    5. tb : Base.bool;
      (*

      b is stored transposed: its tile axes are j, k-major rather than k, j.

      *)
    6. m : Base.int;
    7. n : Base.int;
    8. k : Base.int;
      (*

      Covered block extents (multiples of the backend's intrinsic tile).

      *)
    9. ldd : Base.int;
    10. lda : Base.int;
    11. ldb : Base.int;
      (*

      Leading-dimension strides in elements, recorded by Schedule.optop.Tensorize: the tnode's minor dim in the plain last-two-axes case, larger when interior batch axes sit between the tile roles (gh-ocannl-528).

      *)
    12. lane : Indexing.symbol;
      (*

      The cooperating Workgroup axis (extent = SIMD width).

      *)
    13. tile : Register_tile.t Base.option;
      (*

      The C-tile geometry the register-tiled CPU rendering must use (gh-ocannl-619), carried from Schedule.optop.Tensorize; None lets the renderer choose. Ignored by the intrinsic (tensor-core) emissions.

      *)
    14. fallback : t;
      (*

      Semantically equivalent scalar micro-kernel over fresh serial symbols.

      *)
    }
    (*

    Cooperative tile multiply-accumulate (docs/proposals/tensorize-mma.md): d[i,j] += Σ_{l<k} a[i,l] * b[l,j] for i < m, j < n, relative to the operands' base index vectors, executed jointly by the threads of the lane axis (tensor cores / simdgroup_matrix). Each operand's tile is a 2-D slice: minor tile axis on the tnode's last axis (stride 1), major tile axis at the recorded leading-dimension stride (with ta/tb the stored layout is the role's transpose and emissions use the hardware transpose flag); the base indices must not mention lane. Constructed by schedule transforms only (Schedule.optop.Tensorize), after the optimization pipeline. Backends without an MMA hook render fallback under an if (lane == 0) guard. Validates like Workgroup_barrier plus a write of d for the coverage rule; see validate_parallel.

    *)

Cases: t -- code, scalar_t -- single number at some precision.

and scalar_t =
  1. | Local_scope of {
    1. id : scope_id;
    2. body : t;
    3. orig_indices : Indexing.axis_index Base.array;
    4. mint : scope_mint;
      (*

      Which pass built this scope; see scope_mint.

      *)
    }
    (*

    An inlined sub-computation whose value is body's final Set_local of id.

    Scope purity (gh-ocannl-584): body's only effect is on the locals it owns — its own id, plus ids Declare_locald lexically within it. Never a tensor node, never a sibling or enclosing scope's local, and no barrier or staged callback. A scope body does not execute where it is written: codegen hoists it ahead of the enclosing statement ordered by scope_id rather than by the operand's syntactic position, simplify_llc collapses a single-assignment scope into the expression (moving its reads the other way), and hoist_cross_statement_cse lifts a body shared by sibling statements out of the statement entirely, to run once ahead of the first user. Purity is what makes all three placements unobservable from outside the body — the same assumption Affine.path_before makes when it declines to order sibling Arg positions. (It governs a body's effects, not its inputs: the hoist additionally needs the body's reads — tensor nodes and scope locals alike — untouched across the statements it is lifted over, which is its own hazard check's obligation.) Enforced by validate_scope_bodies; the optimization pipeline satisfies it by construction.

    Scope target (gh-ocannl-681): id.tn must be VIRTUAL for as long as optimize is looking at the scope, since the scope denotes that node's inlined computation. A scope over a materialized node is REJECTED, never rewritten — the optimizer would have to discard the body and read the buffer, which is what used to happen silently. Mind that a node with no setter is decided non-virtual, so a hand-built scope over a freshly created node is rejected too unless the node is declared virtual. AFTER optimize the shape is legal and means the opposite: Schedule's materializing Unroll and Partition mints and C_syntax.try_localize_serial_reduce localize a materialized accumulator this way, and codegen renders it. Localization is codegen's business alone (gh-ocannl-693); IR already in that form reaches a backend through Context.compile ?prelowered, never through optimize. Only a scope optimize MINTED may be retracted to a Get, and only by itself, when a later refusal materializes the node it had inlined — which happens (gh-ocannl-704): the virtualization walk is in source order, so a statement reached after a read that already minted can materialize the node under the scope.

    *)
  2. | Get_local of scope_id
  3. | Get of Tnode.t * Indexing.axis_index Base.array
  4. | Get_dynamic of {
    1. tn : Tnode.t;
      (*

      The gathered table; treated as a read of tn, like Get.

      *)
    2. idcs : Indexing.axis_index Base.array;
      (*

      Static everywhere except dyn_axis.

      *)
    3. dyn_axis : Base.int;
      (*

      Which idcs slot is replaced by dyn_value at codegen time.

      *)
    4. dyn_value : scalar_arg;
      (*

      Integer-valued index spliced into the row-major offset at dyn_axis. gh-343: produced only by rewrite_one_hot_reductions; never escapes low-level / backend codegen.

      *)
    }
  5. | Get_merge_buffer of Tnode.t * Indexing.axis_index Base.array
  6. | Ternop of Ops.ternop * scalar_arg * scalar_arg * scalar_arg
  7. | Binop of Ops.binop * scalar_arg * scalar_arg
  8. | Unop of Ops.unop * scalar_arg
  9. | Constant of Base.float
  10. | Constant_bits of Base.int64
    (*

    Direct bit representation, primarily for uint4x32

    *)
  11. | Embed_index of Indexing.axis_index
and scalar_arg = scalar_t * Ops.prec

The argument precision is preserved in heterogeneous precision operation arguments, and is ignored (overridden) in homogeneous precision operations.

and carried = {
  1. prev : scope_id;
  2. next : scope_id;
  3. init : scalar_t;
}

One loop-carried scalar of a t.Scan_loop: read as prev, written as next, rotated prev := next after each iteration, prev := init before the first. Both ids name the same virtual node, whose precision is the state's precision.

val sexp_of_t : t -> Sexplib0.Sexp.t
val sexp_of_scalar_t : scalar_t -> Sexplib0.Sexp.t
val sexp_of_scalar_arg : scalar_arg -> Sexplib0.Sexp.t
val sexp_of_carried : carried -> Sexplib0.Sexp.t
val equal : t -> t -> Base.bool
val equal_scalar_t : scalar_t -> scalar_t -> Base.bool
val equal_scalar_arg : scalar_arg -> scalar_arg -> Base.bool
val equal_carried : carried -> carried -> Base.bool
val compare : t -> t -> Base.int
val compare_scalar_t : scalar_t -> scalar_t -> Base.int
val compare_scalar_arg : scalar_arg -> scalar_arg -> Base.int
val compare_carried : carried -> carried -> Base.int
module Access_fold : sig ... end

Ordered read-only Low_level traversal. Every policy decision is mandatory. Statements are visited before children and after their reads; Seq is left-to-right. Tile_mma descends through its equivalent fallback. A Prune skips children AND the after hook. Hooks may prune Tile_mma or Scan_loop when the consumer treats them as opaque.

val direct_write : t -> Tnode.t Base.option

The tensor node a statement writes itself (not through a sub-statement): Set, Set_dynamic, Set_from_vec, Zero_out, and a Tile_mma's accumulator. An exhaustive match: with Access_fold.fold's descent, the one place a new statement constructor must be classified for the write-sensitive walks (gh-ocannl-1451).

type write_demand = {
  1. written : Tnode.t Base.list;
    (*

    direct_write of every statement reached, program order.

    *)
  2. binders : Access_fold.binder Base.list;
    (*

    Every loop and scan binder reached, program order.

    *)
}
val write_demand : t -> write_demand

The one statement-position walk (gh-ocannl-1451): loops, scans, guards and sequences, in source order; never an expression (so never a Local_scope body or a scan's synthetic local assignments) and never a Tile_mma's fallback (the statement counts atomically as its write). Scan bodies are entered (gh-ocannl-1134); a scan binder is reported with axis = None. Dead loops are entered: the walk is structural.

val evaluation_reads : blends:Base.bool -> self:Tnode.t -> scalar_t -> Base.Set.M(Ir.Tnode).t

The one read collector (gh-ocannl-1451): the tensor nodes one evaluation of llsc, the right-hand side of a setter of self, loads -- the transitive inline-fanin guard's contributors (gh-573), shared by trace_node_facts' per-setter setter_reads and the guard's replay of stored computations. self and merge-buffer reads are excluded, a gather counts its table, projection-discarded operands are not evaluated, a Where charges its wider arm (the first on a tie) -- or, with blends (a backend whose blends_value_selects fact holds, gh-ocannl-1775), the union of its arms when speculable_select admits it: the blend evaluates both. Admission is judged at the arms' own annotations and ahead of inlining, an over-approximation of the renderer's (which also needs a float compute precision, and on HIP no NaN-literal arm) that errs toward the union only. And Local_scope bodies -- hoisted, unconditional -- charge the statement whichever arm they sit in. Dead loops replay nothing; a Tile_mma is opaque.

val evaluation_reads_of_proc : blends:Base.bool -> self:Tnode.t -> t -> Base.Set.M(Ir.Tnode).t

evaluation_reads of a statement: one replay of a stored computation of self.

type replay_facts = {
  1. replay_reads : Base.Set.M(Ir.Tnode).t;
    (*

    Every tensor node one replay may load, self excluded — all operands, not the cost collector's wider arm (evaluation_reads_of_proc answers a cost question).

    *)
  2. replay_merge_reads : Base.bool;
    (*

    The replay reads a merge buffer (the inliner's own taint predicate, or a merge read in self's setters).

    *)
  3. replay_free_symbols : Base.Set.M(Ir.Indexing.Symbol).t;
    (*

    Index symbols no loop of the replayed code binds: static (launch) indices of the routine that stored the computation.

    *)
}
val replay_facts : self:Tnode.t -> t -> replay_facts

gh-ocannl-777: what one replay of a stored computation of self reads, for Context.observability. Only the statements the inliner's instantiation keeps are visited — loops (a dead one replays nothing), self's own setters, scope-local assignments — so the sibling setters a shared-loop body carries cannot contribute a read or a symbol.

module Canonical_render : sig ... end

gh-563: the one canonical rendering of lowered code, shared by both digest consumers — analysis_cache_stats' cache (keyed inside optimize) and Schedule_cache.canonicalize (schedule replay across sessions).

val is_pure_index_conjunction : scalar_t -> Base.bool

Recognizes And or eager Mul of two signed-index comparisons (or folded 0/1 constants). This identifies the Boolean product used by range guards for narrowing and partitioning; arbitrary numeric multiplication is opaque. It does not prove that index arithmetic is total: the simplifier separately checks the active domains before replacing And with Mul.

val scalar_precision : scalar_t -> Ops.prec
val apply_op : Ops.op -> scalar_t Base.array -> scalar_t
val flat_lines : t Base.list -> t Base.list
val unflat_lines : t Base.list -> t
val loop_over_dims : Base.int Base.array -> body:(Indexing.axis_index Base.array -> t) -> t
val unroll_dims : Base.int Base.array -> body:(Indexing.axis_index Base.array -> offset:Base.int -> t) -> t
val loop_over_padding_region : dims:Base.int Base.array -> padding:Ops.axis_padding Base.array -> body:(Indexing.axis_index Base.array -> t) -> t

Generate loops that iterate only over the padding margins of a tensor. For dimensions with padding, generates separate loops for left margin, middle (recursing), and right margin, each binding its own fresh symbol (gh-ocannl-1384). The middle region continues recursing to find padding in other dimensions.

val has_accumulation : t -> Base.bool

Whether the tree carries a read-modify-write accumulation: some Set (resp. Set_local) reads its own target — a loop-carried dependency through memory when the written cell does not vary with an enclosing loop. Conservative: Local_scope contents count as reading anything, and Tile_mma and (gh-466) Set_dynamic accumulate by construction. Used by the autotune menu and by codegen fallbacks that must not assert iteration independence (e.g. vectorization pragmas) over an accumulating body (gh-ocannl-468).

val scalar_touches_tn : Tnode.t -> scalar_t -> Base.bool

Whether the scalar reads the node anywhere (gh-ocannl-639) — certifies that an accumulation's contribution is free of the accumulator's node, which licenses holding the accumulator out of memory across the reduction.

val code_touches_tn : Tnode.t -> t -> Base.bool

The statement-body counterpart of scalar_touches_tn: any read or write of the node. Blind to a Staged_compilation callback's accesses: pair it with has_staged_compilation wherever a "no" has to mean "never accessed".

val touched_tns_in_order : t -> Tnode.t Base.list

The nodes code_touches_tn reports, each once, in the order the code first touches them (gh-ocannl-1473). A function of the code alone, independent of node uids and of any hash-table traversal, so two fresh builds of one program list corresponding nodes at the same positions -- the order kernel parameters and local declarations are emitted in. Equally blind to Staged_compilation.

val has_staged_compilation : t -> Base.bool

Whether the code holds a Staged_compilation anywhere, scalar positions included: a callback whose emitted code no analysis can inspect, so it may access any node.

val accum_update_parts : tn:Tnode.t -> idcs:Indexing.axis_index Base.array -> scalar_t -> (Ops.binop * scalar_t) Base.option

The accumulation-update statement shape tn[idcs] = op(tn[idcs], contrib) (or its FMA form) over an associative-commutative op, with contrib free of tn; returns (op, contrib). The single source of truth shared by C_syntax's widened renderings and Schedule.Unroll ~materialize:true's scope-form unrolling (gh-ocannl-639), so "what counts as an accumulation" cannot drift between the schedule transform and the emission.

val subst_accum_read : tn:Tnode.t -> idcs:Indexing.axis_index Base.array -> id:scope_id -> scalar_t -> scalar_t

Retarget an accum_update_parts-shaped update's read of tn[idcs] to the scope local id. Raises on any other shape.

val scalar_reads_scope : id:scope_id -> scalar_t -> Base.bool

Whether the scalar reads the scope local id, descending into nested scope bodies — the scope-local counterpart of scalar_touches_tn. Distinguishes a scope-opening init (a Set_local whose value is free of the local) from a self-referential update.

val accum_local_update_op : id:scope_id -> scalar_t -> Ops.binop Base.option

The reduction operator of a scope-local update in either spelling: the plain accum_local_update_parts form, or virtualization's guarded-read form Where (index-only cond, update, Get_local id) (possibly nested per guard condition). The residency classifier's recognizer (gh-ocannl-663): the guarded form is a reduction for accumulator-width purposes but does NOT decompose into an unguarded (op, contrib), so it is deliberately not part of accum_local_update_parts or of scope_updates_reduce_op's hoist license.

val speculable_select : cond:scalar_t -> arms:scalar_arg Base.list -> Base.bool

Whether a value select — a Where, or a gated operator (Relu_gate, Satur01_gate) — may evaluate every one of its arms whatever cond says, so that a backend can render it as a branchless blend (gh-ocannl-1622; C_syntax's select_blend_syntax hook). Pass the arms at the select's result precision.

The invariant it rests on: lowering emits a read that can be out of range only under a guard of its own, and every such guard either compares embedded indices — gh-504's clamp conditions, virtualization's range and unit-solve conditions, Schedule's tile, fragment and pad masks — or protects a Get_dynamic (gh-343's gather guard, whose condition is a data value). So:

  • cond must not mention an Embed_index (read syntactically: a guard is spliced in place, never routed through a Local_scope); a select over an index comparison stays short-circuiting even when it is a user's mask, which is safe and only forgoes the blend;
  • no arm may contain a Get_dynamic, nor a Div / Mod that renders at an integer precision (which can trap) — judged at the precision C_syntax renders each operand at, not at the operand's own annotation, which a homogeneous operator overrides;
  • a nested conditional operator is walked whole, arms included: evaluated eagerly, its own condition can hold where the enclosing one did not, so everything under it is reachable. (An index guard nested inside stays ?: and still protects its read: its condition says exactly whether the read is in range, whatever the enclosing select says.)

Local_scope bodies are hoisted ahead of the statement and run unconditionally already, so a scope reference is a plain local read. Admitting an instance evaluates an arm the classifiers (Ops.ternop_conditionality, Ops.binop_conditionality) call conditional: their floors stay floors, Affine's a_gated stays a conservative over-approximation, and Cost_model's upper walk, which charges both arms, stays an upper bound (flagged approximate: it reads no backend, and CUDA and Metal still short-circuit). The inline fan-in guard's read collector (evaluation_reads) charges a single arm only where the backend short-circuits: on a backend whose blends_value_selects fact holds it charges an admitted select's arms by their union (gh-ocannl-1775).

val float_valued : scalar_t -> Base.bool

gh-ocannl-1785: whether the expression, evaluated at an INTEGER compute precision, is float-valued -- the one reading of Ops.unop_integer_valuation and its siblings over an expression tree. A real-valued operator is; an operator that follows its operands is when one of its value operands is; a projection is its selected operand; everything else (a leaf, a Local_scope read, a truth value, Mod, the heterogeneous and refused operators) is not. C_syntax renders by it -- computing a float-valued subexpression in float and converting it once, at the store or scope holding it or at a Mod operand -- and the simplifier keeps the scope boundaries and folds that conversion depends on. Meaningless at a float precision, where every value is a float.

val mentions_rng_conversion : scalar_t -> Base.bool

Whether an RNG lane conversion occurs anywhere in the expression (not inside a Local_scope). Such a conversion picks its result type and the random bits it consumes from the precision it renders at, so the expression renders at its target's storage precision and an accumulation carrying it keeps a storage-width accumulator.

val accum_update_widens : tn:Tnode.t -> idcs:Indexing.axis_index Base.array -> scalar_t -> Base.bool

Whether code generation holds this update's accumulator at the backend's accumulator residency: an accum_update_parts accumulation, free of RNG conversions (mentions_rng_conversion), with routine logging off. Any other update — a non-reduction recurrence, a mixed-operator update, an RNG-bearing one — narrows to storage at every step in the serial rendering, and a Schedule.Privatize tile taking it over stays at storage precision too (gh-ocannl-1116).

val accum_base_widens : tn:Tnode.t -> idcs:Indexing.axis_index Base.array -> [ `Update of scalar_t | `Scope of scope_id * t Base.list ] -> Base.bool

For a base peel_accum_nest reached: whether code generation holds it at the accumulator residency — accum_update_widens for a raw update; for a scope-form base, routine logging off and no RNG conversion in its scope local's assignments.

val accum_local_update_parts : id:scope_id -> scalar_t -> (Ops.binop * scalar_t) Base.option

The reduce-shaped update of a scope LOCAL, local = op(local, contrib) (or its FMA form) with contrib free of the local — subst_accum_read's output shape; returns (op, contrib). The SIMD reduction rendering uses it to fold vector chains into a widened accumulator's scope local (gh-ocannl-639), and peel_accum_nest's scope-form validation is built on it.

type thread_storage = [
  1. | `Device
  2. | `Shared
  3. | `Thread
]

How a node's storage is shared across the threads a hardware binding creates: device-resident (every thread of the launch), workgroup-shared (the threads of one block), or thread-private (a per-thread local array, which no binding can race). The tags are spelled as the tile-MMA address space C_syntax.mma_space spells them, so the fragment-operand sites answer the same question by coercing this classifier rather than repeating it.

type thread_slot = [ `Grid | `Workgroup ] * Base.int

One coordinate of thread identity: the hardware dimension a bound loop occupies (hardware_axis_info's kind and slot).

val unseparated_thread_write : active:thread_slot Base.list -> thread:(Indexing.symbol -> thread_slot Base.option) -> deferred:(Indexing.symbol -> Base.bool) -> storage:(Tnode.t -> thread_storage) -> t -> (Tnode.t * Indexing.axis_index Base.array * Base.string) Base.option

The legality of a hardware binding, as a query on the written cell (gh-ocannl-959). active is the launch's thread identity, one coordinate per slot some register-bound loop of extent above one occupies; thread names every bound loop and its slot — the register-bound ones and cc's pool-parallel outermost Grid loops, a per-loop binding; deferred the bound loops whose separation another call judges (a reduce lane left to its own rendering: it still covers its slot). The first store in the kernel to storage the threads share that two threads may own, with the reason: an active slot with no bound loop enclosing the store (every coordinate of the slot executes it), or a cell whose index map does not Affine.separates the enclosing thread symbols over every loop symbol in scope — mentioning a bound axis is not separating it (acc[i + j] under two bound axes, acc[i + k] under a bound i and a serial k). Guards narrow the environment as gh-ocannl-566's simplifier narrows it: If (i < 16) shrinks the range the radix argument sees and If (i == 0) pins i to one thread along that axis alone. A Grid axis is not a thread of workgroup-shared storage (one block, one copy). What is judged is the store's own threads: whether other statements' threads touch the cell across a barrier is dependence analysis, the obligation of whoever mints the pin or the binding, as it was for every store before this rule. Reads play no part: a store every thread performs to one cell is a write-write race whatever the values. A dead level and a false guard execute nothing; a Set_dynamic is judged by its static slots, a Set_from_vec by its aligned run blocks, a Zero_out as a store of every cell, a Tile_mma through its fallback with its own lane excused (gh-ocannl-960). The whole kernel is walked. The renderer asks this at every binding it emits: for the kernel's Grid/Workgroup axes before rendering, and for a Workgroup_reduce lane the warp shuffle cannot own, before falling through to the plain binding.

val has_accumulating_cell : t -> Base.bool

Whether the tree holds a SELF-RECURRENCE: some Set whose value reads the very cell it writes (gh-ocannl-733). This is what makes the localizing peel a live question at a level, and it is the predicate the peel census gates on. Narrower than has_accumulation in two ways, both deliberate: the recurrence is on the CELL rather than on the node, and a Local_scope counts only when its body actually reads that cell — has_accumulation must count every scope conservatively, being the predicate that decides whether iteration independence may be asserted, and censusing on it recorded every non-reduction virtualized scope as a declined reduction site. A dynamic gather may recur unless unequal fixed indices outside its runtime axis prove the cells disjoint; symbolic slots remain potentially aliasing. The selector's reads count too.

type peel_guard_verdict =
  1. | Guard_confined
    (*

    Affine.Confined_to_peel: the guard mentions no enclosing loop symbol.

    *)
  2. | Guard_lane_private
    (*

    Affine.Lane_private_if_separated, ADMITTED: the guard mentions enclosing loop symbols and the accumulated cell was shown to separate them (gh-ocannl-721). Appears only in a report that reached a base — separation is checked there, so nothing before it can earn this.

    *)
  3. | Guard_lane_private_unresolved
    (*

    The same guard in a report that REFUSED: the peel stopped before the base could settle whether the cell separates the enclosing symbols — because the cell shares them (Refused_cell_shared) or because the peel never reached a base at all. Distinct from Guard_lane_private so that a refusing report cannot read as an admitted guard beside its own refusal.

    *)

What peel_accum_nest decided about one If it peeled through (gh-ocannl-733). Two shapes that render identically can earn different verdicts here, which is the whole point of reporting them: a test pinning the rendered FORM cannot tell "the guard was confined to the peeled levels" from "the guard mentioned an enclosing lane and the cell separated it".

val sexp_of_peel_guard_verdict : peel_guard_verdict -> Sexplib0.Sexp.t
val peel_guard_verdict_of_sexp : Sexplib0.Sexp.t -> peel_guard_verdict
val equal_peel_guard_verdict : peel_guard_verdict -> peel_guard_verdict -> Base.bool
val compare_peel_guard_verdict : peel_guard_verdict -> peel_guard_verdict -> Base.int
type peel_refusal =
  1. | Refused_not_a_nest
    (*

    The level's body is neither a single peelable loop, nor a pure-index-guarded If, nor an accumulation base: a sibling statement, a data-dependent guard, a non-accumulating Set.

    *)
  2. | Refused_dead_level
    (*

    A level with to_ < from_, which performs no accesses.

    *)
  3. | Refused_guard_fixed of Base.string
    (*

    Affine.Not_peelable with its reason: the guard's truth is fixed for the whole nest.

    *)
  4. | Refused_cell_varies
    (*

    The accumulation base's cell mentions a peeled level, so it is not one cell across them.

    *)
  5. | Refused_cell_shared
    (*

    An admitted Affine.Lane_private_if_separated guard whose cell does not separate the enclosing symbols, or does not stay inside the node's box over their full ranges.

    *)

Why peel_accum_nest stopped, when it did (gh-ocannl-733).

val sexp_of_peel_refusal : peel_refusal -> Sexplib0.Sexp.t
val peel_refusal_of_sexp : Sexplib0.Sexp.t -> peel_refusal
val equal_peel_refusal : peel_refusal -> peel_refusal -> Base.bool
val compare_peel_refusal : peel_refusal -> peel_refusal -> Base.int
type peel_report = {
  1. levels : Base.int;
    (*

    Loop levels peeled before the outcome.

    *)
  2. guards : peel_guard_verdict Base.list;
    (*

    The verdict of each peeled If, outermost first. A report carrying a refusal never carries an ADMITTED lane-private verdict: those resolve to Guard_lane_private_unresolved, since separation is decided at the base the peel did not reach.

    *)
  3. refusal : peel_refusal Base.option;
    (*

    None exactly when the peel reached an accumulation base, i.e. when peel_accum_nest returned Some.

    *)
}

What one peel_accum_nest call decided, beyond whether it succeeded (gh-ocannl-733). Reported to the optional ~report callback on every call, success or refusal; C_syntax turns it into the per-routine peel census, so that a test can pin which DECISION produced a kernel and not only which form was rendered.

val sexp_of_peel_report : peel_report -> Sexplib0.Sexp.t
val peel_accum_nest : ?extra_level:(Indexing.symbol -> axis_type -> Base.bool) -> ?report:(peel_report -> Base.unit) -> loop_ranges:Affine.Loop_ranges.t -> free_of:Indexing.symbol Base.list -> t -> (Tnode.t * Indexing.axis_index Base.array * [ `Update of scalar_t | `Scope of scope_id * t Base.list ] * Base.string * (t -> t)) Base.option

Peel a single-statement reduction nest down to its accumulation base (gh-ocannl-639). Levels are Serial/Unrolled/Vectorized loops, loops extra_level vouches for (codegen passes a predicate accepting hardware-annotated reduction loops its backend serializes — the schedule mints pass nothing, since wrapping a hardware-annotated loop in a scope at transform time would break the schedule on backends that do bind the hardware dimension), and pure-index-guarded Ifs (the gh-490 If (i < s) shape and its constant-bound sibling — data-dependent guards stay opaque), each containing nothing else; the base is a raw accum_update_parts-shaped update (`Update) or the scope form a previous rewrite minted (`Scope: the scope id and the update statements after the opening init, validated to carry ONE reduction operator — mixed-operator sequences are not reductions and keep their per-iteration narrowing), with the accumulated cell invariant across the peeled levels (free_of seeds the invariance check with the caller's own loop). rebuild re-wraps a replacement base statement in the peeled levels. The ONE definition shared by C_syntax's localizing serial fallback and the schedule mints, so transform and emission cannot drift.

A DEAD level (to_ < from_) is never peeled, and that refusal is load-bearing rather than tidiness: the body of a dead loop performs no accesses at all — the routine-interface walk propagates liveness as live && to_ >= from_, so a node reached only under one is absent from the parameters and need not be allocated — while every form this peel licenses reads and writes the accumulated cell OUTSIDE the levels, unconditionally. Peeling a dead level would invent accesses the program does not make, possibly on an identifier the interface never declared.

Whether a GUARD may join the peeled levels is Affine.peel_guard's and Affine.separates's answer, not this function's (gh-ocannl-722) — rebuild keeps the guard around the accumulating update only, so the localized form loads and stores outside it, and the engine is where the two hazards that creates are stated: a guard mentioning no peeled symbol is fixed for the whole nest, so the hoist invents both accesses, and one mentioning an enclosing loop symbol selects among that level's iterations, so lanes the guard excludes would write their unchanged local back over the accumulating lane's result. The second is admitted when the accumulated cell SEPARATES those enclosing symbols — each instance then owning a distinct cell makes the invented load/store pair private and idempotent (gh-ocannl-721) — which is why the cell reaches the decision and the peel defers it to the base.

loop_ranges is loop_ranges of the enclosing program, and supplies both the classification and the ranges: a guard symbol in it that is not peeled is an enclosing level's index, while one outside it is a static index parameter or a runtime extent and is harmless — which is what keeps gh-490's runtime-extent guard (Assignments.extent_guard's i < s, whose bound is a static symbol rather than a constant) peelable. Required rather than defaulted, and derived from the program rather than certified by the caller, so that no call site can forget it: the mints of Schedule need it as much as codegen does, since a refused mint there turns segment seams into narrowing points instead of merely declining an optimization.

report, when given, is called EXACTLY ONCE per call with the peel_report of what was decided — how many levels were peeled, which verdict each peeled guard earned, and the refusal where there was one (gh-ocannl-733). The result alone cannot answer that: two nests differing only in whether the accumulated cell mentions the enclosing index peel a different number of levels under a different guard verdict and render the same localized kernel. Codegen passes a reporter that accumulates the per-routine peel census; the schedule mints pass none.

type accum_skip =
  1. | Skip_debug_logging
    (*

    debug_log_from_routines: the per-step Set form is kept so the trace survives (a Local_scope body renders with log_set_locals:false).

    *)
  2. | Skip_dead_level
    (*

    The level being decided is dead (to_ < from_).

    *)
  3. | Skip_accum_pinned
    (*

    The base is an update mentioning an RNG conversion (mentions_rng_conversion): it renders at storage precision, so localizing it would change the draw rather than move it (gh-ocannl-517, gh-ocannl-682).

    *)

Why decide_accum_width declined a base peel_accum_nest reached.

type accum_decision =
  1. | Accum_not_a_nest of peel_refusal
    (*

    peel_accum_nest reached no accumulation base under the level.

    *)
  2. | Accum_base of {
    1. tn : Tnode.t;
    2. idcs : Indexing.axis_index Base.array;
    3. base : [ `Update of scalar_t | `Scope of scope_id * t Base.list ];
    4. debug : Base.string;
    5. rebuild : t -> t;
      (*

      peel_accum_nest's rebuild of the peeled levels.

      *)
    6. report : peel_report Base.option;
      (*

      What the peel reported (always given on this arm).

      *)
    7. pinned : accum_skip Base.option;
      (*

      None: the serial rendering localizes the accumulator at this level, at the backend's accumulator residency. Some: every rendering of the level keeps the storage cell's width, for this reason.

      *)
    }

How the accumulator of one reduction level is held: decide_accum_width's answer.

val decide_accum_width : extra_level:(Indexing.symbol -> axis_type -> Base.bool) -> loop_ranges:Affine.Loop_ranges.t -> index:Indexing.symbol -> from_:Base.int -> to_:Base.int -> t -> accum_decision

The accumulator-width decision of the reduction level binding index over from_..to_ with body t — peel_accum_nest over the body, then the declines the serial rendering applies to the base it reaches (accum_skip). The ONE owner of that decision (gh-ocannl-1151): code generation asks it at every level that can localize (C_syntax.decide_accum_width, whose warp-shuffle, SIMD-grid and localizing-serial renderings all consume the answer, gh-ocannl-754), and Schedule.Privatize asks it of the program it was handed to choose its tile's width. extra_level is the serialized-slot predicate: a hardware-annotated loop the backend renders serially (no hardware index for its slot) is a peelable level like any other — code generation passes its own (C_syntax.serialized_hardware); loop_ranges is loop_ranges of the enclosing program, as for peel_accum_nest. Pure but for the routine-logging flag it reads.

val accum_decision_widens : accum_decision -> Base.bool

Whether the decision holds the accumulator at the backend's residency: a base the level localizes (pinned = None) that accum_base_widens. A scope-form base whose scope local an RNG conversion reaches is localized at storage width — codegen's census types that scope at storage — so it localizes without widening.

val accumulating_writes : tn:Tnode.t -> t -> t Base.list

The accumulating writes of tn in program order — a Set reading its own cell (raw or scope-form, has_accumulating_cell), a Set_dynamic reading its node — through the statement positions serial_accum_widths walks (loops, guards, sequences; never a Tile_mma or a scope body), so the two lists align entry by entry.

type serial_accum = {
  1. sa_level : Indexing.symbol Base.option;
    (*

    The loop at which the serial rendering localizes the write's accumulator, if any.

    *)
  2. sa_widens : Base.bool;
    (*

    Whether it holds the accumulator at the backend's residency (a localized base that accum_decision_widens, or a scope-form write codegen's census types wide).

    *)
}
val serial_accum_widths : serialized:(Indexing.symbol -> axis_type -> Base.bool) -> tn:Tnode.t -> t -> serial_accum Base.list

For each of accumulating_writes of tn, in the same order, what the program's SERIAL rendering does with its accumulator — decide_accum_width asked at exactly the levels C_syntax.pp_ll asks it at, top-down, with serialized as the serialized-slot predicate: Serial, Unrolled and Vectorized levels always, Workgroup / Workgroup_reduce levels where serialized says the backend renders them serially, Grid levels never. The first level whose decision localizes a base of tn answers for that base; a raw write no level localizes narrows per step, a scope-form one is typed at the residency by codegen's census all the same. Schedule.Privatize reads its tile width off this, asked of the program its schedule was applied to (gh-ocannl-1151).

val fused_accumulation : d:Tnode.t -> di:Indexing.axis_index Base.array -> scalar_t -> scalar_t Base.option

gh-ocannl-998: the accumulation into d[di] in the fused form a tile multiply-accumulate computes. An FMA accumulating into d[di] is returned as is and a floating plain add of a product is returned as the FMA; None for an integer accumulation (gh-ocannl-824) or no product accumulated into d[di]. The sketch seeding asks it whether the register tile can render a site (both forms can be tiled), and fuse_tile_fallback applies it.

val fuse_tile_fallback : d:Tnode.t -> t -> t

A Tile_mma fallback with every accumulation into d in its fused_accumulation form: what C_syntax renders a fallback as unless the backend's compile pins contraction off (C_syntax_config.compiler_contraction_off, gh-ocannl-1457). Schedule.Tensorize leaves the fallback in the form the contraction licence left it, so under a contraction-off compile the tile and its fallback compute the serial loop's two roundings.

val zero_initializer_target : t -> Tnode.t Base.option

A whole-node zero, or its unconditional covering positive-zero Expand_zero loop nest.

val zero_seed_candidate : ?allow_workgroup_barriers:Base.bool -> Tnode.t -> t -> (Tnode.t Affine.access * Indexing.symbol Base.list) Base.option

The zero-forwarding proof shared by code generation and Schedule.Privatize: one unconditional same-cell RMW pair covering the whole node. Returns its write access and every loop that repeats a cell; the accepting accumulator scope must contain all those loops. allow_workgroup_barriers is for Privatize, which retains the barriers while replacing the target accesses with private-tile accesses. Opaque effects and Tile_mma are always refused.

val zero_fill_target : t -> Tnode.t Base.option

zero_initializer_target, also admitting the launch-extent guards of a hardware-mapped zero nest: a guard bounding an enclosing loop's own symbol trims that loop's range before the covering query (gh-ocannl-1257).

val seed_first_write : tn:Tnode.t -> t -> t Base.option

seed_first_write ~tn stmt: stmt, the first statement to touch tn after a covering zero of it, rewritten so that its accumulation into tn opens from zero itself (gh-ocannl-1257) — value-identical, given that zero, down to the accumulator's residency. A constant zero where no loop repeats a cell; the localized scope form where code generation would localize, opening at zero or at the cell guarded by the first iteration of the repeating loops outside the scope; otherwise the opening read guarded by the first iteration of the repeating (Serial/Unrolled) loops. A guard is placed only where the loops it ranges over iterate no more times than the node has cells: it costs an evaluation per iteration per cell to save one pass over the cells. None when zero_seed_candidate refuses stmt, when where code generation localizes depends on which hardware levels it serializes, when a needed guard does not pay, or when no form applies.

val elide_covered_zeros : t Base.list -> t Base.option Base.list

elide_covered_zeros segments, over a routine's kernel bodies in launch order: drops every covering zero (zero_fill_target) whose node's next access is a statement seed_first_write accepts, rewriting that statement (gh-ocannl-1257). Nothing opaque may lie between the two. None marks a segment left with no statement. Post-schedule: the rewritten accumulations are no longer in the shapes schedule transforms and sketch detection expect.

Hardware axis analyses

Phase B of docs/proposals/axis-types-for-loops.md. Hardware slot assignment is positional, not stored in the IR: among a kernel's annotated loops of one kind, the innermost binds .x (slot 0), the next .y, then .z; Workgroup and Workgroup_reduce share the block/threadgroup slot space. Grid slots >= 2 all share the hardware .z dimension by folding (gh-ocannl-643): the launch's .z extent is their per-slot maxima multiplied out (launch_dims) and each such loop binds (z / stride) % cap (grid_fold); Workgroup slots stay capped at 3.

type launch_dims = {
  1. grid : Base.int Base.array;
  2. block : Base.int Base.array;
}

Arrays of length 3 (.x, .y, .z); all-1s for all-Serial code.

val sexp_of_launch_dims : launch_dims -> Sexplib0.Sexp.t
val equal_launch_dims : launch_dims -> launch_dims -> Base.bool
type hardware_axis_info = {
  1. ha_index : Indexing.symbol;
  2. ha_kind : [ `Grid | `Workgroup ];
  3. ha_axis : axis_type;
    (*

    The loop's own annotation: Workgroup and Workgroup_reduce share the kind (and the slot space) but not the rendering, and the binding-legality check reads the difference.

    *)
  4. ha_slot : Base.int;
    (*

    Positional: the innermost same-kind loop binds .x = slot 0.

    *)
  5. ha_from_ : Base.int;
  6. ha_extent : Base.int;
    (*

    to_ - from_ + 1.

    *)
}
val hardware_kind_of_axis : axis_type -> [ `Grid | `Workgroup ] Base.option

The launch slot space a loop annotation binds: Grid the blocks, Workgroup and Workgroup_reduce the threads of a block, the rest none (iterated inside one work item).

val hardware_axes : t -> hardware_axis_info Base.list

All hardware-annotated loops in pre-order, with their positional slots.

val scope_hardware_loop : t -> (Indexing.symbol * axis_type * scope_id) Base.option

The first hardware-annotated (Grid, Workgroup, Workgroup_reduce) loop inside a Local_scope body anywhere in llc — a scope body is hoisted out of its position and is in no kernel's hardware-axis table — with the scope whose body holds it. The one predicate behind validate_parallel's rejection and Schedule's refusal of the op that puts such a loop there (gh-ocannl-1467).

type lane_all_reduce_site = {
  1. lar_index : Indexing.symbol;
    (*

    The Workgroup_reduce loop's symbol.

    *)
  2. lar_extent : Base.int;
  3. lar_local : scope_id;
    (*

    The scope local the reduction accumulates into.

    *)
  4. lar_cross_simdgroup : Base.bool;
    (*

    Whether every lane of the workgroup is sure to reach the site, so it may stage per-simdgroup partials between barriers: the loop starts at 0, covers the whole .x dimension, the workgroup is one-dimensional, no workgroup loop encloses it, and every enclosing If condition reads only constants and index symbols no workgroup loop binds (a scope local or a memory read could carry a lane-dependent value, so either disqualifies the site). Decided per loop SYMBOL: copies of one body (a materializing Unroll, a Partition) share their nested symbols, and a copy failing the check disqualifies every copy, since the renderer meets each loop by its symbol.

    *)
}
val lane_all_reduce_sites : reassociable:(Tnode.t -> Base.bool) -> t -> lane_all_reduce_site Base.list

The lane all-reduce sites of a kernel (gh-ocannl-1168), in pre-order: Workgroup_reduce loops whose body (past their own launch guard) is ONE reduce-shaped update of a scope local reassociable admits (Online_softmax.reassociable_local). C_syntax.try_lane_all_reduce renders a site across simdgroups only where lar_cross_simdgroup holds, and Schedule.workgroup_memory_bytes counts the partials of exactly those sites, so the estimate and the kernel read one predicate.

type inserted_collective_form =
  1. | Lane_all_reduce_partials
    (*

    C_syntax.try_lane_all_reduce's cross-simdgroup form: partials of a reassociable scope local between two barriers.

    *)
  2. | Warp_reduce_partials
    (*

    C_syntax.try_warp_reduce's multi-warp form: per-warp partials of an accumulation into a cell, a barrier, the first warp's combine, a barrier.

    *)
  3. | Lane_all_reduce_butterfly
    (*

    C_syntax.try_lane_all_reduce's single-simdgroup form: an ocannl_shfl_xor butterfly, no barrier (gh-ocannl-1334).

    *)
  4. | Warp_reduce_shuffle
    (*

    C_syntax.try_warp_reduce's single-warp form: an ocannl_shfl_xor tree, no barrier (gh-ocannl-1334).

    *)
val sexp_of_inserted_collective_form : inserted_collective_form -> Sexplib0.Sexp.t
val equal_inserted_collective_form : inserted_collective_form -> inserted_collective_form -> Base.bool
val inserts_barrier : inserted_collective_form -> Base.bool

The two _partials forms; the other two insert warp shuffles only.

type inserted_collective = {
  1. ic_index : Indexing.symbol;
    (*

    The Workgroup_reduce loop's symbol.

    *)
  2. ic_extent : Base.int;
  3. ic_form : inserted_collective_form;
  4. ic_uniform : Base.bool;
    (*

    Whether control uniform across the lanes the collective spans reaches every copy of the loop. For a barrier form, the workgroup: no enclosing workgroup loop, and every enclosing If condition built from constants and index symbols no workgroup loop binds (the judgement lar_cross_simdgroup asks for, gh-ocannl-1168). For a shuffle-only form, the warp: the same rule, except that workgroup loops and symbols at slots above 0 are uniform within a warp, since both renderings span the whole .x dimension in whole warps (gh-ocannl-1334).

    *)
}
val inserted_collectives : reassociable:(Tnode.t -> Base.bool) -> warp_size:Base.int -> t -> inserted_collective Base.list

gh-ocannl-1238 / gh-ocannl-1334: the statements of a kernel whose rendering emits a lane collective the IR does not contain -- a workgroup barrier, or warp shuffles -- on a backend whose warp shuffles span warp_size lanes ([] when warp_size <= 0: the C backends render neither). The ONE notion of "this statement renders a lane collective" shared by validate_parallel, which refuses a non-uniform entry before anything renders, and by the renderers: C_syntax.try_warp_reduce renders either form only for a listed, uniform entry, and C_syntax.try_lane_all_reduce renders its butterfly only for a listed entry and its partials by lane_all_reduce_sites, derived from the same walk.

Lane-form entries are listed where the admission holds: the butterfly at exactly one warp under warp-uniform control, the partials where lar_cross_simdgroup holds past one simdgroup. Both admissions ask for uniformity themselves, so the form declines to the serial loop under divergent control and the validator refuses none of its entries. Cell-form entries are every Workgroup_reduce loop whose body (past its own launch guard) is ONE accumulation statement into a cell invariant in the loop -- C_syntax's shared accumulator decision at a single statement, asked of the same peel_accum_nest -- uniform or not: that form has neither a barrier-free nor a shuffle-free rendering. Past one warp they are Warp_reduce_partials, within one Warp_reduce_shuffle. An entry the renderer refuses for other reasons (partial warps, a narrow residency, ...) is listed too, so the list over-approximates only by statements that never render.

type warp_reduce_cell_site = {
  1. wrc_index : Indexing.symbol;
    (*

    The Workgroup_reduce loop's symbol.

    *)
  2. wrc_extent : Base.int;
  3. wrc_cell : Tnode.t;
    (*

    The node whose cell the reduction accumulates into.

    *)
}
val warp_reduce_cell_sites : reassociable:(Tnode.t -> Base.bool) -> t -> warp_reduce_cell_site Base.list

gh-ocannl-1452: the cell-form statements of inserted_collectives, in pre-order and at every extent, with the node each accumulates into: the Workgroup_reduce loops C_syntax.try_warp_reduce renders as a warp-shuffle tree (a duplicated copy is its own site). Past one warp each declares its per-warp partials in workgroup-shared memory, at the accumulator residency of the cell's storage precision; Schedule.workgroup_memory_bytes charges each through the backend's warp_reduce_workgroup_bytes, which is 0 within one warp and wherever the renderer refuses the site.

val launch_dims : t -> launch_dims

Per-slot maximum extents over the kernel's annotated loops; grid.(2) is the product of the per-slot maxima of grid slots >= 2 (the .z fold, see the section comment).

val grid_fold : hardware_axis_info Base.list -> slot:Base.int -> Base.int * Base.int Base.option

The binding arithmetic of a Grid loop at slot >= 2 under the .z fold: (stride, cap) such that the loop binds (z / stride) % cap — stride the product of the per-slot maxima of grid slots in [2, slot), cap = None (no modulo) when no higher grid slot exists in the kernel. (1, None) — the bare .z register — for the common single-slot-2 case.

val scope_purity_violation : t -> Base.string Base.option

gh-ocannl-584: the predicate form of validate_scope_bodies — None when every Local_scope body in llc is pure, else a description of the first violation. hoist_cross_statement_cse uses it to guard its own precondition (it is the one pass that can move an effect out of a scope, so an impure body reaching it would be laundered past the codegen gate); tests use it to assert that a body was left where it was.

val validate_scope_bodies : t -> Base.unit

gh-ocannl-584: enforces the scope-purity contract stated at scalar_t.Local_scope — a scope body's only effect is on the locals it owns (its own scope id, plus ids Declare_locald lexically within it). Raises Invalid_argument on anything else inside a Local_scope body at any nesting depth: a tensor-node write (Set, Set_from_vec, Set_dynamic, Zero_out, Tile_mma), a Set_local of a sibling or enclosing scope's local, a Workgroup_barrier, or a Staged_compilation. Applied at both ends of the pipeline — optimize on the way in (before a pass can launder a violation out of any body) and C_syntax.compile_proc on the way out. The raw analysis entry points analyze_proc / specialize_proc deliberately do not validate, being the probes that must stay conservative on IR they may not trust. The optimization pipeline satisfies the contract by construction.

val validate_scan_loops : Tnode.Placements.t -> t -> Base.unit

gh-ocannl-696: the well-formedness contract of t.Scan_loop -- one state node DECLARED virtual per carried pair, ids pairwise distinct across the list, inits free of carried state and of the scan's own index, each next written exactly once as a top-level body statement and read only by later statements, no write of a prev, no binder or reference of a carried id beyond its scan, no opaque callback inside, and a non-empty range. Raises Invalid_argument naming the scan and the clause. Like validate_scope_bodies it runs at both ends of the pipeline: optimize on the way in, backend codegen on the way out.

val validate_loop_binders : Indexing.static_symbol Base.list -> t -> Base.unit

gh-ocannl-1402: the loop-binder lint. Each loop symbol is bound by at most one loop (For_loop or Scan_loop) in the whole routine -- sibling loops included, not only nested ones -- and never by a loop when it is one of the routine's static indices; every symbol an index mentions (Iterator, Affine, Concat, Embed_index, a Local_scope's call-site indices, a Tile_mma's lane) is bound by an enclosing loop or is a static index. Raises Invalid_argument naming the symbol and the loops involved. Every flat symbol-keyed consumer -- extent tables, Schedule.statement_mappings, the canonical render behind both digest caches -- relies on this, and misreads a violation silently: concat segments (gh-ocannl-765) and padding-margin strips (gh-ocannl-1384) each broke it once. Applied by optimize on the way in, ahead of the analysis cache, so lowering's generators are held to it; hand-built IR must mint a fresh symbol per loop too. Entry-only by design, unlike validate_scope_bodies: past optimize, schedule rewrites emit sibling nests that share symbols (the staging transfer nests), and the renderer and the post-optimize analyses accept that shape.

val slice_alias_violation : t -> Base.string Base.option

gh-ocannl-1485: the predicate form of validate_no_slice_aliases — None when no storage access in the code names a slice-alias view (Tnode.is_alias), else a description of the first one, naming the construct, the view and its parent.

val validate_no_slice_aliases : ?site:Base.string -> t -> Base.unit

gh-ocannl-1485: a slice-alias view owns no buffer, so every access of one must be rewritten to its parent (Assignments.resolve_alias); a path that forgets the rewrite reads or writes nothing, silently (gh-ocannl-1475). Raises Invalid_argument, prefixed with site, when any construct that addresses a node's buffer names a view: Zero_out, Set, Set_dynamic, Set_from_vec, Get, Get_dynamic, Get_merge_buffer (a view has no buffer to transfer) or any of a Tile_mma's three nodes, at any depth. Scope ids name registers and are not checked. Applied at three points: per assignment in Assignments.to_low_level (naming the assignment whose lowering produced the access), by optimize on the way in (covering the rewrite tier and hand-built IR), and in C_syntax.compile_proc on the way out (covering schedule transforms and a ?prelowered routine that skipped optimize).

val operand_precision_violation : t -> Base.string Base.option

gh-ocannl-1709: the predicate form of validate_operand_precisions — None when every homogeneous operand slot's annotation is descriptive, else a description of the first one that is not.

val validate_operand_precisions : ?site:Base.string -> t -> Base.unit

gh-ocannl-1709: a homogeneous operator (every binop but Uint4x32_to_prec_uniform_lane, every unop but Uint4x32_to_prec_uniform1, FMA, Mul3, and Where's two arms) evaluates its operands at its result precision — the destination node's for a Set, the local's for a Set_local, a Where condition's or a guard's own annotation inside one — so the precision a scalar_arg carries there is never semantics. It may still describe: the evaluation precision itself, or the operand's own precision (scalar_precision), which is what apply_op and the assignments lowering write (an f32 sum of bf16 loads annotates its operands bf16). Any other annotation reads as a request to compute at that precision which the renderer would silently drop, so this raises Invalid_argument, prefixed with site, naming the destination node, the annotation, the evaluation and the operand's own precision, and the remedy (to compute at the annotated precision, store to a node or local of that precision). Heterogeneous slots (a Where condition, a guard, a Get_dynamic index, a Set_from_vec argument, the heterogeneous operators' operands) keep their annotation and are not checked against anything. Applied at two points: by optimize on the way in (covering the assignments lowering, the rewrite tier and hand-built IR routed through optimize), and in C_syntax.compile_proc on the way out (gh-ocannl-1739: covering the optimizer's own output, schedule transforms and a ?prelowered routine that skipped optimize).

val describe_operand_precisions : t -> t

gh-ocannl-1739: rewrites each homogeneous-slot annotation validate_operand_precisions would refuse to the slot's evaluation precision, operands before the operators over them (an operator's own precision is read off its first operand's annotation); every other annotation is left as it is. The renderer never reads such an annotation, so this changes no emitted value. Sound only over code whose builder's annotations already passed the check, so that nothing it rewrites is a request: the optimizer applies it to its output (its input passed the entry gate), and Schedule.apply to a schedule's result when the program the schedule started from passes the check -- a ?prelowered routine carrying a request is left for the exit gate to refuse.

val validate_parallel : ?warp_size:Base.int -> ?reassociable:(Tnode.t -> Base.bool) -> ?fork_join:Base.bool -> Tnode.Placements.t -> t -> Base.unit

Backend-independent well-formedness of hardware annotations (axis-types proposal §2); a no-op for all-Serial code. Raises Invalid_argument on structural violations: nonzero from_, more than 3 slots per kind, annotated loops inside Local_scope bodies, barriers under divergent extents or If guards, writes to materialized nodes not nested under annotated loops covering every active (non-unit) hardware dimension — launch dimensions are global to the kernel, so an uncovered dimension executes the write once per hardware index — and whole-node Zero_out of materialized nodes in multi-threaded kernels (nesting never distributes it). With warp_size (and the lane all-reduce's reassociable, Online_softmax.reassociable_local) it also refuses a lane collective the renderer inserts (inserted_collectives) under control flow that is not uniform across the lanes it spans (gh-ocannl-1238 for barriers, gh-ocannl-1334 for warp shuffles); the default warp_size of 0 is a backend rendering no warp shuffles. fork_join (default false, hardware-thread semantics) states the C backends' execution model (gh-ocannl-1458): each pool-parallel Grid loop is one synchronous dispatch joined before the next statement, and code outside every Grid loop runs once, on the calling thread — so a materialized write or whole-node Zero_out nested under no Grid loop is exempt from the coverage rules, while writes under a Grid loop keep them. Does not prove iteration independence: an annotating pass proves it for the annotations it mints, and the renderer asks unseparated_thread_write of every binding it actually emits.

val validate_parallel_classified : warp_size:Base.int -> reassociable:(Tnode.t -> Base.bool) -> fork_join:Base.bool -> Tnode.Placements.t -> t -> Base.unit

Internal backend-facing variant of validate_parallel; transports a validation Invalid_argument as a typed Schedule_outcome.cause.Illegal_schedule. warp_size and reassociable are required here: the callers are the ones about to render, or to rank candidates for a backend, and know its shuffle width.

val guard_annotated_extents : should_guard:([ `Grid | `Workgroup ] -> Base.bool) -> t -> t

Wraps bodies of annotated loops whose extent is below their slot's launch dimension in If (index < extent) guards, for the kinds the backend binds in hardware.

Optimization

type virtualize_settings = {
  1. mutable enable_device_only : Base.bool;
  2. mutable max_visits : Base.int;
    (*

    Per-cell read multiplicity cap for inlining: a node with a cell read more than this many times (as bounded by the affine access relations, gh-554) is never virtualized unless the computation is simple or a one-hot selector producer.

    *)
  3. mutable max_inline_reduction : Base.int;
    (*

    Recompute-cost cap for inlining: a node whose setters have enclosing reduction loops (loops not appearing in the setter's indices) with a trip-count product exceeding this value is never virtualized. Negative values disable the cap.

    *)
  4. mutable max_inline_fanin : Base.int;
    (*

    Transitive fan-in cap for inlining (gh-573): a node whose fully-inlined computation would read more than this many distinct materialized nodes — accumulated through chains of virtual producers, per setter — is materialized instead. Bounds the per-consumer recomputation of accumulation chains (e.g. a transformer's residual stream), which the per-node visit and reduction caps cannot see. Negative values disable the cap.

    *)
  5. mutable inline_scalar_constexprs : Base.bool;
  6. mutable inline_simple_computations : Base.bool;
  7. mutable inline_complex_computations : Base.bool;
  8. mutable footprint_materialization : Base.bool;
    (*

    gh-ocannl-616 (virtualize_footprint_materialization): whether a node one of the caps above would materialize whole may instead be footprint-scoped when every read of it is an affine sub-image read — a fresh routine-private scratch shaped like the reader's iteration box, filled by a prologue instantiating the node's stored template over that box, right after the producer's last write (ahead of the reader for an inherited template) — and the total scratch cell count is below the node's element count. Also gates the Footprint third of the decision vector (Context.decide_footprint) and the consumer-side form for a node an earlier routine committed Virtual.

    *)
}
val virtualize_settings : virtualize_settings
val footprint_namespace : Base.string

gh-ocannl-616: the namespace of the scratch nodes footprint-scoped materialization mints — how a test tells them from the program's own nodes in an optimized record's traced store.

type setter_read_sets = {
  1. reads_branching : Base.Set.M(Ir.Tnode).t;
    (*

    A select charges its wider arm (it short-circuits: CUDA, Metal).

    *)
  2. reads_blended : Base.Set.M(Ir.Tnode).t;
    (*

    A select speculable_select admits charges the union of its arms (cc, HIP).

    *)
}

One setter's fan-in contributors under both renderings of a value select (gh-ocannl-1775): the analysis is cached across backends, so it keeps both and the guard picks by the backend's blends_value_selects fact.

val sexp_of_setter_read_sets : setter_read_sets -> Sexplib0.Sexp.t
type traced_array = {
  1. tn : Tnode.t;
  2. mutable has_assignment : Base.bool;
    (*

    The code contains a Set or Set_from_vec of the node (Zero_out is tracked separately as zeroed_out). Structural replacement (gh-554) for the retired concrete-index tracer's per-cell assignment table; per-cell facts are answered by affine queries over the access relations instead.

    *)
  3. mutable zero_initialized_by_code : Base.bool;
  4. mutable zeroed_out : Base.bool;
  5. mutable read_before_write : Base.bool;
    (*

    The node is read before it is written (i.e. it is recurrent): its entry values are consumed, so it is an input of the routine (input_and_output_nodes) and not eligible for buffer aliasing. For a node that owns a buffer, the verdict is strict of the read-modify-write exemption (gh-ocannl-618): a read at its enclosing statement's write position still consumes the entry value unless a prior definite write covers its cells — the exemption applies to the visit-counting placement heuristics, not to the interface. The strict classification closes over the SETTLED placements (in reconcile_traced_store, since placement and legality decisions after decide_placements can still flip a candidate non-virtual), and a flipped node is promoted On_device — an entry-consuming node must own a persistent buffer, not Local scratch; a node that stays virtual is exempt by construction — it has no interface, and the virtualizer's partial-write producers depend on that freedom.

    *)
  6. mutable read_only : Base.bool;
    (*

    Surprisingly, the notions of read-only and of constant memory mode come apart: small hosted constants are not read-only because they are initialized on devices by being assigned to; and a volatile memory mode is read-only from the devices' perspective.

    *)
  7. mutable is_scalar_constexpr : Base.bool;
    (*

    True only if the tensor node has all axes of dimension 1, is either zeroed-out or assigned before accessed, is assigned at most once, and from an expression involving only constants or tensor nodes that were at the time is_scalar_constexpr.

    *)
  8. mutable is_accessing : Base.bool;
    (*

    False only if the tensor node is built from index embeddings and scalar constant expressions.

    *)
  9. mutable is_complex : Base.bool;
    (*

    True only if the tensor node is built from a genuinely complex scalar computation (one that accesses other non-constexpr computations). Sharing a loop symbol with another tensor does not, by itself, make a node complex (see #134).

    *)
  10. mutable prefers_virtual_one_hot : Base.bool;
    (*

    True when at least one setter for this tensor is a one-hot selector assignment, i.e. a Cmpeq between the embedded range iterator and a loop-variable-free expression. When has_non_one_hot_setter is false this tensor is exempt from the visit-count Never_virtual rule (task-73617488).

    *)
  11. mutable has_non_one_hot_setter : Base.bool;
    (*

    True when at least one setter is NOT a one-hot selector (including Set_from_vec). A tensor with prefers_virtual_one_hot && not has_non_one_hot_setter is the candidate for the one-hot virtualizer exemption.

    *)
  12. mutable is_range_producer : Base.bool;
    (*

    True when at least one Set assigns this tensor from a bare Embed_index scalar, i.e. the tensor is a Range_over_offsets producer. Used by the indirect arm of is_one_hot_selector_assignment to prove that a Get(rtn, [k]) will inline to Embed_index k rather than arbitrary values (task-73617488).

    *)
  13. mutable inline_reduction_extent : Base.int;
    (*

    The largest product of trip counts of loops that enclose one of the node's setters without appearing in its indices (i.e. reduction loops). Inlining the computation replays these loops at every read site; compared against virtualize_settings.max_inline_reduction.

    *)
  14. mutable read_by_other : Base.bool;
    (*

    True when some statement other than the node's own setters reads the node. Unlike the read-multiplicity metric, same-cell reads count, while a setter's own read-modify-write does not. Gates the recompute-cost guard: a node never read in the routine has no inlining cost, so it must stay eligible for virtual dead-code elimination.

    *)
  15. mutable setter_reads : setter_read_sets Base.list;
    (*

    Per setter statement (Set/Set_from_vec), the tensor nodes its right-hand side reads, under each rendering of a value select (setter_read_sets) — including reads inside Local_scope bodies, excluding the node's own read-modify-write self-reads. Analysis fact behind the transitive inline-fanin guard (gh-573), which takes the per-setter maximum (a read of one cell executes one setter's computation).

    *)
  16. mutable inline_fanin : Base.int;
    (*

    The transitive inline fan-in decide_placements computed for this node under the current placements (at least 1); multiplies into the proxy fa_recompute_cost (gh-ocannl-637: the fallback where the cost model's count is not exact).

    *)
}
val sexp_of_traced_array : traced_array -> Sexplib0.Sexp.t
val get_node : (Tnode.t, traced_array) Base.Hashtbl.t -> Tnode.t -> traced_array
val optimize_integer_pow : Base.bool Base.ref
type traced_store = (Tnode.t, traced_array) Base.Hashtbl.t
val sexp_of_traced_store : traced_store -> Sexplib0.Sexp.t
val copy_traced_store : traced_store -> traced_store

A private copy of a traced store: the table AND its traced_array records (whose fields are mutable), so writes through the copy never reach the original or a sibling copy. Hashtbl.copy alone shares the records.

type optimize_ctx = {
  1. computations : (Tnode.t, (Indexing.axis_index Base.array Base.option * t) Base.list) Base.Hashtbl.t;
    (*

    The computations (of the tensor node) are retrieved for optimization just as they are populated, so that the inlined code corresponds precisely to the changes to the arrays that would happen up till that point. Within the code blocks paired with an index tuple, all assignments and accesses must happen via the index tuple; if this is not the case for some assignment, the node cannot be virtual. Currently, we only allow for-loop symbols in assignment indices of virtual nodes.

    A stored computation is a NAMED COMPUTATION, NOT A SNAPSHOT (gh-617, decided as recompute-at-read): inlining evaluates it at the consumption site with whatever its materialized inputs hold at that moment, so a write to one of its leaves between the deferring routine and a consuming read is observed by the splice — whereas the materialized reading of the same program snapshots the leaf at the deferring routine's execution point. The recompute-vs-materialize semantics is deliberately not fixed at this level; users select the reading via the memory-mode intent and via routine boundaries (routine execution is manual). See "Recompute-at-read" in docs/lowering_and_inlining.md.

    *)
  2. placements : Tnode.Placements.t;
    (*

    Per-compilation-lineage memory-mode resolution (docs/proposals/context-scoped-memory-modes.md): the pipeline's placement decisions (Virtual / Local / On_device) land here, seeded by and never written back to the tnodes' declared intent (Tnode.t.memory_mode_intent).

    *)
  3. alias_candidates : Base.Hash_set.M(Ir.Tnode).t;
    (*

    gh-ocannl-489 liveness-based buffer aliasing: nodes the memory planner may place at overlapping byte ranges within the routine's working pool (decided per compile, before codegen). Codegen must not emit the restrict qualifier for these parameters — whether a candidate pair actually shares bytes is settled only at link time, and an aliased restrict pair is a miscompile.

    *)
  4. inline_preferences : Base.Hash_set.M(Ir.Tnode).t;
    (*

    gh-555: the Inline half of the per-lineage inlining decision vector. A node recorded here is exempt from the heuristic virtualization caps (virtualize_max_visits, virtualize_max_inline_reduction, virtualize_max_inline_fanin) — the caps are priors of the default decision policy, not legality; the legality rejections and observability pessimizations still apply. The Materialize half of the vector is a pre-seeded On_device decision in placements (see Context.decide_materialized / Context.decide_inline).

    *)
  5. footprint_preferences : Base.Hash_set.M(Ir.Tnode).t;
    (*

    gh-ocannl-616: the Footprint third of the decision vector (Context.decide_footprint). A node recorded here is exempt from the heuristic caps like an inline preference, and its reads are served from footprint-scoped scratch where every one of them is footprintable; a read the virtualizer cannot serve that way falls back to inlining, and the legality rejections still apply. Exclusive with inline_preferences per node: record either through prefer_inline / prefer_footprint, which withdraw the other, so the later request wins. Honored only under virtualize_footprint_materialization.

    *)
}
val sexp_of_optimize_ctx : optimize_ctx -> Sexplib0.Sexp.t
val empty_optimize_ctx : Base.unit -> optimize_ctx
val copy_optimize_ctx : optimize_ctx -> optimize_ctx

A shallow-copy fork of the lineage state (computations and placements tables): the copy sees everything decided so far; its later mutations are invisible to the original and to sibling copies. Backend compile forks the incoming context's optimize_ctx through this, so sibling candidate compiles from one frontier are hermetic.

The placement provenances this module reads back

The tags some code interrogates are Ir.Tnode.provenance constructors, so both functions below are exhaustive matches rather than comparisons that can silently stop matching; the one-site tags this module mints are Site literals at their raise sites.

val is_cap_provenance : Tnode.provenance -> Base.bool

Whether the tag is one of decide_placements' three heuristic caps: a flippable policy prior rather than a legality or observability verdict.

val cap_provenance_setting : Tnode.provenance -> Base.string Base.option

The configuration key whose cap forced the decision, when a cap did — the setting to raise to try to keep the node virtual instead of materializing it.

val decide_materialized : ?provenance:Tnode.provenance -> optimize_ctx -> Tnode.t Base.list -> Base.unit

Records an On_device decision for each node this lineage has not already resolved otherwise — the "materialize this node" move of the placement lattice. Nodes already resolved to Virtual / Local / Effectively_constant keep their resolution: decisions are final within a lineage.

Context.decide_materialized is the context-level form (it forks the lineage first) and is what ordinary Assignments callers want. This raw form serves the paths holding an optimize_ctx directly: the analyze-only entry points, and hand-built optimize calls in tests — for which no context-level form can work, since the ?prelowered seam replaces the context's lineage state with the optimized record's own optimize_ctx.

val prefer_inline : optimize_ctx -> Tnode.t Base.list -> Base.unit

Records the Inline preference (gh-555) for the nodes, withdrawing any Footprint preference they carried: the two are one node's mutually exclusive readings, so the later request wins — a search trying a node's sibling flips one after the other must not accumulate them (gh-ocannl-616).

val prefer_footprint : optimize_ctx -> Tnode.t Base.list -> Base.unit

The Footprint counterpart of prefer_inline.

type swizzle_kind =
  1. | Swizzle_elem
    (*

    Element-granularity XOR: P*C + col renders as P*C + (col lxor (P land (C-1))). Spreads same-column scalar reads of consecutive rows across banks; the flavor the scalar and register-blocktiled staged kernels want.

    *)
  2. | Swizzle_b128
    (*

    16-byte-unit XOR: the column's 16-byte-unit index is XORed with the low bits of the row prefix, leaving the offset within the unit alone. This is the CUTLASS-style layout ldmatrix wants — its 8 per-phase row addresses are 16-byte-aligned, so only a remap that keeps 16-byte units intact can both de-conflict them and stay loadable. Requires the row's byte length to be a multiple of 16 and a power of two in 16-byte units.

    *)

Granularity of the XOR remap applied to a swizzled node's minor axis (gh-ocannl-481 item 3, D1). Both flavors are per-row bijections of the minor axis, so the IR-level semantics are identical; they differ in the unit the XOR permutes and therefore in which access pattern they de-conflict.

val sexp_of_swizzle_kind : swizzle_kind -> Sexplib0.Sexp.t
val swizzle_kind_of_sexp : Sexplib0.Sexp.t -> swizzle_kind
val compare_swizzle_kind : swizzle_kind -> swizzle_kind -> Base.int
val equal_swizzle_kind : swizzle_kind -> swizzle_kind -> Base.bool
type reading = [
  1. | `Materialize
  2. | `Inline
  3. | `Footprint
]

gh-ocannl-616/1017: the three placements of a node in the inlining lattice — materialized (a full-domain buffer), inlined (recomputed at every read) and footprint-scoped (a routine-private scratch over the readers' footprint, filled right after the producer's last write). A flip moves a node from its default reading to another one.

val sexp_of_reading : reading -> Sexplib0.Sexp.t
val reading_of_sexp : Sexplib0.Sexp.t -> reading
val __reading_of_sexp__ : Sexplib0.Sexp.t -> reading
val compare_reading : reading -> reading -> Base.int
val equal_reading : reading -> reading -> Base.bool
val reading_to_string : reading -> Base.string

The reading's lowercase name — "materialize", "inline" or "footprint" — for logs, arm labels and test output.

val partition_readings : (Tnode.t * reading) Base.list -> Tnode.t Base.list * Tnode.t Base.list * Tnode.t Base.list

Splits decided readings into the (materialized, inline, footprint) lists that Context.lowered_for_decisions takes, keeping the input order within each list.

type flip_alternative = {
  1. fa_flip : reading;
    (*

    The reading the flip moves the node to.

    *)
  2. fa_recompute_cost : Base.int;
    (*

    The recompute cost of the recompute reading the flip involves: the modeled op count of one instantiation (recompute_pricer, gh-ocannl-637) when fa_modeled, otherwise the traced proxy (reduction extent × transitive inline fan-in), times the number of instantiations that reading performs — the per-cell read multiplicity for an inlined reading, the scratch cell count for a footprint-scoped one (gh-ocannl-616).

    *)
  3. fa_modeled : Base.bool;
  4. fa_refused : Base.string Base.option;
    (*

    gh-ocannl-1093: the virtualizer's rejection code where the world recompute_pricer is handed already shows the flip refused — the virtualizer's own walk refusing to store the node's computation (e.g. 9:sibling-escaping-read-index, a scalar reduction captured at its setter), which refuses an `Inline and a `Footprint flip alike, or a read site the inliner cannot serve (e.g. 13:call-site-index-mismatch), which refuses the `Inline flip only. A refused flip replays to the materialized placement, so it stays on the surface to be inspected but is not a decision: Autotune.rank_flip_candidates and the memory-budget planner exclude it, and its fa_recompute_cost is the traced proxy of a reading that cannot happen. None is not a legality verdict — a flip's legality is settled only when the virtualizer replays — and a `Materialize flip is never refused.

    *)
}

One reading a search can flip a flip_candidate's node to.

val sexp_of_flip_alternative : flip_alternative -> Sexplib0.Sexp.t
type flip_candidate = {
  1. fc_tn : Tnode.t;
  2. fc_default : reading;
    (*

    The reading the default policy decided for the node: `Inline for a node it left virtual, `Footprint for one it footprint-scoped (by this routine, over this routine's own producer or over a template an earlier routine of the lineage left virtual), `Materialize for one a heuristic cap materialized.

    *)
  3. fc_alternatives : flip_alternative Base.list;
    (*

    The readings a search can flip the node to: non-empty, pairwise distinct, none equal to fc_default, most expensive first. Per default reading: `Inline offers `Materialize; `Footprint offers `Materialize and `Inline — `Inline alone when no routine of this specialization writes the node, since nothing would write the buffer; `Materialize offers `Inline, plus `Footprint when every read is footprintable onto a footprint strictly smaller than the node.

    *)
}

gh-555, one record per node since gh-ocannl-1017: one searchable inlining decision dimension of a compile — a node whose placement the default policy decided, with the reading it decided and the readings a search can flip it to. The alternatives are the node's mutually exclusive readings, so a consumer decides a candidate as a whole — keep the default, or take exactly one alternative — and a top-N cut counts candidates. The default is a field rather than an inference from which alternatives exist: an inherited footprint-scoped node and a cap-materialized node with no smaller footprint both offer `Inline alone, from opposite defaults.

Applying a flip: `Materialize via Context.decide_materialized, `Inline via Context.decide_inline, `Footprint via Context.decide_footprint. An `Inline or `Footprint flip's legality is settled only when the virtualizer replays: a rejected flip reproduces the materialized placement.

val sexp_of_flip_candidate : flip_candidate -> Sexplib0.Sexp.t
val recompute_pricer : (static_indices:Indexing.static_symbol Base.list -> (Tnode.t -> (Tnode.Placements.t * (Indexing.axis_index Base.array Base.option * t) Base.list, Base.string) Base.Result.t) -> Tnode.t -> Base.int Base.option) Base.ref

gh-ocannl-637, gh-ocannl-1011: the seam through which the cost model prices flip_candidates' alternatives — given the routine's static indices and, per node, the WORLD a read of it is instantiated in (placements and the node's computations), the exact op count of ONE read of a candidate (None when the model's count is only a bound, or the inliner cannot serve the read). specialize_proc supplies the world: a node's stored computations under the placements as the routine's walk left them (before cleanup commits its own reading), or, for a node a heuristic cap materialized before the walk, the computations the walk itself stores when it is not materialized — the walk re-run over the node's raw setter statements in a scratch copy of the lineage — with the placements that re-run left. Every reading prices that one instantiation, replayed by the inliner at a synthetic read (instantiate_at_synthetic_read), so the guards it emits, its binding choices and the pipeline's rewrites are in the priced code. Exact for the most general in-bounds reader; what a particular reader folds away (a constant substituted for an index, a sub-image collapsing a loop) or shares (sibling readers whose instantiations hoist_cross_statement_cse merges) only lowers what executes, so the product with the per-cell read multiplicity is a bound in the same sense the traced proxy's is. Cost_model registers it at module initialization; the default prices nothing, so every candidate carries the traced proxy. A pricer must be pure and consume only the candidate node's world. specialize_proc resolves and validates each node's private world once per compile, also deciding fa_refused. Its scalar price is memoized across specializations sharing an analysis, keyed by the resolved computations, effective placements and node facts, and current static domains. Replacing the registered function invalidates its prices. The memo retains no world tables.

val post_virtualization_pipeline : Tnode.Placements.t -> traced_store -> input_scopes:Base.Set.M(Scope_id).t -> static_indices:Indexing.static_symbol Base.list -> t -> t

gh-ocannl-1011: the passes virtualized code receives in specialize_proc, in order — cleanup_virtual_llc (placement commitments; a scope over a node committed non-virtual after the scope was minted retracts to a read), simplify_llc under the static indices' interval environment, the hosted-constant link-time conversion, rewrite_one_hot_reductions, eliminate_common_subexpressions and hoist_cross_statement_cse. One function, so the routine and the flip-candidate pricer's synthetic instantiation cannot receive different pipelines. Commits placements in the given table.

val instantiate_at_synthetic_read : placements:Tnode.Placements.t -> static_indices:Indexing.static_symbol Base.list -> Tnode.t -> (Indexing.axis_index Base.array Base.option * t) Base.list -> (t, Base.string) Base.Result.t

gh-ocannl-1011: the code ONE read of the node executes — the given computations (a stored template list, newest first) instantiated by inline_computation's own core at a synthetic read site, wrapped as the scope a reader's statement carries (Set_local of it), and put through post_virtualization_pipeline, all on a copy of placements in which the node is virtual. The read site is the most general in-bounds reader: a fresh symbol per axis, ranging over the axis, except where every value-carrying computation writes the same position free of non-static symbols (a fixed slice, a static index — the only position such a component can serve). Side-effect free: nothing is committed, and the symbols and scope ids it mints are negative, never from the counters generated code is numbered by. Error carries the inliner's rejection code (or the pipeline's message) where the read cannot be served.

type pipelined_tile = {
  1. pt_depth : Base.int;
  2. pt_rotor : Indexing.symbol;
}

gh-487: a software-pipelined (double-buffered) staged tile — codegen allocates pt_depth rotating copies of the tile and renders every access with a buffer-selection term rotated by the pt_rotor loop counter: reads select copy rotor mod depth, writes copy (rotor + 1) mod depth (the schedule emits the loads one iteration ahead), and writes outside the rotor loop (the prologue load) select copy 0. The IR keeps the tile's single-copy dims and indices — the rotation is a physical-layout choice like swizzle_kind, invisible to IR-level semantics — so the pipelined rendering is bitwise identical to the unpipelined one.

val sexp_of_pipelined_tile : pipelined_tile -> Sexplib0.Sexp.t
type optimized = {
  1. traced_store : traced_store;
  2. optimize_ctx : optimize_ctx;
  3. llc : t;
  4. merge_node : Tnode.t Base.option;
  5. workgroup_shared : Base.Set.M(Ir.Tnode).t;
    (*

    Local-memory-mode nodes to be placed in workgroup-shared memory (__shared__ / threadgroup) instead of kernel-local arrays. Populated by schedule transforms; empty for unscheduled code. See docs/proposals/axis-types-for-loops.md.

    *)
  6. simdgroup_fragments : Base.Set.M(Ir.Tnode).t;
    (*

    Local-memory-mode accumulator tiles whose init-load, serial reduction and store-back form one per-simdgroup fragment lifetime. Backends without a fragment rendering ignore the marking and use the ordinary local-array code; Metal maps the marked region to a persistent simdgroup_matrix array.

    *)
  7. masked_fragments : Base.Set.M(Ir.Tnode).t;
    (*

    The pad-masked accumulator tiles Schedule.Tensorize's contraction minted (gh-ocannl-485): the tile whose guarded transfers round-trip the valid region of a padded block through the accumulator. Such a tile is kept out of simdgroup_fragments (renderers read that set as persistent simdgroup matrices, which the guarded transfers are not), so this set is what identifies it -- Schedule.Fuse_epilogue recognizes a fragment store-back by membership here or in simdgroup_fragments, never by the node's label (gh-ocannl-1721). On a staged GPU pipeline's loop-site contraction the tile is a workgroup-shared padded scratch tile staging an unmasked fragment, which is in simdgroup_fragments (gh-ocannl-1736). Read by scheduling only: no renderer consults it.

    *)
  8. swizzled : swizzle_kind Base.Map.M(Ir.Tnode).t;
    (*

    Nodes stored in an XOR-swizzled layout (docs/proposals/tensorize-mma.md, "Swizzled staging"), keyed by the remap's granularity (swizzle_kind): codegen remaps every element access flat = P*C + col (with C the minor dim, P the linearized prefix) to a per-row permutation of col — a bijection on the buffer, so the IR-level semantics are unchanged; only the physical layout differs, spreading same-column accesses across shared-memory banks. Populated by Schedule.Stage ~swizzle. Renderings that assume a row-major layout must decline swizzled nodes; the tile-MMA intrinsic arms decline Swizzle_elem and may consume Swizzle_b128 through ldmatrix-style loads.

    *)
  9. pipelined : pipelined_tile Base.Map.M(Ir.Tnode).t;
    (*

    gh-487: workgroup-shared staged tiles rendered as pt_depth rotating buffer copies (see pipelined_tile). Populated by Schedule.Stage ~pipeline_depth with depth > 1; codegen multiplies the tile's allocation by the depth and rotates a buffer-selection offset with the pt_rotor loop counter. Renderings that assume single-copy storage (vectorized/contiguous multi-element accesses) must decline pipelined nodes.

    *)
  10. zero_fringe : Base.Set.M(Ir.Tnode).t;
    (*

    Schedule-minted staged tiles whose whole index space is safe to read: slots outside the staged source region (edge tiles of a non-dividing or padded staging, gh-ocannl-485) hold 0 — the add-reduce accumulation identity — written by the load nest's Where-form edge guards or by the host-side constant packing. Schedule.Tensorize consults this to discharge pad guards on the intrinsic path.

    *)
  11. flip_candidates : flip_candidate Base.list;
    (*

    gh-555: the searchable inlining decision dimensions of this compile, one per node, most expensive first by each node's most expensive alternative, equal costs in program order (the node's first access in the raw lowered code, gh-ocannl-1777; never node id, which follows construction order), as decided at the whole-routine specialization (schedule-transform copies inherit the whole-routine list). Excluded: nodes never assigned or never read, scalar constexprs and pure one-hot selector producers, and nodes placed by legality, intent or observability rather than the heuristic policy.

    *)
  12. spliced_rbw : Base.Set.M(Ir.Tnode).t;
    (*

    gh-610 review round 6: the nodes whose read_before_write was set by the FINAL-code reconciliation (a spliced read preceding, or not definitely covered by, the routine's own writes) — as opposed to the raw analysis' uncovered-read classification, which also flags every pure input. Backends' prior-context demand keys on this set: a reconcile-flipped node's entry value must already live in the linked context, while raw-classified inputs keep the assignments layer's curated exclusions.

    *)
  13. source : t;
    (*

    The raw lowered code this specialization was decided over (gh-ocannl-786): the routine as Assignments.lower handed it to optimize — rewrites applied, no placement decided, every node still present as the statements that set and read it. Schedule transforms inherit it unchanged (they transform the decided program, not the decision problem). Together with the lineage's prior decisions it is the identity of the placement DECISION PROBLEM, which Schedule_cache.canonicalize_source renders: a decision recorded against that identity addresses its nodes by their first occurrence here, where a node the decision inlined is still visible — in llc it is not.

    *)
}
val sexp_of_optimized : optimized -> Sexplib0.Sexp.t
type footprint = {
  1. fp_total : Base.int;
    (*

    fp_working + fp_constants: the number a memory budget is compared to.

    *)
  2. fp_working : Base.int;
    (*

    Bytes of the working (non-constant) pool as the arena planner would lay it out. Equals fp_dedicated when there is no liveness plan (config buffer_aliasing off, code opaque to the liveness fold, or a layout over the per-pool cap).

    *)
  3. fp_constants : Base.int;
    (*

    Bytes of the constant / read-only pool, always bump-packed.

    *)
  4. fp_dedicated : Base.int;
    (*

    What fp_working would be with every node on its own bytes.

    *)
  5. fp_planned : Base.int;
    (*

    How many working nodes carried a live span, i.e. were arena-eligible.

    *)
  6. fp_nodes : Base.int;
    (*

    In-context nodes scored (working + constants).

    *)
}

gh-ocannl-498: the backend-agnostic byte footprint implied by an optimized routine's placement vector. The allocator-side scorer lives in Backends.score_footprint.

val sexp_of_footprint : footprint -> Sexplib0.Sexp.t
val equal_footprint : footprint -> footprint -> Base.bool
val optimize : ?blends_value_selects:Base.bool -> optimize_ctx -> unoptim_ll_source:(PPrint.document -> Base.unit) Base.option -> ll_source:(PPrint.document -> Base.unit) Base.option -> name:Base.string -> Indexing.static_symbol Base.list -> t -> optimized

blends_value_selects is the rendering backend's fact of that name (Backend_intf.backend_facts; Backends.lower_assignments passes it): where true, the inline fan-in guard charges a select speculable_select admits by the union of its arms, which the backend's blend evaluates (gh-ocannl-1775). Default false, the short-circuiting reading: hand-built IR optimized without a backend in mind.

type analysis

Decision-independent analysis of a lowered routine (gh-555 step 1): the structural per-node facts and the lazily-materialized affine access metrics — everything the optimization pipeline consumes that does not depend on the lineage's placement decisions, plus a bounded memo of scalar recompute prices keyed by validated placement worlds (gh-ocannl-1113).

val analyze_proc : Indexing.static_symbol Base.list -> t -> analysis

Compute the analysis once for a routine. optimize is analyze_proc followed by specialize_proc (plus the pretty-printing callbacks).

val specialize_proc : ?blends_value_selects:Base.bool -> optimize_ctx -> analysis -> optimized

The decision-dependent tail of the pipeline: placement decisions (decide_placements under the given lineage's placements and inline preferences, and the backend's select rendering as in optimize: the analysis keeps both charges, so one analysis serves either), virtualization, cleanup, simplification and CSE. Cheap to replay per candidate over one shared analysis (gh-555): sibling calls with hermetic optimize_ctx forks (see copy_optimize_ctx) produce hermetic optimized results — the traced store is record-copied per call.

val validate_virtualization_decision_coverage : Tnode.Placements.t -> t -> Base.unit

The gh-ocannl-805 seam assertion run after virtualization and before cleanup: every tensor-buffer read (Get / Get_dynamic) in a statement cleanup will keep, including reads inside a kept inlined Local_scope body, must name a node with a placement decision. Reads in a virtual candidate's setter are excluded because cleanup drops the whole setter without visiting its right-hand side. Exposed so hand-built-IR tests can inject the otherwise-unreachable negative control directly at this boundary.

val alpha_digest : ?syms:Indexing.symbol Base.list -> ?tn_token:(Tnode.t -> Base.string) -> t -> Base.string

A structural digest of llc that two lowerings of one routine agree on (gh-ocannl-1526): tensor nodes by identity (uid), or by tn_token when given (to rename the nodes a compile mints); loop binders, local-scope ids and free symbols -- each lowering mints its own -- alpha-renamed by first occurrence, syms pre-registered in order (to name the free index symbols of a stored computation's index tuple). Comments are skipped.

val analysis_cache_stats : Base.unit -> Base.int * Base.int

gh-560: (hits, misses) of the process-global analysis cache consulted by optimize: sibling candidate compiles of one routine share one analyze_proc result — keyed by a canonical digest of the raw lowered code and the static indices (tensor nodes and static symbols by identity, loop binders and local-scope ids alpha-renamed) — and replay only specialize_proc. Cumulative counters, for tests and diagnostics.

val clear_analysis_cache : Base.unit -> Base.unit

Drops the analysis cache's entries (the stats persist). Entries retain their routines' lowered code, hence their tensor nodes; the cache clears itself before accessibility snapshots (Tnode.print_accessible_headers) and callers that tear down a session (Tensor.unsafe_reinitialize) clear it to release the nodes promptly. Never needed for correctness: entries keyed by stale nodes cannot alias fresh ones (Tnode.uid is never reused).

val reads_scope_before_set : scope_id -> t -> Base.bool

reads_scope_before_set id body returns true if id is read (via Get_local) before the first definitely-executed Set_local id in body. Use this at code-generation time to decide whether a Local_scope or Declare_local declaration needs a zero initializer.

module Fp_licence : sig ... end

The licences of the rounding-changing algebra over the reals (gh-ocannl-998), one field per licence, each true when that family of rewrites may run:

val fp_algebra_step : Fp_licence.t -> (scalar_t * Ops.prec) -> (scalar_t * Ops.prec) Base.option

One step of the licensed algebra at the root of the expression: the first licensed rewrite that matches (first-match order matters), not yet re-simplified, or None. A ToPowOf is expected with its operands already simplified.

val simplify_llc : Indexing.static_symbol Base.list -> t -> t

Top-down algebraic simplification with interval-driven comparison folding (in particular, it erases If guards whose conditions the loop extents prove). The interval environment is narrowed by every enclosing If condition that is a conjunction of integer-affine index comparisons (gh-ocannl-566), so a guard the statement guard proves folds too — what is simplified under a condition is valid only where that condition holds. Called internally by optimize; exposed for Schedule.apply, whose transforms construct guards after the pipeline's simplify already ran (docs/proposals/schedule-ir-optops.md §2), and for testing. Pure and idempotent.

Exact: every rewrite preserves IEEE results (the x * 1 and x - (+0) identities, constant folding, guard folds); integer arithmetic is reassociated as fp_algebra_step does with no licence. A single-assignment Local_scope is collapsed into its reader only where its rounding to its own precision is not part of the value (gh-ocannl-1788): the slot evaluates at the scope's precision, or the value is a load whose storage precision the scope's widens exactly (Ops.widens_exactly). A leaf it does not rewrite keeps its annotation, so a ?prelowered routine's request on a load still reaches validate_operand_precisions at codegen. The licensed floating-point algebra runs once, on the raw lowering, in the rewrite tier member Fp_algebra (gh-ocannl-998), so neither this pass nor the schedule transforms that rerun it can contract or reassociate across an inlining boundary.

val without_node_bounds : (Base.unit -> 'a) -> 'a

Runs the thunk with interval folding (simplify_llc and the transforms' trailing simplification) blind to tensor nodes' bounds candidates, on the calling domain: a fold then rests on precisions, loop ranges and static indices alone, so the result is a function of the code (gh-ocannl-1480). A bounds candidate is process state that grows as routines link, and the autotuner's replay of a saved schedule must rebuild the very code it was saved against. Less folding is always sound: a guard left in place is only checked at run time.

val rewrite_one_hot_reductions : ?static_indices:Indexing.static_symbol Base.list -> t -> t

gh-343: rewrites the narrow one-hot embedding pattern -- an Add reduction over a loop variable k whose body selects an embedding-table row via k == index_expr (a logical one-hot) -- into a guarded dynamic gather (Get_dynamic) that reads the table row at index_expr directly, with an in-range guard returning 0 out of [0, vocab_size) to preserve the one-hot semantics. The guard is constructed generically and interval analysis (docs/proposals/interval-analysis-scalar-t.md) erases the conjuncts it can prove -- from the index precision's machine range, loop extents seeded from static_indices, and settled per-tensor bounds (Tnode.bounds_state).

gh-466: also rewrites the transposed one-hot pattern -- the embedding-table gradient for k in [0, V): tn[.., k, ..] += (k == index_expr) * g where the loop variable indexes the written tensor itself -- into a guarded dynamic scatter-accumulate (Set_dynamic): if in_range(index_expr): tn[.., index_expr, ..] += g, dropping the O(V) per-position work (llm.c's deterministic encoder backward, docs/research/llmc-lessons.md B5). The enclosing position loops keep their original serial order and the schedule analyses never parallelize over a dynamically-written node, preserving determinism without atomics.

Unmatched or unsupported reductions are left unchanged. Called internally by optimize between simplify_llc and eliminate_common_subexpressions; exposed for testing.

val scalar_mentions_tn : Tnode.t -> scalar_t -> Base.bool

Whether a scalar reads the node, descending into scope bodies and dynamic index values.

val eliminate_common_subexpressions : t -> t

Eliminates common subexpressions within each statement's scalar expression tree. Replaces duplicate Local_scope nodes (structurally identical modulo scope_id) with Get_local references to the first occurrence. Called internally by optimize; exposed for testing.

val hoist_cross_statement_cse : t -> t

Hoists shared Local_scope computations from sibling statements to the enclosing scope. When two or more sibling statements share an alpha-equivalent Local_scope node, the computation is extracted as a Declare_local + body preceding the first user, and all occurrences are replaced with Get_local.

val input_and_output_nodes : optimized -> (Base.Set.M(Ir.Tnode).t * Base.Set.M(Ir.Tnode).t) * Tnode.t Base.option

Inputs are the materialized read-only and read-before-write (within the code) non-constant non-merge nodes. They are inputs in a broad sense, as they could be recurrent nodes or parameters. Outputs are all the materialized nodes written-to by the code. The last returned component is the input merge node, if used in the code. Reads entering the code only through a cross-routine inlined computation count: the traced store is completed from the final optimized code (gh-610).

val loop_ranges : t -> Affine.Loop_ranges.t

Every For_loop and Scan_loop binding within the code, ONE inclusive range per symbol: the hull of all its bindings. Loop symbols are not unique within a routine: validate_loop_binders holds only at optimize's entry, and schedule rewrites emit sibling nests sharing a symbol by design (staging transfer nests, a materializing unroll's copies), with possibly different bounds. A per-binding list would let a first-binding lookup drop the siblings' values, the hole gh-ocannl-1282 went through (gh-ocannl-1305). The box environment for Affine queries, and what peel_accum_nest classifies its guard symbols against: one absent from it is bound outside every loop (a static index parameter, a runtime extent) and cannot select among any level's iterations.

val scope_value_syms : t -> (Base.int, Indexing.symbol Base.list) Base.Hashtbl.t

The value-dependence symbols of statement-level scalar scope-locals, whole-code: per scope id, the union of the symbols its statement-level assignments depend on, transitively through Get_local references (such a local may be assigned in one statement and read in another). Assignments inside Local_scope bodies are not recorded — a scope id is re-instantiated at every use site with per-site loop symbols, and scope-internal flow is covered lexically by the value scans. Consumed by the setter value scans so a value routed through a scope-local is not laundered of its symbols (gh-494 per-thread value-variance).

val scalar_value_syms : locals:(Base.int, Indexing.symbol Base.list) Base.Hashtbl.t -> scalar_t -> Indexing.symbol Base.list

Loop symbols a scalar expression's value depends on syntactically — index symbols of reads, embedded indices, dynamic-index sub-expressions — resolving scope-locals through locals (from scope_value_syms).

val affine_relations : t -> Tnode.t Affine.access Base.list * Tnode.t Affine.statement_effect Base.list

gh-494 waypoint 1: the routine's tensor-node accesses as explicit affine relations (Affine.access), extracted from (typically optimized) code, in program order (a statement's right-hand-side reads precede its write; Local_scope bodies are descended into at their use site; Tile_mma is traversed through its scalar fallback); read accesses carry their scalar gatedness (a_gated). Beside them, from the same walk (gh-ocannl-1016), the Affine.statement_effect rows: everything the code does that no tensor-node access carries — scope-local writes, reads and declarations (a local named by its node and scope id, so a hazard query can match a read to a write), Local_scope bodies, barriers, opaque Staged_compilation (whose accesses are not enumerated: callers needing exhaustiveness check for its row), Tile_mma as a construct, and merge-buffer reads. A "between two points, what else runs?" question is a query over the pair — never a fresh walk over the raw code.

val affine_accesses : t -> Tnode.t Affine.access Base.list

The access half of affine_relations.

val statement_effects : t -> Tnode.t Affine.statement_effect Base.list

The effect half of affine_relations.

val reads_covered_query : ?write_eligible: (read:Tnode.t Affine.access -> write:Tnode.t Affine.access -> Base.bool) -> Indexing.static_symbol Base.list -> Tnode.t Affine.access Base.list -> Tnode.t -> [ `Covered | `Covered_rmw_exempt of Base.string | `Unknown of Base.string ]

The read-before-write fact behind the placement decisions and the routine interface: whether every read of a node in accesses (typically affine_accesses) is covered by prior writes (`Covered), covered only thanks to the read-modify-write exemption — a read at its enclosing statement's own write position (`Covered_rmw_exempt) — or not covered (`Unknown, with a witness). ?write_eligible restricts which writes may cover which reads. Exposed so hand-built IR tests can assert the verdict itself where the pipeline's later passes would mask a wrong one (gh-ocannl-1244).

val buffer_access_spans : stmt_serial:Base.bool -> t Base.list -> (Tnode.t, Base.int * Base.int) Base.Hashtbl.t Base.option

gh-ocannl-489 liveness-based buffer aliasing: per-tnode access span over the final (post-schedule, post-fission) code of a routine, as a closed interval of positions; the input is the routine's kernels in execution order (a singleton when not fissioned). With stmt_serial:true every top-level statement gets its own position — sound only for backends where consecutive top-level statements of one compiled procedure are fully synchronized (the C backends); with stmt_serial:false all statements of a segment share one position, since GPU kernels have no grid-wide synchronization between top-level statements. Returns None when the code contains Staged_compilation (opaque accesses: no aliasing plan can be trusted).

val sink_zero_outs : t -> t

gh-ocannl-489 follow-up: sinks each top-level Zero_out to just before the first later top-level statement accessing the zeroed node (Train.grad_update's up-front zero_grads block otherwise starts every gradient's live span at that block, nesting the backprop chain's intervals and defeating the arena planner). Sound: a Zero_out commutes with statements not accessing the node; it never crosses such an access, a Staged_compilation, or a Workgroup_barrier. Apply to whole-routine code BEFORE scheduling/fission.

Printing

Both dumps render a Constant through Utils.decimal_float_literal: a floating literal that parses back to exactly the double it names, so the dump distinguishes -0.0 from +0.0 and a constant whose 17th significant digit matters from one whose does not (gh-ocannl-713). The dumps are the surface a constant bug is chased on, so normalizing the values away there costs a session before anyone suspects the view. The C-dialect spellings — INFINITY, the hexadecimal literal for an f32 tie — belong to C_syntax.c_float_literal and are deliberately not carried here: a dump is not C.

val code_hum_margin : Base.int Base.ref
val function_header_doc : ?name:Base.string -> ?static_indices:Indexing.static_symbol Base.list -> Base.unit -> PPrint.document
val get_ident_within_code : ?no_dots:Base.bool -> ?blacklist:Base.string Base.list -> ?reserved_prefixes:Base.string Base.list -> ?reserved:(Base.string -> Base.bool) -> t Base.array -> Tnode.t -> Base.string

The code-name minter for the nodes llcs mention. A node whose label ident is in blacklist, starts with one of reserved_prefixes, or satisfies reserved (a namespace given by rule, such as the implementation's __x / _X, gh-ocannl-1338), is treated as a repeating ident: it takes the disambiguated n<id>_<label> form rather than the bare label, so it cannot equal a name the emitter declares. A prefix reserves a namespace whose members are minted on the fly -- the C-family emitters' ocannl_vec<lanes><suffix> typedefs, one per lane count (gh-ocannl-1102) -- which no finite blacklist can list.

val to_doc_cstyle : ?name:Base.string -> ?static_indices:Indexing.static_symbol Base.list -> Base.unit -> t -> PPrint.document

Adheres more to the C syntax, outputs implicit type casts.

val to_doc : ?name:Base.string -> ?static_indices:Indexing.static_symbol Base.list -> Base.unit -> t -> PPrint.document

Adheres to the %cd syntax.

val fragment_transfer : into_fragment:Base.bool -> t -> (Indexing.symbol * Base.int * Tnode.t * Tnode.t * Indexing.axis_index Base.array * Indexing.symbol Base.list) Base.option

One of the accumulator transfers Schedule.contract_tensorized_accumulator synthesizes around a persistent fragment's reduction — lane { if lane == 0 { serial copy loops { fragment <- target } } } for into_fragment, the store back otherwise — as (lane, lane width, fragment, target, target indices, copy symbols innermost-first). The renderer's fragment-scope recognition and the cost model's operand-load count (gh-ocannl-1254) both read regions through it.

A pad-masked contraction's guarded transfers (the init's Where, the store-back's Ifs) are not transfers, by design: the intrinsic loads and stores move whole tiles, which a padded target does not have. A staged GPU pipeline's masked loop-site contraction instead stages them through a padded scratch tile and wraps the reduction in a plain region whose target is that tile (gh-ocannl-1736, Schedule.contract_tensorized_accumulator); elsewhere the masked tile is never marked as a simdgroup fragment. A marked fragment's region that this function does not recognize is refused by Schedule.fragment_scope_violations.

val lane_selected_writes : t -> Tnode.t Base.list

The nodes llc writes under a lane-selecting guard (if lane == 0 of an enclosing Workgroup loop, the form of fragment_transfer's transfers): one hardware thread writing on behalf of its workgroup, whose cells other threads of the workgroup read after a barrier. Such a node carries values between threads within one kernel, which a no-alias kernel parameter rules out (gh-ocannl-1598): C_syntax leaves it unqualified. In first-write order, deduplicated.

val fragment_target_ld : row_sym:Indexing.symbol Base.option -> (Tnode.t * Indexing.axis_index Base.array) -> Base.int

The transfer target's leading-dimension stride in elements: the stride of the axis carrying the row copy symbol (the outermost of fragment_transfer's copy symbols), the minor dim without one.

type fragment_region = {
  1. fr_init : t;
    (*

    The init transfer, fragment <- target.

    *)
  2. fr_reduction : t;
  3. fr_store : t;
    (*

    The store transfer, target <- fragment.

    *)
  4. fr_rest : t Base.list;
    (*

    Statements following the region in the same body.

    *)
  5. fr_fragment : Tnode.t;
  6. fr_target : Tnode.t;
    (*

    The backing target the transfers read and write.

    *)
  7. fr_target_base : Indexing.axis_index Base.array;
    (*

    The target's base indices at the tile origin.

    *)
  8. fr_target_ld : Base.int;
    (*

    fragment_target_ld of the target.

    *)
  9. fr_tile : t;
    (*

    The one Tile_mma accumulating into the fragment in the reduction.

    *)
}
val fragment_region : t -> fragment_region Base.option

The persistent-fragment region Schedule.contract_tensorized_accumulator synthesizes, recognized at the head of a statement sequence: init transfer, reduction, store transfer agreeing on fragment, target, lane, width and indices, with exactly one Tile_mma into the fragment. The renderer's recognition (C_syntax, which also requires the fragment to be marked in optimized.simdgroup_fragments) and the cost model's operand-load count (gh-ocannl-1254) share it.

val mma_operand_layout : swizzle_kind Base.option -> ld:Base.int -> idcs:Indexing.axis_index Base.array -> dims:Base.int Base.array -> [ `Plain | `Swizzled_b128 | `Decline of Base.string ]

A tile-MMA operand's layout as the emission hooks receive it: `Plain unswizzled, `Swizzled_b128 for a b128-swizzled rank-2 tile addressed from its origin with its minor dim as leading dimension; `Decline (with the reason) for any other swizzle, which renders the scalar fallback without reaching a hook. Shared by the renderer and the cost model (gh-ocannl-1254).