docs menu

on this page

Module Tolk.Ops

The intermediate representation: nodes, patterns and graph rewriting.

A node (t) is an operation (Op.t) applied to source nodes, with a payload (arg) and a data type derived from the three. Nodes form a directed acyclic graph, and one vocabulary serves every stage of compilation: tensor graphs, kernels, programs and command sequences.

Identity. Nodes are hash-consed: building a node equal to a live one returns that node, so two nodes are structurally equal iff they are physically equal (==). A node lives as long as something references it. Nodes can be built and read from any domain.

Construction. v builds any node; the constructors below build the common ones with the defaults and checks of each operation. Arithmetic promotes its operands to a common data type, and has their broadcast shape (Elementwise).

Shapes. A node's shape is a list of sint: integers, or integer nodes for sizes known only when the program runs. Comparing such sizes needs simplification, so shapes, movements and the functions that read them are Shape's.

Patterns. Upat describes nodes, Pattern_matcher pairs patterns with rules, and graph_rewrite rewrites a graph with them to a fixed point.

Errors. A broken precondition raises Invalid_argument.

#type t

The type for nodes.

Axis types

#module Axis_type : sig ... end

Axis types.

Devices

#type device = string

The type for devices. Devices are named by the caller, such as "CPU" or "AMD:1"; only a name starting with "DISK" has a meaning here, a disk.

#val pp_device : Stdlib.Format.formatter -> device -> unit

pp_device formats a device as its quoted name: 'CPU'.

#val is_disk_device : device -> bool

is_disk_device d is true iff d is a disk: its name, up to its first :, is DISK in any case.

Symbolic integers

#type sint =

#| Int of int
#| Sym of t

The type for symbolic integers: an integer, or an integer node whose value is known when the program runs. Functions that return one return Int whenever the value is known.

Arguments

#type param_arg = {

#slot : int;
(*
The parameter's position in its function, or -1 for a named variable.
*)
#dtype : Dtype.t;
(*
The element type.
*)
#vmin_vmax : (Dtype.value * Dtype.value) option;
(*
The least and greatest value of a variable.
*)
#name : string option;
(*
The name of a variable or a named buffer.
*)
#addrspace : Dtype.addr_space;
(*
Where the storage lives; Dtype.addr_space.Alu for a scalar.
*)
#device : device option;
(*
The placement of the storage.
*)
#volatile : bool;
(*
Accesses must not be merged or reordered.
*)
#bind_on_realize : bool;
(*
The storage is bound by whoever realizes the graph, not by a call.
*)
#bound : Dtype.value option;
(*
The value a variable is bound to.
*)
#phase : int;
(*
The bytes by which the storage's first element lies past a multiple of align: from 0 to align - 1, a multiple of the element's size or of align, whichever is less, which a copy that changes the element type must keep. Vector accesses start only where they are aligned to their width.
*)
#align : int;
(*
The power of two, from 1 to 16, the width of the widest vector access, modulo which the start of the storage is known (phase). No vector access is wider.
*)

}

The type for the arguments of Op.t.Param, Op.t.Buffer and Op.t.Alloc: storage, or a scalar variable. The number of elements is the node's first source (size_src).

val param_arg : 
  ?vmin_vmax:(Dtype.value * Dtype.value) ->
  ?name:string ->
  ?addrspace:Dtype.addr_space ->
  ?device:device ->
  ?volatile:bool ->
  ?bind_on_realize:bool ->
  ?bound:Dtype.value ->
  ?phase:int ->
  ?align:int ->
  slot:int ->
  Dtype.t ->
  param_arg

param_arg ~slot dtype is the argument with these fields. addrspace defaults to Dtype.addr_space.Global; the flags default to false, phase to 0, align to 16 and the other fields to None.

Raises Invalid_argument if align is not a power of two from 1 to 16, or phase is not below align and a multiple of dtype's size or of align, whichever is less.

#type estimates = {

#ops : sint;
(*
Arithmetic operations.
*)
#lds : sint;
(*
Bytes loaded and stored.
*)
#mem : sint;
(*
Bytes of memory touched, each counted once.
*)

}

The type for the cost estimates of a kernel.

#val pp_estimates : Stdlib.Format.formatter -> estimates -> unit

pp_estimates formats Estimates(ops=0, lds=0, mem=0).

#type split = {

#iterations : sint;
(*
The loop's iterations, 0 to iterations - 1.
*)
#lo : int;
(*
The slot of the variable that holds a block's first.
*)
#hi : int;
(*
The slot of the variable that holds the one after its last.
*)

}

The type for the loop a host program's launch splits into blocks that the host's cores run at once. Its program runs the iterations from the value of its variable lo up to that of hi.

#type kernel_info = {

#name : string;
(*
The kernel's name.
*)
#applied_opts : Opt.t list;
(*
The optimisations applied, in order.
*)
#opts_to_apply : Opt.t list option;
(*
The optimisations to apply, or None to choose them.
*)
#estimates : estimates option;
(*
The kernel's cost.
*)
#beam : int;
(*
The beam width to search optimisations with, or 0.
*)
#optimized : bool;
(*
Whether the kernel went through the optimiser, which then leaves it as it is.
*)
#kernel : bool;
(*
Whether the runtime launches it; false for a body another kernel calls, as a compiled program.
*)
#split : split option;
(*
The loop its launch splits, if any.
*)

}

The type for the arguments of a kernel's Op.t.Sink.

val kernel_info : 
  ?name:string ->
  ?applied_opts:Opt.t list ->
  ?opts_to_apply:Opt.t list ->
  ?estimates:estimates ->
  ?beam:int ->
  ?optimized:bool ->
  ?kernel:bool ->
  ?split:split ->
  unit ->
  kernel_info

