diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml index e760a43e53..9031af6d04 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.4.7" +version = "0.4.8" [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" @@ -19,7 +19,6 @@ 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" SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" @@ -79,7 +78,6 @@ MacroTools = "0.5" MappedArrays = "0.4" Metal = "1" Octavian = "0.3" -PackageExtensionCompat = "1" Random = "<0.0.1, 1.10" SimpleTraits = "0.9.4" SparseArrays = "<0.0.1, 1.10" diff --git a/NDTensors/src/NDTensors.jl b/NDTensors/src/NDTensors.jl index 4b6abe8ef1..7b2bf2b54a 100644 --- a/NDTensors/src/NDTensors.jl +++ b/NDTensors/src/NDTensors.jl @@ -72,8 +72,7 @@ include("blocksparse/contract.jl") include("blocksparse/contract_utilities.jl") include("blocksparse/contract_generic.jl") include("blocksparse/contract_sequential.jl") -include("blocksparse/contract_folds.jl") -include("blocksparse/contract_threads.jl") +include("blocksparse/contract_threaded.jl") include("blocksparse/diagblocksparse.jl") include("blocksparse/similar.jl") include("blocksparse/combiner.jl") @@ -221,9 +220,4 @@ end function backend_octavian end -using PackageExtensionCompat -function __init__() - @require_extensions -end - end # module NDTensors diff --git a/NDTensors/src/blocksparse/contract.jl b/NDTensors/src/blocksparse/contract.jl index bfe413e7d1..ac8c0771b2 100644 --- a/NDTensors/src/blocksparse/contract.jl +++ b/NDTensors/src/blocksparse/contract.jl @@ -46,8 +46,6 @@ function contract_blockoffsets( alg = Algorithm"sequential"() if using_threaded_blocksparse() && nthreads() > 1 alg = Algorithm"threaded_threads"() - # This code is a bit cleaner but slower: - # alg = Algorithm"threaded_folds"() end return contract_blockoffsets( alg, boffs1, inds1, labels1, boffs2, inds2, labels2, indsR, labelsR diff --git a/NDTensors/src/blocksparse/contract_folds.jl b/NDTensors/src/blocksparse/contract_folds.jl deleted file mode 100644 index 9e483e2eb1..0000000000 --- a/NDTensors/src/blocksparse/contract_folds.jl +++ /dev/null @@ -1,60 +0,0 @@ -# Function barrier to improve type stability, -# since `Folds`/`FLoops` is not type stable: -# https://discourse.julialang.org/t/type-instability-in-floop-reduction/68598 -function contract_blocks!( - alg::Algorithm"threaded_folds", - contraction_plans, - boffs1, - boffs2, - labels1_to_labels2, - labels1_to_labelsR, - labels2_to_labelsR, - ValNR, -) - if nnzblocks(boffs1) > nnzblocks(boffs2) - Folds.foreach(eachnzblock(boffs1).values, ThreadedEx()) do block1 - for block2 in eachnzblock(boffs2) - maybe_contract_blocks!( - contraction_plans[threadid()], - block1, - block2, - labels1_to_labels2, - labels1_to_labelsR, - labels2_to_labelsR, - ValNR, - ) - end - end - else - Folds.foreach(eachnzblock(boffs2).values, ThreadedEx()) do block2 - for block1 in eachnzblock(boffs1) - maybe_contract_blocks!( - contraction_plans[threadid()], - block1, - block2, - labels1_to_labels2, - labels1_to_labelsR, - labels2_to_labelsR, - ValNR, - ) - end - end - end - return nothing -end - -function contract!( - ::Algorithm"threaded_folds", - R::BlockSparseTensor, - labelsR, - tensor1::BlockSparseTensor, - labelstensor1, - tensor2::BlockSparseTensor, - labelstensor2, - contraction_plan, -) - executor = ThreadedEx() - return contract!( - R, labelsR, tensor1, labelstensor1, tensor2, labelstensor2, contraction_plan, executor - ) -end diff --git a/NDTensors/src/blocksparse/contract_generic.jl b/NDTensors/src/blocksparse/contract_generic.jl index e8e59919c7..873e1220f4 100644 --- a/NDTensors/src/blocksparse/contract_generic.jl +++ b/NDTensors/src/blocksparse/contract_generic.jl @@ -11,47 +11,14 @@ function contract_blockoffsets( indsR, labelsR, ) - N1 = length(blocktype(boffs1)) - N2 = length(blocktype(boffs2)) NR = length(labelsR) ValNR = ValLength(labelsR) labels1_to_labels2, labels1_to_labelsR, labels2_to_labelsR = contract_labels( labels1, labels2, labelsR ) - - # Contraction plan element type - T = Tuple{Block{N1},Block{N2},Block{NR}} - - # Thread-local collections of block contractions. - # Could use: - # ```julia - # FLoops.@reduce(contraction_plans = append!(T[], [(block1, block2, blockR)])) - # ``` - # as a simpler alternative but it is slower. - - contraction_plans = Vector{T}[T[] for _ in 1:nthreads()] - - # - # Reserve some capacity - # In theory the maximum is length(boffs1) * length(boffs2) - # but in practice that is too much - #for contraction_plan in contraction_plans - # sizehint!(contraction_plan, max(length(boffs1), length(boffs2))) - #end - # - - contract_blocks!( - alg, - contraction_plans, - boffs1, - boffs2, - labels1_to_labels2, - labels1_to_labelsR, - labels2_to_labelsR, - ValNR, + contraction_plan = contract_blocks( + alg, boffs1, boffs2, labels1_to_labels2, labels1_to_labelsR, labels2_to_labelsR, ValNR ) - - contraction_plan = reduce(vcat, contraction_plans) blockoffsetsR = BlockOffsets{NR}() nnzR = 0 for (_, _, blockR) in contraction_plan @@ -60,7 +27,6 @@ function contract_blockoffsets( nnzR += blockdim(indsR, blockR) end end - return blockoffsetsR, contraction_plan end diff --git a/NDTensors/src/blocksparse/contract_threaded.jl b/NDTensors/src/blocksparse/contract_threaded.jl new file mode 100644 index 0000000000..9a096b2676 --- /dev/null +++ b/NDTensors/src/blocksparse/contract_threaded.jl @@ -0,0 +1,91 @@ +using .Expose: expose +function contract_blocks( + alg::Algorithm"threaded_threads", + boffs1, + boffs2, + labels1_to_labels2, + labels1_to_labelsR, + labels2_to_labelsR, + ValNR::Val{NR}, +) where {NR} + N1 = length(blocktype(boffs1)) + N2 = length(blocktype(boffs2)) + blocks1 = keys(boffs1) + blocks2 = keys(boffs2) + T = Tuple{Block{N1},Block{N2},Block{NR}} + return if length(blocks1) > length(blocks2) + tasks = map( + Iterators.partition(blocks1, max(1, length(blocks1) ÷ nthreads())) + ) do blocks1_partition + @spawn begin + block_contractions = T[] + for block1 in blocks1_partition + for block2 in blocks2 + block_contraction = maybe_contract_blocks( + block1, + block2, + labels1_to_labels2, + labels1_to_labelsR, + labels2_to_labelsR, + ValNR, + ) + if !isnothing(block_contraction) + push!(block_contractions, block_contraction) + end + end + end + return block_contractions + end + end + all_block_contractions = T[] + for task in tasks + append!(all_block_contractions, fetch(task)) + end + return all_block_contractions + else + tasks = map( + Iterators.partition(blocks2, max(1, length(blocks2) ÷ nthreads())) + ) do blocks2_partition + @spawn begin + block_contractions = T[] + for block2 in blocks2_partition + for block1 in blocks1 + block_contraction = maybe_contract_blocks( + block1, + block2, + labels1_to_labels2, + labels1_to_labelsR, + labels2_to_labelsR, + ValNR, + ) + if !isnothing(block_contraction) + push!(block_contractions, block_contraction) + end + end + end + return block_contractions + end + end + all_block_contractions = T[] + for task in tasks + append!(all_block_contractions, fetch(task)) + end + return all_block_contractions + end +end + +function contract!( + ::Algorithm"threaded_folds", + R::BlockSparseTensor, + labelsR, + tensor1::BlockSparseTensor, + labelstensor1, + tensor2::BlockSparseTensor, + labelstensor2, + contraction_plan, +) + executor = ThreadedEx() + return contract!( + R, labelsR, tensor1, labelstensor1, tensor2, labelstensor2, contraction_plan, executor + ) +end diff --git a/NDTensors/src/blocksparse/contract_threads.jl b/NDTensors/src/blocksparse/contract_threads.jl deleted file mode 100644 index e80923ff55..0000000000 --- a/NDTensors/src/blocksparse/contract_threads.jl +++ /dev/null @@ -1,137 +0,0 @@ -using .Expose: expose -# TODO: This seems to be faster than the newer version using `Folds.jl` -# in `contract_folds.jl`, investigate why. -function contract_blocks!( - alg::Algorithm"threaded_threads", - contraction_plans, - boffs1, - boffs2, - labels1_to_labels2, - labels1_to_labelsR, - labels2_to_labelsR, - ValNR, -) - blocks1 = keys(boffs1) - blocks2 = keys(boffs2) - if length(blocks1) > length(blocks2) - @sync for blocks1_partition in - Iterators.partition(blocks1, max(1, length(blocks1) ÷ nthreads())) - @spawn for block1 in blocks1_partition - for block2 in blocks2 - maybe_contract_blocks!( - contraction_plans[threadid()], - block1, - block2, - labels1_to_labels2, - labels1_to_labelsR, - labels2_to_labelsR, - ValNR, - ) - end - end - end - else - @sync for blocks2_partition in - Iterators.partition(blocks2, max(1, length(blocks2) ÷ nthreads())) - @spawn for block2 in blocks2_partition - for block1 in blocks1 - maybe_contract_blocks!( - contraction_plans[threadid()], - block1, - block2, - labels1_to_labels2, - labels1_to_labelsR, - labels2_to_labelsR, - ValNR, - ) - end - end - end - end - return nothing -end - -########################################################################### -# Old version -# TODO: DELETE, keeping around for testing/benchmarking. -function contract!( - ::Algorithm"threaded_threads", - R::BlockSparseTensor{ElR,NR}, - labelsR, - T1::BlockSparseTensor{ElT1,N1}, - labelsT1, - T2::BlockSparseTensor{ElT2,N2}, - labelsT2, - contraction_plan, -) where {ElR,ElT1,ElT2,N1,N2,NR} - # Sort the contraction plan by the output blocks - # This is to help determine which output blocks are the result - # of multiple contractions - sort!(contraction_plan; by=last) - - # Ranges of contractions to the same block - repeats = Vector{UnitRange{Int}}(undef, nnzblocks(R)) - ncontracted = 1 - posR = last(contraction_plan[1]) - posR_unique = posR - for n in 1:(nnzblocks(R) - 1) - start = ncontracted - while posR == posR_unique - ncontracted += 1 - posR = last(contraction_plan[ncontracted]) - end - posR_unique = posR - repeats[n] = start:(ncontracted - 1) - end - repeats[end] = ncontracted:length(contraction_plan) - - contraction_plan_blocks = Vector{Tuple{Tensor,Tensor,Tensor}}( - undef, length(contraction_plan) - ) - for ncontracted in 1:length(contraction_plan) - block1, block2, blockR = contraction_plan[ncontracted] - T1block = T1[block1] - T2block = T2[block2] - Rblock = R[blockR] - contraction_plan_blocks[ncontracted] = (T1block, T2block, Rblock) - end - - indsR = inds(R) - indsT1 = inds(T1) - indsT2 = inds(T2) - - α = one(ElR) - @sync for repeats_partition in - Iterators.partition(repeats, max(1, length(repeats) ÷ nthreads())) - @spawn for ncontracted_range in repeats_partition - # Overwrite the block since it hasn't been written to - # R .= α .* (T1 * T2) - β = zero(ElR) - for ncontracted in ncontracted_range - blockT1, blockT2, blockR = contraction_plan_blocks[ncontracted] - # R .= α .* (T1 * T2) .+ β .* R - - # : - α = compute_alpha( - ElR, labelsR, blockR, indsR, labelsT1, blockT1, indsT1, labelsT2, blockT2, indsT2 - ) - - contract!( - expose(blockR), - labelsR, - expose(blockT1), - labelsT1, - expose(blockT2), - labelsT2, - α, - β, - ) - # Now keep adding to the block, since it has - # been written to - # R .= α .* (T1 * T2) .+ R - β = one(ElR) - end - end - end - return R -end diff --git a/NDTensors/src/blocksparse/contract_utilities.jl b/NDTensors/src/blocksparse/contract_utilities.jl index 219d6ac9b1..f2252ab87e 100644 --- a/NDTensors/src/blocksparse/contract_utilities.jl +++ b/NDTensors/src/blocksparse/contract_utilities.jl @@ -14,18 +14,12 @@ function compute_alpha( return one(ElR) end -function maybe_contract_blocks!( - contraction_plan, - block1, - block2, - labels1_to_labels2, - labels1_to_labelsR, - labels2_to_labelsR, - ValNR, +function maybe_contract_blocks( + block1, block2, labels1_to_labels2, labels1_to_labelsR, labels2_to_labelsR, ValNR ) if are_blocks_contracted(block1, block2, labels1_to_labels2) blockR = contract_blocks(block1, labels1_to_labelsR, block2, labels2_to_labelsR, ValNR) - push!(contraction_plan, (block1, block2, blockR)) + return block1, block2, blockR end return nothing end diff --git a/NDTensors/src/imports.jl b/NDTensors/src/imports.jl index 91f05ed0b9..bbfacd4a87 100644 --- a/NDTensors/src/imports.jl +++ b/NDTensors/src/imports.jl @@ -6,7 +6,6 @@ using Adapt using Base.Threads -using Compat using Dictionaries using Folds using InlineStrings diff --git a/NDTensors/test/Project.toml b/NDTensors/test/Project.toml index 53951dfa18..bfb4c9f816 100644 --- a/NDTensors/test/Project.toml +++ b/NDTensors/test/Project.toml @@ -3,7 +3,6 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a" BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" -Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4" EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" diff --git a/NDTensors/test/backup/arraytensor/Project.toml b/NDTensors/test/backup/arraytensor/Project.toml deleted file mode 100644 index 2238172387..0000000000 --- a/NDTensors/test/backup/arraytensor/Project.toml +++ /dev/null @@ -1,3 +0,0 @@ -[deps] -BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" -NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" diff --git a/NDTensors/test/backup/arraytensor/array.jl b/NDTensors/test/backup/arraytensor/array.jl deleted file mode 100644 index 95c44925a3..0000000000 --- a/NDTensors/test/backup/arraytensor/array.jl +++ /dev/null @@ -1,51 +0,0 @@ -@eval module $(gensym()) -using LinearAlgebra: svd -using NDTensors: array, contract, inds, storage, storagetype, tensor -using Random: randn! -using Test: @test, @testset, @test_broken -@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) - - # TODO: Should this work? Currently broken. - ## I was able to fix this test but labels have to match up - ## If you do U * S * V it fails because (U * S) is (2,2) and V is (3,2) - @test U * S * V' ≈ T1 - # TODO: Should this require labels, or use existing labels? - @test_broken contract(contract(U, (1, -1), S, (-1, 2)), (1, -1), V, (2, -1)) ≈ 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 -end diff --git a/NDTensors/test/backup/arraytensor/blocksparsearray.jl b/NDTensors/test/backup/arraytensor/blocksparsearray.jl deleted file mode 100644 index 5d508608a5..0000000000 --- a/NDTensors/test/backup/arraytensor/blocksparsearray.jl +++ /dev/null @@ -1,52 +0,0 @@ -using NDTensors -using NDTensors.BlockSparseArrays -using BlockArrays: BlockArrays -using LinearAlgebra -using Test - -using NDTensors: storage, storagetype - -@testset "Tensor wrapping BlockSparseArray" begin - is1 = ([1, 1], [1, 2]) - D1 = BlockSparseArray( - [BlockArrays.Block(1, 1), BlockArrays.Block(2, 2)], [randn(1, 1), randn(1, 2)], is1 - ) - - is2 = ([1, 2], [2, 2]) - D2 = BlockSparseArray( - [BlockArrays.Block(1, 1), BlockArrays.Block(2, 2)], [randn(1, 2), randn(2, 2)], 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) <: BlockSparseArray{Float64,2} - @test storage(T1) == D1 - @test eltype(T1) == eltype(D1) - @test inds(T1) == is1 - - @test_broken R = T1 * T2 - @test_broken storagetype(R) <: Matrix{Float64} - @test_broken Array(R) ≈ Array(T1) * Array(T2) - - @test_broken T1r = randn!(similar(T1)) - @test_broken Array(T1r + T1) ≈ Array(T1r) + Array(T1) - @test_broken Array(permutedims(T1, (2, 1))) ≈ permutedims(Array(T1), (2, 1)) - - # TODO: Not implemented yet. - ## U, S, V = svd(T1) - ## @test U * S * V ≈ T1 - - @test_broken T12 = contract(T1, (1, -1), T2, (-1, 2)) - @test_broken T12 ≈ T1 * T2 - - @test_broken D12 = contract(D1, (1, -1), D2, (-1, 2)) - @test_broken D12 ≈ Array(T12) -end diff --git a/NDTensors/test/backup/arraytensor/diagonalarray.jl b/NDTensors/test/backup/arraytensor/diagonalarray.jl deleted file mode 100644 index 1e80d30fd1..0000000000 --- a/NDTensors/test/backup/arraytensor/diagonalarray.jl +++ /dev/null @@ -1,25 +0,0 @@ -@eval module $(gensym()) -using NDTensors: contract, tensor -using NDTensors.SparseArraysBase: densearray -using NDTensors.DiagonalArrays: DiagonalArray -using Test: @test, @testset -@testset "Tensor wrapping DiagonalArray" begin - D = DiagonalArray(randn(3), 3, 4, 5) - Dᵈ = densearray(D) - A = randn(3, 4, 5) - - for convert_to_dense in (true, false) - @test contract(D, (-1, -2, -3), A, (-1, -2, -3); convert_to_dense) ≈ - contract(Dᵈ, (-1, -2, -3), A, (-1, -2, -3)) - @test contract(D, (-1, -2, 1), A, (-1, -2, 2); convert_to_dense) ≈ - contract(Dᵈ, (-1, -2, 1), A, (-1, -2, 2)) - end - - # Tensor tests - Dᵗ = tensor(D, size(D)) - Dᵈᵗ = tensor(Dᵈ, size(D)) - Aᵗ = tensor(A, size(A)) - @test contract(Dᵗ, (-1, -2, -3), Aᵗ, (-1, -2, -3)) ≈ - contract(Dᵈᵗ, (-1, -2, -3), Aᵗ, (-1, -2, -3)) -end -end diff --git a/NDTensors/test/backup/arraytensor/runtests.jl b/NDTensors/test/backup/arraytensor/runtests.jl deleted file mode 100644 index a423524d05..0000000000 --- a/NDTensors/test/backup/arraytensor/runtests.jl +++ /dev/null @@ -1,8 +0,0 @@ -@eval module $(gensym()) -using Test: @testset -@testset "Tensor wrapping AbstractArrays $(f)" for f in [ - "array.jl", "blocksparsearray.jl", "diagonalarray.jl" -] - include(f) -end -end diff --git a/NDTensors/test/lib/Project.toml b/NDTensors/test/lib/Project.toml index fe5a5dbc71..8df49a8ed2 100644 --- a/NDTensors/test/lib/Project.toml +++ b/NDTensors/test/lib/Project.toml @@ -2,7 +2,6 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" -Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"