From 0b786d78e830ca7ab4f078d920f84e3b20cab428 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Thu, 15 May 2025 13:48:05 -0400 Subject: [PATCH 1/8] Unblocked slicing --- .../BlockArraysExtensions.jl | 11 +++-- src/BlockSparseArrays.jl | 1 + src/abstractblocksparsearray/arraylayouts.jl | 11 ++++- .../unblockedsubarray.jl | 44 +++++++++++++++++++ .../blocksparsearrayinterface.jl | 2 +- 5 files changed, 64 insertions(+), 5 deletions(-) create mode 100644 src/abstractblocksparsearray/unblockedsubarray.jl diff --git a/src/BlockArraysExtensions/BlockArraysExtensions.jl b/src/BlockArraysExtensions/BlockArraysExtensions.jl index 46d015a4..3853304f 100644 --- a/src/BlockArraysExtensions/BlockArraysExtensions.jl +++ b/src/BlockArraysExtensions/BlockArraysExtensions.jl @@ -30,13 +30,14 @@ using SparseArraysBase: # A return type for `blocks(array)` when `array` isn't blocked. # Represents a vector with just that single block. -struct SingleBlockView{T,N,Array<:AbstractArray{T,N}} <: AbstractArray{T,N} +struct SingleBlockView{N,Array<:AbstractArray{<:Any,N}} <: AbstractArray{Array,N} array::Array end Base.parent(a::SingleBlockView) = a.array +Base.size(a::SingleBlockView) = ntuple(Returns(1), ndims(a)) blocks_maybe_single(a) = blocks(a) blocks_maybe_single(a::Array) = SingleBlockView(a) -function Base.getindex(a::SingleBlockView{<:Any,N}, index::Vararg{Int,N}) where {N} +function Base.getindex(a::SingleBlockView{N}, index::Vararg{Int,N}) where {N} @assert all(isone, index) return parent(a) end @@ -357,7 +358,11 @@ function blockrange(axis::AbstractUnitRange, r::Base.Slice) end function blockrange(axis::AbstractUnitRange, r::NonBlockedVector) - return Block(1):Block(1) + return Block.(Base.OneTo(1)) +end + +function blockrange(axis::AbstractUnitRange, r::AbstractVector{<:Integer}) + return Block.(Base.OneTo(1)) end function blockrange(axis::AbstractUnitRange, r) diff --git a/src/BlockSparseArrays.jl b/src/BlockSparseArrays.jl index ffbfe937..893bcf0c 100644 --- a/src/BlockSparseArrays.jl +++ b/src/BlockSparseArrays.jl @@ -27,6 +27,7 @@ include("abstractblocksparsearray/abstractblocksparsearray.jl") include("abstractblocksparsearray/abstractblocksparsematrix.jl") include("abstractblocksparsearray/abstractblocksparsevector.jl") include("abstractblocksparsearray/wrappedabstractblocksparsearray.jl") +include("abstractblocksparsearray/unblockedsubarray.jl") include("abstractblocksparsearray/views.jl") include("abstractblocksparsearray/arraylayouts.jl") include("abstractblocksparsearray/sparsearrayinterface.jl") diff --git a/src/abstractblocksparsearray/arraylayouts.jl b/src/abstractblocksparsearray/arraylayouts.jl index 8b771208..b875c5d9 100644 --- a/src/abstractblocksparsearray/arraylayouts.jl +++ b/src/abstractblocksparsearray/arraylayouts.jl @@ -43,11 +43,20 @@ function ArrayLayouts.sub_materialize(layout::BlockLayout{<:SparseLayout}, a, ax return a_dest end +function _similar(arraytype::Type{<:AbstractArray}, size::Tuple) + return similar(arraytype, size) +end +function _similar( + ::Type{<:SubArray{<:Any,<:Any,<:ArrayType}}, size::Tuple +) where {ArrayType} + return similar(ArrayType, size) +end + # Materialize a SubArray view. function ArrayLayouts.sub_materialize( layout::BlockLayout{<:SparseLayout}, a, axes::Tuple{Vararg{Base.OneTo}} ) - a_dest = blocktype(a)(undef, length.(axes)) + a_dest = _similar(blocktype(a), length.(axes)) a_dest .= a return a_dest end diff --git a/src/abstractblocksparsearray/unblockedsubarray.jl b/src/abstractblocksparsearray/unblockedsubarray.jl new file mode 100644 index 00000000..3fb2db51 --- /dev/null +++ b/src/abstractblocksparsearray/unblockedsubarray.jl @@ -0,0 +1,44 @@ +using ArrayLayouts: ArrayLayouts, MemoryLayout +using Base.Broadcast: Broadcast, BroadcastStyle +using BlockArrays: BlockArrays +using TypeParameterAccessors: TypeParameterAccessors, parenttype, similartype + +const UnblockedSubArray{T,N} = SubArray{ + T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{Vector{<:Integer}}} +} + +function BlockArrays.blocks(a::UnblockedSubArray) + return SingleBlockView(a) +end + +function DerivableInterfaces.interface(arraytype::Type{<:UnblockedSubArray}) + return interface(blocktype(parenttype(arraytype))) +end + +function ArrayLayouts.MemoryLayout(arraytype::Type{<:UnblockedSubArray}) + return MemoryLayout(blocktype(parenttype(arraytype))) +end + +function Broadcast.BroadcastStyle(arraytype::Type{<:UnblockedSubArray}) + return BroadcastStyle(blocktype(parenttype(arraytype))) +end + +function TypeParameterAccessors.similartype(arraytype::Type{<:UnblockedSubArray}, elt::Type) + return similartype(blocktype(parenttype(arraytype)), elt) +end + +function Base.map!( + f, a_dest::AbstractArray, a_src1::UnblockedSubArray, a_src_rest::UnblockedSubArray... +) + return invoke( + map!, + Tuple{Any,AbstractArray,AbstractArray,Vararg{AbstractArray}}, + f, + a_dest, + a_src1, + a_src_rest..., + ) +end +function Base.iszero(a::UnblockedSubArray) + return invoke(iszero, Tuple{AbstractArray}, a) +end diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index 1b1af5f5..594470de 100644 --- a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -400,7 +400,7 @@ function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) whe # TODO: Implement this properly. return true end -function SparseArraysBase.eachstoredindex(a::SparseSubArrayBlocks) +function SparseArraysBase.eachstoredindex(::IndexCartesian, a::SparseSubArrayBlocks) return eachstoredindex(view(blocks(parent(a.array)), blockrange(a)...)) end # TODO: Either make this the generic interface or define From 0d11eec40f428dc0627310be01791761415a63a3 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Thu, 15 May 2025 14:47:27 -0400 Subject: [PATCH 2/8] Noncontiguous slicing on GPU --- .../unblockedsubarray.jl | 21 +++++++++++++++++++ test/test_basics.jl | 8 ++++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/src/abstractblocksparsearray/unblockedsubarray.jl b/src/abstractblocksparsearray/unblockedsubarray.jl index 3fb2db51..12e4651d 100644 --- a/src/abstractblocksparsearray/unblockedsubarray.jl +++ b/src/abstractblocksparsearray/unblockedsubarray.jl @@ -27,6 +27,27 @@ function TypeParameterAccessors.similartype(arraytype::Type{<:UnblockedSubArray} return similartype(blocktype(parenttype(arraytype)), elt) end +function Base.similar( + a::UnblockedSubArray, elt::Type, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}} +) + return similar(similartype(blocktype(parenttype(a)), elt), axes) +end +function Base.similar(a::UnblockedSubArray, elt::Type, size::Tuple{Int,Vararg{Int}}) + return similar(a, elt, Base.OneTo.(size)) +end + +function ArrayLayouts.sub_materialize(a::UnblockedSubArray) + a_cpu = adapt(Array, a) + a_cpu′ = similar(a_cpu) + a_cpu′ .= a_cpu + if typeof(a) === typeof(a_cpu) + return a_cpu′ + end + a′ = similar(a) + a′ .= a_cpu′ + return a′ +end + function Base.map!( f, a_dest::AbstractArray, a_src1::UnblockedSubArray, a_src_rest::UnblockedSubArray... ) diff --git a/test/test_basics.jl b/test/test_basics.jl index 0a59f146..67ee8ed9 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -56,7 +56,6 @@ arrayts = (Array, JLArray) a[Block(1, 1)] = dev(randn(elt, 2, 2)) a[Block(2, 2)] = dev(randn(elt, 3, 3)) @test_broken a[:, [2, 4]] - @test_broken a[[3, 5], [2, 4]] # TODO: Fix this and turn it into a proper test. a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) @@ -713,6 +712,13 @@ arrayts = (Array, JLArray) @test a[Block(2, 2)[1:2, 2:3]] == b @test blockstoredlength(a) == 1 + # Non-contiguous slicing. + a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + a[Block(2, 2)] = dev(randn(elt, 3, 3)) + I = ([3, 5], [2, 4]) + @test a[I...] == Array(a)[I...] + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) @views for b in [Block(1, 1), Block(2, 2)] a[b] = randn(elt, size(a[b])) From 11bede8773078b4e5ce8b0c51f4fd48db82f9cda Mon Sep 17 00:00:00 2001 From: mtfishman Date: Thu, 15 May 2025 14:55:39 -0400 Subject: [PATCH 3/8] Fix test on GPU --- test/test_basics.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_basics.jl b/test/test_basics.jl index 67ee8ed9..364bdf9d 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -717,7 +717,7 @@ arrayts = (Array, JLArray) a[Block(1, 1)] = dev(randn(elt, 2, 2)) a[Block(2, 2)] = dev(randn(elt, 3, 3)) I = ([3, 5], [2, 4]) - @test a[I...] == Array(a)[I...] + @test Array(a[I...]) == Array(a)[I...] a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) @views for b in [Block(1, 1), Block(2, 2)] From 3b6cb1e5db3953c32597b3895f0995c61a2bc5e6 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Thu, 15 May 2025 14:57:53 -0400 Subject: [PATCH 4/8] Bump version --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 9e61a2e4..b583b073 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "BlockSparseArrays" uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4" authors = ["ITensor developers and contributors"] -version = "0.5.1" +version = "0.5.2" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" From 5e53555abf15ffe370377512848ab8f14453e8e3 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Thu, 15 May 2025 15:24:30 -0400 Subject: [PATCH 5/8] Start fixing more slicing --- src/abstractblocksparsearray/unblockedsubarray.jl | 8 ++++++-- test/test_basics.jl | 13 +++++++------ 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/src/abstractblocksparsearray/unblockedsubarray.jl b/src/abstractblocksparsearray/unblockedsubarray.jl index 12e4651d..758b24b0 100644 --- a/src/abstractblocksparsearray/unblockedsubarray.jl +++ b/src/abstractblocksparsearray/unblockedsubarray.jl @@ -1,10 +1,14 @@ using ArrayLayouts: ArrayLayouts, MemoryLayout using Base.Broadcast: Broadcast, BroadcastStyle -using BlockArrays: BlockArrays +using BlockArrays: BlockArrays, Block, BlockIndexRange, BlockSlice using TypeParameterAccessors: TypeParameterAccessors, parenttype, similartype +const UnblockedIndices = Union{ + Vector{<:Integer},BlockSlice{<:Block{1}},BlockSlice{<:BlockIndexRange{1}} +} + const UnblockedSubArray{T,N} = SubArray{ - T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{Vector{<:Integer}}} + T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{UnblockedIndices}} } function BlockArrays.blocks(a::UnblockedSubArray) diff --git a/test/test_basics.jl b/test/test_basics.jl index 364bdf9d..395e1109 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -51,12 +51,6 @@ arrayts = (Array, JLArray) a[Block(2, 2)] = dev(randn(elt, 3, 3)) @test_broken a[:, 4] - # TODO: Fix this and turn it into a proper test. - a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a[Block(1, 1)] = dev(randn(elt, 2, 2)) - a[Block(2, 2)] = dev(randn(elt, 3, 3)) - @test_broken a[:, [2, 4]] - # TODO: Fix this and turn it into a proper test. a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) a[Block(1, 1)] = dev(randn(elt, 2, 2)) @@ -719,6 +713,13 @@ arrayts = (Array, JLArray) I = ([3, 5], [2, 4]) @test Array(a[I...]) == Array(a)[I...] + # TODO: Fix this and turn it into a proper test. + a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) + a[Block(1, 1)] = dev(randn(elt, 2, 2)) + a[Block(2, 2)] = dev(randn(elt, 3, 3)) + I = (:, [2, 4]) + @test Array(a[I...]) == Array(a)[I...] + a = BlockSparseArray{elt}(undef, [2, 3], [2, 3]) @views for b in [Block(1, 1), Block(2, 2)] a[b] = randn(elt, size(a[b])) From e80370e4cebefdb3ee3848cbd4ec15bc7e65fbb8 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Thu, 15 May 2025 17:56:57 -0400 Subject: [PATCH 6/8] More general eachstoredindex/isstored for SubArray blocks --- .../blocksparsearrayinterface.jl | 24 +++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index 594470de..a76edbfb 100644 --- a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -364,7 +364,19 @@ end function Base.size(a::SparseSubArrayBlocks) return length.(axes(a)) end -# TODO: Define `isstored`. + +# TODO: Make a faster version for when the slice is blockwise. +function SparseArraysBase.isstored( + a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N} +) where {N} + J = Base.reindex(parentindices(a.array), to_indices(a.array, Block.(I))) + # TODO: Try doing this blockwise when possible rather + # than elementwise. + return any(Iterators.product(J...)) do K + return isstored(parent(a.array), K...) + end +end + # TODO: Define `getstoredindex`, `getunstoredindex` instead. function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where {N} # TODO: Should this be defined as `@view a.array[Block(I)]` instead? @@ -400,9 +412,17 @@ function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) whe # TODO: Implement this properly. return true end + function SparseArraysBase.eachstoredindex(::IndexCartesian, a::SparseSubArrayBlocks) - return eachstoredindex(view(blocks(parent(a.array)), blockrange(a)...)) + return filter(eachindex(a)) do I + return isstored(a, I) + end + + ## # TODO: This only works for blockwise slices, i.e. slices using + ## # `BlockSliceCollection`. + ## return eachstoredindex(view(blocks(parent(a.array)), blockrange(a)...)) end + # TODO: Either make this the generic interface or define # `SparseArraysBase.sparse_storage`, which is used # to defined this. From 3e4a0a17527b6525ec687a3df3177f5f8d9f74f5 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Thu, 15 May 2025 20:57:52 -0400 Subject: [PATCH 7/8] Fix tests --- .../unblockedsubarray.jl | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/src/abstractblocksparsearray/unblockedsubarray.jl b/src/abstractblocksparsearray/unblockedsubarray.jl index 758b24b0..355ae824 100644 --- a/src/abstractblocksparsearray/unblockedsubarray.jl +++ b/src/abstractblocksparsearray/unblockedsubarray.jl @@ -64,6 +64,23 @@ function Base.map!( a_src_rest..., ) end + +# Fix ambiguity and scalar indexing errors with GPUArrays. +using Adapt: adapt +using GPUArraysCore: GPUArraysCore +function Base.map!( + f, + a_dest::GPUArraysCore.AnyGPUArray, + a_src1::UnblockedSubArray, + a_src_rest::UnblockedSubArray..., +) + a_dest_cpu = adapt(Array, a_dest) + a_srcs_cpu = map(adapt(Array), (a_src1, a_src_rest...)) + map!(f, a_dest_cpu, a_srcs_cpu...) + a_dest .= a_dest_cpu + return a_dest +end + function Base.iszero(a::UnblockedSubArray) - return invoke(iszero, Tuple{AbstractArray}, a) + return invoke(iszero, Tuple{AbstractArray}, adapt(Array, a)) end From 38782d8f118293133789dcef25cef75de44eebd6 Mon Sep 17 00:00:00 2001 From: Matt Fishman Date: Fri, 16 May 2025 14:38:19 -0400 Subject: [PATCH 8/8] Update comments in tests --- test/test_basics.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_basics.jl b/test/test_basics.jl index 395e1109..bc220d2b 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -706,14 +706,14 @@ arrayts = (Array, JLArray) @test a[Block(2, 2)[1:2, 2:3]] == b @test blockstoredlength(a) == 1 - # Non-contiguous slicing. + # Noncontiguous slicing. a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) a[Block(1, 1)] = dev(randn(elt, 2, 2)) a[Block(2, 2)] = dev(randn(elt, 3, 3)) I = ([3, 5], [2, 4]) @test Array(a[I...]) == Array(a)[I...] - # TODO: Fix this and turn it into a proper test. + # Noncontiguous slicing. a = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) a[Block(1, 1)] = dev(randn(elt, 2, 2)) a[Block(2, 2)] = dev(randn(elt, 3, 3))