Module Ocannl.Calibrate

Memory-bandwidth calibration (gh-ocannl-578): STREAM-style streaming kernels whose byte counts are exact by construction — one access site per tensor node and direction, no guards — so the calibration rows they append under config autotune_calibration_file satisfy the fitter's per-leg exactness rule and can constrain the envelope's memory leg (model_peak_memory_bandwidth). Matmul-family tuning data cannot: its rows are compute-bound (their bytes-over-time understates achievable bandwidth no matter how exactly counted) and typically carry guards-taken/union upper bounds, so the fit of the memory leg needs exactly this kind of calibration diversity — see tools/calibrate_bandwidth.ml for the command-line entry point and Ir.Cost_model.Calibration for the fit semantics.

val stream : ?elems:int -> ?repeats:int -> Context.t -> (string * Autotune.report) list

stream ?elems ?repeats ctx times the four STREAM kernels — copy c =: a, scale b =: 0.5 *. a, add c =: a + b, triad a =: b + 0.5 *. c — over 1-D single-precision tensors of elems cells each (default 2^26, 256 MiB per tensor: well past any last-level cache, so the measured rates are main-memory rates). Each kernel goes through Autotune.tune with seed candidates only (~rounds:0) and the schedule cache disabled, so every timed candidate appends a calibration row through the ordinary emission path; the search is forced on (~search:true) — timing is the point of a calibration pass, whatever profile the config picked. Choose elems a power of two: an extent every workgroup size divides evenly keeps the parallelized candidates free of range guards, hence their rows bytes-exact.

Returns the kernels' tuning reports in order, (name, report). Raises Utils.User_error when config autotune_calibration_file is not set — the pass would time kernels with nowhere to record them.

Cache-capacity calibration (gh-ocannl-1253)

The ranking score (Ir.Cost_model.ranking_seconds) charges an untiled contraction's operand re-reads once its working set — per-block global read bytes times resident blocks — exceeds a capacity. That capacity is a fitted constant of the proxy, not the device's L2 size, so it is calibrated the way the bandwidth leg is: by a sweep whose kernels sit on both sides of it.

type ws_point = {
  1. wp_k : int;
    (*

    The contraction's reduction extent.

    *)
  2. wp_loops : int list;
    (*

    The lowered loop extents of the nest writing the output, outermost first: [k; m; n] is backprop's weight-gradient order.

    *)
  3. wp_working_set : int;
    (*

    The contraction kernel's working set under the default plan, in bytes (Ir.Cost_model.reuse_traffic).

    *)
  4. wp_requested : int;
    (*

    Its global read bytes under per-block caching.

    *)
  5. wp_distinct : int;
    (*

    Its distinct-cell bytes, read and written (Ir.Cost_model.analyze).

    *)
  6. wp_ms : float;
    (*

    The routine's best synced wall time over the repeats, in milliseconds.

    *)
  7. wp_rank_ms : float option;
    (*

    The kernel's ranking score under the capacity in force (Autotune.ranking_ms).

    *)
  8. wp_bound_ms : float option;
    (*

    Its distinct-cell roofline lower bound.

    *)
}

One point of the working_set sweep.

val working_set : ?m:int -> ?n:int -> ?ks:int list -> ?repeats:int -> Context.t -> ws_point list

working_set ?m ?n ?ks ?repeats ctx times the untiled weight-gradient contraction dw[o,i] = sum_k dy[k,o] * x[k,i] — m rows of output (default 256), n columns (default 1024: a gpt2_mini_train layer's l*_ffn_w2 gradient) — at each reduction extent of ks (default 256 to 65536 in powers of two, batch 2 to 512 at sequence 128), on the plain default pipeline (Context.compile: the plan whose re-reads the score charges, never a seeded or tuned one). The inputs are explicitly initialized host arrays (a standalone fixture: nothing upstream of the kernel is timed or left lazily initialized); the time is the whole routine's, synced, best of repeats (default 10) after one warm-up run. Each point carries the contraction kernel's working set and byte charges as the cost model reads them off the scheduled code.

type capacity_fit =
  1. | Transition of {
    1. cached_ws : int;
      (*

      The largest working set still served at the cache rate.

      *)
    2. spilled_ws : int;
      (*

      The smallest working set served at the DRAM rate.

      *)
    3. capacity : int;
      (*

      Their geometric mean: the fitted model_cache_capacity.

      *)
    4. monotone : bool;
      (*

      Every cached point's working set is below every spilled point's. false means the sweep's rates do not split cleanly by working set, and the capacity is a guess.

      *)
    5. cache_rate : float;
      (*

      The median requested rate of the cached points, bytes/s: the fitted model_cache_bandwidth.

      *)
    }
  2. | All_cached of {
    1. largest_ws : int;
    2. cache_rate : float;
    }
    (*

    Every point is served at the cache rate: the capacity is at least largest_ws, and the sweep says no more than that (no capacity is invented; extend ks to find one).

    *)
  3. | All_spilled of {
    1. smallest_ws : int;
    }
    (*

    Every point is served at the DRAM rate: the capacity is below smallest_ws.

    *)
  4. | No_transition of {
    1. contrast : float;
    }
    (*

    Without a DRAM rate to compare against, the fastest and slowest requested rates differ by contrast < 2: the sweep cannot tell which side of the capacity it sits on.

    *)
  5. | Too_few_points
    (*

    What fit_capacity reads off a sweep.

    *)
val fit_capacity : ?dram_rate:float -> ws_point list -> capacity_fit

The capacity and cache bandwidth the sweep implies. A point's requested rate is its per-block read bytes over its time: within capacity, re-reads are served from cache and the rate sits far above DRAM bandwidth; past it, the rate falls to it. Points with a smaller working set than the fastest one are dropped first: a kernel too short to amortize its launch shows a LOW rate for a reason that is not the cache (tuf's 0.4 ms floor at the smallest extent read as spilled and made the first fit non-monotone). Over the rest, a point is cached when its rate is at least the geometric mean of the fastest rate and dram_rate (bytes/s, the calibrated DRAM bandwidth) — or, without one, of the fastest and the slowest rates, which needs a contrast of at least 2 to mean anything. The capacity is placed at the geometric mean of the working sets either side of the split — the model's decision boundary, between the shapes that demonstrably sit on each side of it — and the cache bandwidth is the cached points' median rate.

