Module Sarek_execute.Launch_epilogue

Epilogue fusion in the launch window (specs/epilogue-fusion.md).

When an element-wise launch C arrives, it is merged into a pending launch P that writes the one vector I it consumes, if all of the following hold. The merged launch replaces P in its group and can itself host a further epilogue or join a horizontal group. The rewrite is Fusion_epilogue.

Conditions (vectors compared by identity, on one device).

The coverage verdict is cached by (P, the producer array, C, their dimensions, the values of the scalars the evaluation read).

module LW = Launch_window

Policy

include module type of struct include Epilogue_policy end
val switched_off : string list Stdlib.Atomic.t
val enabled : string -> bool
val set_backend_enabled : backend:string -> bool -> unit

Switch epilogue fusion off (false) or back on (true) for one backend. Switching on is not a grant: on CUDA, OpenCL and Vulkan fusion still needs Sarek_policy.deviation.Cross_kernel_contraction for both kernels.

val pair_cache : (string * string * string, int * Sarek_policy.decision) Stdlib.Hashtbl.t
val pair_lock : Stdlib.Mutex.t
val policy_decision : backend:string -> producer:string -> consumer:string -> Sarek_policy.decision
val decide_pair : backend:string -> producer:string -> consumer:string -> Sarek_policy.decision

Whether the epilogue of consumer may be merged into producer on backend: the backend is not switched off or inexact, and Sarek_policy.decide allows Sarek_policy.optimisation.Epilogue_fusion for the meet of both kernels' budgets.

val available : string -> bool

Whether fusion can run on backend under the process grants alone: not switched off, exact, and allowed by Sarek_policy.decide for the process budget (HIP and Metal; the others once Cross_kernel_contraction is granted at process scope). Kernel-scoped grants are not seen: the decision per pair is decide_pair.

Per-launch facts

val bound : 'arg LW.t -> 'arg LW.launch -> (string * int) list
val scalars : 'arg LW.t -> 'arg LW.launch -> (string * int32) list
val dims_of : 'a LW.launch -> Index_slice.dims
val shapes : (int, (Fusion_epilogue.consumer, string) Stdlib.result) Stdlib.Hashtbl.t
val shape_of : 'a LW.t -> 'b LW.launch -> (Fusion_epilogue.consumer, string) Stdlib.result

Coverage

val max_enumerated : int
type proof_counts = {
  1. affine : int;
  2. enumerated : int;
  3. over_budget : int;
}

How coverage verdicts were reached, since start-up: by the affine proof (Coverage_affine), by enumeration, or refused as over the enumeration bound.

val n_affine : int Stdlib.Atomic.t
val n_enumerated : int Stdlib.Atomic.t
val n_over_budget : int Stdlib.Atomic.t
val proof_counts : unit -> proof_counts
val threads : 'a LW.launch -> int
exception Stored_twice
val tag : 'a list -> 'a list -> (string * 'a) list
val slice : 'a LW.launch -> is_target:(string -> bool) -> record:(int32 -> unit) -> scalars:(string * int32) list -> string list
val enumerate : p:'a LW.launch -> p_array:Stdlib.String.t -> p_scalars:(string * int32) list -> c:'b LW.launch -> shape:Fusion_epilogue.consumer -> c_scalars:(string * int32) list -> bool * (string * string) list
val covers : p:'a LW.launch -> p_array:Stdlib.String.t -> p_scalars:(string * int32) list -> c:'b LW.launch -> shape:Fusion_epilogue.consumer -> c_scalars:(string * int32) list -> bool * (string * string) list
type verdict = {
  1. used : (string * string) list;
  2. values : int32 list;
  3. ok : bool;
}
val coverage_cache : (int * string * int * Index_slice.dims * Index_slice.dims, verdict list) Stdlib.Hashtbl.t
val lookup_value : p_scalars:('a * 'b) list -> c_scalars:('a * 'b) list -> (string * 'a) -> 'b option
val cached_covers : p:'a LW.launch -> p_array:Stdlib.String.t -> p_scalars:(string * int32) list -> c:'b LW.launch -> shape:Fusion_epilogue.consumer -> c_scalars:(string * int32) list -> bool

Merging

val fatal : exn -> bool
val fused_cache : (int * string * int * (string * string) list, Sarek_ir_types.kernel option) Stdlib.Hashtbl.t
val counter : int Stdlib.Atomic.t
val fused_ir : p:'a LW.launch -> p_array:string -> c:'b LW.launch -> shared:(string * string) list -> Sarek_ir_types.kernel option
val merged : 'arg LW.t -> 'arg LW.launch -> 'arg LW.launch -> fused:Sarek_ir_types.kernel -> p_array:string -> shared:(string * string) list -> tc:int -> 'arg LW.launch
val bufs : pick:(LW.access -> bool) -> LW.access list -> int list
val intermediate : 'a LW.launch -> 'b LW.launch -> int option
val merge_point : 'a LW.t -> 'a LW.group -> 'a LW.launch -> 'a LW.launch -> Fusion_epilogue.consumer -> int -> (string * (string * string) list) option
val permitted : 'a LW.launch -> 'b LW.launch -> bool
val try_member : 'a LW.t -> 'a LW.group -> 'a LW.launch -> 'a LW.launch -> Fusion_epilogue.consumer -> 'a LW.launch option
val one_dimensional : 'a LW.launch -> bool
val join : 'arg LW.t -> 'arg LW.launch -> bool

The window's epilogue hook (install).

val install : 'a LW.t -> unit

Enable epilogue fusion in w (on the backends enabled allows). An unexpected exception while deciding means no merge: the launch then takes the ordinary path.

val clear : unit -> unit

Forget every cached consumer shape, coverage verdict and merged kernel.