Module Tolk_uop.Axis_type

Kernel loop axis classification.

An axis type labels each range (loop) in a kernel schedule with the role it plays after lowering: a hardware execution dimension (threads, warps, workgroups, global grid), a compiler transform (unroll, upcast, reduce), or a plain software loop. The linearizer and renderer consult these labels to decide how each range is materialised in the emitted source.

Invariant. The variant declaration order is load-bearing. compare delegates to Stdlib.compare, which compares by constructor index, and the schedule optimiser relies on this order when sorting ranges by (to_pos, axis_id). Reordering variants changes observable sort order.

Types

type t =
  1. | Global
    (*

    Global work-grid dimension. Materialises as a workgroup index on the target (CUDA block, Metal threadgroup, OpenCL group).

    *)
  2. | Warp
    (*

    Warp / subgroup lane dimension. Used by tensor-core lowerings and cross-lane reductions.

    *)
  3. | Local
    (*

    Workgroup-local thread dimension. Materialises as the in-group thread index and is the natural companion of shared-memory staging.

    *)
  4. | Loop
    (*

    Plain software loop. Emitted as a for in the rendered source; no implicit parallelism.

    *)
  5. | Group_reduce
    (*

    Reduction across a Local workgroup, typically staged through shared memory and a tree reduction.

    *)
  6. | Reduce
    (*

    Sequential reduction axis. Emits an accumulator loop around the reduced body.

    *)
  7. | Upcast
    (*

    Fully unrolled axis that becomes a vector lane. Enables packed loads/stores and SIMD arithmetic in the renderer.

    *)
  8. | Unroll
    (*

    Fully unrolled axis over a reduction, producing an explicit sum of per-iteration expressions.

    *)
  9. | Thread
    (*

    Per-thread private dimension, used by tensor-core lowerings to describe register-held tiles.

    *)
  10. | Placeholder
    (*

    Unassigned axis. Temporary kind used between schedule passes; no range with this kind may reach the renderer.

    *)

The axis kinds assignable to a kernel range.

Exactly one kind is attached to each range; most transforms fix a kind at creation and never mutate it afterwards.

Predicates and ordering

val equal : t -> t -> bool

equal a b is structural equality on variants.

val compare : t -> t -> int

compare a b orders by variant declaration order (see the module invariant).

Formatting

val to_string : t -> string

to_string t is the lowercase constructor name (e.g. "global", "group_reduce"). Stable identifier used in serialised kernel info and debugging output.

val pp : Stdlib.Format.formatter -> t -> unit

pp ppf t prints t using to_string.

Scheduling

val to_pos : t -> int

to_pos t is the schedule priority of t. The postrange pass orders ranges by (to_pos kind, axis_id), which groups them as: loops wrap hardware dimensions, hardware dimensions wrap reductions, reductions wrap unrolls, and placeholders sink to the bottom.

Priorities:

val letter : t -> string

letter t is the single-character tag used when building axis names in the rendered source; the renderer prefixes it to form indices like "gidx0", "lidx1", or "Ridx0".

Mapping: