Skip to content

Commit cabec03

Browse files
lkdvoskshyatt
lkdvos
and
kshyatt
authored
Preparing for GPUs (#392)
* slightly more careful with test setup Co-authored-by: Katharine Hyatt <khyatt@flatironinstitute.org> * removeunit specializations Co-authored-by: Katharine Hyatt <khyatt@flatironinstitute.org> * environment allocation Co-authored-by: Katharine Hyatt <khyatt@flatironinstitute.org> * add `storagetype(::QP)` Co-authored-by: Katharine Hyatt <khyatt@flatironinstitute.org> * careful with MPSTensor constructor Co-authored-by: Katharine Hyatt <khyatt@flatironinstitute.org> * careful with `jordanmpotensortype` Co-authored-by: Katharine Hyatt <khyatt@flatironinstitute.org> * careful with ChepigaAnsatz --------- Co-authored-by: Katharine Hyatt <khyatt@flatironinstitute.org>
1 parent 8480e78 commit cabec03

10 files changed

Lines changed: 62 additions & 52 deletions

File tree

‎Project.toml‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ Accessors = "0.1"
3434
Adapt = "4"
3535
Aqua = "0.8.9"
3636
BlockTensorKit = "0.3.8"
37+
CUDA = "5.9"
3738
Combinatorics = "1"
3839
Compat = "3.47, 4.10"
3940
DocStringExtensions = "0.9.3"
@@ -57,18 +58,21 @@ TensorOperations = "5"
5758
Test = "1"
5859
TestExtras = "0.3"
5960
VectorInterface = "0.2, 0.3, 0.4, 0.5"
61+
cuTENSOR = "2.3"
6062
julia = "1.10"
6163

6264
[extras]
6365
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
6466
Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
67+
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
6568
Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
6669
ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc"
6770
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
6871
Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
6972
TensorKitTensors = "41b62e7d-e9d1-4e23-942c-79a97adf954b"
7073
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
7174
TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
75+
cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
7276

7377
[targets]
74-
test = ["Aqua", "Adapt", "Pkg", "Test", "TestExtras", "Plots", "Combinatorics", "ParallelTestRunner", "TensorKitTensors"]
78+
test = ["Aqua", "Adapt", "CUDA", "cuTENSOR", "Pkg", "Test", "TestExtras", "Plots", "Combinatorics", "ParallelTestRunner", "TensorKitTensors"]

‎src/algorithms/correlators.jl‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,18 +21,17 @@ function correlator(
2121
S₂ == S₁' || throw(ArgumentError("O₂ should end with a trivial leg."))
2222

2323
G = similar(js, scalartype(state))
24-
U = ones(scalartype(state), S₁)
2524

26-
@plansor Vₗ[-1 -2; -3] := state.AC[i][3 4; -3] * conj(U[1]) * O₁[1 2; 4 -2] *
27-
conj(state.AC[i][3 2; -1])
25+
@plansor Vₗ[-1 -2; -3] := state.AC[i][2 3; -3] * removeunit(O₁, 1)[1; 3 -2] *
26+
conj(state.AC[i][2 1; -1])
2827
ctr = i + 1
2928

3029
for (k, j) in enumerate(js)
3130
if j > ctr
3231
Vₗ = Vₗ * TransferMatrix(state.AR[ctr:(j - 1)])
3332
end
34-
G[k] = @plansor Vₗ[2 3; 5] * state.AR[j][5 6; 7] * O₂[3 4; 6 1] * U[1] *
35-
conj(state.AR[j][2 4; 7])
33+
G[k] = @plansor Vₗ[1 2; 4] * state.AR[j][4 5; 6] * removeunit(O₂, 4)[2 3; 5] *
34+
conj(state.AR[j][1 3; 6])
3635
ctr = j
3736
end
3837
return G

‎src/algorithms/excitation/chepigaansatz.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ function excitations(
4242

4343
# add random offset to kickstart Krylov process:
4444
AC = ψ.AC[pos]
45-
AC₀ = add(AC, randn(scalartype(AC), space(AC)), eps(real(scalartype(AC)))^(1 / 4))
45+
AC₀ = add(AC, randn!(similar(AC)), eps(real(scalartype(AC)))^(1 / 4))
4646

4747
H_eff = AC_hamiltonian(pos, ψ, H, ψ, envs)
4848
Es, ACs, info = eigsolve(H_eff, AC₀, num + 1, :SR, alg.alg)
@@ -107,7 +107,7 @@ function excitations(
107107

108108
# add random offset to kickstart Krylov process:
109109
@plansor AC2[-1 -2; -3 -4] := ψ.AC[pos][-1 -2; 1] * ψ.AR[pos + 1][1 -4; -3]
110-
AC2₀ = add(AC2, randn(scalartype(AC2), space(AC2)), eps(real(scalartype(AC2)))^(1 / 4))
110+
AC2₀ = add(AC2, randn!(similar(AC2)), eps(real(scalartype(AC2)))^(1 / 4))
111111

112112
H_eff = AC2_hamiltonian(pos, ψ, H, ψ, envs)
113113
Es, AC2s, info = eigsolve(H_eff, AC2₀, num + 1, :SR, alg.alg)

‎src/environments/abstract_envs.jl‎

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,48 +15,56 @@ Base.unlock(envs::AbstractMPSEnvironments) = unlock(envs.lock);
1515
# ------------------
1616
function allocate_GL(bra::AbstractMPS, mpo::AbstractMPO, ket::AbstractMPS, i::Int)
1717
T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket))
18+
M = TensorKit.promote_storagetype(T, eltype(mpo), eltype(bra), eltype(ket))
19+
S = TensorKit.check_spacetype(bra, mpo, ket)
1820
V = left_virtualspace(bra, i) ⊗ left_virtualspace(mpo, i)' ←
1921
left_virtualspace(ket, i)
2022
if V isa BlockTensorKit.TensorMapSumSpace
21-
TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), T)
23+
TT = blocktensormaptype(S, numout(V), numin(V), M)
2224
else
23-
TT = TensorMap{T}
25+
TT = tensormaptype(S, numout(V), numin(V), M)
2426
end
2527
return TT(undef, V)
2628
end
2729

