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₂₂)
Used by 3

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)
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
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)
Uses 1
4 fully proved

The dependency explorer needs JavaScript and its graph data.

Download the complete graph data.