Lean FFI bindings for the IREE runtime.

Links to `libiree_ffi.so` (thin wrapper) + IREE runtime via the Lean shim
in `ffi/iree_lean_ffi.c`. Exposes:
  - `LowererSession.create` — load a .vmfb, bind to CUDA device
  - `LowererSession.mlpForward` — MLP-specific forward pass (MNIST shape) 

def LowererSession : Type

Equations

• LowererSession = LowererSessionPointed✝.type

instance instNonemptyLowererSession : Nonempty LowererSession

@[extern lean_iree_session_create]

opaque LowererSession.create (path : String) : IO LowererSession

Load a .vmfb bytecode module onto the default CUDA device.

On the XLA backend (libpjrt_ffi.so) the argument is instead the .mlir source — XLA compiles the StableHLO in-process, so there is no separate iree-compile step. Use VerifiedNet.mkSession rather than calling this directly; it picks the right path per backendName.

@[extern lean_iree_backend_name]

opaque LowererSession.backendName : IO String

"iree" or "xla" — which shim this binary was linked against. Detected by probing for a symbol only libpjrt_ffi.so defines, so it cannot disagree with the linked library. See planning/archive/xla_pjrt_ladder.md.

@[extern lean_iree_mlp_forward]

opaque LowererSession.mlpForward (sess : LowererSession) (x W0 b0 W1 b1 W2 b2 : FloatArray) (batch : USize) : IO FloatArray

Run MNIST-MLP forward pass. Shapes are fixed: x is batch×784, W0 is 784×512, b0 is 512, W1 is 512×512, b1 is 512, W2 is 512×10, b2 is 10. Returns the logits as a batch×10 flattened FloatArray.

@[extern lean_iree_mlp_train_step]

opaque LowererSession.mlpTrainStep (sess : LowererSession) (params x : FloatArray) (y : ByteArray) (lr : Float) (batch : USize) : IO FloatArray

Run one SGD training step. Params packed into a single FloatArray of length 669706 in order W0|b0|W1|b1|W2|b2. Labels are a ByteArray of 4*batch bytes (int32 LE). Returns new params + loss as a single FloatArray of length 669707; result[669706] is the loss.

@[extern lean_iree_train_step_packed]

opaque LowererSession.trainStepPacked (sess : LowererSession) (fnName : String) (params : FloatArray) (shapes : ByteArray) (x : FloatArray) (xShape y : ByteArray) (lr : Float) (batch : USize) : IO FloatArray

Generic train step. Shapes are packed ByteArrays (see packShapes).

@[extern lean_iree_train_step_f32]

opaque LowererSession.trainStepF32 (sess : LowererSession) (fnName : String) (params shapes x xShape y : ByteArray) (lr : Float) (batch : USize) : IO ByteArray

Zero-copy f32 train step. All tensors are ByteArray (raw float32 bytes). No Float64↔Float32 conversion at the boundary.

@[extern lean_iree_train_step_adam_f32]

opaque LowererSession.trainStepAdamF32 (sess : LowererSession) (fnName : String) (params shapes x xShape y : ByteArray) (lr t : Float) (bnShapes : ByteArray) (batch : USize) : IO ByteArray

Adam train step (f32). Passes step counter t for bias correction. Params = weights ++ m ++ v. Returns params ++ loss ++ BN stats. bnShapes: packed [n_bn_layers, oc0, oc1, ...] for BN stat output sizes.

@[extern lean_iree_train_step_adam_f32_softlabel]

opaque LowererSession.trainStepAdamF32Soft (sess : LowererSession) (fnName : String) (params shapes x xShape ySoft : ByteArray) (lr t : Float) (bnShapes : ByteArray) (batch nClasses : USize) : IO ByteArray

Soft-label variant: ySoft is a [batch, nClasses] f32 tensor (smoothed + mixed). Routes to the codegen produced with useSoftLabels := true. Used by the mixup/cutmix path.

@[extern lean_iree_train_step_adam_f32_seg]

opaque LowererSession.trainStepAdamF32Seg (sess : LowererSession) (fnName : String) (params shapes x xShape ySeg : ByteArray) (lr t : Float) (bnShapes : ByteArray) (batch H W : USize) : IO ByteArray

Per-pixel segmentation variant: ySeg is an int32 [batch, H, W] per-pixel label tensor. Routes to the codegen produced with useSeg := true.