2830
function allocate_GR(bra::AbstractMPS, mpo::AbstractMPO, ket::AbstractMPS, i::Int)
2931
T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket))
32+
M = TensorKit.promote_storagetype(T, eltype(mpo), eltype(bra), eltype(ket))
33+
S = TensorKit.check_spacetype(bra, mpo, ket)
3034
V = right_virtualspace(ket, i) ⊗ right_virtualspace(mpo, i) ←
3135
right_virtualspace(bra, i)
3236
if V isa BlockTensorKit.TensorMapSumSpace
33-
TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), T)
37+
TT = blocktensormaptype(S, numout(V), numin(V), M)
3438
else
35-
TT = TensorMap{T}
39+
TT = tensormaptype(S, numout(V), numin(V), M)
3640
end
3741
return TT(undef, V)
3842
end
3943

4044
function allocate_GBL(bra::QP, mpo::AbstractMPO, ket::QP, i::Int)
4145
T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket))
46+
M = TensorKit.promote_storagetype(T, eltype(mpo), eltype(bra), eltype(ket))
47+
S = TensorKit.check_spacetype(bra, mpo, ket)
4248
V = left_virtualspace(bra.left_gs, i) ⊗ left_virtualspace(mpo, i)' ←
4349
auxiliaryspace(ket)' ⊗ left_virtualspace(ket.right_gs, i)
4450
if V isa BlockTensorKit.TensorMapSumSpace
45-
TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), T)
51+
TT = blocktensormaptype(S, numout(V), numin(V), M)
4652
else
47-
TT = TensorMap{T}
53+
TT = tensormaptype(S, numout(V), numin(V), M)
4854
end
4955
return TT(undef, V)
5056
end
5157

