diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml index b82d7829dc..8baecda32c 100644 --- a/NDTensors/Project.toml +++ b/NDTensors/Project.toml @@ -18,6 +18,7 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" HalfIntegers = "f0d1745a-41c9-11e9-1dd9-e5d34d218721" InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900" PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -34,20 +35,20 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8" [weakdeps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" -cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1" HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4" TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9" +cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1" [extensions] NDTensorsAMDGPUExt = "AMDGPU" NDTensorsCUDAExt = "CUDA" -NDTensorscuTENSORExt = "cuTENSOR" NDTensorsHDF5Ext = "HDF5" NDTensorsMetalExt = "Metal" NDTensorsOctavianExt = "Octavian" NDTensorsTBLISExt = "TBLIS" +NDTensorscuTENSORExt = "cuTENSOR" [compat] Accessors = "0.1.33" @@ -65,6 +66,7 @@ HDF5 = "0.14, 0.15, 0.16, 0.17" HalfIntegers = "1" InlineStrings = "1" LinearAlgebra = "1.6" +MacroTools = "0.5" MappedArrays = "0.4" PackageExtensionCompat = "1" Random = "1.6" diff --git a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl index da90718dc4..3949abac5e 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl @@ -1,16 +1,24 @@ using Accessors: @set using Dictionaries: Dictionary, set! +using MacroTools: @capture using ..SparseArrayInterface: SparseArrayInterface, AbstractSparseArray, getindex_zero_function # TODO: Parametrize by `data`? struct SparseArrayDOK{T,N,Zero} <: AbstractSparseArray{T,N} data::Dictionary{CartesianIndex{N},T} - dims::NTuple{N,Int} + dims::Ref{NTuple{N,Int}} zero::Zero + function SparseArrayDOK{T,N,Zero}(data, dims::NTuple{N,Int}, zero) where {T,N,Zero} + return new{T,N,Zero}(data, Ref(dims), zero) + end end # Constructors +function SparseArrayDOK(data, dims::Tuple{Vararg{Int}}, zero) + return SparseArrayDOK{eltype(data),length(dims),typeof(zero)}(data, dims, zero) +end + function SparseArrayDOK{T,N,Zero}(dims::Tuple{Vararg{Int}}, zero) where {T,N,Zero} return SparseArrayDOK{T,N,Zero}(default_data(T, N), dims, zero) end @@ -72,7 +80,7 @@ function SparseArrayDOK{T}(::UndefInitializer, dims::Tuple{Vararg{Int}}, zero) w end # Base `AbstractArray` interface -Base.size(a::SparseArrayDOK) = a.dims +Base.size(a::SparseArrayDOK) = a.dims[] SparseArrayInterface.getindex_zero_function(a::SparseArrayDOK) = a.zero function SparseArrayInterface.set_getindex_zero_function(a::SparseArrayDOK, f) @@ -104,3 +112,37 @@ SparseArrayDOK{T}(a::AbstractArray) where {T} = SparseArrayDOK{T,ndims(a)}(a) function SparseArrayDOK{T,N}(a::AbstractArray) where {T,N} return SparseArrayInterface.sparse_convert(SparseArrayDOK{T,N}, a) end + +function Base.resize!(a::SparseArrayDOK{<:Any,N}, new_size::NTuple{N,Integer}) where {N} + a.dims[] = new_size + return a +end + +function setindex_maybe_grow!(a::SparseArrayDOK{<:Any,N}, value, I::Vararg{Int,N}) where {N} + if any(I .> size(a)) + resize!(a, max.(I, size(a))) + end + a[I...] = value + return a +end + +function is_setindex!_expr(expr::Expr) + return is_assignment_expr(expr) && is_getindex_expr(first(expr.args)) +end +is_setindex!_expr(x) = false + +is_getindex_expr(expr::Expr) = (expr.head === :ref) +is_getindex_expr(x) = false + +is_assignment_expr(expr::Expr) = (expr.head === :(=)) +is_assignment_expr(expr) = false + +macro maybe_grow(expr) + if !is_setindex!_expr(expr) + error( + "@maybe_grow must be used with setindex! syntax (as @maybe_grow a[i,j,...] = value)" + ) + end + @capture(expr, array_[indices__] = value_) + return :(setindex_maybe_grow!($(esc(array)), $value, $indices...)) +end diff --git a/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl b/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl index 54cd3e1540..8d786b9525 100644 --- a/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl +++ b/NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl @@ -5,8 +5,10 @@ # Custom zero type # Slicing +using Dictionaries: Dictionary using Test: @test, @testset, @test_broken -using NDTensors.SparseArrayDOKs: SparseArrayDOK, SparseMatrixDOK +using NDTensors.SparseArrayDOKs: + SparseArrayDOKs, SparseArrayDOK, SparseMatrixDOK, @maybe_grow using NDTensors.SparseArrayInterface: storage_indices, nstored using SparseArrays: SparseMatrixCSC, nnz @testset "SparseArrayDOK (eltype=$elt)" for elt in @@ -94,5 +96,35 @@ using SparseArrays: SparseMatrixCSC, nnz end end end + @testset "Maybe Grow Feature" begin + a = SparseArrayDOK{elt,2}((0, 0)) + SparseArrayDOKs.setindex_maybe_grow!(a, 230, 2, 3) + @test size(a) == (2, 3) + @test a[2, 3] == 230 + # Test @maybe_grow macro + @maybe_grow a[5, 5] = 550 + @test size(a) == (5, 5) + @test a[2, 3] == 230 + @test a[5, 5] == 550 + # Test that size remains same + # if we set at an index smaller than + # the maximum size: + @maybe_grow a[3, 4] = 340 + @test size(a) == (5, 5) + @test a[2, 3] == 230 + @test a[5, 5] == 550 + @test a[3, 4] == 340 + # Test vector case + v = SparseArrayDOK{elt,1}((0,)) + @maybe_grow v[5] = 50 + @test size(v) == (5,) + @test v[5] == 50 + end + @testset "Test Lower Level Constructor" begin + d = Dictionary{CartesianIndex{2},elt}() + a = SparseArrayDOK(d, (2, 2), zero(elt)) + a[1, 2] = 12.0 + @test a[1, 2] == 12.0 + end end end