Module Sarek_ir_analysis

Sarek_ir_analysis - Analysis functions for GPU kernel IR

type 'a folder = {
  1. fe : 'a -> Sarek_ir_types.expr -> 'a;
  2. ft : 'a -> Sarek_ir_types.elttype -> 'a;
  3. fs : 'a -> Sarek_ir_types.stmt -> 'a;
    (*

    Combine the accumulator with a STATEMENT node, before the traversal descends into it. The counterpart of fe for the imperative half, added in backlog-62 slice 3 because SCoopmat is the first IR node whose feature content is neither an expression nor an element type: a cooperative-matrix multiply-add carries a Sarek_coopmat_types.config, and a config is what the device gate is keyed on.

    Families that do not inspect statements leave this at the identity, which reproduces the pre-slice-3 behaviour exactly — the hook fires at every statement, so a non-identity value here is a deliberate statement-level detector and never an accident of traversal order.

    *)
  4. fnative : 'a -> 'a;
  5. visit_lvalue : bool;
}

Generic IR traversal

Every "does a kernel use feature X?" detector — and the Float64-intrinsic name collector — shares one traversal skeleton over the IR (expr/lvalue/stmt/decl/helper/kernel). Historically each family duplicated that skeleton, so adding an IR node meant editing ~7 copies and any omission silently under-reported a feature.

A single polymorphic fold now carries the skeleton; each family supplies only its per-node behaviour through a folder record. This is also the requirement-extraction primitive a future capability/affinity model reuses, hence the fully general 'a accumulator rather than a fixed boolean.

The four hooks capture every axis along which the old families differed:

  • fe: combine the accumulator with a single expression node. This is the only place a family's "leaf" fires; the traversal always recurses into the node's sub-expressions afterwards. A rich leaf may inspect an expression's embedded types here (e.g. ECast/EArrayCreate element types, an EVar's var_type) — Float64 detection does exactly this, so the traversal never forces a lowest-common-denominator leaf.
  • ft: combine the accumulator with an element type occurring at a binder or declaration (SFor/SLet/SLetMut binder, DParam/DShared types, helper return/param types, kernel record/variant field types). Families that do not inspect types leave this at the identity, which makes those positions contribute nothing — reproducing the old "types ignored" behaviour exactly.
  • fnative: combine the accumulator at an SNative node. Inline native GPU code is opaque text, so its polarity is asymmetric across families (atomics/int_mod/copysign/generic-intrinsic conservatively assume the feature is present; float64/nonfinite/collector treat it as absent). This is carried explicitly and never flattened away.
  • visit_lvalue: whether SAssign descends into its l-value. The float64 detector deliberately ignores assignment l-values; the others recurse (an index expression can hide the feature).
val expr_fold : 'a folder -> 'a -> Sarek_ir_types.expr -> 'a
val lvalue_fold : 'a folder -> 'a -> Sarek_ir_types.lvalue -> 'a
val stmt_fold : 'a folder -> 'a -> Sarek_ir_types.stmt -> 'a
val decl_fold : 'a folder -> 'a -> Sarek_ir_types.decl -> 'a
val helper_fold : 'a folder -> 'a -> Sarek_ir_types.helper_func -> 'a
val kernel_fold : 'a folder -> 'a -> Sarek_ir_types.kernel -> 'a

Fold the whole kernel: params, locals, body, helper functions, and record /variant field types. The type positions (ft) contribute nothing for detectors that do not inspect types, so families that historically skipped kern_types/kern_variants are unaffected by visiting them here.

val exists_folder : leaf:(Sarek_ir_types.expr -> bool) -> ?type_leaf:(Sarek_ir_types.elttype -> bool) -> ?stmt_leaf:(Sarek_ir_types.stmt -> bool) -> native:bool -> ?visit_lvalue:bool -> unit -> bool folder

A boolean folder for an "exists" detector. leaf fires per expression node (and may inspect embedded types); type_leaf fires per binder/decl type; native is the SNative verdict; visit_lvalue controls whether SAssign descends into its l-value.

Numeric-width feature detection

ONE parameterised detector family, not one family per width. Adding a width (bf16 is next) is a constructor in feature, an arm in elttype_uses, and a line in folder — not a fresh copy of a rich leaf, a folder and five wrappers. The previous shape had float64 and float16 as two structurally identical families whose own docstring said so.

