Optimizing the block-style broadcasting (#66)

* Optimize block array broadcast

* Define unsafe_view for BlockIndexRange

* Flatten Broadcasted

* Allow `ones(2)[Block(1)[1:1], Block(1)[1:1]]`

* Disambiguate unsafe_view(::SubArray, ::BlockSlice...)

* Use BlockSlice's indices in unsafe_view

* Restrict an unsafe_view definition to BlockArray

Author: tkf

src/BlockArrays.jl

@@ -23,7 +23,7 @@ import Base: @propagate_inbounds, Array, to_indices, to_index,
             RangeIndex, Int, Integer, Number,
             +, -, min, max, *, isless, in, copy, copyto!, axes, @deprecate,
             BroadcastStyle
-using Base: dataids
+using Base: ReshapedArray, dataids
 
 
 

src/blockbroadcast.jl

@@ -25,11 +25,6 @@ BroadcastStyle(::BlockStyle{M}, ::PseudoBlockStyle{N}) where {M,N} = BlockStyle(
 BroadcastStyle(::PseudoBlockStyle{M}, ::BlockStyle{N}) where {M,N} = BlockStyle(Val(max(M,N)))
 
 
-####
-# Default to standard Array broadcast
-####
-
-
 # following code modified from julia/base/broadcast.jl
 broadcast_cumulsizes(::Number) = ()
 broadcast_cumulsizes(A::AbstractArray) = cumulsizes(blocksizes(A))
@@ -48,11 +43,84 @@ blocksizes(A::Broadcasted{<:AbstractArrayStyle{N}}) where N =
     BlockSizes(combine_cumulsizes(broadcast_cumulsizes.(A.args)...))
 
 
-copyto!(dest::AbstractArray, bc::Broadcasted{<:AbstractBlockStyle{N}}) where N =
-   copyto!(dest, Broadcasted{DefaultArrayStyle{N}}(bc.f, bc.args, bc.axes))
-
 similar(bc::Broadcasted{<:AbstractBlockStyle{N}}, ::Type{T}) where {T,N} =
     BlockArray{T,N}(undef, blocksizes(bc))
 
 similar(bc::Broadcasted{PseudoBlockStyle{N}}, ::Type{T}) where {T,N} =
     PseudoBlockArray{T,N}(undef, blocksizes(bc))
+
+
+subblocks(::Any, bs::BlockSizes, dim::Integer) =
+    (nothing for _ in 1:nblocks(bs, dim))
+
+function subblocks(arr::AbstractArray, bs::BlockSizes, dim::Integer)
+    if size(arr, dim) == 1
+        return (BlockIndexRange(Block(1), 1:1) for _ in 1:nblocks(bs, dim))
+    end
+    j = 1
+    next = 1
+    arrstops = cumulsizes(arr, dim)
+    return (
+        let n = blocksize(bs, dim, i)
+            start = next
+            next = start + n
+            j0 = j
+            if arrstops[j + 1] == next
+                j += 1
+            end
+            BlockIndexRange(Block(j0), (start:next - 1) .- (arrstops[j0] - 1))
+        end
+        for i in 1:nblocks(bs, dim))
+end
+
+@inline _bview(arg, ::Vararg) = arg
+@inline _bview(A::AbstractArray, I...) = view(A, I...)
+
+@generated function copyto!(
+        dest::AbstractArray,
+        bc::Broadcasted{<:AbstractBlockStyle{NDims}, <:Any, <:Any, Args},
+        ) where {NDims, Args <: Tuple}
+
+    NArgs = length(Args.parameters)
+
+    # `bvar(0, dim)` is a variable for BlockIndexRange of `dim`-th dimension
+    # of `dest` array.  `bvar(i, dim)` is a similar variable of `i`-th
+    # argument in `bc.args`.
+    bvar(i, dim) = Symbol("blockindexrange_", i, "_", dim)
+
+    function forloop(dim)
+        if dim > 0
+            quote
+                for ($(bvar(0, dim)), $(bvar.(1:NArgs, dim)...),) in zip(
+                        subblocks(dest, bs, $dim),
+                        subblocks.(bc.args, Ref(bs), Ref($dim))...)
+                    $(forloop(dim - 1))
+                end
+            end
+        else
+            bview(a, i) = :(_bview($a, $([bvar(i, d) for d in 1:NDims]...)))
+            destview = bview(:dest, 0)
+            argblocks = [bview(:(bc.args[$i]), i) for i in 1:NArgs]
+            quote
+                broadcast!(bc.f, $destview, $(argblocks...))
+            end
+        end
+    end
+
+    quote
+        bs = blocksizes(bc)
+        if blocksizes(dest) ≠ bs
+            copyto!(PseudoBlockArray(dest, bs), bc)
+            return dest
+        end
+
+        $(forloop(NDims))
+        return dest
+    end
+end
+
+@inline function Broadcast.instantiate(
+        bc::Broadcasted{Style}) where {Style <:AbstractBlockStyle}
+    bcf = Broadcast.flatten(Broadcasted{Nothing}(bc.f, bc.args, bc.axes))
+    return Broadcasted{Style}(bcf.f, bcf.args, bcf.axes)
+end

src/blockindexrange.jl

@@ -79,7 +79,45 @@ reindex(V, idxs::Tuple{BlockSlice{<:BlockRange}, Vararg{Any}},
                                             idxs[1].indices[subidxs[1].indices]),
                                     reindex(V, tail(idxs), tail(subidxs))...))
 