5258
function allocate_GBR(bra::QP, mpo::AbstractMPO, ket::QP, i::Int)
5359
T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket))
60+
M = TensorKit.promote_storagetype(T, eltype(mpo), eltype(bra), eltype(ket))
61+
S = TensorKit.check_spacetype(bra, mpo, ket)
5462
V = right_virtualspace(ket.left_gs, i) ⊗ right_virtualspace(mpo, i) ←
5563
auxiliaryspace(ket)' ⊗ right_virtualspace(bra.right_gs, i)
5664
if V isa BlockTensorKit.TensorMapSumSpace
57-
TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), T)
65+
TT = blocktensormaptype(S, numout(V), numin(V), M)
5866
else
59-
TT = TensorMap{T}
67+
TT = tensormaptype(S, numout(V), numin(V), M)
6068
end
6169
return TT(undef, V)
6270
end

‎src/operators/jordanmpotensor.jl‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -121,20 +121,20 @@ function JordanMPOTensor(W::SparseBlockTensorMap{TT, E, S, 2, 2}) where {TT, E,
121121
)
122122
end
123123

124-
function jordanmpotensortype(::Type{S}, ::Type{E}) where {S <: VectorSpace, E <: Number}
125-
TA = Union{tensormaptype(S, 2, 2, E), BraidingTensor{E, S}}
124+
function jordanmpotensortype(::Type{S}, ::Type{E}) where {S <: VectorSpace, E}
125+
TA = tensormaptype(S, 2, 2, E)
126+
T = scalartype(TA)
127+
Tτ = BraidingTensor{T, S}
126128
TB = tensormaptype(S, 2, 1, E)
127129
TC = tensormaptype(S, 1, 2, E)
128130
TD = tensormaptype(S, 1, 1, E)
129-
return JordanMPOTensor{E, S, TA, TB, TC, TD}
131+
return JordanMPOTensor{T, S, Union{TA, Tτ}, TB, TC, TD}
130132
end
131-
function jordanmpotensortype(::Type{O}) where {O <: MPOTensor}
132-
return jordanmpotensortype(spacetype(O), scalartype(O))
133-
end
134-
135-
function Base.similar(W::JordanMPOTensor, ::Type{T}) where {T <: Number}
136-
return JordanMPOTensor{T}(undef, space(W))
133+
function jordanmpotensortype(::Type{O}) where {O <: AbstractTensorMap}
134+
return jordanmpotensortype(spacetype(O), storagetype(O))
137135
end
136+
Base.similar(W::JordanMPOTensor, ::Type{TorA}) where {TorA} =
137+
jordanmpotensortype(spacetype(W), TorA)(undef, space(W))
138138

139139
# Properties
140140
# ----------

