Module Sarek_execute.Fusion_epilogue

Epilogue fusion: an element-wise consumer inlined after the producer's stores (specs/epilogue-fusion.md).

Producer P stores elements of an intermediate array I; consumer C is element-wise (let i = global_thread_id in body) and reads I only at its own index i. In the fused kernel, every store I[e] <- v of P becomes

  let t = e in
  I[t] <- v ;
  if t < tc then (let i = t in body)

where tc is the number of threads C was launched with. The thread that stored I[t] then computes C's element t, reading I[t] back from memory just after writing it: the same operations in the same order, with the same values. What makes this equal to running P then C is checked elsewhere: P stores each element of C's domain exactly once (Index_slice, per launch), C reads nothing else P writes, and C writes nothing P or C reads elsewhere (Launch_window). The one remaining difference is that the device compiler now sees P's store and C's load in one function: a compiler allowed to contract a*b+c across them could fuse P's last multiply into C's first add. It can also share a subexpression of P and C that the unfused C contracts. That is why this fusion is on by default only where contraction is off in every kernel (-ffp-contract=off on HIP).

Names: P's are prefixed p_, C's c_, the new ones ep_, so no name can capture another. C's parameters bound to the same vector as one of P's are mapped onto P's parameter (in particular I itself).

Consumer shape

type consumer = {
  1. index : Sarek_ir_types.var;
    (*

    C's i

    *)
  2. body : Sarek_ir_types.stmt;
    (*

    the body under let i = global id in

    *)
  3. reads_at_index : string list;
    (*

    array params that C reads only at index i (or not at all)

    *)
  4. written : string list;
    (*

    array params C writes (always at i)

    *)
}
exception Refuse of string
val refuse : ('a, unit, string, 'b) Stdlib.format4 -> 'a
val global_ids : string list
val thread_intrinsics : string list
val accesses_stmt : (string * Sarek_ir_types.expr option * bool) list -> Sarek_ir_types.stmt -> (string * Sarek_ir_types.expr option * bool) list
val accesses_expr : (string * Sarek_ir_types.expr option * bool) list -> Sarek_ir_types.expr -> (string * Sarek_ir_types.expr option * bool) list
val is_index : Sarek_ir_types.var -> Sarek_ir_types.expr option -> bool
val consumer_shape : Sarek_ir_types.kernel -> consumer
val consumer : Sarek_ir_types.kernel -> (consumer, string) Stdlib.result

Rewrite

val ep_var : int -> string -> Sarek_ir_types.elttype -> Sarek_ir_types.var
type ids = {
  1. c_shift : int;
  2. tmp_id : int;
  3. tc_id : int;
}
val ids : producer:Sarek_ir_types.kernel -> consumer:Sarek_ir_types.kernel -> ids
val producer_names : Ir_rename.ctx
val consumer_names : shared:(string * string) list -> c_shift:int -> Ir_rename.ctx
val insert : array:string -> tmp_id:int -> after:(Sarek_ir_types.expr -> Sarek_ir_types.stmt) -> Sarek_ir_types.stmt -> Sarek_ir_types.stmt
exception Incompatible of string
val fused_body : producer:Sarek_ir_types.kernel -> p_array:string -> c:consumer -> cn:Ir_rename.ctx -> ids:ids -> tc:Sarek_ir_types.var -> Sarek_ir_types.stmt
val merged_declarations : Sarek_ir_types.kernel -> Sarek_ir_types.kernel -> (string * (string * Sarek_ir_types.elttype) list) list * (string * (string * Sarek_ir_types.elttype list) list) list * Sarek_ir_types.helper_func list
val fuse : name:string -> producer:Sarek_ir_types.kernel -> p_array:string -> consumer:Sarek_ir_types.kernel -> shared:(string * string) list -> Sarek_ir_types.kernel

fuse ~name ~producer ~p_array ~consumer ~shared inlines consumer after producer's stores to p_array. shared maps every consumer array parameter bound to the same vector as a producer parameter (in particular the one bound to p_array's vector) to that producer parameter. Parameters: P's, then C's that are not shared, then ep_tc.