[BlockSparseArrays] GPU support

Author: mtfishman

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/abstractblocksparsearray.jl

@@ -12,7 +12,15 @@ abstract type AbstractBlockSparseArray{T,N} <: AbstractBlockArray{T,N} end
 
 Base.axes(::AbstractBlockSparseArray) = error("Not implemented")
 
-blockstype(::Type{<:AbstractBlockSparseArray}) = error("Not implemented")
+# TODO: Add some logic to unwrapping wrapped arrays.
+# TODO: Decide what a good default is.
+blockstype(arraytype::Type{<:AbstractBlockSparseArray}) = SparseArrayDOK{AbstractArray}
+function blockstype(arraytype::Type{<:AbstractBlockSparseArray{T}}) where {T}
+  return SparseArrayDOK{AbstractArray{T}}
+end
+function blockstype(arraytype::Type{<:AbstractBlockSparseArray{T,N}}) where {T,N}
+  return SparseArrayDOK{AbstractArray{T,N},N}
+end
 
 ## # Specialized in order to fix ambiguity error with `BlockArrays`.
 function Base.getindex(a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Int,N}) where {N}

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/arraylayouts.jl

@@ -1,6 +1,8 @@
 using ArrayLayouts: ArrayLayouts, MemoryLayout, MulAdd
 using BlockArrays: BlockLayout
 using ..SparseArrayInterface: SparseLayout
+# TODO: Move to `NDTensors.TypeParameterAccessors`.
+using ..NDTensors: similartype
 
 function ArrayLayouts.MemoryLayout(arraytype::Type{<:BlockSparseArrayLike})
   outer_layout = typeof(MemoryLayout(blockstype(arraytype)))
@@ -9,9 +11,14 @@ function ArrayLayouts.MemoryLayout(arraytype::Type{<:BlockSparseArrayLike})
 end
 
 function Base.similar(
-  ::MulAdd{<:BlockLayout{<:SparseLayout},<:BlockLayout{<:SparseLayout}}, elt::Type, axes
-)
-  return similar(BlockSparseArray{elt}, axes)
+  mul::MulAdd{<:BlockLayout{<:SparseLayout},<:BlockLayout{<:SparseLayout},<:Any,<:Any,A,B},
+  elt::Type,
+  axes,
+) where {A,B}
+  # TODO: Check that this equals `similartype(blocktype(B), elt, axes)`,
+  # or maybe promote them?
+  output_blocktype = similartype(blocktype(A), elt, axes)
+  return similar(BlockSparseArray{elt,length(axes),output_blocktype}, axes)
 end
 
 # Materialize a SubArray view.

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -9,6 +9,8 @@ using BlockArrays:
   mortar,
   unblock
 using SplitApplyCombine: groupcount
+# TODO: Move to `NDTensors.TypeParameterAccessors`.
+using ..NDTensors: similartype
 
 const WrappedAbstractBlockSparseArray{T,N} = WrappedArray{
   T,N,AbstractBlockSparseArray,AbstractBlockSparseArray{T,N}
@@ -187,28 +189,27 @@ function Base.similar(
   return similar(arraytype, eltype(arraytype), axes)
 end
 
+function blocksparse_similar(a, elt::Type, axes::Tuple)
+  return BlockSparseArray{elt,length(axes),similartype(blocktype(a), axes)}(undef, axes)
+end
+
 # Needed by `BlockArrays` matrix multiplication interface
 # TODO: Define a `blocksparse_similar` function.
 function Base.similar(
   arraytype::Type{<:BlockSparseArrayLike},
   elt::Type,
   axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}},
 )
-  # TODO: Make generic for GPU, maybe using `blocktype`.
-  # TODO: For non-block axes this should output `Array`.
-  return BlockSparseArray{elt}(undef, axes)
+  return blocksparse_similar(arraytype, elt, axes)
 end
 
 # TODO: Define a `blocksparse_similar` function.
 function Base.similar(
   a::BlockSparseArrayLike, elt::Type, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}
 )
-  # TODO: Make generic for GPU, maybe using `blocktype`.
-  # TODO: For non-block axes this should output `Array`.
-  return BlockSparseArray{elt}(undef, axes)
+  return blocksparse_similar(a, elt, axes)
 end
 
-# TODO: Define a `blocksparse_similar` function.
 # Fixes ambiguity error with `BlockArrays`.
 function Base.similar(
   a::BlockSparseArrayLike,
@@ -217,21 +218,16 @@ function Base.similar(
     AbstractBlockedUnitRange{<:Integer},Vararg{AbstractBlockedUnitRange{<:Integer}}
   },
 )
-  # TODO: Make generic for GPU, maybe using `blocktype`.
-  # TODO: For non-block axes this should output `Array`.
-  return BlockSparseArray{elt}(undef, axes)
+  return blocksparse_similar(a, elt, axes)
 end
 
-# TODO: Define a `blocksparse_similar` function.
 # Fixes ambiguity error with `OffsetArrays`.
 function Base.similar(
   a::BlockSparseArrayLike,
   elt::Type,
   axes::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
 )
-  # TODO: Make generic for GPU, maybe using `blocktype`.
-  # TODO: For non-block axes this should output `Array`.
-  return BlockSparseArray{elt}(undef, axes)
+  return blocksparse_similar(a, elt, axes)
 end
 
 # Fixes ambiguity error with `BlockArrays`.
