Lean-source custom operations

Program.of elaborates a Lean function before recognizing its supported computations. The result retains that function and a kernel-checked equality with the expression evaluator. Unsupported functions are rejected, rather than executed on the host or silently treated as GPU primitives. Captured scalar, unsigned-index and Boolean values become constants shared by all output lanes; calculations involving the reader or any lane-local variable remain in the generated expression. Fixed-parameter scalar recurrences use the equation compiler's theorems to produce checked loops.

def NN.Kernel.programStx : Lean.ParserDescr

Compile an elementwise scalar function, or an indexed function with checked input reads.

Program.of (fun (x : Float32) => x * x + 1) reads each element of input zero. Indexed functions take a Reader and UInt64 index and return Except Error explicitly. Both forms retain a kernel-checked equality with their reference calculation; unsupported operations are rejected. Captured values can parameterize scalar arithmetic, branches and bounded-fold lengths. They are evaluated on the host and emitted as constants when GPU source is rendered. Named recursive functions may have a final natural-number argument, with zero and successor equations and recursive calls on the predecessor at fixed parameters. The count must be a literal below 2^64 or a UInt64.toNat expression. A proved recurrence bridge preserves read failures and the arithmetic order while reusing the loop backend.

def NN.Kernel.Frontend.equationBody (source : Lean.Expr) (equations : Array Lean.Name) : Lean.Elab.TermElabM (Lean.Expr × Lean.Name)

Use the equation compiler's theorem, rather than unfolding its recursive implementation.

def NN.Kernel.Frontend.elaborate (source : Lean.Syntax) (expectedType? : Option Lean.Expr) : Lean.Elab.TermElabM Lean.Expr

Recognize a source function and construct its checked expression representation.

def NN.Kernel.Frontend.elabProgram : Lean.Elab.Term.TermElab

Elaborate the explicit indexed-program constructor.