Documentation

LeanMlir.Proofs.Codegen.CnnRender

CNN + CIFAR render half — conv train-step text as a name-threaded render of proven graphs #

The peer of MlpRender.lean (mlpTrainStepStructured) for the Chapter-3 MNIST CNN (cnnTrainStepStructured) and the Chapter-4 CIFAR CNN (cifarTrainStepStructured). The MLP render was all-flat, so pretty's flat result names fed the backward/param-grad templates directly. The CNN forward graph (cnnFwdGraph) is also rendered all-flat — each .flatConvF/.maxPoolF token reshapes flat→NCHW internally and back to flat at its boundary (emitTok, StableHLO.lean), so the names pretty exposes are flat. But the conv-specific tail ops (convWGrad, selMask4, select_and_scatter) consume the 4-D NCHW activations. We bridge that with explicit reshape glue in the tail: capture the flat pre-acts/acts from pretty (cnnFwdGraph …) (proof-rendered), reshape the four the conv tail needs (%hc1,%ac1,%hc2,%ac2) plus %xr back to [B,c,H,W], then emit the GPU-validated backward/param-grad/SGD templates around the captured names. reshape is a semantic/GPU no-op (flat and NCHW are the same buffer).

The forward pieces (flatConvF/reluF/maxPoolF/denseF, loss cotangent) are denotable and proven faithful (flatConvF_faithful/reluF_faithful/maxPoolF_faithful/ denseF_faithful/lossCotGraph-style); the denotation-side close (each conv/dense SGD output denotes θ − lr·certified) is cnn_render_conv{W,b}_certified + the M2 dense bridges (CnnTrainStep.lean). See planning/archive/render_close_handoff.md §1.

def Proofs.StableHLO.cnnTrainStepStructured (B ic c h w d1 nClasses kH kW : ) (lr : String) (W₁ : Kernel4 c ic kH kW) (b₁ : Vec c) (W₂ : Kernel4 c c kH kW) (b₂ : Vec c) (W₃ : Mat (c * h * w) d1) (b₃ : Vec d1) (W₄ : Mat d1 d1) (b₄ : Vec d1) (W₅ : Mat d1 nClasses) (b₅ : Vec nClasses) (x : Vec (ic * (2 * h) * (2 * w))) :

