[BlockSparseArrays] Define in-place view that may instantiate blocks

[mtfishman]

To-do:

• Investigate issue with too many blocks being allocated when performing block sparse matrix multiplication.

As discussed in ITensor/BlockSparseArrays.jl#2. For example:

julia> using BlockArrays: Block

julia> using NDTensors.BlockSparseArrays: BlockSparseArray, @view!

julia> a = BlockSparseArray{Float64}([2, 3], [2, 3])
typeof(axes) = Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}

Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.

2×2-blocked 5×5 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}}, Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}:
 0.0  0.0  │  0.0  0.0  0.0
 0.0  0.0  │  0.0  0.0  0.0
 ──────────┼───────────────
 0.0  0.0  │  0.0  0.0  0.0
 0.0  0.0  │  0.0  0.0  0.0
 0.0  0.0  │  0.0  0.0  0.0

julia> b = @view! a[Block(2, 2)] # or: `b = view!(a, Block(2, 2))`
3×3 Matrix{Float64}:
 0.0  0.0  0.0
 0.0  0.0  0.0
 0.0  0.0  0.0

julia> b[2, 2] = 2
2

julia> a
typeof(axes) = Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}

Warning: To temporarily circumvent a bug in printing BlockSparseArrays with mixtures of dual and non-dual axes, the types of the dual axes printed below might not be accurate. The types printed above this message are the correct ones.

2×2-blocked 5×5 BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}}, Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}:
 0.0  0.0  │  0.0  0.0  0.0
 0.0  0.0  │  0.0  0.0  0.0
 ──────────┼───────────────
 0.0  0.0  │  0.0  0.0  0.0
 0.0  0.0  │  0.0  2.0  0.0
 0.0  0.0  │  0.0  0.0  0.0

julia> b
3×3 Matrix{Float64}:
 0.0  0.0  0.0
 0.0  2.0  0.0
 0.0  0.0  0.0

This can be better to work with than view(a, Block(2, 2))/@view a[Block(2, 2)] since as of #1481, view(::BlockSparseArray, ...) creates a SubArray wrapper around the entire BlockSparseArray instead of returning the block data directly, in order to uniformly handle stored and unstored blocks:

julia> @view a[Block(2, 2)]
3×3 view(::BlockSparseArray{Float64, 2, Matrix{Float64}, NDTensors.SparseArrayDOKs.SparseArrayDOK{Matrix{Float64}, 2, NDTensors.BlockSparseArrays.BlockZero{Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}}, Tuple{BlockArrays.BlockedOneTo{Int64, Vector{Int64}}, BlockArrays.BlockedOneTo{Int64, Vector{Int64}}}}, BlockSlice(Block(2),3:5), BlockSlice(Block(2),3:5)) with eltype Float64:
 0.0  0.0  0.0
 0.0  2.0  0.0
 0.0  0.0  0.0

Also note that sub-slicing of blocks is also supported with the syntax @view! a[Block(2, 2)[1:2, 1:2]] or view!(a, Block(2, 2)[1:2, 1:2]), in which case it outputs a view of the block data (either the existing block data or the newly allocated block data):

julia> @view! a[Block(2, 2)[1:2, 1:2]]
2×2 view(::Matrix{Float64}, 1:2, 1:2) with eltype Float64:
 0.0  0.0
 0.0  2.0

@ogauthe @emstoudenmire

[mtfishman]

@ogauthe hopefully by using view!/@view! that can circumvent the issues you were seeing passing blocks/block slices to TensorAlgebra.contract and Strided.@strided (which for future reference are listed in ITensor/BlockSparseArrays.jl#2). Of course it would be ideal in the longer run to fix whatever is practical to fix when view/@view is used but this should provide a simple and reliable workaround in the short term.

This change introduced an issue that more more blocks than necessary are being allocated in the result of block sparse matrix multiplication, I'll have to investigate that before merging this. EDIT: It turned out that fill!(a::BlockSparseArray, 0) an a::BlockSparseArray .= 0 were both allocating blocks filled with zeros, and ArrayLayouts.jl, which we are using for a matrix multiplication interface in BlockSparseArrays and SparseArrayDOKs, calls fill!(a, 0) at some point in the matrix multiplication call stack. Now both operations drop all blocks from a which fixed that issue.

[emstoudenmire]

Looks really nice