kernel_info () is the argument with these fields. name defaults to "test", the lists to [] and None, beam to 0, optimized to false, kernel to true and split to None.

#val function_name : kernel_info -> string

function_name k is k.name as an identifier (Helpers.to_function_name).

#type program_info = {

#global_size : sint list;
(*
The number of workgroups on each axis.
*)
#local_size : sint list;
(*
The threads of a workgroup on each axis.
*)
#params : t list;
(*
The parameters, buffers and scalars, one per slot, by slot.
*)
#globals : int list;
(*
The slots of the buffers.
*)
#outs : int list;
(*
The slots of the buffers written.
*)
#ins : int list;
(*
The slots of the buffers read.
*)
#target : Helpers.Target.t;
(*
The target compiled for.
*)

}

The type for the arguments of Op.t.Program.

#type keep =

#| Removable
(*
Drop the axes its value does not vary along, and inline it back where that costs little.
*)
#| Broadcast
(*
As Removable, for a value read where it is broadcast: it is not inlined back where computing it runs a transcendental function, which its readers would compute again for each element of the ranges it does not vary along.
*)
#| Whole
(*
Nothing: a value the user materialises is stored as it is.
*)
#| Argument
(*
As Whole, for a value a custom kernel reads: a store through the same view of storage writes that storage, not this buffer.
*)

What later passes may do to a buffer an Op.t.Stage makes.

#type bufferize_opts = {

#device : device option;
(*
Where the new buffer lives.
*)
#addrspace : Dtype.addr_space;
(*
Its address space.
*)
#keep : keep;
(*
What later passes may do to it.
*)

}

The type for the arguments of an Op.t.Stage that makes a buffer.

#type custom_function = {

#name : string;
(*
The symbol of the function, or with fmt the code that the call writes.
*)
#dtype : Dtype.t;
(*
The type it returns, Dtype.t.Void for none.
*)
#fmt : bool;
(*
name is code in the target's language: a format string of the call's arguments, where {} is the next argument, {i} the ith from 0, and {{ and }} are braces, as Python's str.format reads positional fields; any other field is refused.
*)

}

The type for the arguments of Op.t.Custom_function.

#type call_info = {

#name : string option;
(*
The name of the function called.
*)
#precompile : bool;
(*
Compile the body on its own.
*)
#inline : bool;
(*
Write the call of code in the target's language into its reader, as an expression.
*)
#write : (int * int) list;
(*
The disjoint writes (after_disjoint) this kernel call is a kernel of, each as the write's number in its schedule and the position, among the call's arguments, of the state it writes. A schedule sets it and reads it alone: a program's key and the call's printed form (pp_arg) leave it out.
*)

}

The type for the arguments of Op.t.Call.

#type wmma_axis = (Axis_type.t * int list) * int

The type for an axis a tensor core folds: its range's role and identity, and its size.

#type wmma = {

#dims : int * int * int;
(*
The matrix dimensions N, M and K.
*)
#dtype_in : Dtype.t;
(*
The type of the multiplied operands.
*)
#threads : int;
(*
The threads that cooperate on one product.
*)
#upcast_axes : (wmma_axis list * wmma_axis list * wmma_axis list) option;
(*
For each operand, the axes still to fold into it, or None once folded.
*)

}

The type for the arguments of Op.t.Wmma.

#type arg =

#| No_arg
#| Const of Dtype.const
(*
Op.t.Const: its value.
*)
#| Dtype of Dtype.t
(*
Op.t.Cast, Op.t.Bitcast: the target type.
*)
#| Param of param_arg
(*
Op.t.Param, Op.t.Buffer, Op.t.Alloc.
*)
#| Range of {
1. #axis_type : Axis_type.t;
2. #axis_id : int list;
}
(*
Op.t.Range: its role and identity.
*)
#| Reduce of {
1. #op : Op.t;
2. #num_axes : int;
}
(*
Op.t.Reduce: the operation, and how many leading axes it reduces (0 for a kernel's reduction over range sources).
*)
#| Device of device
(*
Op.t.Copy: the target.
*)
#| String of string
(*
Op.t.Special, Op.t.Linear: a name; Op.t.Source: the source text.
*)
#| Bytes of string
(*
Op.t.Binary: the bytes.
*)
#| Code of {
1. #code : string;
2. #dtype : Dtype.t;
}
(*
Op.t.Ins: an instruction, with the type produced.
*)
#| Bufferize of bufferize_opts
(*
Op.t.Stage.
*)
#| Kernel of kernel_info
(*
Op.t.Sink: the root of a kernel.
*)
#| Program of program_info
(*
Op.t.Program.
*)
#| Custom_function of custom_function
(*
Op.t.Custom_function.
*)
#| Call of call_info
(*
Op.t.Call.
*)
#| Wmma of wmma
(*
Op.t.Wmma.
*)
#| Disjoint
(*
An Op.t.After of disjoint stores (after_disjoint).
*)

The type for node arguments. Each operation carries one shape of argument, given by v's table; the others carry No_arg.

#val equal_arg : arg -> arg -> bool

equal_arg a0 a1 is structural equality, with constants compared by Dtype.equal_const and nodes by ==.

#val pp_arg : Stdlib.Format.formatter -> arg -> unit

pp_arg formats an argument as the literal that denotes it: None, 42, ConstFloat(1.5), True, Invalid, dtypes.f32, (AxisType.WEAK, 0), (Ops.ADD, 0), 'CPU', b'ab\x00'.

The records print as calls:

