Module Sarek_policy

One typed store for every precision grant and consent (specs/optimisation-policy.md).

The program states, once, which deviations from the interpreter's result it accepts (deviation) and which rewrites of its own code it allows (consent); this module records those statements and answers "is this grant in force, and why". It sits in the lowest layer (sarek_ir, no dependencies) so that code generation, the capability gate and the launch path can all ask it.

What it is not. It is not a switch board: nothing here turns an optimisation on or off, and no environment variable grants anything (SPOC and Sarek environment variables only restrict). It does not change what any caller computes: Sarek_coopmat.accept_relaxed_accumulation and the other opt-ins are thin wrappers over grant, with the same refusals and the same messages as before it existed.

Who consults which scope (milestone M2). Every decision point that knows the kernel it is deciding for reads the grants of that kernel through budget (or effective): the f32 to f16 narrowing, the coopmat verdicts (Sarek_coopmat.policy_verdict_for), the matrix rewrite and the epilogue fusion. A Kernel-scoped grant therefore changes the decision for that kernel and for no other. Callers with no kernel read the process scope. The text a code generator emits, and what the launch memo stores, depend on the policy: a grant given after a kernel's source was generated takes effect at its next launch (the memo compares version, then the digest of the kernel's effective policy).

Restrictions. restrict is the one way to take an optimisation away: process-wide and one-way, like a grant. It is a diagnostic switch for SPOC's own environment parsing and for tests, not a scoped API.

The environment (SAREK_OPT_DISABLE, SAREK_EXPLAIN) is read once, when this module is initialised (at program start, before any query), by Env. SAREK_OPT_DISABLE only restricts: it disables optimisations and revokes grants for the whole process (a revoked grant is still recorded by grant, but no budget, reader or decision sees it). It never grants. SAREK_EXPLAIN=1 installs a reporter that writes one line per decision to stderr (Env.explain_line).

Vocabulary

type operand_precision =
  1. | Operands_f16
  2. | Operands_bf16
  3. | Operands_bf16x3

The operand type a float32 matrix product is rounded to before it runs on matrix units. Operands_bf16x3 splits each operand into three bfloat16 terms for near-float32 accuracy. Re-exported, with its constructors, by Sarek_coopmat.

type deviation =
  1. | Accumulation_order
    (*

    the hardware chooses the order of the additions (Sarek_coopmat.accept_relaxed_accumulation)

    *)
  2. | Operand_rounding of operand_precision
    (*

    float32 operands are rounded to a narrower type on load (Sarek_coopmat.accept_operand_rounding)

    *)
  3. | Single_rounding
    (*

    an f32 to f16 narrowing may round once where the interpreter rounds twice (Sarek_f16_policy.accept_single_rounding)

    *)
  4. | Cross_kernel_contraction
    (*

    a device compiler may contract a multiply of one fused kernel with an add of the next. Grantable now; read by epilogue fusion from M2.

    *)

A named way a result may differ from the interpreter's. Each carries its own documented bound where one exists; there is no common numeric scale.

A permission to rewrite the program's own code, separate from accepting any accuracy deviation: granting one never grants the other.

type grant =
  1. | Deviation of deviation
  2. | Consent of consent
type optimisation =
  1. | Launch_window
  2. | Epilogue_fusion
  3. | Matrix_rewrite
  4. | Narrowing_proof
  5. | Auto_unroll

Every transformation a policy decision can apply to a launch. Vocabulary for the explain report; the decision points are wired to it from M2.

type scope =
  1. | Process
  2. | Kernel of string

Where a grant applies. Kernel name matches Sarek_ir_types.kernel.kern_name exactly; a kernel built by a fusion pass has the fused name, which is never matched by a grant: its budget is the meet over its parts (register_fused).

type refusal =
  1. | Blank_reason
  2. | Default_kernel_name

Why grant (or Budget.add) refused. Default_kernel_name: Kernel "sarek_kern", the name every unnamed [%kernel] gets, would cover all of them.

type evidence =
  1. | Proof of string
  2. | Granted of {
    1. grants : grant list;
    2. reason : string;
    3. lowering : string option;
    }
  3. | Exact_by_construction

What justifies an applied optimisation, for the explain report. Granted: the grants it needed (each held by the kernel), the program's reasons for them (join_reasons) and, for a rewrite, a description of what it did (lowering, printed as a trailing field).

type decision =
  1. | Applied of evidence
  2. | Refused of string

