Autotune (arrayjit.Autotune)

Empirical schedule search (autotuning)

tinygrad-style beam search over Ir.Schedule transforms, timed on the real device (docs/proposals/schedule-ir-optops.md; the search-harness half of the OptOps port). tune is a drop-in replacement for Context.compile: it compiles candidate schedules through the ?lowered_transform / ?lowered_transforms seams, times each on the context's device, and returns the routine of the fastest one. Every candidate (and the winner replay) derives from a hermetic copy of the one base lowering captured at the start — each candidate compile's own fresh lowering is ignored, because timing runs settle tensor-node value bounds and later lowerings can fold guards or re-segment fission differently, silently corrupting digest comparisons and replays. Winning schedules are persisted, in the structurally-rebindable saved form of Ir.Schedule_cache, to a disk cache keyed by the code's canonical digest and the backend, so a re-run of the same program skips the search (cross-process replay is guarded by digest equality against that process's own base lowering).

The candidate space:

Whole-routine presets: the serial baseline, the default annotator, and a block-size sweep through Ir.Schedule.default_gpu.
Fissioned candidates: the kernel-fission pipeline (Ir.Schedule.fission_scheduled) with per-segment schedules — the same preset sweep per segment, and beam rounds that extend one segment at a time. Per-segment schedules are cached keyed by the pre-schedule segment's canonical digest. `Zeros segments keep the default zero-expansion; `Solo segments stay unscheduled. One seed uses the config-default thresholds, reproducing the untuned default pipeline exactly — so the winner is never worse than not tuning, even on launch-overhead-bound workloads where every aggressive preset loses to it. Each preset is additionally seeded in a privatized variant (extend_with_privatize): per segment, every materialized read-modify-write accumulator is contracted into a per-thread register tile (Ir.Schedule.optop.Privatize) over its serial reduction loop where the op's preconditions permit — a routine-local accumulator beats a device-memory RMW, and on Metal it sidesteps the volatile scalar-RMW workaround.
Matmul sketches: when a matmul micro-kernel is detected, parameterized instantiations of the composed pipelines pinned by the schedule tests — register blocktiling (Split + Swap + shared Stage + Privatize + materializing Unroll) on GPU backends, operand packing (non-shared Stage + Privatize) on CPU backends — with dividing tile sizes.
Beam-round menu actions on the incumbents: dividing serial Splits, Swaps of perfect serial pairs, Unrolls, Retype-Vectorized on innermost loops (explicit SIMD on CPU including the reduction-chains rendering of accumulations — gh-ocannl-468 — while GPU accumulations stay excluded; 128-bit packed loads/stores on GPU — gh-ocannl-463), and Tensorize role permutations when the backend reports an mma capability — including the CPU backends, whose Tile_mma renders as the register-tiled vector micro-kernel (gh-ocannl-469).

Caveats (v1):

Timing runs execute the routine several times, mutating its outputs (and accumulators — a non-idempotent routine, e.g. gradient accumulation, will accumulate the timing runs). Initialize inputs before tuning, and tune before meaningful state exists, or re-initialize afterwards.
Timing uses wall clock around a device sync, so it includes queue overhead; times are min-of-N, where fast routines get extra runs beyond repeats until ~25 ms of total measured time — on sub-millisecond kernels a min-of-3 is launch-jitter roulette and can crown the wrong candidate. Static indices are bound to the midpoint of their declared ranges during timing and restored afterwards.

val extend_with_privatize : 
  static_indices:Ir.Indexing.static_symbol list ->
  Ir.Schedule.schedule ->
  Ir.Low_level.optimized ->
  Ir.Schedule.schedule

The privatized preset extension used by the fissioned candidates: appends a Schedule.Privatize { target; over } for every materialized read-modify-write accumulator detected in the schedule's application to the segment — over being the outermost enclosing Serial loop whose symbol the access vector does not mention and whose subtree contains no hardware-typed loop. Each proposal is validated by try-applying the grown schedule against a hermetic copy of the segment (proposals violating the op's preconditions are dropped), so the result always applies cleanly where the input schedule does. Exposed for tests.

type report = {

cache_hit : bool;
(*
The schedule came from the disk cache; no search ran.
*)
candidates_timed : int;
(*
Including the serial baseline.
*)
candidates_failed : int;
(*
Candidates rejected by op preconditions, hardware limits, or backend compilation.
*)
rounds_run : int;
(*
Beam-expansion rounds actually executed (0 = seeds only).
*)
sketch_candidates : int;
(*
Matmul-sketch instantiations seeded (0 when no matmul micro-kernel was detected or no tile sizes divide the extents). Deterministic given the computation and backend.
*)
fissioned : bool;
(*
The winning candidate compiles as multiple fissioned kernels.
*)
baseline_ms : float;
best_ms : float;
best_schedule : Ir.Schedule_cache.saved_schedule;
(*
The winner's schedule; for a fissioned winner, the concatenation of the per-segment schedules (informational).
*)

}

val tune : 
  ?beam_width:int ->
  ?rounds:int ->
  ?repeats:int ->
  ?seed_block_sizes:int list ->
  ?cache_dir:string ->
  ?timing_ctx:Context.t ->
  ?report:(report -> unit) ->
  Context.t ->
  Ir.Assignments.comp ->
  Ir.Indexing.unit_bindings ->
  Context.t * Context.routine

Like Context.compile, but returns the empirically fastest of the searched schedule candidates. The returned context/routine come from an ordinary sibling compile of ctx, so execution-dependency tracking behaves as if the winning compile were the only one. Raises like Context.run would (e.g. uninitialized inputs) — tune in the same state you would run in.