Param: the slot, the type and the size, then the fields that differ from their defaults, by name: ParamArg(-1, dtypes.weakint, vmin_vmax=(1, 10), name='i', ...);
Kernel: KernelInfo(name='test', applied_opts=(), opts_to_apply=None, estimates=None, beam=0);
Program: ProgramInfo(global_size=(1, 1, 1), local_size=(1, 1, 1), params=(), ...);
Bufferize: BufferizeOpts(device='CPU', addrspace=AddrSpace.GLOBAL, removable=True, broadcast=False), removable False for Whole and Argument only and broadcast True for Broadcast only;
Custom_function: CustomFunction(name='f', dtype=dtypes.i32, fmt=False);
Call: CallInfo(None, 'f', False, False, False, None).

Raises Invalid_argument on an integer that Bigint.to_string refuses.

Nodes

#val v : ?src:t list -> ?arg:arg -> Op.t -> t

v ~src ~arg op is the node op of src (default []) with argument arg (default No_arg).

Its data type is derived from op, src and arg (dtype_of). The operations whose type comes from their argument need its shape: Op.t.Const a Const, Op.t.Cast and Op.t.Bitcast a Dtype, Op.t.Param, Op.t.Buffer and Op.t.Alloc a Param, and Op.t.Ins a Code.

An integer literal among src holds its value at the integer type the operation fixes for it, wrapped to that type's width, so every rule and fold reads the machine value of a literal without knowing to wrap it: a comparison at its operands' Dtype.least_upper, another binary or ternary operation at its own type, a cast at its target. On a Dtype.t.Uint8, 300 is 44 and -1 is 255. The node stays a bare weak literal. A shift's count and a power's exponent (is_number) are numbers and keep their value. A float type leaves a literal alone.

Raises Invalid_argument if no type can be derived, as for a Op.t.Where whose condition is not boolean.

An exception raised asynchronously inside v, as Sys.Break on Ctrl-C or a test runner's timeout is, passes through it: later calls on any domain find each node built before.

#val op : t -> Op.t

op u is u's operation.

#val dtype : t -> Dtype.t

dtype u is u's data type.

#val src : t -> t list

src u is u's sources, in order.

#val nth : t -> int -> t

nth u i is the ith source of u, from 0.

Raises Invalid_argument if u has no ith source.

#val arg : t -> arg

arg u is u's argument.

#val replace : ?op:Op.t -> ?src:t list -> ?arg:arg -> t -> t

replace u is u with the given fields replaced; u itself if none changes.

#val equal : t -> t -> bool

equal u0 u1 is u0 == u1: structural equality, since nodes are hash-consed.

#val compare : t -> t -> int

compare is a total order compatible with equal. It depends on the order in which nodes were built; see compare_structure for one that does not.

#val hash : t -> int

hash u is a hash of u, compatible with equal.

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

pp formats a node and its sources, one per line and indented, as the calls that build them:

UOp(Ops.ADD, arg=None, src=(
  x0:=UOp(Ops.CONST, arg=42, src=()),
  x0,))

A node reached more than once is named xN:= where first printed and xN afterwards. The argument prints as pp_arg, and raises as it does.

#val compare_structure : t -> t -> int

compare_structure u0 u1 orders by operation, then argument as pp_arg prints it, then data type, then sources, recursively. The order does not depend on how or when the nodes were built.

Raises as pp_arg does where it prints an argument.

#val structure_ranks : t -> t -> int

structure_ranks u numbers u and the nodes it reaches, call bodies included, in the order of compare_structure, and is the function that gives each its number: of two of them n0 and n1, n0's number is less than n1's exactly when compare_structure n0 n1 < 0, and the two are equal exactly when it is 0. Numbering a graph costs about its nodes times their logarithm, where compare_structure walks the graphs it compares.

The function raises Not_found on a node u does not reach. Raises as pp_arg does where it prints an argument.

#val key : t -> string

key u is a digest of u's operation, type, argument and sources, recursively: equal for equal graphs, whenever and wherever they are built, and different for NaN constants of different bits.

Raises as pp_arg does on an argument of u.

#val key_with : arg:(t -> arg option) -> t -> string

key_with ~arg u is the key of the graph u would be with each node n for which arg n is Some a taking the argument a, computed without building that graph. Each such a must be an argument n's operation takes and keep n's data type, as a storage parameter's slot does: the key is not that of a graph whose types it would change. arg is called once on each node of u.

Raises as key does.

#module Tbl : Stdlib.Hashtbl.S with type key = t

Hash tables keyed by nodes.

#module Nodes : sig ... end

Sets of nodes that remember the order in which nodes joined them.

Data types

#val dtype_of : Op.t -> t list -> arg -> Dtype.t

dtype_of op src arg is the data type of the node op of src with argument arg. Unary operations and movements keep their source's type, comparisons are Dtype.t.Bool, transcendentals widen to a float, broadcastable operations promote (promo_dtype), and effects are Dtype.t.Void.

Raises Invalid_argument if there is none.

#val promo_dtype : t list -> Dtype.t

promo_dtype us is the common type of us: theirs if they share it, their least upper bound (Dtype.least_upper) otherwise.

#val identity_element : Op.t -> Dtype.t -> Dtype.const

identity_element op dt is the value e of type dt with op e x = x for every x of type Dtype.strong dt (Dtype.strong), the width an accumulator of dt takes: 0 for Op.t.Add, 1 for Op.t.Mul, and that type's least value for Op.t.Max.

Raises Invalid_argument for any other op.

Graphs

A walk of a graph, by toposort, topovisit, graph_rewrite or a function built on them, that an exception ends, one raised asynchronously as Sys.Break is included, leaves later walks on any domain their results.

#type calls =

#| Enter
(*
A call's body is walked with its arguments.
*)
#| Skip
(*
A call's body is left out: the call is its arguments alone.
*)

The type for how a walk of a graph treats a call (Op.t.Call), whose first source is its body.

#val toposort : ?gate:(t -> bool) -> calls:calls -> t -> t list

