Preparing for GPUs

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,18 +58,21 @@ 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"
 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"]

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

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)

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

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
 # ----------

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,)))
 

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
 
 """

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)

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

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)