update to CUDA 4.4.1, TensorOperations 4 and TensorCast 0.4. SEE README

Project.toml

@@ -1,7 +1,7 @@
 name = "SpinGlassTensors"
 uuid = "7584fc6a-5a23-4eeb-8277-827aab0146ea"
 authors = ["Anna Maria Dziubyna <[email protected]>", "Tomasz Śmierzchalski <[email protected]>", "Bartłomiej Gardas <[email protected]>", "Konrad Jałowiecki <[email protected]>", "Łukasz Pawela <[email protected]>", "Marek M. Rams <[email protected]>"]
-version = "1.0"
+version = "1.0.0"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -11,17 +11,24 @@ LowRankApprox = "898213cb-b102-5a47-900c-97e73b919f73"
 MKL = "33e6dc65-8f57-5167-99aa-e5a354878fb2"
 Memoization = "6fafb56a-5788-4b4e-91ca-c0cea6611c73"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 TSVD = "9449cd9e-2762-5aa3-a617-5413e99d722e"
 TensorCast = "02d47bb6-7ce6-556a-be16-bb1710789e2b"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
+TransmuteDims = "24ddb15e-299a-5cc3-8414-dbddc482d9ca"
+cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [compat]
+CUDA = "4.4.1"
+DocStringExtensions = "0.9.3"
+LowRankApprox = "0.5.5"
 MKL = "0.4.2"
+Memoization = "0.2.1"
+SparseArrays = "1.9"
 TensorCast = "0.4"
-TensorOperations = "3.0.1"
-julia = "1.8, 1.9"
+TensorOperations = "4"
+cuTENSOR = "1.1.0"
+julia = "1.9, 1.10"
 
 [extras]
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"

README.md

