Module Sarek_warp

Sarek_warp - the one definition of warp shuffle semantics.

The four shuffle primitives (warp_shuffle, warp_shuffle_down, warp_shuffle_up, warp_shuffle_xor) and the two warp introspection constants (warp_size, lane_id) reach the IR as EIntrinsic nodes. The shuffles are TYPED in the IR: the front end appends the value type to the name (warp_shuffle_down_f32) so a backend never has to infer the type of an expression to choose a spelling.

Contract (docs/plans/warp-shuffle-intake-2026-10-17.md, section 4.3), for lane i of a warp of width W (a power of two) and operand x read as a two's-complement int32:

source_lane is that table, executable. Native and the interpreter use it directly; the source backends emit equivalent guards, and the parity tests compare every backend with it.

type op =
  1. | Idx
  2. | Down
  3. | Up
  4. | Xor
type value_type =
  1. | I32
  2. | F32
  3. | F64
type call =
  1. | Shuffle of op * value_type
  2. | Warp_size
  3. | Lane_id
val op_base : op -> string
val type_suffix : value_type -> string
val all_ops : op list
val all_types : value_type list
val ir_name : op -> value_type -> string

The IR name of a typed shuffle, e.g. "warp_shuffle_down_f32".

val typed_names : (string * call) list
val classify : string -> call option

Classify an IR intrinsic name. None for anything that is not a warp shuffle or a warp introspection constant.

val user_shuffle_names : string list

The user-facing names of the four shuffles (before typing).

val is_user_shuffle_name : string -> bool

Is name one of the four user-facing shuffle names?

val op_of_user_name : string -> op option

The op of a user-facing shuffle name.

val is_shuffle_name : string -> bool

Does this IR intrinsic name denote a shuffle (a warp collective)?

val is_warp_name : string -> bool

Does this IR intrinsic name depend on the warp structure (shuffle, warp_size, lane_id)? Passes that change which threads run, or the block shape, must leave such a kernel alone.

val source_lane : op -> width:int -> lane:int -> int -> int

Index of the lane read by lane lane for operand x in a warp of width lanes. The result is always in 0, width). [x] is the int32 operand sign-extended to an OCaml int.

Simulated warp width

Native and the interpreter have no hardware warp. They simulate one of simulated_width () lanes (default 32, the width of the CUDA and Apple devices and of RDNA). The width is a launch-time parameter in the sense that it is read at each launch; with_width is the test hook that lets the parity tests run the simulation at a GPU's width.

val default_width : int
val is_valid_width : int -> bool
val simulated : int Stdlib.Atomic.t
val simulated_width : unit -> int
val with_width : int -> (unit -> 'a) -> 'a

Run f with the simulated width set to w, restoring the previous width afterwards. Raises Invalid_argument when w is not a power of two in 1, 1024. Not re-entrant across domains: it is a test hook.

val linear_thread_id : block:(int * int * 'a) -> (int * int * int) -> int

Linear thread id of (tx, ty, tz) in a block of (bx, by, _) (x fastest).

val lane_of_thread : width:int -> int -> int

Lane of a linear thread id in a warp of width lanes.

val block_is_whole_warps : block:(int * int * int) -> width:int -> bool

Is a block of bx * by * bz threads a whole number of warps of width?