[SparseArrayInterface] [BlockSparseArrays] Sparse and block sparse dot products (#1577)

* [SparseArrayInterface] [BlockSparseArrays] Sparse and block sparse dot products

* [NDTensors] Bump to v0.3.62

Author: mtfishman

src/abstractblocksparsearray/arraylayouts.jl

@@ -1,14 +1,27 @@
-using ArrayLayouts: ArrayLayouts, MemoryLayout, MulAdd
+using ArrayLayouts: ArrayLayouts, DualLayout, MemoryLayout, MulAdd
 using BlockArrays: BlockLayout
 using ..SparseArrayInterface: SparseLayout
-using ..TypeParameterAccessors: similartype
+using ..TypeParameterAccessors: parenttype, similartype
 
 function ArrayLayouts.MemoryLayout(arraytype::Type{<:BlockSparseArrayLike})
   outer_layout = typeof(MemoryLayout(blockstype(arraytype)))
   inner_layout = typeof(MemoryLayout(blocktype(arraytype)))
   return BlockLayout{outer_layout,inner_layout}()
 end
 
+# TODO: Generalize to `BlockSparseVectorLike`/`AnyBlockSparseVector`.
+function ArrayLayouts.MemoryLayout(
+  arraytype::Type{<:Adjoint{<:Any,<:AbstractBlockSparseVector}}
+)
+  return DualLayout{typeof(MemoryLayout(parenttype(arraytype)))}()
+end
+# TODO: Generalize to `BlockSparseVectorLike`/`AnyBlockSparseVector`.
+function ArrayLayouts.MemoryLayout(
+  arraytype::Type{<:Transpose{<:Any,<:AbstractBlockSparseVector}}
+)
+  return DualLayout{typeof(MemoryLayout(parenttype(arraytype)))}()
+end
+
 function Base.similar(
   mul::MulAdd{<:BlockLayout{<:SparseLayout},<:BlockLayout{<:SparseLayout},<:Any,<:Any,A,B},
   elt::Type,

src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -158,6 +158,13 @@ function Base.similar(
   return similar(arraytype, eltype(arraytype), axes)
 end
 
+# Fixes ambiguity error.
+function Base.similar(
+  arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}}
+)
+  return similar(arraytype, eltype(arraytype), axes)
+end
+
 # Needed by `BlockArrays` matrix multiplication interface
 # TODO: This fixes an ambiguity error with `OffsetArrays.jl`, but
 # is only appears to be needed in older versions of Julia like v1.6.
@@ -221,7 +228,7 @@ function Base.similar(
 end
 
 # Fixes ambiguity error.
-function Base.similar(a::BlockSparseArrayLike{<:Any,0}, elt::Type, axes::Tuple{})
+function Base.similar(a::BlockSparseArrayLike, elt::Type, axes::Tuple{})
   return blocksparse_similar(a, elt, axes)
 end
 

src/blocksparsearrayinterface/arraylayouts.jl

@@ -1,23 +1,48 @@
-using ArrayLayouts: ArrayLayouts, MatMulMatAdd
+using ArrayLayouts: ArrayLayouts, Dot, MatMulMatAdd, MatMulVecAdd, MulAdd
 using BlockArrays: BlockLayout
 using ..SparseArrayInterface: SparseLayout
-using LinearAlgebra: mul!
+using LinearAlgebra: dot, mul!
 
 function blocksparse_muladd!(
-  α::Number, a1::AbstractMatrix, a2::AbstractMatrix, β::Number, a_dest::AbstractMatrix
+  α::Number, a1::AbstractArray, a2::AbstractArray, β::Number, a_dest::AbstractArray
 )
   mul!(blocks(a_dest), blocks(a1), blocks(a2), α, β)
   return a_dest
 end
 
+function blocksparse_matmul!(m::MulAdd)
+  α, a1, a2, β, a_dest = m.α, m.A, m.B, m.β, m.C
+  blocksparse_muladd!(α, a1, a2, β, a_dest)
+  return a_dest
+end
+
 function ArrayLayouts.materialize!(
   m::MatMulMatAdd{
     <:BlockLayout{<:SparseLayout},
     <:BlockLayout{<:SparseLayout},
     <:BlockLayout{<:SparseLayout},
   },
 )
-  α, a1, a2, β, a_dest = m.α, m.A, m.B, m.β, m.C
-  blocksparse_muladd!(α, a1, a2, β, a_dest)
-  return a_dest
+  blocksparse_matmul!(m)
+  return m.C
+end
+function ArrayLayouts.materialize!(
+  m::MatMulVecAdd{
+    <:BlockLayout{<:SparseLayout},
+    <:BlockLayout{<:SparseLayout},
+    <:BlockLayout{<:SparseLayout},
+  },
+)
+  blocksparse_matmul!(m)
+  return m.C
+end
+
+function blocksparse_dot(a1::AbstractArray, a2::AbstractArray)
+  # TODO: Add a check that the blocking of `a1` and `a2` are
+  # the same, or the same up to a reshape.
+  return dot(blocks(a1), blocks(a2))
+end
+
+function Base.copy(d::Dot{<:BlockLayout{<:SparseLayout},<:BlockLayout{<:SparseLayout}})
+  return blocksparse_dot(d.A, d.B)
 end

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -184,7 +184,7 @@ reverse_index(index::CartesianIndex) = CartesianIndex(reverse(Tuple(index)))
 
 # Represents the array of arrays of a `Transpose`
 # wrapping a block spare array, i.e. `blocks(array)` where `a` is a `Transpose`.
-struct SparseTransposeBlocks{T,BlockType<:AbstractMatrix{T},Array<:Transpose{T}} <:
+struct SparseTransposeBlocks{T,BlockType<:AbstractArray{T},Array<:Transpose{T}} <:
        AbstractSparseMatrix{BlockType}
   array::Array
 end
@@ -219,7 +219,7 @@ end
 
 # Represents the array of arrays of a `Adjoint`
 # wrapping a block spare array, i.e. `blocks(array)` where `a` is a `Adjoint`.
-struct SparseAdjointBlocks{T,BlockType<:AbstractMatrix{T},Array<:Adjoint{T}} <:
+struct SparseAdjointBlocks{T,BlockType<:AbstractArray{T},Array<:Adjoint{T}} <:
        AbstractSparseMatrix{BlockType}
   array::Array
 end

test/test_basics.jl

@@ -16,7 +16,7 @@ using BlockArrays:
   mortar
 using Compat: @compat
 using GPUArraysCore: @allowscalar
-using LinearAlgebra: Adjoint, mul!, norm
+using LinearAlgebra: Adjoint, dot, mul!, norm
 using NDTensors.BlockSparseArrays:
   @view!,
   BlockSparseArray,
@@ -575,7 +575,11 @@ using .NDTensorsTestUtils: devices_list, is_supported_eltype
       a[b] = randn(elt, size(a[b]))
     end
     @test isassigned(a, 1, 1)
+    @test isassigned(a, 1, 1, 1)
+    @test !isassigned(a, 1, 1, 2)
     @test isassigned(a, 5, 7)
+    @test isassigned(a, 5, 7, 1)
+    @test !isassigned(a, 5, 7, 2)
     @test !isassigned(a, 0, 1)
     @test !isassigned(a, 5, 8)
     @test isassigned(a, Block(1), Block(1))
@@ -852,6 +856,16 @@ using .NDTensorsTestUtils: devices_list, is_supported_eltype
       @allowscalar @test Array(a_dest) ≈ Array(a1′) * Array(a2′)
     end
   end
+  @testset "Dot product" begin
+    a1 = dev(BlockSparseArray{elt}([2, 3, 4]))
+    a1[Block(1)] = dev(randn(elt, size(@view(a1[Block(1)]))))
+    a1[Block(3)] = dev(randn(elt, size(@view(a1[Block(3)]))))
+    a2 = dev(BlockSparseArray{elt}([2, 3, 4]))
+    a2[Block(2)] = dev(randn(elt, size(@view(a1[Block(2)]))))
+    a2[Block(3)] = dev(randn(elt, size(@view(a1[Block(3)]))))
+    @test a1' * a2 ≈ Array(a1)' * Array(a2)
+    @test dot(a1, a2) ≈ a1' * a2
+  end
   @testset "TensorAlgebra" begin
     a1 = dev(BlockSparseArray{elt}([2, 3], [2, 3]))
     a1[Block(1, 1)] = dev(randn(elt, size(@view(a1[Block(1, 1)]))))