Merge branch 'main' into graded_svd

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.5.2"
+version = "0.5.4"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -30,13 +30,14 @@ using SparseArraysBase:
 
 # A return type for `blocks(array)` when `array` isn't blocked.
 # Represents a vector with just that single block.
-struct SingleBlockView{T,N,Array<:AbstractArray{T,N}} <: AbstractArray{T,N}
+struct SingleBlockView{N,Array<:AbstractArray{<:Any,N}} <: AbstractArray{Array,N}
   array::Array
 end
 Base.parent(a::SingleBlockView) = a.array
+Base.size(a::SingleBlockView) = ntuple(Returns(1), ndims(a))
 blocks_maybe_single(a) = blocks(a)
 blocks_maybe_single(a::Array) = SingleBlockView(a)
-function Base.getindex(a::SingleBlockView{<:Any,N}, index::Vararg{Int,N}) where {N}
+function Base.getindex(a::SingleBlockView{N}, index::Vararg{Int,N}) where {N}
   @assert all(isone, index)
   return parent(a)
 end
@@ -289,13 +290,6 @@ function blockrange(axis::AbstractUnitRange, r::Int)
   return error("Slicing with integer values isn't supported.")
 end
 
-function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Block{1}})
-  for b in r
-    @assert b ∈ blockaxes(axis, 1)
-  end
-  return r
-end
-
 # This handles changing the blocking, for example:
 # a = BlockSparseArray{Float64}([2, 2, 2, 2], [2, 2, 2, 2])
 # I = blockedrange([4, 4])
@@ -314,13 +308,20 @@ end
 # I = BlockedVector([Block(4), Block(3), Block(2), Block(1)], [2, 2])
 # I = BlockVector([Block(4), Block(3), Block(2), Block(1)], [2, 2])
 # a[I, I]
-function blockrange(axis::BlockedOneTo{<:Integer}, r::AbstractBlockVector{<:Block{1}})
+function blockrange(axis::AbstractUnitRange, r::AbstractBlockVector{<:Block{1}})
   for b in r
     @assert b ∈ blockaxes(axis, 1)
   end
   return only(blockaxes(r))
 end
 
+function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Block{1}})
+  for b in r
+    @assert b ∈ blockaxes(axis, 1)
+  end
+  return r
+end
+
 using BlockArrays: BlockSlice
 function blockrange(axis::AbstractUnitRange, r::BlockSlice)
   return blockrange(axis, r.block)
@@ -357,7 +358,11 @@ function blockrange(axis::AbstractUnitRange, r::Base.Slice)
 end
 
 function blockrange(axis::AbstractUnitRange, r::NonBlockedVector)
-  return Block(1):Block(1)
+  return Block.(Base.OneTo(1))
+end
+
+function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Integer})
+  return Block.(Base.OneTo(1))
 end
 
 function blockrange(axis::AbstractUnitRange, r)

src/BlockSparseArrays.jl

@@ -27,6 +27,7 @@ include("abstractblocksparsearray/abstractblocksparsearray.jl")
 include("abstractblocksparsearray/abstractblocksparsematrix.jl")
 include("abstractblocksparsearray/abstractblocksparsevector.jl")
 include("abstractblocksparsearray/wrappedabstractblocksparsearray.jl")
+include("abstractblocksparsearray/unblockedsubarray.jl")
 include("abstractblocksparsearray/views.jl")
 include("abstractblocksparsearray/arraylayouts.jl")
 include("abstractblocksparsearray/sparsearrayinterface.jl")

src/abstractblocksparsearray/arraylayouts.jl