The family has a rich leaf: it inspects element types, not just constructors, at every binder, declaration, cast and array construction, plus record/variant field types at the kernel level. Two properties are shared by every width and are deliberate:

The one per-width asymmetry is CONSTANTS: float64 has CFloat64 literals, float16 has no literal and hence no CFloat16 constant (see Sarek_ir_types.elttype). An f16 value always enters through ECast (TFloat16, _) or an f16-typed binder/parameter, both of which the leaf sees. const_uses expresses that directly rather than by omitting an arm.

Consumers: kernel_uses Float64 drives the OpenCL/GLSL fp64 pragma/extension, and kernel_uses Float16 drives both the CUDA/HIP conditional #include <cuda_fp16.h> and the slice-2 rejection gate at every backend's generate entry (see Sarek_ir_codegen.reject_feature).

kernel_requirements is the set-valued form: it is what a future Kernel.requirements capability field reduces to, and it lives in the right layer already (spoc/ir, no backend dependencies).

type feature =
  1. | Float64
  2. | Float16
  3. | Int64
  4. | Coopmat
val all_features : feature list
val feature_name : feature -> string
val elttype_uses : feature -> Sarek_ir_types.elttype -> bool

Does element type t mention the width f, transitively through records, variants, arrays and vectors?

val const_uses : feature -> Sarek_ir_types.const -> bool

Is constant c a literal of width f? Float64 and Int64 each have one (CFloat64, CInt64); f16 has no literal form, so this is false for Float16 by construction rather than by a missing case.

val feature_leaf : feature -> Sarek_ir_types.expr -> bool
val feature_stmt_leaf : feature -> Sarek_ir_types.stmt -> bool

A statement mentions f when it is an SCoopmat and f is Coopmat.

Written as an explicit match on the feature rather than f = Coopmat && is_coopmat s so that a future statement-level feature (a barrier class, a printf) is a compile error here rather than a silent false.

val folder_of : feature -> bool folder
val float64_folder : bool folder
val float16_folder : bool folder
val int64_folder : bool folder
val coopmat_folder : bool folder
val folder : feature -> bool folder
val expr_uses : feature -> Sarek_ir_types.expr -> bool

Does expression e use width f?

val stmt_uses : feature -> Sarek_ir_types.stmt -> bool

Does statement s use width f?

val decl_uses : feature -> Sarek_ir_types.decl -> bool

Does declaration d use width f?

val helper_uses : feature -> Sarek_ir_types.helper_func -> bool

Does helper function hf use width f?

val kernel_uses : feature -> Sarek_ir_types.kernel -> bool

Does kernel k use width f anywhere — params, locals, body, helper params and return types, and record/variant field types?

val kernel_requirements : Sarek_ir_types.kernel -> feature list

The set of numeric-width features kernel k requires. The set-valued form of kernel_uses; a future Kernel.requirements is this.

Per-width aliases

Thin, so existing call sites and the analysis test suite do not churn. New code should prefer kernel_uses Float16 over the alias.

val elttype_uses_float64 : Sarek_ir_types.elttype -> bool
val const_uses_float64 : Sarek_ir_types.const -> bool
val expr_uses_float64 : Sarek_ir_types.expr -> bool
val stmt_uses_float64 : Sarek_ir_types.stmt -> bool
val decl_uses_float64 : Sarek_ir_types.decl -> bool
val helper_uses_float64 : Sarek_ir_types.helper_func -> bool
val kernel_uses_float64 : Sarek_ir_types.kernel -> bool
val elttype_uses_float16 : Sarek_ir_types.elttype -> bool
val expr_uses_float16 : Sarek_ir_types.expr -> bool
val stmt_uses_float16 : Sarek_ir_types.stmt -> bool
val decl_uses_float16 : Sarek_ir_types.decl -> bool
val helper_uses_float16 : Sarek_ir_types.helper_func -> bool
val kernel_uses_float16 : Sarek_ir_types.kernel -> bool
val is_atomic_intrinsic_name : string -> bool

Atomic-operation detection

Atomic intrinsics have no dedicated IR constructor: the PPX lowers every atomic primitive (see the category = "atomic" entries registered in sarek/ppx/Sarek_core_primitives.ml, and the %sarek_intrinsic atomics in sarek/Sarek_stdlib/Gpu.ml) to a plain EIntrinsic (path, name, args) node, e.g. "atomic_add_int32", "atomic_cas_int32", "atomic_add_global_int32", ... All such names share the "atomic_" prefix by registration convention.

