Module Tolk.Program_spec

Compile-time kernel descriptions extracted from a linearized uop program.

A t is the runtime-facing description of a lowered kernel before device-specific preparation. It captures the lowered program, kernel name, launch metadata, scalar variables, buffer reads and writes, and cost estimates.

Types

type program = Tolk_uop.Uop.t list

The linearized uop program produced by the linearizer.

type var = {
  1. name : string;
    (*

    Variable name matching the IR definition.

    *)
  2. lo : int;
    (*

    Inclusive lower bound.

    *)
  3. hi : int;
    (*

    Inclusive upper bound.

    *)
  4. dtype : Tolk_uop.Dtype.t;
    (*

    Scalar data type.

    *)
}

Bounded scalar Tolk_uop.Uop.Param_arg kernel parameter.

type core_id = {
  1. var_index : int;
  2. lo : int;
  3. hi : int;
}

Runtime-managed "core_id" variable for multi-core dispatch.

val thread_count : core_id -> int

thread_count cid is cid.hi - cid.lo + 1.

type launch_kind =
  1. | Serial
  2. | Thread_groups
  3. | Threads
    (*

    Kernel launch model.

    *)

Cost estimates

module Estimates : sig ... end

Kernel specifications

type t

Compile-time kernel description.

val of_program : name:string -> src:string -> device:string -> ?lib:bytes -> ?applied_opts:Tolk_uop.Uop.Opt.t list -> ?estimates:Estimates.t -> ?aux:string list -> program -> t

of_program ~name ~src ~device ?lib ?applied_opts ?estimates ?aux program extracts a kernel description from program. If estimates is omitted, estimates are computed from program. aux is copied to Tolk_uop.Uop.program_info.

val with_lib : bytes -> t -> t

with_lib lib spec is spec with lib set to Some lib.

val with_estimates : Estimates.t -> t -> t

with_estimates e spec is spec with estimates replaced by e.

val with_global_dims : int array -> t -> t

with_global_dims dims spec is spec with the global launch dimensions replaced by constant values dims.

val name : t -> string
val src : t -> string
val device : t -> string
val program : t -> program
val lib : t -> bytes option
val applied_opts : t -> Tolk_uop.Uop.Opt.t list
val vars : t -> var list
val outs : t -> int list
val ins : t -> int list
val globals : t -> int list
val core_id : t -> core_id option
val launch_kind : t -> launch_kind
val estimates : t -> Estimates.t
val global_size : t -> Tolk_uop.Uop.t array
val local_size : t -> Tolk_uop.Uop.t array option
val program_info : t -> Tolk_uop.Uop.program_info

program_info spec is the tinygrad-shaped program metadata carried by spec. Symbolic global dimensions are preserved as launch expressions; local dimensions are present only when all local dimensions are fixed integers, and backend auxiliary metadata is preserved.

val launch_dims : t -> (string * int) list -> int array * int array option

launch_dims spec var_vals evaluates launch dimensions by replacing symbolic variables with the values in var_vals.