toposort ~gate ~calls u is u and the nodes it reaches, each after its sources, in the order a depth-first walk of the sources, in order, finishes them. The walk enters only the nodes gate accepts (default all), and treats calls as calls says.

#val topovisit : t -> (t -> 'a) -> 'a Tbl.t -> 'a

topovisit u f cache is f u, after f has been applied to each node u reaches, each after its sources, and stored in cache. Nodes already in cache are not visited again.

#val backward_slice : calls:calls -> t -> Nodes.t

backward_slice ~calls u is the nodes u reaches, without u, in toposort ~calls's order.

#val backward_slice_with_self : calls:calls -> t -> Nodes.t

backward_slice_with_self ~calls u is u followed by backward_slice ~calls u.

#val op_in_backward_slice_with_self : calls:calls -> t -> Op.t list -> bool

op_in_backward_slice_with_self ~calls u ops is true iff u or a node of backward_slice ~calls u has an operation in ops.

#val reaches : ?gate:(t -> bool) -> calls:calls -> t -> t -> bool

reaches ~gate ~calls u x is true iff x is u or a node u reaches, treating calls as calls says, through nodes other than x that gate accepts (default all).

#val split_uop : t -> Op.t -> t list

split_uop u op is the operands of the tree of op nodes rooted at u, left to right: [u] if u's operation is not op.

Shapes

#val broadcast_shape : sint list list -> sint list

broadcast_shape shapes is the shape shapes broadcast to: aligned to the right, each axis is the size of its sources that is not 1, or 1 if all are.

Raises Invalid_argument if an axis has two different sizes other than 1.

#val to_max_shape : sint list -> int list

to_max_shape s is s with each symbolic size replaced by its greatest value.

#val sint_to_uop : ?dtype:Dtype.t -> sint -> t

sint_to_uop ~dtype s is s as a node of type dtype, by default Dtype.t.Weak_int.

Ranges

#val range_start : Op.t -> int option

range_start op is the index of the first source of an op node that is a range it ends: 1 for Op.t.Stage, Op.t.Reduce and Op.t.End, 0 for Op.t.Linear, and None otherwise: a call ends none of the ranges its arguments run inside.

#val ended_ranges : t -> t list

ended_ranges u is the ranges that u closes: those among its range sources, those a Op.t.Barrier or an Op.t.After's effects close, the loop a Op.t.Backedge repeats.

#val ranges : t -> Nodes.t

ranges u is the ranges u runs inside: u itself if it is a range, then those its sources run inside, minus those it closes.

#val axis_id : t -> int list

axis_id r is the identity of the range r.

Raises Invalid_argument if r is not a range.

#val axis_type : t -> Axis_type.t

axis_type r is the role of the range r.

Raises Invalid_argument if r is not a range.

#val bound_loads : t -> bool

bound_loads r is true iff the range r has a bound that, its ordering after effects (Op.t.After) aside, reads memory, as j < x[i]: how many trips the loop runs is what memory holds.

Raises Invalid_argument if r is not a range.

#val known_at_launch : t -> bool

known_at_launch r is true iff the range r has a bound that, its ordering after effects (Op.t.After) aside, reads no range and no memory (bound_loads), as j < n for a variable n and unlike j < x[i] or j < i + 1: a launch computes it from the run's variables alone, before the kernel runs, as it does a count of workgroups or of blocks.

Raises Invalid_argument if r is not a range.

#val compare_range_arg : t -> t -> int

compare_range_arg r0 r1 orders the ranges r0 and r1 by argument: axis type (Axis_type.compare), then identity, lexicographically, a prefix before its extensions.

Raises Invalid_argument if r0 or r1 is not a range.

#val range_str : ?color:bool -> t -> string

range_str ~color r is r's identity written with underscores, a negative part written m and its magnitude: 0_1, m1. It is in r's axis colour if color (default false).

#val multirange_str : ?color:bool -> ?pad:int -> t list -> string

multirange_str ~color ~pad rs is the range_str of each of rs, sorted by argument, separated by commas and padded with spaces to pad printed columns.

Evaluating

#val vmin : t -> Dtype.value

vmin u is a lower bound of every value u can take.

#val vmax : t -> Dtype.value

vmax u is an upper bound of every value u can take. Bounds follow interval arithmetic through the integer operations, the less-than and equality comparisons, selections, casts and loads from constant tables, a gated load's alternative included; elsewhere they are the bounds of u's type. A committed integer wraps at its width, and so does a weak integer operand of an operation on one (operand_bounds): an interval that leaves its type is one value wrapped, or the type's bounds. A table holding a NaN has its type's bounds. A comparison of floats is never bounded to true: a NaN compares false.

#val operand_bounds : t -> t -> Dtype.value * Dtype.value

operand_bounds u s is the bounds of u's source s as u reads it: an operation on a committed integer commits a weak integer operand to its type, which wraps it, so bounds that leave the type are one value wrapped, or the type's; an operation on a committed float commits a weak float operand to its type, which rounds it, and its bounds round with it. A shift's count is a number of bits and keeps its own bounds.

#val overflows : t -> Dtype.t -> bool

overflows u dt is true iff u's bounds reach outside dt's.

#val fits_no_int : t -> bool

fits_no_int u is true iff u's bounds are known and neither Dtype.t.Int64 nor Dtype.t.Uint64 holds them: no machine integer can carry every value of u. Bounds at a weak integer's own range, as of a node nothing bounds, are not known.

#val exact : Dtype.t -> Dtype.value list -> bool

exact dt vs is true iff dt is not a committed integer type, or each of vs is one of its values. Arithmetic on a committed integer type wraps at its width, so it computes what unbounded integers do only where every value it takes, vs, is exact.

