Module Tolk.Shape

shape_opt u is u's shape, or None for nodes that have none, such as effects and program structure. Movements check their argument against their source's shape; broadcastable operations broadcast their sources' shapes (Ops.broadcast_shape).

Raises Invalid_argument if a movement does not fit its source's shape, or if sources cannot be broadcast.

shape u is u's shape.

Raises Invalid_argument if u has none.

ndim u is the length of u's shape.

numel u is the product of u's shape.

max_shape u is u's shape with each symbolic size replaced by its greatest value (Ops.to_max_shape).

max_numel u is the product of max_shape u.

broadcast_axes src out is the axes of out that broadcasting src to out adds or expands.

Raises Invalid_argument if src has more axes than out.

resolve ~default u is the value of the boolean node u if its simplification (simplify) has one possible value, and default (default true) otherwise.

Raises Invalid_argument if u is not boolean.

simplify u is u rewritten with symbolic to a fixed point. A constant is itself, and so is a sink of constants and stacks of constants, which the rules leave as they are.

ssimplify u is simplify u as an Int if it is an integer constant, a Sym otherwise.

smax ss is the greatest of ss, a Op.t.Max node if some are symbolic, simplified.

Raises Invalid_argument if ss is empty.

smin ss is the least of ss, as smax.

to_bool u is the value of the boolean node u.

Raises Invalid_argument if u is not boolean or its simplification has more than one possible value.

to_z u is the value of the integer node u, as to_bool.

to_float u is the value of the float node u, as to_bool.

Symbolic integers.

sym_infer s vars is the value of s with each variable named in vars bound to its value. Integer arithmetic is exact, divisions round as their operations say, and casts convert without truncating.

Raises Invalid_argument if s reads a variable vars lacks, or if its value does not fit an int (sym_compile_z computes one of any size).

sym_compile s var is the function computing s in an environment, each variable v of s read by var v: sym_compile s var env is sym_infer s vars for vars binding each variable v to var v env. s is simplified once, and integer arithmetic that fits an int is computed on ints, so that a call costs a few operations: a value that a schedule computes on each run, such as the offset of a view that moves with a range.

The function raises as sym_infer does, and as var does.

sym_compile_z s var is sym_compile's function with values of any size: sym_compile_z s var env is the exact value of s, an integer whether or not it fits an int, such as a kernel's estimated operations over a loop bounded by a value loaded from memory (Ops.estimates). It costs what sym_compile does where values fit an int.

The function raises as sym_compile's does, but never for a value past an int.

consts ~dtype cs is the stack of const ~dtype c for each c of cs. dtype defaults to the committed type of the literals cs (Dtype.of_consts).

