Module Tolk.Renderer

GPU kernel renderer.

A renderer converts Program_spec.program values to backend-specific source code and owns its Compiler.t. The abstract type t encapsulates target capabilities (memory hierarchy, grid limits, supported operations), a rendering function, and an optional compiler.

Backends construct renderers via make, supplying only the fields that differ from the defaults. The compiler is typically attached by the device backend via with_compiler or the ?compiler parameter of make.

See Cstyle for C-family language backends (CUDA, Metal, OpenCL, HIP, Clang).

Types

type code_op =
  1. | Sqrt
    (*

    Square root.

    *)
  2. | Recip
    (*

    Reciprocal (1/x).

    *)
  3. | Neg
    (*

    Arithmetic negation.

    *)
  4. | Exp2
    (*

    Base-2 exponential.

    *)
  5. | Log2
    (*

    Base-2 logarithm.

    *)
  6. | Sin
    (*

    Sine.

    *)
  7. | Trunc
    (*

    Truncation to integer.

    *)
  8. | And
    (*

    Bitwise AND.

    *)
  9. | Xor
    (*

    Bitwise XOR.

    *)
  10. | Or
    (*

    Bitwise OR.

    *)
  11. | Add
    (*

    Addition.

    *)
  12. | Sub
    (*

    Subtraction.

    *)
  13. | Mul
    (*

    Multiplication.

    *)
  14. | Cmod
    (*

    C-style integer remainder.

    *)
  15. | Cdiv
    (*

    C-style truncating integer division.

    *)
  16. | Cmpne
    (*

    Not-equal comparison.

    *)
  17. | Shr
    (*

    Bitwise right shift.

    *)
  18. | Shl
    (*

    Bitwise left shift.

    *)
  19. | Cmplt
    (*

    Less-than comparison.

    *)
  20. | Where
    (*

    Ternary select (cond ? a : b).

    *)
  21. | Cmpeq
    (*

    Equality comparison.

    *)
  22. | Fdiv
    (*

    Floating-point division.

    *)
  23. | Max
    (*

    Maximum.

    *)
  24. | Mulacc
    (*

    Fused multiply-accumulate.

    *)
  25. | Threefry
    (*

    Threefry 2x32 PRNG mixing function.

    *)

ALU operations that a backend can provide custom rendering for.

The decomposition pass uses supported_ops to decide which composite operations to lower; code_for_op lists the operations a renderer handles natively.

Operations without a corresponding flag in supported_ops (Add, Sub, Mul, Cmod, Cdiv, Cmpne, Xor, Where, Trunc) are always required. Floor division and modulo are lowered before rendering, matching tinygrad's C-style renderers.

Supported operations

val all_supported_ops : Decomp_op.supported_ops

all_supported_ops marks every decomposable operation as natively supported.

val supported_ops_of_code_for_op : code_op list -> Decomp_op.supported_ops

supported_ops_of_code_for_op ops derives capability flags from a list of natively rendered operations. An operation absent from ops is marked unsupported.

Renderer

type t

The type for renderers.

Properties

val name : t -> string

name r is the renderer name (e.g., "metal", "cuda").

val device : t -> string

device r is the target device identifier (e.g., "NV", "HIP", "CPU"). Passed as context to codegen rewrite passes for device-specific transformations.

val compiler : t -> Compiler.t option

compiler r is r's compiler, or None if the renderer has no associated compiler (e.g., interpreter backends).

val has_local : t -> bool

has_local r is true iff r supports local thread IDs.

val has_threads : t -> bool

has_threads r is true iff r supports host-side threading instead of GPU grid dimensions.

val has_shared : t -> bool

has_shared r is true iff r supports shared memory.

val global_max : t -> int list option

global_max r is the maximum global grid dimensions [x; y; z], or None when unconstrained. The list has exactly three elements when present.

val global_prod_max : t -> int list option

global_prod_max r is the per-axis product limit for global dimensions, or None when unconstrained. When present, each global dimension is capped at min(global_max.(i), global_prod_max.(i) / local_hw.(i)).

val local_max : t -> int list option

local_max r is the maximum local workgroup dimensions [x; y; z], or None when unconstrained. The list has exactly three elements when present.