val exec_alu : 
  ?truncate_output:bool ->
  Op.t ->
  Dtype.t ->
  Dtype.const list ->
  Dtype.const

exec_alu ~truncate_output op dt args computes the arithmetic operation op on the scalars args, one lane of a vector operation, exactly, then truncates the result to dt (Dtype.truncate) if truncate_output (default true). Integers are unbounded before truncation, and an integer to a negative integer power truncates toward 0: it is 0 but for 1 and -1. A shift's count and a power's exponent are numbers, read at no width: truncated to a committed integer type a left shift by a count of its width or more is 0, a right shift is the floor of its division by 2^count, whatever the count, and a power is computed in the type's bits, whatever the exponent. Out-of-domain values follow IEEE: log2 0. is -inf, sqrt of a negative is NaN, 1. /. 0. is an infinity of the operand's sign, sin of an infinity is NaN, and an overflowing exp2 or pow is the infinity IEEE gives. A float log2, exp2 and pow are correctly rounded to a double, ties to even, and sin is but where rounding it needs more than 192 bits past its argument's reduction; each is the same on every host. The NaN of an operation none of whose operands is NaN is Dtype.nan, whatever the host gives. Op.t.Fdiv is IEEE division: a zero divisor gives an infinity of the quotient's sign, and 0. /. 0. is NaN. Op.t.Trunc of a float is IEEE's, a float of its operand's sign: trunc (-0.5) is -0.. Op.t.Max is y if x < y and x otherwise, as compiled code computes it, so a NaN wins only as the first operand. A binary operation on `Invalid is `Invalid. Op.t.Where is the branch it picks, `Invalid included; its condition must not be `Invalid. Op.t.Where and Op.t.Neg convert nothing: truncated, a float NaN they give is Dtype.const of dt, which keeps the operand's sign, flipped by Op.t.Neg, and its payload, a signalling NaN staying one, in a float of 16 bits or more.

Raises Invalid_argument if op is not an arithmetic operation or the arguments do not fit it, as for a shift by a negative count, which has no value, a left shift of a non-zero integer by more than max_exact_shift bits that nothing truncates, or a power that nothing truncates and exact_power leaves out: at a weak dt, or without truncate_output.

#val is_number : t -> int -> bool

is_number u i is true iff source i of u is a number and no value of its peer's type, so that it never wraps at that type's width: the count of a Op.t.Shl or Op.t.Shr, or the weak integer exponent of an Op.t.Pow of an integer base.

#val exact_power : Dtype.value -> Dtype.value -> bool

exact_power b e is true iff b ** e, for integers b and e, can be computed exactly in exec_alu when nothing truncates: a base of 0, 1 or -1, a negative e, or a power of at most max_exact_shift bits.

#val max_exact_shift : int

max_exact_shift is the most bits a left shift that nothing truncates computes exactly in exec_alu. Past every committed integer's width, it bounds what one constant in a program may make the compiler allocate.

Construction

#val const : ?dtype:Dtype.t -> Dtype.const -> t

const ~dtype c is the constant c of type dtype. Without dtype, or for `Invalid, it is the weak literal of c's kind (Dtype.of_const). A constant of any other type is a Op.t.Cast of the literal, whose value is c converted to dtype (Dtype.const).

#val int : ?dtype:Dtype.t -> int -> t

int ~dtype n is const ~dtype (`Int n).

#val float : ?dtype:Dtype.t -> float -> t

float ~dtype x is const ~dtype (`Float x).

#val bool : ?dtype:Dtype.t -> bool -> t

bool ~dtype b is const ~dtype (`Bool b).

#val invalid : t

invalid is const `Invalid.

#val value : t -> Dtype.const

value c is the value of the constant c, read through a cast.

Raises Invalid_argument if c is neither a constant nor a cast of one.

#val is_literal : t -> bool

is_literal u is true iff u is a constant or a cast of one, whose value value reads. A constant of a type other than the weak ones is a cast of one: int ~dtype:Int32 16 is a cast of the weak 16.

#val is_invalid : t -> bool

is_invalid u is true iff u is the constant `Invalid.

#val ccast : t -> Dtype.t -> t

ccast u dt is u as a dt constant if u is a constant, and cast u dt otherwise.

#val cconst : Dtype.t -> Dtype.const -> t

cconst dt c is the cast to dt of the literal c, even where const would fold the cast away.

#val sink : ?kernel:kernel_info -> t list -> t

sink ~kernel us is the Op.t.Sink of us, the root of the kernel kernel if given.

#val group : t list -> t

group us is the empty Op.t.Stack ordered after us (after): the effects us merged without a value. When us is one node, it is that node.

#val is_group : t -> bool

is_group u is true if and only if u is an Op.t.After of the empty Op.t.Stack, as group makes of several effects.

#val broadcast : t -> int -> t

broadcast u n is the Op.t.Stack of n copies of u, or u if n is 1.

#val index : t -> t list -> t

index u idxs is the Op.t.Index of u by idxs; the ith source of a stack if u is a stack and idxs the constant i.

#val load : t -> t list -> t

load p rest is the Op.t.Load through p.

#val store : ?gate:t -> t -> t -> t

store p x is the Op.t.Store of x through p, guarded by gate. A tensor graph that stores through a broadcast is refused when it is prepared (Prepare.prepare_rangeify).

#val end_ : t -> t list -> t

end_ u rs is the Op.t.End of the ranges rs around u, or u if rs is empty.

#val backedge : t -> loop:t -> cond:t -> t

In a kernel, backedge u ~loop ~cond runs u, then repeats loop while cond holds.

Around a call of a schedule, or an Op.t.Linear of calls, it tests cond before each trip: loop is a Axis_type.t.Loop range, or the 0 that a range of one trip simplifies to, cond is storage of one boolean, and the calls run in order once per trip of loop while cond holds, so they may run no trip.

