[BlockSparseArrays] Sparse and block sparse matrix multiplication (#1278)

Author: mtfishman

src/DiagonalArrays.jl

@@ -6,7 +6,9 @@ include("diaginterface/diagindices.jl")
 include("abstractdiagonalarray/abstractdiagonalarray.jl")
 include("abstractdiagonalarray/sparsearrayinterface.jl")
 include("abstractdiagonalarray/diagonalarraydiaginterface.jl")
+include("abstractdiagonalarray/arraylayouts.jl")
 include("diagonalarray/diagonalarray.jl")
 include("diagonalarray/diagonalmatrix.jl")
 include("diagonalarray/diagonalvector.jl")
+include("diagonalarray/arraylayouts.jl")
 end

src/abstractdiagonalarray/arraylayouts.jl

@@ -0,0 +1,7 @@
+using ArrayLayouts: ArrayLayouts
+using ..SparseArrayInterface: AbstractSparseLayout
+
+abstract type AbstractDiagonalLayout <: AbstractSparseLayout end
+struct DiagonalLayout <: AbstractDiagonalLayout end
+
+ArrayLayouts.MemoryLayout(::Type{<:AbstractDiagonalArray}) = DiagonalLayout()

src/diagonalarray/arraylayouts.jl

@@ -0,0 +1,11 @@
+using ArrayLayouts: MulAdd
+
+# Default sparse array type for `AbstractDiagonalLayout`.
+default_diagonalarraytype(elt::Type) = DiagonalArray{elt}
+
+# TODO: Preserve GPU memory! Implement `CuSparseArrayLayout`, `MtlSparseLayout`?
+function Base.similar(
+  ::MulAdd{<:AbstractDiagonalLayout,<:AbstractDiagonalLayout}, elt::Type, axes
+)
+  return similar(default_diagonalarraytype(elt), axes)
+end

test/runtests.jl

@@ -1,6 +1,8 @@
 @eval module $(gensym())
-using Test: @test, @testset
-using NDTensors.DiagonalArrays: DiagonalArrays
+using Test: @test, @testset, @test_broken
+using NDTensors.DiagonalArrays: DiagonalArrays, DiagonalArray, DiagonalMatrix, diaglength
+using NDTensors.SparseArrayDOKs: SparseArrayDOK
+using NDTensors.SparseArrayInterface: nstored
 @testset "Test NDTensors.DiagonalArrays" begin
   @testset "README" begin
     @test include(
@@ -9,12 +11,49 @@ using NDTensors.DiagonalArrays: DiagonalArrays
       ),
     ) isa Any
   end
-  @testset "Basics" begin
-    using NDTensors.DiagonalArrays: diaglength
-    a = fill(1.0, 2, 3)
-    @test diaglength(a) == 2
-    a = fill(1.0)
-    @test diaglength(a) == 1
+  @testset "DiagonalArray (eltype=$elt)" for elt in
+                                             (Float32, Float64, ComplexF32, ComplexF64)
+    @testset "Basics" begin
+      a = fill(one(elt), 2, 3)
+      @test diaglength(a) == 2
+      a = fill(one(elt))
+      @test diaglength(a) == 1
+    end
+    @testset "Matrix multiplication" begin
+      a1 = DiagonalArray{elt}(undef, (2, 3))
+      a1[1, 1] = 11
+      a1[2, 2] = 22
+      a2 = DiagonalArray{elt}(undef, (3, 4))
+      a2[1, 1] = 11
+      a2[2, 2] = 22
+      a2[3, 3] = 33
+      a_dest = a1 * a2
+      # TODO: Use `densearray` to make generic to GPU.
+      @test Array(a_dest) ≈ Array(a1) * Array(a2)
+      # TODO: Make this work with `ArrayLayouts`.
+      @test nstored(a_dest) == 2
+      @test a_dest isa DiagonalMatrix{elt}
+
+      # TODO: Make generic to GPU, use `allocate_randn`?
+      a2 = randn(elt, (3, 4))
+      a_dest = a1 * a2
+      # TODO: Use `densearray` to make generic to GPU.
+      @test Array(a_dest) ≈ Array(a1) * Array(a2)
+      @test nstored(a_dest) == 8
+      @test a_dest isa Matrix{elt}
+
+      a2 = SparseArrayDOK{elt}(3, 4)
+      a2[1, 1] = 11
+      a2[2, 2] = 22
+      a2[3, 3] = 33
+      a_dest = a1 * a2
+      # TODO: Use `densearray` to make generic to GPU.
+      @test Array(a_dest) ≈ Array(a1) * Array(a2)
+      # TODO: Define `SparseMatrixDOK`.
+      # TODO: Make this work with `ArrayLayouts`.
+      @test nstored(a_dest) == 2
+      @test a_dest isa SparseArrayDOK{elt,2}
+    end
   end
 end
 end