REGISTRATION POINT: this is the single source of truth for recognizing an atomic intrinsic from IR. If a future atomic primitive is registered under a name that does not start with "atomic_", update is_atomic_intrinsic_name below (and consider exporting the name list from Sarek_core_primitives.ml instead of relying on the prefix convention). Do not duplicate this check elsewhere.

Inline native GPU code (SNative) is opaque; fusion must not assume it is atomic-free, so the detector is conservative there.

val atomics_leaf : Sarek_ir_types.expr -> bool
val atomics_folder : bool folder
val expr_uses_atomics : Sarek_ir_types.expr -> bool

Check if an expression contains an atomic intrinsic call

val lvalue_uses_atomics : Sarek_ir_types.lvalue -> bool

Check if an l-value contains an atomic intrinsic call (in its index/base expression). LVar has no sub-expression; LRecordField recurses into the inner l-value.

val stmt_uses_atomics : Sarek_ir_types.stmt -> bool

Check if a statement contains an atomic intrinsic call

val decl_uses_atomics : Sarek_ir_types.decl -> bool

Check if a declaration contains an atomic intrinsic call

val helper_uses_atomics : Sarek_ir_types.helper_func -> bool

Check if a helper function contains an atomic intrinsic call

val kernel_uses_atomics : Sarek_ir_types.kernel -> bool

Check if a kernel uses atomic operations anywhere: params/locals initializers, body, and helper functions called from the kernel. Helper bodies are walked explicitly — a body-only check would miss atomics hidden inside a called helper function.

val int_mod_leaf : Sarek_ir_types.expr -> bool

Integer-remainder detection

EBinop (Mod, _, _) is always integer remainder — float mod is lowered to the fmod/mod intrinsic (an EIntrinsic), never to Ir.Mod. Backends that cannot lower % directly (e.g. GLSL, whose % is undefined for negative operands) use this to decide whether to emit a remainder helper.

L-values are recursed (an array index may carry a mod, e.g. arr.(j mod n).field <- v); SNative is conservatively assumed to contain a remainder so any helper it references is still emitted.

val int_mod_folder : bool folder
val expr_uses_int_mod : Sarek_ir_types.expr -> bool
val lvalue_uses_int_mod : Sarek_ir_types.lvalue -> bool
val stmt_uses_int_mod : Sarek_ir_types.stmt -> bool
val decl_uses_int_mod : Sarek_ir_types.decl -> bool
val helper_uses_int_mod : Sarek_ir_types.helper_func -> bool
val kernel_uses_int_mod : Sarek_ir_types.kernel -> bool

Check if a kernel uses integer remainder anywhere: locals initializers, body, and helper functions.

val is_copysign_intrinsic_name : Stdlib.String.t -> bool

copysign detection

copysign is not a dedicated IR node (unlike Mod); it is an ordinary EIntrinsic (path, "copysign", [x; y]) emitted for Float32.copysign and Float64.copysign. GLSL has no copysign builtin under any name, and abs(x)*sign(y) is wrong for y=0 (GLSL sign(0)=0 zeroes the result, whereas C copysign(x, ±0) = ±|x|) and for the x=0/NaN sign-transfer edge cases. The GLSL backend therefore lowers it to a bit-level sarek_copysign helper emitted in the preamble; this predicate decides whether that helper is emitted. L-values are recursed (the round-3 LRecordField lesson) and SNative is conservatively assumed to reference the helper.

val copysign_leaf : Sarek_ir_types.expr -> bool
val copysign_folder : bool folder
val expr_uses_copysign : Sarek_ir_types.expr -> bool
val lvalue_uses_copysign : Sarek_ir_types.lvalue -> bool
val stmt_uses_copysign : Sarek_ir_types.stmt -> bool
val decl_uses_copysign : Sarek_ir_types.decl -> bool
val helper_uses_copysign : Sarek_ir_types.helper_func -> bool
val kernel_uses_copysign : Sarek_ir_types.kernel -> bool

Check if a kernel uses copysign anywhere: locals initializers, body, and helper functions.

val path_is_float64 : string list -> bool

Float64 intrinsic detection