#val after : t -> t list -> t

after u deps is u ordered after deps, or u if deps is empty.

Outside a kernel, an Op.t.After of storage names a state of it. One whose deps list calls, loops of calls or stores is the state those effects leave the storage in, after the state u names: they land in the order listed, each reading the states the graph names, older states as they were, so a later store wins where stores overlap. A call loading through an argument it stores through loads what the effects listed before it left there, with each store the call runs after. An element two Op.t.Afters of one state both write, which Schedule.create_schedule admits for views of the state, holds an unspecified value. A state over the whole of its storage holds its own effects, those of the states under it, and each store those effects run after, by reading it directly or through other storage, or by listing it. A store of an Op.t.After over a view of one of those states that they do not run after is no part of it. A store listed over a later state of its storage than its destination names stores into that later state. One listing only stores into other storage leaves u's state as it is, read once they have landed. One whose deps are all Op.t.Afters is a join (is_join): it names no state of its own, holds the value of the state u names, and is read once the states deps name are reached.

#val after_disjoint : t -> t list -> t

after_disjoint u stores is after u stores with a promise from its maker: stores, the stores of one write of u's storage, write pairwise disjoint elements, none reads an element another writes, and every kernel of them may run at once. A schedule then runs the kernels of this write as one launch, where it can, without testing which elements each writes (Schedule.create_linear_with_vars).

Stores that break the promise give undefined values: a after of them, unmarked, runs them in the order listed.

#val is_disjoint : t -> bool

is_disjoint u is true iff u is the Op.t.After of an after_disjoint.

#val is_join : t -> bool

is_join u is true iff u is an Op.t.After whose dependencies are all Op.t.Afters: a join, which orders a read of its first source after them and writes nothing.

#val without_after : t -> t

without_after u is u without its Op.t.Afters.

#val barrier : t -> t list -> t

barrier u rest is the Op.t.Barrier of u and rest.

#val ins : ?src:t list -> ?dtype:Dtype.t -> t -> string -> t

ins u i is the instruction i with u's sources and type, unless given.

#val range : axis_type:Axis_type.t -> ?dtype:Dtype.t -> sint -> int list -> t

range ~axis_type ~dtype end_ axis_id is the loop variable counting from 0 to end_, excluded, of type dtype (default Dtype.t.Weak_int) and role axis_type, which has no default: a Axis_type.t.Weak range promises that its trips are independent, and a Axis_type.t.Loop runs them in turn. A loop ordered after effects has its bound ordered after them (after).

#val loop : int -> t

loop axis is the unbounded loop of identity [axis], which a backedge (Op.t.Backedge) repeats. It runs its trips in turn, whatever its type, and is no optimisation's axis.

#val special : sint -> string -> t

special end_ name is the hardware index name, bounded by end_.

val wmma : 
  ?upcast_axes:(wmma_axis list * wmma_axis list * wmma_axis list) ->
  t ->
  t ->
  acc:t ->
  dims:(int * int * int) ->
  threads:int ->
  t

wmma a b ~acc ~dims ~threads is the matrix multiply-accumulate of a and b into acc.

#val reduce : t -> Op.t -> t list -> t

reduce u op ranges is the kernel reduction of u with op over ranges.

#val get_idx : t -> t

get_idx u is u without its validity condition.

#val get_valid : t -> t

get_valid u is u's validity condition: true unless u is `Invalid.

#val bufferize : ?opts:bufferize_opts -> t -> t list -> t

bufferize ~opts u ranges is the Op.t.Stage of u over ranges, into a buffer placed by opts if given.

Placement

#val device : t -> device option

device u is where u lives: its argument's device for storage and copies, its sources' first otherwise.

#val on_disk : t -> bool

on_disk u is true iff u lives on a disk device.

#val copy_to_device : t -> device -> t

copy_to_device u d is the copy of u to d.

Raises Invalid_argument if d is a disk or u's type is weak.

#val size_src : int option -> t

size_src n is the first source of a Op.t.Param, Op.t.Buffer or Op.t.Alloc of n elements: the weak integer constant n, or the empty Op.t.Stack for a scalar (None). The size is never symbolic: storage of a symbolic size is a larger one shrunk.

Movement

#val base : t -> t

base u is u without its movements.

#val storage_base : t -> t

storage_base u is the storage u targets: base u without bitcasts and Op.t.Afters.

#type movement =

#| Reshape of sint list
(*
The new shape.
*)
#| Expand of sint list
(*
The axes added in front.
*)
#| Pad of (sint * sint) list
(*
Where each axis starts in the result, and the result's size.
*)
#| Shrink of (sint * sint) list
(*
Where each axis starts in the source, and the result's size.
*)
#| Permute of int list
(*
The source axis of each result axis.
*)

The type for the arguments of movements.

Storage

Sizes of storage are ints: a shape of more than max_int elements raises Invalid_argument.

#val addrspace : t -> Dtype.addr_space option

addrspace u is the address space u lives in. A node of type Void holds no memory and lives in Alu.

#val whole : t -> t

whole p is the parameter of the storage the Op.t.Param p views, for a view: a parameter of its own that a run of a buffer takes where a call stores into it or a loop reads it, so that the kernel's program does not depend on where the run starts. Any other node is its own whole.

#val view_state : t -> t

view_state p is the state of storage the view p names, its second source: its storage's parameter, or the Op.t.After whose elements a read of the run takes. Any other node is its own view_state.

#val buf_uop : t -> t

buf_uop u is the storage node u accesses.

#val is_virtual : t -> bool

is_virtual u is true iff u cannot back a buffer as it is: it has no device, or its type is weak.

#val has_buffer_identity : ?after_ok:bool -> t -> bool

