Functionality for slicing with unit ranges that preserves block information

[mtfishman]

Indexing with UnitRange does not preserve the blocking of blocked arrays:

julia> using BlockArrays

julia> r = blockedrange([2, 4])
2-blocked 6-element BlockedUnitRange{Int64, Vector{Int64}}:
 1
 2
 ─
 3
 4
 5
 6

julia> r[2:4]
2:4

and:

julia> a = BlockArray(randn(4, 4), [2, 2], [2, 2])
2×2-blocked 4×4 BlockMatrix{Float64}:
 0.372344   -0.735109  │   0.316133  -0.935035
 2.37641     1.30943   │  -0.931113   0.339143
 ──────────────────────┼──────────────────────
 0.422144    0.839572  │  -1.4721    -0.74205 
 0.0436636  -0.722946  │   0.799941  -0.192523

julia> a[2:3, 2:3]
2×2 Matrix{Float64}:
 1.30943   -0.931113
 0.839572  -1.4721

It would be nice to have convenient functionality for indexing with UnitRange that does preserve blocking.

[jishnub]

Generally, the result of an indexing operation has the same axes as that of the indices. This lets us use results such as A[ax][i] == A[ax[i]], which are quite convenient. Perhaps we should enforce this more rigorously, so that indexing with a blocked unitrange would lead to a blocked result.

[mtfishman]

@jishnub that seems like a reasonable rule, but to clarify do you think blockedrange([2, 4])[2:4] should output a BlockedUnitRange and BlockArray(randn(4, 4), [2, 2], [2, 2])[2:4, 2:4] should output a BlockArray? I think that follows from your proposal A[ax][i] == A[ax[i]].

blockedrange([2, 4])[2:4]::BlockedUnitRange doesn't exactly follow that "the result of an indexing operation has the same axes as that of the indices" which I think is a good guideline to follow (both for type stability and as a design philosophy), but I think it is also ok if the axes pick up some information from the axes of the object that is being indexed into, in this case some blocking information.

[mtfishman]

Responding to https://github.com/JuliaArrays/BlockArrays.jl/pull/356#issuecomment-2024751896.

It seems like we could have both worlds, where the rule is that if you slice with a UnitRange, it picks up the blocking information based on slicing the axes of the block array, and then if you slice with a BlockedUnitRange, then it changes the blocking of the block array based on the blocks of the slice. I realize there is a slight inconsistency there because a UnitRange right now is treated as a BlockedUnitRange with a single block, but it seems reasonable to make UnitRange act in a special way for the sake of slicing (as a kind of "blank slate" with potentially any kind of blocking).

The use case I have in mind is designing a BlockSparseArray type based around the BlockArrays interface, as an AbstractBlockArray subtype. We have a low density of non-zero blocks (thus the need for a BlockSparseArray type), and also our plan is to potentially have a mixture blocks of blocks on CPU or GPU, across multiple CPUs and GPUs, etc. I would like to follow the BlockArrays interface as closely as possible when designing our type, and therefore would like to keep the slicing behavior consistent. However, I think the current convention that slicing with UnitRange doesn't preserve the block structure would make that operation pretty useless and very surprising for that application. If a slicing operation like B[1:end, 1:end] on a BlockSparseArray with blocks distributed across multiple CPUs and GPUs is supposed to instantiate a dense array, where should that dense array live? Should it be instantiated on CPU, GPU, etc., and if so which one? It could also mean that a seemingly benign slicing operation like B[1:end, 1:end] leads us to instantiate a dense array that can't fit into memory.

The alternative is that if we want to perform a slice with a UnitRange and preserve the blocking structure, we would have to convert it to a BlockedUnitRange ourselves, however there doesn't seem to be a convenient way to do that with BlockArrays right now.

[jishnub]

To me, it seems like we need two things here:

1. We need an easy way to convert a UnitRange to a BlockedUnitRange with specified block sizes. blockedrange is almost that function, but we need a simpler interface.
2. We need a macro (tentatively named @blocked) that acts on an indexing operation and translates the block size information from the parent array to the indices. This way, @blocked A[1:3, 1:3] would create a blocked array by retaining the block sizes and types from A, even when A[1:3, 1:3] would discard the block sizes of A.

Meanwhile, we may also be able to improve the default indexing behavior of an indexing operation like A[1:3, 1:3] where the blocks of A are sparse. Currently, we always allocate a dense array, but we may do better by concatenating the blocks instead, which should preserve sparsity.

[mtfishman]

That could be a good way to go. Then @blocked blockedrange([3, 4])[2:5] would be a convenient way to make blockedrange(2, [2, 2]), in the notation of #348. So @blocked B[2:5, 2:5] performs @blocked slices of the axes of B, and then uses that blocking.

