Module Sarek_execute.Launch_window

Deferred launch window with horizontal fusion (specs/launch-window.md).

Execute submits JIT launches here instead of issuing them. Pending launches are kept in groups, in submission order. A new launch joins the newest pending group it is compatible with, provided it conflicts with no launch in that group or in any later one. Otherwise it opens a group of its own. A group of several launches is issued as one horizontally fused kernel (Fusion_horizontal); a group of one is the original launch.

Conflicts. Each launch carries the buffers its vector arguments bind and whether the kernel reads or writes each one (Kernel_access). A buffer is named by its vector's id: every device buffer belongs to exactly one vector (only Transfer.ensure_buffer installs one, in that vector's own table), so on one device two arguments share memory exactly when they are the same vector. Device addresses are not used: Vulkan reports none and OpenCL reports object handles. Two launches on the same device conflict when they bind the same buffer and at least one of the two accesses writes. A launch is never moved before a launch it conflicts with, so every conflicting pair is issued in submission order.

Flushing. flush issues every pending group. Execute installs it as the Spoc_core.Launch_hooks flusher, so every synchronisation point flushes before it acts. At most max_groups groups are pending: a new group beyond that issues the oldest one first.

Extension point. A launch is a plain record. A later stage that specialises a launch (for instance a launch-time kernel rewrite) can map launch values before submit without touching grouping or flushing.

This module depends only on the IR and Stdlib, so it builds for jsoo too.

Launch records

type access = {
  1. buf : int;
  2. read : bool;
  3. write : bool;
}

One bound vector argument: its buffer (the vector's id) and how the kernel uses it.

type 'arg launch = {
  1. device_id : int;
  2. backend : string;
  3. ir : Sarek_ir_types.kernel;
  4. ir_id : int;
    (*

    ir_info.ir_id of ir

    *)
  5. shared_bytes : int option;
    (*

    shared memory the kernel declares when it can be fused horizontally (ir_info.fusable); None when it cannot

    *)
  6. max_grid_x : int;
    (*

    the device's grid-size limit in x

    *)
  7. max_shared : int;
    (*

    the device's shared memory per work-group

    *)
  8. max_arg_slots : int;
    (*

    the backend's kernel argument slots (arg_slot_limit); max_int when unbounded

    *)
  9. args : 'arg list;
  10. block : Dims.dims;
  11. grid : Dims.dims;
  12. shared_mem : int;
  13. accesses : access list;
  14. issue : unit -> unit;
    (*

    issue this launch now, unfused

    *)
  15. issue_fused : Sarek_ir_types.kernel -> 'arg list -> Dims.dims -> unit;
    (*

    issue a fused kernel on this launch's device with this launch's block, the given arguments and grid

    *)
}

Per-kernel facts

type ir_info = {
  1. ir_id : int;
  2. param_accesses : (string * Kernel_access.access) list;
  3. fusable : (int, string) Stdlib.result;
    (*

    shared bytes, or why not fusable

    *)
}

Window state