val shared_max : t -> int

shared_max r is the maximum shared memory size in bytes.

  • 0 when shared memory is unsupported (has_shared is false).
  • For GPU backends, a conservative default (e.g., 32 KB for OpenCL, 48 KB for CUDA). Actual limits may vary by device.
val tensor_cores : t -> Tc.t list

tensor_cores r is the list of tensor_core configurations supported by r. Empty when the backend has no hardware matrix-multiply support.

Capabilities

val code_for_op : t -> code_op list

code_for_op r is the list of ALU operations that r provides custom rendering for.

See also supported_ops.

val supported_ops : t -> Decomp_op.supported_ops

supported_ops r is the backend capability flags for the decomposition pass, derived from code_for_op unless explicitly overridden via make.

val supports_dtype : t -> Tolk_uop.Dtype.t -> bool

supports_dtype r dt is true iff the backend natively supports dt. When false, the decomposition pass emulates dt using supported types.

val emulated_float_dtypes : t -> (Tolk_uop.Dtype.scalar * Tolk_uop.Dtype.scalar) list

emulated_float_dtypes r is the list of (from, to) dtype pairs for float emulation. Each from float is promoted to to (typically f32). Empty for backends that natively support all float types.

val pre_matcher : t -> (Tolk_uop.Uop.t -> Tolk_uop.Uop.t option) option

pre_matcher r is an optional device-specific rewrite rule applied before decompositions.

val extra_matcher : t -> (Tolk_uop.Uop.t -> Tolk_uop.Uop.t option) option

extra_matcher r is an optional device-specific rewrite rule composed into the final rewrite fixpoint.

Load/store policy

val supports_float4 : t -> bool

supports_float4 r is true iff r supports vectorized (float4/float2) load and store operations. Codegen lowering uses this to decide whether wide accesses can be folded. Defaults to true.

val image_pitch_alignment : t -> int option

image_pitch_alignment r is the image row pitch alignment in scalar pixels used by late image selection, when the target exposes one.

Rendering

val render : t -> ?name:string -> Tolk_uop.Uop.t list -> string

render r ~name program converts program to backend-specific source code.

name defaults to "kernel".

val aux : t -> Tolk_uop.Uop.t list -> string list

aux r program is backend-specific program metadata derived from program. Empty when the backend has no auxiliary runtime payload.

Construction

val make : ?tensor_cores:Tc.t list -> ?supports_float4:bool -> ?image_pitch_alignment:int -> ?has_threads:bool -> ?global_max:int list -> ?global_prod_max:int list -> ?local_max:int list -> ?code_for_op:code_op list -> ?supported_ops:Decomp_op.supported_ops -> ?compiler:Compiler.t -> ?pre_matcher:(Tolk_uop.Uop.t -> Tolk_uop.Uop.t option) -> ?extra_matcher:(Tolk_uop.Uop.t -> Tolk_uop.Uop.t option) -> ?supports_dtype:(Tolk_uop.Dtype.t -> bool) -> ?aux:(Tolk_uop.Uop.t list -> string list) -> ?emulated_floats:(Tolk_uop.Dtype.scalar * Tolk_uop.Dtype.scalar) list -> name:string -> device:string -> has_local:bool -> has_shared:bool -> shared_max:int -> render:(?name:string -> Tolk_uop.Uop.t list -> string) -> unit -> t

make ~name ~device ~has_local ~has_shared ~shared_max ~render () is a renderer with the given capabilities.

Optional parameters and their defaults:

  • tensor_cores: [] (none).
  • supports_float4: true.
  • image_pitch_alignment: None.
  • has_threads: false.
  • global_max: Some [0x8FFFFFFF; 0x8FFFFFFF; 0x8FFFFFFF].
  • global_prod_max: None.
  • local_max: Some [0x8FFFFFFF; 0x8FFFFFFF; 0x8FFFFFFF].
  • code_for_op: [] (no custom ops).
  • supported_ops: derived from code_for_op via supported_ops_of_code_for_op. When code_for_op is [], defaults to all_supported_ops.
  • supports_dtype: fun _ -> true.
  • compiler: None.
val with_compiler : Compiler.t -> t -> t

with_compiler c r is r with compiler set to Some c.