Module Tolk_uop.Dtype

Data types for tensor computations.

This module defines two levels of data type and their union:

  • Val.t — value dtypes (a scalar identity with a vector width).
  • Ptr.t — pointer dtypes (a base value dtype plus address space, buffer element count, pointer vector width, and optional image metadata).
  • Image — constructors and accessors for image pointer dtypes.
  • t — the union of both, for IR nodes whose dtype can be either (e.g., Index).

Val-specific operations live in Val, pointer-specific operations in Ptr, and dispatching operations at the top level. Kernel view fields use the most precise type: Val.t for value-only nodes, Ptr.t for pointer-only nodes, and t where either is possible.

Promotion. Scalar types are organized in a promotion lattice based on JAX JEP-9407. Promotion is total over numeric types: any pair has a common supertype, at the cost of some lossy edges (e.g., Uint64 promotes through fp8 to reach floats). See Val.least_upper_dtype.

Vectorization. A value dtype represents count lanes of a scalar. Most predicates and properties work on the scalar component. Use Val.vec and Val.scalarize to convert.

Weakint. The Weakint scalar is a weakly-typed integer used for loop counters, shape vectors, and untyped integer literals. It has priority 0 and promotes to any concrete integer type (Weakint + Int8 = Int8). Carries a sentinel bitsize of 800 so it cannot be confused with a machine width; backends lower it to int32 or int64.

Scalar types

type scalar =
  1. | Void
    (*

    Absence of a value. Has zero bitsize.

    *)
  2. | Weakint
    (*

    Weakly-typed integer. Priority 0. Used for loop counters, shape vectors, and untyped integer literals. Promotes to any concrete integer type. Carries a sentinel bitsize of 800.

    *)
  3. | Bool
    (*

    Boolean. 1 bit.

    *)
  4. | Int8
    (*

    Signed 8-bit integer.

    *)
  5. | Int16
    (*

    Signed 16-bit integer.

    *)
  6. | Int32
    (*

    Signed 32-bit integer.

    *)
  7. | Int64
    (*

    Signed 64-bit integer.

    *)
  8. | Uint8
    (*

    Unsigned 8-bit integer.

    *)
  9. | Uint16
    (*

    Unsigned 16-bit integer.

    *)
  10. | Uint32
    (*

    Unsigned 32-bit integer.

    *)
  11. | Uint64
    (*

    Unsigned 64-bit integer.

    *)
  12. | Uint128
    (*

    Non-portable virtual 128-bit storage helper. Tinygrad keeps this as a private dtype; it is not classified as an integer or unsigned dtype. Environment dtype parsing accepts only tinygrad's private spelling "_uint128".

    *)
  13. | Uint256
    (*

    Non-portable virtual 256-bit storage helper. Tinygrad keeps this as a private dtype; it is not classified as an integer or unsigned dtype. Environment dtype parsing accepts only tinygrad's private spelling "_uint256".

    *)
  14. | Float16
    (*

    IEEE 754 binary16 (half precision).

    *)
  15. | Bfloat16
    (*

    Brain floating-point 16: 8-bit exponent, 7-bit mantissa.

    *)
  16. | Float32
    (*

    IEEE 754 binary32 (single precision).

    *)
  17. | Float64
    (*

    IEEE 754 binary64 (double precision).

    *)
  18. | Fp8e4m3
    (*

    8-bit float: 4-bit exponent, 3-bit mantissa.

    *)
  19. | Fp8e5m2
    (*

    8-bit float: 5-bit exponent, 2-bit mantissa.

    *)
  20. | Fp8e4m3fnuz
    (*

    8-bit float, FNUz variant: 4 exp / 3 mant bits, no signed zero, NaN encoded as 0x80.

    *)
  21. | Fp8e5m2fnuz
    (*

    8-bit float, FNUz variant: 5 exp / 2 mant bits.

    *)

Scalar data type identity.

Promotion priority ranges from -1 (Void) to 14 (Float64). Bool and Weakint share priority 0. See priority and Val.least_upper_dtype for the full ordering.

