Skip to content

Commit b657bb4

Browse files
kshyattkshyatt
committed
More tests for CUDA + MPS
1 parent 888de7e commit b657bb4

3 files changed

Lines changed: 198 additions & 1 deletion

File tree

‎ext/MPSKitAdaptExt.jl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,5 +35,7 @@ end
3535
MPSKit.JordanMPOTensor(space(W), adapt(to, W.A), adapt(to, W.B), adapt(to, W.C), adapt(to, W.D))
3636
@inline Adapt.adapt_structure(to, mpo::MPOHamiltonian) =
3737
MPOHamiltonian(map(x -> adapt(to, x), mpo.W))
38+
@inline Adapt.adapt_structure(to, ml::MPSKit.Multiline) = MPSKit.Multiline(map(adapt(to), ml.data))
39+
@inline Adapt.adapt_structure(to, pa::MPSKit.PeriodicArray) = MPSKit.PeriodicArray(map(adapt(to), pa.data))
3840

3941
end

‎test/cuda/operators.jl‎

Lines changed: 156 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,15 @@ using MPSKit: GeometryStyle, FiniteChainStyle, InfiniteChainStyle, OperatorStyle
55
using TensorKit
66
using MatrixAlgebraKit
77
using TensorKit: ℙ, tensormaptype, TensorMapWithStorage
8-
using Adapt, CUDA, cuTENSOR
8+
using Adapt, CUDA, cuTENSOR, CUDA.CUBLAS
99

1010
# TODO revisit this once https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/issues/176
1111
# is resolved
1212
MPSKit.Defaults.alg_svd() = CUSOLVER_QRIteration()
1313

