Changes for non-CPU array support

Fix sources

Don't use a storagetype as input for MPS's

Remove duplicates in extension

more comments

Apply suggestions from code review

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

Fix sources

Remove some extra storagetypes

Author: kshyatt

Project.toml

@@ -76,3 +76,6 @@ cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [targets]
 test = ["Aqua", "Adapt", "CUDA", "cuTENSOR", "Pkg", "Test", "TestExtras", "Plots", "Combinatorics", "ParallelTestRunner", "TensorKitTensors"]
+
+[sources]
+TensorKit = {url="https://github.com/QuantumKitHub/TensorKit.jl", rev="ksh/cuda_tweaks"}

src/algorithms/toolbox.jl

@@ -19,11 +19,9 @@ function entropy(spectrum::TensorKit.SectorVector{T}) where {T}
     S = zero(T)
     tol = eps(T)
     for (c, b) in pairs(spectrum)
-        s = zero(S)
-        for x in b
-            x < tol && break
+        s = sum(b; init = zero(S)) do x
             x² = x^2
-            s += x² * log(x²)
+            return x < tol ? zero(x) : x² * log(x²)
         end
         S += oftype(S, dim(c) * s)
     end

src/operators/mpohamiltonian.jl

@@ -32,6 +32,7 @@ struct MPOHamiltonian{TO <: JordanMPOTensor, V <: AbstractVector{TO}} <: Abstrac
     W::V
 end
 OperatorStyle(::Type{<:MPOHamiltonian}) = HamiltonianStyle()
+TensorKit.storagetype(::Type{MPOHamiltonian{O, V}}) where {O, V} = storagetype(O)
 
 const FiniteMPOHamiltonian{O <: MPOTensor} = MPOHamiltonian{O, Vector{O}}
 Base.isfinite(::Type{<:FiniteMPOHamiltonian}) = true

src/states/abstractmps.jl

@@ -33,15 +33,15 @@ Construct an `MPSTensor` with given physical and virtual spaces.
 - `right_D::Int`: right virtual dimension
 """
 function MPSTensor(
-        ::UndefInitializer, eltype, P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ
+        ::UndefInitializer, T, P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ
     ) where {S <: ElementarySpace}
-    TT = tensormaptype(S, 1 + (P isa S ? 1 : length(P)), 1, eltype)
+    TT = tensormaptype(S, 1 + (P isa S ? 1 : length(P)), 1, T)
     return TT(undef, Vₗ ⊗ P ← Vᵣ)
 end
 function MPSTensor(
-        f, eltype, P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ
+        f, T, P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ
     ) where {S <: ElementarySpace}
-    A = MPSTensor(undef, eltype, P, Vₗ, Vᵣ)
+    A = MPSTensor(undef, T, P, Vₗ, Vᵣ)
     if f === rand
         return rand!(A)
     elseif f === randn
@@ -71,15 +71,15 @@ Construct an `MPSTensor` with given physical and virtual dimensions.
 - `Dₗ::Int`: left virtual dimension
 - `Dᵣ::Int`: right virtual dimension
 """
-MPSTensor(f, eltype, d::Int, Dₗ::Int, Dᵣ::Int = Dₗ) = MPSTensor(f, eltype, ℂ^d, ℂ^Dₗ, ℂ^Dᵣ)
+MPSTensor(f, T, d::Int, Dₗ::Int, Dᵣ::Int = Dₗ) = MPSTensor(f, T, ℂ^d, ℂ^Dₗ, ℂ^Dᵣ)
 MPSTensor(d::Int, Dₗ::Int; Dᵣ::Int = Dₗ) = MPSTensor(ℂ^d, ℂ^Dₗ, ℂ^Dᵣ)
 
 """
     MPSTensor(A::AbstractArray)
 
 Convert an array to an `MPSTensor`.
 """
-function MPSTensor(A::AbstractArray{T}) where {T <: Number}
+function MPSTensor(A::AbstractArray{<:Number})
     @assert ndims(A) > 2 "MPSTensor should have at least 3 dims, but has $ndims(A)"
     sz = size(A)
     V = foldl(⊗, ComplexSpace.(sz[1:(end - 1)])) ← ℂ^sz[end]

src/states/infinitemps.jl