Address spaces

type addr_space =
  1. | Global
    (*

    Global device memory.

    *)
  2. | Local
    (*

    Shared/local memory (workgroup scope).

    *)
  3. | Reg
    (*

    Register storage.

    *)
  4. | Alu
    (*

    Scalar ALU value space, used for symbolic variables and loaded values.

    *)

GPU memory address space.

type image_kind =
  1. | Imageh
  2. | Imagef

Image storage kind. Mirrors tinygrad's imageh and imagef dtype constructors.

Value dtypes

module Val : sig ... end

Pointer dtypes

module Ptr : sig ... end

Image pointer dtypes

module Image : sig ... end

Unified dtype

type t =
  1. | Val of Val.t
  2. | Ptr of Ptr.t
    (*

    A dtype that is either a value or a pointer. Used in IR nodes whose dtype can be either (e.g., Index nodes that may or may not carry pointer semantics).

    *)

Dispatching accessors

val scalar : t -> scalar

scalar dt is the scalar identity. scalar (Val v) is Val.scalar v. scalar (Ptr p) is Ptr.scalar p.

val count : t -> int

count dt is the value vector width. count (Val v) is Val.count v. count (Ptr p) is Ptr.count p.

val vcount : t -> int

vcount dt is the outer vector width: lane count for values, pointer vector width for pointers. vcount (Val v) is Val.count v. vcount (Ptr p) is Ptr.v p.

val is_ptr : t -> bool

is_ptr dt is true iff dt is Ptr _.

val val_of : t -> Val.t

val_of dt is the value component of dt: v for Val v and Ptr.base p for Ptr p. Pointer vector width v is discarded.

Dispatching transformers

val scalarize : t -> t

scalarize dt dispatches to Val.scalarize or Ptr.scalarize, preserving the Val/Ptr wrapper.

val vec : int -> t -> t

vec n dt dispatches to Val.vec or Ptr.vec, preserving the Val/Ptr wrapper.

Predicates

val is_float : t -> bool

is_float dt is true iff dt is a value dtype whose scalar is a floating-point type (Float16, Bfloat16, Float32, Float64, Fp8e4m3, or Fp8e5m2) or an image pointer dtype. Ordinary pointer dtypes are not classified by their pointee scalar.

val is_int : t -> bool

is_int dt is true iff dt is a value dtype whose scalar is a signed or unsigned integer type, including Weakint. Uint128 and Uint256 are private storage helpers and are not classified as integers.

val is_unsigned : t -> bool

is_unsigned dt is true iff dt is a value dtype whose scalar is an unsigned integer dtype. Uint128 and Uint256 are private storage helpers and are not classified as unsigned.

val is_bool : t -> bool

is_bool dt is true iff dt is a value dtype whose scalar is Bool.

val is_fp8 : t -> bool

is_fp8 dt is true iff dt is a value dtype whose scalar is an fp8 dtype.

Promotion

val least_upper_dtype : t list -> t

least_upper_dtype ts is the least upper bound of ts.

If ts contains image pointer dtypes, the first image dtype is returned. Ordinary pointer dtypes are rejected. Otherwise this dispatches to Val.least_upper_dtype and wraps the result in Val.

Raises Invalid_argument if ts is empty or contains an ordinary pointer dtype.

val least_upper_float : t -> t

least_upper_float dt is dt for image pointer dtypes. Ordinary pointer dtypes are rejected. Floating-point value dtypes are returned unchanged; other value dtypes promote with Val.default_float.

Raises Invalid_argument if dt is an ordinary pointer dtype.

Properties

val bitsize : t -> int

bitsize dt is the total size of dt in bits: the scalar bit width multiplied by count dt. Void has bitsize 0; Weakint has sentinel bitsize 800. For ordinary pointers, this measures the pointed-to base width and ignores the pointer vector width v. For image pointers, this is 16 for Imageh and 32 for Imagef.

val itemsize : t -> int

