Module Sarek_convergence

module StringSet : sig ... end
type exec_mode =
  1. | Converged
    (*

    All threads in workgroup execute this code

    *)
  2. | Diverged
    (*

    Threads may have taken different branches

    *)

Execution mode for convergence analysis

type sync_kind =
  1. | Barrier_sync
    (*

    a block barrier or a superstep

    *)
  2. | Warp_sync of string
    (*

    a warp collective, by name

    *)

The synchronisation point a helper function contains: what a call to it reaches.

type helper_def = {
  1. hd_self : string;
    (*

    the name it was defined under

    *)
  2. hd_params : string list;
  3. hd_body : Sarek_typed_ast.texpr;
  4. hd_defs : (string * helper_def) list;
    (*

    The helpers in scope at its definition (itself excluded): a call resolves names lexically, as the generated code does, not in the caller's scope.

    *)
  5. hd_captured : StringSet.t;
    (*

    Thread-varying variables a nested helper captures from its definition site; empty for a module-level function.

    *)
}

A helper function in scope: what a call to it can be analysed against.

type variance = {
  1. strict : bool;
  2. defs : (string * helper_def) list;
    (*

    helpers in scope, innermost first

    *)
  3. visiting : string list;
    (*

    helpers being unfolded: recursion guard

    *)
}

How variance is computed. The legacy analysis is the first, cheap pass. The strict analysis is applied to every synchronisation point (warp collectives, whose missed divergence is a silent wrong answer, and barriers, whose missed divergence is a deadlock or undefined behaviour): it also follows the value of an if/match, a helper result that reads the thread, and a mutable assigned under thread-varying control flow (specs/convergence-semantics.md section 3).

val legacy_variance : variance
type assigned_entry = {
  1. ae_body : Sarek_typed_ast.texpr;
  2. ae_env : StringSet.t;
  3. ae_vi : variance;
  4. ae_result : bool;
}

Memo of assigned_varying, per mutable name: the scan of an inner let mut body is repeated by every enclosing let mut, which is exponential in their nesting. An entry is valid for one body (physical equality), one variance environment and one set of visible helpers. Cleared at each kernel check.

val assigned_memo : (string, assigned_entry list) Stdlib.Hashtbl.t
val find_assigned : variance -> StringSet.t -> string -> Sarek_typed_ast.texpr -> assigned_entry option
type ctx = {
  1. mode : exec_mode;
  2. varying_vars : StringSet.t;
    (*

    Variables known to carry thread-varying values at this point

    *)
  3. sync_helpers : (string * sync_kind) list;
    (*

    Helper functions in scope that contain a synchronisation point, transitively (specs/convergence-semantics.md, rule H-2): a call to one in diverged flow is an error at the call site.

    *)
  4. variance : variance;
  5. call_stack : string list;
    (*

    Helpers whose body is being re-checked at a call (strict analysis)

    *)
}

Analysis context

val init_ctx : ctx

Initial context - start converged

val diverge : ctx -> ctx

Enter diverged mode, preserving known-varying variable set

val is_thread_varying_intrinsic : Sarek_env.intrinsic_ref -> bool

Check if an intrinsic ref is thread-varying using core primitives

val is_barrier_ref : Sarek_env.intrinsic_ref -> bool

Check if an intrinsic ref is a barrier/convergence point

val has_reachable_return : Sarek_typed_ast.texpr -> bool

True if texpr e contains a TEReturn reachable on some execution path

The direct sub-expressions of te, in evaluation order.

val helper_name : Sarek_typed_ast.texpr -> string option
val body_defs : helper_def -> (string * helper_def) list

The helpers in scope inside the body of def: those of its definition site plus itself.

val helper_call_env : helper_def -> bool list -> StringSet.t

The environment a helper body is evaluated in at a call: what it captured, minus its own parameters, plus the parameters whose argument is varying.

val varying : variance -> StringSet.t -> Sarek_typed_ast.texpr -> bool

Check if an expression's value varies per-thread.

Thread-varying (different per thread):

  • thread_idx_x, thread_idx_y, thread_idx_z
  • global_thread_id
  • Expressions derived from above
  • Array accesses with thread-varying index

Uniform (same for all threads in workgroup):

  • Constants and literals
  • block_idx_x/y/z (same for whole block)
  • block_dim_x/y/z, grid_dim_x/y/z
  • Expressions derived only from uniform values

With vi.strict the value of control flow, helper results and mutables assigned in divergence are followed too (see variance).

val alias_def : variance -> string -> Sarek_typed_ast.texpr -> variance

let g = f: g is another name for the helper f (strict analysis).

Register a nested helper.

val helper_result_varying : variance -> StringSet.t -> Sarek_typed_ast.texpr -> Sarek_typed_ast.texpr list -> bool

Strict analysis: does a call to a known helper return a thread-dependent value, given which arguments vary? The body is unfolded once per call (a recursion guard stops cycles).

val assigned_varying : variance -> StringSet.t -> string -> Sarek_typed_ast.texpr -> bool

Strict analysis: is the mutable name assigned, anywhere in body, a thread-varying value or under thread-varying control flow? Then its value after the assignment differs per thread.

