[Performance] TreeTransformer refactor + multithreading

Project.toml

@@ -1,14 +1,13 @@
 name = "TensorKit"
 uuid = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 authors = ["Jutho Haegeman"]
-version = "0.14.6"
+version = "0.14.7"
 
 [deps]
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
@@ -33,7 +32,6 @@ LRUCache = "1.0.2"
 LinearAlgebra = "1"
 PackageExtensionCompat = "1"
 Random = "1"
-SparseArrays = "1"
 Strided = "2"
 TensorKitSectors = "0.1"
 TensorOperations = "5.1"

benchmark/TensorKitBenchmarks/TensorKitBenchmarks.jl

@@ -4,7 +4,7 @@ using BenchmarkTools
 using TensorKit
 using TOML
 
-BenchmarkTools.DEFAULT_PARAMETERS.seconds = 1.0
+BenchmarkTools.DEFAULT_PARAMETERS.seconds = 20.0
 BenchmarkTools.DEFAULT_PARAMETERS.samples = 10000
 BenchmarkTools.DEFAULT_PARAMETERS.time_tolerance = 0.15
 BenchmarkTools.DEFAULT_PARAMETERS.memory_tolerance = 0.01

benchmark/TensorKitBenchmarks/tensornetworks/TensorNetworkBenchmarks.jl

@@ -45,9 +45,8 @@ function benchmark_mpo!(bench; sigmas=nothing, T="Float64", I="Trivial", dims)
 end
 
 if haskey(all_parameters, "mpo")
-    g = addgroup!(SUITE, "mpo")
     for params in all_parameters["mpo"]
-        benchmark_mpo!(g, params)
+        benchmark_mpo!(SUITE, params)
     end
 end
 
@@ -90,9 +89,8 @@ function benchmark_pepo!(bench; sigmas=nothing, T="Float64", I="Trivial", dims)
 end
 
 if haskey(all_parameters, "pepo")
-    g = addgroup!(SUITE, "pepo")
     for params in all_parameters["pepo"]
-        benchmark_pepo!(g, params)
+        benchmark_pepo!(SUITE, params)
     end
 end
 
@@ -136,9 +134,8 @@ function benchmark_mera!(bench; sigmas=nothing, T="Float64", I="Trivial", dims)
 end
 
 if haskey(all_parameters, "mera")
-    g = addgroup!(SUITE, "mera")
     for params in all_parameters["mera"]
-        benchmark_mera!(g, params)
+        benchmark_mera!(SUITE, params)
     end
 end
 

benchmark/TensorKitBenchmarks/tensornetworks/benchparams.toml

@@ -19,6 +19,12 @@ I = "U1Irrep"
 dims = [[40, 5, 3], [160, 5, 3], [640, 5, 3], [2560, 5, 3], [6120, 5, 3], [200, 20, 20], [400, 20, 20], [400, 40, 40]]
 sigmas = [0.5, 0.5, 0.5]
 
+[[mpo]]
+T = ["Float64"]
+I = "SU2Irrep"
+dims = [[40, 5, 3], [160, 5, 3], [640, 5, 3], [2560, 5, 3], [6120, 5, 3], [200, 20, 20], [400, 20, 20], [400, 40, 40]]
+sigmas = [2, 2, 2]
+
 # PEPO
 # ----
 # dims = [peps, pepo, phys, env]
@@ -40,6 +46,12 @@ I = "U1Irrep"
 dims = [[4, 2, 2, 100], [4, 4, 4, 200], [6, 2, 2, 100], [6, 3, 4, 200], [8, 2, 2, 100], [8, 2, 4, 200], [10, 2, 2, 50], [10, 3, 2, 100]]
 sigmas = [0.5, 0.5, 0.5, 0.5]
 
+[[pepo]]
+T = ["Float64"]
+I = "SU2Irrep"
+dims = [[4, 2, 2, 100], [4, 4, 4, 200], [6, 2, 2, 100], [6, 3, 4, 200], [8, 2, 2, 100], [8, 2, 4, 200], [10, 2, 2, 50], [10, 3, 2, 100]]
+sigmas = [2.0, 2.0, 2.0, 2.0]
+
 # MERA
 # ----
 # dims = mera
@@ -60,3 +72,9 @@ T = ["Float64"]
 I = "U1Irrep"
 dims = [4, 8, 12, 16, 22, 28]
 sigmas = [0.5]
