Loop over colsupport(vector) in matrix-vector multiplication (#390)

* Loop over colsupport(vector) in matrix-vector multiplication

* Bump version to v0.17.36

* Add tests

Project.toml

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

src/generic/matmul.jl

@@ -99,8 +99,10 @@ end
     checkdimensions(M)
     α,A,B,β,C = M.α,M.A,M.B,M.β,M.C
     _fill_lmul!(β, C)
-    @inbounds for j = rowsupport(A), k = colrange(A,j)
-        C[k] += α*inbands_getindex(A,k,j)*B[j]
+    @inbounds for j = intersect(rowsupport(A), colsupport(B))
+        for k = colrange(A,j)
+            C[k] += α*inbands_getindex(A,k,j)*B[j]
+        end
     end
     C
 end
@@ -111,8 +113,10 @@ end
     A = transpose(At)
     _fill_lmul!(β, C)
 
-    @inbounds for j = rowsupport(A), k = colrange(A,j)
-        C[j] +=  α*transpose(inbands_getindex(A,k,j))*B[k]
+    @inbounds for j = rowsupport(A)
+        for k = intersect(colrange(A,j), colsupport(B))
+            C[j] +=  α*transpose(inbands_getindex(A,k,j))*B[k]
+        end
     end
     C
 end
@@ -122,8 +126,10 @@ end
     α,Ac,B,β,C = M.α,M.A,M.B,M.β,M.C
     A = Ac'
     _fill_lmul!(β, C)
-    @inbounds for j = rowsupport(A), k = colrange(A,j)
-        C[j] += α*inbands_getindex(A,k,j)'*B[k]
+    @inbounds for j = rowsupport(A)
+        for k = intersect(colrange(A,j), colsupport(B))
+            C[j] += α*inbands_getindex(A,k,j)'*B[k]
+        end
     end
     C
 end
@@ -276,4 +282,4 @@ function materialize!(M::MulAdd{<:AbstractColumnMajor, <:BandedColumns, <:Abstra
     end
 
     return C
-end
+end

test/test_linalg.jl

@@ -1,7 +1,25 @@
-using BandedMatrices, ArrayLayouts, LinearAlgebra, FillArrays, Test
+using ArrayLayouts
+using BandedMatrices
+using FillArrays
+using LinearAlgebra
+using Test
+
 import Base.Broadcast: materialize, broadcasted
 import BandedMatrices: BandedColumns, _BandedMatrix
 
+# wrap a OneElement to dispatch without type-piracy
+
+struct MyOneElement{T,N,A<:OneElement{T,N}} <: AbstractArray{T,N}
+    arr :: A
+end
+Base.size(M::MyOneElement) = size(M.arr)
+Base.axes(M::MyOneElement) = axes(M.arr)
+Base.getindex(M::MyOneElement{<:Any,N}, inds::Vararg{Int,N}) where {N} =
+    getindex(M.arr, inds...)
+
+ArrayLayouts.colsupport(::UnknownLayout, A::MyOneElement{<:Any,1}, _) =
+    intersect(axes(A,1), A.arr.ind[1]:A.arr.ind[1])
+
 @testset "Linear Algebra" begin
     @testset "Matrix types" begin
         A = brand(5,5,1,2)
@@ -90,6 +108,17 @@ import BandedMatrices: BandedColumns, _BandedMatrix
         end
     end
 
+    @testset "BandedMatrix * sparse" begin
+        B = brand(6,6,2,2)
+        x = MyOneElement(OneElement(2, 4, 6))
+        y = Array(x)
+        @test B * x ≈ B * y
+        @test B' * x ≈ B' * y
+
+        B = brand(Complex{Int8}, 6,6,2,2)
+        @test B' * x ≈ B' * y
+    end
+
     @testset "gbmm!" begin
         @testset "gbmm! subpieces step by step and column by column" begin
             for n in (1,5,50), ν in (1,5,50), m in (1,5,50),