@@ -1,2 +1,3 @@
 [![Coverage Status](https://coveralls.io/repos/github/iitis/SpinGlassTensors.jl/badge.svg?branch=master)](https://coveralls.io/github/iitis/SpinGlassTensors.jl?branch=master)
 # SpinGlassTensors.jl
+This works with CUDA v4.4.1. You need to manually `]add TransmuteDims#strided2`

bench_mm.jl

@@ -0,0 +1,29 @@
+using TensorCast, TensorOperations
+function time_mm()
+    M = rand(100, 100, 100)
+    L = rand(100, 100)
+    R = rand(100, 100)
+    @time begin
+        @matmul M1[x, σ, α] := sum(β) L[x, β] * M[β, σ, α] 
+        @matmul MM[x, σ, y] := sum(α) M1[x, σ, α] * R[α, y]
+    end
+end
+
+function time_tensor()
+    M = rand(100, 100, 100)
+    L = rand(100, 100)
+    R = rand(100, 100)
+
+    @time begin
+        @tensor M̃[x, σ, y] := L[x, β] * M[β, σ, α] * R[α, y] order = (α, β)
+        # @cast B[(x, σ), y] |= M̃[x, σ, y]
+    end
+end
+
+println("matmul")
+time_mm()
+time_mm()
+
+println("\n tensor")
+time_tensor()
+time_tensor()

src/SpinGlassTensors.jl

@@ -3,7 +3,8 @@ module SpinGlassTensors
     using TensorOperations, TensorCast
     using LowRankApprox, TSVD
     using CUDA, CUDA.CUSPARSE
-    using NNlib, NNlibCUDA
+    using cuTENSOR
+    using NNlib
     using Memoization
     using SparseArrays
     using DocStringExtensions

src/contractions/dense.jl

@@ -26,7 +26,7 @@ end
         -- B --
 """
 function update_env_left(LE::S, A::S, M::T, B::S) where {S <: Tensor{R, 3}, T <: Tensor{R, 4}} where R <: Real
-    @tensor LE[nb, nt, nc] := LE[ob, ot, oc] * A[ot, nt, α] * M[oc, α, nc, β] * B[ob, nb, β] order = (ot, α, oc, β, ob) # TODO: split the line
+    @tensor order = (ot, α, oc, β, ob) LE[nb, nt, nc] := LE[ob, ot, oc] * A[ot, nt, α] * M[oc, α, nc, β] * B[ob, nb, β] # TODO: split the line
 end
 
 """
@@ -37,7 +37,7 @@ end
         -- B --
 """
 function update_env_left(LE::T, A::S, B::S) where {S <: Tensor{R, 3}, T <: Tensor{R, 2}} where R <: Real
-    @tensor LE[nb, nt] := LE[ob, ot] * A[ot, nt, α] * B[ob, nb, α] order = (ot, α, ob)
+    @tensor order = (ot, α, ob) LE[nb, nt] := LE[ob, ot] * A[ot, nt, α] * B[ob, nb, α]
 end
 
 """
@@ -59,7 +59,7 @@ end
       -- B --
 """
 function update_env_right(RE::S, A::S, M::T, B::S) where {T <: Tensor{R, 4}, S <: Tensor{R, 3}} where R <: Real
-    @tensor RE[nb, nt, nc] := RE[ob, ot, oc] * A[nt, ot, α] * M[nc, α, oc, β] * B[nb, ob, β] order = (ot, α, oc, β, ob)
+    @tensor order = (ot, α, oc, β, ob) RE[nb, nt, nc] := RE[ob, ot, oc] * A[nt, ot, α] * M[nc, α, oc, β] * B[nb, ob, β]
 end
 
 """
@@ -70,7 +70,7 @@ end
       -- B --
 """
 function update_env_right(RE::T, A::S, B::S) where {T <: Tensor{R, 2}, S <: Tensor{R, 3}} where R <: Real
-    @tensor RE[nb, nt] := RE[ob, ot] * A[nt, ot, α] * B[nb, ob, α] order = (ot, α, ob)
+    @tensor order = (ot, α, ob) RE[nb, nt] := RE[ob, ot] * A[nt, ot, α] * B[nb, ob, α]
 end
 
 """
@@ -81,7 +81,7 @@ end
 
 """
 function update_env_right(RE::S, C::S) where {S <: Tensor{R, 3}} where R <: Real
-    @tensor RR[nb, nt] := RE[nb, ot, oc] * C[nt, ot, oc] order = (ot, oc)
+    @tensor order = (ot, oc) RR[nb, nt] := RE[nb, ot, oc] * C[nt, ot, oc]
 end
 
 """
@@ -91,7 +91,7 @@ end
      -- B --
 """
 function project_ket_on_bra(LE::S, B::S, M::T, RE::S) where {T <: Tensor{R, 4}, S <: Tensor{R, 3}} where R <: Real
-    @tensor A[nl, nr, nc] := LE[ol, nl, lc] * B[ol, or, oc] * M[lc, nc, rc, oc] * RE[or, nr, rc] order = (ol, lc, oc, or, rc)
+    @tensor order = (ol, lc, oc, or, rc) A[nl, nr, nc] := LE[ol, nl, lc] * B[ol, or, oc] * M[lc, nc, rc, oc] * RE[or, nr, rc]
 end
 
 """
@@ -100,7 +100,7 @@ end
      -- B --
 """
 function project_ket_on_bra(LE::T, B::S, RE::T) where {T <: Tensor{R, 2}, S <: Tensor{R, 3}} where R <: Real
-    @tensor A[nl, nr, nc] := LE[ol, nl] * B[ol, or, nc] * RE[or, nr] order = (ol, or)
+    @tensor order = (ol, or) A[nl, nr, nc] := LE[ol, nl] * B[ol, or, nc] * RE[or, nr]
 end
 
 """
@@ -135,7 +135,7 @@ end
 function update_reduced_env_right(
     K::Tensor{R, 1}, RE::Tensor{R, 2}, M::Tensor{R, 4}, B::Tensor{R, 3}
 ) where R <: Real
-    @tensor RE[x, y] := K[d] * M[y, d, β, γ] * B[x, α, γ] * RE[α, β] order = (d, β, γ, α)
+    @tensor order = (d, β, γ, α) RE[x, y] := K[d] * M[y, d, β, γ] * B[x, α, γ] * RE[α, β]
 end
 
 function update_reduced_env_right(RR::S, M0::S) where S <: Tensor{<:Real, 2}
@@ -147,5 +147,5 @@ function contract_tensors43(B::Tensor{R, 4}, A::Tensor{R, 3}) where R <: Real
 end
 
 function corner_matrix(C::S, M::T, B::S) where {S <: Tensor{R, 3}, T <: Tensor{R, 4}} where R <: Real
-    @tensor Cnew[ll, ml, tt, mt] := M[ml, mt, mr, mb] * B[ll, rr, mb] * C[rr, tt, mr] order = (rr, mb, mr)
+    @tensor order = (rr, mb, mr) Cnew[ll, ml, tt, mt] := M[ml, mt, mr, mb] * B[ll, rr, mb] * C[rr, tt, mr] 
 end

src/contractions/virtual.jl

@@ -378,7 +378,7 @@ function contract_tensors43(M::VirtualTensor{R, 4}, B::Tensor{R, 3}) where R <:
     else
         pls = contract_tensor3_matrix(pls, M.con)
     end
-    @tensor MB[l, lt, r, rt] := pls[lt, lb, c] * prs[rt, rb, c] * B[l, r, lb, rb]  order=(lb, c, rb)
+    @tensor order=(lb, c, rb) MB[l, lt, r, rt] := pls[lt, lb, c] * prs[rt, rb, c] * B[l, r, lb, rb]
     MB = reshape(MB, slb, slct, slcp, srb, srct, srcp)
     MB = permutedims(MB, (1, 3, 4, 6, 2, 5))
     reshape(MB, (slb * slcp, srb * srcp, slct * srct))

src/gauges.jl

@@ -122,7 +122,7 @@ function _overlap_forward(ϕ::QMps{T}, ψ::QMps{T}, k::Site) where T <: Real
     i = ψ.sites[1]
     while i < k
         A, B = ψ[i], ϕ[i]
-        @tensor C[x, y] := conj(B)[β, σ, x] * C[β, α] * A[α, σ, y] order = (α, β, σ)
+        @tensor order = (α, β, σ) C[x, y] := conj(B)[β, σ, x] * C[β, α] * A[α, σ, y]
         i += 1
     end
     C
@@ -133,7 +133,7 @@ function _overlap_backwards(ϕ::QMps{T}, ψ::QMps{T}, k::Site) where T <: Real
     i = ψ.sites[end]
     while i > k
         A, B = ψ[i], ϕ[i]
-        @tensor D[x, y] := conj(B)[x, σ, β] * D[β, α] * A[y, σ, α] order = (α, β, σ)
+        @tensor order = (α, β, σ) D[x, y] := conj(B)[x, σ, β] * D[β, α] * A[y, σ, α]
         i -= 1
     end
     D

src/linear_algebra_ext.jl

@@ -10,7 +10,7 @@ export
 @inline phase(d::T; atol=eps()) where T <: Real = isapprox(d, zero(T), atol=atol) ? one(T) : d / abs(d)
 @inline phase(d::AbstractArray; atol=eps()) = map(x -> phase(x; atol=atol), d)
 
-function svd_fact(A::AbstractMatrix{T}, Dcut::Int=typemax(Int), tol::T=eps(), kwargs...) where T <: Real
+function svd_fact(A::AbstractMatrix{T}, Dcut::Int=typemax(Int), tol=eps(T); kwargs...) where T <: Real
     U, Σ, V = svd(A; kwargs...)
     δ = min(Dcut, sum(Σ .> Σ[1] * max(eps(), tol)))
     U, Σ, V = U[:, 1:δ], Σ[1:δ], V[:, 1:δ]

src/mps/canonise.jl

@@ -52,7 +52,7 @@ function canonise_truncate!(ψ::QMps, type::Symbol, Dcut::Int=typemax(Int), tolS
     end
 end
 
-function _right_sweep!(ψ::QMps{T}, Dcut::Int=typemax(Int), tolS::T=eps(); kwargs...) where T <: Real
+function _right_sweep!(ψ::QMps{T}, Dcut::Int=typemax(Int), tolS::T=eps(T); kwargs...) where T <: Real
     R = ψ.onGPU ? CUDA.ones(T, 1, 1) : ones(T, 1, 1)
     for i ∈ ψ.sites
         A = ψ[i]
@@ -66,7 +66,7 @@ function _right_sweep!(ψ::QMps{T}, Dcut::Int=typemax(Int), tolS::T=eps(); kwarg
     end
 end
 
-function _left_sweep!(ψ::QMps{T}, Dcut::Int=typemax(Int), tolS::T=eps(); kwargs...) where T <: Real
+function _left_sweep!(ψ::QMps{T}, Dcut::Int=typemax(Int), tolS::T=eps(T); kwargs...) where T <: Real
     R = ψ.onGPU ? CUDA.ones(T, 1, 1) : ones(T, 1, 1)
     for i ∈ reverse(ψ.sites)
         B = permutedims(ψ[i], (1, 3, 2)) # [x, σ, α]

src/mps/dot.jl

@@ -11,7 +11,7 @@ function LinearAlgebra.dot(ψ::QMps{T}, ϕ::QMps{T}) where T <: Real
     C = ψ.onGPU && ϕ.onGPU ? CUDA.ones(T, 1, 1) : ones(T, 1, 1)
     for i ∈ ϕ.sites
         A, B = ϕ[i], ψ[i]
-        @tensor C[x, y] := conj(B)[β, x, σ] * C[β, α] * A[α, y, σ] order = (α, β, σ)
+        @tensor order = (α, β, σ) C[x, y] := conj(B)[β, x, σ] * C[β, α] * A[α, y, σ]
     end
     tr(C)
 end

src/mps/utils.jl

@@ -29,19 +29,35 @@ function is_consistent(ψ::QMps)
 end
 
 function eye(::Type{T}, dim; toGPU::Bool=false) where T
-    id = Diagonal(ones(T, dim))
-    toGPU && return CuArray(id)
-    id
+    v = ones(T, dim)
+    toGPU && return cu(spdiagm(v))
+    Diagonal(v)
 end
 
-function is_left_normalized(ψ::QMps)
+function is_left_normalized(ψ::QMps, ::Val{false})
     all(
-        eye(eltype(ψ), size(A, 2); toGPU = ψ.onGPU) ≈ @tensor Id[x, y] := A[α, x, σ] * A[α, y, σ] order = (α, σ) for A ∈ values(ψ.tensors) # TODO: split the line
+        eye(eltype(ψ), size(A, 2); toGPU = false) ≈ @tensor(Id[x, y] := A[α, x, σ] * A[α, y, σ];  order = (α, σ)) for A ∈ values(ψ.tensors) # TODO: split the line
         )
 end
 
-function is_right_normalized(ψ::QMps)
+function is_left_normalized(ψ::QMps, ::Val{true})
     all(
-        eye(eltype(ψ), size(B, 1); toGPU = ψ.onGPU) ≈ @tensor Id[x, y] := B[x, α, σ] * B[y, α, σ] order = (α, σ) for B ∈ values(ψ.tensors) # TODO: split the line
+        eye(eltype(ψ), size(A, 2); toGPU = true) ≈ @cutensor(Id[x, y] := A[α, x, σ] * A[α, y, σ]) for A ∈ values(ψ.tensors) # TODO: split the line
         )
 end
+
+is_left_normalized(ψ::QMps) = is_left_normalized(ψ, Val(ψ.onGPU))
+
+function is_right_normalized(ψ::QMps, ::Val{false})
+    all(
+        eye(eltype(ψ), size(B, 1); toGPU = false) ≈ @tensor(Id[x, y] := B[x, α, σ] * B[y, α, σ];  order = (α, σ)) for B ∈ values(ψ.tensors) # TODO: split the line
+        )
+end
+
+function is_right_normalized(ψ::QMps, ::Val{true})
+    all(
+        eye(eltype(ψ), size(B, 1); toGPU = true) ≈ @cutensor(Id[x, y] := B[x, α, σ] * B[y, α, σ]) for B ∈ values(ψ.tensors) # TODO: split the line
+        )
+end
+
+is_right_normalized(ψ::QMps) = is_right_normalized(ψ, Val(ψ.onGPU))

test/attic/compressions.jl

@@ -10,27 +10,26 @@
     tol = 1E-10
 
     ψ = rand(QMps{T}, sites, D, d)
-    #W = rand(QMpo{T}, [1,2,3,4], 2, 4)
+    W = rand(QMpo{T}, [1,2,3,4], 2, 4)
 
     bra = ψ
     ket = ψ
-    #mpo = W
+    mpo = W
 
     @testset "Two mps representations are compressed to the same state" begin
-        #χ = W * ψ
-        #@test is_left_normalized(χ)
+        χ = W * ψ
+        @test is_left_normalized(χ)
 
-        #ϕ = copy(ψ)
+        ϕ = copy(ψ)
         @test bond_dimension(bra) == max(D, d)
         @test bond_dimensions(bra) == [(1, d, D), (D, d, D), (D, d, D), (D, d, 1)]
         canonise!(bra, :left)
-        #bra = QMps(ψ)
+        bra = QMps(ψ)
 
-        #@time overlap, env = variational_compress!(bra, mpo, ket, tol, max_sweeps)
-        #println(overlap)
+        @time overlap, env = variational_compress!(bra, mpo, ket, tol, max_sweeps)
 
-        #ϕ = MPS(bra)
-        #@time is_right_normalized(ϕ)
-        #@test norm(χ) ≈ norm(bra) ≈ 1
+        ϕ = MPS(bra)
+        @time is_right_normalized(ϕ)
+        @test norm(χ) ≈ norm(bra) ≈ 1
     end
 end

test/attic/linear_algebra_ext.jl

@@ -42,8 +42,5 @@ end
     r1 = U1 * Diagonal(Σ1) * V1'
     r2 = U2 * Diagonal(Σ2) * V2'
 
-    println(norm(a-r2))
-    println(norm(a-r1))
-
     @test norm(r1-r2) < tol
 end

test/canonise.jl

@@ -1,12 +1,11 @@
-T = Float64
 D = 16
 
 sites = [1, 3//2, 2, 5//2, 3, 7//2, 4]
 d = [1, 2, 2, 2, 4, 2, 2]
 
 id = Dict(j => d[i] for (i, j) in enumerate(sites))
 
-@testset "Random QMps" begin
+@testset "Random QMps ($T)" for T in (Float32, Float64)
     for toCUDA ∈ (true, false)
         ψ = rand(QMps{T}, id, D)
         ϕ = rand(QMps{T}, id, D)
@@ -36,10 +35,10 @@ id = Dict(j => d[i] for (i, j) in enumerate(sites))
     end
 end
 
-@testset "Measure spectrum" begin
+@testset "Measure spectrum($T)" for T in (Float32, Float64)
     svd_mps = TensorMap{T}(
-        1 => [-0.694389933025747 -0.7195989305943268;;; 0.7195989305943268 -0.6943899330257469],
-        2 => [0.7071067811865477; 0.0;;; -7.850461536237973e-17; 0.7071067811865477]
+        1 => T[-0.694389933025747 -0.7195989305943268;;; 0.7195989305943268 -0.6943899330257469],
+        2 => T[0.7071067811865477; 0.0;;; -7.850461536237973e-17; 0.7071067811865477]
     )
     ψ = QMps(svd_mps)
     @test is_left_normalized(ψ)

test/variational.jl

@@ -28,23 +28,10 @@ end
 
 @testset "Compressions for sparse mps and mpo works" begin
     W = rand(QMpo{S}, mpomap)
-    println(W)
-    println(bond_dimensions(W))
     ψ = rand(QMps{S}, id, D)
-    println(ψ)
     canonise!(ψ, :left)
-    println(bond_dimensions(ψ))
     ϕ = rand(QMps{S}, id, D)
     canonise!(ϕ, :left)
 
-    Dcut = 8
-    max_sweeps = 100
-    tol = 1E-10
 
-    # χ = W * ψ
-    # @time overlap, env = variational_compress!(ψ, W, ϕ, tol, max_sweeps)
 end
-
-
-# A = Dict(1.0 => Dict(0.0 => rand(1,2,2, 1)), 2.0 => Dict(0.0 =>rand(2,2,2,1)), 3.0 => Dict(0.0 =>rand(2,2,1,1)))
-# QMpo(A)