MLIR (StableHLO) code generator: emit MLIR modules from NetSpec.

Supports MLPs (`.dense`) and CNNs (`.conv2d`, `.maxPool`, `.flatten`).
Walks layers tracking the current activation shape. For CNNs with a flat
input, emits a reshape to (batch, ic, imageH, imageW) at the head.

Layout: NCHW tensors, OIHW kernels (matches the JAX codegen convention).
Params are interleaved as (W0, b0, W1, b1, ...) in the function signature,
with `idx` advancing for every conv or dense layer. 

def MlirCodegen.inputFlatDim (spec : NetSpec) : Nat

Total flat input size: for CNN, ic * imageH * imageW; for MLP, first dense layer's fanIn.

Equations

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

def MlirCodegen.inputChannels (spec : NetSpec) : Option Nat

If the first layer is conv/convBn, returns the NCHW input channels.

Equations

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

def MlirCodegen.sanitize (s : String) : String

Sanitize a string for use as an MLIR identifier. Public so trainers can derive eval-call names from spec.name instead of hardcoding them (and getting them wrong — see e.g. the EffNetV2 eval-name bug).

Equations

• MlirCodegen.sanitize s = String.map (fun (c : Char) => if c.isAlphanum = true then c else '_') s.toLower

FPN neck (top-down multi-scale merge) — detection-infra brick #3

Composes 1×1 lateral convs + the already-verified `bilinearUpsample` + adds
into the RetinaNet top-down pyramid. No new primitives: the merge is a DAG of
existing ops. Forward/backward FD-verified in numpy first
(`scripts/fpn_neck_check.py`), then these emitters are checked against that
oracle by `scripts/fpn_neck_probe_check.py` (the `fpn-neck-probe` exe). 

def MlirCodegen.generate (spec : NetSpec) (batchSize : Nat) : String

Generate a StableHLO MLIR module for forward-only inference.

Equations

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

def MlirCodegen.collectBnLayers (spec : NetSpec) : Array (Nat × Nat)

Collect (pidx, oc) pairs for every convBn layer (including those inside residual/bottleneck blocks).

Equations

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

def MlirCodegen.generateEval (spec : NetSpec) (batchSize : Nat) : String

Generate a StableHLO MLIR module for eval (fixed BN using running mean/var).

Equations

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

def MlirCodegen.preGAPShape (spec : NetSpec) : Option (Nat × Nat × Nat)

Shape walker: returns (C, H, W) of the activation immediately before the LAST .globalAvgPool in the spec — i.e. the last conv feature map, the input to the GAP+dense head. Returns none for specs without a globalAvgPool (those are not CAM-eligible). Mirrors the shape transitions in emitForwardEvalSig; kept focused (no param bookkeeping) so it stays readable.

Equations

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

def MlirCodegen.generateForwardCam (spec : NetSpec) (batchSize : Nat) : String

Generate a StableHLO MLIR module for forward-with-capture (CAM mode). Module name <sanitized>_cam, function @forward_cam. Walks the network up through the last conv feature map, reshapes it to flat [batch, C*H*W], and returns. The Lean side recomputes logits and the CAM heatmap from the dense-head weights.

Equations

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

@[implicit_reducible]

instance MlirCodegen.instInhabitedFwdRec : Inhabited MlirCodegen.FwdRec✝

Equations

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

def MlirCodegen.generateTrainStep (spec : NetSpec) (batchSize : Nat) (moduleName : String := "jit_train_step") (labelSmoothing : Float := 0.1) (weightDecay : Float := 1e-4) (useAdam : Bool := true) (useSoftLabels useFocal : Bool := false) (focalGamma : Float := 2.0) (useSeg useDdpm : Bool := false) (ddpmOutShape : List Nat := []) (useYolov1 : Bool := false) (yoloGridH yoloGridW : Nat := 7) (yoloNumBoxes : Nat := 2) (yoloNumClasses : Nat := 20) (gradClipNorm : Float := 0.0) (headLrMult : Float := 1.0) (useMuon useShampoo : Bool := false) (segLoss : SegLoss := SegLoss.ce) (useDiouBox : Bool := false) (yoloAnchors : List (Float × Float) := []) (fpnScales : List (Nat × List (Float × Float)) := []) (yoloClsWeights : List Float := []) (yoloClsFocalGamma : Float := 0.0) : String

Generate a full train_step MLIR module with VJPs. When useSoftLabels is true, the function takes a %y_soft : [B, NC] f32 tensor (mixup/ cutmix-style) instead of an int32 %y : [B] label vector; label smoothing is then expected to be applied by the caller.

Equations

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

def MlirCodegen.segLossProbeModule (B NC H W : Nat) (segLoss : SegLoss) (labelSmoothing : Float := 0.0) : String

Standalone segmentation-loss module — exercises emitSegLossBlock in isolation so the loss and its gradient can be checked against finite differences: @main(logits : [B,NC,H,W], y : [B,H,W] i32) -> (loss, d_logits).

