-
Notifications
You must be signed in to change notification settings - Fork 125
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[ITensorGPU] [ENHANCEMENT] Improve functionality for transferring dat…
…a between CPU and GPU (#956) * Fix `expect` and `correlation_matrix` for MPS on GPU. * Make sure QR and SVD preserve element type (i.e. single precision) in more cases.
- Loading branch information
Showing
48 changed files
with
803 additions
and
533 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,25 @@ | ||
| # | ||
| # Used to adapt `EmptyStorage` types | ||
| # | ||
|
|
||
| NDTensors.cu(eltype::Type{<:Number}, x) = fmap(x -> adapt(CuArray{eltype}, x), x) | ||
| NDTensors.cu(x) = fmap(x -> adapt(CuArray, x), x) | ||
|
|
||
| NDTensors.to_vector_type(arraytype::Type{CuArray}) = CuVector | ||
| NDTensors.to_vector_type(arraytype::Type{CuArray{T}}) where {T} = CuVector{T} | ||
|
|
||
| function NDTensors.set_eltype_if_unspecified( | ||
| arraytype::Type{CuVector{T}}, eltype::Type | ||
| ) where {T} | ||
| return arraytype | ||
| end | ||
| function NDTensors.set_eltype_if_unspecified(arraytype::Type{CuVector}, eltype::Type) | ||
| return CuVector{eltype} | ||
| end | ||
|
|
||
| # Overload `CUDA.cu` for convenience | ||
| const ITensorType = Union{ | ||
| TensorStorage,Tensor,ITensor,Array{ITensor},Array{<:Array{ITensor}},MPS,MPO | ||
| } | ||
| CUDA.cu(x::ITensorType) = NDTensors.cu(x) | ||
| CUDA.cu(eltype::Type{<:Number}, x::ITensorType) = NDTensors.cu(eltype, x) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,136 +1,13 @@ | ||
| function cuMPS(psi::MPS) | ||
| phi = copy(psi) | ||
| for site in 1:length(psi) | ||
| phi.data[site] = cuITensor(psi.data[site]) | ||
| end | ||
| return phi | ||
| end | ||
| # cu(ψ::Union{MPS,MPO}) = map(cu, ψ) | ||
| # cpu(ψ::Union{MPS,MPO}) = map(cpu, ψ) | ||
|
|
||
| cu(ψ::MPS) = cuMPS(ψ) | ||
| cuMPS(ψ::MPS) = cu(ψ) | ||
| cuMPS(args...; kwargs...) = cu(MPS(args...; kwargs...)) | ||
| randomCuMPS(args...; kwargs...) = cu(randomMPS(args...; kwargs...)) | ||
|
|
||
| cuMPS() = MPS() | ||
| function cuMPS(::Type{T}, sites::Vector{<:Index}; linkdims::Integer=1) where {T<:Number} | ||
| N = length(sites) | ||
| v = Vector{ITensor}(undef, N) | ||
| if N == 1 | ||
| v[1] = emptyITensor(T, sites[1]) | ||
| return MPS(v) | ||
| end | ||
| # For backwards compatibility | ||
| productCuMPS(args...; kwargs...) = cuMPS(args...; kwargs...) | ||
|
|
||
| space = if hasqns(sites) | ||
| [QN() => linkdims] | ||
| else | ||
| linkdims | ||
| end | ||
|
|
||
| l = [Index(1, "Link,l=$ii") for ii in 1:(N - 1)] | ||
| for ii in eachindex(sites) | ||
| s = sites[ii] | ||
| if ii == 1 | ||
| v[ii] = cuITensor(l[ii], s) | ||
| elseif ii == N | ||
| v[ii] = cuITensor(l[ii - 1], s) | ||
| else | ||
| v[ii] = cuITensor(l[ii - 1], s, l[ii]) | ||
| end | ||
| end | ||
| return MPS(v, 0, N + 1) | ||
| end | ||
| function cuMPS(sites::Vector{<:Index}, args...; kwargs...) | ||
| return cuMPS(Float64, sites, args...; kwargs...) | ||
| end | ||
| function cuMPS(N::Int; ortho_lims::UnitRange=1:N) | ||
| return MPS(Vector{ITensor}(undef, N); ortho_lims=ortho_lims) | ||
| end | ||
|
|
||
| function randomCuMPS(sites) | ||
| M = cuMPS(sites) | ||
| for i in eachindex(sites) | ||
| randn!(M[i]) | ||
| normalize!(M[i]) | ||
| end | ||
| M.llim = 1 | ||
| M.rlim = length(M) | ||
| return M | ||
| end | ||
| function randomCuMPS(N::Int; ortho_lims::UnitRange=1:N) | ||
| return randomCuMPS(Vector{ITensor}(undef, N); ortho_lims=ortho_lims) | ||
| end | ||
|
|
||
| const productCuMPS = cuMPS | ||
|
|
||
| function cuMPS(::Type{T}, ivals::Vector{<:Pair{<:Index}}) where {T<:Number} | ||
| N = length(ivals) | ||
| As = Vector{ITensor}(undef, N) | ||
| links = Vector{Index}(undef, N) | ||
| for n in 1:N | ||
| s = ITensors.ind(ivals[n]) | ||
| links[n] = Index(1, "Link,l=$n") | ||
| if n == 1 | ||
| A = ITensor(T, s, links[n]) | ||
| A[ivals[n], links[n](1)] = 1.0 | ||
| elseif n == N | ||
| A = ITensor(T, links[n - 1], s) | ||
| A[links[n - 1](1), ivals[n]] = 1.0 | ||
| else | ||
| A = ITensor(T, links[n - 1], s, links[n]) | ||
| A[links[n - 1](1), ivals[n], links[n](1)] = 1.0 | ||
| end | ||
| As[n] = cuITensor(A) | ||
| end | ||
| return MPS(As, 0, 2) | ||
| end | ||
| cuMPS(ivals::Vector{Pair{<:Index}}) = cuMPS(Float64, ivals) | ||
|
|
||
| function cuMPS(::Type{T}, sites::Vector{<:Index}, states_) where {T<:Number} | ||
| if length(sites) != length(states_) | ||
| throw(DimensionMismatch("Number of sites and and initial vals don't match")) | ||
| end | ||
| N = length(states_) | ||
| M = cuMPS(N) | ||
|
|
||
| if N == 1 | ||
| M[1] = state(sites[1], states_[1]) | ||
| return M | ||
| end | ||
|
|
||
| states = [state(sites[j], states_[j]) for j in 1:N] | ||
|
|
||
| if hasqns(states[1]) | ||
| lflux = QN() | ||
| for j in 1:(N - 1) | ||
| lflux += flux(states[j]) | ||
| end | ||
| links = Vector{QNIndex}(undef, N - 1) | ||
| for j in (N - 1):-1:1 | ||
| links[j] = Index(lflux => 1; tags="Link,l=$j", dir=In) | ||
| lflux -= flux(states[j]) | ||
| end | ||
| else | ||
| links = [Index(1; tags="Link,l=$n") for n in 1:N] | ||
| end | ||
|
|
||
| M[1] = cuITensor(T, sites[1], links[1]) | ||
| M[1] += cuITensor(states[1] * state(links[1], 1)) | ||
| for n in 2:(N - 1) | ||
| M[n] = cuITensor(T, dag(links[n - 1]), sites[n], links[n]) | ||
| M[n] += cuITensor(state(dag(links[n - 1]), 1) * states[n] * state(links[n], 1)) | ||
| end | ||
| M[N] = cuITensor(T, dag(links[N - 1]), sites[N]) | ||
| M[N] += cuITensor(state(dag(links[N - 1]), 1) * states[N]) | ||
|
|
||
| return M | ||
| end | ||
|
|
||
| function cuMPS( | ||
| ::Type{T}, sites::Vector{<:Index}, state::Union{String,Integer} | ||
| ) where {T<:Number} | ||
| return cuMPS(T, sites, fill(state, length(sites))) | ||
| end | ||
|
|
||
| function cuMPS(::Type{T}, sites::Vector{<:Index}, states::Function) where {T<:Number} | ||
| states_vec = [states(n) for n in 1:length(sites)] | ||
| return cuMPS(T, sites, states_vec) | ||
| end | ||
|
|
||
| cuMPS(sites::Vector{<:Index}, states) = cuMPS(Float64, sites, states) | ||
| cuMPO(M::MPO) = cu(M) | ||
| cuMPO(args...; kwargs...) = cu(MPO(args...; kwargs...)) | ||
| randomCuMPO(args...; kwargs...) = cu(randomMPO(args...; kwargs...)) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.