@@ -116,35 +116,42 @@ function InfiniteMPS(
         convert(PeriodicVector{B}, C), convert(PeriodicVector{A}, AC)
     )
 end
-
+function InfiniteMPS(
+        T::Type,
+        pspaces::AbstractVector{S}, Dspaces::AbstractVector{S};
+        kwargs...
+    ) where {S <: IndexSpace}
+    return InfiniteMPS(MPSTensor.(rand, T, pspaces, circshift(Dspaces, 1), Dspaces); kwargs...)
+end
 function InfiniteMPS(
         pspaces::AbstractVector{S}, Dspaces::AbstractVector{S};
         kwargs...
     ) where {S <: IndexSpace}
-    return InfiniteMPS(MPSTensor.(pspaces, circshift(Dspaces, 1), Dspaces); kwargs...)
+    return InfiniteMPS(MPSTensor.(rand, ComplexF64, pspaces, circshift(Dspaces, 1), Dspaces); kwargs...)
 end
 function InfiniteMPS(
-        f, elt::Type{<:Number}, pspaces::AbstractVector{S}, Dspaces::AbstractVector{S};
+        f, T::Type, pspaces::AbstractVector{S}, Dspaces::AbstractVector{S};
         kwargs...
     ) where {S <: IndexSpace}
     return InfiniteMPS(
-        MPSTensor.(f, elt, pspaces, circshift(Dspaces, 1), Dspaces);
+        MPSTensor.(f, T, pspaces, circshift(Dspaces, 1), Dspaces);
         kwargs...
     )
 end
-InfiniteMPS(d::S, D::S) where {S <: Union{Int, <:IndexSpace}} = InfiniteMPS([d], [D])
+InfiniteMPS(T::Type, d::S, D::S; kwargs...) where {S <: Union{Int, <:IndexSpace}} = InfiniteMPS(T, [d], [D]; kwargs...)
+InfiniteMPS(d::S, D::S; kwargs...) where {S <: Union{Int, <:IndexSpace}} = InfiniteMPS([d], [D]; kwargs...)
 function InfiniteMPS(
-        f, elt::Type{<:Number}, d::S, D::S
+        f, T::Type, d::S, D::S; kwargs...
     ) where {S <: Union{Int, <:IndexSpace}}
-    return InfiniteMPS(f, elt, [d], [D])
+    return InfiniteMPS(f, T, [d], [D]; kwargs...)
 end
-function InfiniteMPS(ds::AbstractVector{Int}, Ds::AbstractVector{Int})
-    return InfiniteMPS(ComplexSpace.(ds), ComplexSpace.(Ds))
+function InfiniteMPS(ds::AbstractVector{Int}, Ds::AbstractVector{Int}; kwargs...)
+    return InfiniteMPS(ComplexSpace.(ds), ComplexSpace.(Ds); kwargs...)
 end
 function InfiniteMPS(
-        f, elt::Type{<:Number}, ds::AbstractVector{Int}, Ds::AbstractVector{Int}, kwargs...
+        f, T::Type, ds::AbstractVector{Int}, Ds::AbstractVector{Int}, kwargs...
     )
-    return InfiniteMPS(f, elt, ComplexSpace.(ds), ComplexSpace.(Ds); kwargs...)
+    return InfiniteMPS(f, T, ComplexSpace.(ds), ComplexSpace.(Ds); kwargs...)
 end
 
 function InfiniteMPS(A::AbstractVector{<:GenericMPSTensor}; kwargs...)

src/states/multilinemps.jl

@@ -81,16 +81,6 @@ for f in (:r_RR, :r_RL, :r_LR, :r_LL)
     @eval $f(t::MultilineMPS, i, j = size(t, 2)) = $f(t[i], j)
 end
 
-site_type(::Type{Multiline{S}}) where {S} = site_type(S)
-bond_type(::Type{Multiline{S}}) where {S} = bond_type(S)
-site_type(st::Multiline) = site_type(typeof(st))
-bond_type(st::Multiline) = bond_type(typeof(st))
-VectorInterface.scalartype(::Multiline{T}) where {T} = scalartype(T)
-TensorKit.sectortype(t::Multiline) = sectortype(typeof(t))
-TensorKit.sectortype(::Type{Multiline{T}}) where {T} = sectortype(T)
-TensorKit.spacetype(t::Multiline) = spacetype(typeof(t))
-TensorKit.spacetype(::Type{Multiline{T}}) where {T} = spacetype(T)
-
 function TensorKit.dot(a::MultilineMPS, b::MultilineMPS; kwargs...)
     return sum(dot.(parent(a), parent(b); kwargs...))
 end