Issue-rate calibration (gh-ocannl-1442)

The ranking score prices a kernel's scalar operations at an effective issue rate (model_issue_rate, the backends' advisory issue_rate): what a register-blocktiled kernel achieves, its operand staging included — far below the envelope's peak_flops, which stays the bound's ceiling. It is fitted the way the capacity is, on a sweep: the seeded sketch of the weight-gradient contraction over reduction extents, timed, against its scalar operation count.

type issue_point = {
  1. ip_k : int;
    (*

    The contraction's reduction extent.

    *)
  2. ip_scalar_ops : int;
    (*

    The contraction kernel's scalar operations (Ir.Cost_model.analyze's flops less the tile-MMA work the backend accepts): the count the score divides by the issue rate.

    *)
  3. ip_mma_ops : int;
    (*

    Its accepted tile-MMA operations (Ir.Cost_model.accepted_tile_mma_flops): 0 for the scalar sketch the sweep seeds.

    *)
  4. ip_kernels : int;
    (*

    The routine's shipped kernels.

    *)
  5. ip_ms : float;
    (*

    The routine's best synced wall time over the repeats, in milliseconds.

    *)
  6. ip_rank_ms : float option;
    (*

    The contraction kernel's ranking score under the configuration in force (Autotune.ranking_ms).

    *)
  7. ip_bound_ms : float option;
    (*

    Its distinct-cell roofline lower bound.

    *)
  8. ip_contexts : Ir.Cost_model.launch_context list;
    (*

    The contraction kernel's parallel contexts (Ir.Cost_model.launch_contexts): what a launch-fill model needs to predict the point (gh-ocannl-1479).

    *)
}

One point of the issue_sweep.

val issue_sweep : ?m:int -> ?n:int -> ?ks:int list -> ?repeats:int -> Context.t -> issue_point list