val assigned_varying_scan : variance -> StringSet.t -> string -> Sarek_typed_ast.texpr -> bool
val has_varying_return_with : variance -> StringSet.t -> Sarek_typed_ast.texpr -> bool

True if e contains a return reachable under a thread-varying condition (some threads exit early).

val is_thread_varying_env : StringSet.t -> Sarek_typed_ast.texpr -> bool
val has_varying_return : StringSet.t -> Sarek_typed_ast.texpr -> bool

True if e contains a TEReturn that is reachable under a thread-varying condition. Such a return causes SOME threads (those that take the varying branch) to exit early, while other threads continue, creating an implicit divergence point at every subsequent barrier in the enclosing sequence.

val is_thread_varying : Sarek_typed_ast.texpr -> bool

Backward-compatible wrapper using an empty environment

val is_barrier_intrinsic : Sarek_env.intrinsic_ref -> bool

Check if an intrinsic reference is a barrier

val is_warp_convergence_ref : Sarek_env.intrinsic_ref -> bool

Check if an intrinsic ref requires warp convergence

val first_sync : helpers:(string * sync_kind) list -> Sarek_typed_ast.texpr -> sync_kind option

The first synchronisation point te contains: a barrier, a superstep, a warp collective, or a call to a helper of helpers (which contains one).

The error a call reaching kind gives in diverged flow.

val is_sync_error : Sarek_error.error -> bool
val check_expr : ctx -> Sarek_typed_ast.texpr -> Sarek_error.error list

Collect errors from convergence analysis

val check_helper_call : ctx -> Sarek_typed_ast.texpr -> Sarek_typed_ast.texpr list -> Sarek_error.error list

Strict analysis: a call to a helper whose body is checked with the parameters that the call makes thread-varying. The definition itself is checked with uniform parameters, so a shuffle under if param > k is only seen here. Only synchronisation errors are reported; a helper already being re-checked is not entered again.

val check_module_item : ctx -> Sarek_typed_ast.tmodule_item -> Sarek_error.error list

Check a module item

val check_kernel_with : ctx -> Sarek_typed_ast.tkernel -> Sarek_error.error list

Run check over the module items and the body of kernel with ctx.

val add_strict_sync_errors : legacy:Sarek_error.error list -> strict:Sarek_error.error list -> Sarek_error.error list

The synchronisation errors only the strict analysis sees in legacy's blind spots: strict errors that legacy did not already report.

val has_sync_point : Sarek_typed_ast.texpr -> bool

Does te contain a synchronisation point (barrier, superstep or warp collective)? The strict analysis is skipped for a kernel without one (it only reports synchronisation errors).

val check_toplevel_expr : Sarek_typed_ast.texpr -> Sarek_error.error list

Check an expression as a whole kernel body: the legacy analysis, then the strict analysis for synchronisation points.

val check_kernel : Sarek_typed_ast.tkernel -> (unit, Sarek_error.error list) Stdlib.result

Check a kernel for convergence safety

val expr_uses_barriers : Sarek_typed_ast.texpr -> bool

Check if a kernel uses any barriers (explicit or implicit). Used for compile-time optimization of native kernel execution.

val kernel_uses_barriers : Sarek_typed_ast.tkernel -> bool

Check if a kernel uses barriers

Kernel Dimensionality Analysis

Detect which thread/block dimensions a kernel uses to enable optimized native CPU execution. Simple kernels that only use global_idx_x can be run with a simple parallel for loop without the thread_state overhead.

type dim_usage = {
  1. uses_x : bool;
  2. uses_y : bool;
  3. uses_z : bool;
  4. uses_block_dim : bool;
    (*

    Uses block_dim_x/y/z

    *)
  5. uses_grid_dim : bool;
    (*

    Uses grid_dim_x/y/z

    *)
  6. uses_thread_idx : bool;
    (*

    Uses thread_idx_x/y/z directly

    *)
  7. uses_block_idx : bool;
    (*

    Uses block_idx_x/y/z directly

    *)
  8. uses_shared_mem : bool;
    (*

    Uses shared memory

    *)
}

Dimension usage record

val empty_dim_usage : dim_usage
val merge_dim_usage : dim_usage -> dim_usage -> dim_usage
val dim_usage_of_name : string -> dim_usage

Check if an intrinsic name affects dimension usage

val dim_usage_of_intrinsic_ref : Sarek_env.intrinsic_ref -> dim_usage

Get dimension usage from an intrinsic reference

val expr_dim_usage : Sarek_typed_ast.texpr -> dim_usage

Analyze expression for dimension usage

type exec_strategy =
  1. | Simple1D
    (*

    Only uses global_idx_x - can use simple parallel for

    *)
  2. | Simple2D
    (*

    Only uses global_idx_x/y - can use nested loops

    *)
  3. | Simple3D
    (*

    Uses all three dimensions with simple loops

    *)
  4. | FullState
    (*

    Uses block/thread indices or shared memory - needs full state

    *)

Kernel execution strategy based on dimension analysis

val kernel_exec_strategy : Sarek_typed_ast.tkernel -> exec_strategy

Determine the optimal execution strategy for a kernel