Module Ir.Schedule_cache

Process-independent schedule identities and the schedule disk cache

Support for persisting and replaying Schedule.schedule values (docs/proposals: the autotune companion of schedule-ir-optops.md). Schedules embed Indexing.symbols and Tnode.ts whose identities are process-local (global counters), so a schedule value is only meaningful against the one lowering it was built for — and every backend compile lowers afresh. This module gives both a canonical, structural identity:

The same traversal renders the code into a canonical string and digests it. The digest is the safety guarantee: a schedule saved against a digest is only ever replayed onto code with an equal digest, which makes the canonical numbering total and unambiguous by construction — nondeterministic lowering degrades to cache misses, never to a schedule applied to the wrong loop.

type mint_role =
  1. | Split_outer
  2. | Split_inner
  3. | Expand_axis of Base.int
    (*

    The i-th fresh symbol of an Expand_zero.

    *)
  4. | Tensorize_lane

Which fresh symbol of a schedule op a Minted reference names.

val sexp_of_mint_role : mint_role -> Sexplib0.Sexp.t
val mint_role_of_sexp : Sexplib0.Sexp.t -> mint_role
val compare_mint_role : mint_role -> mint_role -> Base.int
val equal_mint_role : mint_role -> mint_role -> Base.bool
type sym_ref =
  1. | Base of Base.int
  2. | Static of Base.int
  3. | Minted of Base.int * mint_role

A process-independent name for a symbol occurring in a schedule. Base i is the i-th For_loop binder in preorder of the optimized code the schedule applies to; Static k the k-th static index; Minted (op, role) the fresh symbol in role of the op-th (0-based) op of the same saved schedule.

val sexp_of_sym_ref : sym_ref -> Sexplib0.Sexp.t
val sym_ref_of_sexp : Sexplib0.Sexp.t -> sym_ref
val compare_sym_ref : sym_ref -> sym_ref -> Base.int
val equal_sym_ref : sym_ref -> sym_ref -> Base.bool
type saved_optop =
  1. | Split of {
    1. axis : sym_ref;
    2. factor : Base.int;
    3. outer : Low_level.axis_type;
    4. inner : Low_level.axis_type;
    }
  2. | Swap of {
    1. outer : sym_ref;
    2. inner : sym_ref;
    }
  3. | Retype of {
    1. axis : sym_ref;
    2. ty : Low_level.axis_type;
    }
  4. | Unroll of {
    1. axis : sym_ref;
    2. materialize : Base.bool;
    }
  5. | Stage of {
    1. source : Base.int;
    2. tile_loops : sym_ref Base.list;
    3. shared : Base.bool;
    4. cooperative : Base.int Base.option;
    5. hoisted : Base.bool;
    }
  6. | Privatize of {
    1. target : Base.int;
    2. over : sym_ref;
    }
  7. | Expand_zero of {
    1. tn : Base.int;
    }
  8. | Tensorize of {
    1. i : sym_ref;
    2. j : sym_ref;
    3. k : sym_ref;
    4. simd_width : Base.int;
    }

Schedule.optop with symbols replaced by sym_refs and tensor nodes by their canonical first-occurrence index.

val sexp_of_saved_optop : saved_optop -> Sexplib0.Sexp.t
val saved_optop_of_sexp : Sexplib0.Sexp.t -> saved_optop
val compare_saved_optop : saved_optop -> saved_optop -> Base.int
val equal_saved_optop : saved_optop -> saved_optop -> Base.bool
type saved_schedule = saved_optop Base.list
val sexp_of_saved_schedule : saved_schedule -> Sexplib0.Sexp.t
val saved_schedule_of_sexp : Sexplib0.Sexp.t -> saved_schedule
val compare_saved_schedule : saved_schedule -> saved_schedule -> Base.int
val equal_saved_schedule : saved_schedule -> saved_schedule -> Base.bool
type canonical

The canonical identity of one Low_level.optimized value: the digest, the loop-binder and static-symbol numbering, and the tensor-node numbering.

val canonicalize : ?static_indices:Indexing.static_symbol Base.list -> ?with_placements:Base.bool -> Low_level.optimized -> canonical

