Fix 5-term mul with non-zero beta (#395)

* Fix 5-term mul with non-zero beta

* Add tests

* import ArrayLayouts

Author: jishnub

Project.toml

@@ -1,6 +1,6 @@
 name = "BandedMatrices"
 uuid = "aae01518-5342-5314-be14-df237901396f"
-version = "0.17.36"
+version = "0.17.37"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/generic/matmul.jl

@@ -40,7 +40,6 @@ end
 
 function _banded_muladd!(α, A, x::AbstractVector, β, y)
     m, n = size(A)
-    (length(y) ≠ m || length(x) ≠ n) && throw(DimensionMismatch("*"))
     l, u = bandwidths(A)
     if -l > u # no bands
         _fill_lmul!(β, y)
@@ -51,40 +50,43 @@ function _banded_muladd!(α, A, x::AbstractVector, β, y)
     elseif u < 0 # with -l <= u < 0, that is, all bands lie below the diagnoal.
         # E.g. (l,u) = (2,-1)
         # set lview = l + u >= 0 and uview = 0
-        y[1:-u] .= zero(eltype(y))
+        _fill_lmul!(β, @view(y[1:-u]))
         _banded_gbmv!('N', α, view(A, 1-u:m, :), x, β, view(y, 1-u:m))
         y
     else
         _banded_gbmv!('N', α, A, x, β, y)
     end
 end
 
-materialize!(M::BlasMatMulVecAdd{<:BandedColumnMajor,<:AbstractStridedLayout,<:AbstractStridedLayout,T}) where T<:BlasFloat =
+function materialize!(M::BlasMatMulVecAdd{<:BandedColumnMajor,<:AbstractStridedLayout,<:AbstractStridedLayout,<:BlasFloat})
+    checkdimensions(M)
     _banded_muladd!(M.α, M.A, M.B, M.β, M.C)
+end
 
 function _banded_muladd_row!(tA, α, At, x, β, y)
     n, m = size(At)
-    (length(y) ≠ m || length(x) ≠ n) && throw(DimensionMismatch("*"))
     u, l = bandwidths(At)
     if -l > u # no bands
         _fill_lmul!(β, y)
     elseif l < 0
         _banded_gbmv!(tA, α, view(At, 1-l:n, :,), view(x, 1-l:n), β, y)
     elseif u < 0
-        y[1:-u] .= zero(eltype(y))
+        _fill_lmul!(β, @view(y[1:-u]))
         _banded_gbmv!(tA, α, view(At, :, 1-u:m), x, β, view(y, 1-u:m))
         y
     else
         _banded_gbmv!(tA, α, At, x, β, y)
     end
 end
 
-function materialize!(M::BlasMatMulVecAdd{<:BandedRowMajor,<:AbstractStridedLayout,<:AbstractStridedLayout,T}) where T<:BlasFloat
+function materialize!(M::BlasMatMulVecAdd{<:BandedRowMajor,<:AbstractStridedLayout,<:AbstractStridedLayout,<:BlasFloat})
+    checkdimensions(M)
     α, A, x, β, y = M.α, M.A, M.B, M.β, M.C
     _banded_muladd_row!('T', α, transpose(A), x, β, y)
 end
 
-function materialize!(M::BlasMatMulVecAdd{<:ConjLayout{<:BandedRowMajor},<:AbstractStridedLayout,<:AbstractStridedLayout,T}) where T<:BlasComplex
+function materialize!(M::BlasMatMulVecAdd{<:ConjLayout{<:BandedRowMajor},<:AbstractStridedLayout,<:AbstractStridedLayout,<:BlasComplex})
+    checkdimensions(M)
     α, A, x, β, y = M.α, M.A, M.B, M.β, M.C
     _banded_muladd_row!('C', α, A', x, β, y)
 end
@@ -173,25 +175,17 @@ const ConjOrBandedLayout = Union{AbstractBandedLayout,ConjLayout{<:AbstractBande
 const ConjOrBandedColumnMajor = Union{<:BandedColumnMajor,ConjLayout{<:BandedColumnMajor}}
 
 function _banded_muladd!(α::T, A, B::AbstractMatrix, β, C) where T
-    Am, An = size(A)
-    Bm, Bn = size(B)
-    if An != Bm || size(C, 1) != Am || size(C, 2) != Bn
-        throw(DimensionMismatch("*"))
-    end
-
-    Al, Au = bandwidths(A)
-    Bl, Bu = bandwidths(B)
-
     gbmm!('N', 'N', α, A, B, β, C)
-
     C
 end
 
 materialize!(M::BlasMatMulMatAdd{<:AbstractBandedLayout,<:AbstractBandedLayout,<:BandedColumnMajor}) =
     materialize!(MulAdd(M.α, convert(DefaultBandedMatrix,M.A), convert(DefaultBandedMatrix,M.B), M.β, M.C))
 
-materialize!(M::BlasMatMulMatAdd{<:BandedColumnMajor,<:BandedColumnMajor,<:BandedColumnMajor}) =
+function materialize!(M::BlasMatMulMatAdd{<:BandedColumnMajor,<:BandedColumnMajor,<:BandedColumnMajor})
+    checkdimensions(M)
     _banded_muladd!(M.α, M.A, M.B, M.β, M.C)
+end
 
 
 # function generally_banded_matmatmul!(C::AbstractMatrix{T}, tA::Val, tB::Val, A::AbstractMatrix{U}, B::AbstractMatrix{V}) where {T, U, V}
@@ -246,13 +240,9 @@ end
 ### BandedMatrix * dense matrix
 
 function materialize!(M::MatMulMatAdd{<:BandedColumns, <:AbstractColumnMajor, <:AbstractColumnMajor})
+    checkdimensions(M)
     α, β, A, B, C = M.α, M.β, M.A, M.B, M.C
 
-    mA, nA = size(A)
-    mB, nB = size(B)
-    mC, nC = size(C)
-    (nA == mB && mC == mA && nC == nB) || throw(DimensionMismatch("A has size ($mA, $nA), B has size ($mB, $nB), C has size ($mC, $nC)"))
-
     if iszero(α)
         lmul!(β, C)
     else
@@ -265,11 +255,8 @@ function materialize!(M::MatMulMatAdd{<:BandedColumns, <:AbstractColumnMajor, <:
 end
 
 function materialize!(M::MatMulMatAdd{<:AbstractColumnMajor, <:BandedColumns, <:AbstractColumnMajor})
+    checkdimensions(M)
     α, β, A, B, C = M.α, M.β, M.A, M.B, M.C
-    mA, nA = size(A)
-    mB, nB = size(B)
-    mC, nC = size(C)
-    (nA == mB && mC == mA && nC == nB) || throw(DimensionMismatch("A has size ($mA, $nA), B has size ($mB, $nB), C has size ($mC, $nC)"))
 
     if iszero(α)
         lmul!(β, C)

test/test_banded.jl

@@ -1,6 +1,10 @@
-using BandedMatrices, FillArrays, Test, LinearAlgebra, SparseArrays
+using ArrayLayouts
+using BandedMatrices
 import BandedMatrices: _BandedMatrix
-
+using FillArrays
+using LinearAlgebra
+using SparseArrays
+using Test
 
 # used to test general matrix backends
 struct MyMatrix{T} <: AbstractMatrix{T}
@@ -105,24 +109,38 @@ Base.similar(::MyMatrix, ::Type{T}, m::Int, n::Int) where T = MyMatrix{T}(undef,
     end
 
     @testset "BandedMatrix * Vector" begin
-        let A=brand(10,12,2,3), v=rand(12), w=rand(10)
-            @test A*v ≈ Matrix(A)*v
-            @test A'*w ≈ Matrix(A)'*w
+        let v=rand(12), w=rand(10)
+            for (l,u) in ((2,3), (-2,2), (2,-2), (2,-3))
+                A=brand(length(w),length(v),l,u)
+                @test A*v ≈ Matrix(A)*v
+                # the left-side uses BLAS, while the right doesn't
+                @test mul!(ones(size(A,1)), A, v, 1.0, 2.0) ≈ mul!(ones(size(A,1)), A, v, 1, 2)
+                @test A'*w ≈ Matrix(A)'*w
+                @test mul!(ones(size(A,2)), A', w, 1.0, 2.0) ≈ mul!(ones(size(A,2)), A', w, 1, 2)
+                # explicitly test materialize!
+                @test materialize!(MulAdd(1.0, A', w, 2.0, ones(size(A,2)))) ≈ materialize!(MulAdd(1, A', w, 2, ones(size(A,2))))
+            end
         end
 
         let A=brand(Float64,5,3,2,2), v=rand(ComplexF64,3), w=rand(ComplexF64,5)
             @test A*v ≈ Matrix(A)*v
+            @test mul!(ones(ComplexF64,size(A,1)), A, v, 1.0, 2.0) ≈ mul!(ones(ComplexF64,size(A,1)), A, v, 1, 2)
             @test A'*w ≈ Matrix(A)'*w
+            @test mul!(ones(ComplexF64,size(A,2)), A', w, 1.0, 2.0) ≈ mul!(ones(ComplexF64,size(A,2)), A', w, 1, 2)
         end
 
         let A=brand(ComplexF64,5,3,2,2), v=rand(ComplexF64,3), w=rand(ComplexF64,5)
             @test A*v ≈ Matrix(A)*v
+            @test mul!(ones(ComplexF64,size(A,1)), A, v, 1.0, 2.0) ≈ mul!(ones(ComplexF64,size(A,1)), A, v, 1, 2)
             @test A'*w ≈ Matrix(A)'*w
+            @test mul!(ones(ComplexF64,size(A,2)), A', w, 1.0, 2.0) ≈ mul!(ones(ComplexF64,size(A,2)), A', w, 1, 2)
         end
 
         let A=brand(ComplexF64,5,3,2,2), v=rand(Float64,3), w=rand(Float64,5)
             @test A*v ≈ Matrix(A)*v
+            @test mul!(ones(ComplexF64,size(A,1)), A, v, 1.0, 2.0) ≈ mul!(ones(ComplexF64,size(A,1)), A, v, 1, 2)
             @test A'*w ≈ Matrix(A)'*w
+            @test mul!(ones(ComplexF64,size(A,2)), A', w, 1.0, 2.0) ≈ mul!(ones(ComplexF64,size(A,2)), A', w, 1, 2)
         end
 
         @testset "empty" begin