[Performance] TreeTransformer refactor + multithreading (#251)

This PR features a change in the implementation of the treetransformers in the generic case of non-unique fusion.

The core idea is to avoid having to access data non-contiguously as much as possible, while making use of BLAS to facilitate the recoupling.
In order to achieve this, we now loop over the set of uncoupled charges and process all fusion trees that are attached to this simultaneously.
The recoupling is done through BLAS, the copying into a temporary buffer with an optimized copy and the final permutation with Strided.
Additionally, I added and centralized the implementation of multithreading for the treetransformers.

This is a result of many discussions and inspiration from @ogauthe.

Co-authored-by: Olivier Gauthe <[email protected]>

Author: lkdvos

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