Module Sarek_execute.Index_slice

Host evaluation of a kernel's store indices for one launch (specs/epilogue-fusion.md).

stores answers "which elements of array a does this launch write, and how often", by running, for every thread of the launch, only the integer computations that the stores to a depend on: the index expressions, the conditions and loop bounds around the stores, and, transitively, the variables those read. These may use thread and work-group indices, launch dimensions, int32 scalar parameters and int32 constants. Anything else (an element read, a float, a match, an int64) makes the evaluation fail, which callers must treat as "unknown": never as "no store".

The slice. A variable is relevant when a store index, a condition or loop bound enclosing a store or a return, or an assignment to a relevant variable reads it (fixpoint). Statements that neither contain a store to a, a return, nor assign a relevant variable are skipped: they cannot change which stores execute. A return stops the thread, so the conditions around it are relevant too.

Values follow int32 two's-complement arithmetic, as the device does. Division or remainder by zero fails the evaluation.

exception Unknown of string
val unknown : ('a, unit, string, 'b) Stdlib.format4 -> 'a

Relevance

module Ids : sig ... end
val expr_vars : Ids.t -> Sarek_ir_types.expr -> Ids.t
val has_target : is_target:(string -> bool) -> Sarek_ir_types.stmt -> bool
val assigned : Ids.t -> Sarek_ir_types.stmt -> Ids.t
val matters : is_target:(string -> bool) -> Ids.t -> Sarek_ir_types.stmt -> bool
val grow : is_target:(string -> bool) -> Ids.t -> Sarek_ir_types.stmt -> Ids.t
val relevant : is_target:(string -> bool) -> Sarek_ir_types.stmt -> Ids.t

Compilation to closures

type state = {
  1. slots : int32 array;
  2. mutable tx : int;
  3. mutable ty : int;
  4. mutable tz : int;
  5. mutable bx : int;
  6. mutable by : int;
  7. mutable bz : int;
  8. dims : int array;
    (*

    block x y z, grid x y z

    *)
  9. mutable stop : bool;
    (*

    the thread executed return

    *)
  10. mutable steps : int;
    (*

    loop iterations and threads so far, this run

    *)
  11. budget : int;
    (*

    bound on steps

    *)
}
type cenv = {
  1. slot_of : (int, int) Stdlib.Hashtbl.t;
    (*

    var_id -> slot

    *)
  2. scalars : (string, int32) Stdlib.Hashtbl.t;
    (*

    int32 scalar parameters

    *)
  3. mutable used : string list;
    (*

    scalar parameters the slice read

    *)
}
val b2i : bool -> int32
val i32 : int -> int32
val position : string -> state -> int32
val check_div : int32 -> int32
val binop : Sarek_ir_types.binop -> (state -> int32) -> (state -> int32) -> state -> int32
val cexpr : cenv -> Sarek_ir_types.expr -> state -> int32
type cstmt = state -> unit
val seq : cstmt list -> cstmt
val max_steps : int
val tick : state -> unit
val cstmt : is_target:(string -> bool) -> record:(int32 -> unit) -> cenv -> Ids.t -> Sarek_ir_types.stmt -> cstmt

Running a launch

type dims = {
  1. block : int * int * int;
  2. grid : int * int * int;
}
val run : ?bind:(Sarek_ir_types.var * int32) list -> ?budget:int -> is_target:(string -> bool) -> record:(int32 -> unit) -> scalars:(string * int32) list -> dims:dims -> Sarek_ir_types.stmt -> string list

Run body for every thread of a launch of dims. is_target selects the arrays whose stores are recorded; record receives each stored index (and, for stores in a thread, in program order). scalars gives the int32 scalar parameters by name; bind pre-binds variables (var, value) before each thread. Returns the scalar parameters the evaluation read. Raises Unknown when a relevant computation is not an integer function of the launch, or when more than budget threads plus loop iterations (default max_steps) would run.