Documentation

LeanMlir.Proofs.Float.Bf16Fold

The bf16-mixed MNIST-linear render-tie #

The companion to E4M3Fold.lean. The fp8 render-tie is proven but not lowered (the IREE CUDA backend could not emit f8; see historical/upstream-issues/2026-06-iree-cuda-fp8-nvptx-lowering/). bf16 lowers on CUDA (a bf16-in / f32-accumulate dot_general compiles for sm_86/sm_89), and its mixed-precision accuracy bound is dense_close_mixed (u_leaf = 2⁻⁸). This file proves the structural half: the bf16-mixed linear graph denotes the rounded-operand linear, for any rounding rnd (bf16 being one instance). No accuracy claim — purely "the bf16 cast + matmul + fp32 accumulate computes dense of the rounded operands".

Why it is simpler than the fp8 tie. The deployed bf16-mixed kernel casts operands to bf16 and multiply-accumulates in fp32. The fp32 accumulate makes the reduction the exact ∑ (in ℝ), so — as the fp8 tie treats the int accumulate — the only deviation from exact is in the operands. There is no block scale to factor through the sum (fp8's QuantFold.dequant_factors), so the tie falls straight out of the den-faithful operand/dotIn/addBcast ops.

No SHlo surgery at depth 1. The bf16 cast on the input activations is baked into the operand value rnd ∘ x — the byte preparation the kernel does before the GEMM — mirroring E4M3Fold's host-side codes. Depth 2 rounds the intermediate activation in the graph with the convertF round node (den (convertF rnd e) = rnd ∘ den e), below.

Accuracy is separate. Instantiating this tie's abstract rnd at bf16 round-to-nearest and feeding |rnd x − x| ≤ 2⁻⁸|x| into dense_close_mixed (u_leaf = 2⁻⁸, u_acc = 2⁻²⁴) would give an accuracy bound for the tied graph; that composition is not stated in the repo (there is no bf16 FloatModel).

The theorems here are rfl; tests/AuditAxioms.lean prints all four.

source
noncomputable def Proofs.Bf16Fold.actBf16 {m : ℕ} (rnd : ℝ → ℝ) (x : Vec m) :
Vec m

The bf16-rounded input activations rnd(xᵢ) — the bf16 bytes the kernel feeds the GEMM (the leaf cast, baked into the operand like fp8's actCode).

Equations
Instances For
    source
    noncomputable def Proofs.Bf16Fold.wBf16 {m n : ℕ} (rnd : ℝ → ℝ) (W : Mat m n) :
    Mat m n

    The bf16-rounded weights rnd(Wᵢⱼ) (cast once; the other GEMM operand).

    Equations
    Instances For
      source
      noncomputable def Proofs.Bf16Fold.bf16LinearGraph {m n : ℕ} (rnd : ℝ → ℝ) (W : Mat m n) (b : Vec n) (x : Vec m) :
      StableHLO.SHlo n

      The emitted bf16-mixed linear graph. Built only from den-faithful SHlo ops: operand (the bf16 activation bytes) → dotIn (bf16 weights; the den ∑ is the fp32 accumulate) → addBcast (the fp32 bias). The "bf16 leaf, fp32 accumulate" linear. No block scale (cf. e4m3LinearGraph's layerScaleF).

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        source
        noncomputable def Proofs.Bf16Fold.bf16Linear {m n : ℕ} (rnd : ℝ → ℝ) (W : Mat m n) (b : Vec n) (x : Vec m) :
        Vec n

        The intended algorithm: the exact-ℝ linear map on the rounded operands, mnistLinear (rnd∘W) b (rnd∘x). fp32 accumulate ⇒ the ∑ is exact; the only deviation from mnistLinear W b x lives in rnd.

        Equations
        Instances For
          source
          theorem Proofs.Bf16Fold.bf16_render_faithful {m n : ℕ} (rnd : ℝ → ℝ) (W : Mat m n) (b : Vec n) (x : Vec m) :
          StableHLO.den (bf16LinearGraph rnd W b x) = bf16Linear rnd W b x

          bf16 render-tie (structural faithfulness). The emitted bf16-mixed graph denotes exactly the rounded-operand linear, for any rnd (bf16 round-to-nearest being one instance). No accuracy claim — purely "the bf16 cast + fp32-accumulate matmul implements dense of the rounded operands". Strictly simpler than e4m3_render_faithful: no scale to factor, so it is the den of the faithful ops unfolded.

          The EMITTABLE graph — same tie, one dotInBf16 node #

          bf16LinearGraph above pre-rounds its operands (wBf16/actBf16), which is right for the proof and unemittable in practice: expressing "round, then multiply" as separate nodes means a convert PAIR, and XLA deletes those (measured — see convertF's comment). The dotInBf16 node bundles the casts into the matmul so the value stays bf16 across the op, which is the only shape that survives the optimizer and reaches tensor cores.

          The point of this section is that bundling costs no proof: the tie is the SAME rfl.

          source
          noncomputable def Proofs.Bf16Fold.bf16LinearGraphEmit {m n : ℕ} (rnd : ℝ → ℝ) (W : Mat m n) (b : Vec n) (x : Vec m) :
          StableHLO.SHlo n

          The bf16 linear graph as actually emitted: raw operands, one dotInBf16.

          Equations
          Instances For
            source
            theorem Proofs.Bf16Fold.bf16_render_faithful_emit {m n : ℕ} (rnd : ℝ → ℝ) (W : Mat m n) (b : Vec n) (x : Vec m) :
            StableHLO.den (bf16LinearGraphEmit rnd W b x) = bf16Linear rnd W b x

            The emittable graph denotes the same thing as the pre-rounded one. So bf16_render_faithful transfers to the dotInBf16 node, the form that can be lowered.

            Depth 2 — the convertF round node #

            Rounding an intermediate activation uses the in-graph convertF round node (den (convertF rnd e) = rnd ∘ den e, in StableHLO). The tie composes: a rounded activation feeding the next layer is still exactly dense of rounded operands, with no cross-layer error term to track, because the fp32 accumulate keeps every ∑ exact in ℝ. So the two-layer statement is the one-layer statement applied twice.

            source
            noncomputable def Proofs.Bf16Fold.bf16Depth2Graph {m n : ℕ} (rnd : ℝ → ℝ) {p : ℕ} (W₀ : Mat m n) (b₀ : Vec n) (W₁ : Mat n p) (b₁ : Vec p) (x : Vec m) :
            StableHLO.SHlo p

            A depth-2 bf16-mixed graph: rounded operands at BOTH layers, and — the new part — the intermediate activation rounded in the graph by convertF, exactly where a bf16 kernel would hand its f32 accumulator back as bf16 bytes for the next GEMM.

            Equations
            • One or more equations did not get rendered due to their size.
            Instances For
              source
              noncomputable def Proofs.Bf16Fold.bf16Depth2 {m n : ℕ} (rnd : ℝ → ℝ) {p : ℕ} (W₀ : Mat m n) (b₀ : Vec n) (W₁ : Mat n p) (b₁ : Vec p) (x : Vec m) :
              Vec p

              The intended depth-2 algorithm: bf16Linear applied to itself. No explicit rnd ∘ between the layers — bf16Linear already rounds its input operand via actBf16, and that is precisely the rounding the in-graph convertF performs. The layer-1 leaf cast and the layer-0 output round are THE SAME CAST, seen from the two sides, which is why the composition needs no glue.

              Equations
              Instances For
                source
                theorem Proofs.Bf16Fold.bf16_render_faithful_depth2 {m n : ℕ} (rnd : ℝ → ℝ) {p : ℕ} (W₀ : Mat m n) (b₀ : Vec n) (W₁ : Mat n p) (b₁ : Vec p) (x : Vec m) :
                StableHLO.den (bf16Depth2Graph rnd W₀ b₀ W₁ b₁ x) = bf16Depth2 rnd W₀ b₀ W₁ b₁ x

                bf16 render-tie at depth 2. The emitted graph — including the in-graph round — denotes the rounded-operand two-layer linear, for any rnd. Still rfl: no accuracy reasoning, no error propagation, purely denotational.

                No depth-k theorem is stated; the depth-2 proof has no depth-dependent constant, which is the structural content of "bf16 has no block scale to factor through the sum". The round node in isolation is Proofs.StableHLO.convertF_faithful; it is not restated here.