This is the FD harness for .dice / .diceCE, whose gradient (unlike CE's p - y) is a real softmax Jacobian-vector product and therefore worth checking numerically rather than by inspection. Driven by scripts/seg_loss_probe_check.py; CPU is fine, no GPU needed.

Equations

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

def MlirCodegen.diouProbeModule (B gH gW : Nat) (anchorW anchorH : Float := 1.0) : String

Standalone DIoU box-loss probe: @main(pred,tgt : [B,4,gH,gW], mask : [B,gH,gW]) -> (loss, d_pred). Compiled + run on CPU by scripts/diou_probe_check.py — the emitted forward is checked against numpy and the emitted backward against central finite differences, before wiring into the YOLOv1 train step.

Equations

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

def MlirCodegen.anchorLossProbeModule (B gH gW A : Nat) : String

Standalone anchor-YOLO-loss probe: @main(pred,tgt : [B,A·15,gH,gW], mask : [B,A,gH,gW]) -> (loss, d_pred). Deterministic test anchors (w=0.02+0.03·i, h=0.03+0.04·i), γ=2, λ_box=5 — mirrored in scripts/anchor_loss_probe_check.py, which checks the emitted forward against numpy and the emitted backward against finite differences.

Equations

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

def MlirCodegen.fpnNeckProbeModule (B oc c3 c4 c5 g5 : Nat) : String

Standalone FPN-neck probe: @main(C3,C4,C5, W3,W4,W5, dP3,dP4,dP5) -> (P3,P4,P5, dC3,dC4,dC5, dW3,dW4,dW5). Compiled + run on CPU by scripts/fpn_neck_probe_check.py, which checks the emitted forward against the numpy fpn_forward and the emitted backward against the f64-FD-verified fpn_grad oracle (brick #3). Cotangents are explicit inputs — no scalar-loss reduce — so the probe mirrors the real train-step wiring (head backwards feed dPn tensors straight into the neck VJP).

Equations

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

def MlirCodegen.fpnLossProbeModule (B : Nat) (scaleGrids : List Nat) (A : Nat) (clsWeights : List Float := []) (clsFocalGamma : Float := 0.0) : String

Standalone FPN multi-scale-loss probe: @main(logits [B,Ntot], T0,T1,...) -> (loss, grad [B,Ntot]), one target block per scale. Exercises the concat split

• per-scale anchor loss + grad re-concat (bites 4+6) in isolation — conv-free, so it CPU-compiles for FD checking (the conv heads feeding this are verified convBn, validated separately on ROCm). Deterministic per-scale anchors (w=0.02+0.03·i, h=0.03+0.04·i), matching scripts/fpn_loss_probe_check.py.

Equations

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

def MlirCodegen.fpnDetectProbeModule (B oc c3 c4 c5 g5 A : Nat) (tower : Nat := 0) : String

Standalone whole-FPN-detector probe: neck + optional RetinaNet tower + 1×1-conv heads + concat + multi-scale loss + full DAG backward. Run at focal γ=0 so the objectness weight is a genuine constant (not just detached) and EVERY input/param grad is exactly finite-differenceable — this validates the DAG backward (all 9 + 6·tower param grads) before the emitTrainStepBody wiring.

@main(C3,C4,C5, Wn3,Wn4,Wn5, [Wt_i,bt_i × 3 levels × tower], Wh3,Wh4,Wh5, bh3,bh4,bh5, T3,T4,T5) -> (loss, dC3,dC4,dC5, <the same param grads in the same order>)

Param args after the neck follow fpnDetectParamShapes exactly, so the probe is also a check that that ordering is what the emitters actually consume. tower = 0 keeps the module conv1×1-only (pure dot_general), which is what lets it CPU-compile; tower > 0 introduces real 3×3 stablehlo.convolution ops, so CPU-compilability there is not guaranteed (see the probe checker). Deterministic per-scale anchors (w=0.02+0.03·i, h=0.03+0.04·i).

Equations

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

def MlirCodegen.flashProbeModule (b heads n dh bk : Nat) (causal : Bool) : String

Standalone FlashAttention forward module — exercises emitFlashAttnSdpa in isolation for validation against dense attention: @main(Q,K,V : [b,h,n,dh]) -> O. Compiled + run in the flash-attn probe (planning/archive/flash_attention.md rung 2-3).

Equations

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

def MlirCodegen.denseSdpaProbeModule (b heads n dh : Nat) (causal : Bool) : String

Standalone DENSE SDPA forward module @main(Q,K,V) -> O (materializes the full [b,h,n,n] scores) — the memory-comparison baseline for the flash probe. Compile both with --iree-scheduling-dump-statistics-file at increasing n to see dense peak allocation grow O(n²) while flash stays O(n·bk).

Equations

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

def MlirCodegen.ropeProbeModule (b heads n dh : Nat) (backward : Bool) : String

Standalone RoPE module @main(x) -> rope(x) (or the VJP with backward), for op-level validation against jax/demos/rope_ref.py.

Equations

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

def MlirCodegen.flashBwdProbeModule (b heads n dh bk : Nat) (causal : Bool) : String

Standalone FlashAttention forward+backward module: @main(Q,K,V,dO) -> (dQ,dK,dV). Runs the forward emit to get O/Lse, then the backward emit; validated against the dense-attention VJP.

Equations

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