+# De-reference blocks before creating a view to avoid taking `global2blockindex`
+# path in `AbstractBlockStyle` broadcasting.
+@inline function Base.unsafe_view(
+        A::BlockArray{<:Any, N},
+        I::Vararg{BlockSlice{<:BlockIndexRange{1}}, N}) where {N}
+    @_propagate_inbounds_meta
+    B = A[map(x -> x.block.block, I)...]
+    return view(B, _splatmap(x -> x.block.indices, I)...)
+end
+
+@inline function Base.unsafe_view(
+        A::PseudoBlockArray{<:Any, N},
+        I::Vararg{BlockSlice{<:BlockIndexRange{1}}, N}) where {N}
+    @_propagate_inbounds_meta
+    return view(A.blocks, map(x -> x.indices, I)...)
+end
+
+@inline function Base.unsafe_view(
+        A::ReshapedArray{<:Any, N, <:AbstractBlockArray{<:Any, M}},
+        I::Vararg{BlockSlice{<:BlockIndexRange{1}}, N}) where {N, M}
+    @_propagate_inbounds_meta
+    # Note: assuming that I[M+1:end] are verified to be singletons
+    return reshape(view(A.parent, I[1:M]...), Val(N))
+end
+
+@inline function Base.unsafe_view(
+        A::AbstractArray{<:Any, N},
+        I::Vararg{BlockSlice{<:BlockIndexRange{1}}, N}) where {N}
+    @_propagate_inbounds_meta
+    return view(A, map(x -> x.indices, I)...)
+end
 
+# Disambiguation
+@inline function Base.unsafe_view(
+        A::SubArray,
+        I::Vararg{BlockSlice{<:BlockIndexRange{1}}, N}) where {N}
+    @_propagate_inbounds_meta
+    return view(A, map(x -> x.indices, I)...)
+end
 
 
 # #################

src/views.jl

@@ -71,8 +71,15 @@ blocksizes(V::SubArray) = BlockSizes(_sub_cumul_sizes(cumulsizes(parent(V)), par
 
 Returns the indices associated with a block as a `BlockSlice`.
 """
-unblock(A::AbstractArray{T,N}, inds, I) where {T, N} = _unblock(cumulsizes(A, N - length(inds) + 1), I)
-
+function unblock(A::AbstractArray{T,N}, inds, I) where {T, N}
+    if length(inds) == 0
+        # Allow `ones(2)[Block(1)[1:1], Block(1)[1:1]]` which is
+        # similar to `ones(2)[1:1, 1:1]`.
+        _unblock(Base.OneTo(2), I)
+    else
+        _unblock(cumulsizes(A, N - length(inds) + 1), I)
+    end
+end
 
 
 to_index(::Block) = throw(ArgumentError("Block must be converted by to_indices(...)"))