@@ -43,11 +43,20 @@ function ArrayLayouts.sub_materialize(layout::BlockLayout{<:SparseLayout}, a, ax
   return a_dest
 end
 
+function _similar(arraytype::Type{<:AbstractArray}, size::Tuple)
+  return similar(arraytype, size)
+end
+function _similar(
+  ::Type{<:SubArray{<:Any,<:Any,<:ArrayType}}, size::Tuple
+) where {ArrayType}
+  return similar(ArrayType, size)
+end
+
 # Materialize a SubArray view.
 function ArrayLayouts.sub_materialize(
   layout::BlockLayout{<:SparseLayout}, a, axes::Tuple{Vararg{Base.OneTo}}
 )
-  a_dest = blocktype(a)(undef, length.(axes))
+  a_dest = _similar(blocktype(a), length.(axes))
   a_dest .= a
   return a_dest
 end

src/abstractblocksparsearray/unblockedsubarray.jl

@@ -0,0 +1,86 @@
+using ArrayLayouts: ArrayLayouts, MemoryLayout
+using Base.Broadcast: Broadcast, BroadcastStyle
+using BlockArrays: BlockArrays, Block, BlockIndexRange, BlockSlice
+using TypeParameterAccessors: TypeParameterAccessors, parenttype, similartype
+
+const UnblockedIndices = Union{
+  Vector{<:Integer},BlockSlice{<:Block{1}},BlockSlice{<:BlockIndexRange{1}}
+}
+
+const UnblockedSubArray{T,N} = SubArray{
+  T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{UnblockedIndices}}
+}
+
+function BlockArrays.blocks(a::UnblockedSubArray)
+  return SingleBlockView(a)
+end
+
+function DerivableInterfaces.interface(arraytype::Type{<:UnblockedSubArray})
+  return interface(blocktype(parenttype(arraytype)))
+end
+
+function ArrayLayouts.MemoryLayout(arraytype::Type{<:UnblockedSubArray})
+  return MemoryLayout(blocktype(parenttype(arraytype)))
+end
+
+function Broadcast.BroadcastStyle(arraytype::Type{<:UnblockedSubArray})
+  return BroadcastStyle(blocktype(parenttype(arraytype)))
+end
+
+function TypeParameterAccessors.similartype(arraytype::Type{<:UnblockedSubArray}, elt::Type)
+  return similartype(blocktype(parenttype(arraytype)), elt)
+end
+
+function Base.similar(
+  a::UnblockedSubArray, elt::Type, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}}
+)
+  return similar(similartype(blocktype(parenttype(a)), elt), axes)
+end
+function Base.similar(a::UnblockedSubArray, elt::Type, size::Tuple{Int,Vararg{Int}})
+  return similar(a, elt, Base.OneTo.(size))
+end
+
+function ArrayLayouts.sub_materialize(a::UnblockedSubArray)
+  a_cpu = adapt(Array, a)
+  a_cpu′ = similar(a_cpu)
+  a_cpu′ .= a_cpu
+  if typeof(a) === typeof(a_cpu)
+    return a_cpu′
+  end
+  a′ = similar(a)
+  a′ .= a_cpu′
+  return a′
+end
+
+function Base.map!(
+  f, a_dest::AbstractArray, a_src1::UnblockedSubArray, a_src_rest::UnblockedSubArray...
+)
+  return invoke(
+    map!,
+    Tuple{Any,AbstractArray,AbstractArray,Vararg{AbstractArray}},
+    f,
+    a_dest,
+    a_src1,
+    a_src_rest...,
+  )
+end
+
+# Fix ambiguity and scalar indexing errors with GPUArrays.
+using Adapt: adapt
+using GPUArraysCore: GPUArraysCore
+function Base.map!(
+  f,
+  a_dest::GPUArraysCore.AnyGPUArray,
+  a_src1::UnblockedSubArray,
+  a_src_rest::UnblockedSubArray...,
+)
+  a_dest_cpu = adapt(Array, a_dest)
+  a_srcs_cpu = map(adapt(Array), (a_src1, a_src_rest...))
+  map!(f, a_dest_cpu, a_srcs_cpu...)
+  a_dest .= a_dest_cpu
+  return a_dest
+end
+
+function Base.iszero(a::UnblockedSubArray)
+  return invoke(iszero, Tuple{AbstractArray}, adapt(Array, a))
+end

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -364,7 +364,19 @@ end
 function Base.size(a::SparseSubArrayBlocks)
   return length.(axes(a))
 end
-# TODO: Define `isstored`.
+
+# TODO: Make a faster version for when the slice is blockwise.
+function SparseArraysBase.isstored(
+  a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}
+) where {N}
+  J = Base.reindex(parentindices(a.array), to_indices(a.array, Block.(I)))
+  # TODO: Try doing this blockwise when possible rather
+  # than elementwise.
+  return any(Iterators.product(J...)) do K
+    return isstored(parent(a.array), K...)
+  end
+end
+
 # TODO: Define `getstoredindex`, `getunstoredindex` instead.
 function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where {N}
   # TODO: Should this be defined as `@view a.array[Block(I)]` instead?
@@ -400,9 +412,17 @@ function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) whe
   # TODO: Implement this properly.
   return true
 end
-function SparseArraysBase.eachstoredindex(a::SparseSubArrayBlocks)
-  return eachstoredindex(view(blocks(parent(a.array)), blockrange(a)...))
+
+function SparseArraysBase.eachstoredindex(::IndexCartesian, a::SparseSubArrayBlocks)
+  return filter(eachindex(a)) do I
+    return isstored(a, I)
+  end
+
+  ## # TODO: This only works for blockwise slices, i.e. slices using
+  ## # `BlockSliceCollection`.
+  ## return eachstoredindex(view(blocks(parent(a.array)), blockrange(a)...))
 end
+
 # TODO: Either make this the generic interface or define
 # `SparseArraysBase.sparse_storage`, which is used
 # to defined this.
@@ -425,7 +445,7 @@ end
 
 to_blocks_indices(I::BlockSlice{<:BlockRange{1}}) = Int.(I.block)
 to_blocks_indices(I::BlockIndices{<:Vector{<:Block{1}}}) = Int.(I.blocks)
