Define simpler codepath for blockwise map

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.3"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -290,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])
@@ -315,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)

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -445,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)