Walks the optimized code once in preorder, numbering For_loop binders, first-occurrence tensor nodes (their dims, precision, hoisted-packing eligibility Schedule.hoistable_constant — schedule validity depends on operand constancy, gh-ocannl-470 — and effective placement class from the compile's Ir.Low_level.optimize_ctx — identical code over Local scratch vs an On_device buffer generates different kernels, so same-code different-placement programs must not share cache keys — all enter the digest), and rendering the canonical form. with_placements = false omits the placement classes, giving the structural identity: placement classes can render differently across compilation lineages on byte-identical code, so per-segment schedule matching in fissioned replays keys on structure only. The binder/tensor-node numbering is identical either way. static_indices must be the same list the code was lowered with (Indexing.bound_symbols of the compile's bindings).

val digest : canonical -> Base.string

Hex digest of the canonical rendering. Equal digests mean structurally identical code, hence interchangeable canonical numberings.

val complete : canonical -> Base.bool

false when the code contains constructs the canonical rendering cannot capture (Staged_compilation closures, unbound or shadowed loop symbols). Incomplete canonical forms must not be used as disk-cache keys (distinct programs could collide); within one process they still support to_saved/of_saved round-trips.

val tn_of_ref : canonical -> Base.int -> Tnode.t

The tensor node at a canonical index. Raises Invalid_argument on out-of-range.

Symbol resolution registries

A registry resolves the symbols of a particular compile's code to sym_refs: base and static symbols through its canonical, schedule-minted symbols through entries recorded by to_saved / of_saved. Use it to translate loops of transformed code (base code with a schedule prefix applied) into references a schedule extension can persist.

type registry
val base_registry : canonical -> registry
val resolve : registry -> Indexing.symbol -> sym_ref Base.option
val resolve_tn : registry -> Tnode.t -> Base.int Base.option

Serializes a schedule built against the registry's compile (e.g. by Schedule.default_gpu), recording each op's minted symbols in the returned registry (op indices continue from the number of ops already recorded in the input registry, so extensions of replayed prefixes stay consistent). Raises Invalid_argument when an op references a symbol or tensor node the registry cannot resolve.

Replays a saved schedule against a (fresh) compile's canonical form: base and static references resolve through canonical, minting ops go through the Schedule builders and their fresh symbols are recorded for later references. Raises Invalid_argument on dangling references (canonical mismatch — always guard with digest equality first).

The disk cache

type entry = {
  1. version : Base.int;
  2. backend : Base.string;
  3. source_digest : Base.string;
  4. saved : saved_schedule;
  5. segments : (Base.string * saved_schedule) Base.list Base.option;
    (*

    A fissioned winner (docs: per-fission-segment tuning): per-segment schedules keyed by the pre-schedule segment's canonical digest — replay routes each of Schedule.fission_scheduled's `Normal segments through this association (saved is then empty; unmatched segments degrade to the empty schedule). None for whole-routine schedules.

    *)
  6. best_ms : Base.float;
    (*

    The winning candidate's measured time, for diagnostics.

    *)
  7. baseline_ms : Base.float;
    (*

    The unscheduled baseline's measured time, for diagnostics.

    *)
}
val sexp_of_entry : entry -> Sexplib0.Sexp.t
val entry_of_sexp : Sexplib0.Sexp.t -> entry
val entry_version : Base.int

Bumped when the canonical rendering or the saved-schedule format changes; stale entries are ignored by lookup.

val cache_key : canonical -> backend:Base.string -> Base.string

Filename-safe cache key: the digest plus the backend name. Callers time kernels on a concrete device, so include anything else that distinguishes performance environments in backend (e.g. a device id) if needed.

val store : dir:Base.string -> key:Base.string -> entry -> Base.unit

Writes the entry to dir/key.sexp, creating dir (and parents) if missing. Tolerates concurrent writers (last write wins; writes go through a temp file + rename).

val lookup : dir:Base.string -> key:Base.string -> entry Base.option

None on missing file, unparsable content, or version/digest mismatch.