Strassen algorithm
DefinitionStrassen's matrix accumulator#✓fully proved
Let \(R\) be a type with a zero, an addition, a subtraction and a multiplication (no ring axioms are assumed). For \(k \ge 0\) and \(2^k \times 2^k\) matrices \(A\), \(B\), \(C\) over \(R\), this procedure runs Strassen's seven-product recursion for adding \(AB\) to \(C\). It returns a \(2^k \times 2^k\) matrix together with a pair of counts: the number of scalar multiplications and the number of scalar additions and subtractions it performs. For \(k = 0\), with \(a\), \(b\), \(c\) the single entries of \(A\), \(B\), \(C\), it performs one multiplication and one addition and returns the \(1 \times 1\) matrix with entry \(c + ab\). For \(k + 1\), write \(X_{11}, X_{12}, X_{21}, X_{22}\) for the top-left, top-right, bottom-left and bottom-right \(2^k \times 2^k\) blocks of a matrix \(X\). The procedure forms, entrywise, \(S_1 = B_{12} - B_{22}\), \(S_2 = A_{11} + A_{12}\), \(S_3 = A_{21} + A_{22}\), \(S_4 = B_{21} - B_{11}\), \(S_5 = A_{11} + A_{22}\), \(S_6 = B_{11} + B_{22}\), \(S_7 = A_{12} - A_{22}\), \(S_8 = B_{21} + B_{22}\), \(S_9 = A_{11} - A_{21}\) and \(S_{10} = B_{11} + B_{12}\); obtains \(P_1, \dots, P_7\) by seven recursive calls with the zero matrix as accumulator on the pairs \((A_{11}, S_1)\), \((S_2, B_{22})\), \((S_3, B_{11})\), \((A_{22}, S_4)\), \((S_5, S_6)\), \((S_7, S_8)\) and \((S_9, S_{10})\); and returns the matrix with blocks \(C_{11} + P_5 + P_4 - P_2 + P_6\), \(C_{12} + P_1 + P_2\), \(C_{21} + P_3 + P_4\) and \(C_{22} + P_5 + P_1 - P_3 - P_7\), each sum formed entrywise from left to right. Each entrywise scalar addition or subtraction adds one to the second count and each base-case multiplication adds one to the first; indexing, allocation, copying and zero-fill are uncharged.
Cslib.Algorithms.Lean.Matrix.strassenMatrixAccumulate · 48 lines
public def strassenMatrixAccumulate {R : Type u} [Zero R] [Add R] [Sub R] [Mul R] :
(k : Nat) → (A B C : Vector (Vector R (2 ^ k)) (2 ^ k)) →
TimeM (Nat × Nat) (Vector (Vector R (2 ^ k)) (2 ^ k))
| 0, A, B, C => do
TimeM.tick (1, 0)
let p := (A.get 0).get 0 * (B.get 0).get 0
TimeM.tick (0, 1)
pure #v[#v[(C.get 0).get 0 + p]]
| k + 1, A, B, C => do
let a₁₁ := quadrant k A false false
let a₁₂ := quadrant k A false true
let a₂₁ := quadrant k A true false
let a₂₂ := quadrant k A true true
let b₁₁ := quadrant k B false false
let b₁₂ := quadrant k B false true
let b₂₁ := quadrant k B true false
let b₂₂ := quadrant k B true true
let s₁ ← blockMap (· - ·) b₁₂ b₂₂
let s₂ ← blockMap (· + ·) a₁₁ a₁₂
let s₃ ← blockMap (· + ·) a₂₁ a₂₂
let s₄ ← blockMap (· - ·) b₂₁ b₁₁
let s₅ ← blockMap (· + ·) a₁₁ a₂₂
let s₆ ← blockMap (· + ·) b₁₁ b₂₂
let s₇ ← blockMap (· - ·) a₁₂ a₂₂
let s₈ ← blockMap (· + ·) b₂₁ b₂₂
let s₉ ← blockMap (· - ·) a₁₁ a₂₁
let s₁₀ ← blockMap (· + ·) b₁₁ b₁₂
let z := Vector.replicate (2 ^ k) (Vector.replicate (2 ^ k) (0 : R))
let p₁ ← strassenMatrixAccumulate k a₁₁ s₁ z
let p₂ ← strassenMatrixAccumulate k s₂ b₂₂ z
let p₃ ← strassenMatrixAccumulate k s₃ b₁₁ z
let p₄ ← strassenMatrixAccumulate k a₂₂ s₄ z
let p₅ ← strassenMatrixAccumulate k s₅ s₆ z
let p₆ ← strassenMatrixAccumulate k s₇ s₈ z
let p₇ ← strassenMatrixAccumulate k s₉ s₁₀ z
let c₁₁ ← blockMap (· + ·) (quadrant k C false false) p₅
let c₁₁ ← blockMap (· + ·) c₁₁ p₄
let c₁₁ ← blockMap (· - ·) c₁₁ p₂
let c₁₁ ← blockMap (· + ·) c₁₁ p₆
let c₁₂ ← blockMap (· + ·) (quadrant k C false true) p₁
let c₁₂ ← blockMap (· + ·) c₁₂ p₂
let c₂₁ ← blockMap (· + ·) (quadrant k C true false) p₃
let c₂₁ ← blockMap (· + ·) c₂₁ p₄
let c₂₂ ← blockMap (· + ·) (quadrant k C true true) p₅
let c₂₂ ← blockMap (· + ·) c₂₂ p₁
let c₂₂ ← blockMap (· - ·) c₂₂ p₃
let c₂₂ ← blockMap (· - ·) c₂₂ p₇
pure (joinBlocks k c₁₁ c₁₂ c₂₁ c₂₂)
@[no_expose] private def quadrant {R : Type u} (k : Nat)
(M : Vector (Vector R (2 ^ (k + 1))) (2 ^ (k + 1))) (r c : Bool) :
Vector (Vector R (2 ^ k)) (2 ^ k) :=
Vector.ofFn fun i => Vector.ofFn fun j =>
(M.get (halfEquiv k (if r then Sum.inr i else Sum.inl i))).get
(halfEquiv k (if c then Sum.inr j else Sum.inl j))
CSLibExt/Algorithms/Lean/Matrix/Strassen/Basic.lean:53
@[no_expose] private def blockMap {R : Type u} {n : Nat} (op : R → R → R)
(A B : Vector (Vector R n) n) :
TimeM (Nat × Nat) (Vector (Vector R n) n) :=
(A.zipWith (fun a b => a.zipWith Prod.mk b) B).mapM fun row =>
row.mapM fun ab => do
TimeM.tick (0, 1)
pure (op ab.1 ab.2)
CSLibExt/Algorithms/Lean/Matrix/Strassen/Basic.lean:42
@[no_expose] private def joinBlocks {R : Type u} (k : Nat)
(a b c d : Vector (Vector R (2 ^ k)) (2 ^ k)) :
Vector (Vector R (2 ^ (k + 1))) (2 ^ (k + 1)) :=
Vector.ofFn fun i => Vector.ofFn fun j =>
match (halfEquiv k).symm i, (halfEquiv k).symm j with
| Sum.inl x, Sum.inl y => (a.get x).get y
| Sum.inl x, Sum.inr y => (b.get x).get y
| Sum.inr x, Sum.inl y => (c.get x).get y
| Sum.inr x, Sum.inr y => (d.get x).get y
CSLibExt/Algorithms/Lean/Matrix/Strassen/Basic.lean:60
TheoremCorrectness of Strassen's matrix accumulator#✓fully proved
Let \(R\) be a non-unital, non-associative ring. For every \(k \ge 0\) and all \(2^k \times 2^k\) matrices \(A\), \(B\), \(C\) over \(R\), the matrix returned by Strassen's accumulator on \(A\), \(B\), \(C\) equals \(C + AB\), where \(AB\) is the usual matrix product, \((AB)_{ij} = \sum_t A_{it} B_{tj}\).
Multiplication in \(R\) need not be commutative or associative.
public theorem strassenMatrixAccumulate_ret {R : Type u} [NonUnitalNonAssocRing R]
(k : Nat) (A B C : Vector (Vector R (2 ^ k)) (2 ^ k)) :
(Matrix.of fun i j => ((strassenMatrixAccumulate k A B C).ret.get i).get j) =
(Matrix.of fun i j => (C.get i).get j) +
(Matrix.of fun i j => (A.get i).get j) * (Matrix.of fun i j => (B.get i).get j)
public def strassenMatrixAccumulate {R : Type u} [Zero R] [Add R] [Sub R] [Mul R] :
(k : Nat) → (A B C : Vector (Vector R (2 ^ k)) (2 ^ k)) →
TimeM (Nat × Nat) (Vector (Vector R (2 ^ k)) (2 ^ k))
| 0, A, B, C => do
TimeM.tick (1, 0)
let p := (A.get 0).get 0 * (B.get 0).get 0
TimeM.tick (0, 1)
pure #v[#v[(C.get 0).get 0 + p]]
| k + 1, A, B, C => do
let a₁₁ := quadrant k A false false
let a₁₂ := quadrant k A false true
let a₂₁ := quadrant k A true false
⋯
CSLibExt/Algorithms/Lean/Matrix/Strassen/Basic.lean:72
Uses 1 · Used by 1
TheoremEntries of Strassen's matrix accumulator#✓fully proved
Let \(R\) be a non-unital, non-associative ring. For every \(k \ge 0\), all \(2^k \times 2^k\) matrices \(A\), \(B\), \(C\) over \(R\) and all indices \(i, j\), the \((i, j)\) entry of the matrix returned by Strassen's accumulator on \(A\), \(B\), \(C\) is \(C_{ij} + \sum_{t} A_{it} B_{tj}\), the sum running over all \(2^k\) indices \(t\).
public theorem strassenMatrixAccumulate_apply {R : Type u} [NonUnitalNonAssocRing R]
(k : Nat) (A B C : Vector (Vector R (2 ^ k)) (2 ^ k)) (i j : Fin (2 ^ k)) :
((strassenMatrixAccumulate k A B C).ret.get i).get j =
(C.get i).get j + ∑ t : Fin (2 ^ k), (A.get i).get t * (B.get t).get j
public def strassenMatrixAccumulate {R : Type u} [Zero R] [Add R] [Sub R] [Mul R] :
(k : Nat) → (A B C : Vector (Vector R (2 ^ k)) (2 ^ k)) →
TimeM (Nat × Nat) (Vector (Vector R (2 ^ k)) (2 ^ k))
| 0, A, B, C => do
TimeM.tick (1, 0)
let p := (A.get 0).get 0 * (B.get 0).get 0
TimeM.tick (0, 1)
pure #v[#v[(C.get 0).get 0 + p]]
| k + 1, A, B, C => do
let a₁₁ := quadrant k A false false
let a₁₂ := quadrant k A false true
let a₂₁ := quadrant k A true false
⋯
CSLibExt/Algorithms/Lean/Matrix/Strassen/Basic.lean:72
Uses 2
TheoremOperation counts of Strassen's matrix accumulator#✓fully proved
Let \(R\) be a type with a zero, an addition, a subtraction and a multiplication. For every \(k \ge 0\) and all \(2^k \times 2^k\) matrices \(A\), \(B\), \(C\) over \(R\), Strassen's accumulator on \(A\), \(B\), \(C\) performs exactly \(7^k\) scalar multiplications and exactly \((25 \cdot 7^k - 22 \cdot 4^k)/3\) scalar additions and subtractions.
This count does not charge allocation or bit operations.
public theorem strassenMatrixAccumulate_time {R : Type u}
[Zero R] [Add R] [Sub R] [Mul R] (k : Nat)
(A B C : Vector (Vector R (2 ^ k)) (2 ^ k)) :
(strassenMatrixAccumulate k A B C).time =
(7 ^ k, (25 * 7 ^ k - 22 * 4 ^ k) / 3)
public def strassenMatrixAccumulate {R : Type u} [Zero R] [Add R] [Sub R] [Mul R] :
(k : Nat) → (A B C : Vector (Vector R (2 ^ k)) (2 ^ k)) →
TimeM (Nat × Nat) (Vector (Vector R (2 ^ k)) (2 ^ k))
| 0, A, B, C => do
TimeM.tick (1, 0)
let p := (A.get 0).get 0 * (B.get 0).get 0
TimeM.tick (0, 1)
pure #v[#v[(C.get 0).get 0 + p]]
| k + 1, A, B, C => do
let a₁₁ := quadrant k A false false
let a₁₂ := quadrant k A false true
let a₂₁ := quadrant k A true false
⋯
CSLibExt/Algorithms/Lean/Matrix/Strassen/Basic.lean:72
Uses 1
Loading dependency explorer…
The dependency explorer needs JavaScript and its graph data.