diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml index 56ed0df36b..ece7615a0d 100644 --- a/NDTensors/Project.toml +++ b/NDTensors/Project.toml @@ -1,7 +1,7 @@ name = "NDTensors" uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" authors = ["Matthew Fishman "] -version = "0.3.28" +version = "0.3.29" [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" diff --git a/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl b/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl index 2f24823317..1c058a8efa 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl @@ -335,3 +335,34 @@ function blocked_cartesianindices(axes::Tuple, subaxes::Tuple, blocks) return cartesianindices(subaxes, block) end end + +function view!(a::BlockSparseArray{<:Any,N}, index::Block{N}) where {N} + return view!(a, Tuple(index)...) +end +function view!(a::AbstractArray{<:Any,N}, index::Vararg{Block{1},N}) where {N} + blocks(a)[Int.(index)...] = blocks(a)[Int.(index)...] + return blocks(a)[Int.(index)...] +end + +function view!(a::AbstractArray{<:Any,N}, index::BlockIndexRange{N}) where {N} + # TODO: Is there a better code pattern for this? + indices = ntuple(N) do dim + return Tuple(Block(index))[dim][index.indices[dim]] + end + return view!(a, indices...) +end +function view!(a::AbstractArray{<:Any,N}, index::Vararg{BlockIndexRange{1},N}) where {N} + b = view!(a, Block.(index)...) + r = map(index -> only(index.indices), index) + return @view b[r...] +end + +using MacroTools: @capture +using NDTensors.SparseArrayDOKs: is_getindex_expr +macro view!(expr) + if !is_getindex_expr(expr) + error("@view must be used with getindex syntax (as `@view! a[i,j,...]`)") + end + @capture(expr, array_[indices__]) + return :(view!($(esc(array)), $(esc.(indices)...))) +end diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index 3533366b41..8b3d37bd36 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -186,6 +186,11 @@ function Base.setindex!(a::BlockSparseArrayLike{<:Any,1}, value, I::Block{1}) return a end +function Base.fill!(a::BlockSparseArrayLike, value) + blocksparse_fill!(a, value) + return a +end + # `BlockArrays` interface # TODO: Is this needed if `blocks` is defined? function BlockArrays.viewblock(a::BlockSparseArrayLike{<:Any,N}, I::Block{N,Int}) where {N} diff --git a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index f2e9326219..d1fd59ebf5 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -9,7 +9,7 @@ using BlockArrays: blocklengths, findblockindex using LinearAlgebra: Adjoint, Transpose -using ..SparseArrayInterface: perm, iperm, nstored +using ..SparseArrayInterface: perm, iperm, nstored, sparse_zero! ## using MappedArrays: mappedarray blocksparse_blocks(a::AbstractArray) = error("Not implemented") @@ -95,6 +95,18 @@ function blocksparse_setindex!( return a end +function blocksparse_fill!(a::AbstractArray, value) + if iszero(value) + # This drops all of the blocks. + sparse_zero!(blocks(a)) + return a + end + for b in BlockRange(a) + a[b] .= value + end + return a +end + function blocksparse_view(a::AbstractArray{<:Any,N}, I::Block{N}) where {N} return blocksparse_view(a, Tuple(I)...) end diff --git a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl index a1dfeb79fd..adc54453b8 100644 --- a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl +++ b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl @@ -11,7 +11,8 @@ using BlockArrays: blocksize, mortar using LinearAlgebra: mul! -using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored, block_reshape +using NDTensors.BlockSparseArrays: + @view!, BlockSparseArray, block_nstored, block_reshape, view! using NDTensors.SparseArrayInterface: nstored using NDTensors.TensorAlgebra: contract using Test: @test, @test_broken, @test_throws, @testset @@ -403,6 +404,80 @@ include("TestBlockSparseArraysUtils.jl") @test !isassigned(a, Block(2)[3], Block(2)[5]) @test !isassigned(a, Block(1)[1], Block(0)[1]) @test !isassigned(a, Block(3)[3], Block(2)[4]) + + a = BlockSparseArray{elt}([2, 3], [3, 4]) + @test iszero(a) + @test iszero(block_nstored(a)) + fill!(a, 0) + @test iszero(a) + @test iszero(block_nstored(a)) + fill!(a, 2) + @test !iszero(a) + @test all(==(2), a) + @test block_nstored(a) == 4 + fill!(a, 0) + @test iszero(a) + @test iszero(block_nstored(a)) + + a = BlockSparseArray{elt}([2, 3], [3, 4]) + @test iszero(a) + @test iszero(block_nstored(a)) + a .= 0 + @test iszero(a) + @test iszero(block_nstored(a)) + a .= 2 + @test !iszero(a) + @test all(==(2), a) + @test block_nstored(a) == 4 + a .= 0 + @test iszero(a) + @test iszero(block_nstored(a)) + end + @testset "view!" begin + for blk in ((Block(2, 2),), (Block(2), Block(2))) + a = BlockSparseArray{elt}([2, 3], [2, 3]) + b = view!(a, blk...) + x = randn(elt, 3, 3) + b .= x + @test b == x + @test a[blk...] == x + @test @view(a[blk...]) == x + @test view!(a, blk...) == x + @test @view!(a[blk...]) == x + end + for blk in ((Block(2, 2),), (Block(2), Block(2))) + a = BlockSparseArray{elt}([2, 3], [2, 3]) + b = @view! a[blk...] + x = randn(elt, 3, 3) + b .= x + @test b == x + @test a[blk...] == x + @test @view(a[blk...]) == x + @test view!(a, blk...) == x + @test @view!(a[blk...]) == x + end + for blk in ((Block(2, 2)[2:3, 1:2],), (Block(2)[2:3], Block(2)[1:2])) + a = BlockSparseArray{elt}([2, 3], [2, 3]) + b = view!(a, blk...) + x = randn(elt, 2, 2) + b .= x + @test b == x + @test a[blk...] == x + @test @view(a[blk...]) == x + @test view!(a, blk...) == x + @test @view!(a[blk...]) == x + end + for blk in ((Block(2, 2)[2:3, 1:2],), (Block(2)[2:3], Block(2)[1:2])) + a = BlockSparseArray{elt}([2, 3], [2, 3]) + b = @view! a[blk...] + x = randn(elt, 2, 2) + b .= x + @test b == x + @test a[blk...] == x + @test @view(a[blk...]) == x + @test view!(a, blk...) == x + @test @view!(a[blk...]) == x + end end @testset "LinearAlgebra" begin a1 = BlockSparseArray{elt}([2, 3], [2, 3]) diff --git a/NDTensors/src/lib/SparseArrayDOKs/test/Project.toml b/NDTensors/src/lib/SparseArrayDOKs/test/Project.toml index 9b1d5ccd25..ade2efc435 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/test/Project.toml +++ b/NDTensors/src/lib/SparseArrayDOKs/test/Project.toml @@ -1,2 +1,3 @@ [deps] +Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4" NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" diff --git a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/arraylayouts.jl b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/arraylayouts.jl index e9c80050f2..038689be8a 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/arraylayouts.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/arraylayouts.jl @@ -1,5 +1,13 @@ -using ArrayLayouts: ArrayLayouts +using ArrayLayouts: ArrayLayouts, MatMulMatAdd function ArrayLayouts.MemoryLayout(arraytype::Type{<:SparseArrayLike}) return SparseLayout() end + +function ArrayLayouts.materialize!( + m::MatMulMatAdd{<:AbstractSparseLayout,<:AbstractSparseLayout,<:AbstractSparseLayout} +) + α, a1, a2, β, a_dest = m.α, m.A, m.B, m.β, m.C + sparse_mul!(a_dest, a1, a2, α, β) + return a_dest +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/baseinterface.jl b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/baseinterface.jl index bd9dd326af..323289246e 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/baseinterface.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/baseinterface.jl @@ -29,3 +29,7 @@ end function Base.setindex!(a::AbstractSparseArray, I...) return SparseArrayInterface.sparse_setindex!(a, I...) end + +function Base.fill!(a::AbstractSparseArray, value) + return SparseArrayInterface.sparse_fill!(a, value) +end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/sparsearrayinterface.jl b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/sparsearrayinterface.jl index fa8d1ac340..897a77b960 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/sparsearrayinterface.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/abstractsparsearray/sparsearrayinterface.jl @@ -13,5 +13,6 @@ end function SparseArrayInterface.setindex_notstored!( a::AbstractSparseArray{<:Any,N}, value, I::CartesianIndex{N} ) where {N} + iszero(value) && return a return error("Setting the specified unstored index is not supported.") end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/base.jl b/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/base.jl index 474e3bd129..16144e66bc 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/base.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/base.jl @@ -44,13 +44,16 @@ end # but we don't use that here since `sparse_fill!` # is used inside of `sparse_map!`. function sparse_fill!(a::AbstractArray, x) - ## TODO: Add this back? - ## if iszero(x) - ## # TODO: Check that `typeof(x)` is compatible - ## # with `eltype(a)`. - ## dropall!(a) - ## end - fill!(sparse_storage(a), x) + if iszero(x) + # This checks that `x` is compatible + # with `eltype(a)`. + x = convert(eltype(a), x) + sparse_zero!(a) + return a + end + for I in eachindex(a) + a[I] = x + end return a end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/indexing.jl b/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/indexing.jl index 6ab1f77a88..1a9a73060c 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/indexing.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/indexing.jl @@ -153,9 +153,7 @@ function sparse_setindex!(a::AbstractArray, value, I::StoredIndex) end function sparse_setindex!(a::AbstractArray, value, I::NotStoredIndex) - if !iszero(value) - setindex_notstored!(a, value, index(I)) - end + setindex_notstored!(a, value, index(I)) return a end diff --git a/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/interface_optional.jl b/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/interface_optional.jl index 122995dd34..710cd2570d 100644 --- a/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/interface_optional.jl +++ b/NDTensors/src/lib/SparseArrayInterface/src/sparsearrayinterface/interface_optional.jl @@ -15,6 +15,7 @@ end # where there is not a stored value. # Some types (like `Diagonal`) may not support this. function setindex_notstored!(a::AbstractArray, value, I) + iszero(value) && return a return throw(ArgumentError("Can't set nonzero values of $(typeof(a)).")) end diff --git a/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/AbstractSparseArrays.jl b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/AbstractSparseArrays.jl index deb9348d42..1455b65c15 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/AbstractSparseArrays.jl +++ b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/AbstractSparseArrays.jl @@ -1,6 +1,6 @@ module AbstractSparseArrays +using ArrayLayouts: ArrayLayouts, MatMulMatAdd, MemoryLayout, MulAdd using NDTensors.SparseArrayInterface: SparseArrayInterface, AbstractSparseArray -using ArrayLayouts: ArrayLayouts, MemoryLayout, MulAdd struct SparseArray{T,N} <: AbstractSparseArray{T,N} data::Vector{T} @@ -26,6 +26,13 @@ ArrayLayouts.MemoryLayout(::Type{<:SparseArray}) = SparseLayout() function Base.similar(::MulAdd{<:SparseLayout,<:SparseLayout}, elt::Type, axes) return similar(SparseArray{elt}, axes) end +function ArrayLayouts.materialize!( + m::MatMulMatAdd{<:SparseLayout,<:SparseLayout,<:SparseLayout} +) + α, a1, a2, β, a_dest = m.α, m.A, m.B, m.β, m.C + SparseArrayInterface.sparse_mul!(a_dest, a1, a2, α, β) + return a_dest +end # AbstractArray interface Base.size(a::SparseArray) = a.dims diff --git a/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrays.jl b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrays.jl index ddf8b2f8cc..abb531303c 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrays.jl +++ b/NDTensors/src/lib/SparseArrayInterface/test/SparseArrayInterfaceTestUtils/SparseArrays.jl @@ -1,4 +1,5 @@ module SparseArrays +using LinearAlgebra: LinearAlgebra using NDTensors.SparseArrayInterface: SparseArrayInterface struct SparseArray{T,N} <: AbstractArray{T,N} @@ -19,6 +20,18 @@ function SparseArray{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}) where {T} end SparseArray{T}(dims::Vararg{Int}) where {T} = SparseArray{T}(dims) +# LinearAlgebra interface +function LinearAlgebra.mul!( + a_dest::AbstractMatrix, + a1::SparseArray{<:Any,2}, + a2::SparseArray{<:Any,2}, + α::Number, + β::Number, +) + SparseArrayInterface.sparse_mul!(a_dest, a1, a2, α, β) + return a_dest +end + # AbstractArray interface Base.size(a::SparseArray) = a.dims function Base.similar(a::SparseArray, elt::Type, dims::Tuple{Vararg{Int}}) @@ -28,8 +41,11 @@ end function Base.getindex(a::SparseArray, I...) return SparseArrayInterface.sparse_getindex(a, I...) end -function Base.setindex!(a::SparseArray, I...) - return SparseArrayInterface.sparse_setindex!(a, I...) +function Base.setindex!(a::SparseArray, value, I...) + return SparseArrayInterface.sparse_setindex!(a, value, I...) +end +function Base.fill!(a::SparseArray, value) + return SparseArrayInterface.sparse_fill!(a, value) end # Minimal interface diff --git a/NDTensors/src/lib/SparseArrayInterface/test/test_abstractsparsearray.jl b/NDTensors/src/lib/SparseArrayInterface/test/test_abstractsparsearray.jl index 78bab9aecf..cd233f4ff1 100644 --- a/NDTensors/src/lib/SparseArrayInterface/test/test_abstractsparsearray.jl +++ b/NDTensors/src/lib/SparseArrayInterface/test/test_abstractsparsearray.jl @@ -280,7 +280,12 @@ using Test: @test, @testset a′ = copy(a) a′ .+= b @test a′ == a + b - @test SparseArrayInterface.nstored(a′) == 2 + # TODO: Should this be: + # ```julia + # @test SparseArrayInterface.nstored(a′) == 2 + # ``` + # ? I.e. should it only store the nonzero values? + @test SparseArrayInterface.nstored(a′) == 6 # Matrix multiplication a1 = SparseArray{elt}(2, 3)