Module Sarek_execute.Coverage_affine

Epilogue coverage without enumeration, for affine store indices (specs/epilogue-fusion.md).

Answers the question Launch_epilogue asks of Index_slice (does P store each element of I at most once, at a non-negative index, and every element C writes?) in time independent of the launch size, when the relevant slice (Index_slice.relevant) of both kernels has this form:

Digits. The six thread and work-group indices start as one digit each, ranging over [0, extent) (a launch with an extent of 0 runs no thread and stores nothing). Dividing a value that is non-negative over the whole launch (guards are not taken into account, so a guarded (tid - 1) / 2 is Unknown) by 2^k, taking it modulo 2^k (or land (2^k - 1), or shifting it right by k) is exact when every coefficient but one digit's is a multiple of 2^k, that digit's coefficient is 1, the constant is a multiple of 2^k, and the digit's extent divides by 2^k or is at most 2^k: the digit d is then split into d = lo + 2^k * hi, lo in [0, 2^k), and the quotient and remainder are affine in the new digits. A warp-per-row index row = block * 8 + tid / 32 with the guard tid mod 32 == 0 has this form.

Solving. The digits with a non-zero coefficient in one guard must form a composite: sorted by magnitude, the coefficients are g times 1, e1, e1*e2, ... where e_i are the digits' extents. The composite L ranges bijectively over [0, product of extents) and the guard is an interval of L. Composites of different guards are equal or share no digit. Every other digit, and every composite whose interval is all of it, counts as free. P's store set is the image of this box. The index must be sum k_u * u + c over these units u. The image has each element exactly once and is one contiguous interval when the units with more than one value, sorted by k, form a mixed radix (first k is 1, each next is the previous times the previous unit's size). A k smaller than the span below it (zero included) repeats an element: P stores something twice, and the answer is false. C's store sets only need bounds.

Integer range. A digit of extent e >= 2 whose coefficient is a in a checked value satisfies |a| * (e - 1) <= 2^32 (every value is checked to stay in int32 over the launch), so |a| * e <= 2^33: products of a coefficient and a digit value, and composite weights, stay far below OCaml's 63-bit max_int. scale refuses a constant product that could overflow.

Any other form is Unknown: the caller enumerates instead.

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

Digits and affine values

type aff = {
  1. co : (int * int) list;
  2. c : int;
}
type space = {
  1. dims : int array;
    (*

    block x y z, grid x y z

    *)
  2. extent : (int, int) Stdlib.Hashtbl.t;
    (*

    of each digit; positions are 0..5

    *)
  3. split : (int, int * int * int) Stdlib.Hashtbl.t;
    (*

    digit -> (lo, hi, m): digit = lo + m * hi

    *)
  4. mutable next : int;
    (*

    next fresh digit

    *)
  5. env : env;
}
and env = {
  1. vars : (int, aff) Stdlib.Hashtbl.t;
    (*

    let-bound relevant variables by id

    *)
  2. scalars : (string * int32) list;
  3. mutable used : string list;
    (*

    scalar parameters read

    *)
}
val lo32 : int
val hi32 : int
val const : int -> aff
val extent : space -> int -> int
val merge : ('a * int) list -> ('a * int) list
val norm : space -> aff -> aff
val range : space -> aff -> int * int
val checked : space -> aff -> aff
val add : space -> aff -> aff -> aff
val scale : space -> int -> aff -> aff
val as_const : space -> aff -> int option
val mul : space -> aff -> aff -> aff
val fresh : space -> int -> int
val digit : int -> aff
val divmod : space -> aff -> int -> aff * aff
val power_of_two : space -> aff -> int
val shift : space -> aff -> int
val mask : space -> aff -> aff -> aff
val position : space -> int -> aff
val global : space -> int -> aff
val intrinsic : space -> string -> aff
val scalar : space -> string -> aff
val eval : space -> Sarek_ir_types.expr -> aff

Guards

val guard : space -> Sarek_ir_types.expr -> aff list

Store sites

type site = {
  1. guards : aff list;
  2. index : aff;
}
val sites : space -> is_target:(string -> bool) -> Index_slice.Ids.t -> aff list -> Sarek_ir_types.stmt -> site list -> site list
val space : dims:Index_slice.dims -> scalars:(string * int32) list -> space
val store_sites : is_target:(string -> bool) -> scalars:(string * int32) list -> dims:Index_slice.dims -> Sarek_ir_types.stmt -> space * site list

Composites

type composite = {
  1. digits : (int * int) list;
  2. lo : int;
  3. hi : int;
}
val size_of : space -> (int * 'a) list -> int
val composite_of : space -> aff -> ((int * int) list * int) option
val cdiv : int -> int -> int
val fdiv : int -> int -> int
val shares : ('a * 'b) list -> ('a * 'c) list -> bool
val constrain : space -> composite list -> aff -> composite list option
val box : space -> site -> composite list option
val live : space -> int list
val units : space -> site -> ((int * int * int) list * int) option

Store sets

val bounds_of : ((int * int * int) list * int) -> int * int
val bounds : space -> site -> (int * int) option
type stores =
  1. | Exactly_once of (int * int) option
  2. | Repeats
val exactly_once : space -> site -> stores
val covers : p_body:Sarek_ir_types.stmt -> p_array:string -> p_scalars:(string * int32) list -> p_dims:Index_slice.dims -> c_body:Sarek_ir_types.stmt -> c_written:string list -> c_scalars:(string * int32) list -> c_dims:Index_slice.dims -> bool * (string * string) list

covers returns the verdict and the scalar parameters it read (("p", n) for P's, ("c", n) for C's), as Launch_epilogue's enumeration does. Raises Unknown when either kernel is outside the affine form.