has_buffer_identity ~after_ok u is true iff u names storage, through reshapes and, if after_ok, an Op.t.After.

#val needs_storage : t -> bool

needs_storage u is true iff realizing u allocates.

#val unique_num : unit -> int

unique_num () is a slot number no other call returns.

#val take_unique : int -> unit

take_unique n makes every number unique_num returns afterwards greater than n, up to max_int, past which no number is left.

Raises Invalid_argument if n is max_int.

#val new_buffer : ?slot:int -> ?phase:int -> device -> int -> Dtype.t -> t

new_buffer ~slot ~phase d size dt is a Op.t.Buffer of size elements of dt on d, in slot (default unique_num), whose first element lies phase bytes past a 16-byte boundary (default 0; see param_arg). A given slot is taken (take_unique), so unique_num returns a greater number afterwards. slot must not be a number unique_num returned for another buffer: buffers equal in slot, size, type, device and phase are one node and one storage.

Raises Invalid_argument if dt is weak or slot is max_int.

#val set : ?ends:t list -> t -> t -> t

set p x stores x through p, ends ends, and is p's storage after the store.

Variables

#val variable : ?dtype:Dtype.t -> string -> Dtype.value -> Dtype.value -> t

variable ~dtype name lo hi is the scalar variable name of type dtype (default Dtype.t.Weak_int) ranging over lo to hi.

#val expr : t -> string

expr u is the name of the parameter or buffer u.

Raises Invalid_argument if it has none.

Divisibility

#val const_factor : t -> Bigint.t

const_factor u is a known integer that divides every value of u.

#val pop_const : ?op:Op.t -> t -> t * Dtype.const

pop_const ~op u is (x, c) for u = op x c with c a constant, and (u, identity_element op) otherwise. op defaults to Op.t.Add.

Calls

#val body : t -> t

body c is the body the call c calls.

Raises Invalid_argument if c is not a call.

#val src_without_body : t -> t list

src_without_body u is a call's arguments, and any other node's sources.

#val opaque_call_bodies : Op.Set.t

opaque_call_bodies is the operations a call body can have: Op.t.Sink, Op.t.Program, Op.t.Linear, Op.t.Store and Op.t.Custom_function.

#val is_inline_call : t -> bool

is_inline_call u is true iff u calls a plain sink that is not compiled on its own.

#val has_unbound_outputs : t -> bool

has_unbound_outputs c is true iff the call c still has outputs that no buffer is bound to.

#val unbound_outputs : t -> t list

unbound_outputs c is those outputs, each ordered after c.

val custom_function : 
  ?dtype:Dtype.t ->
  ?shape:sint list ->
  ?fmt:bool ->
  string ->
  t list ->
  t

custom_function ~dtype ~shape ~fmt name src is the external function name, or with fmt (default false) the code name in the target's language (custom_function), which returns dtype (default Dtype.t.Void). Its value has no shape when dtype is Dtype.t.Void, and is of shape (default []) otherwise, which its first source states.

#val call : ?name:string -> ?precompile:bool -> ?inline:bool -> t -> t list -> t

call ~inline body args calls body on args; a call of code in the target's language is written into its reader if inline (default false). The call has the type and the shape of body. Given one state of storage at an argument body writes (Shape.written_params) and at another, it reads that state at the other as it was before the call, as memmove reads its source. A call compiled on its own takes storage, a run of its elements from the first, and a row a loop moves along as they are, and reads a copy of most other arguments, made once before a loop around the call; scheduling refuses one it would copy where the call writes it, and a constant (Indexing.run_rangeify). In a loop, an argument the call only reads that reads a state the loop writes, or a later one, is computed as each trip starts, each read in it taking its own state: a state older than the one the loop is listed over as it was. A value that moves with the loop and is no row of storage, or that reads, at a state the loop writes or a later one, a join over other writes of that storage, is refused.

Raises Invalid_argument if body cannot be called (opaque_call_bodies) or a range leaks out of it.

Elementwise

#module type Elementwise = sig ... end

Elementwise operations.

include Elementwise with type t := t

#val alu : t -> Op.t -> t list -> t

alu x op rest is the operation op on x :: rest as given.

#val cast : t -> Dtype.t -> t

cast x dt is x converted to dt, or x if it is of type dt.

Arithmetic

add, mul, lt, eq, maximum, bitwise_and, bitwise_or, bitwise_xor, shl and shr are Op.t.Add, Op.t.Mul, Op.t.Cmplt, Op.t.Cmpeq, Op.t.Max, Op.t.And, Op.t.Or, Op.t.Xor, Op.t.Shl and Op.t.Shr of their promoted operands. reciprocal, trunc, sqrt, exp2 and log2 are Op.t.Recip, Op.t.Trunc, Op.t.Sqrt, Op.t.Exp2 and Op.t.Log2 of their operand. The others are built from these.

#val add : t -> t -> t

#val mul : t -> t -> t

#val lt : t -> t -> t

#val eq : t -> t -> t

#val maximum : t -> t -> t

#val bitwise_and : t -> t -> t

#val bitwise_or : t -> t -> t

#val bitwise_xor : t -> t -> t

#val shl : t -> t -> t

#val shr : t -> t -> t

#val reciprocal : t -> t

#val trunc : t -> t

#val sqrt : t -> t

#val exp2 : t -> t

#val log2 : t -> t

#val sub : t -> t -> t

sub x y is x + neg y.

#val neg : t -> t

neg x is logical_not x for booleans and x * -1 otherwise.

#val div : ?rounding:[ `Trunc | `Floor ] -> t -> t -> t

div ~rounding x y divides integers by Op.t.Truncdiv or Op.t.Div when rounding; otherwise it multiplies x, as a float, by y's reciprocal, then rounds as asked.

