Refactor and add multithreading

Author: lkdvos

src/TensorKit.jl

@@ -183,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
@@ -218,6 +233,8 @@ include("auxiliary/deprecate.jl")
 # ----------
 function __init__()
     @require_extensions
+    set_num_transformer_threads(Threads.nthreads())
+    return nothing
 end
 
 end

src/tensors/indexmanipulations.jl

@@ -473,38 +473,68 @@ function add_transform!(tdst::AbstractTensorMap,
     return tdst
 end
 
-function add_transform_kernel!(tdst::TensorMap,
-                               tsrc::TensorMap,
-                               (p₁, p₂)::Index2Tuple,
-                               ::TrivialTreeTransformer,
-                               α::Number,
-                               β::Number,
-                               backend::AbstractBackend...)
-    return TO.tensoradd!(tdst[], tsrc[], (p₁, p₂), false, α, β, backend...)
+function use_threaded_transform(t::TensorMap, transformer::TreeTransformer)
+    # TODO: heuristic for not threading over small tensors
+    return get_num_transformer_threads() > 1
 end
 
 function add_transform_kernel!(tdst::TensorMap,
                                tsrc::TensorMap,
                                p::Index2Tuple,
-                               transformer::AbelianTreeTransformer,
+                               transformer::TreeTransformer,
                                α::Number,
                                β::Number,
                                backend::AbstractBackend...)
-    # TODO: this could be multithreaded
+    if use_threaded_transform(tsrc, transformer)
+        _add_transform_threaded!(tdst, tsrc, p, transformer, α, β, backend...)
+    else
+        _add_transform_nonthreaded!(tdst, tsrc, p, transformer, α, β, backend...)
+    end
+
+    return nothing
+end
+
+# Trivial implementation
+# ----------------------
+# Hijack before threading is used
+function add_transform_kernel!(tdst::TensorMap, tsrc::TensorMap, (p₁, p₂)::Index2Tuple,
+                               ::TrivialTreeTransformer,
+                               α::Number, β::Number, backend::AbstractBackend...)
+    TO.tensoradd!(tdst[], tsrc[], (p₁, p₂), false, α, β, backend...)
+    return nothing
+end
+
+# Abelian implementations
+# -----------------------
+function _add_transform_nonthreaded!(tdst, tsrc, p, transformer::AbelianTreeTransformer,
+                                     α, β, backend...)
     for subtransformer in transformer.data
-        _add_transform_single!(tdst, tsrc, p, α, β, subtransformer, backend...)
+        _add_transform_single!(tdst, tsrc, p, subtransformer, α, β, backend...)
     end
+    return nothing
+end
 
+function _add_transform_threaded!(tdst, tsrc, p, transformer::AbelianTreeTransformer, α, β,
+                                  backend...; ntasks::Int=get_num_transformer_threads())
+    nblocks = length(transformer.data)
+    counter = Threads.Atomic{Int}(1)
+    Threads.@sync for _ in 1:min(ntasks, nblocks)
+        Threads.@spawn begin
+            while true
+                local_counter = Threads.atomic_add!(counter, 1)
+                local_counter > nblocks && break
+                @inbounds subtransformer = transformer.data[local_counter]
+                _add_transform_single!(tdst, tsrc, p, subtransformer, α, β, backend...)
+            end
+        end
+    end
     return nothing
 end
 
-function add_transform_kernel!(tdst::TensorMap,
-                               tsrc::TensorMap,
-                               p::Index2Tuple,
-                               transformer::GenericTreeTransformer,
-                               α::Number,
-                               β::Number,
-                               backend::AbstractBackend...)
+# Non-abelian implementations
+# ---------------------------
+function _add_transform_nonthreaded!(tdst, tsrc, p, transformer::GenericTreeTransformer,
+                                     α, β, backend...)
     # preallocate buffers
     buffersize = maximum(transformer.data) do (_, structures_dst, _)
         return prod(structures_dst[1][1])
@@ -522,24 +552,59 @@ function add_transform_kernel!(tdst::TensorMap,
                                   (buffer1, buffer2), backend...)
         end
     end
-    return tdst
+    return nothing
 end
 
-function _add_transform_single!(tdst, tsrc, p, α, β,
+function _add_transform_threaded!(tdst, tsrc, p, transformer::GenericTreeTransformer,
+                                  α, β, backend...;
+                                  ntasks::Int=get_num_transformer_threads())
+    buffersize = maximum(transformer.data) do (_, structures_dst, _)
+        return prod(structures_dst[1][1])
+    end
+    nblocks = length(transformer.data)
+
+    counter = Threads.Atomic{Int}(1)
+    Threads.@sync for _ in 1:min(ntasks, nblocks)
+        Threads.@spawn begin
+            # preallocate buffers for each task
+            buffer1 = similar(tsrc.data, buffersize)
+            buffer2 = similar(tdst.data, buffersize)
+
+            while true
+                local_counter = Threads.atomic_add!(counter, 1)
+                local_counter > nblocks && break
+                @inbounds subtransformer = transformer.data[local_counter]
+                if length(subtransformer[1]) == 1
+                    _add_transform_single!(tdst, tsrc, p, subtransformer, α, β, backend...)
+                else
+                    _add_transform_multi!(tdst, tsrc, p, subtransformer, (buffer1, buffer2),
+                                          α, β, backend...)
+                end
+            end
+        end
+    end
+
+    return nothing
+end
+
+# Kernels
+# -------
+function _add_transform_single!(tdst, tsrc, p,
                                 (basistransform, structures_dst, structures_src),
-                                backend...)
+                                α, β, backend...)
     structure_dst = structures_dst isa Vector ? only(structures_dst) : structures_dst
     structure_src = structures_src isa Vector ? only(structures_src) : structures_src
     coeff = basistransform isa Number ? basistransform : only(basistransform)
+
     subblock_dst = StridedView(tdst.data, structure_dst...)
     subblock_src = StridedView(tsrc.data, structure_src...)
     TO.tensoradd!(subblock_dst, subblock_src, p, false, α * coeff, β, backend...)
     return nothing
 end
 
-function _add_transform_multi!(tdst, tsrc, p, α, β,
+function _add_transform_multi!(tdst, tsrc, p,
                                (basistransform, structures_dst, structures_src),
-                               (buffer1, buffer2), backend...)
+                               (buffer1, buffer2), α, β, backend...)
     rows, cols = size(basistransform)
     sz_src = first(first(structures_src))
     blocksize = prod(sz_src)
@@ -553,7 +618,7 @@ function _add_transform_multi!(tdst, tsrc, p, α, β,
 
     # Resummation into a second buffer using BLAS
     buffer_dst = StridedView(buffer2, (blocksize, rows), (1, blocksize), 0)
-    mul!(buffer_dst, buffer_src, basistransform, α)
+    mul!(buffer_dst, buffer_src, basistransform, α, Zero())
 
     # Filling up the output
     for (i, structure_dst) in enumerate(structures_dst)
@@ -565,6 +630,9 @@ function _add_transform_multi!(tdst, tsrc, p, α, β,
     return nothing
 end
 
+# Other implementations
+# ---------------------
+
 function add_transform_kernel!(tdst::AbstractTensorMap,
                                tsrc::AbstractTensorMap,
                                (p₁, p₂)::Index2Tuple,

src/tensors/treetransformers.jl

@@ -8,7 +8,7 @@ abstract type TreeTransformer end
 struct TrivialTreeTransformer <: TreeTransformer end
 
 struct AbelianTreeTransformer{T,N} <: TreeTransformer
-    data::Vector{Tuple{StridedStructure{N},StridedStructure{N},T}}
+    data::Vector{Tuple{T,StridedStructure{N},StridedStructure{N}}}
 end
 
 function AbelianTreeTransformer(transform, p, Vsrc, Vdst)
@@ -19,15 +19,15 @@ function AbelianTreeTransformer(transform, p, Vsrc, Vdst)
     L = length(structure_src.fusiontreelist)
     T = sectorscalartype(sectortype(Vdst))
     N = numind(Vsrc)
-    data = Vector{Tuple{StridedStructure{N},StridedStructure{N},T}}(undef, L)
+    data = Vector{Tuple{T,StridedStructure{N},StridedStructure{N}}}(undef, L)
 
     for i in 1:L
         f₁, f₂ = structure_src.fusiontreelist[i]
         (f₃, f₄), coeff = only(transform(f₁, f₂))
         j = structure_dst.fusiontreeindices[(f₃, f₄)]
         stridestructure_dst = structure_dst.fusiontreestructure[j]
         stridestructure_src = structure_src.fusiontreestructure[i]
-        data[i] = (stridestructure_dst, stridestructure_src, coeff)
+        data[i] = (coeff, stridestructure_dst, stridestructure_src)
     end
 
     return AbelianTreeTransformer(data)