I think if I was designing things from scratch I would make it the other way around, and have @blocked as the default behavior of B[2:5, 2:5], and some other notation for the version of B[2:5, 2:5] that flattens the block structure (@flatten B[2:5, 2:5] or @reaxis B[2:5, 2:5]?), since I imagine using @blocked B[2:5, 2:5] a lot more than B[2:5, 2:5]. Of course, that's how I feel about the choice of @view B[2:5, 2:5] vs. B[2:5, 2:5] in Base Julia. But I see why it might be designed that way in the interest of using slicing as a way to change the blocking of a block array, and biasing towards preserving information about the input indices. Also a compelling reason is to have B[2:5, 2:5] and B[Vector(2:5), Vector(2:5)] perform the same operation.

I suppose @blocked will also require a corresponding view version, perhaps called @blockedview B[r, r]. I guess that can still construct a SubArray, but the axes of that SubArray would be constructed from @blocked slices of the axes of B, i.e. @blocked axes(B, i)[r].

Agreed that a generic way to implement B[1:3, 1:3] would be through concatenation, in which case in the case of mixed CPU and GPU blocks, rules about concatenation of CPU and GPU arrays would apply (whatever they are, I assume that just isn't defined so B[1:3, 1:3] would error).

[mtfishman]

@jishnub what do you think @blocked should do when indexing into a blocked array with mismatched blocked indices? For example:

julia> using BlockArrays

julia> a = BlockArray(randn(4, 4), [2, 2], [2, 2])
2×2-blocked 4×4 BlockMatrix{Float64}:
  0.381752   0.999573  │   0.861305  -0.114894
 -0.4414    -1.44202   │   0.906593  -0.383571
 ──────────────────────┼──────────────────────
  0.575925  -1.25182   │   0.672723  -1.18554 
  0.480034   0.899287  │  -0.519823  -1.36788 

julia> r = blockedrange([1, 2, 1])
3-blocked 4-element BlockedUnitRange{Vector{Int64}}:
 1
 ─
 2
 3
 ─
 4

julia> @blocked a[r, r]
# ...

Should it preserve the blocks of a, so equivalent to:

julia> a[blockedrange([2, 2]), blockedrange([2, 2])]
2×2-blocked 4×4 BlockMatrix{Float64}:
  0.381752   0.999573  │   0.861305  -0.114894
 -0.4414    -1.44202   │   0.906593  -0.383571
 ──────────────────────┼──────────────────────
  0.575925  -1.25182   │   0.672723  -1.18554 
  0.480034   0.899287  │  -0.519823  -1.36788 

or should it combine the blocking of the input indices and the axes of a, i.e. equivalent to:

julia> a[blockedrange([1, 1, 1, 1]), blockedrange([1, 1, 1, 1])]
4×4-blocked 4×4 BlockMatrix{Float64}:
  0.381752  │   0.999573  │   0.861305  │  -0.114894
 ───────────┼─────────────┼─────────────┼───────────
 -0.4414    │  -1.44202   │   0.906593  │  -0.383571
 ───────────┼─────────────┼─────────────┼───────────
  0.575925  │  -1.25182   │   0.672723  │  -1.18554 
 ───────────┼─────────────┼─────────────┼───────────
  0.480034  │   0.899287  │  -0.519823  │  -1.36788 

which could be based on BlockArrays.combine_blockaxes:

julia> foreach(display, BlockArrays.combine_blockaxes.((r, r), axes(a)))
4-blocked 4-element BlockedUnitRange{Vector{Int64}}:
 1
 ─
 2
 ─
 3
 ─
 4
4-blocked 4-element BlockedUnitRange{Vector{Int64}}:
 1
 ─
 2
 ─
 3
 ─
 4

I'm asking because I may start implementing @blocked privately for my own AbstractBlockArrays and see how the design works out, and I'd like to align it with how you might want it designed.

[jishnub]

I think it makes sense to combine the blocks. This way, we obtain:

• @blocked blocked array[non-blocked indices]: blocksizes of the parent
• @blocked blocked array[blocked indices]: combined blocksizes of the parent and the indices
• @blocked non-blocked array[blocked indices]: blocksizes of the indices (consistent with the standard result without @blocked. Otherwise, this wouldn't be blocked at all)

The alternative doesn't provide an easy way to combine the blocks, which may be useful as well. Since it's usually easy to go from a blocked to a non-blocked index (e.g. UnitRange(blocked_index)), it should be straightforward to go from the second case to the first.

[mtfishman]

Makes sense to me.