diff --git a/Project.toml b/Project.toml index 7be06e0c5..cbd3a1015 100644 --- a/Project.toml +++ b/Project.toml @@ -34,6 +34,7 @@ Accessors = "0.1" Adapt = "4" Aqua = "0.8.9" BlockTensorKit = "0.3.8" +CUDA = "5.9" Combinatorics = "1" Compat = "3.47, 4.10" DocStringExtensions = "0.9.3" @@ -57,11 +58,13 @@ TensorOperations = "5" Test = "1" TestExtras = "0.3" VectorInterface = "0.2, 0.3, 0.4, 0.5" +cuTENSOR = "2.3" julia = "1.10" [extras] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" @@ -69,6 +72,7 @@ Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" TensorKitTensors = "41b62e7d-e9d1-4e23-942c-79a97adf954b" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" +cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1" [targets] -test = ["Aqua", "Adapt", "Pkg", "Test", "TestExtras", "Plots", "Combinatorics", "ParallelTestRunner", "TensorKitTensors"] +test = ["Aqua", "Adapt", "CUDA", "cuTENSOR", "Pkg", "Test", "TestExtras", "Plots", "Combinatorics", "ParallelTestRunner", "TensorKitTensors"] diff --git a/src/algorithms/correlators.jl b/src/algorithms/correlators.jl index 29d6d73f2..4413ef63c 100644 --- a/src/algorithms/correlators.jl +++ b/src/algorithms/correlators.jl @@ -21,18 +21,17 @@ function correlator( S₂ == S₁' || throw(ArgumentError("O₂ should end with a trivial leg.")) G = similar(js, scalartype(state)) - U = ones(scalartype(state), S₁) - @plansor Vₗ[-1 -2; -3] := state.AC[i][3 4; -3] * conj(U[1]) * O₁[1 2; 4 -2] * - conj(state.AC[i][3 2; -1]) + @plansor Vₗ[-1 -2; -3] := state.AC[i][2 3; -3] * removeunit(O₁, 1)[1; 3 -2] * + conj(state.AC[i][2 1; -1]) ctr = i + 1 for (k, j) in enumerate(js) if j > ctr Vₗ = Vₗ * TransferMatrix(state.AR[ctr:(j - 1)]) end - G[k] = @plansor Vₗ[2 3; 5] * state.AR[j][5 6; 7] * O₂[3 4; 6 1] * U[1] * - conj(state.AR[j][2 4; 7]) + G[k] = @plansor Vₗ[1 2; 4] * state.AR[j][4 5; 6] * removeunit(O₂, 4)[2 3; 5] * + conj(state.AR[j][1 3; 6]) ctr = j end return G diff --git a/src/algorithms/excitation/chepigaansatz.jl b/src/algorithms/excitation/chepigaansatz.jl index 719073ff3..e212861e0 100644 --- a/src/algorithms/excitation/chepigaansatz.jl +++ b/src/algorithms/excitation/chepigaansatz.jl @@ -42,7 +42,7 @@ function excitations( # add random offset to kickstart Krylov process: AC = ψ.AC[pos] - AC₀ = add(AC, randn(scalartype(AC), space(AC)), eps(real(scalartype(AC)))^(1 / 4)) + AC₀ = add(AC, randn!(similar(AC)), eps(real(scalartype(AC)))^(1 / 4)) H_eff = AC_hamiltonian(pos, ψ, H, ψ, envs) Es, ACs, info = eigsolve(H_eff, AC₀, num + 1, :SR, alg.alg) @@ -107,7 +107,7 @@ function excitations( # add random offset to kickstart Krylov process: @plansor AC2[-1 -2; -3 -4] := ψ.AC[pos][-1 -2; 1] * ψ.AR[pos + 1][1 -4; -3] - AC2₀ = add(AC2, randn(scalartype(AC2), space(AC2)), eps(real(scalartype(AC2)))^(1 / 4)) + AC2₀ = add(AC2, randn!(similar(AC2)), eps(real(scalartype(AC2)))^(1 / 4)) H_eff = AC2_hamiltonian(pos, ψ, H, ψ, envs) Es, AC2s, info = eigsolve(H_eff, AC2₀, num + 1, :SR, alg.alg) diff --git a/src/environments/abstract_envs.jl b/src/environments/abstract_envs.jl index 14aaacf16..612a37207 100644 --- a/src/environments/abstract_envs.jl +++ b/src/environments/abstract_envs.jl @@ -15,48 +15,56 @@ Base.unlock(envs::AbstractMPSEnvironments) = unlock(envs.lock); # ------------------ function allocate_GL(bra::AbstractMPS, mpo::AbstractMPO, ket::AbstractMPS, i::Int) T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket)) + M = TensorKit.promote_storagetype(T, eltype(mpo), eltype(bra), eltype(ket)) + S = TensorKit.check_spacetype(bra, mpo, ket) V = left_virtualspace(bra, i) ⊗ left_virtualspace(mpo, i)' ← left_virtualspace(ket, i) if V isa BlockTensorKit.TensorMapSumSpace - TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), T) + TT = blocktensormaptype(S, numout(V), numin(V), M) else - TT = TensorMap{T} + TT = tensormaptype(S, numout(V), numin(V), M) end return TT(undef, V) end function allocate_GR(bra::AbstractMPS, mpo::AbstractMPO, ket::AbstractMPS, i::Int) T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket)) + M = TensorKit.promote_storagetype(T, eltype(mpo), eltype(bra), eltype(ket)) + S = TensorKit.check_spacetype(bra, mpo, ket) V = right_virtualspace(ket, i) ⊗ right_virtualspace(mpo, i) ← right_virtualspace(bra, i) if V isa BlockTensorKit.TensorMapSumSpace - TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), T) + TT = blocktensormaptype(S, numout(V), numin(V), M) else - TT = TensorMap{T} + TT = tensormaptype(S, numout(V), numin(V), M) end return TT(undef, V) end function allocate_GBL(bra::QP, mpo::AbstractMPO, ket::QP, i::Int) T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket)) + M = TensorKit.promote_storagetype(T, eltype(mpo), eltype(bra), eltype(ket)) + S = TensorKit.check_spacetype(bra, mpo, ket) V = left_virtualspace(bra.left_gs, i) ⊗ left_virtualspace(mpo, i)' ← auxiliaryspace(ket)' ⊗ left_virtualspace(ket.right_gs, i) if V isa BlockTensorKit.TensorMapSumSpace - TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), T) + TT = blocktensormaptype(S, numout(V), numin(V), M) else - TT = TensorMap{T} + TT = tensormaptype(S, numout(V), numin(V), M) end return TT(undef, V) end function allocate_GBR(bra::QP, mpo::AbstractMPO, ket::QP, i::Int) T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket)) + M = TensorKit.promote_storagetype(T, eltype(mpo), eltype(bra), eltype(ket)) + S = TensorKit.check_spacetype(bra, mpo, ket) V = right_virtualspace(ket.left_gs, i) ⊗ right_virtualspace(mpo, i) ← auxiliaryspace(ket)' ⊗ right_virtualspace(bra.right_gs, i) if V isa BlockTensorKit.TensorMapSumSpace - TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), T) + TT = blocktensormaptype(S, numout(V), numin(V), M) else - TT = TensorMap{T} + TT = tensormaptype(S, numout(V), numin(V), M) end return TT(undef, V) end diff --git a/src/operators/jordanmpotensor.jl b/src/operators/jordanmpotensor.jl index 4a3dbe798..04c4ee1cb 100644 --- a/src/operators/jordanmpotensor.jl +++ b/src/operators/jordanmpotensor.jl @@ -121,20 +121,20 @@ function JordanMPOTensor(W::SparseBlockTensorMap{TT, E, S, 2, 2}) where {TT, E, ) end -function jordanmpotensortype(::Type{S}, ::Type{E}) where {S <: VectorSpace, E <: Number} - TA = Union{tensormaptype(S, 2, 2, E), BraidingTensor{E, S}} +function jordanmpotensortype(::Type{S}, ::Type{E}) where {S <: VectorSpace, E} + TA = tensormaptype(S, 2, 2, E) + T = scalartype(TA) + Tτ = BraidingTensor{T, S} TB = tensormaptype(S, 2, 1, E) TC = tensormaptype(S, 1, 2, E) TD = tensormaptype(S, 1, 1, E) - return JordanMPOTensor{E, S, TA, TB, TC, TD} + return JordanMPOTensor{T, S, Union{TA, Tτ}, TB, TC, TD} end -function jordanmpotensortype(::Type{O}) where {O <: MPOTensor} - return jordanmpotensortype(spacetype(O), scalartype(O)) -end - -function Base.similar(W::JordanMPOTensor, ::Type{T}) where {T <: Number} - return JordanMPOTensor{T}(undef, space(W)) +function jordanmpotensortype(::Type{O}) where {O <: AbstractTensorMap} + return jordanmpotensortype(spacetype(O), storagetype(O)) end +Base.similar(W::JordanMPOTensor, ::Type{TorA}) where {TorA} = + jordanmpotensortype(spacetype(W), TorA)(undef, space(W)) # Properties # ---------- diff --git a/src/operators/mpohamiltonian.jl b/src/operators/mpohamiltonian.jl index a938c172e..fe3ffdb1c 100644 --- a/src/operators/mpohamiltonian.jl +++ b/src/operators/mpohamiltonian.jl @@ -437,7 +437,6 @@ function FiniteMPOHamiltonian(lattice::AbstractArray{<:VectorSpace}, local_opera # construct the sparse MPO T = _find_tensortype(nonzero_opps) - E = scalartype(T) S = spacetype(T) # avoid using one(S) @@ -465,7 +464,7 @@ function FiniteMPOHamiltonian(lattice::AbstractArray{<:VectorSpace}, local_opera end # construct the tensor - TW = jordanmpotensortype(S, E) + TW = jordanmpotensortype(T) Os = map(1:length(lattice)) do site V = virtualsumspaces[site] * lattice[site] ← lattice[site] * virtualsumspaces[site + 1] @@ -476,7 +475,7 @@ function FiniteMPOHamiltonian(lattice::AbstractArray{<:VectorSpace}, local_opera key_R = key_R′ == 0 ? length(virtualsumspaces[site + 1]) : key_R′ O[key_L, 1, 1, key_R] += if o isa Number iszero(o) && continue - τ = BraidingTensor{E}(eachspace(O)[key_L, 1, 1, key_R]) + τ = BraidingTensor{scalartype(TW)}(eachspace(O)[key_L, 1, 1, key_R]) isone(o) ? τ : τ * o else o @@ -520,7 +519,6 @@ function InfiniteMPOHamiltonian(lattice′::AbstractArray{<:VectorSpace}, local_ # construct the sparse MPO T = _find_tensortype(nonzero_opps) - E = scalartype(T) S = spacetype(T) # construct the virtual spaces @@ -588,7 +586,7 @@ function InfiniteMPOHamiltonian(lattice′::AbstractArray{<:VectorSpace}, local_ end # construct the tensor - TW = jordanmpotensortype(S, E) + TW = jordanmpotensortype(T) Os = map(1:length(lattice)) do site V = virtualsumspaces[site - 1] * lattice[site] ← lattice[site] * virtualsumspaces[site] @@ -599,7 +597,7 @@ function InfiniteMPOHamiltonian(lattice′::AbstractArray{<:VectorSpace}, local_ key_R = key_R′ == 0 ? length(virtualspaces[site]) : key_R′ O[key_L, 1, 1, key_R] += if o isa Number iszero(o) && continue - τ = BraidingTensor{E}(eachspace(O)[key_L, 1, 1, key_R]) + τ = BraidingTensor{scalartype(TW)}(eachspace(O)[key_L, 1, 1, key_R]) isone(o) ? τ : τ * o else o @@ -824,8 +822,7 @@ function Base.:*(H::FiniteMPOHamiltonian, mps::FiniteMPS) ) ) # left to middle - U = ones(scalartype(H), left_virtualspace(H, 1)) - @plansor a[-1 -2; -3 -4] := A[1][-1 2; -3] * H[1][1 -2; 2 -4] * conj(U[1]) + @plansor a[-1 -2; -3 -4] := A[1][-1 1; -3] * removeunit(H[1], 1)[-2; 1 -4] Q, R = qr_compact!(a) A′[1] = TensorMap(Q) @@ -836,8 +833,7 @@ function Base.:*(H::FiniteMPOHamiltonian, mps::FiniteMPS) end # right to middle - U = ones(scalartype(H), right_virtualspace(H, N)) - @plansor a[-1 -2; -3 -4] := A[end][-1 2; -3] * H[end][-2 -4; 2 1] * U[1] + @plansor a[-1 -2; -3 -4] := A[end][-1 1; -3] * removeunit(H[end], 4)[-2 -4; 1] L, Q = lq_compact!(a) A′[end] = transpose(TensorMap(Q), ((1, 3), (2,))) diff --git a/src/states/abstractmps.jl b/src/states/abstractmps.jl index e442100db..a8b704127 100644 --- a/src/states/abstractmps.jl +++ b/src/states/abstractmps.jl @@ -35,7 +35,8 @@ Construct an `MPSTensor` with given physical and virtual spaces. function MPSTensor( ::UndefInitializer, eltype, P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ ) where {S <: ElementarySpace} - return TensorMap{eltype}(undef, Vₗ ⊗ P ← Vᵣ) + TT = tensormaptype(S, 1 + (P isa S ? 1 : length(P)), 1, eltype) + return TT(undef, Vₗ ⊗ P ← Vᵣ) end function MPSTensor( f, eltype, P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ @@ -81,9 +82,8 @@ Convert an array to an `MPSTensor`. function MPSTensor(A::AbstractArray{T}) where {T <: Number} @assert ndims(A) > 2 "MPSTensor should have at least 3 dims, but has $ndims(A)" sz = size(A) - t = TensorMap(undef, T, foldl(⊗, ComplexSpace.(sz[1:(end - 1)])) ← ℂ^sz[end]) - t[] .= A - return t + V = foldl(⊗, ComplexSpace.(sz[1:(end - 1)])) ← ℂ^sz[end] + return TensorMap(A, V) end """ diff --git a/src/states/quasiparticle_state.jl b/src/states/quasiparticle_state.jl index 2b10cd762..7e9d78e46 100644 --- a/src/states/quasiparticle_state.jl +++ b/src/states/quasiparticle_state.jl @@ -216,6 +216,7 @@ GeometryStyle(::Type{<:QP{S, T1, T2}}) where {S, T1, T2} = GeometryStyle(S) TensorKit.spacetype(::Union{QP{S}, Type{<:QP{S}}}) where {S} = spacetype(S) TensorKit.sectortype(::Union{QP{S}, Type{<:QP{S}}}) where {S} = sectortype(S) +TensorKit.storagetype(::Type{<:QP{S, T1, T2}}) where {S, T1, T2} = storagetype(T2) physicalspace(state::QP, i::Int) = physicalspace(state.left_gs, i) physicalspace(state::QP) = physicalspace(state.left_gs) diff --git a/test/cuda/cu_adapt.jl b/test/cuda/cu_adapt.jl index 24360f637..1bbab9e69 100644 --- a/test/cuda/cu_adapt.jl +++ b/test/cuda/cu_adapt.jl @@ -20,11 +20,11 @@ using CUDA, cuTENSOR, Adapt @test isfinite(mpo₁) @test isfinite(typeof(mpo₁)) - @test GeometryStyle(typeof(mpo₁)) == FiniteChainStyle() - @test GeometryStyle(mpo₁) == FiniteChainStyle() - @test OperatorStyle(typeof(mpo₁)) == MPOStyle() - @test TensorKit.storagetype(mpo₁) == CuVector{T, 1, CUDA.DeviceMemory} - @test TensorKit.storagetype(mpo₂) == CuVector{T, 1, CUDA.DeviceMemory} - @test TensorKit.storagetype(mpo₃) == CuVector{T, 1, CUDA.DeviceMemory} + @test MPSKit.GeometryStyle(typeof(mpo₁)) == MPSKit.FiniteChainStyle() + @test MPSKit.GeometryStyle(mpo₁) == MPSKit.FiniteChainStyle() + @test MPSKit.OperatorStyle(typeof(mpo₁)) == MPSKit.MPOStyle() + @test TensorKit.storagetype(mpo₁) == CuVector{T, CUDA.DeviceMemory} + @test TensorKit.storagetype(mpo₂) == CuVector{T, CUDA.DeviceMemory} + @test TensorKit.storagetype(mpo₃) == CuVector{real(T), CUDA.DeviceMemory} end end diff --git a/test/runtests.jl b/test/runtests.jl index c1c525ade..6fe83d898 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,13 +11,15 @@ init_code = quote using .TestSetup end -args = parse_args(ARGS) +# only run CUDA if on buildkite is_buildkite = get(ENV, "BUILDKITE", "false") == "true" -if is_buildkite - empty!(testsuite) - gpu_testsuite = find_tests(joinpath(@__DIR__, "cuda")) - append!(testsuite, gpu_testsuite) -else +is_buildkite && filter!(startswith("cuda") ∘ first, testsuite) + +# only run CUDA/cuTENSOR if available +using CUDA, cuTENSOR +(CUDA.functional() && cuTENSOR.has_cutensor()) || filter!(!(startswith("cuda") ∘ first), testsuite) -end + +# run tests +args = parse_args(ARGS) runtests(MPSKit, args; testsuite, init_code)