LeanMlir.MlirCodegen

MLIR (StableHLO) code generator for the REFERENCE path: emit MLIR modules from a NetSpec at run time. Unverified — the verified path renders verified_mlir/ from the proof side (LeanMlir/Proofs/Codegen/) instead.

Emits dense, conv, BN, depthwise, MBConv/SE, UIB, attention, UNet, ConvNeXt, detection and
token-model layers. The constructors `noEmitter` names have no emitter here, and
`unsupported` refuses any spec that contains one: the walks would skip such a layer and the
net would train without it. Walks layers tracking the current activation shape; for a flat
input it 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.

## References

- Howard et al. 2019, *Searching for MobileNetV3* (h-swish / h-sigmoid). <https://arxiv.org/abs/1905.02244>
- Su et al. 2021, *RoFormer: Enhanced Transformer with Rotary Position Embedding* (RoPE). <https://arxiv.org/abs/2104.09864>
- Ronneberger, Fischer, Brox 2015, *U-Net*. <https://arxiv.org/abs/1505.04597>
- Zheng et al. 2020, *Distance-IoU Loss*. <https://arxiv.org/abs/1911.08287>

source

def MlirCodegen.noEmitter :

Layer → Option String

The constructor's name when this emitter has no lowering for it, else none. The forward and train walks have no arm for these: they write an UNSUPPORTED comment and leave the activation and the parameter list as they were, so the layer would silently become the identity (or a shape error downstream). layerNorm and convNextStem are lowered only by the JAX emitter.

Equations

MlirCodegen.noEmitter (Layer.layerNorm dim) = some "layerNorm"
MlirCodegen.noEmitter (Layer.separableConv ic oc stride) = some "separableConv"
MlirCodegen.noEmitter (Layer.fireModule ic squeeze expand1x1 expand3x3) = some "fireModule"
MlirCodegen.noEmitter (Layer.mambaBlock dim stateSize expand nBlocks) = some "mambaBlock"
MlirCodegen.noEmitter (Layer.swinStage dim heads mlpDim windowSize nBlocks) = some "swinStage"
MlirCodegen.noEmitter (Layer.patchMerging inDim outDim) = some "patchMerging"
MlirCodegen.noEmitter (Layer.transformerDecoder dim heads mlpDim nBlocks nQueries) = some "transformerDecoder"
MlirCodegen.noEmitter (Layer.detrHeads dim nClasses) = some "detrHeads"
MlirCodegen.noEmitter (Layer.shuffleBlock ic oc groups nUnits) = some "shuffleBlock"
MlirCodegen.noEmitter (Layer.shuffleV2Block ic oc nUnits) = some "shuffleV2Block"
MlirCodegen.noEmitter (Layer.evoformerBlock msaChannels pairChannels nBlocks) = some "evoformerBlock"
MlirCodegen.noEmitter (Layer.structureModule singleChannels pairChannels nBlocks) = some "structureModule"
MlirCodegen.noEmitter (Layer.mobileVitBlock ic dim heads mlpDim nTxBlocks) = some "mobileVitBlock"
MlirCodegen.noEmitter (Layer.convNextStem ic oc patch) = some "convNextStem"
MlirCodegen.noEmitter (Layer.waveNetBlock residualCh skipCh nLayers) = some "waveNetBlock"
MlirCodegen.noEmitter (Layer.positionalEncoding inputDim numFrequencies) = some "positionalEncoding"
MlirCodegen.noEmitter (Layer.nerfMLP encodedPosDim encodedDirDim hiddenDim) = some "nerfMLP"
MlirCodegen.noEmitter (Layer.darknetBlock channels nBlocks) = some "darknetBlock"
MlirCodegen.noEmitter (Layer.cspBlock ic oc nBlocks) = some "cspBlock"
MlirCodegen.noEmitter (Layer.inceptionModule ic b1out b2reduce b2out b3reduce b3out b4out) = some "inceptionModule"
MlirCodegen.noEmitter (Layer.asppModule ic oc) = some "asppModule"
MlirCodegen.noEmitter (Layer.fpnModule c2 c3 c4 c5 target) = some "fpnModule"
MlirCodegen.noEmitter (Layer.denseBlock ic growthRate nLayers) = some "denseBlock"
MlirCodegen.noEmitter (Layer.transitionLayer ic oc) = some "transitionLayer"
MlirCodegen.noEmitter x✝ = none

Instances For

source

def MlirCodegen.unsupported (spec : NetSpec) :

Option String

What this emitter cannot render faithfully, or none:

a layer with no emitter (noEmitter);
pad := .valid on a .conv2d/.convBn, which this emitter would pad to keep its spatial size;
a bare .convBn with NetSpec.convBnAct other than ReLU, which this emitter would apply ReLU to anyway.

The JAX emitter implements all three; no spec on this path uses them, so they are refused rather than implemented.

Equations

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

Instances For

source

def MlirCodegen.checkSupported (spec : NetSpec) :

IO Unit

Refuse a spec this emitter cannot render faithfully (unsupported).

Equations

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

Instances For

source

def MlirCodegen.inputFlatDim (spec : NetSpec) :

Nat

Total flat input size: ic · imageH · imageW for an image input (inputChannels); otherwise the first dense layer's fanIn, a token model's sequence (or one-hot) width, or imageH · imageW.

Equations

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

Instances For

source

def MlirCodegen.sanitize (s : String) :

String

Sanitize a string for use as an MLIR identifier: lower-case, every non-alphanumeric character replaced by _. NetSpec.sanitizedName applies it to a spec's name; the trainer applies it to a spec's buildTag.

Equations

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

Instances For

source

def NetSpec.sanitizedName (spec : NetSpec) :

String

The spec's name as an MLIR identifier: the base of every module name the codegen emits (@<sanitizedName>, @<sanitizedName>_eval, jit_<sanitizedName>_train_step, …).

Equations

spec.sanitizedName = MlirCodegen.sanitize spec.name

Instances For

FPN neck (top-down multi-scale merge) #

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/probes/fpn_neck_check.py`), then these emitters are checked against that
oracle by `scripts/probes/fpn_neck_probe_check.py` (the `fpn-neck-probe` exe).

source

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.

Instances For

source

def NetSpec.bnLayers (spec : NetSpec) :

Array (Nat × Nat)

BN-layer (pidx, oc) pairs, one per convBn layer (including those inside residual/bottleneck blocks), in the order the codegen threads their running stats. ViT-style transformer specs return an empty array.

Equations

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

Instances For

source

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.

Instances For

source

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.

Instances For

source

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.

Instances For

source

@[instance_reducible]

instance MlirCodegen.instInhabitedFwdRec :

Inhabited MlirCodegen.FwdRec✝

Equations

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

source

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) (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.

Instances For

source

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/probes/seg_loss_probe_check.py; CPU is fine, no GPU needed.

Instances For

source

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/probes/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.

Instances For

source

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/probes/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.

Instances For

source

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/probes/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. 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).

Instances For

source

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 in isolation — conv-free, so it CPU-compiles for FD checking (the conv layers that feed it in the train step are not part of this probe). Deterministic per-scale anchors (w=0.02+0.03·i, h=0.03+0.04·i), matching scripts/probes/fpn_loss_probe_check.py.

Equations

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

Instances For

source

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).

Instances For

source

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.

Equations

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

Instances For

source

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.

Instances For

source

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.

Instances For

source

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.

Instances For