‎src/operators/mpohamiltonian.jl‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -437,7 +437,6 @@ function FiniteMPOHamiltonian(lattice::AbstractArray{<:VectorSpace}, local_opera
437437

438438
# construct the sparse MPO
439439
T = _find_tensortype(nonzero_opps)
440-
E = scalartype(T)
441440
S = spacetype(T)
442441

443442
# avoid using one(S)
@@ -465,7 +464,7 @@ function FiniteMPOHamiltonian(lattice::AbstractArray{<:VectorSpace}, local_opera
465464
end
466465

467466
# construct the tensor
468-
TW = jordanmpotensortype(S, E)
467+
TW = jordanmpotensortype(T)
469468
Os = map(1:length(lattice)) do site
470469
V = virtualsumspaces[site] * lattice[site] ←
471470
lattice[site] * virtualsumspaces[site + 1]
@@ -476,7 +475,7 @@ function FiniteMPOHamiltonian(lattice::AbstractArray{<:VectorSpace}, local_opera
476475
key_R = key_R′ == 0 ? length(virtualsumspaces[site + 1]) : key_R′
477476
O[key_L, 1, 1, key_R] += if o isa Number
478477
iszero(o) && continue
479-
τ = BraidingTensor{E}(eachspace(O)[key_L, 1, 1, key_R])
478+
τ = BraidingTensor{scalartype(TW)}(eachspace(O)[key_L, 1, 1, key_R])
480479
isone(o) ? τ : τ * o
481480
else
482481
o
@@ -520,7 +519,6 @@ function InfiniteMPOHamiltonian(lattice′::AbstractArray{<:VectorSpace}, local_
520519

521520
# construct the sparse MPO
522521
T = _find_tensortype(nonzero_opps)
523-
E = scalartype(T)
524522
S = spacetype(T)
525523

526524
# construct the virtual spaces
@@ -588,7 +586,7 @@ function InfiniteMPOHamiltonian(lattice′::AbstractArray{<:VectorSpace}, local_
588586
end
589587

590588
# construct the tensor
591-
TW = jordanmpotensortype(S, E)
589+
TW = jordanmpotensortype(T)
592590
Os = map(1:length(lattice)) do site
593591
V = virtualsumspaces[site - 1] * lattice[site] ←
594592
lattice[site] * virtualsumspaces[site]
@@ -599,7 +597,7 @@ function InfiniteMPOHamiltonian(lattice′::AbstractArray{<:VectorSpace}, local_
599597
key_R = key_R′ == 0 ? length(virtualspaces[site]) : key_R′
600598
O[key_L, 1, 1, key_R] += if o isa Number
601599
iszero(o) && continue
602-
τ = BraidingTensor{E}(eachspace(O)[key_L, 1, 1, key_R])
600+
τ = BraidingTensor{scalartype(TW)}(eachspace(O)[key_L, 1, 1, key_R])
603601
isone(o) ? τ : τ * o
604602
else
605603
o
@@ -824,8 +822,7 @@ function Base.:*(H::FiniteMPOHamiltonian, mps::FiniteMPS)
824822
)
825823
)
826824
# left to middle
827-
U = ones(scalartype(H), left_virtualspace(H, 1))
828-
@plansor a[-1 -2; -3 -4] := A[1][-1 2; -3] * H[1][1 -2; 2 -4] * conj(U[1])
825+
@plansor a[-1 -2; -3 -4] := A[1][-1 1; -3] * removeunit(H[1], 1)[-2; 1 -4]
829826
Q, R = qr_compact!(a)
830827
A′[1] = TensorMap(Q)
831828

@@ -836,8 +833,7 @@ function Base.:*(H::FiniteMPOHamiltonian, mps::FiniteMPS)
836833
end
837834

838835
# right to middle
839-
U = ones(scalartype(H), right_virtualspace(H, N))
840-
@plansor a[-1 -2; -3 -4] := A[end][-1 2; -3] * H[end][-2 -4; 2 1] * U[1]
836+
@plansor a[-1 -2; -3 -4] := A[end][-1 1; -3] * removeunit(H[end], 4)[-2 -4; 1]
841837
L, Q = lq_compact!(a)
842838
A′[end] = transpose(TensorMap(Q), ((1, 3), (2,)))
843839

‎src/states/abstractmps.jl‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,8 @@ Construct an `MPSTensor` with given physical and virtual spaces.
3535
function MPSTensor(
3636
::UndefInitializer, eltype, P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ
3737
) where {S <: ElementarySpace}
38-
return TensorMap{eltype}(undef, Vₗ ⊗ P ← Vᵣ)
38+
TT = tensormaptype(S, 1 + (P isa S ? 1 : length(P)), 1, eltype)
39+
return TT(undef, Vₗ ⊗ P ← Vᵣ)
3940
end
4041
function MPSTensor(
4142
f, eltype, P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ
@@ -81,9 +82,8 @@ Convert an array to an `MPSTensor`.
8182
function MPSTensor(A::AbstractArray{T}) where {T <: Number}
8283
@assert ndims(A) > 2 "MPSTensor should have at least 3 dims, but has $ndims(A)"
8384
sz = size(A)
84-
t = TensorMap(undef, T, foldl(⊗, ComplexSpace.(sz[1:(end - 1)])) ← ℂ^sz[end])
85-
t[] .= A
86-
return t
85+
V = foldl(⊗, ComplexSpace.(sz[1:(end - 1)])) ← ℂ^sz[end]
86+
return TensorMap(A, V)
8787
end
8888

8989
"""

‎src/states/quasiparticle_state.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,7 @@ GeometryStyle(::Type{<:QP{S, T1, T2}}) where {S, T1, T2} = GeometryStyle(S)
216216

217217
TensorKit.spacetype(::Union{QP{S}, Type{<:QP{S}}}) where {S} = spacetype(S)
218218
TensorKit.sectortype(::Union{QP{S}, Type{<:QP{S}}}) where {S} = sectortype(S)
219+
TensorKit.storagetype(::Type{<:QP{S, T1, T2}}) where {S, T1, T2} = storagetype(T2)
219220

220221
physicalspace(state::QP, i::Int) = physicalspace(state.left_gs, i)
221222
physicalspace(state::QP) = physicalspace(state.left_gs)

‎test/cuda/cu_adapt.jl‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,11 @@ using CUDA, cuTENSOR, Adapt
2020

2121
@test isfinite(mpo₁)
2222
@test isfinite(typeof(mpo₁))
23-
@test GeometryStyle(typeof(mpo₁)) == FiniteChainStyle()
24-
@test GeometryStyle(mpo₁) == FiniteChainStyle()
25-
@test OperatorStyle(typeof(mpo₁)) == MPOStyle()
26-
@test TensorKit.storagetype(mpo₁) == CuVector{T, 1, CUDA.DeviceMemory}
27-
@test TensorKit.storagetype(mpo₂) == CuVector{T, 1, CUDA.DeviceMemory}
28-
@test TensorKit.storagetype(mpo₃) == CuVector{T, 1, CUDA.DeviceMemory}
23+
@test MPSKit.GeometryStyle(typeof(mpo₁)) == MPSKit.FiniteChainStyle()
24+
@test MPSKit.GeometryStyle(mpo₁) == MPSKit.FiniteChainStyle()
25+
@test MPSKit.OperatorStyle(typeof(mpo₁)) == MPSKit.MPOStyle()
26+
@test TensorKit.storagetype(mpo₁) == CuVector{T, CUDA.DeviceMemory}
27+
@test TensorKit.storagetype(mpo₂) == CuVector{T, CUDA.DeviceMemory}
28+
@test TensorKit.storagetype(mpo₃) == CuVector{real(T), CUDA.DeviceMemory}
2929
end
3030
end

‎test/runtests.jl‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,15 @@ init_code = quote
1111
using .TestSetup
1212
end
1313

14-
args = parse_args(ARGS)
14+
# only run CUDA if on buildkite
1515
is_buildkite = get(ENV, "BUILDKITE", "false") == "true"
16-
if is_buildkite
17-
empty!(testsuite)
18-
gpu_testsuite = find_tests(joinpath(@__DIR__, "cuda"))
19-
append!(testsuite, gpu_testsuite)
20-
else
16+
is_buildkite && filter!(startswith("cuda") ∘ first, testsuite)
17+
18+
# only run CUDA/cuTENSOR if available
19+
using CUDA, cuTENSOR
20+
(CUDA.functional() && cuTENSOR.has_cutensor()) ||
2121
filter!(!(startswith("cuda") ∘ first), testsuite)
22-
end
22+
23+
# run tests
24+
args = parse_args(ARGS)
2325
runtests(MPSKit, args; testsuite, init_code)

0 commit comments

Comments
 (0)