10. Block-Float Arithmetic
The datapath multiplies integer mantissas and applies the shared block scale at the boundary of a block. The theorem in this chapter is deliberately exact: for every block length and every starting offset, the integer MAC followed by one scale multiplication equals the direct scaled dot.
namespace Janus
theorem scalePull (ea eb x y : Int) :
(ea * x) * (eb * y) = (ea * eb) * (x * y) := ea:Inteb:Intx:Inty:Int⊢ ea * x * (eb * y) = ea * eb * (x * y)
All goals completed! 🐙
theorem scaled_dot (ea eb : Int) (a b : Nat -> Int) :
forall n wp sp,
fsumFrom (fun i => ea * a i) (fun i => eb * b i) wp sp n =
(ea * eb) * fsumFrom a b wp sp n
| 0, wp, sp => ea:Inteb:Inta:Nat → Intb:Nat → Intwp:Natsp:Nat⊢ fsumFrom (fun i => ea * a i) (fun i => eb * b i) wp sp 0 = ea * eb * fsumFrom a b wp sp 0 by ea:Inteb:Inta:Nat → Intb:Nat → Intwp:Natsp:Nat⊢ fsumFrom (fun i => ea * a i) (fun i => eb * b i) wp sp 0 = ea * eb * fsumFrom a b wp sp 0
simp [fsumFrom] All goals completed! 🐙
| n + 1, wp, sp => ea:Inteb:Inta:Nat → Intb:Nat → Intn:Natwp:Natsp:Nat⊢ fsumFrom (fun i => ea * a i) (fun i => eb * b i) wp sp (n + 1) = ea * eb * fsumFrom a b wp sp (n + 1) by ea:Inteb:Inta:Nat → Intb:Nat → Intn:Natwp:Natsp:Nat⊢ fsumFrom (fun i => ea * a i) (fun i => eb * b i) wp sp (n + 1) = ea * eb * fsumFrom a b wp sp (n + 1)
simp [fsumFrom, scaled_dot ea eb a b n] ea:Inteb:Inta:Nat → Intb:Nat → Intn:Natwp:Natsp:Nat⊢ ea * a wp * (eb * b sp) + ea * eb * fsumFrom a b (wp + 1) (sp + 1) n =
ea * eb * (a wp * b sp + fsumFrom a b (wp + 1) (sp + 1) n)
rw [Int.mul_add, ea:Inteb:Inta:Nat → Intb:Nat → Intn:Natwp:Natsp:Nat⊢ ea * a wp * (eb * b sp) + ea * eb * fsumFrom a b (wp + 1) (sp + 1) n =
ea * eb * (a wp * b sp) + ea * eb * fsumFrom a b (wp + 1) (sp + 1) n scalePull ea eb (a wp) (b sp) ea:Inteb:Inta:Nat → Intb:Nat → Intn:Natwp:Natsp:Nat⊢ ea * eb * (a wp * b sp) + ea * eb * fsumFrom a b (wp + 1) (sp + 1) n =
ea * eb * (a wp * b sp) + ea * eb * fsumFrom a b (wp + 1) (sp + 1) n] All goals completed! 🐙
def blockMac (a b : Nat -> Int) (wp sp n : Nat) : Int :=
fsumFrom a b wp sp n
def exactDot (ea eb : Int) (a b : Nat -> Int)
(wp sp n : Nat) : Int :=
fsumFrom (fun i => ea * a i) (fun i => eb * b i) wp sp n
def applyScale (ea eb acc : Int) : Int :=
(ea * eb) * acc
theorem blockMac_correct (ea eb : Int) (a b : Nat -> Int)
(wp sp n : Nat) :
applyScale ea eb (blockMac a b wp sp n) =
exactDot ea eb a b wp sp n := by ea:Inteb:Inta:Nat → Intb:Nat → Intwp:Natsp:Natn:Nat⊢ applyScale ea eb (blockMac a b wp sp n) = exactDot ea eb a b wp sp n
unfold applyScale blockMac exactDot ea:Inteb:Inta:Nat → Intb:Nat → Intwp:Natsp:Natn:Nat⊢ ea * eb * fsumFrom a b wp sp n = fsumFrom (fun i => ea * a i) (fun i => eb * b i) wp sp n
exact (scaled_dot ea eb a b n wp sp).symm All goals completed! 🐙
The default Janus block size is 32, but the theorem is not tied to that value. The concrete block-size-32 example below is a build-time smoke test.
def b32a : Nat -> Int := fun i => if i < 32 then 1 else 0
def b32b : Nat -> Int := fun i => if i < 32 then 2 else 0
theorem blockMac32_demo :
applyScale 3 5 (blockMac b32a b32b 0 0 32) = 960 := by ⊢ applyScale 3 5 (blockMac b32a b32b 0 0 32) = 960
native_decide All goals completed! 🐙
theorem blockMac_demo :
applyScale 2 3 (blockMac wDemo aDemo 0 0 3) = 192 := by ⊢ applyScale 2 3 (blockMac wDemo aDemo 0 0 3) = 192
native_decide All goals completed! 🐙
end Janus