@@ -240,9 +236,7 @@ function Base.similar(
   elt::Type,
   axes::Tuple{AbstractBlockedUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
 )
-  # TODO: Make generic for GPU, maybe using `blocktype`.
-  # TODO: For non-block axes this should output `Array`.
-  return BlockSparseArray{elt}(undef, axes)
+  return blocksparse_similar(a, elt, axes)
 end
 
 # Fixes ambiguity errors with BlockArrays.
@@ -255,15 +249,12 @@ function Base.similar(
     Vararg{AbstractUnitRange{<:Integer}},
   },
 )
-  return BlockSparseArray{elt}(undef, axes)
+  return blocksparse_similar(a, elt, axes)
 end
 
-# TODO: Define a `blocksparse_similar` function.
 # Fixes ambiguity error with `StaticArrays`.
 function Base.similar(
   a::BlockSparseArrayLike, elt::Type, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}}
 )
-  # TODO: Make generic for GPU, maybe using `blocktype`.
-  # TODO: For non-block axes this should output `Array`.
-  return BlockSparseArray{elt}(undef, axes)
+  return blocksparse_similar(a, elt, axes)
 end

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearray/blocksparsearray.jl

@@ -35,6 +35,18 @@ function BlockSparseArray(
   return BlockSparseArray(Dictionary(block_indices, block_data), axes)
 end
 
+function BlockSparseArray{T,N,A,Blocks}(
+  blocks::AbstractArray{<:AbstractArray{T,N},N}, axes::Tuple{Vararg{AbstractUnitRange,N}}
+) where {T,N,A<:AbstractArray{T,N},Blocks<:AbstractArray{A,N}}
+  return BlockSparseArray{T,N,A,Blocks,typeof(axes)}(blocks, axes)
+end
+
+function BlockSparseArray{T,N,A}(
+  blocks::AbstractArray{<:AbstractArray{T,N},N}, axes::Tuple{Vararg{AbstractUnitRange,N}}
+) where {T,N,A<:AbstractArray{T,N}}
+  return BlockSparseArray{T,N,A,typeof(blocks)}(blocks, axes)
+end
+
 function BlockSparseArray{T,N}(
   blocks::AbstractArray{<:AbstractArray{T,N},N}, axes::Tuple{Vararg{AbstractUnitRange,N}}
 ) where {T,N}
@@ -49,9 +61,15 @@ function BlockSparseArray{T,N}(
   return BlockSparseArray{T,N}(blocks, axes)
 end
 
+function BlockSparseArray{T,N,A}(
+  axes::Tuple{Vararg{AbstractUnitRange,N}}
+) where {T,N,A<:AbstractArray{T,N}}
+  blocks = default_blocks(A, axes)
+  return BlockSparseArray{T,N,A}(blocks, axes)
+end
+
 function BlockSparseArray{T,N}(axes::Tuple{Vararg{AbstractUnitRange,N}}) where {T,N}
-  blocks = default_blocks(T, axes)
-  return BlockSparseArray{T,N}(blocks, axes)
+  return BlockSparseArray{T,N,default_arraytype(T, axes)}(axes)
 end
 
 function BlockSparseArray{T,N}(dims::Tuple{Vararg{Vector{Int},N}}) where {T,N}
@@ -74,6 +92,12 @@ function BlockSparseArray{T}(axes::Vararg{AbstractUnitRange}) where {T}
   return BlockSparseArray{T}(axes)
 end
 
+function BlockSparseArray{T,N,A}(
+  ::UndefInitializer, dims::Tuple
+) where {T,N,A<:AbstractArray{T,N}}
+  return BlockSparseArray{T,N,A}(dims)
+end
+
 # undef
 function BlockSparseArray{T,N}(
   ::UndefInitializer, axes::Tuple{Vararg{AbstractUnitRange,N}}
@@ -109,7 +133,23 @@ Base.axes(a::BlockSparseArray) = a.axes
 blocksparse_blocks(a::BlockSparseArray) = a.blocks
 
 # TODO: Use `TypeParameterAccessors`.
-blockstype(::Type{<:BlockSparseArray{<:Any,<:Any,<:Any,B}}) where {B} = B
+function blockstype(
+  arraytype::Type{<:BlockSparseArray{T,N,A,Blocks}}
+) where {T,N,A<:AbstractArray{T,N},Blocks<:AbstractArray{A,N}}
+  return Blocks
+end
+function blockstype(
+  arraytype::Type{<:BlockSparseArray{T,N,A}}
+) where {T,N,A<:AbstractArray{T,N}}
+  return SparseArrayDOK{A,N}
+end
+function blockstype(arraytype::Type{<:BlockSparseArray{T,N}}) where {T,N}
+  return SparseArrayDOK{AbstractArray{T,N},N}
+end
+function blockstype(arraytype::Type{<:BlockSparseArray{T}}) where {T}
+  return SparseArrayDOK{AbstractArray{T}}
+end
+blockstype(arraytype::Type{<:BlockSparseArray}) = SparseArrayDOK{AbstractArray}
 
 ## # Base interface
 ## function Base.similar(

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearray/defaults.jl

@@ -29,8 +29,8 @@ function default_arraytype(elt::Type, axes::Tuple{Vararg{AbstractUnitRange}})
   return Array{elt,length(axes)}
 end
 
-function default_blocks(elt::Type, axes::Tuple{Vararg{AbstractUnitRange}})
-  block_data = Dictionary{Block{length(axes),Int},default_arraytype(elt, axes)}()
+function default_blocks(blocktype::Type, axes::Tuple{Vararg{AbstractUnitRange}})
+  block_data = Dictionary{Block{length(axes),Int},blocktype}()
   return default_blocks(block_data, axes)
 end