Define simpler codepath for blockwise map

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

src/blocksparsearrayinterface/map.jl

@@ -1,6 +1,21 @@
+using BlockArrays: BlockRange, blockisequal
 using DerivableInterfaces: @interface, AbstractArrayInterface, interface
 using GPUArraysCore: @allowscalar
 
+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
+    ∪(map(eachblockstoredindex, (a_dest, a_srcs...))...)
+  else
+    BlockRange(a_dest)
+  end
+  for I in Is
+    map!(f, view(a_dest, I), map(Base.Fix2(view, I), a_srcs)...)
+  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 +31,16 @@ using GPUArraysCore: @allowscalar
     @interface interface map_zero_dim!(f, a_dest, a_srcs...)
     return a_dest
   end
+  blockwise = all(
+    ntuple(ndims(a_dest)) do dim
+      ax = map(Base.Fix2(axes, dim), (a_dest, a_srcs...))
+      return blockisequal(ax...)
+    end,
+  )
+  if blockwise
+    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)