gizai:tensor/lazy

fused elementwise chains (Lazy): a small tensor IR run 1024 elements at a time in L1, and products (Product) that a chain feeds or reads.

import * as Lazy from "gizai:tensor/lazy";

Types

Step

type Step = { op: i64; a: i64; b: i64; }

One operation of a chain: its op code and its operands (inputs or earlier steps).

• op: i64
• a: i64
• b: i64

Lazy

type Lazy

An elementwise expression over tensors, evaluated by eval. With a _product, input 0 is that product, not yet computed: on a GPU the expression runs as the product is stored (its epilogue), so C is written once.

• shape: Vec<i64>
• inputs: Vec<_gizlang_lib_ai__tensor__storage.Tensor>
• steps: Vec<M49.Step>

Product

type Product = { a: Tensor; ta: bool; b: Tensor; tb: bool; m: i64; k: i64; n: i64; compute: i64; low: Vec<Tensor>; lowAt: i64; patch: Vec<i64>; }

A times B as matmulT takes them: m x k times k x n, on a GPU with the operands rounded to compute (-1: as they are; lazyMatmulIn). low is [D, U] of a low-rank product D Uᵀ (D m x r, U n x r, both F16 or BF16 on a GPU) the chain adds as its input lowAt: on NVIDIA's matrix cores the product's epilogue computes it from the stage memory the main loop is done with, so D Uᵀ is never stored (a LoRA's up projection); elsewhere lowApart.

• a: _gizlang_lib_ai__tensor__storage.Tensor
• ta: bool
• b: _gizlang_lib_ai__tensor__storage.Tensor
• tb: bool
• m: i64
• k: i64
• n: i64
• compute: i64
• low: Vec<_gizlang_lib_ai__tensor__storage.Tensor>
• lowAt: i64
• patch: Vec<i64>

Functions

withProduct

function withProduct(x: Lazy, c: Tensor): Lazy

The chain with the product's tensor c as input 0.

Effects: fail, read, write, console, io.

lazy

function lazy(t: Tensor): Lazy

T as a lazy expression with no steps yet.

Effects: fail.

lazyMatmul

function lazyMatmul(x: Tensor, transposeA: bool, y: Tensor, transposeB: bool): Lazy

x times y as matmulT, not yet computed: the steps on it run as the product is stored.

Effects: fail, read, write, console, io.

lazyMatmulIn

function lazyMatmulIn(x: Tensor, transposeA: bool, y: Tensor, transposeB: bool, compute: i64): Lazy

As lazyMatmul, with operands on a GPU rounded to compute (F16, BF16 or FP8) as Tensor.cast rounds them; in host memory they stay f32. The matrix cores' kernel rounds an F32 operand of an F16 or BF16 product as it copies it, so the cast takes no pass of its own (the bits are those of the cast and then the product). FP8 operands are scaled to its range first (fp8Scaled), and the product's store multiplies by both scales.

Effects: fail, read, write, console, io.

lazyProduct

function lazyProduct(shape: Vec<i64>, product: Product): Lazy

The expression of PRODUCT alone, of shape: its input 0 stands for the product and has no values.

Effects: fail.

lowApart

function lowApart(x: Lazy): Lazy

x with its Product.low computed apart (F32), as an input like any other.

Effects: fail, read, write, console, io.

run

function run(op: i64, a: i64, b: i64, out: i64, lanes: i64, regs: mut Vec<f32>): void

One step over a block of lanes elements (a multiple of 16): the operation chosen once per block, then one vector loop for it.

Effects: fail.

fp8Scaled

function fp8Scaled(t: Tensor): Vec<Tensor>

t scaled to the range of FP8 as tensorwise scaling (Transformer Engine) does it: [t / s as FP8, s], with s = max |t| / 448 a tensor of shape [] where t is.

Effects: fail, read, write, console, io.

fp4Blocks

function fp4Blocks(t: Tensor, across: bool): Vec<Tensor>

t (on a GPU) in NVFP4 blocks with its rows along k, as NVIDIA's Transformer Engine makes them: [q, s]. q is of type FP4 and shape [rows, k]: rows x k / 2 bytes of E2M1 pairs (row-major, the first in the low half), then k / 16 E4M3 block scales for each of lines (rows rounded up to 4; zeros past rows), in words of 4 (64 of k) by columns: word w of row r at w lines + r, so the words of a warp's mma lie together and copy in 16 bytes. s (F32, shape []) is max |t| / (6 x 448). The value at row r, p is its E2M1 times the scale of its block times s. across: t is [k, rows], else its last axis is k.

Effects: fail, read, write, console, io.

fp4Lines

function fp4Lines(rows: i64): i64

The rows of FP4 scales for rows rows (fp4Blocks).

Effects: fail.

fresh

function fresh(): Option<Buffer>

No output memory given: an evaluation takes new memory.

evalAs

function evalAs(self: Lazy, dtype: i64, into: Option<Buffer>): Tensor

The expression's tensor as dtype, or (-1) as eval gives it.

Effects: fail, read, write, console, io.

Methods

Lazy.pending

function Lazy.pending(self: Lazy): bool

Whether input 0 is a product not yet evaluated (_product): a step over it evaluates that product first.

Lazy.expand

function Lazy.expand(self: Lazy, shape: Vec<i64>): Lazy

The expression broadcast to SHAPE.

Effects: fail, read, write, console, io.

Lazy.add

function Lazy.add(self: Lazy, y: Lazy): Lazy

SELF + Y, element by element, with broadcasting.

Effects: fail, read, write, console, io.

Lazy.sub

function Lazy.sub(self: Lazy, y: Lazy): Lazy

SELF - Y, element by element, with broadcasting.

Effects: fail, read, write, console, io.

Lazy.mul

function Lazy.mul(self: Lazy, y: Lazy): Lazy

SELF * Y, element by element, with broadcasting.

Effects: fail, read, write, console, io.

Lazy.div

function Lazy.div(self: Lazy, y: Lazy): Lazy

SELF / Y, element by element, with broadcasting.

Effects: fail, read, write, console, io.

Lazy.maximum

function Lazy.maximum(self: Lazy, y: Lazy): Lazy

The larger of SELF and Y, element by element, with broadcasting.

Effects: fail, read, write, console, io.

Lazy.minimum

function Lazy.minimum(self: Lazy, y: Lazy): Lazy

The smaller of SELF and Y, element by element, with broadcasting.

Effects: fail, read, write, console, io.

Lazy.greater

function Lazy.greater(self: Lazy, y: Lazy): Lazy

1 where self > y, else 0.

Effects: fail, read, write, console, io.

Lazy.neg

function Lazy.neg(self: Lazy): Lazy

-SELF, element by element.

Effects: fail.

Lazy.abs

function Lazy.abs(self: Lazy): Lazy

The absolute value of SELF, element by element.

Effects: fail.

Lazy.sqrt

function Lazy.sqrt(self: Lazy): Lazy

The square root of SELF, element by element.

Effects: fail.

Lazy.exp

function Lazy.exp(self: Lazy): Lazy

e to the power SELF, element by element.

Effects: fail.

Lazy.log

function Lazy.log(self: Lazy): Lazy

The natural logarithm of SELF, element by element.

Effects: fail.

Lazy.tanh

function Lazy.tanh(self: Lazy): Lazy

The hyperbolic tangent of SELF, element by element.

Effects: fail.

Lazy.sigmoid

function Lazy.sigmoid(self: Lazy): Lazy

The sigmoid 1 / (1 + e^-x) of SELF, element by element.

Effects: fail.

Lazy.relu

function Lazy.relu(self: Lazy): Lazy

max(SELF, 0), element by element.

Effects: fail.

Lazy.square

function Lazy.square(self: Lazy): Lazy

SELF squared, element by element.

Effects: fail.

Lazy.roundedBf16

function Lazy.roundedBf16(self: Lazy): Lazy

Each value rounded to BF16 (nearest even) and kept as F32: where a pipeline stores BF16 between two steps that a product's epilogue now runs together.

Effects: fail.

Lazy.gelu

function Lazy.gelu(self: Lazy): Lazy

GELU, tanh form (as GPT-2 and torch's approximate="tanh").

Effects: fail.

Lazy.geluGrad

function Lazy.geluGrad(self: Lazy, dy: Lazy): Lazy

The gradient through gelu: self is x, dy the output's gradient.

Effects: fail, read, write, console, io.

Lazy.geluErf

function Lazy.geluErf(self: Lazy): Lazy

GELU, erf form (torch's default and BERT's "gelu").

Effects: fail.

Lazy.geluErfGrad

function Lazy.geluErfGrad(self: Lazy, dy: Lazy): Lazy

The gradient through geluErf: self is x, dy the output's gradient.

Effects: fail, read, write, console, io.

Lazy.on

function Lazy.on(self: Lazy, gpu: Gpu): Lazy

The expression with its inputs moved to gpu, so eval runs it there.

Effects: fail, read, write, console, io.

Lazy.like

function Lazy.like(self: Lazy): Tensor

A tensor of the expression's shape where it is evaluated (a GPU or host memory), with no values: for its shape and place (Tensor.filled, Tensor.gpu), never to read.

Effects: fail.

Lazy.eval

function Lazy.eval(self: Lazy): Tensor

The expression's tensor: one pass, on every core when it is large, or one kernel on the GPU of an input. It is f32 but over a product, where it is of the product's type (productType) when every other input is of that type or has one element, as torch's type promotion; else f32.

Effects: fail, read, write, console, io.

Lazy.evalIn

function Lazy.evalIn(self: Lazy, dtype: i64): Tensor

The expression's tensor stored as dtype (F32, F16, BF16 or FP8) on a GPU, so a cast takes no pass of its own; in host memory it is f32.

Effects: fail, read, write, console, io.

Lazy.evalInto

function Lazy.evalInto(self: Lazy, target: Tensor): void

Stores the expression's values into GPU tensor target (of its shape) as target's type: an update in place, so the kernels of a recording (Gpu.record) that read target read them. Its kernel writes there where nothing it reads can see the write (it reads target only at the element it writes); else it computes a new tensor and copies it.

Effects: fail, read, write, console, io.

Lazy.gluIn

function Lazy.gluIn(self: Lazy, dtype: i64): Tensor

A product chain's rows [m, n] as SwiGLU pairs, dtype elements [m, n / 2]: column c is silu(v[2c]) times v[2c + 1], each after the chain (a gate and an up projection with their rows interleaved). On a GPU the product's store makes them, so C's n columns are never stored.

Effects: fail, read, write, console, io.

Lazy.gluInto

function Lazy.gluInto(self: Lazy, target: Tensor): void

Lazy.gluIn into GPU tensor target ([m, n / 2]).

Effects: fail, read, write, console, io.

Lazy.reads

function Lazy.reads(self: Lazy, t: Tensor): bool

Whether the expression reads GPU tensor t's memory (an input or a product's operand).

Effects: fail.

Lazy.total

function Lazy.total(self: Lazy, policy: Reduction): Tensor

The sum of the expression's elements (Tensor.total) without its tensor: on a GPU, a strict sum of a small one is one kernel that computes each element as it adds it.

Effects: fail, read, write, console, io.

Lazy.average

function Lazy.average(self: Lazy, policy: Reduction): Tensor

The mean of the expression's elements (Tensor.average), as Lazy.total.

Effects: fail, read, write, console, io.

Constants

OP_INPUT

const OP_INPUT = 0

The step of a lazy expression that reads one of its inputs.

OP_GELU

const OP_GELU = 1

The step: tanh-approximated GELU of a.

OP_ADD

const OP_ADD = 2

The step: a + b.

OP_SUB

const OP_SUB = 3

The step: a - b.

OP_MUL

const OP_MUL = 4

The step: a * b.

OP_DIV

const OP_DIV = 5

The step: a / b.

OP_MAX

const OP_MAX = 6

The step: the larger of a and b.

OP_MIN

const OP_MIN = 7

The step: the smaller of a and b.

OP_NEG

const OP_NEG = 8

The step: -a.

OP_ABS

const OP_ABS = 9

The step: |a|.

OP_SQRT

const OP_SQRT = 10

The step: the square root of a.

OP_EXP

const OP_EXP = 11

The step: e to the a.

OP_LOG

const OP_LOG = 12

The step: the natural logarithm of a.

OP_TANH

const OP_TANH = 13

The step: tanh a.

OP_SIGMOID

const OP_SIGMOID = 14

The step: the logistic sigmoid of a.

OP_RELU

const OP_RELU = 15

The step: max(a, 0).

OP_SQUARE

const OP_SQUARE = 16

The step: a * a.

OP_GREATER

const OP_GREATER = 17

The step: 1 when a > b, else 0.

OP_GELU_GRAD

const OP_GELU_GRAD = 18

The step: the derivative of the tanh GELU at a.

OP_GELU_ERF

const OP_GELU_ERF = 19

The step: the exact (erf) GELU of a.

OP_GELU_ERF_GRAD

const OP_GELU_ERF_GRAD = 20

The step: the derivative of the exact GELU at a.

OP_BF16

const OP_BF16 = 21

The step: a rounded to bfloat16.