itemsize dt is bitsize rounded up to whole bytes.

val priority : t -> int

priority dt is the promotion priority of scalar dt. Higher priorities absorb lower ones in Val.least_upper_dtype. Ranges from -1 (Void) through 0 (Bool, Weakint) to 14 (Float64); image pointer dtypes have priority 100.

Bounds

type bound = [
  1. | `Bool of bool
  2. | `SInt of int64
  3. | `UInt of int64
  4. | `Float of float
]

Numeric bounds for dtypes. Returned by min and max.

  • `Bool b for boolean bounds.
  • `SInt n for signed integer bounds, including Weakint, which reports the Int64 range as an approximation.
  • `UInt n for unsigned integer bounds. Values are raw 64-bit unsigned bit patterns stored in an int64 (e.g. Uint64's max is `UInt Int64.minus_one).
  • `Float f for floating-point bounds (neg_infinity and infinity).
val min : t -> bound

min dt is the smallest value representable by dt.

Value dtypes use their scalar bounds. Image pointer dtypes are float-like and return `Float neg_infinity. Ordinary pointer dtypes are not classified by their pointee scalar and return `Bool false.

Raises Invalid_argument if dt is a value dtype whose scalar is Void.

See also max.

val max : t -> bound

max dt is the largest value representable by dt.

Value dtypes use their scalar bounds. Image pointer dtypes are float-like and return `Float infinity. Ordinary pointer dtypes are not classified by their pointee scalar and return `Bool true.

Raises Invalid_argument if dt is a value dtype whose scalar is Void.

See also min.

val finfo : t -> int * int

finfo dt is (exponent_bits, mantissa_bits) for the floating-point dtype dt. Mantissa bits exclude the implicit leading 1. For example, finfo (Val float32) is (8, 23) and finfo (Val float16) is (5, 10).

Raises Invalid_argument if dt is a pointer dtype or if dt is a value dtype whose scalar is not floating-point.

See also is_float.

Comparison

val equal : t -> t -> bool

equal a b is true iff a and b are both Val or both Ptr with equal components.

val compare : t -> t -> int

compare a b is a total order over t. Val sorts before Ptr; within each kind, compares by Val.compare or Ptr.compare.

Formatting

val to_string : t -> string

to_string dt formats dt using Val.to_string or Ptr.to_string.

val repr : t -> string

repr dt formats dt using tinygrad-style dtype syntax.

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

pp formats a dtype using to_string.

Scalar formatting

val scalar_to_string : scalar -> string

scalar_to_string s is the short name of s (e.g., "f32", "i64", "bool", "void", "weakint").

val pp_scalar : Stdlib.Format.formatter -> scalar -> unit

pp_scalar formats a scalar using scalar_to_string.

val addr_space_to_string : addr_space -> string

addr_space_to_string a is "global", "local", "reg", or "alu".

val pp_addr_space : Stdlib.Format.formatter -> addr_space -> unit

pp_addr_space formats an address space.

Convenience constructors

Value dtype constants wrapped as t for direct use in any context expecting a unified dtype. For the unwrapped Val.t versions, use the Val module directly.

val of_scalar : scalar -> t

of_scalar s is Val (Val.of_scalar s).

The named constants below are Val-wrapped versions of the matching Val constants. default_float follows DEFAULT_FLOAT as described by Val.default_float; default_int is int32.

val void : t
val bool : t
val int8 : t
val int16 : t
val int32 : t
val int64 : t
val uint8 : t
val uint16 : t
val uint32 : t
val uint64 : t
val float16 : t
val bfloat16 : t
val float32 : t
val float64 : t
val fp8e4m3 : t
val fp8e5m2 : t
val fp8e4m3fnuz : t
val fp8e5m2fnuz : t
val weakint : t
val default_float : t
val default_int : t
val imageh : int list -> t

imageh shape is Ptr (Image.imageh shape).

val imagef : int list -> t

imagef shape is Ptr (Image.imagef shape).

Floating-point conversion

Bit-exact precision-narrowing utilities used for constant folding. Rounding is round-to-nearest-even for all formats; conversions are implemented in OCaml rather than delegated to hardware so results are reproducible across backends.

val float_to_fp16 : float -> float

float_to_fp16 x rounds x to IEEE 754 binary16 (half) precision using round-to-nearest-even. The result is a float holding a value exactly representable in half precision. Overflow produces an infinity; underflow produces zero or a subnormal; infinities pass through; NaNs produce a NaN (bit pattern not preserved).

val float_to_bf16 : float -> float

float_to_bf16 x rounds x to bfloat16 precision using round-to-nearest-even. Non-finite values pass through unchanged.

val float_to_fp8 : scalar -> float -> int

float_to_fp8 s x encodes x as an fp8 byte value in 0..255.

Raises Invalid_argument if s is not an fp8 dtype.

See also fp8_to_float.

val fp8_to_float : scalar -> int -> float

fp8_to_float s byte decodes the fp8 byte into a float.

Raises Invalid_argument if s is not an fp8 dtype.

See also float_to_fp8.

val truncate_float : Val.t -> float -> float

truncate_float dt x rounds x to the precision of floating-point dtype dt. Float64 is the identity. Float32 round-trips through Int32.bits_of_float. Narrower types delegate to float_to_fp16, float_to_bf16, or the fp8 conversion pair. The count of dt is ignored.

Raises Invalid_argument if dt is not floating-point.

See also truncate_int.

val truncate_int : Val.t -> int -> int

truncate_int dt x reduces integer x into the range of integer dtype dt. Unsigned types mask to the low bitsize dt bits; signed types and Weakint apply the same mask and sign-extend. When the target width equals or exceeds Sys.int_size (notably for Int64, Uint64, and Weakint at its 800-bit sentinel) x is returned unchanged. Bool maps 0 to 0 and any other value to 1.

Raises Invalid_argument if dt is not an integer or bool type.

See also truncate_float.

Storage conversion

Tinygrad-compatible scalar storage helpers. These functions operate on one scalar lane; callers handling vectors should apply them lane-by-lane.

type storage_scalar = [
  1. | `Bool of bool
  2. | `Float of float
  3. | `Int of int64
]

