Module Ir.Affine

module Idx = Indexing

Affine legality queries over access index vectors (gh-ocannl-494, manifesto §6 waypoints 1-2).

Loop nests are generated from einsum projections, so every access in optimized Low_level.t code carries its affine index map (Indexing.axis_index array) natively. This module is the query engine over those maps: the procedural schedule-legality analyses (the default annotator's per-nest hazard agreement rule, the cross-nest alignment rule, C_syntax.parallel_grid_safe's shared rule, the covering checks) are conservative special cases of the decision procedures here, and are being re-derived as queries — one implementation of "is this access pattern race-free" instead of three.

Scope: linear-integer reasoning over box domains — per-axis linear Diophantine equations with gcd/interval infeasibility and forced-equality derivation via the mixed-radix injectivity criterion (the same criterion as Indexing.affine_injective). No full Presburger machinery: every query form needed so far is decided (or conservatively declined) at this level. Any component the engine cannot interpret (Sub_axis, Concat, dynamic indices) contributes no information, which errs on the side of declining — soundness is preserved by construction.

The pair-conflict query

Can two accesses touch a common cell from different "threads"? Thread identity is a tuple of parallel loop indices; the query works over two copies of the iteration space (the `Left and `Right access's enclosing loops), sharing the symbols that are equal across concurrently executing threads (static indices, loops enclosing the parallel region — a join separates their iterations).

type var =
  1. | Shared of Idx.symbol
  2. | Left_v of Idx.symbol
  3. | Right_v of Idx.symbol
val compare_var : var -> var -> Base.int
val sexp_of_var : var -> Sexplib0.Sexp.t
module Var : sig ... end
type verdict =
  1. | Disjoint
    (*

    The two accesses never touch a common cell at all.

    *)
  2. | Same_thread
    (*

    Common cells occur only when every paired parallel symbol is equal — conflicts are confined to a single thread, where program order applies.

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

    A cross-thread conflict is possible, or the engine cannot rule one out; the payload is the witness/explanation (the axis or symbol pair that failed).

    *)
val sexp_of_verdict : verdict -> Sexplib0.Sexp.t
val equal_verdict : verdict -> verdict -> bool
val gcd : Base__Int.t -> Base__Int.t -> Base__Int.t
val terms_of : range:(Idx.symbol -> (Base__Int.t * int) option) -> dup:(Idx.symbol -> bool) -> side:(Idx.symbol -> var) -> Idx.axis_index -> ((Base.int * var) Base.list * Base.int) Base.option
val equation_of : range:(Idx.symbol -> (Base__Int.t * int) option) -> dup_left:(Idx.symbol -> bool) -> dup_right:(Idx.symbol -> bool) -> Idx.axis_index -> Idx.axis_index -> ((Base.int * var) Base.list * Base.int) Base.option
val range_of_var : range:(Idx.symbol -> 'a) -> var -> 'a
val infeasible : range:(Idx.symbol -> (Base__Int.t * Base__Int.t) option) -> ((Base__Int.t * var) Base.List.t * Base__Int.t) -> Base.bool
val forced_pairs : range:(Idx.symbol -> (Base__Int.t * Base__Int.t) option) -> ((Base__Int.t * var) Base.List.t * int) -> (Idx.symbol * Idx.symbol) Base.list
val axis_index_to_string : Idx.axis_index -> string
val pair_conflict : range:(Idx.symbol -> (Base__Int.t * Base__Int.t) option) -> dup_left:(Idx.symbol -> bool) -> dup_right:(Idx.symbol -> bool) -> pairs:(Idx.symbol * Idx.symbol) Base.list -> left:Idx.axis_index Base.array -> right:Idx.axis_index Base.array -> verdict

pair_conflict ~range ~dup_left ~dup_right ~pairs ~left ~right: verdict on whether the accesses with index vectors left and right (over the same tensor node; rank-padded with Fixed_idx 0) can touch a common cell from different threads. range s gives the inclusive iteration bounds of loop symbol s (None for static/unknown symbols). dup_left/dup_right select the symbols iterated independently by each side (its enclosing loops within the analyzed parallel region); other symbols are shared — equal across concurrently executing threads. pairs is the thread identity: the parallel symbols of the left copy paired with those of the right copy (for same-nest analyses, pairs of the form (p, p)).

Sound and conservative: Disjoint and Same_thread are proven; everything else is Cross_thread.

The covering query

val covers_box : range:(Idx.symbol -> (int * Base__Int.t) option) -> dims:Base.int Base.array -> Idx.axis_index Base.array -> Base.bool

covers_box ~range ~dims idcs: whether the index vector idcs, as its symbols range over their (loop) bounds, enumerates every cell of the dims box exactly once — a bijection onto the box. This is the write-dominance building block: a covering unguarded write rewrites the whole array. Requirements: each symbol used at most once across the vector; per axis, a zero-based full-extent iterator, a mixed-radix affine combination of zero-based symbols whose radix chain exactly composes to the axis dimension, or Fixed_idx 0 on a unit axis. Generalizes (and is checked against) the procedural per-axis rule of C_syntax.first_access_standalone_covering.

val fiber_cardinality : domain:(Idx.symbol * Base.int) Base.list -> Idx.axis_index Base.array -> [ `At_least of Base.int | `Exact of Base.int ]

Counting

fiber_cardinality ~domain idcs: how many points of the loop box domain (symbol, width pairs) map to one given cell in the image of the access map idcs — the per-cell visit count of a read access, and the recompute cost per read site of inlining a setter (the retired Low_level concrete tracer's reduction extent, virtualize_max_inline_reduction's subject). Domain symbols absent from the map contribute the product of their widths; when the map is injective on its mentioned symbols (Indexing.affine_injective) that product is the exact fiber size of every image cell (cells outside the image have zero), otherwise it is a lower bound.

val fiber_cardinality_ub : domain:(Idx.symbol * Base.int) Base.list -> Idx.axis_index Base.array -> [ `At_most of Base.int | `Exact of Base.int ]

Upper-bound companion of fiber_cardinality: at most how many points of the loop box domain map to any single cell of the image of idcs. Domain symbols absent from the map contribute the product of their widths exactly; when the map is injective on its mentioned symbols that is the whole fiber and the bound is exact. Otherwise the mentioned symbols' contribution is bounded per component: the solutions of Σ c_k·s_k = const over a box are pinned once all symbols but one are fixed, so one component admits at most the product of its symbols' widths divided by the largest of them, times the widths of the mentioned symbols it does not constrain; the smallest component-wise bound is taken. Uninterpretable components (Sub_axis, Concat) constrain nothing.

Projection-level predicates

Moved from Indexing (they are queries about the affine LHS map of a projection, this module's home turf): is_surjective decides whether every LHS position is written — used to elide zero-initialization before assignments; is_injective whether no LHS position is written twice — used with is_surjective to elide initialization entirely.

val is_surjective : Idx.projections -> bool
val is_injective : Idx.projections -> bool

Access records

The extraction target for Low_level.affine_accesses (gh-494 waypoint 1): each tensor-node access as an explicit affine relation — the enclosing loop box, the index map into the node's cells, and the program placement. 'tn abstracts the tensor-node type to keep this module below Tnode in the dependency order.

type path_comp =
  1. | Stmt of Base.int
    (*

    Statement index at one Seq nesting level.

    *)
  2. | Arg of Base.int
    (*

    A Local_scope occurrence's per-statement evaluation position: two scope bodies inlined into one statement's scalar tree extend distinct bases, so their interior components never interleave (bare-bodied and Seq-bodied siblings used to collide — a later operand's Stmt sorted before an earlier operand's Rhs). Sibling positions are deliberately incomparable in the visibility rule (path_before): evaluation order among one statement's operands is not modeled, so no cross-operand ordering is claimed.

    *)
  3. | Cond
    (*

    Inside an If statement's condition.

    *)
  4. | Body
    (*

    Inside an If statement's guarded body.

    *)
  5. | Rhs
    (*

    Inside a Set-family statement's right-hand side (or a Set_local's).

    *)
  6. | Write
    (*

    The Set-family or Zero_out statement's own write.

    *)

gh-561: one component of an access's program position (a_path). Stmt components are per-Seq statement indices as before; the other constructors encode intra-statement order, which bare statement positions cannot express — an If body that is not a Seq used to share its condition's path, so a consumer ordering reads against writes by path aliased a guarded body's write with its own condition's read (the gh-554 round-3 bug). Constructor order is execution order, so the derived compare makes lexicographic path comparison program order: a statement's Cond (If condition) evaluates before its Body, and a Set-family statement's Rhs (right-hand side and dynamic-index reads, including Local_scope bodies inlined there) executes before its Write. Every access path ends in Cond, Rhs or Write, and nothing extends a path past Write (writes have no interior), so a write's path is never a proper prefix of a read's.

val sexp_of_path_comp : path_comp -> Sexplib0.Sexp.t
val compare_path_comp : path_comp -> path_comp -> Base.int
val equal_path_comp : path_comp -> path_comp -> Base.bool
val same_statement : path_comp Base.List.t -> path_comp Base.List.t -> bool

Whether two accesses sit in the same Set-family statement — their paths agree above their final component (the write's Write against the statement's own direct Rhs/Cond reads; reads nested deeper, e.g. in a Local_scope body's statements, have longer paths and do not match). This is the path-level counterpart of the a_stmt_write subordination.

val stmt_head : path_comp list -> Base.int

The top-level statement index of a path, -1 when the whole routine is a single statement (its accesses' paths start with an intra-statement component).

type 'tn access = {
  1. a_tn : 'tn;
  2. a_map : Idx.axis_index Base.array;
    (*

    The affine map from the loop box into the node's cells. Empty and standing for every cell when a_whole.

    *)
  3. a_write : Base.bool;
  4. a_dynamic : Base.bool;
    (*

    The effective cell is not statically known (dynamic gather/scatter): the map has a placeholder component, so queries must not interpret it.

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

    A whole-node access (Zero_out).

    *)
  6. a_vec_last : Base.bool;
    (*

    A vectorized write (Set_from_vec): the last map component is the base of a run along the minor axis, not a single cell — queries must treat that component as opaque.

    *)
  7. a_vec_len : Base.int;
    (*

    The run length of a vectorized write along the minor axis; 0 unless a_vec_last.

    *)
  8. a_guarded : Base.bool;
    (*

    Under an If guard: executes conditionally, never a definite write.

    *)
  9. a_rmw : Base.bool;
    (*

    The statement also reads a_tn on its right-hand side (an accumulation): the write carries a reduction dependence — an order-sensitive legality dimension (the determinism contract): a loop carrying only reduction edges may be reassociated (vectorized) under an explicit license, but never parallelized.

    *)
  10. a_val_syms : Idx.symbol Base.list;
    (*

    Writes only: loop symbols the written value depends on syntactically (index symbols of rhs reads, embedded indices, dynamic-index sub-expressions). Direct dependence only — a chain through another node's cells is not tracked.

    *)
  11. a_stmt_write : Idx.axis_index Base.array Base.option;
    (*

    Reads only: the index map of the enclosing Set/Set_from_vec/Set_dynamic statement's write when the read occurs in that statement's right-hand side; None elsewhere (If conditions and Local_scope inner statements carry their own statements' writes). The subject of the read-modify-write exemption (Low_level.rmw_exempt): matching by statement subordination rather than by program path, so a guarded body's write cannot alias its If condition's read (they share a path).

    *)
  12. a_loops : (Idx.symbol * (Base.int * Base.int)) Base.list;
    (*

    Enclosing loops, outermost first, with inclusive iteration bounds.

    *)
  13. a_path : path_comp Base.list;
    (*

    Lexicographic program-order position: statement indices per Seq nesting level, interleaved with intra-statement components (path_comp) at each statement the traversal descends into.

    *)
}
val sexp_of_access : 'tn. ('tn -> Sexplib0.Sexp.t) -> 'tn access -> Sexplib0.Sexp.t
val may_touch_same_cell : ?static_range:(Idx.symbol -> (Base__Int.t * Base__Int.t) option) -> 'tn access -> 'tn access -> Base.bool

Whether two accesses (of the same node) can touch a common cell, each access taken over its whole loop box — the two sides' iterations paired independently, including iterations of loops the sides share (the accesses need not be simultaneous, so a shared loop's symbol varies independently between one side's visit and the other's). Symbols bound by neither side's loops (static indices) are shared parameters, equal on both sides, bounded by static_range when known. Conservative: false only when pair_conflict proves disjointness; uninterpretable access kinds (dynamic, whole-node, vectorized) count as overlapping.

The containment query

read_covered_before ~read ~writes (): is every cell the read access can touch necessarily written before the read executes — the dominance side of dependence analysis, and the fifth decision procedure (gh-494 waypoint 2). Unlike the ∃-flavored pair_conflict (negated to prove disjointness), containment is a ∀∃ query — for every read instance there must exist a covering write instance — so the variable treatment differs: the read's own loop symbols are universal, a write's own loop symbols (below the loops common with the read) are existential, and symbols shared by both sides (common enclosing loops, thread identity under ?thread, static indices) are parameters that must cancel per axis. A residual parameter on the write side declines that write (it would pin the write to one parameter value, or — for a common loop — refer to other iterations of that loop); a residual common-loop or static parameter on the read side is soundly universalized over its range; a residual thread parameter on the read side declines against a thread-bound write (the thread would read a cell another thread wrote), but universalizes against a write not under that thread loop — such a write executes redundantly on every thread, so each thread's copy receives it.

Visibility is same-common-iteration program order: a write is usable when its statement path is lexicographically before the read's, and coverage is proven within the current iteration of the common enclosing loops (the write's whole subtree, including its own inner loops, has then executed). Loop-carried coverage — a read covered only by earlier iterations of a shared loop — is declined, conservatively.

With ?thread naming the parallel (thread-identity) symbols, `Covered proves the cell side of the per-thread-copy transform: the thread reads only cells it wrote itself, earlier in its own serial chunk. For reads covered within the same top-level statement that also proves the VALUE correct — chunks are serially contiguous there, so the thread's own write is the serial last-writer. A cross-statement covering write needs a side condition the caller must supply: other chunks of the writing statement run serially after the reader's own chunk, so the values coincide only when the written value cannot vary across the chunks writing the same cell — every thread symbol feeding the value (a_val_syms) must also pin the written cell (appear in the write's map).

Guarded writes are the caller's choice: include them to mirror guards-taken analyses (Low_level.trace_node_facts and the coverage queries take guards unconditionally), pre-filter a_guarded for execution-accurate coverage. writes must be accesses of the same node as read.

type ap = {
  1. ap_lo : Base.int;
  2. ap_hi : Base.int;
  3. ap_step : Base.int;
}
val sexp_of_ap : ap -> Sexplib0.Sexp.t
val floor_div : Base__Int.t -> Base__Int.t -> Base__Int.t
val ceil_div : Base__Int.t -> Base__Int.t -> Base__Int.t
val ap_of_form : exact:Base.bool -> (Base__Int.t * (Base__Int.t * Base__Int.t)) Base.List.t -> Base.int -> ap Base.option
val ap_subset : ap -> ap -> Base.bool
val ap_covered_chunk : ap -> ap -> (Base.int * Base.int) Base.option
val linear_terms : Idx.axis_index -> ((Base.int * Idx.symbol) Base.list * Base.int) Base.option
val read_covered_before : ?thread:(Idx.symbol -> Base.bool) -> ?static_range:(Idx.symbol -> (Base.int * Base.int) option) -> read:'tn access -> writes:'tn access Base.list -> unit -> [ `Covered | `Unknown of Base.string ]

Crosscheck

Config legality_crosscheck: when enabled, the call sites swapped onto the queries also run the legacy procedural analysis and compare. A query stricter than the procedural answer raises — either a query precision regression or a latent unsoundness of the procedural rule, both needing eyes. A query more permissive than the procedural answer is the expected precision gain, logged to stderr for review.

val crosscheck_enabled : bool lazy_t
val crosscheck : site:string -> context:string -> procedural_safe:(Base.unit -> Base.bool) -> query_safe:Base.Bool.t -> witness:string -> Base.unit