#val mod_ : t -> t -> t

mod_ x y is the remainder of div ~rounding:`Floor x y.

#val fmod : t -> t -> t

fmod x y is the remainder of div ~rounding:`Trunc x y.

#val floor : t -> t

floor x is trunc x, less one where that is above x.

#val pow : t -> t -> t

pow x y is Op.t.Pow of x and y: promoted to a common type, but for an integer x and a weak integer y, which keeps its value at no width (Op.t.Pow's exponent is a number).

Raises Invalid_argument if the promoted type is not a float and y is a negative integer constant.

#val gt : t -> t -> t

gt x y is lt y x.

#val le : t -> t -> t

le x y is logical_not (gt x y).

#val ge : t -> t -> t

ge x y is logical_not (lt x y).

#val ne : t -> t -> t

ne x y is logical_not (eq x y).

#val logical_not : t -> t

logical_not x is eq (cast x Bool) false.

#val bitwise_not : t -> t

bitwise_not x is logical_not x for booleans, and x xor all ones otherwise.

#val where : t -> t -> t -> t

where c x y is x where c holds and y elsewhere.

#val minimum : t -> t -> t

minimum x y is neg (maximum (neg x) (neg y)) for floats; for integers, the same through an order-reversing xor.

#module O : sig ... end

Operators, for local opening: O.(x + int 4 * y). Each is the function of the same meaning: + is add, // is div ~rounding:`Floor, % is mod_, <> is ne, land, lor, lxor and lnot are the bitwise operations, lsl and lsr the shifts.

#val bitcast : t -> Dtype.t -> t

bitcast x dt reinterprets x's bits as dt, or is x if it is of type dt.

Raises Invalid_argument if either type is weak.

#val commit_dtype : ?default_int:Dtype.t -> t -> Dtype.t

commit_dtype x is the type x is stored as: a weak integer commits to the narrowest type holding its bounds (Dtype.commit_int), a weak float to its default (Dtype.strong).

#val element_size : t -> int

element_size x is the size in bytes of one of x's elements.

Raises Invalid_argument if x's type is weak.

#val contiguous : t -> t

contiguous x is x materialised into its own buffer: x itself if it is weak, already staged, placed nowhere, or storage.

#val usum : t -> t list -> t

usum x ys is the sum of x :: ys, or their disjunction if x is boolean.

#val uprod : t -> t list -> t

uprod x ys is the product of x :: ys, or their conjunction if x is boolean.

Patterns

#module Upat : sig ... end

Patterns.

#module Pattern_matcher : sig ... end

Pattern matchers.

Rewriting

#exception Bottom_up_gate

Raised by a rule before the sources (Before_sources, Around_sources's before) under Fixed_point to keep the node it last produced and leave its sources unvisited. Raised by a rule after the sources, or under Once, it escapes graph_rewrite.

#type pass =

#| Fixed_point
(*
A rule before the sources rewrites a node until none applies, and the sources of the result are then rewritten. A node that a rule after the sources rewrites, or that is rebuilt on rewritten sources, is rewritten again in turn, to a fixed point.
*)
#| Once
(*
A rule rewrites a node at most once, and its result is never rewritten again: a node that a rule before the sources rewrites is replaced and its sources are left as they are.
*)

The type for how often a rewrite applies its rules to a node and its results.

#type 'ctx rules =

#| After_sources of ('ctx, t) Pattern_matcher.t
(*
Rewrites each node once its sources are rewritten.
*)
#| Before_sources of ('ctx, t) Pattern_matcher.t
(*
Rewrites each node before its sources.
*)
#| Around_sources of {
1. #before : ('ctx, t) Pattern_matcher.t;
2. #after : ('ctx, t) Pattern_matcher.t;
}
(*
Rewrites each node with before before its sources and with after once they are rewritten.
*)

The type for the rules of a rewrite and when they meet a node.

val graph_rewrite : 
  calls:calls ->
  pass:pass ->
  ctx:'ctx ->
  t ->
  'ctx rules ->
  t

graph_rewrite ~calls ~pass ~ctx u rules is u rewritten by rules as pass says: each node is rebuilt on its rewritten sources and rewritten by the first rule that applies, and shared nodes stay shared. A call's body is rewritten with its arguments under Enter, and left alone under Skip.

Raises Invalid_argument if a Fixed_point rewrite does not terminate: a rule before the sources cycles, the work list exceeds Setting.rewrite_stack_limit, or a node's rewrite depends on itself, and Bottom_up_gate as that exception's description says.

val substitute : 
  ?extra_pm:(t Tbl.t, t) Pattern_matcher.t ->
  calls:calls ->
  pass:pass ->
  t ->
  (t * t) list ->
  t

substitute ~calls ~pass u subs is u with each node of subs replaced by its pair, top-first, rewriting with extra_pm as well, which reads the substitution, and treating calls and results as graph_rewrite does. Nodes paired with themselves are ignored.

#val assuming : t -> (t * t) list -> t

assuming u subs is u with each node of subs replaced by its pair, as substitute does outside calls, for a fact that holds only where u is read: it leaves each range it does not replace whole, bound included, since a loop is one loop wherever it is read.

#val pm_substitute : (t Tbl.t, t) Pattern_matcher.t

pm_substitute replaces each node its context maps by its image.

#val pm_drop_after : (unit, t) Pattern_matcher.t

pm_drop_after replaces each Op.t.After with its first source.

#val resolve_returned_after : t -> t -> t option

resolve_returned_after r effects is the value stored into the output r by the one store of effects into r's storage: r after that store if r is a parameter of the enclosing function.

#val gate_kernel_sink : t -> bool

gate_kernel_sink u is false for a linear program and a kernel's sink, so walks gated by it do not enter kernels.