Computations inside custom tensor operations

Tensor graphs describe dependencies between operations. A kernel expression describes the calculation of one output element inside an operation. Variables are typed, inputs are read-only, branches are lazy, and folds accumulate in increasing index order. Arithmetic is supplied explicitly; no field laws or floating-point reassociation are assumed.

Indices use unsigned 64-bit arithmetic rather than silently compiling unbounded natural numbers to machine integers. The frontend must reject constants outside that range. Tensor reads return errors instead of inventing a value for an invalid address.

inductive NN.Kernel.Ty : Type

The three runtime value types supported inside a kernel.

• scalar : Ty
• index : Ty
• predicate : Ty

@[instance_reducible]

instance NN.Kernel.instDecidableEqTy : DecidableEq Ty

@[instance_reducible]

instance NN.Kernel.instReprTy : Repr Ty

def NN.Kernel.instReprTy.repr : Ty → Nat → Std.Format

@[reducible, inline]

abbrev NN.Kernel.Ty.Value (α : Type) : Ty → Type

Scalar arithmetic remains a parameter of the semantics.

inductive NN.Kernel.Var : List Ty → Ty → Type

A variable can only refer to a value of its declared type.

• zero {t : Ty} {Γ : List Ty} : Var (t :: Γ) t
• succ {Γ : List Ty} {t s : Ty} : Var Γ t → Var (s :: Γ) t

structure NN.Kernel.Env (α : Type) (Γ : List Ty) : Type

Values of the variables currently in scope.

• get {t : Ty} : Var Γ t → Ty.Value α t

def NN.Kernel.Env.push {α : Type} {Γ : List Ty} {t : Ty} (value : Ty.Value α t) (env : Env α Γ) : Env α (t :: Γ)

Introduce a local binding without changing the values of older variables.

inductive NN.Kernel.ScalarOp : Type

Floating-point operations are kept separate: multiplication followed by addition is not FMA.

• add : ScalarOp
• sub : ScalarOp
• mul : ScalarOp
• div : ScalarOp

@[instance_reducible]

instance NN.Kernel.instReprScalarOp : Repr ScalarOp

def NN.Kernel.instReprScalarOp.repr : ScalarOp → Nat → Std.Format

@[instance_reducible]

instance NN.Kernel.instDecidableEqScalarOp : DecidableEq ScalarOp

inductive NN.Kernel.IndexOp : Type

Index operations follow Lean's UInt64 semantics, including wraparound.

• add : IndexOp
• sub : IndexOp
• mul : IndexOp
• div : IndexOp
• mod : IndexOp

def NN.Kernel.instReprIndexOp.repr : IndexOp → Nat → Std.Format

@[instance_reducible]

instance NN.Kernel.instReprIndexOp : Repr IndexOp

@[instance_reducible]

instance NN.Kernel.instDecidableEqIndexOp : DecidableEq IndexOp

inductive NN.Kernel.Compare : Type

Comparisons produce a predicate, not an integer mask.

• eq : Compare
• lt : Compare
• le : Compare

@[instance_reducible]

instance NN.Kernel.instReprCompare : Repr Compare

def NN.Kernel.instReprCompare.repr : Compare → Nat → Std.Format

@[instance_reducible]

instance NN.Kernel.instDecidableEqCompare : DecidableEq Compare

structure NN.Kernel.Arithmetic (α : Type) : Type

An arithmetic implementation fixes scalar operations and comparisons for the whole kernel.

• binary : ScalarOp → α → α → α
• neg : α → α
• compare : Compare → α → α → Bool

inductive NN.Kernel.Error : Type

Evaluation failures retain the offending input, address, or unsupported output size.

• input (operand : Nat) : Error
• bounds (operand : Nat) (index : UInt64) (size : Nat) : Error
• size (count : Nat) : Error

def NN.Kernel.instReprError.repr : Error → Nat → Std.Format

@[instance_reducible]

instance NN.Kernel.instReprError : Repr Error

def NN.Kernel.instDecidableEqError.decEq (x✝ x✝¹ : Error) : Decidable (x✝ = x✝¹)

@[instance_reducible]

instance NN.Kernel.instDecidableEqError : DecidableEq Error

@[reducible, inline]

abbrev NN.Kernel.Reader (α : Type) : Type

Input access is shared by reference execution and lowered execution.

inductive NN.Kernel.Expr (α : Type) : List Ty → Ty → Type

A typed expression with explicitly ordered arithmetic and bounded sequential folds.

The fold body binds the accumulator first and the iteration index second. A count of zero leaves the initial value unchanged. Only the selected conditional branch is evaluated.