Structured CNN train-step renderer (@cnn_train_step): forward conv/relu/maxpool/dense pre-acts/activations/logits/cotangent from the proven cnnFwdGraph pieces (name-threaded via pretty), then the backward + conv/dense param-grad + SGD ops referencing the captured names. The four 4-D activations the conv tail needs (%hc1,%ac1,%hc2,%ac2) and %xr are recovered by explicit flat→NCHW reshapes. Dim convention matches cnnFwdGraph/cnnFwdModuleV: h,w are the POST-pool spatial sizes, the image is 2h × 2w (so H := 2h, W := 2w, flattened map flat := c·h·w).

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    def Proofs.StableHLO.cnnTrainStepFaithfulV (B ic c h w d1 nClasses kH kW : ) (lrStr : String) (W₁ : Kernel4 c ic kH kW) (b₁ : Vec c) (W₂ : Kernel4 c c kH kW) (b₂ : Vec c) (W₃ : Mat (c * h * w) d1) (b₃ : Vec d1) (W₄ : Mat d1 d1) (b₄ : Vec d1) (W₅ : Mat d1 nClasses) (b₅ : Vec nClasses) (x : Vec (ic * (2 * h) * (2 * w))) :

    MNIST-CNN train step rendered ENTIRELY from the verified AST. The peer of mlpTrainStepFaithfulV for the conv net: like cnnTrainStepStructured for the forward, but the backward chain (dotOut/selectPos/maxPoolBack/convBack) and ALL ten parameter SGD updates are now pretty of denoted SHlo nodes too — the dense head via weightSgd/biasSgd, the conv layers via the new convWeightSgd/convBiasSgd ops. So every emitted line is pretty(provenNode), and CnnFold proves each output's den = the certified loss-descent step. Cotangents (%dy/dy4/dy3/dac2/dhc2/dac1/dhc1) are rendered once and shared as operand leaves; operand/lr/weight VALUES are skel-erased, so these placeholders print identically to the live graphs the den theorems use. Dim convention matches cnnTrainStepStructured (h,w post-pool, image 2h×2w).

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      def Proofs.StableHLO.cifarTrainStepStructured (B ic c1 c2 h w d1 nClasses kH kW : ) (lr : String) (W₁ : Kernel4 c1 ic kH kW) (b₁ : Vec c1) (W₂ : Kernel4 c1 c1 kH kW) (b₂ : Vec c1) (W₃ : Kernel4 c2 c1 kH kW) (b₃ : Vec c2) (W₄ : Kernel4 c2 c2 kH kW) (b₄ : Vec c2) (W₅ : Mat (c2 * h * w) d1) (b₅ : Vec d1) (W₆ : Mat d1 d1) (b₆ : Vec d1) (W₇ : Mat d1 nClasses) (b₇ : Vec nClasses) (x : Vec (ic * (2 * (2 * h)) * (2 * (2 * w)))) :

      Structured CIFAR CNN train-step renderer (@cifar_train_step): the Chapter-4 peer of cnnTrainStepStructured, a re-parameterization across two conv→conv→pool stages at two spatial scales (channels ic→c1→c1 then c1→c2→c2; spatial H×W → H/2 → H/4). Forward rendered all-flat from the proven cifarFwdGraph; the conv tail's 4-D consumers are recovered by ten flat→NCHW reshapes (%hc{1..4},%ac{1..4},%pool1,%xr, at the two scales). Tail = the GPU-validated backward/grad/SGD templates wired to the captured names. Dim convention matches cifarFwdGraph/cifarFwdModuleV: h,w are the FINAL pooled spatial sizes, so the image is 4h × 4w (H := 4h, stage-2 H2 := 2h, flattened map flat := c2·h·w). The close is cnn_render_conv{W,b}_certified (generic in dims — covers all four conv layers) + the M2 dense bridges. See planning/archive/render_close_handoff.md §2a.

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        def Proofs.StableHLO.cifarTrainStepFaithfulV (B ic c1 c2 h w d1 nClasses kH kW : ) (lrStr : String) (W₁ : Kernel4 c1 ic kH kW) (b₁ : Vec c1) (W₂ : Kernel4 c1 c1 kH kW) (b₂ : Vec c1) (W₃ : Kernel4 c2 c1 kH kW) (b₃ : Vec c2) (W₄ : Kernel4 c2 c2 kH kW) (b₄ : Vec c2) (W₅ : Mat (c2 * h * w) d1) (b₅ : Vec d1) (W₆ : Mat d1 d1) (b₆ : Vec d1) (W₇ : Mat d1 nClasses) (b₇ : Vec nClasses) (x : Vec (ic * (2 * (2 * h)) * (2 * (2 * w)))) :

        CIFAR-CNN (Chapter 4, no-BN) train step rendered ENTIRELY from the verified AST. The deeper, two-spatial-scale peer of cnnTrainStepFaithfulV: like cifarTrainStepStructured for the forward, but the backward chain (dotOut/selectPos/maxPoolBack/convBack, twice through) and all 14 parameter SGD updates are now pretty of denoted SHlo nodes — the dense head via weightSgd/biasSgd, the four conv layers via the convWeightSgd/convBiasSgd ops (reused from cnn, NO new ops). Every emitted line is pretty(provenNode), and CifarFold proves each output's den = the certified loss-descent step. Dim convention matches cifarTrainStepStructured (h,w final pooled; image 4h×4w, stage-2 spatial 2h×2w).

        Equations
        • One or more equations did not get rendered due to their size.
        Instances For
          def Proofs.StableHLO.cifar8TrainStepFaithfulV (B ic c1 c2 c3 c4 h w d1 nClasses kH kW : ) (lrStr : String) (W₁ : Kernel4 c1 ic kH kW) (b₁ : Vec c1) (W₂ : Kernel4 c1 c1 kH kW) (b₂ : Vec c1) (W₃ : Kernel4 c2 c1 kH kW) (b₃ : Vec c2) (W₄ : Kernel4 c2 c2 kH kW) (b₄ : Vec c2) (W₅ : Kernel4 c3 c2 kH kW) (b₅ : Vec c3) (W₆ : Kernel4 c3 c3 kH kW) (b₆ : Vec c3) (W₇ : Kernel4 c4 c3 kH kW) (b₇ : Vec c4) (W₈ : Kernel4 c4 c4 kH kW) (b₈ : Vec c4) (W₉ : Mat (c4 * h * w) d1) (b₉ : Vec d1) (Wa : Mat d1 d1) (ba : Vec d1) (Wb : Mat d1 nClasses) (bb : Vec nClasses) (x : Vec (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w)))))) (bf16 : Bool := false) :

          Deeper 8-conv CIFAR (cifar8, no-BN) train step rendered ENTIRELY from the verified AST. The 4-stage peer of cifarTrainStepFaithfulV ((conv→relu)×2→pool ×4, 3 dense; 22 params). Backward chain (dotOut/selectPos/maxPoolBack/convBack, four stages) and all 22 param SGD ops are pretty of denoted nodes — conv via convWeightSgd/ convBiasSgd, dense via weightSgd/biasSgd (NO new ops). Cifar8Fold proves each output's den = certified. h,w are the final pooled sizes; stage spatials build up ×2 per pool (s4=2h, s3=4h, s2=8h, s1=16h; image 16h×16w).

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

            Which optimizer tail the cifar8 render emits (handoff §2i). All three share ONE forward, backward and un-fused-gradient body, and one packed [θ|m|v] signature — 71 in / 69 out for every variant — so the ablation section's "SGD several ways" is genuinely the same net with the optimizer swapped, and a reader can diff the artifacts to see only the tail move.

            • adamw : CifarOpt

              θ' = θ − lr·(m̂/(√v̂+ε)) − lr·wd·θ, and both moments live.

            • sgd : CifarOpt

              θ' = θ − lr·g; m/v ride through untouched.

            • nesterov : CifarOpt

              v' = μ·v + g, θ' = θ − lr·(g + μ·v'); velocity in the v slot, m untouched.

            Instances For
              @[instance_reducible]
              Equations
              Equations
              • One or more equations did not get rendered due to their size.
              Instances For
                def Proofs.StableHLO.cifar8AdamTrainStepFaithfulV (B ic c1 c2 c3 c4 h w d1 nClasses kH kW : ) (invBStr b1Str ob1Str b2Str ob2Str epsStr wdStr : String) (W₁ : Kernel4 c1 ic kH kW) (b₁ : Vec c1) (W₂ : Kernel4 c1 c1 kH kW) (b₂ : Vec c1) (W₃ : Kernel4 c2 c1 kH kW) (b₃ : Vec c2) (W₄ : Kernel4 c2 c2 kH kW) (b₄ : Vec c2) (W₅ : Kernel4 c3 c2 kH kW) (b₅ : Vec c3) (W₆ : Kernel4 c3 c3 kH kW) (b₆ : Vec c3) (W₇ : Kernel4 c4 c3 kH kW) (b₇ : Vec c4) (W₈ : Kernel4 c4 c4 kH kW) (b₈ : Vec c4) (W₉ : Mat (c4 * h * w) d1) (b₉ : Vec d1) (Wa : Mat d1 d1) (ba : Vec d1) (Wb : Mat d1 nClasses) (bb : Vec nClasses) (x : Vec (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w)))))) (replicas : := 1) (opt : CifarOpt := CifarOpt.adamw) (bf16 : Bool := false) :

                cifar8 AdamW train step rendered ENTIRELY from the verified AST — the optimizer half of planning/archive/xla_pjrt_handoff.md §2a. Identical forward/backward to cifar8TrainStepFaithfulV; the 22 fused SGD ops are replaced by 22 un-fused param gradients (convWeightGrad/convBiasGrad/weightGrad/biasGrad) each feeding the three proven AdamW ops (adamWParamF/adamMNextF/adamVNextF, denoting Proofs.adamWStep).

                Signature matches the packed trainAdamSched protocol byte for byte: (x, θ×22, m×22, v×22, %lr, %bc1, %bc2, onehot) → (θ'×22, m'×22, v'×22, loss, bc1, bc2). β₁/β₂/ε/wd are baked as in-body constants (%b1 %ob1 %b2 %ob2 %eps %wd), so the conv biases are named %cb1…%cb8%b1 is β₁.

                Two lines are outside the proven surface, both marked in the emitted text: the scalar %loss (report-only; the kit has no rank-0 loss op and it does not feed the update) and the %bc1/%bc2 passthroughs. The mean-loss 1/B on the cotangent IS proven — it is scaleF, not hand-written text. Unlike the SGD render it cannot be folded into lr, because lr is a runtime scalar here.

                Instances For
                  def Proofs.StableHLO.cifar8AdamTrainStepFaithfulB (B ic c1 c2 c3 c4 h w d1 nClasses kH kW : ) (invBStr b1Str ob1Str b2Str ob2Str epsStr wdStr : String) (W₁ : Kernel4 c1 ic kH kW) (b₁ : Vec c1) (W₂ : Kernel4 c1 c1 kH kW) (b₂ : Vec c1) (W₃ : Kernel4 c2 c1 kH kW) (b₃ : Vec c2) (W₄ : Kernel4 c2 c2 kH kW) (b₄ : Vec c2) (W₅ : Kernel4 c3 c2 kH kW) (b₅ : Vec c3) (W₆ : Kernel4 c3 c3 kH kW) (b₆ : Vec c3) (W₇ : Kernel4 c4 c3 kH kW) (b₇ : Vec c4) (W₈ : Kernel4 c4 c4 kH kW) (b₈ : Vec c4) (W₉ : Mat (c4 * h * w) d1) (b₉ : Vec d1) (Wa : Mat d1 d1) (ba : Vec d1) (Wb : Mat d1 nClasses) (bb : Vec nClasses) (x : Vec (B * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) (replicas : := 1) (opt : CifarOpt := CifarOpt.adamw) (bf16 fp8 : Bool := false) :

                  cifar8AdamTrainStepFaithfulB — the batched peer of cifar8AdamTrainStepFaithfulV.

                  Same net, same three optimizers, same packed [θ|m|v] signature. The difference is the op FAMILY: this render carries the batch in the Lean type (SHlo (B*(c*h*w))) and uses the batched constructors, where the …V render is per-example (SHlo (c*h*w)) with pretty B broadcasting. Naming follows the ImageNet renderers, where …RenderB is exactly this migration done once per net (ResNet34RenderB, MobileNetV2RenderB, …).

                  ⭐⭐ Why it exists: bf16, and rehearsal. The 27 bf16 ops were built for ImageNet, which is entirely on the batched family — so bf16 twins exist for convBackBatched/denseRowBack and do NOT exist for the per-example convBack/dotOut that …V uses. Rather than write two CIFAR-only bf16 ops (twins of convBack and dotOut) that ImageNet would never run, this moves CIFAR onto the ops ImageNet already uses. bf16 then drops in for the whole step, forward AND backward, with zero new verified ops — and CIFAR becomes a real rehearsal for ImageNet instead of a parallel dialect. See planning/archive/cifar_lowprec_stability.md §4.1.

                  The migration is semantically free, not a re-derivation. Both families denote the SAME proven VJP — StableHLO.lean l.2016 vs l.2200 are (conv2d_has_vjp3 W b).backward v … and batchMap N (… (conv2d_has_vjp3 W b)).backward (fun _ => 0) …. The only difference is the primal argument, and l.2990 records why it is free: conv is linear, so this is a global VJP — the input-VJP ignores the primal. That is why .convBack's primal argument is simply dropped below rather than threaded.

                  ⚠ Faithful by CONSTRUCTION, like every render here: the AST is built only from verified constructors, so pretty(provenGraph) needs no new proof. Nothing in this function is hand-written MLIR except the report-only %loss, exactly as in the …V peer.

                  ⚠ Parameters are NOT batched — only activations are. The optimizer tail (optTail) is therefore untouched and shared verbatim with …V.

                  Instances For
                    def Proofs.StableHLO.cifar8BnTrainStepFaithfulV (B ic c1 c2 c3 c4 h w d1 nClasses kH kW : ) (epsStr lrStr : String) (W₁ : Kernel4 c1 ic kH kW) (b₁ : Vec c1) (W₂ : Kernel4 c1 c1 kH kW) (b₂ : Vec c1) (W₃ : Kernel4 c2 c1 kH kW) (b₃ : Vec c2) (W₄ : Kernel4 c2 c2 kH kW) (b₄ : Vec c2) (W₅ : Kernel4 c3 c2 kH kW) (b₅ : Vec c3) (W₆ : Kernel4 c3 c3 kH kW) (b₆ : Vec c3) (W₇ : Kernel4 c4 c3 kH kW) (b₇ : Vec c4) (W₈ : Kernel4 c4 c4 kH kW) (b₈ : Vec c4) (W₉ : Mat (c4 * h * w) d1) (b₉ : Vec d1) (Wa : Mat d1 d1) (ba : Vec d1) (Wb : Mat d1 nClasses) (bb : Vec nClasses) (x : Vec (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w)))))) (opt : Option CifarOpt := none) (invBStr : String := "0.0078125") (b1Str : String := "0.9") (ob1Str : String := "0.1") (b2Str : String := "0.999") (ob2Str : String := "0.001") (aEpsStr : String := "1.0e-8") (wdStr : String := "0.0001") :

                    Deeper 8-conv CIFAR-BN (cifar8-bn) train step rendered ENTIRELY from the verified AST. The per-channel-BatchNorm peer of cifar8TrainStepFaithfulV ((conv→BN→relu)×2→pool ×4, 3 dense; 38 params). Pure reuse — NO new ops and NO new proof: conv via convWeightSgd/convBiasSgd, BN via bnGammaSgd/bnBetaSgd, dense via weightSgd/ biasSgd; every output's den = certified by the existing generic lemmas (CifarPoC.conv{W,B}_den, CifarBnPoC.bn{Gamma,Beta}_den, Cifar8PoC.dense{W,B}_den) instantiated per layer. Forward + BN-back proof-rendered via bnPerChannelF/ bnPerChannelBack. h,w final pooled; stage spatials s4=2h…s1=16h.

                    opt selects the optimizer tail (handoff §2i), and it changes the INTERFACE, unlike the no-BN cifar8AdamTrainStepFaithfulV where all three variants share one packed signature:

                    optentryinterfacetail
                    none@cifar8_bn_train_step40 in / 38 outthe 38 fused *Sgd ops, lr a baked literal
                    some o@cifar8_bn_{adam,mom,sgd}_train_step119 in / 117 outun-fused *Grad + optTail o, packed [θ|m|v], %lr runtime

                    Three things branch, not just the tail — this is why §2i scoped it as more than a tail swap: the cotangent (fused folds the batch mean into lrStr; packed cannot, lr is a runtime arg, so it emits an explicit scaleF invB plus a report-only %loss), the conv bias names (%b1..%b8 fused, but AdamW bakes β₁/β₂ as %b1/%b2, so packed renames to %cb1..%cb8 — the same collision-free naming the retired hand-written emitter used), and the signature/return. Everything else — the whole forward, the whole backward, all 38 gradients — is shared verbatim, and at none the render is byte-identical to the incumbent (gate 1), which is what proves the threading inert.

                    Instances For
                      def Proofs.StableHLO.cifar8BnTrainStepFaithfulB (B ic c1 c2 c3 c4 h w d1 nClasses kH kW : ) (bnEpsStr invBStr b1Str ob1Str b2Str ob2Str aEpsStr wdStr : String) (W₁ : Kernel4 c1 ic kH kW) (b₁ : Vec c1) (W₂ : Kernel4 c1 c1 kH kW) (b₂ : Vec c1) (W₃ : Kernel4 c2 c1 kH kW) (b₃ : Vec c2) (W₄ : Kernel4 c2 c2 kH kW) (b₄ : Vec c2) (W₅ : Kernel4 c3 c2 kH kW) (b₅ : Vec c3) (W₆ : Kernel4 c3 c3 kH kW) (b₆ : Vec c3) (W₇ : Kernel4 c4 c3 kH kW) (b₇ : Vec c4) (W₈ : Kernel4 c4 c4 kH kW) (b₈ : Vec c4) (W₉ : Mat (c4 * h * w) d1) (b₉ : Vec d1) (Wa : Mat d1 d1) (ba : Vec d1) (Wb : Mat d1 nClasses) (bb : Vec nClasses) (x : Vec (B * (ic * (2 * (2 * (2 * (2 * h)))) * (2 * (2 * (2 * (2 * w))))))) (replicas : := 1) (opt : CifarOpt := CifarOpt.adamw) (bf16 : Bool := false) :

                      cifar8BnTrainStepFaithfulB — the batched peer of cifar8BnTrainStepFaithfulV.

                      Same net, same three optimizers, same packed [θ|m|v] 38-parameter signature. What moves is the op FAMILY for the convolutions and the dense head: this render carries the batch in the Lean type (SHlo (B*(c*h*w))) and uses the batched constructors, where the …V render is per-example with pretty B broadcasting. It is cifar8AdamTrainStepFaithfulB with BatchNorm spliced in, and it exists for one reason: bf16 on the NORMALIZED net.

                      ⭐⭐ Why the migration is what unlocks bf16. The 27 bf16 ops were built for ImageNet, which is entirely batched, so bf16 twins exist for convBackBatched/convWeightGradB and do NOT exist for the per-example convBack/dotOut the …V render uses. Chapter 4's precision lever could therefore only measure bf16 on the un-normalized net — the one that NaNs under AdamW at fp32 — which is the hardest case and the wrong one for licensing Chapter 5. This render puts bf16 on the net that actually has BatchNorm, with zero new verified ops.

                      BatchNorm stays PER-EXAMPLE, and that is deliberate. bnPerChannelF/bnPerChannelBack reduce over [2,3], not [0,2,3]: each example is normalized by its own statistics. Three consequences, all of them the point:

                      • It is the same net as …V. Swapping to the batch-coupled bnBatchF would change the FUNCTION, not just the emit, and the f32-vs-bf16 comparison would be confounded by a BatchNorm semantics change. The op family is the only variable here.
                      • The eval forward is unchanged. Per-example BN needs no running statistics — train and eval normalize identically — so cifar8w_bn_fwd.mlir is reused verbatim and cannot drift from what the f32 arm evaluates against. Batch BN would need μ/var returned from the step, a host-side EMA, and a frozen-stat eval render; the packed protocol has nowhere to put them.
                      • It matches mixed precision as everyone practises it. Every bf16 net in this repo — R34, R50, MNv2, MNv4, EfficientNet — keeps BN in f32 and rounds only the convolutions and the dense layers (ResNet34RenderB.lean:141 is the pattern: the conv backward branches on bf16, the bnBatchBack above it does not). A bf16 BatchNorm would measure a recipe Chapter 5 does not use.

                      ⚠ The BN nodes are per-example SHlo trees inside a pretty B render, exactly as the head's rows := 1 ops are in cifar8AdamTrainStepFaithfulB: each pretty node is an independent tree, linked to the next only by the SSA name, so a per-example BN node and a batched conv node compose in the emitted text without composing in the Lean types. The emitted BN fragment is byte-identical to the one …V emits.

                      ⚠ Faithful by CONSTRUCTION, like every render here: the AST is built only from verified constructors, so pretty(provenGraph) needs no new proof. Nothing is hand-written MLIR except the report-only %loss and the %bc1/%bc2 passthroughs, exactly as in both peers.

                      ⚠ Parameters are NOT batched — only activations are — so the optimizer tail (optTail) is untouched and shared verbatim with both peers. Conv biases are %cb1…%cb8: %b1 is β₁.

                      Instances For

                        The §2i plain-SGD cifar8 render: cifar8AdamTrainStepFaithfulV with opt := .sgd, so the forward, backward and all 22 un-fused gradients are shared verbatim with the AdamW render and only the tail differs. Entry @cifar8_sgd_train_step, 71 in / 69 out.

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

                          The §2i Nesterov cifar8 render (opt := .nesterov), μ baked at 0.9. Entry @cifar8_mom_train_step, 71 in / 69 out.

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