src/states/ortho.jl

@@ -21,7 +21,7 @@ $(TYPEDFIELDS)
     verbosity::Int = VERBOSE_WARN
 
     "algorithm used for orthogonalization of the tensors"
-    alg_orth = LAPACK_HouseholderQR(; positive = true)
+    alg_orth = Defaults.alg_qr()
     "algorithm used for the eigensolver"
     alg_eigsolve = _GAUGE_ALG_EIGSOLVE
     "minimal amount of iterations before using the eigensolver steps"
@@ -46,7 +46,7 @@ $(TYPEDFIELDS)
     verbosity::Int = VERBOSE_WARN
 
     "algorithm used for orthogonalization of the tensors"
-    alg_orth = LAPACK_HouseholderLQ(; positive = true)
+    alg_orth = Defaults.alg_lq()
     "algorithm used for the eigensolver"
     alg_eigsolve = _GAUGE_ALG_EIGSOLVE
     "minimal amount of iterations before using the eigensolver steps"
@@ -73,16 +73,22 @@ end
 
 function MixedCanonical(;
         tol::Real = Defaults.tolgauge, maxiter::Int = Defaults.maxiter,
-        verbosity::Int = VERBOSE_WARN, alg_orth = LAPACK_HouseholderQR(; positive = true),
+        verbosity::Int = VERBOSE_WARN, alg_orth = Defaults.alg_qr(),
         alg_eigsolve = _GAUGE_ALG_EIGSOLVE,
         eig_miniter::Int = 10, order::Symbol = :LR
     )
     if alg_orth isa LAPACK_HouseholderQR
         alg_leftorth = alg_orth
         alg_rightorth = LAPACK_HouseholderLQ(; alg_orth.kwargs...)
+    elseif alg_orth isa CUSOLVER_HouseholderQR
+        alg_leftorth = alg_orth
+        alg_rightorth = LQViaTransposedQR(CUSOLVER_HouseholderQR(; alg_orth.kwargs...))
     elseif alg_orth isa LAPACK_HouseholderLQ
         alg_leftorth = LAPACK_HouseholderQR(; alg_orth.kwargs...)
         alg_rightorth = alg_orth
+    elseif alg_orth isa LQViaTransposedQR
+        alg_leftorth = alg_orth
+        alg_rightorth = alg_orth.qr_alg
     else
         alg_leftorth = alg_rightorth = alg_orth
     end

src/states/quasiparticle_state.jl

@@ -1,7 +1,7 @@
 #=
 Should not be constructed by the user - acts like a vector (used in eigsolve)
 I think it makes sense to see these things as an actual state instead of return an array of B tensors (what we used to do)
-This will allow us to plot energy density (finite qp) and measure observeables.
+This will allow us to plot energy density (finite qp) and measure observables.
 =#
 
 struct LeftGaugedQP{S, T1, T2, E <: Number}

src/utility/multiline.jl

@@ -106,3 +106,14 @@ function VectorInterface.inner(x::Multiline, y::Multiline)
 end
 
 LinearAlgebra.norm(x::Multiline) = sqrt(real(inner(x, x)))
+
+# TensorKit
+#----------
+
+site_type(::Type{Multiline{S}}) where {S} = site_type(S)
+bond_type(::Type{Multiline{S}}) where {S} = bond_type(S)
+site_type(st::Multiline) = site_type(typeof(st))
+bond_type(st::Multiline) = bond_type(typeof(st))
+TensorKit.sectortype(::Type{Multiline{T}}) where {T} = sectortype(T)
+TensorKit.spacetype(::Type{Multiline{T}}) where {T} = spacetype(T)
+TensorKit.storagetype(::Type{Multiline{T}}) where {T} = storagetype(T)

test/cuda/operators.jl