+
+[[mera]]
+T = ["Float64"]
+I = "SU2Irrep"
+dims = [4, 8, 12, 16, 22, 28]
+sigmas = [2.0]

benchmark/TensorKitBenchmarks/utils/BenchUtils.jl

@@ -55,13 +55,15 @@ function generate_space(::Type{U1Irrep}, D::Int, sigma::Real=0.5)
     return U1Space((s => d for (s, d) in zip(sectors, dims))...)
 end
 function generate_space(::Type{SU2Irrep}, D::Int, sigma::Real=0.5)
-    poisson_pdf(x) = ceil(Int, D * exp(-sigma) * sigma^x / factorial(x + 1))
+    normal_pdf = let D = D
+        x -> D * exp(-0.5 * (x / sigma)^2) / (sigma * sqrt(2π))
+    end
 
     sectors = SU2Irrep[]
     dims = Int[]
 
     for sector in values(SU2Irrep)
-        d = poisson_pdf(Int(sector.j * 2))
+        d = ceil(Int, normal_pdf(sector.j) / dim(sector))
         push!(sectors, sector)
         push!(dims, d)
         D -= d * dim(sector)

src/TensorKit.jl

@@ -126,8 +126,6 @@ using LinearAlgebra: norm, dot, normalize, normalize!, tr,
                      isposdef, isposdef!, ishermitian, rank, cond,
                      Diagonal, Hermitian
 
-using SparseArrays: SparseMatrixCSC, sparse, nzrange, rowvals, nonzeros
-
 import Base.Meta
 
 using Random: Random, rand!, randn!
@@ -185,6 +183,21 @@ include("fusiontrees/fusiontrees.jl")
 #-------------------------------------------
 include("spaces/vectorspaces.jl")
 
+# Multithreading settings
+#-------------------------
+const TRANSFORMER_THREADS = Ref(1)
+
+get_num_transformer_threads() = TRANSFORMER_THREADS[]
+
+function set_num_transformer_threads(n::Int)
+    N = Base.Threads.nthreads()
+    if n > N
+        n = N
+        Strided._set_num_threads_warn(n)
+    end
+    return TRANSFORMER_THREADS[] = n
+end
+
 # Definitions and methods for tensors
 #-------------------------------------
 # general definitions

src/auxiliary/auxiliary.jl

@@ -55,3 +55,27 @@ _interleave(::Tuple{}, ::Tuple{}) = ()
 function _interleave(a::NTuple{N}, b::NTuple{N}) where {N}
     return (a[1], b[1], _interleave(tail(a), tail(b))...)
 end
+
+# Low-overhead implementation of `copyto!` for specific case of `stride(B, 1) < stride(B, 2)`
+# used in indexmanipulations: avoids the overhead of Strided.jl
+function _copyto!(A::StridedView{<:Any,1}, B::StridedView{<:Any,2})
+    length(A) == length(B) || throw(DimensionMismatch())
+
+    Adata = parent(A)
+    Astr = stride(A, 1)
+    IA = A.offset
+
+    Bdata = parent(B)
+    Bstr = strides(B)
+
+    IB_1 = B.offset
+    @inbounds for _ in axes(B, 2)
+        IB = IB_1
+        for _ in axes(B, 1)
+            Adata[IA += Astr] = Bdata[IB += Bstr[1]]
+        end
+        IB_1 += Bstr[2]
+    end
+
+    return A
+end

src/spaces/homspace.jl

@@ -225,11 +225,15 @@ end
 
 # Block and fusion tree ranges: structure information for building tensors
 #--------------------------------------------------------------------------
+
+# sizes, strides, offset
+const StridedStructure{N} = Tuple{NTuple{N,Int},NTuple{N,Int},Int}
+
 struct FusionBlockStructure{I,N,F₁,F₂}
     totaldim::Int
     blockstructure::SectorDict{I,Tuple{Tuple{Int,Int},UnitRange{Int}}}
     fusiontreelist::Vector{Tuple{F₁,F₂}}
-    fusiontreestructure::Vector{Tuple{NTuple{N,Int},NTuple{N,Int},Int}}
+    fusiontreestructure::Vector{StridedStructure{N}}
     fusiontreeindices::FusionTreeDict{Tuple{F₁,F₂},Int}
 end