• scalar {α : Type} {Γ : List Ty} (value : α) : Expr α Γ Ty.scalar
• index {α : Type} {Γ : List Ty} (value : UInt64) : Expr α Γ Ty.index
• predicate {α : Type} {Γ : List Ty} (value : Bool) : Expr α Γ Ty.predicate
• var {α : Type} {Γ : List Ty} {t : Ty} (v : Var Γ t) : Expr α Γ t
• load {α : Type} {Γ : List Ty} (operand : Nat) (index : Expr α Γ Ty.index) : Expr α Γ Ty.scalar
• binary {α : Type} {Γ : List Ty} (op : ScalarOp) (x y : Expr α Γ Ty.scalar) : Expr α Γ Ty.scalar
• neg {α : Type} {Γ : List Ty} (x : Expr α Γ Ty.scalar) : Expr α Γ Ty.scalar
• indexBinary {α : Type} {Γ : List Ty} (op : IndexOp) (x y : Expr α Γ Ty.index) : Expr α Γ Ty.index
• compare {α : Type} {Γ : List Ty} (op : Compare) (x y : Expr α Γ Ty.scalar) : Expr α Γ Ty.predicate
• indexCompare {α : Type} {Γ : List Ty} (op : Compare) (x y : Expr α Γ Ty.index) : Expr α Γ Ty.predicate
• cond {α : Type} {Γ : List Ty} {t : Ty} (p : Expr α Γ Ty.predicate) (yes no : Expr α Γ t) : Expr α Γ t
• letIn {α : Type} {Γ : List Ty} {s t : Ty} (value : Expr α Γ s) (body : Expr α (s :: Γ) t) : Expr α Γ t
• fold {α : Type} {Γ : List Ty} (count : Expr α Γ Ty.index) (initial : Expr α Γ Ty.scalar) (body : Expr α (Ty.scalar :: Ty.index :: Γ) Ty.scalar) : Expr α Γ Ty.scalar

def NN.Kernel.IndexOp.eval : IndexOp → UInt64 → UInt64 → UInt64

Unsigned division and remainder by zero retain Lean's specified behavior.

def NN.Kernel.Compare.index : Compare → UInt64 → UInt64 → Bool

Unsigned comparisons do not reinterpret indices as signed integers.

def NN.Kernel.iterate {α : Type} (step : UInt64 → α → Except Error α) : Nat → UInt64 → α → Except Error α

Execute precisely count iterations, stopping at the first failed input read.

theorem NN.Kernel.iterate_congr {α : Type} {f g : UInt64 → α → Except Error α} (h : ∀ (i : UInt64) (acc : α), f i acc = g i acc) (n : Nat) (i : UInt64) (acc : α) : iterate f n i acc = iterate g n i acc

Equality of step functions preserves the full fold, including its failure behavior.

theorem NN.Kernel.iterate_succ {α : Type} (step : UInt64 → α → Except Error α) (n : Nat) (i : UInt64) (acc : α) : iterate step (n + 1) i acc = do let value ← iterate step n i acc step (i + n.toUInt64) value

Splitting off the last step preserves the order of all reads and arithmetic. The index conversion is modular, just like the unsigned counter used by iterate.

theorem NN.Kernel.iterate_eq_recurrence {α : Type} (f : Nat → Except Error α) (step : UInt64 → α → Except Error α) (initial : Except Error α) (hzero : f 0 = initial) (hsucc : ∀ (n : Nat), f (n + 1) = do let previous ← f n step n.toUInt64 previous) (n : Nat) : (do let value ← initial iterate step n 0 value) = f n

A recursive calculation can use the loop backend when its equations give a scalar recurrence. This includes failing input reads: neither later steps nor the final result mask an earlier error. No laws about the scalar arithmetic, and no reassociation of operations, are assumed.

def NN.Kernel.Expr.eval {α : Type} (arithmetic : Arithmetic α) (read : Reader α) {Γ : List Ty} {t : Ty} (env : Env α Γ) : Expr α Γ t → Except Error (Ty.Value α t)

Reference semantics: operands are evaluated left to right and folds are never reassociated.

@[reducible, inline]

abbrev NN.Kernel.Renaming (Γ Δ : List Ty) : Type

A type-preserving change of variable addresses.

def NN.Kernel.Renaming.lift {Γ Δ : List Ty} {s : Ty} (ρ : Renaming Γ Δ) : Renaming (s :: Γ) (s :: Δ)

Extend a renaming under a binder, keeping the newly bound variable fixed.

def NN.Kernel.Expr.rename {α : Type} {Γ Δ : List Ty} {t : Ty} (ρ : Renaming Γ Δ) : Expr α Γ t → Expr α Δ t

Rename local variables without changing the expression's operation order.

def NN.Kernel.Env.pull {α : Type} {Γ Δ : List Ty} (ρ : Renaming Γ Δ) (env : Env α Δ) : Env α Γ

Pull an environment back along variable addresses.

theorem NN.Kernel.Env.pull_push {α : Type} {Γ Δ : List Ty} {s : Ty} (ρ : Renaming Γ Δ) (env : Env α Δ) (value : Ty.Value α s) : pull (fun {t : Ty} => Renaming.lift fun {t : Ty} => ρ) (push value env) = push value (pull (fun {t : Ty} => ρ) env)

theorem NN.Kernel.Expr.eval_rename {α : Type} {Γ : List Ty} {t : Ty} (expr : Expr α Γ t) (arithmetic : Arithmetic α) (read : Reader α) {Δ : List Ty} (ρ : Renaming Γ Δ) (env : Env α Δ) : eval arithmetic read env (rename (fun {t : Ty} => ρ) expr) = eval arithmetic read (Env.pull (fun {t : Ty} => ρ) env) expr

Moving local variables is valid for arbitrary scalar arithmetic and failing input reads.