type 'arg group = {
  1. g_device : int;
  2. mutable rev_members : 'arg launch list;
    (*

    newest first

    *)
  3. mutable size : int;
  4. mutable grid_x : int;
    (*

    sum of the members' grid x

    *)
  5. mutable shared : int;
    (*

    sum of the members' shared memory

    *)
}
val members : 'a group -> 'a launch list
val key_of : 'a group -> int list
type stats = {
  1. submitted : int;
    (*

    JIT launches submitted through Execute

    *)
  2. issued : int;
    (*

    backend launches actually made

    *)
  3. fused_groups : int;
    (*

    groups of two or more launches issued fused

    *)
  4. fused_built : int;
    (*

    distinct fused kernels built

    *)
  5. fused_refused : int;
    (*

    member sequences refused (build or launch failed)

    *)
}
type 'arg t = {
  1. lock : Stdlib.Mutex.t;
  2. mutable groups : 'arg group list;
    (*

    oldest first

    *)
  3. mutable epilogue : 'arg t -> 'arg launch -> bool;
    (*

    called under lock before grouping: true when the launch was merged into a pending launch as its epilogue (Launch_epilogue)

    *)
  4. int32_arg : int32 -> 'arg;
  5. arg_buf : 'arg -> int option;
    (*

    the vector id a vector argument binds

    *)
  6. arg_int32 : 'arg -> int32 option;
    (*

    the value of an int32 scalar

    *)
  7. set_pending : int -> unit;
  8. guard : (unit -> unit) -> unit;
  9. info_key : ir_info Ir_memo.key;
    (*

    this window's slot of each IR's facts

    *)
  10. mutable next_ir_id : int;
  11. cache : (int list, Sarek_ir_types.kernel) Stdlib.Hashtbl.t;
    (*

    member ir_ids -> fused kernel

    *)
  12. refused : (int list, unit) Stdlib.Hashtbl.t;
    (*

    member sequences never fused

    *)
  13. c_submitted : int Stdlib.Atomic.t;
  14. c_issued : int Stdlib.Atomic.t;
  15. c_groups : int Stdlib.Atomic.t;
  16. c_built : int Stdlib.Atomic.t;
    (*

    never reset: it also names fused kernels without a content name

    *)
  17. c_refused : int Stdlib.Atomic.t;
}
val create : int32_arg:(int32 -> 'arg) -> arg_buf:('arg -> int option) -> arg_int32:('arg -> int32 option) -> set_pending:(int -> unit) -> guard:((unit -> unit) -> unit) -> 'arg t

create ~int32_arg ~set_pending ~guard makes an empty window. int32_arg builds a scalar argument for the fused kernel's position parameters. set_pending is told the number of pending launches after each change. guard f runs f (the issuing of pending groups) with the synchronisation hooks disabled on this domain. Each window has its own slot in the Ir_memo record of every IR it sees; the slot lives as long as that IR, not as long as the window.

Configuration

val max_groups : int Stdlib.ref
val max_members : int Stdlib.ref
val max_cached : int
include module type of struct include Window_gate end
type refusal = Window_gate.refusal =
  1. | Switched_off
    (*

    Execute.Window.set_enabled false

    *)
  2. | Restricted
    (*

    SAREK_OPT_DISABLE=window

    *)
  3. | Cpu_device
    (*

    host-pointer buffers alias host memory, and distinct vectors can alias each other through it, so buffer identity is no longer the vector's

    *)
  4. | Not_on of string
    (*

    the backend (framework name) is not in backends

    *)
  5. | Streamed
    (*

    a launch on a stream

    *)
  6. | Soa_params
    (*

    structure-of-arrays arguments

    *)

Why a launch is not deferred, in the order the reasons are checked.

val enabled : bool Stdlib.Atomic.t
val default_backends : string list

Backends deferred by default (specs/launch-window.md); Native, Interpreter and WebGPU stay immediate. set_backend_enabled opts in or out.

val backends : string list Stdlib.Atomic.t
val refusal : framework:string -> is_cpu:bool -> streamed:bool -> soa_params:'a list -> refusal option

refusal ~framework ~is_cpu ~streamed ~soa_params: None when a JIT launch on a device of framework is deferred.

val defers : string -> bool

true when launches on backend are deferred.

val defers_device : framework:string -> is_cpu:bool -> bool

true when JIT launches on a device of framework are deferred: its backend's window is on and it is not a CPU device.

include module type of struct include Launch_limits end
val slot_limits : (string * int) list Stdlib.Atomic.t
val arg_slot_limit : string -> int

Kernel argument slots a launch's signature may use on backend: Metal binds every argument at one of 31 [buffer(n)] indices. A backend with no limit configured here is unbounded as far as the window is concerned.

val set_limit : ('a * 'b) list Stdlib.Atomic.t -> backend:'a -> 'b -> unit
val set_arg_slot_limit : backend:string -> int -> unit

Set the argument slot limit of backend (for tests and for backends whose limit becomes known).

val push_limits : (string * int) list Stdlib.Atomic.t
val push_constant_limit : string -> int

Bytes of push constants a launch's signature may use on backend: Vulkan declares one 128-byte range in every pipeline layout and refuses a launch whose block is larger. A backend with no limit configured here is unbounded as far as the window is concerned.

val set_push_constant_limit : backend:string -> int -> unit

Set the push-constant byte limit of backend (for tests and for backends whose limit becomes known).

val push_bytes_ok : backend:string -> Sarek_ir_types.decl list list -> bool

true when a group whose members have these parameter lists keeps its kernel's push-constant block within backend's limit: the fused kernel's block for two or more members, the launch's own for one. Answers at once when the backend has no limit.

Per-IR facts

val ir_info_locked : 'a t -> Sarek_ir_types.kernel -> ir_info

As ir_info, for a caller that already holds w.lock.

val ir_info : 'a t -> Sarek_ir_types.kernel -> ir_info

Grouping

val accesses_conflict : access -> access -> bool
val conflicts : 'a launch -> 'b launch -> bool
val group_conflicts : 'a group -> 'b launch -> bool
val same_dims : Dims.dims -> Dims.dims -> bool
val joinable : 'a launch -> bool
val arg_slots : 'a t -> 'a launch -> int

Slots one launch's arguments take in a kernel signature: a vector is a buffer plus its length, a scalar one slot.

val fused_slots : 'a t -> 'a launch list -> int

Slots of the signature of the fused kernel of members: every member's arguments plus (offset, grid size) per member.

val slots_ok : 'a t -> 'a launch list -> bool

true when members, issued as one group, stay within the backend's argument slots: the fused signature for two or more, the launch's own signature for one.

val params_of : 'a launch -> Sarek_ir_types.decl list
val push_ok : 'a launch list -> bool

The group's push-constant block fits the backend (push_bytes_ok).

val signature_ok : 'a t -> 'a launch list -> bool

The backend can issue members as one group: slots_ok and push_ok.

val fits : 'a t -> 'a group -> 'a launch -> bool
val compatible : 'a t -> 'a group -> 'a launch -> bool
val target : 'a t -> 'a launch -> 'a group option

Issuing

val bump : int Stdlib.Atomic.t -> unit
val refuse : 'a t -> int list -> unit
val fallback_names : int Stdlib.Atomic.t
val next_fallback_name : unit -> int
val fused_kernel : 'a t -> int list -> 'b launch list -> Sarek_ir_types.kernel option
val fused_args : 'a t -> 'a launch list -> 'a list * int
val fatal : exn -> bool
val each : ('a -> unit) -> 'a list -> unit
val issue_singly : 'a t -> 'b launch list -> unit
val issue_fused : 'a t -> int list -> 'a launch list -> 'a launch -> unit
val issue_group : 'a t -> 'a group -> unit
val count_pending : 'a group list -> int
val issue_all : 'a t -> 'a group list -> unit

Public operations

val flush : 'a t -> unit

Issue every pending launch, in order. The pending count stays non-zero until they are all issued, so a synchronisation point reached meanwhile on another domain waits for this flush (on w.lock) instead of going ahead (FR-007).

val submit : 'arg t -> 'arg launch -> unit

Defer l. When l needs a new group and the window is full, the oldest group is issued first; if that raises, l is not submitted and the exception propagates (the launching call fails before its own launch is accepted, so its vectors' state is unchanged).

val touches : 'a t -> device_id:int -> buf:int -> bool

true when a pending launch on device_id binds vector buf.

val count_immediate : 'a t -> unit

Count a launch that bypassed the window (issued immediately).

val stats : 'a t -> stats
val reset_stats : 'a t -> unit
val clear_cache : ?also:(unit -> unit) -> 'a t -> unit

Forget every fused kernel and refusal (the next fused group rebuilds), and run also in the same critical section.

val set_enabled : 'a t -> bool -> unit
val set_backend_enabled : 'a t -> backend:string -> bool -> unit