Module Sarek_codegen.Sarek_ir_inline_vec

val has_vec_param : Sarek_ir_types.helper_func -> bool

Does hf take at least one vector parameter? Only such helpers are inlined; scalar-only helpers stay real device functions on both backends.

type sink =
  1. | Assign of Sarek_ir_types.lvalue
    (*

    store the returned value into lvalue

    *)
  2. | Return
    (*

    return the value from the enclosing function

    *)
  3. | Discard
    (*

    unit-typed call kept only for its side effects

    *)

Where the result of an inlined call must go.

type ctx = {
  1. vec_helpers : (string, Sarek_ir_types.helper_func) Stdlib.Hashtbl.t;
  2. backend : string;
  3. mutable stack : string list;
  4. mutable temp_counter : int;
}

Inliner state: the vector-parameter helpers keyed by name, the backend name for error messages, and the stack of helpers currently being inlined (the recursion guard).

val fail : ctx -> string -> string -> 'a
val fresh_temp_name : ctx -> string
val fresh_local_name : ctx -> string

Fresh name for a helper local that must be alpha-renamed to avoid capturing a substituted-in buffer reference. sarek_-prefixed (a reserved, collision-proof namespace on every backend).

val as_vec_call : ctx -> Sarek_ir_types.expr -> (string * Sarek_ir_types.expr list) option

Is e a call to a vector-parameter helper? Returns its name + args.

val sub : ('a * 'a) list -> 'a -> 'a

subst maps a name to its replacement (identity if absent). Two uses:

  • vector-parameter substitution: maps a helper's vector-parameter name to the call site's buffer name, so every buffer access (EArrayRead, EArrayLen, LArrayElem, and any bare EVar) reads the caller's global buffer directly. Here ~rename_binders:false — the vector parameter has no binder in the body, and a shadowing local of the same name must NOT be touched.
  • alpha-renaming: maps a helper local that would collide with a substituted-in buffer name to a fresh name. Here ~rename_binders:true so the binder occurrences (SLet/SLetMut/SFor variables and match pattern bindings) are rewritten too, not only the uses.
val sub_pattern : rename_binders:bool -> (string * string) list -> Sarek_ir_types.pattern -> Sarek_ir_types.pattern
val subst_expr : rename_binders:bool -> (string * string) list -> Sarek_ir_types.expr -> Sarek_ir_types.expr
val subst_lvalue : rename_binders:bool -> (string * string) list -> Sarek_ir_types.lvalue -> Sarek_ir_types.lvalue
val sub_binder : rename_binders:bool -> (string * string) list -> Sarek_ir_types.var -> Sarek_ir_types.var

Rewrite a binder's variable, renaming it only when ~rename_binders.

val subst_stmt : rename_binders:bool -> (string * string) list -> Sarek_ir_types.stmt -> Sarek_ir_types.stmt
val collect_binders : Sarek_ir_types.stmt -> string list

All local binder names introduced anywhere in s: SLet/SLetMut/SFor variables and match-pattern bindings. Used to detect names that would capture a substituted-in buffer reference.

val rewrite_returns : ctx -> sink -> Sarek_ir_types.stmt -> Sarek_ir_types.stmt

Replace each (tail-position) SReturn e in the spliced helper body with the action dictated by sink. A return found in a non-tail position (e.g. the middle of a sequence) is refused — it would fall through into the code after the splice, changing semantics.

val assert_no_return : ctx -> Sarek_ir_types.stmt -> unit

Verify s contains no SReturn anywhere (used for non-tail sub-statements of an inlined body).

val splice_call : ctx -> sink -> string -> Sarek_ir_types.expr list -> Sarek_ir_types.stmt

Build the statement that computes f args and routes the result to sink. Scalar parameters are bound with SLet; vector parameters are substituted by name; the whole body is wrapped in an SBlock so its locals never leak into (or collide across) call sites. The spliced body is itself run through inline_stmt so a vector helper that calls another vector helper is fully resolved.

val hoist_expr : ctx -> Sarek_ir_types.expr -> (Sarek_ir_types.var * string * Sarek_ir_types.expr list) list * Sarek_ir_types.expr

Lift every vector-helper call nested inside expression e into a preceding result temporary, so it can be spliced as a statement. Returns the list of (temp var, helper name, args) hoisted (in evaluation order) and the rewritten expression referring to the temporaries. Direct-position calls (handled by inline_stmt) are not reached here.

val with_hoisted : ctx -> (Sarek_ir_types.var * string * Sarek_ir_types.expr list) list -> Sarek_ir_types.stmt -> Sarek_ir_types.stmt

Wrap core (a statement using the hoisted temporaries) with a mutable result temporary and an inlined splice for each hoisted call, in order.

Rewrite one statement, inlining vector-helper calls. Direct-position calls (the whole RHS of an assignment / let / return / expression statement) are spliced without a temporary; calls nested inside other expressions are hoisted first.

val inline_vec_helpers : backend:string -> Sarek_ir_types.kernel -> Sarek_ir_types.kernel

Inline all vector-parameter helpers into the kernel body and into the remaining (scalar-only) helper bodies, then drop the inlined helpers. A no-op when the kernel has no vector-parameter helper.