@[extern lean_iree_train_step_adam_f32_ddpm]

opaque LowererSession.trainStepAdamF32Ddpm (sess : LowererSession) (fnName : String) (params shapes x xShape yDdpm : ByteArray) (lr t : Float) (bnShapes : ByteArray) (batch outC outH outW : USize) : IO ByteArray

DDPM variant: yDdpm is a [batch, C, H, W] f32 tensor — the target ε noise the model learns to predict. Routes to the codegen produced with useDdpm := true. Loss is per-pixel MSE.

@[extern lean_iree_train_step_adam_f32_yolov1]

opaque LowererSession.trainStepAdamF32Yolov1 (sess : LowererSession) (fnName : String) (params shapes x xShape yYolo mYolo : ByteArray) (lr t : Float) (bnShapes : ByteArray) (batch gridH gridW perCell : USize) : IO ByteArray

YOLOv1 variant. yYolo is a [batch, perCell, gridH, gridW] f32 target tensor (NCHW); mYolo is a [batch, gridH, gridW] f32 per-cell objectness mask (1.0 where a GT box's center falls in the cell, 0.0 otherwise). Routes to the codegen produced with useYolov1 := true. Loss is the 5-term masked MSE described in planning/archive/yolo_demo_v2.md Phase 1.

perCell = numBoxes * 5 + numClasses. For VOC this is 2*5 + 20 = 30; gridH = gridW = 7.

@[extern lean_iree_forward_f32]

opaque LowererSession.forwardF32 (sess : LowererSession) (fnName : String) (params shapes x xShape : ByteArray) (batch nClasses : USize) (nResident gen : USize := 0) : IO ByteArray

Zero-copy f32 forward pass. Pushes x then param tensors, returns logits. For inference/eval — no y, lr, or velocity inputs.

nResident / gen — device residency in HOLD mode (§2d.3), and it is a different mechanism from the train step's. This graph returns logits, not parameters, so there is nothing to retain from the output; instead the whole parameter set is seeded once and reused across every eval batch, rather than pushed 79-123 times per epoch. Measured on the MNIST MLP, 73% of an eval step was the parameter push (0.6 ms of 0.8 — compute is 0.1).

⚠ gen is what makes holding safe, and it must change whenever params does. Pass the epoch number. A held set that went stale would score the previous epoch's weights silently, which reads as a training plateau rather than as an error — a nastier failure than anything the update mode has. The shim re-seeds the moment the token differs.

Defaults (0, 0) = the copying path, so every inference demo that calls this is unaffected.

@[extern lean_iree_linear_train_step]

opaque LowererSession.linearTrainStepV (sess : LowererSession) (fnName : String) (x W0 b0 y : ByteArray) (batch d0 d1 : USize) (nResident : USize := 0) : IO ByteArray

Drive the verified-renderer @linear_train_step (StableHLO.linearTrainStepModuleV) through the generic IREE invoke. Inputs are raw f32 ByteArrays: x is batch×d₀, W0 is d₀×d₁, b0 is d₁; y is int32 [batch] (the one-hot is built in the C shim). Returns W0n (d₀·d₁ f32) ++ b0n (d₁ f32).

nResident: see mlpTrainStepV. Here it is 2 — W0 and b0, i.e. the whole parameter set, since this graph returns exactly its two param inputs.

@[extern lean_iree_mlp_train_step_v_dp2]

opaque LowererSession.mlpTrainStepVDP (sess : LowererSession) (fnName : String) (x params shapes y : ByteArray) (batch d0 d3 replicas : USize) (nResident nShardTail : USize := 0) : IO ByteArray

Data-parallel @<slug>_train_step: same packed-params protocol as mlpTrainStepV, but batch is the GLOBAL batch and the XLA shim splits x and the labels across replicas devices while replicating the parameters. The emitted graph all-reduces every gradient before the optimizer consumes it (ViTRender.emitAdamVDP), so all replicas produce identical parameters and the result is read back from replica 0.

Only the XLA shim exports the underlying entry point; on the IREE build this raises rather than silently running single-device.

nResident: see mlpTrainStepV. Each replica keeps its own retained set on its own device, which is where the bigger half of the win is — today the full [θ|m|v] is pushed to every replica every step, an O(N−1) cost against O(1) compute (§2d.3a: 4 GPUs currently buy 1.46×).

nShardTail: how many TRAILING entries of the param list are PER-EXAMPLE and must be sharded like x rather than replicated like the parameters. Today that is exactly the stochastic-depth drop masks. ⚠ It is a COUNT supplied by the driver rather than something the shim infers: an index would be per-net and a shape test ("outer dim == batch") would sweep up any parameter that happens to be batch-sized. Default 0, so every existing call site is unchanged.

⚠ The extern is _dp2, not _dp, because this ADDED AN ARGUMENT — §4's rule for pjrt_ffi_invoke_f32_resident_v2: a stale .so against a new binary shifts every argument, which is garbage rather than a link error. A rename makes it a link error.

@[extern lean_iree_mlp_train_step_v]

opaque LowererSession.mlpTrainStepV (sess : LowererSession) (fnName : String) (x params shapes y : ByteArray) (batch d0 d3 : USize) (nResident : USize := 0) : IO ByteArray

Drive the verified-renderer @mlp_train_step (StableHLO.mlpTrainStepText) through the generic IREE invoke. params is the packed f32 weights (sliced per shapes, same layout as forwardF32); x is batch×d₀; y is int32 [batch] (one-hot built in the C shim with d₃ classes). Returns the updated params, packed in the same layout.

nResident — how many LEADING param tensors may stay on the device between steps (handoff §2d.3). The driver is the only place that knows the packed layout is [θ|m|v | lr,bc₁,bc₂ | bn stats], and hence that the first 3×P tensors are exactly the ones the host writes once and thereafter only feeds straight back; so it states the count and the shim checks that input i+1 and output i really are the same tensor before retaining anything.

It is a request, not a mode. The transport is chosen in C — residency engages only under $PJRT_FFI_RESIDENT=1 on the XLA build, and is inert everywhere else — precisely so that this driver keeps no backend branch to drift (§2d.3, "the design decision that protects every existing gate"). Default 0 = the copying path, which is what every tie and DP-check harness wants: those read the whole returned blob, and a retained prefix would leave it unwritten.

@[extern lean_iree_read_params]

opaque LowererSession.readParams (sess : LowererSession) (packed : ByteArray) (nBytes : USize) : IO ByteArray

Read the authoritative leading nBytes of the packed parameter blob.

Without residency this is packed.extract 0 nBytes and nothing more — which is what it must be on IREE, where the weak read-back symbol does not exist. With residency live the [θ|m|v] prefix of packed is unwritten (it never came back from the device), and this performs the one d2h that still happens.

Either way it is a per-epoch call: the eval pass and the checkpoint are the only things that want the whole blob, and that call site was already once-per-epoch before any of this.

def MlpLayout.nW0 : Nat

Equations

• MlpLayout.nW0 = 784 * 512

def MlpLayout.nb0 : Nat

Equations

• MlpLayout.nb0 = 512

def MlpLayout.nW1 : Nat

Equations

• MlpLayout.nW1 = 512 * 512

def MlpLayout.nb1 : Nat

Equations

• MlpLayout.nb1 = 512

def MlpLayout.nW2 : Nat

Equations

• MlpLayout.nW2 = 512 * 10

def MlpLayout.nb2 : Nat

Equations

• MlpLayout.nb2 = 10

def MlpLayout.nParams : Nat

Equations

• MlpLayout.nParams = MlpLayout.nW0 + MlpLayout.nb0 + MlpLayout.nW1 + MlpLayout.nb1 + MlpLayout.nW2 + MlpLayout.nb2

def MlpLayout.lossIdx : Nat

Equations

• MlpLayout.lossIdx = MlpLayout.nParams

def packShapes (shapes : Array (Array Nat)) : ByteArray

Pack param shape descriptors: [nParams, rank0, d0..., rank1, d1..., ...] as int32 LE.

Equations

• One or more equations did not get rendered due to their size.

def packXShape (dims : Array Nat) : ByteArray

Pack a single shape: [rank, d0, d1, ...] as int32 LE (for x input).

Equations

• One or more equations did not get rendered due to their size.

def CnnLayout.paramShapes : Array (Array Nat)

Equations

• CnnLayout.paramShapes = #[#[32, 1, 3, 3], #[32], #[32, 32, 3, 3], #[32], #[6272, 512], #[512], #[512, 512], #[512], #[512, 10], #[10]]

def CnnLayout.nParams : Nat

Equations

• CnnLayout.nParams = 32 * 1 * 3 * 3 + 32 + 32 * 32 * 3 * 3 + 32 + 6272 * 512 + 512 + 512 * 512 + 512 + 512 * 10 + 10

def CnnLayout.lossIdx : Nat

Equations

• CnnLayout.lossIdx = CnnLayout.nParams

def CnnLayout.shapesBA : ByteArray

Equations

• CnnLayout.shapesBA = packShapes CnnLayout.paramShapes

def CnnLayout.xShape (batch : Nat) : ByteArray

Equations

• CnnLayout.xShape batch = packXShape #[batch, 784]

def CifarLayout.paramShapes : Array (Array Nat)

Equations

• CifarLayout.paramShapes = #[#[32, 3, 3, 3], #[32], #[32, 32, 3, 3], #[32], #[64, 32, 3, 3], #[64], #[64, 64, 3, 3], #[64], #[4096, 512], #[512], #[512, 512], #[512], #[512, 10], #[10]]

def CifarLayout.nParams : Nat

Equations

• CifarLayout.nParams = 32 * 3 * 3 * 3 + 32 + 32 * 32 * 3 * 3 + 32 + 64 * 32 * 3 * 3 + 64 + 64 * 64 * 3 * 3 + 64 + 4096 * 512 + 512 + 512 * 512 + 512 + 512 * 10 + 10

def CifarLayout.lossIdx : Nat

Equations

• CifarLayout.lossIdx = CifarLayout.nParams

def CifarLayout.shapesBA : ByteArray

Equations

• CifarLayout.shapesBA = packShapes CifarLayout.paramShapes

def CifarLayout.xShape (batch : Nat) : ByteArray

Equations

• CifarLayout.xShape batch = packXShape #[batch, 3072]

def CifarBnLayout.paramShapes : Array (Array Nat)

BN-CIFAR params: each conv layer carries per-channel γ/β [c] after its bias, interleaved as W|b|γ|β. 22 params (4×{W,b,γ,β} + 3×{W,b}). Order MUST match @cifar_bn_train_step's signature.

Equations

• One or more equations did not get rendered due to their size.

def CifarBnLayout.nParams : Nat

Equations

• One or more equations did not get rendered due to their size.

def CifarBnLayout.lossIdx : Nat

Equations

• CifarBnLayout.lossIdx = CifarBnLayout.nParams

def CifarBnLayout.shapesBA : ByteArray

Equations

• CifarBnLayout.shapesBA = packShapes CifarBnLayout.paramShapes

def CifarBnLayout.xShape (batch : Nat) : ByteArray

Equations

• CifarBnLayout.xShape batch = packXShape #[batch, 3072]

def ResNet34Layout.specs : Array (Array Nat × Nat)

(dims, initKind) for every param, in func-arg order.

Equations

• One or more equations did not get rendered due to their size.

def ResNet34Layout.paramShapes : Array (Array Nat)

Equations

• ResNet34Layout.paramShapes = Array.map (fun (x : Array Nat × Nat) => x.fst) ResNet34Layout.specs

def ResNet34Layout.nParams : Nat

Equations

• ResNet34Layout.nParams = Array.foldl (fun (x1 x2 : Nat) => x1 + x2) 0 (Array.map (fun (s : Array Nat × Nat) => Array.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 s.fst) ResNet34Layout.specs)

def ResNet34Layout.shapesBA : ByteArray

Equations

• ResNet34Layout.shapesBA = packShapes ResNet34Layout.paramShapes

def ResNet34Layout.xShape (batch : Nat) : ByteArray

Equations

• ResNet34Layout.xShape batch = packXShape #[batch, 3 * 224 * 224]

def MobileNetV2Layout.specs : Array (Array Nat × Nat)

(dims, initKind) for every param, in func-arg order.

Equations

• One or more equations did not get rendered due to their size.

def MobileNetV2Layout.paramShapes : Array (Array Nat)

Equations

• MobileNetV2Layout.paramShapes = Array.map (fun (x : Array Nat × Nat) => x.fst) MobileNetV2Layout.specs

def MobileNetV2Layout.nParams : Nat

Equations

• One or more equations did not get rendered due to their size.

def MobileNetV2Layout.shapesBA : ByteArray

Equations

• MobileNetV2Layout.shapesBA = packShapes MobileNetV2Layout.paramShapes

def MobileNetV2Layout.xShape (batch : Nat) : ByteArray

Equations

• MobileNetV2Layout.xShape batch = packXShape #[batch, 3 * 224 * 224]

def EfficientNetLayout.specs : Array (Array Nat × Nat)

(dims, initKind) for every param, in func-arg order — generated from the B0 stage spec exactly as tests/TestEfficientNet*.lean blocks (stem out 32, prev threading).

Equations

• One or more equations did not get rendered due to their size.

def EfficientNetLayout.paramShapes : Array (Array Nat)

Equations

• EfficientNetLayout.paramShapes = Array.map (fun (x : Array Nat × Nat) => x.fst) EfficientNetLayout.specs

def EfficientNetLayout.nParams : Nat

Equations

• One or more equations did not get rendered due to their size.

def EfficientNetLayout.shapesBA : ByteArray

Equations

• EfficientNetLayout.shapesBA = packShapes EfficientNetLayout.paramShapes

def EfficientNetLayout.xShape (batch : Nat) : ByteArray

Equations

• EfficientNetLayout.xShape batch = packXShape #[batch, 3 * 224 * 224]

def ConvNeXtLayout.specs : Array (Array Nat × Nat)

(dims, initKind) for every param, in func-arg order.

⚠ §2m moved three things at once: every LN affine went rank-0 #[] → per-channel #[c], the stem LN appeared, and the head LN went away. The first two were right; the third was not, and the note that used to sit here — "the last two nearly cancel … so a matching parameter count is a decomposition test, not an architecture check" — was the correct warning drawn at the wrong conclusion. The residue is not noise, it IS the missing layer: 28,587,592 against timm.create_model('convnext_tiny')'s 28,589,128 is short by exactly 2×768 = 1,536.

⭐ The head LN is back (2026-08-30, §7.1), so the head is GAP → LN(768) → dense as in both the paper (self.norm(x.mean([-2,-1])), nn.LayerNorm(dims[-1], eps=1e-6)) and timm (NormMlpClassifierHead). 182 param tensors; the floats are 27,827,818 at K = 10, i.e. 28,589,128 at K = 1000 — timm's count exactly.

Equations

• One or more equations did not get rendered due to their size.

def ConvNeXtLayout.paramShapes : Array (Array Nat)

Equations

• ConvNeXtLayout.paramShapes = Array.map (fun (x : Array Nat × Nat) => x.fst) ConvNeXtLayout.specs

def ConvNeXtLayout.nParams : Nat

Equations

• ConvNeXtLayout.nParams = Array.foldl (fun (x1 x2 : Nat) => x1 + x2) 0 (Array.map (fun (s : Array Nat × Nat) => Array.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 s.fst) ConvNeXtLayout.specs)

def ConvNeXtLayout.shapesBA : ByteArray

Equations

• ConvNeXtLayout.shapesBA = packShapes ConvNeXtLayout.paramShapes

def ConvNeXtLayout.xShape (batch : Nat) : ByteArray

Equations

• ConvNeXtLayout.xShape batch = packXShape #[batch, 3 * 224 * 224]

def ViTLayout.specs : Array (Array Nat × Nat)

(dims, initKind) for every param, in @vit_train_step func-arg order.

Equations

• One or more equations did not get rendered due to their size.

def ViTLayout.paramShapes : Array (Array Nat)

Equations

• ViTLayout.paramShapes = Array.map (fun (x : Array Nat × Nat) => x.fst) ViTLayout.specs

def ViTLayout.nParams : Nat

Equations

• ViTLayout.nParams = Array.foldl (fun (x1 x2 : Nat) => x1 + x2) 0 (Array.map (fun (s : Array Nat × Nat) => Array.foldl (fun (x1 x2 : Nat) => x1 * x2) 1 s.fst) ViTLayout.specs)

def ViTLayout.shapesBA : ByteArray

Equations

• ViTLayout.shapesBA = packShapes ViTLayout.paramShapes

def ViTLayout.xShape (batch : Nat) : ByteArray

Equations

• ViTLayout.xShape batch = packXShape #[batch, 3 * 224 * 224]

def MlpLayout.paramShapes : Array (Array Nat)

Equations

• MlpLayout.paramShapes = #[#[784, 512], #[512], #[512, 512], #[512], #[512, 10], #[10]]

def MlpLayout.shapesBA : ByteArray

Equations

• MlpLayout.shapesBA = packShapes MlpLayout.paramShapes