The outcome of one decision point, for the explain report.

val grant_name : grant -> string

Stable names, used in reports and in cache keys: "accumulation-order", "operand-rounding-f16", "operand-rounding-bf16", "operand-rounding-bf16x3", "single-rounding", "cross-kernel-contraction", "automatic-float-rewrite".

val optimisation_name : optimisation -> string

"window", "epilogue", "matrix", "narrowing-proof", "unroll".

val all_grants : grant list

The seven grants, in the order Budget.grants lists them.

val all_optimisations : optimisation list

Budgets

module type BUDGET = sig ... end
module Budget : BUDGET

The process-wide store

val grant : scope -> grant -> reason:string -> (unit, refusal) Stdlib.result

grant scope g ~reason records that the program accepts g for scope. Process-wide and one-way: there is no withdrawal, so a check that passed (and that a caller memoised) can never become stale. Thread-safe (an Atomic updated by compare-and-set). The first non-empty reason per (scope, g) is kept; later calls for the same pair are no-ops. A blank reason is Error Blank_reason and Kernel "sarek_kern" is Error Default_kernel_name; a refusal records nothing.

val reason : ?kernel:string -> grant -> string option

reason ?kernel g: with ~kernel, that kernel's own grant first, then the process grant; without, the process grant only. None when not granted.

val budget : kernels:string list -> Budget.t

budget ~kernels: the budget of a launch whose kernel is built from kernels. For each kernel, its process grants plus its own; the result is the Budget.meet over the kernels, so a grant held by one member only is absent. [] gives the process grants only. Reads one consistent snapshot of the store.

val refusal_message : refusal -> string

The sentence for a refusal of grant.

Fused kernels

val register_fused : name:string -> members:string list -> unit

register_fused ~name ~members: a fusion pass built the kernel name from the kernels members (themselves possibly fused kernels). From then on budget ~kernels:[name] is the Budget.meet over the budgets of the members, recursively, so a fused kernel holds a grant only if every member does, and a grant given to a member after the registration is seen.

A Kernel grant on the fused kernel's own name is not counted: fusion names are generated (content digests, counters) and are not a statement of the program. Registering a name again adds the new members to those already registered (a content-derived name can come from different members; it is then covered only by what all of them allow). A registration with no member is ignored, and a name that would contain itself contributes no grant. Thread-safe; bumps version when it changes the registry.

val register_derived : name:string -> from:string -> unit

register_derived ~name ~from: the kernel name is from rewritten (the matrix rewrite renames its output <from>_sarek_mm), not a statement of the program, and has exactly from's budget (recursively, if from is itself fused or derived) so that a Kernel grant of from is read for the rewritten kernel at the capability gate and by the interpreter. register_fused ~name ~members:[from]; same registry, same rules.

Restrictions

val restrict : optimisation -> unit

restrict opt takes opt away for the rest of the process. One-way and thread-safe; restricting twice is a no-op. Only SPOC's environment parsing (SAREK_OPT_DISABLE) and tests call it. A restriction never invalidates a check a caller memoised, because a memoised check is re-keyed on version. Bumps version when opt was not yet restricted.

val restricted : optimisation -> bool
val revoked : grant -> bool