@@ -0,0 +1,78 @@
+using .TestSetup
+using Test, TestExtras
+using MPSKit
+using MPSKit: GeometryStyle, FiniteChainStyle, InfiniteChainStyle, OperatorStyle, MPOStyle
+using TensorKit
+using TensorKit: ℙ, tensormaptype, TensorMapWithStorage
+using Adapt, CUDA, cuTENSOR
+
+@testset "CuFiniteMPO" for V in (ℂ^2, U1Space(0 => 1, 1 => 1))
+    # start from random operators
+    L = 4
+    T = ComplexF64
+
+    O₁ = rand(T, V^L, V^L)
+    O₂ = rand(T, space(O₁))
+    O₃ = rand(real(T), space(O₁))
+
+    mpo₁ = adapt(CuArray, FiniteMPO(O₁))
+    mpo₂ = adapt(CuArray, FiniteMPO(O₂))
+    mpo₃ = adapt(CuArray, FiniteMPO(O₃))
+
+    @test isfinite(mpo₁)
+    @test isfinite(typeof(mpo₁))
+    @test GeometryStyle(typeof(mpo₁)) == FiniteChainStyle()
+    @test GeometryStyle(mpo₁) == FiniteChainStyle()
+    @test OperatorStyle(typeof(mpo₁)) == MPOStyle()
+
+    @test @constinferred physicalspace(mpo₁) == fill(V, L)
+    Vleft = @constinferred left_virtualspace(mpo₁)
+    Vright = @constinferred right_virtualspace(mpo₂)
+    for i in 1:L
+        @test Vleft[i] == left_virtualspace(mpo₁, i)
+        @test Vright[i] == right_virtualspace(mpo₁, i)
+    end
+
+    TM = TensorMapWithStorage{T, CuVector{T, CUDA.DeviceMemory}}
+    #@test convert(TM, mpo₁) ≈ O₁
+    #@test convert(TM, -mpo₂) ≈ -O₂
+    #@test convert(TM, @constinferred complex(mpo₃)) ≈ complex(O₃)
+
+
+    # test scalar multiplication
+    α = rand(T)
+    #@test convert(TM, α * mpo₁) ≈ α * O₁
+    #@test convert(TM, mpo₁ * α) ≈ O₁ * α
+    @test α * mpo₃ ≈ α * complex(mpo₃) atol = 1.0e-6
+
+    # test addition and multiplication
+    #@test convert(TM, mpo₁ + mpo₂) ≈ O₁ + O₂
+    #@test convert(TM, mpo₁ + mpo₃) ≈ O₁ + O₃
+    #@test convert(TM, mpo₁ * mpo₂) ≈ O₁ * O₂
+    #@test convert(TM, mpo₁ * mpo₃) ≈ O₁ * O₃
+
+    # test application to a state
+    ψ₁ = adapt(CuArray, rand(T, domain(O₁)))
+    #ψ₂ = adapt(CuArray, rand(real(T), domain(O₂))) # not allowed due to cuTENSOR
+    mps₁ = adapt(CuArray, FiniteMPS(ψ₁))
+    #mps₂ = adapt(CuArray, FiniteMPS(ψ₂))
+
+    @test @constinferred GeometryStyle(mps₁, mpo₁, mps₁) == GeometryStyle(mps₁)
+
+    #@test convert(TM, mpo₁ * mps₁) ≈ O₁ * ψ₁
+    @test mpo₁ * ψ₁ ≈ O₁ * ψ₁
+    #@test convert(TM, mpo₃ * mps₁) ≈ O₃ * ψ₁
+    @test mpo₃ * ψ₁ ≈ O₃ * ψ₁
+    #@test convert(TM, mpo₁ * mps₂) ≈ O₁ * ψ₂
+    #@test mpo₁ * ψ₂ ≈ O₁ * ψ₂
+
+    @test dot(mps₁, mpo₁, mps₁) ≈ dot(ψ₁, O₁, ψ₁)
+    @test dot(mps₁, mpo₁, mps₁) ≈ dot(mps₁, mpo₁ * mps₁)
+    # test conversion to and from mps
+    mpomps₁ = convert(FiniteMPS, mpo₁)
+    mpompsmpo₁ = convert(FiniteMPO, mpomps₁)
+
+    @test convert(FiniteMPO, mpomps₁) ≈ mpo₁ rtol = 1.0e-6
+
+    @test dot(mpomps₁, mpomps₁) ≈ dot(mpo₁, mpo₁)
+end