Matrix-vector multiplication for compatible block sizes (#262)

* matrix-vector multiplication for compatible block sizes

* _block_muladd for AbstractVector

* use MatMulVecAdd alias

* mul instead of muladd for block components

Author: jishnub

Project.toml

@@ -1,6 +1,6 @@
 name = "BlockArrays"
 uuid = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
-version = "0.16.28"
+version = "0.16.29"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/blocklinalg.jl

@@ -199,10 +199,18 @@ function materialize!(M::MatMulVecAdd{<:AbstractBlockLayout,<:AbstractStridedLay
     y_in
 end
 
-function _block_muladd!(α, A, X, β, Y)
+@inline function _block_muladd!(α, A, X::AbstractVector, β, Y::AbstractVector)
+    _fill_lmul!(β, Y)
+    for N = blockcolsupport(X), K = blockcolsupport(A,N)
+        mul!(view(Y,K), view(A,K,N), view(X,N), α, one(β))
+    end
+    Y
+end
+
+@inline function _block_muladd!(α, A, X, β, Y)
     _fill_lmul!(β, Y)
     for J = blockaxes(X,2), N = blockcolsupport(X,J), K = blockcolsupport(A,N)
-        muladd!(α, view(A,K,N), view(X,N,J), one(α), view(Y,K,J))
+        mul!(view(Y,K,J), view(A,K,N), view(X,N,J), α, one(α))
     end
     Y
 end
@@ -211,7 +219,7 @@ mul_blockscompatible(A, B, C) = blockisequal(axes(A,2), axes(B,1)) &&
     blockisequal(axes(A,1), axes(C,1)) &&
     blockisequal(axes(B,2), axes(C,2))
 
-function materialize!(M::MatMulMatAdd{<:AbstractBlockLayout,<:AbstractBlockLayout,<:AbstractBlockLayout})
+@inline function _matmul(M)
     α, A, B, β, C = M.α, M.A, M.B, M.β, M.C
     if mul_blockscompatible(A,B,C)
         _block_muladd!(α, A, B, β, C)
@@ -220,6 +228,14 @@ function materialize!(M::MatMulMatAdd{<:AbstractBlockLayout,<:AbstractBlockLayou
     end
 end
 
+function materialize!(M::MatMulMatAdd{<:AbstractBlockLayout,<:AbstractBlockLayout,<:AbstractBlockLayout})
+    _matmul(M)
+end
+
+function materialize!(M::MatMulVecAdd{<:AbstractBlockLayout, <:AbstractBlockLayout, <:AbstractBlockLayout})
+    _matmul(M)
+end
+
 function materialize!(M::MatMulMatAdd{<:AbstractBlockLayout,<:AbstractBlockLayout,<:AbstractColumnMajor})
     α, A, X, β, Y_in = M.α, M.A, M.B, M.β, M.C
     Y = PseudoBlockArray(Y_in, (axes(A,1), axes(X,2)))

test/test_blocklinalg.jl

@@ -92,6 +92,20 @@ import ArrayLayouts: DenseRowMajor, ColumnMajor, StridedLayout
                     Matrix(A)*PseudoBlockArray(B) .=== Matrix(A)*Matrix(B))
     end
 
+    @testset "BlockMatrix * BlockVector" begin
+        A = BlockArray(randn(6,6), fill(2,3), 1:3)
+        v = BlockArray(rand(6), 1:3)
+        w = A * v
+        @test w isa BlockArray
+        @test blocksizes(w,1) == fill(2, 3)
+        @test w ≈ Array(A) * v ≈ A * Array(v) ≈ Array(A) * Array(v)
+
+        z = A * w
+        @test z isa BlockArray
+        @test blocksizes(z,1) == fill(2, 3)
+        @test z ≈ Array(A) * w ≈ A * Array(w) ≈ Array(A) * Array(w)
+    end
+
     @testset "adjoint" begin
         A = BlockArray(randn(6,6), fill(2,3), 1:3)
         B = BlockArray(randn(6,6), 1:3, 1:3)