revoked g: SAREK_OPT_DISABLE revoked g for the process. A revoked grant is still recorded by grant (the program's statement), but reason, budget, effective and decide never see it. There is no way to revoke in code: only the environment does, once, at start-up.

The effective policy of a launch

type effective = {
  1. budget : Budget.t;
  2. disabled : optimisation list;
}

What a decision point reads: the grants in force for the kernel being decided, and the optimisations taken away. disabled is in all_optimisations order.

val effective : kernel:string -> effective

effective ~kernel: budget ~kernels:[kernel] and the current restrictions.

val effective_for : kernels:string list -> effective

As effective for a launch built from kernels (the meet of their budgets): the epilogue decision of a producer and its consumer.

val effective_digest : effective -> string

A stable string of the whole effective policy: Budget.digest and the disabled optimisations. Equal iff the policies are equal.

val requires : optimisation -> backend:string -> grant list option

What an optimisation needs on a backend, as grants that must be held; Some [] is exact there and None is not offered there.

  • Epilogue_fusion: Some [] on "HIP" and "Metal" (every kernel is compiled without contraction); Some [Deviation Cross_kernel_contraction] on "CUDA/PTX", "CUDA/C", "CUDA", "OpenCL" and "Vulkan"; None elsewhere (Native, Interpreter, WebGPU never defer).
  • Launch_window, Narrowing_proof, Auto_unroll: Some [] everywhere.
  • Matrix_rewrite: Some [] everywhere; its float lowerings check their own grants (the accuracy opt-ins and the consent), unchanged.

A Metal that marked itself inexact (Fp_exactness) is refused by the epilogue even with the grant: that fact lives above this module.

val disabled_reason : string

"disabled (SAREK_OPT_DISABLE)": the reason of a Refused by a restriction, for decision points that refuse without decide.

val decide : effective -> optimisation -> backend:string -> decision

The common decision, in this order:

  • Refused "disabled (SAREK_OPT_DISABLE)" if opt is in e.disabled;
  • Refused ("not on " ^ backend) if requires is None;
  • Refused ("no grant " ^ name ^ " (kernel or process)") for the first required grant e.budget does not hold, or Refused ("grant " ^ name ^ " revoked (SAREK_OPT_DISABLE)") when it does not hold because the environment revoked it;
  • Applied Exact_by_construction if nothing is required;
  • Applied (Granted {grants; reason}) otherwise: the required grants and their reasons joined by "; ".

Change counter and decision reporting

val version : unit -> int

Increases with every change of the policy: a grant that was not yet recorded, a restriction not yet in force, a fused registration that adds a member. Refusals, duplicates and reads leave it alone. A cache that memoised something derived from the policy stores version () read before it read the policy; one atomic read tells it whether anything changed since.

val join_reasons : string list -> string

The reasons joined by "; ", each once, in order of first appearance.

type record = {
  1. kernel : string;
  2. target : string;
  3. optimisation : optimisation;
  4. decision : decision;
}

One decision at a decision point, for the explain report. target says where it was taken, in one scheme: "<framework>:<device id>" for a decision about a device (window, matrix), "<framework>" for one about a backend (epilogue), "gen:<generator>" for code generation (narrowing proof, unroll: "gen:cuda-c", "gen:opencl", "gen:metal", "gen:glsl").

type ignored =
  1. | Generated_name
  2. | Revoked

Why a grant the program made has no effect. Generated_name: a Kernel grant whose name is a fused, derived or static-fusion kernel's (generated, never matched). Revoked: the environment revoked the grant.

type event =
  1. | Decision of record
  2. | Gate of {
    1. kernel : string;
    2. target : string;
    3. decision : decision;
    }
  3. | Ignored_grant of {
    1. kernel : string;
    2. grant : grant;
    3. cause : ignored;
    }

What the sink of report receives: a decision, the outcome of the explicit cooperative-matrix gate, or a statement of the program that the policy does not obey. Gate is an event kind of its own, not an optimisation: nothing restricts it and SAREK_OPT_DISABLE has no name for it; its target is "<framework>:<device id>". Ignored_grant comes right after the report of a kernel, once per (kernel, grant, cause), also when the decisions of that kernel settled earlier; a grant revoked at Process scope is reported once, with the kernel name "(process)".

val set_reporter : (event -> unit) option -> unit

Install (or, with None, remove) the sink of report. The default is no sink, except that SAREK_EXPLAIN=1 installs one at start-up that writes Env.explain_line to stderr; a sink installed by the program replaces it. Installing one forgets which outcomes were already reported. The sink is called outside the policy's lock, possibly from several domains: lines of different domains may interleave, and the sink must be thread-safe.

val reporting : unit -> bool

true when a sink is installed: lets a decision point skip building a record (and its strings) when nobody listens.

val report : record -> unit

report r: with a sink installed, calls it with Decision r the first time this decision is seen for (r.kernel, r.target, r.optimisation): a decision that comes back after another one is not reported again, so a window that is switched off and on per prompt chunk prints two lines, not one per switch. The memory of what was seen is bounded (4096 distinct decisions, then forgotten). Without a sink it is one atomic read. Safe to call from any domain.

val report_gate : kernel:string -> target:string -> decision -> unit

report_gate ~kernel ~target decision: the same for the explicit cooperative-matrix gate of a kernel (Gate event).

The environment

module Env : sig ... end

SAREK_OPT_DISABLE, SAREK_EXPLAIN. The pure part, so that tests need no environment; the module reads the variables once at initialisation and applies them with restrict and an internal revocation. Nothing here can grant.