The type for scalar values at the storage boundary.

`Int n carries raw signed or unsigned integer storage bits in an int64. `Float f carries a host floating-point value. `Bool b carries a boolean storage value.

val storage_fmt_for_dtype : Val.t -> char option

storage_fmt_for_dtype dt is the Python struct format character used by tinygrad for one scalar lane of dt.

Returns None for dtypes without a portable storage format: Void, Weakint, Uint128, and Uint256. Bfloat16 stores as 'H' and fp8 dtypes store as 'B'. The count of dt is ignored.

val to_storage_scalar : Val.t -> storage_scalar -> storage_scalar

to_storage_scalar dt x maps x to the scalar value written by tinygrad's storage path for dt.

Bfloat16 returns the high 16 bits of the rounded float as `Int. Fp8 dtypes return the encoded byte as `Int. Float16 returns a precision-rounded `Float. Other scalar dtypes return the corresponding boolean, integer, or float representation.

Raises Invalid_argument if dt is Void. The count of dt is ignored.

val from_storage_scalar : storage_scalar -> Val.t -> storage_scalar

from_storage_scalar x dt maps one stored scalar value for dt back to a host scalar.

Bfloat16 and fp8 storage integers decode to `Float. Other scalar dtypes return the corresponding boolean, integer, or float representation.

Raises Invalid_argument if dt is Void. The count of dt is ignored.

val truncate : Val.t -> storage_scalar -> storage_scalar

truncate dt x coerces x to the value domain of scalar dtype dt.

Floating-point dtypes use truncate_float. Integer dtypes use two's complement wrapping with the target width, including unsigned raw bit patterns. Bool maps zero-like values to `Bool false and all other values to `Bool true.

Raises Invalid_argument if dt is Void. The count of dt is ignored.