14+
pspaces = (ℙ^4, Rep[U₁](0 => 2), Rep[SU₂](1 => 1))
15+
vspaces = (ℙ^10, Rep[U₁]((0 => 20)), Rep[SU₂](1 // 2 => 10, 3 // 2 => 5, 5 // 2 => 1))
16+
1417
@testset "CuFiniteMPO" for V in (ℂ^2, U1Space(0 => 1, 1 => 1))
1518
# start from random operators
1619
L = 4
@@ -87,3 +90,155 @@ MPSKit.Defaults.alg_svd() = CUSOLVER_QRIteration()
8790

8891
@test dot(mpomps₁, mpomps₁) ≈ dot(mpo₁, mpo₁)
8992
end
93+
94+
@testset "Finite CuMPOHamiltonian" begin
95+
L = 3
96+
T = ComplexF64
97+
for T in (Float64, ComplexF64), V in (ℂ^2, U1Space(-1 => 1, 0 => 1, 1 => 1))
98+
lattice = fill(V, L)
99+
O₁ = randn(T, V, V)
100+
O₁ += O₁'
101+
E = id(storagetype(O₁), domain(O₁))
102+
O₂ = randn(T, V^2 ← V^2)
103+
O₂ += O₂'
104+
105+
H1 = adapt(CuVector{T, CUDA.DeviceMemory}, FiniteMPOHamiltonian(lattice, i => O₁ for i in 1:L))
106+
H2 = adapt(CuVector{T, CUDA.DeviceMemory}, FiniteMPOHamiltonian(lattice, (i, i + 1) => O₂ for i in 1:(L - 1)))
107+
H3 = adapt(CuVector{T, CUDA.DeviceMemory}, FiniteMPOHamiltonian(lattice, 1 => O₁, (2, 3) => O₂, (1, 3) => O₂))
108+
@test TensorKit.storagetype(H1) == CuVector{T, CUDA.DeviceMemory}
109+
@test TensorKit.storagetype(H2) == CuVector{T, CUDA.DeviceMemory}
110+
@test TensorKit.storagetype(H3) == CuVector{T, CUDA.DeviceMemory}
111+
112+
@test scalartype(H1) == scalartype(H2) == scalartype(H3) == T
113+
#=if !(T <: Complex)
114+
for H in (H1, H2, H3)
115+
Hc = @constinferred complex(H)
116+
@test scalartype(Hc) == complex(T)
117+
@test TensorKit.storagetype(Hc) == CuVector{complex(T), CUDA.DeviceMemory}
118+
end
119+
end=# # TODO
120+
121+
# check if constructor works by converting back to tensormap
122+
#= H1_tm = convert(TensorMap, H1)
123+
operators = vcat(fill(E, L - 1), O₁)
124+
@test H1_tm ≈ mapreduce(+, 1:L) do i
125+
return reduce(⊗, circshift(operators, i))
126+
end
127+
operators = vcat(fill(E, L - 2), O₂)
128+
@test convert(TensorMap, H2) ≈ mapreduce(+, 1:(L - 1)) do i
129+
return reduce(⊗, circshift(operators, i))
130+
end
131+
@test convert(TensorMap, H3) ≈
132+
O₁ ⊗ E ⊗ E + E ⊗ O₂ + permute(O₂ ⊗ E, ((1, 3, 2), (4, 6, 5))) =# # needs a fix in BlockTensorKit
133+
134+
# check if adding terms on the same site works
135+
single_terms = Iterators.flatten(Iterators.repeated((i => O₁ / 2 for i in 1:L), 2))
136+
H4 = adapt(CuArray, FiniteMPOHamiltonian(lattice, single_terms))
137+
@test TensorKit.storagetype(H4) == CuVector{T, CUDA.DeviceMemory}
138+
@test H4 ≈ H1 atol = 1.0e-6
139+
double_terms = Iterators.flatten(
140+
Iterators.repeated(((i, i + 1) => O₂ / 2 for i in 1:(L - 1)), 2)
141+
)
142+
H5 = adapt(CuArray, FiniteMPOHamiltonian(lattice, double_terms))
143+
@test TensorKit.storagetype(H5) == CuVector{T, CUDA.DeviceMemory}
144+
@test H5 ≈ H2 atol = 1.0e-6
145+
146+
# test linear algebra
147+
@test H1 ≈
148+
adapt(CuArray, FiniteMPOHamiltonian(lattice, 1 => O₁)) +
149+
adapt(CuArray, FiniteMPOHamiltonian(lattice, 2 => O₁)) +
150+
adapt(CuArray, FiniteMPOHamiltonian(lattice, 3 => O₁))
151+
@test TensorKit.storagetype(H1) == CuVector{T, CUDA.DeviceMemory}
152+
#@test 0.8 * H1 + 0.2 * H1 ≈ H1 atol = 1.0e-6 # broken due to JordanMPOTensorMap
153+
#=@test convert(TensorMap, H1 + H2) ≈ convert(TensorMap, H1) + convert(TensorMap, H2) atol = 1.0e-6
154+
H1_trunc = changebonds(H1, SvdCut(; trscheme = truncrank(0)))
155+
@test H1_trunc ≈ H1
156+
@test all(left_virtualspace(H1_trunc) .== left_virtualspace(H1))=# # needs fix in BlockTensorKit
157+
158+
# test dot and application
159+
state = rand(T, prod(lattice))
160+
mps = adapt(CuArray, FiniteMPS(state))
161+
162+
#=@test convert(TensorMap, H1 * mps) ≈ H1_tm * state # needs fix in BlockTensorKit
163+
@test H1 * state ≈ H1_tm * state
164+
@test dot(mps, H2, mps) ≈ dot(mps, H2 * mps)=#
165+
166+
# test constructor from dictionary with mixed linear and Cartesian lattice indices as keys
167+
grid = square = fill(V, 3, 3)
168+
169+
local_operators = Dict((I,) => O₁ for I in eachindex(grid))
170+
I_vertical = CartesianIndex(1, 0)
171+
vertical_operators = Dict(
172+
(I, I + I_vertical) => O₂ for I in eachindex(IndexCartesian(), square) if I[1] < size(square, 1)
173+
)
174+
I_horizontal = CartesianIndex(0, 1)
175+
horizontal_operators = Dict(
176+
(I, I + I_horizontal) => O₂ for I in eachindex(IndexCartesian(), square) if I[2] < size(square, 1)
177+
)
178+
operators = merge(local_operators, vertical_operators, horizontal_operators)
179+
H4 = adapt(CuArray, FiniteMPOHamiltonian(grid, operators))
180+
@test TensorKit.storagetype(H4) == CuVector{T, CUDA.DeviceMemory}
181+
182+
@test H4 ≈
183+
adapt(CuArray, FiniteMPOHamiltonian(grid, local_operators)) +
184+
adapt(CuArray, FiniteMPOHamiltonian(grid, vertical_operators)) +
185+
adapt(CuArray, FiniteMPOHamiltonian(grid, horizontal_operators)) atol = 1.0e-4
186+
@test TensorKit.storagetype(H4) == CuVector{T, CUDA.DeviceMemory}
187+
188+
#H4′= H4 / 3 + 2H4 / 3
189+
#@test TensorKit.storagetype(H4′) == CuVector{T, CUDA.DeviceMemory}
190+
#H5 = changebonds(H4′, SvdCut(; trscheme = trunctol(; atol = 1.0e-12)))
191+
#@test TensorKit.storagetype(H5) == CuVector{T, CUDA.DeviceMemory} # more problems with arithmetic operations...
192+
#psi = adapt(CuArray, FiniteMPS(physicalspace(H5), V ⊕ rightunitspace(V)))
193+
#@test expectation_value(psi, H4) ≈ expectation_value(psi, H5)
194+
end
195+
end
196+
197+
@testset "CuInfiniteMPOHamiltonian $(sectortype(pspace))" for (pspace, Dspace) in zip(pspaces, vspaces)
198+
# generate a 1-2-3 body interaction
199+
operators = ntuple(3) do i
200+
O = rand(ComplexF64, pspace^i, pspace^i)
201+
return O += O'
202+
end
203+
204+
H1 = adapt(CuVector{ComplexF64, CUDA.DeviceMemory}, InfiniteMPOHamiltonian(operators[1]))
205+
H2 = adapt(CuVector{ComplexF64, CUDA.DeviceMemory}, InfiniteMPOHamiltonian(operators[2]))
206+
H3 = adapt(CuVector{ComplexF64, CUDA.DeviceMemory}, repeat(InfiniteMPOHamiltonian(operators[3]), 2))
207+
208+
@test TensorKit.storagetype(H1) == CuVector{ComplexF64, CUDA.DeviceMemory}
209+
@test TensorKit.storagetype(H2) == CuVector{ComplexF64, CUDA.DeviceMemory}
210+
@test TensorKit.storagetype(H3) == CuVector{ComplexF64, CUDA.DeviceMemory}
211+
# make a teststate to measure expectation values for
212+
ψ1 = adapt(CuVector{ComplexF64, CUDA.DeviceMemory}, InfiniteMPS([pspace], [Dspace]))
213+
ψ2 = adapt(CuVector{ComplexF64, CUDA.DeviceMemory}, InfiniteMPS([pspace, pspace], [Dspace, Dspace]))
214+
@test TensorKit.storagetype(ψ1) == CuVector{ComplexF64, CUDA.DeviceMemory}
215+
@test TensorKit.storagetype(ψ2) == CuVector{ComplexF64, CUDA.DeviceMemory}
216+
217+
#=
218+
e1 = expectation_value(ψ1, H1)
219+
e2 = expectation_value(ψ1, H2)
220+
=# # broken due to BraidingTensor
221+
222+
# H1 = 2 * H1 - CuArray([1]) # scalar indexing
223+
# @test TensorKit.storagetype(H1) == CuVector{ComplexF64, CUDA.DeviceMemory}
224+
# @test e1 * 2 - 1 ≈ expectation_value(ψ1, H1) atol = 1.0e-10 # broken due to BraidingTensor
225+
226+
H1 = H1 + H2
227+
@test TensorKit.storagetype(H1) == CuVector{ComplexF64, CUDA.DeviceMemory}
228+
229+
# @test e1 * 2 + e2 - 1 ≈ expectation_value(ψ1, H1) atol = 1.0e-10 # broken due to BraidingTensor
230+
231+
H1 = repeat(H1, 2)
232+
@test TensorKit.storagetype(H1) == CuVector{ComplexF64, CUDA.DeviceMemory}
233+
234+
#=e1 = expectation_value(ψ2, H1)
235+
e3 = expectation_value(ψ2, H3)
236+
237+
@test e1 + e3 ≈ expectation_value(ψ2, H1 + H3) atol = 1.0e-10=# # broken due to BraidingTensor
238+
239+
#H4 = H1 + H3 # broken due to BraidingTensor
240+
#@test TensorKit.storagetype(H4) == CuVector{ComplexF64, CUDA.DeviceMemory}
241+
#h4 = H4 * H4
242+
#@test TensorKit.storagetype(h4) == CuVector{ComplexF64, CUDA.DeviceMemory}
243+
#@test real(expectation_value(ψ2, H4)) >= 0 # broken due to BraidingTensor
244+
end

‎test/cuda/states.jl‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,3 +50,43 @@ using Adapt, CUDA, cuTENSOR
5050

5151
@test norm(2 * ψ + ψ - 3 * ψ) ≈ 0.0 atol = sqrt(eps(real(ComplexF64)))
5252
end
53+
54+
@testset "CuMultilineMPS ($(sectortype(D)), $elt)" for (D, d, elt) in
55+
[(ℙ^10, ℙ^2, ComplexF64), (Rep[U₁](1 => 3), Rep[U₁](0 => 1), ComplexF32)]
56+
tol = Float64(eps(real(elt)) * 100)
57+
ψ = adapt(
58+
CuVector{elt, CUDA.DeviceMemory}, MultilineMPS(
59+
[
60+
rand(elt, D * d, D) rand(elt, D * d, D)
61+
rand(elt, D * d, D) rand(elt, D * d, D)
62+
]; tol
63+
)
64+
)
65+
66+
@test GeometryStyle(typeof(ψ)) == InfiniteChainStyle()
67+
@test GeometryStyle(ψ) == InfiniteChainStyle()
68+
@test TensorKit.storagetype(ψ) == CuVector{elt, CUDA.DeviceMemory}
69+
@test TensorKit.sectortype(ψ) == sectortype(D)
70+
71+
@test !isfinite(typeof(ψ))
72+
73+
@test physicalspace(ψ) == fill(d, 2, 2)
74+
@test all(x -> x ≾ D, left_virtualspace(ψ))
75+
@test all(x -> x ≾ D, right_virtualspace(ψ))
76+
77+
for i in 1:size(ψ, 1), j in 1:size(ψ, 2)
78+
@plansor difference[-1 -2; -3] := ψ.AL[i, j][-1 -2; 1] * ψ.C[i, j][1; -3] -
79+
ψ.C[i, j - 1][-1; 1] * ψ.AR[i, j][1 -2; -3]
80+
@test norm(difference, Inf) < tol * 10
81+
82+
@test l_LL(ψ, i, j) * TransferMatrix(ψ.AL[i, j], ψ.AL[i, j]) ≈ l_LL(ψ, i, j + 1)
83+
@test l_LR(ψ, i, j) * TransferMatrix(ψ.AL[i, j], ψ.AR[i, j]) ≈ l_LR(ψ, i, j + 1)
84+
@test l_RL(ψ, i, j) * TransferMatrix(ψ.AR[i, j], ψ.AL[i, j]) ≈ l_RL(ψ, i, j + 1)
85+
@test l_RR(ψ, i, j) * TransferMatrix(ψ.AR[i, j], ψ.AR[i, j]) ≈ l_RR(ψ, i, j + 1)
86+
87+
@test TransferMatrix(ψ.AL[i, j], ψ.AL[i, j]) * r_LL(ψ, i, j) ≈ r_LL(ψ, i, j + 1)
88+
@test TransferMatrix(ψ.AL[i, j], ψ.AR[i, j]) * r_LR(ψ, i, j) ≈ r_LR(ψ, i, j + 1)
89+
@test TransferMatrix(ψ.AR[i, j], ψ.AL[i, j]) * r_RL(ψ, i, j) ≈ r_RL(ψ, i, j + 1)
90+
@test TransferMatrix(ψ.AR[i, j], ψ.AR[i, j]) * r_RR(ψ, i, j) ≈ r_RR(ψ, i, j + 1)
91+
end
92+
end

0 commit comments

Comments
 (0)