-to_blocks_indices(I::Base.Slice{<:BlockedOneTo}) = Base.OneTo(blocklength(I.indices))
+to_blocks_indices(I::Base.Slice) = Base.OneTo(blocklength(I.indices))
 
 @interface ::AbstractBlockSparseArrayInterface function BlockArrays.blocks(
   a::SubArray{<:Any,<:Any,<:Any,<:Tuple{Vararg{BlockSliceCollection}}}

src/blocksparsearrayinterface/map.jl

@@ -1,6 +1,52 @@
+using BlockArrays: BlockRange, blockisequal
 using DerivableInterfaces: @interface, AbstractArrayInterface, interface
 using GPUArraysCore: @allowscalar
 
+# Check if the block structures are the same.
+function same_block_structure(as::AbstractArray...)
+  isempty(as) && return true
+  return all(
+    ntuple(ndims(first(as))) do dim
+      ax = map(Base.Fix2(axes, dim), as)
+      return blockisequal(ax...)
+    end,
+  )
+end
+
+# Find the common stored blocks, assuming the block structures are the same.
+function union_eachblockstoredindex(as::AbstractArray...)
+  return ∪(map(eachblockstoredindex, as)...)
+end
+
+function map_blockwise!(f, a_dest::AbstractArray, a_srcs::AbstractArray...)
+  # TODO: This assumes element types are numbers, generalize this logic.
+  f_preserves_zeros = f(zero.(eltype.(a_srcs))...) == zero(eltype(a_dest))
+  Is = if f_preserves_zeros
+    union_eachblockstoredindex(a_dest, a_srcs...)
+  else
+    BlockRange(a_dest)
+  end
+  for I in Is
+    # TODO: Use:
+    # block_dest = @view a_dest[I]
+    # or:
+    # block_dest = @view! a_dest[I]
+    block_dest = blocks_maybe_single(a_dest)[Int.(Tuple(I))...]
+    # TODO: Use:
+    # block_srcs = map(a_src -> @view(a_src[I]), a_srcs)
+    block_srcs = map(a_srcs) do a_src
+      return blocks_maybe_single(a_src)[Int.(Tuple(I))...]
+    end
+    # TODO: Use `map!!` to handle immutable blocks.
+    map!(f, block_dest, block_srcs...)
+    # Replace the entire block, handles initializing new blocks
+    # or if blocks are immutable.
+    # TODO: Use `a_dest[I] = block_dest`.
+    blocks(a_dest)[Int.(Tuple(I))...] = block_dest
+  end
+  return a_dest
+end
+
 # TODO: Rewrite this so that it takes the blocking structure
 # made by combining the blocking of the axes (i.e. the blocking that
 # is used to determine `union_stored_blocked_cartesianindices(...)`).
@@ -16,6 +62,10 @@ using GPUArraysCore: @allowscalar
     @interface interface map_zero_dim!(f, a_dest, a_srcs...)
     return a_dest
   end
+  if same_block_structure(a_dest, a_srcs...)
+    map_blockwise!(f, a_dest, a_srcs...)
+    return a_dest
+  end
   # TODO: This assumes element types are numbers, generalize this logic.
   f_preserves_zeros = f(zero.(eltype.(a_srcs))...) == zero(eltype(a_dest))
   a_dest, a_srcs = reblock(a_dest), reblock.(a_srcs)

test/test_basics.jl

@@ -51,13 +51,6 @@ arrayts = (Array, JLArray)
     a[Block(2, 2)] = dev(randn(elt, 3, 3))
     @test_broken a[:, 4]
 
-    # TODO: Fix this and turn it into a proper test.
-    a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3]))
-    a[Block(1, 1)] = dev(randn(elt, 2, 2))
-    a[Block(2, 2)] = dev(randn(elt, 3, 3))
-    @test_broken a[:, [2, 4]]
-    @test_broken a[[3, 5], [2, 4]]
-
     # TODO: Fix this and turn it into a proper test.
     a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3]))
     a[Block(1, 1)] = dev(randn(elt, 2, 2))
@@ -713,6 +706,20 @@ arrayts = (Array, JLArray)
     @test a[Block(2, 2)[1:2, 2:3]] == b
     @test blockstoredlength(a) == 1
 
+    # Noncontiguous slicing.
+    a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3]))
+    a[Block(1, 1)] = dev(randn(elt, 2, 2))
+    a[Block(2, 2)] = dev(randn(elt, 3, 3))
+    I = ([3, 5], [2, 4])
+    @test Array(a[I...]) == Array(a)[I...]
+
+    # Noncontiguous slicing.
+    a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3]))
+    a[Block(1, 1)] = dev(randn(elt, 2, 2))
+    a[Block(2, 2)] = dev(randn(elt, 3, 3))
+    I = (:, [2, 4])
+    @test Array(a[I...]) == Array(a)[I...]
+
     a = BlockSparseArray{elt}(undef, [2, 3], [2, 3])
     @views for b in [Block(1, 1), Block(2, 2)]
       a[b] = randn(elt, size(a[b]))