Module Ir.Cost_model

Analytic cost model, the extraction half (gh-ocannl-491 task 1): per-kernel compulsory memory footprints, arithmetic-op counts, arithmetic intensity, and the roofline lower-bound time under advisory envelope constants (Backend_intf.hardware_limits's peak_flops / peak_memory_bandwidth).

The analysis is pure and backend-free: it reads a (typically optimized) Low_level.t through the affine-access artifacts of gh-ocannl-494 (Low_level.affine_accesses, Affine.fiber_cardinality) and returns numbers; envelope constants are plain floats supplied by the caller. The model ranks candidate schedules — it does not predict runtimes.

Approximation contract (each direction is explicit, all biases point the same way):

type node_footprint = {
  1. fp_read_bytes : int;
    (*

    Distinct cells read (including accumulation reads) times byte width.

    *)
  2. fp_write_bytes : int;
    (*

    Distinct cells written times byte width.

    *)
  3. fp_rmw_bytes : int;
    (*

    The subset of fp_write_bytes written by read-modify-write accumulations (a_rmw).

    *)
  4. fp_approx : bool;
    (*

    true when any contributing count is an upper bound rather than exact (see the module contract for the causes).

    *)
}
val sexp_of_node_footprint : node_footprint -> Sexplib0.Sexp.t
type summary = {
  1. per_node : (Tnode.t * node_footprint) list;
    (*

    In first-access program order.

    *)
  2. read_bytes : int;
    (*

    Sum of fp_read_bytes over per_node.

    *)
  3. write_bytes : int;
    (*

    Sum of fp_write_bytes over per_node.

    *)
  4. flops : int;
    (*

    Whole-kernel arithmetic-op count (loop extents times per-statement ops).

    *)
  5. flops_approx : bool;
    (*

    true when guarded (If) code contributed (guards-taken bound), or a short-circuiting form was charged in full while executing only in part — a Where whose arms carry any inline cost (only one arm's expression executes), a gated right operand with nonzero inline cost (gh-ocannl-578). A cost residing in a hoisted Local_scope body executes unconditionally — every renderer emits scope definitions before the statement's expression — so it is charged exactly and never flags (gh-ocannl-637).

    *)
  6. opaque : bool;
    (*

    true when the code contains Staged_compilation or merge-buffer reads: some traffic and ops are invisible to the analysis, so counts may UNDER-estimate.

    *)
}
val sexp_of_summary : summary -> Sexplib0.Sexp.t
val analyze : Low_level.t -> summary

The whole-kernel extraction: footprints from Low_level.affine_accesses image cardinalities, op counts from a loop-nest walk. Input is typically post-optimize code — the analysis charges whatever materialization the code actually performs, so virtual/inlined nodes never appear and recomputation is charged as ops, not bytes.

val total_bytes : summary -> int

read_bytes + write_bytes: the "bytes moved" denominator of arithmetic intensity — an accumulation's cells are charged once in each direction.

val arithmetic_intensity : summary -> float

FLOPs per byte moved, flops / max 1 (total_bytes).

val footprint_approximate : summary -> bool

true when any per-node footprint is an upper bound rather than exact (fp_approx) — the byte counts over-approximate while the op count may still be exact.

type floor = {
  1. fr_flops : int;
  2. fr_bytes : int;
  3. fr_exact : bool;
}

Lower bounds on arithmetic work and compulsory traffic — the dual extraction (gh-ocannl-514 phase 3), for bounding every completion of a partial placement vector where analyze upper-bounds one concrete candidate. fr_exact is false when some flooring-to-zero occurred (a guarded body, a non-exact access image, opaque code): the floor is then sound but not tight.

val sexp_of_floor : floor -> Sexplib0.Sexp.t
val completion_floor : ?open_placement:(Tnode.t -> bool) -> Low_level.t -> floor

The floor of the roofline legs over placement completions, with every approximation biased down — the exact dual of analyze's contract:

  • fr_flops: guarded (If) bodies count zero (guards-never-taken, dual to guards-taken); the short-circuiting forms count only their certain part — Where its condition plus the cheaper arm's inline work (rendered as ?:), And/Or the left operand's (rendered as &&/||), both sides' hoisted scope bodies in either case (they execute regardless, gh-ocannl-637), the Arg1/Arg2 projections only the selected operand (the discarded one is never rendered, hoisted definitions included); opaque code counts zero (an under-count is sound in this direction); Tile_mma keeps the lane-cooperative attribution, exact when the lane binding is in scope. Statements producing an open_placement node — the Set family and Tile_mma with an open accumulator — count zero: an inline completion instantiates the producer only at surviving consumer sites, possibly fewer cells than the setter loop covers, so "recomputation only adds ops" does not hold and the producer's whole effect attributes to the open placement. Call this on the all-materialized specialization of the decision surface, where every open node's work sits in its own producer statement.
  • fr_bytes: per node and direction, the sum of the exact images when they are pairwise provably disjoint (a disjoint union attains its sum — both extractions then agree, gh-ocannl-578), otherwise the largest exact image (a union is at least its largest member — dual to the upper extraction's capped sum; a second nonzero contribution then marks the floor loose); guarded, non-exact, dead-loop-enclosed, conditionally-evaluated (a Where arm's or an And/Or right operand's inline read — not a hoisted scope body's) and open-producer-operand accesses contribute zero — their execution is not certain in every completion. Nodes with open_placement contribute zero.

Committing a placement decision is re-evaluation with the narrowed open_placement: the suppression sets only shrink, so the floor is monotone in refinement — the property that lets a bound prune as commitments accumulate. (An incremental per-node delta cannot be sound in isolation: a Materialize commitment also makes its producer's operations and operand reads certain, not just the node's own traffic.) Committing to Inline tightens nothing for now: a nonzero recompute floor needs lower-bound multiplicity metrics, deferred until the driver proves the need.

A max (fr_flops/peak_flops) (fr_bytes/peak_bandwidth) roofline over these floors (roofline_seconds) lower-bounds every completion — the two legs' minima need not be simultaneously achievable for the max of the two to be sound.

val approximate : summary -> bool

flops_approx || footprint_approximate: some count is an upper bound rather than exact. Such counts are still upper bounds (unlike under opaque), but a candidate whose guards mostly fail can carry counts far above the work it performs — quantities derived from an approximate count (achieved throughput, the roofline bound vs. a measurement) are not evidence about the hardware. Exactness is per leg: op counts (flops_approx) and byte counts (footprint_approximate) go approximate independently.

val roofline_seconds : ?peak_flops:float -> ?peak_memory_bandwidth:float -> flops:int -> bytes:int -> unit -> float option

The roofline lower-bound time: max (flops/peak_flops) (bytes/peak_memory_bandwidth) over the envelope constants present (peak_flops in FLOP/s, peak_memory_bandwidth in bytes/s — the advisory Backend_intf.hardware_limits fields); None when neither is given. Monotone: raising either constant never increases the bound. A lower bound only up to the model's upper-bound byte/op counts — rank with it, do not predict.

The cost of one inlined computation (gh-ocannl-637 Part 2, gh-ocannl-1011)

The cost model's account of what a node costs to recompute at ONE read site — the flops and bytes of one instantiation of its computation, with the exactness the extraction tracks. The reader's multiplicity stays outside, as in the virtualizer's traced proxy (inline_reduction_extent × read multiplicity × inline_fanin); these queries replace the extent × fan-in factor with the modeled arithmetic, and Low_level.recompute_pricer feeds them to the flip-candidate ordering (gh-ocannl-555), the memory-budget planner and, through fa_recompute_cost, footprint-scoped materialization (gh-ocannl-616).

What is priced is code, not a rewrite of it: the inliner's own instantiation of the node's computations at a synthetic read site, after the pipeline the emitted code receives (Low_level.instantiate_at_synthetic_read). The guards the inliner emits (a multi-setter node's per-component range guards and selects, a diagonal producer's consistency guards, a unit-solved affine position's range guard and the loop it keeps), the lane-extract form of a packed-uniform producer, a scope retracted to a read, the simplifier's collapses and the CSE are all in that code, so their cost — and whether each leg is exact — is analyze's verdict on it: a guarded arm with inline work makes the op count a guards-taken bound there as everywhere, a read under a guard's arm makes the byte count one, and a component whose work hoists or whose arm only reads keeps the op count exact.

The exactness contract: a count is exact for the most general in-bounds reader — per axis a fresh index spanning the axis. What a particular reader does on top of that — folding arithmetic once a constant stands for an index, binding an affine position structurally where the generic reader's unit solving keeps a loop, collapsing a loop over a sub-image, or sharing one instantiation between sibling statements through hoist_cross_statement_cse — only LOWERS what executes, so a modeled count times a read multiplicity is an upper bound in the same sense the traced proxy's is; the reader's side of the account stays outside these queries.

type recompute = {
  1. rc_flops : int;
    (*

    Operations of one instantiation.

    *)
  2. rc_bytes : int;
    (*

    Bytes one instantiation reads from nodes other than the computed one — its own cells are a scope local once inlined. Distinct cells within the instantiation's code.

    *)
  3. rc_flops_approx : bool;
    (*

    The op count is an upper bound rather than exact (analyze's flops_approx): a guarded arm with inline work, charged guards-taken. Low_level.recompute_pricer prices only operations, so this flag alone decides whether it returns a count.

    *)
  4. rc_bytes_approx : bool;
    (*

    The byte count is an upper bound (analyze's per-node footprint flags): notably a read under a guard's arm, of which one executes per read. Exactness is per leg, as in summary.

    *)
  5. rc_opaque : bool;
    (*

    Opaque code: the counts may under-estimate.

    *)
}
val sexp_of_recompute : recompute -> Sexplib0.Sexp.t
val instantiation_cost : ?static_indices:Indexing.static_symbol list -> ?placements:Tnode.Placements.t -> self:Tnode.t -> (Indexing.axis_index array option * Low_level.t) list -> recompute option

One read of a node whose computations are given (an optimize_ctx.computations entry, newest first): analyze over Low_level.instantiate_at_synthetic_read's code, the node's own traffic excluded. placements (default: empty) is the lineage the read is instantiated against — it decides which nested scopes the pipeline retracts to reads and whether a packed-uniform producer's counter can serve the lane-extract form; static_indices are the routine's. None when the inliner refuses the read.

val recompute_cost : ?static_indices:Indexing.static_symbol list -> Low_level.optimize_ctx -> Tnode.t -> recompute option

One read of a node the lineage stored computations for, through instantiation_cost under the lineage's placements; None for a node with none. Nested virtual producers need no expansion: a stored computation carries every producer inlined when it was stored as a nested scope, exactly what a read replays, and a producer read that stayed a read is a read in the emitted code too. The flip-candidate pricer (Low_level.recompute_pricer) reaches the same instantiation through the world specialize_proc hands it per node — including, for a node a heuristic cap materialized, the computations the virtualizer's own walk stores once the node is not materialized. Memoized per lineage: partially apply to the context once per compile.

module Calibration : sig ... end

The calibration TSV schema (config autotune_calibration_file, gh-ocannl-491 task 4) and the envelope fitter over it (gh-ocannl-514 phase 0): one row per timed candidate, the model's inputs and roofline score next to the measured time. This module is the schema's single owner — rows are emitted through to_line (by Autotune) and read back through of_line (by tools/fit_envelope.exe), so writer and reader cannot drift apart.