Module Sarek_ir_affine

Affine index forms for access-pattern analysis.

An index expression such as (row * k) + i + off is put in the form Σ cᵢ·aᵢ + c₀ where each atom aᵢ is a loop variable or a thread-varying sub-expression (thread_idx_x, global_thread_id / n) and each coefficient cᵢ is a uniform polynomial: an integer combination of products of uniform factors (scalar parameters, block_dim_*, array lengths, or an opaque uniform sub-expression such as (k + 15) / 16).

Uniform means "the same for every thread of a launch and every iteration of every loop", which is what a stride must be for a matrix unit to load a tile with one base address and one leading dimension.

The form is purely syntactic: two atoms are the same atom when their let-inlined expressions are structurally equal. That is sound for the use made of it (deciding which index varies with which loop) because every atom is a pure function of the thread and of the loop variables. It is not a decision procedure for index equality: tid / n * n + tid mod n is not recognised as tid.

Used by Sarek_matrix_pattern (specs/matrix-pattern-advice.md).

Uniform polynomials

type monomial = {
  1. coef : int;
  2. factors : Sarek_ir_types.expr list;
}
type poly = monomial list

A normalised polynomial: monomials sorted by factor list, no zero coefficient, no repeated factor list. [] is zero.

val poly_const : int -> poly
val poly_is_zero : poly -> bool
val poly_is_one : poly -> bool
val poly_equal : poly -> poly -> bool
val poly_mul : poly -> poly -> poly
val poly_div : poly -> poly -> poly option

poly_div p d is the exact quotient of p by the single-monomial d, or None (also when d is not a single monomial).

val poly_to_expr : poly -> Sarek_ir_types.expr
val poly_constant : poly -> int option

The value of a constant polynomial (Some 0 for zero), None otherwise.

Affine forms

type t = {
  1. terms : (Sarek_ir_types.expr * poly) list;
  2. const : poly;
}

terms is sorted by atom and has no zero coefficient.

type leaf =
  1. | Uniform_leaf
    (*

    a scalar kernel parameter or another uniform value

    *)
  2. | Atom_leaf
    (*

    a loop variable: may appear linearly, never non-linearly

    *)
  3. | Rejected_leaf
    (*

    a mutable variable, or one the caller cannot classify

    *)

How a variable that the analysis meets in an index is to be read.

type env = {
  1. classify_var : Sarek_ir_types.var -> leaf;
  2. definition : Sarek_ir_types.var -> Sarek_ir_types.expr option;
}

classify_var classifies a variable that is not let-bound; definition v is the defining expression of an immutable let-bound variable, which is inlined before analysis.

type error =
  1. | Not_affine of Sarek_ir_types.expr
    (*

    a product, division, modulo, shift or bit operation involving a loop variable

    *)
  2. | Data_dependent of Sarek_ir_types.expr
    (*

    the index reads memory

    *)
  3. | Unknown_var of Sarek_ir_types.var
    (*

    a Rejected_leaf variable

    *)
  4. | Wide_index of Sarek_ir_types.expr
    (*

    the expression computes in 64 bits (an int64 constant, variable or cast): the form is int-sized and the rewrite addresses with int32, so it is refused rather than narrowed

    *)

The canonical atom of a variable (the lowering marks a loop variable mutable at its binder and immutable at its uses). Use it to look up a loop variable's coefficient.

val is_pure_intrinsic : string -> bool

Intrinsics that are pure functions of their arguments (uniform on uniform arguments). Any other intrinsic is treated as thread-varying.

The let-inlined expression, as atoms are compared.

val involves_int64 : Sarek_ir_types.expr -> bool

Whether an (inlined) expression computes in 64 bits: an int64 constant, an int64-typed variable, a cast to int64, or a 64-bit shift. of_expr refuses such an expression (Wide_index). The test is conservative: a deliberate ECast (TInt32, e) of an int64 e is refused too, because the narrowing is what the refusal exists to prevent.

val of_expr : env -> Sarek_ir_types.expr -> (t, error) Stdlib.result
val coefficient : t -> Sarek_ir_types.expr -> poly

The coefficient of an atom (poly_is_zero when absent).

val atoms : t -> Sarek_ir_types.expr list
val of_terms : (Sarek_ir_types.expr * poly) list -> t

The form Σ cᵢ·aᵢ with no constant term.

val without : t -> Sarek_ir_types.expr list -> t

The form without the terms of the given atoms.

val is_uniform : t -> bool

No atom: the value is uniform.

val to_expr : t -> Sarek_ir_types.expr
val error_expr : error -> Sarek_ir_types.expr option

The offending sub-expression of an error, when there is one.