Collects the names of every path-qualified Float64 math intrinsic invoked anywhere in a kernel — an EIntrinsic (path, name, _) whose path carries a "Float64" component (matching the four registry-exposing paths ["Float64"], ["Math"; "Float64"] and their Sarek_stdlib_meta twins, exactly the test the GLSL polyfill already uses).

A backend with no native f64 transcendental (GLSL core has no double overload for sin/cos/exp/log/pow/… — see Sarek_ir_glsl) uses this to decide which software helper family (Sarek_ir_softmath) to emit per kernel. Names are returned deduplicated; the caller filters to the subset it routes to helpers and maps the composed cases (exp2/log2/cbrt). This is a collector rather than a boolean detector, so it uses the generic fold with a string list accumulator: it ignores types (ft = identity) and treats SNative as contributing nothing.

val kernel_float64_intrinsics : Sarek_ir_types.kernel -> string list

Deduplicated names of the Float64 math intrinsics a kernel invokes.

val const_is_nonfinite_float64 : Sarek_ir_types.const -> bool

Non-finite Float64 constant detection

A CFloat64 whose value is ±inf or NaN cannot be spelled as a GLSL literal (GLSL has no inf/nan literal), so a backend targeting GLSL reconstructs it from its bit pattern via int64BitsToDouble — which needs GL_ARB_gpu_shader_int64. Such a constant can occur independently of any transcendental (e.g. a user-written Float64.infinity), so the int64 extension must be gated on this too, not only on the software helper family. SNative is treated as non-finite-free (native code carries its own literals).

val nonfinite_f64_leaf : Sarek_ir_types.expr -> bool
val nonfinite_f64_folder : bool folder
val kernel_uses_nonfinite_float64 : Sarek_ir_types.kernel -> bool

Whether the kernel contains a non-finite Float64 constant anywhere.

val kernel_uses_intrinsic : Stdlib.String.t -> Sarek_ir_types.kernel -> bool

Generic intrinsic-usage detection

Whether a kernel calls a named EIntrinsic anywhere. Generalizes the bespoke kernel_uses_copysign / kernel_uses_int_mod walkers for backends that must conditionally emit a helper for one intrinsic (e.g. the GLSL sarek_fmod helper for Float32.fmod/Float64.fmod, which GLSL has no builtin for). Matches on the intrinsic name only, ignoring the module path, so both the Float32 and Float64 spellings are detected. Inline native GPU code (SNative) is opaque text and is conservatively assumed to reference the intrinsic, mirroring the copysign/int_mod detectors.

Cooperative-matrix extraction — backlog-62 slice 3

kernel_uses Coopmat answers a yes/no question, and a yes/no answer is not enough for the launch gate. Device_optional capabilities are decided per CONFIGURATION: the RX 7900 XTX advertises fourteen and the same device that permits u8 x u8 + s32 -> s32 refuses f16 x f16 -> f16 saturating. So the gate needs the configurations a kernel actually asks for, and this is where they are collected — in spoc/ir, with no backend dependency, on the same traversal skeleton as every other detector rather than a bespoke walk.

val kernel_coopmat_ops : Sarek_ir_types.kernel -> Sarek_ir_types.coopmat_op list

Every cooperative-matrix operation in the kernel, in traversal order.

Uses the folder.fs hook, which is why that hook exists: a coopmat operation is neither an expression nor an element type, so no pre-slice-3 hook could see one.

val kernel_has_coopmat_op : Sarek_ir_types.kernel -> bool

Whether the kernel contains a cooperative-matrix OPERATION.

Strictly stronger than kernel_uses Coopmat, which is also true for a kernel that merely declares a TUint8 buffer. The distinction is what lets the GLSL backend emit the 8-bit extension without the cooperative-matrix one: a kernel that never reaches a multiply-add must not carry a shader requirement the device gate was never asked about.

val kernel_coopmat_configs : Sarek_ir_types.kernel -> Sarek_coopmat_types.config list

The distinct configurations a kernel's multiply-adds require.

Only Sarek_ir_types.coopmat_op.CM_muladd carries a configuration, and deliberately: a load or a store constrains the fragment's component type and shape but says nothing about which multiply-add the device must provide, and it is the multiply-add that VK_KHR_cooperative_matrix enumerates. A kernel that loads and stores a fragment without ever multiplying needs no advertised configuration at all, and reporting one would refuse it on a device that can run it perfectly well.