From d56b4a798606908ada78c31961695236da0d6624 Mon Sep 17 00:00:00 2001 From: Matt Fishman Date: Thu, 28 Sep 2023 13:57:13 -0400 Subject: [PATCH 1/2] [NDTensor] Allow general objects are storage backends for Tensor (#1206) --- NDTensors/src/NDTensors.jl | 7 ++ NDTensors/src/arraytensor/array.jl | 66 +++++++++++ NDTensors/src/arraytensor/arraytensor.jl | 110 ++++++++++++++++++ NDTensors/src/tensor/tensor.jl | 9 +- .../generic_tensor_operations.jl | 2 +- NDTensors/test/arraytensor/arraytensor.jl | 45 +++++++ NDTensors/test/runtests.jl | 1 + src/itensor.jl | 4 +- 8 files changed, 236 insertions(+), 8 deletions(-) create mode 100644 NDTensors/src/arraytensor/array.jl create mode 100644 NDTensors/src/arraytensor/arraytensor.jl create mode 100644 NDTensors/test/arraytensor/arraytensor.jl diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl index ae138d11fc..d2ffcafb59 100644 --- a/NDTensors/src/NDTensors.jl +++ b/NDTensors/src/NDTensors.jl @@ -109,6 +109,13 @@ include("empty/EmptyTensor.jl") include("empty/tensoralgebra/contract.jl") include("empty/adapt.jl") +##################################### +# Array Tensor (experimental) +# TODO: Move to `Experimental` module. +# +include("arraytensor/arraytensor.jl") +include("arraytensor/array.jl") + ##################################### # Deprecations # diff --git a/NDTensors/src/arraytensor/array.jl b/NDTensors/src/arraytensor/array.jl new file mode 100644 index 0000000000..40e89c5181 --- /dev/null +++ b/NDTensors/src/arraytensor/array.jl @@ -0,0 +1,66 @@ +# Combiner +promote_rule(::Type{<:Combiner}, arraytype::Type{<:MatrixOrArrayStorage}) = arraytype + +# Generic AbstractArray code +function contract( + array1::MatrixOrArrayStorage, + labels1, + array2::MatrixOrArrayStorage, + labels2, + labelsR=contract_labels(labels1, labels2), +) + output_array = contraction_output(array1, labels1, array2, labels2, labelsR) + contract!(output_array, labelsR, array1, labels1, array2, labels2) + return output_array +end + +function contraction_output( + array1::MatrixOrArrayStorage, array2::MatrixOrArrayStorage, indsR +) + arraytypeR = contraction_output_type(typeof(array1), typeof(array2), indsR) + return NDTensors.similar(arraytypeR, indsR) +end + +function contraction_output_type( + arraytype1::Type{<:MatrixOrArrayStorage}, arraytype2::Type{<:MatrixOrArrayStorage}, inds +) + return similartype(promote_type(arraytype1, arraytype2), inds) +end + +function contraction_output( + array1::MatrixOrArrayStorage, + labelsarray1, + array2::MatrixOrArrayStorage, + labelsarray2, + labelsoutput_array, +) + # TODO: Maybe use `axes` here to be more generic, for example for BlockArrays? + indsoutput_array = contract_inds( + size(array1), labelsarray1, size(array2), labelsarray2, labelsoutput_array + ) + output_array = contraction_output(array1, array2, indsoutput_array) + return output_array +end + +# Required interface for specific AbstractArray types +function contract!( + arrayR::MatrixOrArrayStorage, + labelsR, + array1::MatrixOrArrayStorage, + labels1, + array2::MatrixOrArrayStorage, + labels2, +) + props = ContractionProperties(labels1, labels2, labelsR) + compute_contraction_properties!(props, array1, array2, arrayR) + # TODO: Change this to just `contract!`, or maybe `contract_ttgt!`? + _contract!(arrayR, array1, array2, props) + return arrayR +end + +function permutedims!( + output_array::MatrixOrArrayStorage, array::MatrixOrArrayStorage, perm, f::Function +) + @strided output_array .= f.(output_array, permutedims(array, perm)) + return output_array +end diff --git a/NDTensors/src/arraytensor/arraytensor.jl b/NDTensors/src/arraytensor/arraytensor.jl new file mode 100644 index 0000000000..d91b4a6d01 --- /dev/null +++ b/NDTensors/src/arraytensor/arraytensor.jl @@ -0,0 +1,110 @@ +# Used for dispatch to distinguish from Tensors wrapping TensorStorage. +# Remove once TensorStorage is removed. +const ArrayStorage{T,N} = Union{ + Array{T,N},ReshapedArray{T,N},SubArray{T,N},PermutedDimsArray{T,N},StridedView{T,N} +} +const MatrixStorage{T} = Union{ + ArrayStorage{T,2}, + Transpose{T}, + Adjoint{T}, + Symmetric{T}, + Hermitian{T}, + UpperTriangular{T}, + LowerTriangular{T}, + UnitUpperTriangular{T}, + UnitLowerTriangular{T}, + Diagonal{T}, +} +const MatrixOrArrayStorage{T} = Union{MatrixStorage{T},ArrayStorage{T}} + +const ArrayStorageTensor{T,N,S,I} = Tensor{T,N,S,I} where {S<:ArrayStorage{T,N}} +const MatrixStorageTensor{T,S,I} = Tensor{T,2,S,I} where {S<:MatrixStorage{T}} +const MatrixOrArrayStorageTensor{T,S,I} = + Tensor{T,N,S,I} where {N,S<:MatrixOrArrayStorage{T}} + +function Tensor(storage::MatrixOrArrayStorageTensor, inds::Tuple) + return Tensor(NeverAlias(), storage, inds) +end + +function Tensor(as::AliasStyle, storage::MatrixOrArrayStorage, inds::Tuple) + return Tensor{eltype(storage),length(inds),typeof(storage),typeof(inds)}( + as, storage, inds + ) +end + +function getindex(tensor::MatrixOrArrayStorageTensor, I::Integer...) + return storage(tensor)[I...] +end + +function setindex!(tensor::MatrixOrArrayStorageTensor, v, I::Integer...) + storage(tensor)[I...] = v + return tensor +end + +function contraction_output( + tensor1::MatrixOrArrayStorageTensor, tensor2::MatrixOrArrayStorageTensor, indsR +) + tensortypeR = contraction_output_type(typeof(tensor1), typeof(tensor2), indsR) + return NDTensors.similar(tensortypeR, indsR) +end + +function contract!( + tensorR::MatrixOrArrayStorageTensor, + labelsR, + tensor1::MatrixOrArrayStorageTensor, + labels1, + tensor2::MatrixOrArrayStorageTensor, + labels2, +) + contract!(storage(tensorR), labelsR, storage(tensor1), labels1, storage(tensor2), labels2) + return tensorR +end + +function permutedims!( + output_tensor::MatrixOrArrayStorageTensor, + tensor::MatrixOrArrayStorageTensor, + perm, + f::Function, +) + permutedims!(storage(output_tensor), storage(tensor), perm, f) + return output_tensor +end + +# Linear algebra (matrix algebra) +function Base.adjoint(tens::MatrixStorageTensor) + return tensor(adjoint(storage(tens)), reverse(inds(tens))) +end + +function LinearAlgebra.mul!( + C::MatrixStorageTensor, A::MatrixStorageTensor, B::MatrixStorageTensor +) + mul!(storage(C), storage(A), storage(B)) + return C +end + +function LinearAlgebra.svd(tens::MatrixStorageTensor) + F = svd(storage(tens)) + U, S, V = F.U, F.S, F.Vt + i, j = inds(tens) + # TODO: Make this more general with a `similar_ind` function, + # so the dimension can be determined from the length of `S`. + min_ij = dim(i) ≤ dim(j) ? i : j + α = sim(min_ij) # similar_ind(i, space(S)) + β = sim(min_ij) # similar_ind(i, space(S)) + Utensor = tensor(U, (i, α)) + # TODO: Remove conversion to `Diagonal` to make more general, or make a generic `Diagonal` concept that works for `BlockSparseArray`. + # Used for now to avoid introducing wrapper types. + Stensor = tensor(Diagonal(S), (α, β)) + Vtensor = tensor(V, (β, j)) + return Utensor, Stensor, Vtensor, Spectrum(nothing, 0.0) +end + +array(tensor::MatrixOrArrayStorageTensor) = storage(tensor) + +# Combiner +function contraction_output( + tensor1::MatrixOrArrayStorageTensor, tensor2::CombinerTensor, indsR +) + tensortypeR = contraction_output_type(typeof(tensor1), typeof(tensor2), indsR) + return NDTensors.similar(tensortypeR, indsR) +end diff --git a/NDTensors/src/tensor/tensor.jl b/NDTensors/src/tensor/tensor.jl index 84699222f1..7e4af3f403 100644 --- a/NDTensors/src/tensor/tensor.jl +++ b/NDTensors/src/tensor/tensor.jl @@ -1,11 +1,10 @@ - """ Tensor{StoreT,IndsT} A plain old tensor (with order independent interface and no assumption of labels) """ -struct Tensor{ElT,N,StoreT<:TensorStorage,IndsT} <: AbstractArray{ElT,N} +struct Tensor{ElT,N,StoreT,IndsT} <: AbstractArray{ElT,N} storage::StoreT inds::IndsT @@ -21,8 +20,8 @@ struct Tensor{ElT,N,StoreT<:TensorStorage,IndsT} <: AbstractArray{ElT,N} and tensor(store::TensorStorage, inds) constructors. """ function Tensor{ElT,N,StoreT,IndsT}( - ::AllowAlias, storage::TensorStorage, inds::Tuple - ) where {ElT,N,StoreT<:TensorStorage,IndsT} + ::AllowAlias, storage, inds::Tuple + ) where {ElT,N,StoreT,IndsT} @assert ElT == eltype(StoreT) @assert length(inds) == N return new{ElT,N,StoreT,IndsT}(storage, inds) @@ -74,7 +73,7 @@ end # already (like a Vector). In the future this may be lifted # to allow for very large tensor orders in which case Tuple # operations may become too slow. -function Tensor(as::AliasStyle, storage::TensorStorage, inds) +function Tensor(as::AliasStyle, storage, inds) return Tensor(as, storage, Tuple(inds)) end diff --git a/NDTensors/src/tensoralgebra/generic_tensor_operations.jl b/NDTensors/src/tensoralgebra/generic_tensor_operations.jl index 170b472bd8..db0a740e22 100644 --- a/NDTensors/src/tensoralgebra/generic_tensor_operations.jl +++ b/NDTensors/src/tensoralgebra/generic_tensor_operations.jl @@ -18,7 +18,7 @@ end function permutedims!(output_tensor::Tensor, tensor::Tensor, perm, f::Function=(r, t) -> t) Base.checkdims_perm(output_tensor, tensor, perm) error( - "`perutedims!(output_tensor::Tensor, tensor::Tensor, perm, f::Function=(r, t) -> t)` not implemented for `typeof(output_tensor) = $(typeof(output_tensor))`, `typeof(tensor) = $(typeof(tensor))`, `perm = $perm`, and `f = $f`.", + "`permutedims!(output_tensor::Tensor, tensor::Tensor, perm, f::Function=(r, t) -> t)` not implemented for `typeof(output_tensor) = $(typeof(output_tensor))`, `typeof(tensor) = $(typeof(tensor))`, `perm = $perm`, and `f = $f`.", ) return output_tensor end diff --git a/NDTensors/test/arraytensor/arraytensor.jl b/NDTensors/test/arraytensor/arraytensor.jl new file mode 100644 index 0000000000..39809d0db3 --- /dev/null +++ b/NDTensors/test/arraytensor/arraytensor.jl @@ -0,0 +1,45 @@ +using NDTensors +using LinearAlgebra +using Test + +using NDTensors: storage, storagetype + +@testset "Tensor wrapping Array" begin + is1 = (2, 3) + D1 = randn(is1) + + is2 = (3, 4) + D2 = randn(is2) + + T1 = tensor(D1, is1) + T2 = tensor(D2, is2) + + @test T1[1, 1] == D1[1, 1] + + x = rand() + T1[1, 1] = x + + @test T1[1, 1] == x + @test array(T1) == D1 + @test storagetype(T1) <: Matrix{Float64} + @test storage(T1) == D1 + @test eltype(T1) == eltype(D1) + @test inds(T1) == is1 + + R = T1 * T2 + @test storagetype(R) <: Matrix{Float64} + @test Array(R) ≈ Array(T1) * Array(T2) + + T1r = randn!(similar(T1)) + @test Array(T1r + T1) ≈ Array(T1r) + Array(T1) + @test Array(permutedims(T1, (2, 1))) ≈ permutedims(Array(T1), (2, 1)) + + U, S, V = svd(T1) + @test U * S * V ≈ T1 + + T12 = contract(T1, (1, -1), T2, (-1, 2)) + @test T12 ≈ T1 * T2 + + D12 = contract(D1, (1, -1), D2, (-1, 2)) + @test D12 ≈ Array(T12) +end diff --git a/NDTensors/test/runtests.jl b/NDTensors/test/runtests.jl index 90aeeea118..cdd0d6cf35 100644 --- a/NDTensors/test/runtests.jl +++ b/NDTensors/test/runtests.jl @@ -28,6 +28,7 @@ end "emptynumber.jl", "emptystorage.jl", "combiner.jl", + "arraytensor/arraytensor.jl", ] println("Running $filename") include(filename) diff --git a/src/itensor.jl b/src/itensor.jl index 19286f065c..c26786e5e2 100644 --- a/src/itensor.jl +++ b/src/itensor.jl @@ -82,7 +82,7 @@ NDTensors.Dense{Float64,Array{Float64,1}} ## Accessor Functions, Index Functions and Operations mutable struct ITensor tensor::Tensor - function ITensor(::AllowAlias, T::Tensor{<:Any,<:Any,<:TensorStorage,<:Tuple}) + function ITensor(::AllowAlias, T::Tensor{<:Any,<:Any,<:Any,<:Tuple}) @debug_check begin is = inds(T) if !allunique(is) @@ -761,7 +761,7 @@ end Return a view of the TensorStorage of the ITensor. """ -storage(T::ITensor)::TensorStorage = storage(tensor(T)) +storage(T::ITensor) = storage(tensor(T)) storagetype(x::ITensor) = storagetype(tensor(x)) From ea0f602cd24307913b797c23b398914f8cd40e51 Mon Sep 17 00:00:00 2001 From: Matt Fishman Date: Thu, 28 Sep 2023 20:34:05 -0400 Subject: [PATCH 2/2] [NDTensors] Add `SmallVectors` submodule (#1202) --- NDTensors/src/NDTensors.jl | 2 + NDTensors/src/SmallVectors/README.md | 73 +++++ .../src/SmallVectors/src/SmallVectors.jl | 16 + .../abstractsmallvector.jl | 28 ++ .../src/abstractsmallvector/deque.jl | 295 ++++++++++++++++++ .../src/msmallvector/msmallvector.jl | 75 +++++ .../src/smallvector/smallvector.jl | 71 +++++ .../src/subsmallvector/subsmallvector.jl | 80 +++++ NDTensors/src/SmallVectors/test/runtests.jl | 155 +++++++++ NDTensors/test/SmallVectors.jl | 4 + NDTensors/test/runtests.jl | 1 + 11 files changed, 800 insertions(+) create mode 100644 NDTensors/src/SmallVectors/README.md create mode 100644 NDTensors/src/SmallVectors/src/SmallVectors.jl create mode 100644 NDTensors/src/SmallVectors/src/abstractsmallvector/abstractsmallvector.jl create mode 100644 NDTensors/src/SmallVectors/src/abstractsmallvector/deque.jl create mode 100644 NDTensors/src/SmallVectors/src/msmallvector/msmallvector.jl create mode 100644 NDTensors/src/SmallVectors/src/smallvector/smallvector.jl create mode 100644 NDTensors/src/SmallVectors/src/subsmallvector/subsmallvector.jl create mode 100644 NDTensors/src/SmallVectors/test/runtests.jl create mode 100644 NDTensors/test/SmallVectors.jl diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl index d2ffcafb59..38d5f53a05 100644 --- a/NDTensors/src/NDTensors.jl +++ b/NDTensors/src/NDTensors.jl @@ -19,6 +19,8 @@ using TupleTools include("SetParameters/src/SetParameters.jl") using .SetParameters +include("SmallVectors/src/SmallVectors.jl") +using .SmallVectors using Base: @propagate_inbounds, ReshapedArray, DimOrInd, OneTo diff --git a/NDTensors/src/SmallVectors/README.md b/NDTensors/src/SmallVectors/README.md new file mode 100644 index 0000000000..71ee904af7 --- /dev/null +++ b/NDTensors/src/SmallVectors/README.md @@ -0,0 +1,73 @@ +# SmallVectors + +## Introduction + +A module that defines small (mutable and immutable) vectors with a maximum length. Externally they have a dynamic/runtime length, but internally they are backed by a statically sized vector. This makes it so that operations can be performed faster because they can remain on the stack, but it provides some more convenience compared to StaticArrays.jl where the length is encoded in the type. + +## Examples + +For example: +```julia +using NDTensors.SmallVectors + +mv = MSmallVector{10}([1, 2, 3]) # Mutable vector with length 3, maximum length 10 +push!(mv, 4) +mv[2] = 12 +sort!(mv; rev=true) + +v = SmallVector{10}([1, 2, 3]) # Immutable vector with length 3, maximum length 10 +v = SmallVectors.push(v, 4) +v = SmallVectors.setindex(v, 12, 2) +v = SmallVectors.sort(v; rev=true) +``` +This also has the advantage that you can efficiently store collections of `SmallVector`/`MSmallVector` that have different runtime lengths, as long as they have the same maximum length. + +## List of functionality + +`SmallVector` and `MSmallVector` are subtypes of `AbstractVector` and therefore can be used in `Base` `AbstractVector` functions, though `SmallVector` will fail for mutating functions like `setindex!` because it is immutable. + +`MSmallVector` has specialized implementations of `Base` functions that involve resizing such as: +- `resize!` +- `push!` +- `pushfirst!` +- `pop!` +- `popfirst!` +- `append!` +- `prepend!` +- `insert!` +- `deleteat!` +which are guaranteed to not realocate memory, and instead just use the memory buffer that already exists, unlike Base's `Vector` which may have to reallocate memory depending on the operation. However, they will error if they involve operations that resize beyond the maximum length of the `MSmallVector`, which you can access with `SmallVectors.maxlength(v)`. + +In addition, `SmallVector` and `MSmallVector` implement basic non-mutating operations such as: +- `SmallVectors.setindex` +, non-mutating resizing operations: +- `SmallVector.resize` +- `SmallVector.push` +- `SmallVector.pushfirst` +- `SmallVector.pop` +- `SmallVector.popfirst` +- `SmallVector.append` +- `SmallVector.prepend` +- `SmallVector.insert` +- `SmallVector.deleteat` +which output a new vector. In addition, it implements: +- `SmallVectors.circshift` +- `sort` (overloaded from `Base`). + +Finally, it provides some new helpful functions that are not in `Base`: +- `SmallVectors.insertsorted[!]` +- `SmallVectors.insertsortedunique[!]` +- `SmallVectors.mergesorted[!]` +- `SmallVectors.mergesortedunique[!]` + +## TODO + +Add specialized overloads for: +- `splice[!]` +- `union[!]` (`∪`) +- `intersect[!]` (`∩`) +- `setdiff[!]` +- `symdiff[!]` +- `unique[!]` + +Please let us know if there are other operations that would warrant specialized implmentations for `AbstractSmallVector`. diff --git a/NDTensors/src/SmallVectors/src/SmallVectors.jl b/NDTensors/src/SmallVectors/src/SmallVectors.jl new file mode 100644 index 0000000000..e3f7083795 --- /dev/null +++ b/NDTensors/src/SmallVectors/src/SmallVectors.jl @@ -0,0 +1,16 @@ +module SmallVectors +using StaticArrays + +export SmallVector, MSmallVector, SubSmallVector + +struct NotImplemented <: Exception + msg::String +end +NotImplemented() = NotImplemented("Not implemented.") + +include("abstractsmallvector/abstractsmallvector.jl") +include("abstractsmallvector/deque.jl") +include("msmallvector/msmallvector.jl") +include("smallvector/smallvector.jl") +include("subsmallvector/subsmallvector.jl") +end diff --git a/NDTensors/src/SmallVectors/src/abstractsmallvector/abstractsmallvector.jl b/NDTensors/src/SmallVectors/src/abstractsmallvector/abstractsmallvector.jl new file mode 100644 index 0000000000..014f0e0d9b --- /dev/null +++ b/NDTensors/src/SmallVectors/src/abstractsmallvector/abstractsmallvector.jl @@ -0,0 +1,28 @@ +""" +A vector with a fixed maximum length, backed by a fixed size buffer. +""" +abstract type AbstractSmallVector{T} <: AbstractVector{T} end + +# Required buffer interface +buffer(vec::AbstractSmallVector) = throw(NotImplemented()) + +similar_type(vec::AbstractSmallVector) = typeof(vec) + +# Required buffer interface +maxlength(vec::AbstractSmallVector) = length(buffer(vec)) + +# Required AbstractArray interface +Base.size(vec::AbstractSmallVector) = throw(NotImplemented()) + +# Derived AbstractArray interface +function Base.getindex(vec::AbstractSmallVector, index::Integer) + return throw(NotImplemented()) +end +function Base.setindex!(vec::AbstractSmallVector, item, index::Integer) + return throw(NotImplemented()) +end +Base.IndexStyle(::Type{<:AbstractSmallVector}) = IndexLinear() + +function Base.convert(::Type{T}, a::AbstractArray) where {T<:AbstractSmallVector} + return a isa T ? a : T(a)::T +end diff --git a/NDTensors/src/SmallVectors/src/abstractsmallvector/deque.jl b/NDTensors/src/SmallVectors/src/abstractsmallvector/deque.jl new file mode 100644 index 0000000000..96c1e80640 --- /dev/null +++ b/NDTensors/src/SmallVectors/src/abstractsmallvector/deque.jl @@ -0,0 +1,295 @@ +# TODO: add +# splice[!] +# union[!] (∪) +# intersect[!] (∩) +# setdiff[!] +# symdiff[!] +# unique[!] + +# unionsorted[!] +# setdiffsorted[!] +# deletesorted[!] (delete all or one?) +# deletesortedfirst[!] (delete all or one?) + +Base.resize!(vec::AbstractSmallVector, len) = throw(NotImplemented()) + +@inline function resize(vec::AbstractSmallVector, len) + mvec = Base.copymutable(vec) + resize!(mvec, len) + return convert(similar_type(vec), mvec) +end + +@inline function Base.empty!(vec::AbstractSmallVector) + resize!(vec, 0) + return vec +end + +@inline function empty(vec::AbstractSmallVector) + mvec = Base.copymutable(vec) + empty!(mvec) + return convert(similar_type(vec), mvec) +end + +@inline function StaticArrays.setindex(vec::AbstractSmallVector, item, index::Integer) + @boundscheck checkbounds(vec, index) + mvec = Base.copymutable(vec) + @inbounds mvec[index] = item + return convert(similar_type(vec), mvec) +end + +@inline function Base.push!(vec::AbstractSmallVector, item) + resize!(vec, length(vec) + 1) + @inbounds vec[length(vec)] = item + return vec +end + +@inline function StaticArrays.push(vec::AbstractSmallVector, item) + mvec = Base.copymutable(vec) + push!(mvec, item) + return convert(similar_type(vec), mvec) +end + +@inline function Base.pop!(vec::AbstractSmallVector) + resize!(vec, length(vec) - 1) + return vec +end + +@inline function StaticArrays.pop(vec::AbstractSmallVector) + mvec = Base.copymutable(vec) + pop!(mvec) + return convert(similar_type(vec), mvec) +end + +@inline function Base.pushfirst!(vec::AbstractSmallVector, item) + insert!(vec, firstindex(vec), item) + return vec +end + +# Don't `@inline`, makes it slower. +function StaticArrays.pushfirst(vec::AbstractSmallVector, item) + mvec = Base.copymutable(vec) + pushfirst!(mvec, item) + return convert(similar_type(vec), mvec) +end + +@inline function Base.popfirst!(vec::AbstractSmallVector) + circshift!(vec, -1) + resize!(vec, length(vec) - 1) + return vec +end + +# Don't `@inline`, makes it slower. +function StaticArrays.popfirst(vec::AbstractSmallVector) + mvec = Base.copymutable(vec) + popfirst!(mvec) + return convert(similar_type(vec), mvec) +end + +# This implementation of `midpoint` is performance-optimized but safe +# only if `lo <= hi`. +# TODO: Replace with `Base.midpoint`. +midpoint(lo::T, hi::T) where {T<:Integer} = lo + ((hi - lo) >>> 0x01) +midpoint(lo::Integer, hi::Integer) = midpoint(promote(lo, hi)...) + +@inline function Base.reverse!(vec::AbstractSmallVector) + start, stop = firstindex(vec), lastindex(vec) + r = stop + @inbounds for i in start:midpoint(start, stop - 1) + vec[i], vec[r] = vec[r], vec[i] + r -= 1 + end + return vec +end + +@inline function Base.reverse!( + vec::AbstractSmallVector, start::Integer, stop::Integer=lastindex(v) +) + reverse!(smallview(vec, start, stop)) + return vec +end + +@inline function Base.circshift!(vec::AbstractSmallVector, shift::Integer) + start, stop = firstindex(vec), lastindex(vec) + n = length(vec) + n == 0 && return vec + shift = mod(shift, n) + shift == 0 && return vec + reverse!(smallview(vec, start, stop - shift)) + reverse!(smallview(vec, stop - shift + 1, stop)) + reverse!(smallview(vec, start, stop)) + return vec +end + +@inline function Base.insert!(vec::AbstractSmallVector, index::Integer, item) + resize!(vec, length(vec) + 1) + circshift!(smallview(vec, index, lastindex(vec)), 1) + @inbounds vec[index] = item + return vec +end + +# Don't @inline, makes it slower. +function StaticArrays.insert(vec::AbstractSmallVector, index::Integer, item) + mvec = Base.copymutable(vec) + insert!(mvec, index, item) + return convert(similar_type(vec), mvec) +end + +@inline function Base.deleteat!(vec::AbstractSmallVector, index::Integer) + circshift!(smallview(vec, index, lastindex(vec)), -1) + resize!(vec, length(vec) - 1) + return vec +end + +# Don't @inline, makes it slower. +function StaticArrays.deleteat(vec::AbstractSmallVector, index::Integer) + mvec = Base.copymutable(vec) + deleteat!(mvec, index) + return convert(similar_type(vec), mvec) +end + +# InsertionSortAlg +# https://github.com/JuliaLang/julia/blob/bed2cd540a11544ed4be381d471bbf590f0b745e/base/sort.jl#L722-L736 +# https://en.wikipedia.org/wiki/Insertion_sort#:~:text=Insertion%20sort%20is%20a%20simple,%2C%20heapsort%2C%20or%20merge%20sort. +# Alternatively could use `TupleTools.jl` or `StaticArrays.jl` for out-of-place sorting. +@inline function Base.sort!( + vec::AbstractSmallVector; lt=isless, by=identity, rev::Bool=false +) + lo, hi = firstindex(vec), lastindex(vec) + lo_plus_1 = (lo + 1) + @inbounds for i in lo_plus_1:hi + j = i + x = vec[i] + jmax = j + for _ in jmax:-1:lo_plus_1 + y = vec[j - 1] + if !(lt(by(x), by(y)) != rev) + break + end + vec[j] = y + j -= 1 + end + vec[j] = x + end + return vec +end + +# Don't @inline, makes it slower. +function Base.sort(vec::AbstractSmallVector; kwargs...) + mvec = Base.copymutable(vec) + sort!(mvec; kwargs...) + return convert(similar_type(vec), mvec) +end + +@inline function insertsorted!(vec::AbstractSmallVector, item; kwargs...) + insert!(vec, searchsortedfirst(vec, item; kwargs...), item) + return vec +end + +function insertsorted(vec::AbstractSmallVector, item; kwargs...) + mvec = Base.copymutable(vec) + insertsorted!(mvec, item; kwargs...) + return convert(similar_type(vec), mvec) +end + +@inline function insertsortedunique!(vec::AbstractSmallVector, item; kwargs...) + r = searchsorted(vec, item; kwargs...) + if length(r) == 0 + insert!(vec, first(r), item) + end + return vec +end + +# Code repeated since inlining doesn't work. +function insertsortedunique(vec::AbstractSmallVector, item; kwargs...) + r = searchsorted(vec, item; kwargs...) + if length(r) == 0 + vec = insert(vec, first(r), item) + end + return vec +end + +@inline function mergesorted!(vec::AbstractSmallVector, item::AbstractVector; kwargs...) + for x in item + insertsorted!(vec, x; kwargs...) + end + return vec +end + +function mergesorted(vec::AbstractSmallVector, item; kwargs...) + mvec = Base.copymutable(vec) + mergesorted!(mvec, item; kwargs...) + return convert(similar_type(vec), mvec) +end + +@inline function mergesortedunique!( + vec::AbstractSmallVector, item::AbstractVector; kwargs... +) + for x in item + insertsortedunique!(vec, x; kwargs...) + end + return vec +end + +# Code repeated since inlining doesn't work. +function mergesortedunique(vec::AbstractSmallVector, item; kwargs...) + for x in item + vec = insertsortedunique(vec, x; kwargs...) + end + return vec +end + +Base.@propagate_inbounds function Base.copyto!( + vec::AbstractSmallVector, item::AbstractVector +) + for i in eachindex(item) + vec[i] = item[i] + end + return vec +end + +# Don't @inline, makes it slower. +function Base.circshift(vec::AbstractSmallVector, shift::Integer) + mvec = Base.copymutable(vec) + circshift!(mvec, shift) + return convert(similar_type(vec), mvec) +end + +@inline function Base.append!(vec::AbstractSmallVector, item::AbstractVector) + l = length(vec) + r = length(item) + resize!(vec, l + r) + @inbounds copyto!(smallview(vec, l + 1, l + r + 1), item) + return vec +end + +# Missing from `StaticArrays.jl`. +# Don't @inline, makes it slower. +function append(vec::AbstractSmallVector, item::AbstractVector) + mvec = Base.copymutable(vec) + append!(mvec, item) + return convert(similar_type(vec), mvec) +end + +@inline function Base.prepend!(vec::AbstractSmallVector, item::AbstractVector) + l = length(vec) + r = length(item) + resize!(vec, l + r) + circshift!(vec, length(item)) + @inbounds copyto!(vec, item) + return vec +end + +# Missing from `StaticArrays.jl`. +# Don't @inline, makes it slower. +function prepend(vec::AbstractSmallVector, item::AbstractVector) + mvec = Base.copymutable(vec) + prepend!(mvec, item) + return convert(similar_type(vec), mvec) +end + +# Don't @inline, makes it slower. +function Base.vcat(vec1::AbstractSmallVector, vec2::AbstractVector) + mvec1 = Base.copymutable(vec1) + append!(mvec1, vec2) + return convert(similar_type(vec1), mvec1) +end diff --git a/NDTensors/src/SmallVectors/src/msmallvector/msmallvector.jl b/NDTensors/src/SmallVectors/src/msmallvector/msmallvector.jl new file mode 100644 index 0000000000..434d1e4d48 --- /dev/null +++ b/NDTensors/src/SmallVectors/src/msmallvector/msmallvector.jl @@ -0,0 +1,75 @@ +""" +MSmallVector + +TODO: Make `buffer` field `const` (new in Julia 1.8) +""" +mutable struct MSmallVector{S,T} <: AbstractSmallVector{T} + buffer::MVector{S,T} + length::Int +end + +# Constructors +function MSmallVector{S}(buffer::AbstractVector, len::Int) where {S} + return MSmallVector{S,eltype(buffer)}(buffer, len) +end +function MSmallVector(buffer::AbstractVector, len::Int) + return MSmallVector{length(buffer),eltype(buffer)}(buffer, len) +end + +""" +`MSmallVector` constructor, uses `MVector` as a buffer. +```julia +MSmallVector{10}([1, 2, 3]) +MSmallVector{10}(SA[1, 2, 3]) +``` +""" +function MSmallVector{S,T}(vec::AbstractVector) where {S,T} + buffer = MVector{S,T}(undef) + copyto!(buffer, vec) + return MSmallVector(buffer, length(vec)) +end + +# Derive the buffer length. +MSmallVector(vec::AbstractSmallVector) = MSmallVector{length(buffer(vec))}(vec) + +function MSmallVector{S}(vec::AbstractVector) where {S} + return MSmallVector{S,eltype(vec)}(vec) +end + +function MSmallVector{S,T}(::UndefInitializer, dims::Tuple{Integer}) where {S,T} + return MSmallVector{S,T}(undef, prod(dims)) +end +function MSmallVector{S,T}(::UndefInitializer, length::Integer) where {S,T} + return MSmallVector{S,T}(MVector{S,T}(undef), length) +end + +# Buffer interface +buffer(vec::MSmallVector) = vec.buffer + +# Accessors +Base.size(vec::MSmallVector) = (vec.length,) + +# Required Base overloads +@inline function Base.getindex(vec::MSmallVector, index::Integer) + @boundscheck checkbounds(vec, index) + return @inbounds buffer(vec)[index] +end + +@inline function Base.setindex!(vec::MSmallVector, item, index::Integer) + @boundscheck checkbounds(vec, index) + @inbounds buffer(vec)[index] = item + return vec +end + +@inline function Base.resize!(vec::MSmallVector, len::Integer) + len < 0 && throw(ArgumentError("New length must be ≥ 0.")) + len > maxlength(vec) && + throw(ArgumentError("New length $len must be ≤ the maximum length $(maxlength(vec)).")) + vec.length = len + return vec +end + +# `similar` creates a `MSmallVector` by default. +function Base.similar(vec::AbstractSmallVector, elt::Type, dims::Dims) + return MSmallVector{length(buffer(vec)),elt}(undef, dims) +end diff --git a/NDTensors/src/SmallVectors/src/smallvector/smallvector.jl b/NDTensors/src/SmallVectors/src/smallvector/smallvector.jl new file mode 100644 index 0000000000..3976480f47 --- /dev/null +++ b/NDTensors/src/SmallVectors/src/smallvector/smallvector.jl @@ -0,0 +1,71 @@ +""" +SmallVector +""" +struct SmallVector{S,T} <: AbstractSmallVector{T} + buffer::SVector{S,T} + length::Int +end + +# Accessors +# TODO: Use `Accessors.jl`. +@inline setbuffer(vec::SmallVector, buffer) = SmallVector(buffer, vec.length) +@inline setlength(vec::SmallVector, length) = SmallVector(vec.buffer, length) + +# Constructors +function SmallVector{S}(buffer::AbstractVector, len::Int) where {S} + return SmallVector{S,eltype(buffer)}(buffer, len) +end +function SmallVector(buffer::AbstractVector, len::Int) + return SmallVector{length(buffer),eltype(buffer)}(buffer, len) +end + +""" +`SmallVector` constructor, uses `SVector` as buffer storage. +```julia +SmallVector{10}([1, 2, 3]) +SmallVector{10}(SA[1, 2, 3]) +``` +""" +function SmallVector{S,T}(vec::AbstractVector) where {S,T} + return SmallVector{S,T}(MSmallVector{S,T}(vec)) +end +# Special optimization codepath for `MSmallVector` +# to avoid a copy. +function SmallVector{S,T}(vec::MSmallVector) where {S,T} + return SmallVector{S,T}(buffer(vec), length(vec)) +end + +function SmallVector{S}(vec::AbstractVector) where {S} + return SmallVector{S,eltype(vec)}(vec) +end + +# Specialized constructor +function MSmallVector{S,T}(vec::SmallVector) where {S,T} + return MSmallVector{S,T}(buffer(vec), length(vec)) +end + +# Derive the buffer length. +SmallVector(vec::AbstractSmallVector) = SmallVector{length(buffer(vec))}(vec) + +# Empty constructor +(smallvector_type::Type{SmallVector{S,T}} where {S,T})() = smallvector_type(undef, 0) +function SmallVector{S,T}(::UndefInitializer, length::Integer) where {S,T} + return SmallVector{S,T}(SVector{S,T}(MVector{S,T}(undef)), length) +end + +# Buffer interface +buffer(vec::SmallVector) = vec.buffer + +# AbstractArray interface +Base.size(vec::SmallVector) = (vec.length,) + +# Base overloads +@inline function Base.getindex(vec::SmallVector, index::Integer) + @boundscheck checkbounds(vec, index) + return @inbounds buffer(vec)[index] +end + +Base.copy(vec::SmallVector) = vec + +# Optimization, default uses `similar`. +Base.copymutable(vec::SmallVector) = MSmallVector(vec) diff --git a/NDTensors/src/SmallVectors/src/subsmallvector/subsmallvector.jl b/NDTensors/src/SmallVectors/src/subsmallvector/subsmallvector.jl new file mode 100644 index 0000000000..922eaa3521 --- /dev/null +++ b/NDTensors/src/SmallVectors/src/subsmallvector/subsmallvector.jl @@ -0,0 +1,80 @@ +abstract type AbstractSubSmallVector{T} <: AbstractSmallVector{T} end + +""" +SubSmallVector +""" +struct SubSmallVector{T,P} <: AbstractSubSmallVector{T} + parent::P + start::Int + stop::Int +end + +mutable struct SubMSmallVector{T,P<:AbstractVector{T}} <: AbstractSubSmallVector{T} + parent::P + start::Int + stop::Int +end + +# TODO: Use Accessors.jl +Base.parent(vec::SubSmallVector) = vec.parent +Base.parent(vec::SubMSmallVector) = vec.parent + +# buffer interface +buffer(vec::AbstractSubSmallVector) = buffer(parent(vec)) + +function smallview(vec::SmallVector, start::Integer, stop::Integer) + return SubSmallVector(vec, start, stop) +end +function smallview(vec::MSmallVector, start::Integer, stop::Integer) + return SubMSmallVector(vec, start, stop) +end + +function smallview(vec::SubSmallVector, start::Integer, stop::Integer) + return SubSmallVector(parent(vec), vec.start + start - 1, vec.start + stop - 1) +end +function smallview(vec::SubMSmallVector, start::Integer, stop::Integer) + return SubMSmallVector(parent(vec), vec.start + start - 1, vec.start + stop - 1) +end + +# Constructors +function SubSmallVector(vec::AbstractVector, start::Integer, stop::Integer) + return SubSmallVector{eltype(vec),typeof(vec)}(vec, start, stop) +end +function SubMSmallVector(vec::AbstractVector, start::Integer, stop::Integer) + return SubMSmallVector{eltype(vec),typeof(vec)}(vec, start, stop) +end + +# Accessors +Base.size(vec::AbstractSubSmallVector) = (vec.stop - vec.start + 1,) + +Base.@propagate_inbounds function Base.getindex(vec::AbstractSubSmallVector, index::Integer) + return parent(vec)[index + vec.start - 1] +end + +Base.@propagate_inbounds function Base.setindex!( + vec::AbstractSubSmallVector, item, index::Integer +) + buffer(vec)[index + vec.start - 1] = item + return vec +end + +function SubSmallVector{T,P}(vec::SubMSmallVector) where {T,P} + return SubSmallVector{T,P}(P(parent(vec)), vec.start, vec.stop) +end + +function Base.convert(smalltype::Type{<:SubSmallVector}, vec::SubMSmallVector) + return smalltype(vec) +end + +@inline function Base.resize!(vec::SubMSmallVector, len::Integer) + len < 0 && throw(ArgumentError("New length must be ≥ 0.")) + len > maxlength(vec) - vec.start + 1 && + throw(ArgumentError("New length $len must be ≤ the maximum length $(maxlength(vec)).")) + vec.stop = vec.start + len - 1 + return vec +end + +# Optimization, default uses `similar`. +function Base.copymutable(vec::SubSmallVector) + return SubMSmallVector(Base.copymutable(parent(vec)), vec.start, vec.stop) +end diff --git a/NDTensors/src/SmallVectors/test/runtests.jl b/NDTensors/src/SmallVectors/test/runtests.jl new file mode 100644 index 0000000000..4b1daac154 --- /dev/null +++ b/NDTensors/src/SmallVectors/test/runtests.jl @@ -0,0 +1,155 @@ +using NDTensors.SmallVectors +using Test + +using NDTensors.SmallVectors: + setindex, + resize, + push, + pushfirst, + pop, + popfirst, + append, + prepend, + insert, + deleteat, + circshift, + insertsorted, + insertsorted!, + insertsortedunique, + insertsortedunique!, + mergesorted, + mergesorted!, + mergesortedunique, + mergesortedunique! + +function test_smallvectors() + return @testset "SmallVectors" begin + x = SmallVector{10}([1, 3, 5]) + mx = MSmallVector(x) + + @test x isa SmallVector{10,Int} + @test mx isa MSmallVector{10,Int} + @test eltype(x) === Int + @test eltype(mx) === Int + + # TODO: Test construction has zero allocations. + # TODO: Extend construction to arbitrary collections, like tuple. + + # conversion + @test @inferred(SmallVector(x)) == x + @test @allocated(SmallVector(x)) == 0 + @test @inferred(SmallVector(mx)) == x + @test @allocated(SmallVector(mx)) == 0 + + # length + @test @inferred(length(x)) == 3 + @test @allocated(length(x)) == 0 + @test @inferred(length(SmallVectors.buffer(x))) == 10 + @test @allocated(length(SmallVectors.buffer(x))) == 0 + + nalloc_limit = 128 + + item = 115 + no_broken = (false, false, false, false) + for ( + f!, + f, + ans, + args, + nalloc, + f!_impl_broken, + f!_noalloc_broken, + f_impl_broken, + f_noalloc_broken, + ) in [ + (:push!, :push, [1, 3, 5, item], (item,), nalloc_limit, no_broken...), + (:append!, :append, [1, 3, 5, item], ([item],), nalloc_limit, no_broken...), + (:prepend!, :prepend, [item, 1, 3, 5], ([item],), nalloc_limit, no_broken...), + (:pushfirst!, :pushfirst, [item, 1, 3, 5], (item,), nalloc_limit, no_broken...), + (:setindex!, :setindex, [1, item, 5], (item, 2), nalloc_limit, no_broken...), + (:pop!, :pop, [1, 3], (), nalloc_limit, no_broken...), + (:popfirst!, :popfirst, [3, 5], (), nalloc_limit, no_broken...), + (:insert!, :insert, [1, item, 3, 5], (2, item), nalloc_limit, no_broken...), + (:deleteat!, :deleteat, [1, 5], (2,), nalloc_limit, no_broken...), + (:circshift!, :circshift, [5, 1, 3], (1,), nalloc_limit, no_broken...), + (:sort!, :sort, [1, 3, 5], (), nalloc_limit, no_broken...), + (:insertsorted!, :insertsorted, [1, 2, 3, 5], (2,), nalloc_limit, no_broken...), + (:insertsorted!, :insertsorted, [1, 3, 3, 5], (3,), nalloc_limit, no_broken...), + ( + :insertsortedunique!, + :insertsortedunique, + [1, 2, 3, 5], + (2,), + nalloc_limit, + no_broken..., + ), + ( + :insertsortedunique!, + :insertsortedunique, + [1, 3, 5], + (3,), + nalloc_limit, + no_broken..., + ), + (:mergesorted!, :mergesorted, [1, 2, 3, 3, 5], ([2, 3],), nalloc_limit, no_broken...), + ( + :mergesortedunique!, + :mergesortedunique, + [1, 2, 3, 5], + ([2, 3],), + nalloc_limit, + no_broken..., + ), + ] + mx_tmp = copy(mx) + @eval begin + if VERSION < v"1.7" + # broken kwarg wasn't added to @test yet + if $f!_impl_broken + @test_broken @inferred($f!(copy($mx), $args...)) == $ans + else + @test @inferred($f!(copy($mx), $args...)) == $ans + end + if $f!_noalloc_broken + @test_broken @allocated($f!($mx_tmp, $args...)) ≤ $nalloc + else + @test @allocated($f!($mx_tmp, $args...)) ≤ $nalloc + end + if $f_impl_broken + @test_broken @inferred($f($x, $args...)) == $ans + else + @test @inferred($f($x, $args...)) == $ans + end + if $f_noalloc_broken + @test_broken @allocated($f($x, $args...)) ≤ $nalloc + else + @test @allocated($f($x, $args...)) ≤ $nalloc + end + else + @test @inferred($f!(copy($mx), $args...)) == $ans broken = $f!_impl_broken + @test @allocated($f!($mx_tmp, $args...)) ≤ $nalloc broken = $f!_noalloc_broken + @test @inferred($f($x, $args...)) == $ans broken = $f_impl_broken + @test @allocated($f($x, $args...)) ≤ $nalloc broken = $f_noalloc_broken + end + end + end + + # Separated out since for some reason it breaks the `@inferred` + # check when `kwargs` are interpolated into `@eval`. + ans, kwargs = [5, 3, 1], (; rev=true) + mx_tmp = copy(mx) + @test @inferred(sort!(copy(mx); kwargs...)) == ans + @test @allocated(sort!(mx_tmp; kwargs...)) == 0 + @test @inferred(sort(x; kwargs...)) == ans + @test @allocated(sort(x; kwargs...)) ≤ nalloc_limit + + ans, args = [1, 3, 5, item], ([item],) + @test @inferred(vcat(x, args...)) == ans + @test @allocated(vcat(x, args...)) ≤ nalloc_limit + end +end + +# TODO: switch to: +# @testset "SmallVectors" test_smallvectors() +# (new in Julia 1.9) +test_smallvectors() diff --git a/NDTensors/test/SmallVectors.jl b/NDTensors/test/SmallVectors.jl new file mode 100644 index 0000000000..62b552dc72 --- /dev/null +++ b/NDTensors/test/SmallVectors.jl @@ -0,0 +1,4 @@ +using Test +using NDTensors + +include(joinpath(pkgdir(NDTensors), "src", "SmallVectors", "test", "runtests.jl")) diff --git a/NDTensors/test/runtests.jl b/NDTensors/test/runtests.jl index cdd0d6cf35..8e4f1733f8 100644 --- a/NDTensors/test/runtests.jl +++ b/NDTensors/test/runtests.jl @@ -20,6 +20,7 @@ end @safetestset "NDTensors" begin @testset "$filename" for filename in [ "SetParameters.jl", + "SmallVectors.jl", "linearalgebra.jl", "dense.jl", "blocksparse.jl",