const_like ~dtype u c is the constant c of type dtype (default u's), expanded to u's shape.

vconst_like u c is the constant c of u's type repeated max_numel u times.

stack ~axis us is the Op.t.Stack of us, each converted to their common type (Ops.ccast) except `Invalid constants, with the new axis moved to axis (default 0).

Raises Invalid_argument if us is empty or its nodes differ in shape.

rop u op axes reduces the axes axes of u with op: axes of size 1 are reshaped away, and the others are permuted to the front and reduced by one Op.t.Reduce.

valid u cond is u where cond holds and `Invalid elsewhere.

contract u rs is the stack of u with the upcast ranges rs substituted by each of their values, the last range varying fastest.

Raises Invalid_argument if a range of rs is not Ops.Axis_type.t.Upcast.

marg u is the argument of the movement u.

Raises Invalid_argument if u is not a movement.

as_shape u is the shape the node u denotes: its constant, its stack's sources, or u itself.

mop u m applies the movement m to u as one node, its shape arguments stored as simplified sources. An empty expand is u, and so is an empty pad or shrink of a scalar.

Raises Invalid_argument on an empty pad or shrink of a node that is not a scalar, or where shape_opt does.

reshape u shape is u with shape shape, or u if unchanged; a size of -1 is inferred.

Raises Invalid_argument if the element counts differ or -1 appears twice.

expand u shape broadcasts u to shape; a size of -1 keeps u's.

permute u order is u with its axes in order, or u; negative axes count from the end.

Raises Invalid_argument if order is not a permutation.

flip u axes reverses axes of u, negative axes counting from the end.

Raises Invalid_argument if an axis is out of range or given twice, or if u is a scalar and axes is not empty.

shrink u bounds keeps, on each axis, the elements from its start to its end, excluded; None keeps the axis whole. It is u if it keeps every element and each start is an integer: a symbolic start, even one of 0, stays, with the effects it waits on.

shrink_to u shape keeps the first elements of each axis.

pad ~value u padding adds, on each axis, the given numbers of elements of value (default 0) before and after it; a negative number removes elements.

pad_to ~value u shape pads the end of each axis to shape.

flatten ~start ~stop u merges the axes from start (default 0) to stop (default -1), included.

unflatten u axis sizes splits axis into sizes.

squeeze ~axis u removes axis if its size is 1, or every axis of size 1.

unsqueeze u axis is u with an axis of size 1 inserted at axis, which counts from the end of the new shape when negative.

Raises Invalid_argument if axis is out of range.

transpose u a b exchanges the axes a and b of u.

Raises Invalid_argument if an axis is out of range.

split ~axis u sizes is the consecutive slices of u along axis (default 0) of sizes elements.

Raises Invalid_argument if axis is out of range or of symbolic size, or if sizes does not sum to its size.

repeat u repeats tiles u repeats times along each axis, the axes aligned to the right: a repeats longer than u's shape adds leading axes.

Raises Invalid_argument if a repeat is negative, or if a movement does.

pool ~stride ~dilation u kernel is the windows of kernel over the last axes of u, each stride apart (default 1) and dilation between its elements (default 1): u's leading axes, then the number of windows along each pooled axis, then the kernel's axes. Only movements build it: u is repeated and read back in rows one element longer, so that windows overlap without padding.

Raises Invalid_argument if u has fewer axes than kernel, if stride or dilation do not have one entry per kernel axis, or if a dilated kernel is longer than its axis.

cat ~axis u rest concatenates u :: rest along axis (default 0).

cumalu u axis op is the inclusive running op (Op.t.Add, Op.t.Mul or Op.t.Max) of u along axis, at u's type: each element reduces the window of the elements up to it. Past 16 elements it runs within chunks of 16 from the axis's first element, then along the chunks' last elements the same way, and combines each chunk with the result before it: an axis of n elements takes O(n) operations in about log₁₆ n stages. An element's operations and their order depend on its index alone, but for operations with op's identity: the running op of a prefix of the axis is the prefix of its running op, bit for bit, but for a zero sum's sign and a float maximum's result at a NaN, which Op.t.Max gives in its operands' order. It is u if an axis of u is empty.

Raises Invalid_argument if axis has a symbolic size.

arange ~start ~step ~dtype stop is the vector of the integers from start (default 0) up to stop, excluded, by step (default 1), down to it for a negative step, of type dtype (default the first of Dtype.default_int, Dtype.t.Int32, Dtype.t.Int64 and Dtype.t.Uint64 that holds them). It is empty if stop is not beyond start. Element i is the integer start + i * step, exact in an integer type. In a float type it is computed in that type, or Dtype.t.Float32 if narrower, up to 1518500249 elements, and past that it is the integer rounded once.

Raises Invalid_argument if step is 0, dtype does not hold the integers, or there are more than max_int - 2^30 of them.

nbytes u is the number of bytes of u's elements.

empty ~device shape dt is uninitialised storage of shape for whoever realizes the graph to bind.

Raises Invalid_argument if dt is weak.

empty_like u is uninitialised storage of u's shape, type and placement.

clone ~device u is a copy of u into new storage on device (default u's).

Raises Invalid_argument if device is a disk.

alloc ~slot shape dt is call-local storage of shape, in slot (default Ops.unique_num). A given slot is taken (Ops.take_unique), so Ops.unique_num returns a greater number afterwards. slot must not be a number Ops.unique_num returned for other storage: allocs equal in slot, size, type and device are one node and one storage.

Raises Invalid_argument if slot is max_int.

alloc_like ~slot u is alloc of u's greatest shape and type.

placeholder ~slot shape dt is a parameter of shape and dt (dt committed, Dtype.strong), or workgroup or register storage for those address spaces, in slot (default Ops.unique_num). A given slot is taken (Ops.take_unique), so Ops.unique_num returns a greater number afterwards.

Raises Invalid_argument if the address space is the ALU's or slot is max_int.

placeholder_like u slot is a placeholder in slot of u's greatest shape and type.

Raises Invalid_argument if u's shape has a symbolic size or slot is max_int.

param ~shape slot dt is the parameter slot of type dt: a scalar without shape, flat storage of its greatest size viewed as shape otherwise, whose first element lies phase bytes past a multiple of align (defaults 0 and 16; see Ops.param_arg).

Raises Invalid_argument if dt is weak.

param_like ~name u slot is a parameter in slot, named name (default none), that u can be passed to: a scalar variable without its value, or storage of u's shape. Storage has the phase and alignment of the storage u views (storage_phase).

Raises Invalid_argument if the phase does not fit u's type.

storage_phase u is the alignment and phase (align, phase) (Ops.param_arg) of the storage u views: its storage's, moved by the bytes a shrink skips when it shrinks the storage seen whole, a buffer or a stage the schedule allocates, in row-major order. A shrink by a symbolic start known only to multiples of fewer bytes than the storage's alignment, such as a window that moves with a range, lowers the alignment to the largest power of two those bytes are a multiple of, and a symbolic start into a view that reorders or pads the storage keeps only the element's size. Any other view keeps its storage's: a reordering, a bitcast or an ordering. A stage of a view of a buffer through movements and bitcasts is the view when scheduling finds it contiguous, and storage of its own on a 16-byte boundary otherwise, so it has what both hold: phase 0 and the largest power of two up to the buffer's alignment that the view's first byte is a multiple of; (1, 0) when a size is symbolic. A view whose first or last element is padding, or whose last element does not lie as many elements past its first as a run of its size, is no run, and its stage has (16, 0). Storage on a disk, which no vector access reads, and anything else that is not storage or a view of it have (16, 0).

fits_start ~param start is true iff storage that starts at start (storage_phase) starts where a parameter that assumes param takes it: known modulo at least the parameter's alignment, at its phase.

allocated_start is (16, 0), the start of storage that scheduling allocates: on a 16-byte boundary.

param_start body slot is the start (align, phase) that the storage parameter in slot of the call body body assumes (Ops.param_arg), or None if body takes no storage there.

view_as u shape views the flat storage u as shape.

is_variable u is true iff u is a scalar variable with a range.

is_bound_var u is true iff u is a variable bound to a value.

bind v x is v bound to x.

Raises Invalid_argument if v is not an unbound variable, or x is out of its range.

unbound v is v without its value and tag.

unbind v is (unbound v, x) for v bound to x.

Raises Invalid_argument if v is not bound.

unbind_all u is u with each bound variable unbound, and each variable with its value.

variables u is the unbound scalar variables u reads, sorted by name and slot.

known_size s is s's value when it is an integer: an Ops.sint.Int or a literal constant, which may lie past an int.

divides u n is u / n if u is known to be a multiple of n.

gcd us is a common divisor of us: their common factors times the greatest common divisor of their constant coefficients, a scalar constant whatever the shape of us.

Raises Invalid_argument if us is empty.

divide_exact u d is u / d if it divides exactly.

store_call dst src is the call that copies src into dst.

call_with_outputs values args calls a body that computes values from args, and is the outputs, each ordered after the call. Each output is new storage passed to the call at its position in output_pos (default after args); args fill the other positions in order.

Raises Invalid_argument if output_pos is not strictly ascending within the argument list.

call_with_output value args is the one output of call_with_outputs [value] args.

written_params body is the slots of the storage parameters the call body body writes, in increasing order: those a store of a Op.t.Sink or an Op.t.Store stores into, through indices, movements and Op.t.Afters, and, for a call the body makes, the parameters it passes where that call's body writes, at any depth; a call whose body does not show its writes writes every parameter it is passed. A body of an Op.t.Linear, a schedule's calls or the lines a body inside a kernel is lowered to, writes what its calls and stores write, in loops too. None when body does not show what it writes: a Op.t.Program or a Op.t.Custom_function, the body of a call of code in the target's language among them.

reads_storage_written c i is true iff the call c, whose body shows its writes (written_params), only reads its argument i and writes the storage it holds (Ops.storage_base, Ops.whole) through another.

reads_state_written c n is true iff the node n, storage or a state or view of it, names a state the call c writes through another argument or a later one.

custom_kernel args f calls the kernel f builds on placeholders of args, and is args, each that the kernel writes (written_params) ordered after it: an argument it only reads is itself. The placeholder of the ith argument takes slot i: a variable's is the variable, unbound, and any other's is a buffer of its shape (placeholder_like). An argument the kernel writes is storage, a run of its elements from the first, or a row a loop moves along, on a 16-byte boundary as its placeholder assumes, and so is one a loop moves along that it reads: scheduling refuses any other, which the kernel would write a copy of or access off its boundary (Indexing.run_rangeify).

is_launch u is true iff u is a launch: a kernel's Op.t.Sink whose sources are calls of kernels compiled on their own (Ops.call_info), its members. A launch runs its members at once, each lowered as a kernel of its own (Codegen.to_program), so a member must touch no element that another writes.

launch calls is the call of the launch of the kernel calls calls: its arguments are those of calls, each once, storage first, and each member calls its kernel, compiled on its own, on the launch's parameters at those arguments, a variable its kernel holds in an argument's slot becoming a nameless parameter of that slot. The launch's kernel information is named "launch" and optimized, since each member keeps its own optimisations.

Raises Invalid_argument if a call is not of a kernel's sink, or is compiled on its own.

launch_members c is the calls of the members of the launch c calls, each on c's arguments: launch_members (launch cs) is cs.

Raises Invalid_argument if c is not a call of a launch.

program_info_of_sink ~target sink reads the program's launch dimensions from its Op.t.Special nodes, each named variable they read at slot -1, as a launch binds it by name; its parameters and buffers from its Op.t.Param nodes; and which buffers it reads and writes from its loads and stores, and those of the bodies of the calls it makes at the arguments they pass. When none is found reading or writing, every buffer is taken to do both. target defaults to the empty target.

kernel_params prg is the parameters of the linearized Op.t.Program prg in slot order, the order its kernel takes them, as parameters or packed into one: those of Ops.program_info.params its linear order holds, then each other one whose slot none of those takes.

Raises Invalid_argument if prg is not a linearized program.

launch_dims p vars is p's global and local sizes with each variable named in vars replaced by its value (sym_infer).

invalid_gate matches where cond x i with i the constant Ops.invalid, naming its sources "cond", "x" and "i".

pm_remove_invalid replaces each Ops.invalid that a gate or a stack holds by 0 of the holder's type, once the gates have been read.

fits x v is true unless x is of a signed integer type that does not hold v: the rules that make a constant from others make it only where it fits (symbolic_simple, symbolic).

symbolic_simple folds nodes whose value follows from their operation and their sources alone:

• invalid values: an operation on a gated value moves inside the gate; an operation on Ops.invalid is invalid; a gate on a reduction's operand that does not depend on its ranges moves out of the reduction, and one that does selects the reduction's identity where it fails; a store to an invalid index does nothing, and a load from one is its alternative value, or 0;
• identities: x + 0, x lxor 0, x lor 0, x lsl 0, x lsr 0, x * 1 and x // 1 are x, a float x + 0 only for -0.; x // -1 is -x for a signed or weak x; x // x is 1 where x's bounds hold no 0; (x lxor y) lxor y is x for a weak y; (x % y) % y is x % y; a boolean x land c and x lor c with c constant are x or c; x = true, a double negation, where x true false and an idempotent operation of x with itself are x; where x false true is the negation of x; a boolean cast to an integer and equal to 0 or 1 is the boolean's negation or the boolean, and to any other integer false; the truncation of an integer is itself;
• recombination: in a weak integer sum, (b % d) * m and a term q * (d * m), where q is b' // d for some b' congruent to b modulo d up to constants, recombine into b' * m; where q is (b' // d) % k with k positive, into (b' % (d * k)) * m;
• zeros: x < x is false, x = x is true for integers and booleans, x % x, x lxor x and x land 0 are 0, and so is x * 0 for integers and booleans; a mask that clears only bits that a right shift drops is removed;
• constants: an arithmetic operation on constants is its value (Ops.exec_alu), except Op.t.Threefry; weak constants keep their mathematical value, a committed integer holds its type's value (a weak operand of an operation on one is read, and the result written, at its width, but a shift's count and a power's exponent, numbers, are read as they are), and an operation mixing weak and committed constants commits the weak ones to the promoted type; a cast of a constant is the constant of the cast's type; 0 / 0 is NaN;
• booleans: a boolean * is land, and a boolean + or maximum is lor;
• casts: a cast or bitcast to its operand's type is its operand; a bitcast of a constant is the constant with the same bits; a cast through a type that holds every value of the result's type, back to that type, is its operand; two bitcasts are one; a cast to a boolean is x <> 0;
• powers: x ** c, for a constant c that is an integer or a half integer, 0 or at least 1 in magnitude, is a product of powers of x, its reciprocal and its square root, a float half-integer power being +0. at -0. and +inf at -inf; c ** x is 1 of the power's type if c = 1, and exp2 (x * log2 c) for positive finite c, log2 c correctly rounded (Ops.exec_alu); a 64-bit integer packed from two 32-bit halves and unpacked again is the half read;
• selections: a selection between equal values is that value, and a selection by a constant is the branch it picks, keeping the selection's type;

then cleans up movements (mop_cleanup).

commutative orders the two operands of each commutative weak integer operation by Ops.compare_structure, least first, so that two sums of the same terms become the same node.

symbolic is symbolic_simple and commutative, followed by rewrites that match deeper:

• terms: x lor not x is true; x + x is x * 2; for integers, like terms combine where the factor they make fits, x * c0 + x * c1 into x * (c0 + c1) and y + x + x into y + x * 2, also as the last two terms of a longer sum; c * (x + c') is c * x + c * c' for a weak integer x;
• selections: a selection by a negation swaps its branches; within where c t f, c is true in t and false in f, unless c is a constant through movements, which is one node for every use of that constant, c, t or f reaches an Op.t.Index or an Op.t.Load, or c reaches the operand of a conversion of a float to an integer in a branch, where it may be the guard of a value that conversion is undefined for; a range's bound keeps c whatever the branch, since a loop is one loop wherever it is read; where g x 0 <> 0 is g land (x <> 0); nested selections sharing a branch merge their conditions with land or lor, or with land of the inner one's negation where the outer's other branch is the inner's first; an operation on two selections by the same condition, one of whose branch pairs is constant, selects between the operations on the branches, also as the last two terms of an integer sum; for integers, where c t 0 + where c 0 f is where c t f;
• bounds: a less-than or equality comparison, division, remainder, variable, Op.t.After, Op.t.Special or range with a constant end whose bounds are equal is that constant; an integer maximum of two operands whose bounds, as it reads them (Ops.operand_bounds), do not overlap is the greater, committed to its type; an integer selection that computes a maximum is Ops.maximum: where (a < b) b a with a or b a constant, and where (a < b) c (a + (c - b)) for a signed a and constants b and c, where a + (c - b) does not wrap;
• constants: two applications of an associative operation to constants fold the constants together, sums, products and maxima for integers only, a signed result only where it fits, and a weak result on an unsigned integer wrapped to its width, as (x * 65536) * 65536 on uint32 is x * 0; (x // c1) // c2 is x // (c1 * c2) for positive c2 where c1 * c2 does not wrap (Ops.exact); constants move to the end of integer sums and products, and a factor -1 to the end of float products;
• comparisons, on integers: c0 + x < c1 is x < c1 - c0 where neither side wraps; c0 * x < c1 divides both sides by c0, rounding up, and flips x's sign if c0 is negative; x // d < c is x < c * d for positive d; in x < c with c positive, a divisor d common to c and the coefficients of some terms of x, whose other terms stay within 0 and d - 1, divides out; -x < -y is y < x; not (x < 1) with x a sum of non-negative terms with positive coefficients drops the coefficients;
• ranges: a range modulo its end is the range, and divided by its end is 0;
• casts: a cast through a type that holds every value of the operand is one cast; a cast of an integer that fits the intermediate integer type is one cast; a cast of an unsigned x land y, or x lsr k with k below x's width, to a wider unsigned type is the operation on the widened operands; a binary operation on 64-bit or weak integers whose values fit 32 bits computes in Dtype.t.Int32 and casts back, but for a right shift whose count can reach 32 and a remainder of a dividend that can be the least Dtype.t.Int32 by a divisor that can be -1, which have no value there; a cast of x + c to a signed integer is the cast of x plus c;
• ordering: an Op.t.After waits only on effects (Op.Set.effectful), ranges, ends, backedges, linear programs, stages and Op.t.Afters of global storage, on a reduction's source with the reduction's ranges ended, and on the sources of any other node it names; an Op.t.After or Op.t.End of nothing is its value, and an Op.t.End drops the ranges that became constants;

then simplifies division and remainder (div_and_mod_symbolic), and restores bare constant operands (pm_uncast_const).

div_and_mod_symbolic rewrites divisions and remainders. With x and y nodes and a, c and d constants:

• (x // c + a) // d is (x + a * c) // (c * d) when d is positive and, for a committed integer x, no value wraps (Ops.exact);
• for a weak integer x and c % d other than c, (x + c) // d is (x + c % d) // d + c // d, and (x + c) % d is (x + c % d) % d.

Any other weak integer x // y or x % y takes the first of the following forms that applies. When the quotient has one possible value q, x // y is q and x % y is x - q * y. When x is a parameter declared a multiple of m, and m is a multiple of the constant y, x % y is 0 and x // y stays as it is.

When y is a positive constant c, with x a sum of terms and a constant:

• (z % (k * c)) // c is (z // c) % k for a positive k;
• in x % c, a term z % m, with m a multiple of c, is z;
• when replacing each term's factor by one of its residues modulo c leaves a sum within one period of c, that sum gives the remainder, and the factors less their residues give the quotient;
• a divisor common to c and every term's factor is divided out of both, when the quotient stays non-negative;
• for each factor f of a term that divides c, x // c is (x // f) // (c / f), and x % c its reconstruction, keeping the smallest result.

Otherwise:

• a common divisor of y and of the terms of x is divided out of both;
• for x and y non-negative, the terms of x that are multiples of y leave the division.

Raises Division_by_zero if a weak integer division or remainder has a divisor that is always 0, or a positive constant divisor and a dividend with a term that is the constant 0 or a product by it. The symbolic rules fold such a term away before these rules see it.

The exception number raises on `Invalid.

number c is the value of c.

Raises Not_a_number if c is `Invalid.

rule p f is Ops.Pattern_matcher.rule p f, which declines where f raises Not_a_number.

equals u v is true iff the constant u holds the number v.

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

const_v u v is v as a constant of u's type (const_like).

conj us is the conjunction of us.

Raises Invalid_argument on the empty list.

invalid_pat matches the constant `Invalid, named "i".

rule_ctx p f is Ops.Pattern_matcher.rule_ctx p f. Unlike rule, it lets Not_a_number through.

with_ctx m is Ops.Pattern_matcher.with_ctx m.

mop_cleanup merges and removes movements and indexing:

• a shrink of a shrink is one shrink of the inner one's source, starting at the sum of their starts, of the outer one's sizes;
• a reshape of a reshape is one reshape of the inner one's source, and a reshape to its source's shape is its source;
• a permutation of a permutation is their composition;
• a stack of the elements 0, 1, ... of a node x, in order and of x's shape, is x;
• indexing a stack by a constant i and further indices is its ith source indexed by those indices;
• indexing an index, when every index of both is a scalar, is one index by the inner indices followed by the outer ones;
• indexing the storage b indexed by i with as many indices as i has axes is indexing b by the element of i they select.

is_weak dt is true iff dt is one of Dtype.weaks.

weak u is is_weak (Ops.dtype u).

is_const u is true iff u is a constant node.

unchanged src u is true iff src is u's sources, node for node.

derived_dtypes u src is the operands' meet and u's own type when u is broadcastable and sources src commit both to a width.

pm_uncast_const removes the cast from each committed constant source of a broadcastable node, leaving the bare weak literal, when the node's sources still derive the same least upper type and the node the same type. Rules keyed on constant values then match the literal.