issue_sweep ?m ?n ?ks ?repeats ctx times the contraction dw[o,i] = sum_k dy[k,o] * x[k,i] — m rows (default 512), n columns (default 512: deliberately not the 256 x 1024 layer the calibration's validation cells time, so that predicting those is not reading the fit back) — at each reduction extent of ks (default 256 to 32768 in powers of two), compiled through the seeded default (Autotune.seeded_transform) at threshold 32 with the split and the reduction sink off: the one register-blocktiled kernel whose rate the score prices, whatever gate the configuration picks. Raises Utils.User_error where the seeded default does not apply (a backend other than CUDA or HIP, scheduling or fission off) or a point did not seed — a fallback is the default's kernel, never timed under the sketch's name. Inputs, timing and the kernel's selection (the shipped segment with the most operations) are working_set's.

type issue_fit =
  1. | Issue_fit of {
    1. rate : float;
      (*

      Scalar operations per second: the slope's reciprocal, the fitted model_issue_rate.

      *)
    2. latency_ms : float;
      (*

      The fit's intercept: time not proportional to the operations (launch, ramp-up). The score does not price it; it is reported so a prediction's residual can be read against it.

      *)
    3. max_residual : float;
      (*

      The largest relative error of the fitted line over the points.

      *)
    4. points : int;
      (*

      The points fitted.

      *)
    }
  2. | Issue_unfittable of string
    (*

    Why no line fits: too few points, or no growth.

    *)
val fit_issue_rate : ?fill:Ir.Cost_model.fill_model -> issue_point list -> issue_fit

The least-squares line time = latency + ops / rate over the timed points with operations; under fill (gh-ocannl-1479), each point's operations divided by its kernel's scalar launch fill (Ir.Cost_model.launch_fill of ip_contexts): the fill-normalized rate, equal to the plain one where the calibration kernel fills the device.

Launch-fill calibration (gh-ocannl-1479)

The ranking score prices a kernel's scalar work at the issue rate times its launch fill (Ir.Cost_model.context_fill): min 1 (B * min T t / (P * t)) for B blocks of T threads on P multiprocessors saturating at t threads each. P is the device's; t and the full-fill issue rate are fitted together, on a sweep that leaves full fill: one PINNED scalar blocktile geometry per series (never the seeding's size-dependent choice, which would mix geometries of different rates), over output sizes from a handful of blocks to many, at two block widths so a saturation above the narrower block's threads is identifiable.

type launch_point = {
  1. lp_geometry : string;
    (*

    The pinned leaf's bm x bn x bk / tm x tn.

    *)
  2. lp_m : int;
  3. lp_n : int;
    (*

    The output's extents.

    *)
  4. lp_k : int;
    (*

    The reduction extent, sized so the kernel's work is about the sweep's target.

    *)
  5. lp_contexts : Ir.Cost_model.launch_context list;
    (*

    The shipped kernels' parallel contexts (Ir.Cost_model.launch_contexts).

    *)
  6. lp_ms : float;
    (*

    The routine's best synced wall time over the repeats, in milliseconds.

    *)
}
val launch_sweep : ?geometries:(int * int * int * int * int) list -> ?sizes:int list -> ?target_ops:int -> ?max_k:int -> ?repeats:int -> Context.t -> launch_point list

launch_sweep ctx times dw[o,i] = sum_k dy[k,o] * x[k,i] at each square output size of sizes (default 32 to 1024 in powers of two) a geometry of geometries tiles exactly (default 32x32x8/4x4, 64 threads a block, and 64x64x8/4x4, 256), compiled through Autotune.pinned_sketch_transform with that geometry, at the power-of-two reduction extent (at most max_k, default 2^20) putting about target_ops (default 2^31) operations in the kernel. Raises Utils.User_error on a backend without the GPU sketch family, or where a geometry is no leaf of the family at a size.

type fill_form =
  1. | Capped
    (*

    Ir.Cost_model.context_fill: each block counts at most t threads.

    *)
  2. | Uncapped
    (*

    min 1 (B * T / (P * t)): the alternative the fit reports for comparison.

    *)
val effective_ops : ?form:fill_form -> multiprocessors:int -> fill_threads:int -> Ir.Cost_model.launch_context list -> float

The scalar contexts' operations, each divided by its fill under form (default Capped).

type fill_fit =
  1. | Fill_fit of {
    1. cores : int;
      (*

      The P the fit was given, the model_multiprocessors.

      *)
    2. saturation : int;
      (*

      The fitted t, the model_fill_threads.

      *)
    3. rate : float;
      (*

      The full-fill scalar operation rate, operations/s: the normalized model_issue_rate.

      *)
    4. max_residual : float;
      (*

      The largest relative error of the fitted line over the points.

      *)
    5. partial_points : int;
      (*

      Points below 0.9 fill under the fitted t.

      *)
    6. points : int;
      (*

      The points fitted.

      *)
    7. starved_points : int;
      (*

      Points left out as latency-starved: below starved of their geometry's best throughput.

      *)
    8. raw_rate : float option;
      (*

      The same points' proportional fit over RAW scalar operations — what a fit blind to the launch reads off them.

      *)
    9. raw_residual : float option;
    10. uncapped : (int * float) option;
      (*

      The best t and its largest residual under the Uncapped form, beside the capped fit: which form the data prefers.

      *)
    }
  2. | Fill_unfittable of string
    (*

    Too few points, a point with tile-MMA work, no finite fit, no geometry whose own points both reach and fall below full fill under the best t (each series' rate would absorb any saturation), or a sweep that never leaves full fill under the best t (fewer than two points below 0.9 fill): no saturation is identified.

    *)
val default_fill_candidates : int list
val fit_fill : ?candidates:int list -> ?starved:float -> multiprocessors:int -> launch_point list -> fill_fit

For each candidate t (default default_fill_candidates), per geometry, the proportional fit time = effective_ops / rate minimizing the squared relative residuals — per geometry, as a wider tile reuses more and runs at its own rate while the saturation is the device's; relative, as a launch sweep's times span orders of magnitude; through the origin, as the sweep holds the operations fixed and an intercept would trade against the rate (the score prices no latency). The candidate with the smallest summed squared relative residuals wins; rate is the lead geometry's (the first point's: the seeded sketch's). Points below starved (default 1/8) of their geometry's best throughput are left out: a few warps per multiprocessor with no latency hidden sit on a regime the fill does not model, and fitting them inflates t until well-filled launches read as partial. The capped and uncapped forms coincide on points whose blocks are no wider than t.