Overhaul indexmanipulations to reuse multithreading criterion

Author: lkdvos

src/tensors/indexmanipulations.jl

@@ -453,68 +453,71 @@ end
 
 function add_transform!(tdst::AbstractTensorMap,
                         tsrc::AbstractTensorMap,
-                        (p₁, p₂)::Index2Tuple,
+                        p::Index2Tuple,
                         transformer,
                         α::Number,
                         β::Number,
                         backend::AbstractBackend...)
     @boundscheck begin
-        permute(space(tsrc), (p₁, p₂)) == space(tdst) ||
+        permute(space(tsrc), p) == space(tdst) ||
             throw(SpaceMismatch("source = $(codomain(tsrc))←$(domain(tsrc)),
-            dest = $(codomain(tdst))←$(domain(tdst)), p₁ = $(p₁), p₂ = $(p₂)"))
+            dest = $(codomain(tdst))←$(domain(tdst)), p₁ = $(p[1]), p₂ = $(p[2])"))
     end
 
-    if p₁ === codomainind(tsrc) && p₂ === domainind(tsrc)
+    if p[1] === codomainind(tsrc) && p[2] === domainind(tsrc)
         add!(tdst, tsrc, α, β)
     else
-        add_transform_kernel!(tdst, tsrc, (p₁, p₂), transformer, α, β, backend...)
+        I = sectortype(tdst)
+        if I === Trivial
+            _add_trivial_kernel!(tdst, tsrc, p, transformer, α, β, backend...)
+        elseif FusionStyle(I) === UniqueFusion()
+            if use_threaded_transform(tdst, transformer)
+                _add_abelian_kernel_threaded!(tdst, tsrc, p, transformer, α, β, backend...)
+            else
+                _add_abelian_kernel_nonthreaded!(tdst, tsrc, p, transformer, α, β,
+                                                 backend...)
+            end
+        else
+            if use_threaded_transform(tdst, transformer)
+                _add_general_kernel_threaded!(tdst, tsrc, p, transformer, α, β, backend...)
+            else
+                _add_general_kernel_nonthreaded!(tdst, tsrc, p, transformer, α, β,
+                                                 backend...)
+            end
+        end
     end
 
     return tdst
 end
 
-function use_threaded_transform(t::TensorMap, transformer::TreeTransformer)
+function use_threaded_transform(t::TensorMap, transformer)
     return get_num_transformer_threads() > 1 && length(t.data) > Strided.MINTHREADLENGTH
 end
-
-function add_transform_kernel!(tdst::TensorMap,
-                               tsrc::TensorMap,
-                               p::Index2Tuple,
-                               transformer::TreeTransformer,
-                               α::Number,
-                               β::Number,
-                               backend::AbstractBackend...)
-    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
+function use_threaded_transform(t::AbstractTensorMap, transformer)
+    return get_num_transformer_threads() > 1 && dim(space(t)) > Strided.MINTHREADLENGTH
 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...)
+# Trivial implementations
+# -----------------------
+function _add_trivial_kernel!(tdst, tsrc, p, transformer, α, β, backend...)
+    TO.tensoradd!(tdst[], tsrc[], p, false, α, β, backend...)
     return nothing
 end
 
 # Abelian implementations
 # -----------------------
-function _add_transform_nonthreaded!(tdst, tsrc, p, transformer::AbelianTreeTransformer,
-                                     α, β, backend...)
+function _add_abelian_kernel_nonthreaded!(tdst, tsrc, p,
+                                          transformer::AbelianTreeTransformer,
+                                          α, β, backend...)
     for subtransformer in transformer.data
         _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())
+function _add_abelian_kernel_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)
@@ -530,14 +533,49 @@ function _add_transform_threaded!(tdst, tsrc, p, transformer::AbelianTreeTransfo
     return nothing
 end
 
+function _add_transform_single!(tdst, tsrc, p,
+                                (basistransform, structures_dst, structures_src),
+                                α, β, 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_abelian_kernel_nonthreaded!(tdst, tsrc, p, transformer, α, β, backend...)
+    for (f₁, f₂) in fusiontrees(tsrc)
+        _add_abelian_block!(tdst, tsrc, p, transformer, f₁, f₂, α, β, backend...)
+    end
+    return nothing
+end
+
+function _add_abelian_kernel_threaded!(tdst, tsrc, p, transformer, α, β, backend...)
+    Threads.@sync for (f₁, f₂) in fusiontrees(tsrc)
+        Threads.@spawn _add_abelian_block!(tdst, tsrc, p, transformer, f₁, f₂, α, β,
+                                           backend...)
+    end
+    return nothing
+end
+
+function _add_abelian_block!(tdst, tsrc, p, transformer, f₁, f₂, α, β, backend...)
+    (f₁′, f₂′), coeff = first(transformer(f₁, f₂))
+    @inbounds TO.tensoradd!(tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff, β,
+                            backend...)
+    return nothing
+end
+
 # Non-abelian implementations
 # ---------------------------
-function _add_transform_nonthreaded!(tdst, tsrc, p, transformer::GenericTreeTransformer,
-                                     α, β, backend...)
+function _add_general_kernel_nonthreaded!(tdst, tsrc, p,
+                                          transformer::GenericTreeTransformer,
+                                          α, β, backend...)
     # preallocate buffers
     buffers = allocate_buffers(tdst, tsrc, transformer)
 
-    # TODO: this could be multithreaded
     for subtransformer in transformer.data
         # Special case without intermediate buffers whenever there is only a single block
         if length(subtransformer[1]) == 1
@@ -549,9 +587,9 @@ function _add_transform_nonthreaded!(tdst, tsrc, p, transformer::GenericTreeTran
     return nothing
 end
 
-function _add_transform_threaded!(tdst, tsrc, p, transformer::GenericTreeTransformer,
-                                  α, β, backend...;
-                                  ntasks::Int=get_num_transformer_threads())
+function _add_general_kernel_threaded!(tdst, tsrc, p, transformer::GenericTreeTransformer,
+                                       α, β, backend...;
+                                       ntasks::Int=get_num_transformer_threads())
     nblocks = length(transformer.data)
 
     counter = Threads.Atomic{Int}(1)
@@ -577,18 +615,18 @@ function _add_transform_threaded!(tdst, tsrc, p, transformer::GenericTreeTransfo
     return nothing
 end
 
-# Kernels
-# -------
-function _add_transform_single!(tdst, tsrc, p,
-                                (basistransform, structures_dst, structures_src),
-                                α, β, 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...)
+function _add_general_kernel_nonthreaded!(tdst, tsrc, p, transformer, α, β, backend...)
+    if iszero(β)
+        tdst = zerovector!(tdst)
+    elseif !isone(β)
+        tdst = scale!(tdst, β)
+    end
+    for (f₁, f₂) in fusiontrees(tsrc)
+        for ((f₁′, f₂′), coeff) in transformer(f₁, f₂)
+            @inbounds TO.tensoradd!(tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff,
+                                    One(), backend...)
+        end
+    end
     return nothing
 end
 
@@ -620,88 +658,25 @@ function _add_transform_multi!(tdst, tsrc, p,
     return nothing
 end
 
-# Other implementations
-# ---------------------
-
-function add_transform_kernel!(tdst::AbstractTensorMap,
-                               tsrc::AbstractTensorMap,
-                               (p₁, p₂)::Index2Tuple,
-                               fusiontreetransform::Function,
-                               α::Number,
-                               β::Number,
-                               backend::AbstractBackend...)
-    I = sectortype(spacetype(tdst))
-
-    if I === Trivial
-        _add_trivial_kernel!(tdst, tsrc, (p₁, p₂), fusiontreetransform, α, β, backend...)
-    elseif FusionStyle(I) isa UniqueFusion
-        _add_abelian_kernel!(tdst, tsrc, (p₁, p₂), fusiontreetransform, α, β, backend...)
-    else
-        _add_general_kernel!(tdst, tsrc, (p₁, p₂), fusiontreetransform, α, β, backend...)
-    end
-
-    return nothing
-end
-
-# internal methods: no argument types
-function _add_trivial_kernel!(tdst, tsrc, p, fusiontreetransform, α, β, backend...)
-    TO.tensoradd!(tdst[], tsrc[], p, false, α, β, backend...)
-    return nothing
-end
-
-function _add_abelian_kernel!(tdst, tsrc, p, fusiontreetransform, α, β, backend...)
-    if Threads.nthreads() > 1
-        Threads.@sync for (f₁, f₂) in fusiontrees(tsrc)
-            Threads.@spawn _add_abelian_block!(tdst, tsrc, p, fusiontreetransform,
-                                               f₁, f₂, α, β, backend...)
-        end
-    else
-        for (f₁, f₂) in fusiontrees(tsrc)
-            _add_abelian_block!(tdst, tsrc, p, fusiontreetransform,
-                                f₁, f₂, α, β, backend...)
-        end
-    end
-    return nothing
-end
-
-function _add_abelian_block!(tdst, tsrc, p, fusiontreetransform, f₁, f₂, α, β, backend...)
-    (f₁′, f₂′), coeff = first(fusiontreetransform(f₁, f₂))
-    @inbounds TO.tensoradd!(tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff, β,
-                            backend...)
-    return nothing
-end
-
-function _add_general_kernel!(tdst, tsrc, p, fusiontreetransform, α, β, backend...)
+function _add_general_kernel_threaded!(tdst, tsrc, p, transformer, α, β, backend...)
     if iszero(β)
         tdst = zerovector!(tdst)
-    elseif β != 1
+    elseif !isone(β)
         tdst = scale!(tdst, β)
     end
-    β′ = One()
-    if Threads.nthreads() > 1
-        Threads.@sync for s₁ in sectors(codomain(tsrc)), s₂ in sectors(domain(tsrc))
-            Threads.@spawn _add_nonabelian_sector!(tdst, tsrc, p, fusiontreetransform, s₁,
-                                                   s₂, α, β′, backend...)
-        end
-    else
-        for (f₁, f₂) in fusiontrees(tsrc)
-            for ((f₁′, f₂′), coeff) in fusiontreetransform(f₁, f₂)
-                @inbounds TO.tensoradd!(tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff,
-                                        β′, backend...)
-            end
-        end
+    Threads.@sync for s₁ in sectors(codomain(tsrc)), s₂ in sectors(domain(tsrc))
+        Threads.@spawn _add_nonabelian_sector!(tdst, tsrc, p, transformer, s₁, s₂, α,
+                                               backend...)
     end
     return nothing
 end
 
-# TODO: β argument is weird here because it has to be 1
-function _add_nonabelian_sector!(tdst, tsrc, p, fusiontreetransform, s₁, s₂, α, β,
-                                 backend...)
+function _add_nonabelian_sector!(tdst, tsrc, p, fusiontreetransform, s₁, s₂, α, backend...)
     for (f₁, f₂) in fusiontrees(tsrc)
         (f₁.uncoupled == s₁ && f₂.uncoupled == s₂) || continue
         for ((f₁′, f₂′), coeff) in fusiontreetransform(f₁, f₂)
-            @inbounds TO.tensoradd!(tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff, β,
-                                    backend...)
+            @inbounds TO.tensoradd!(tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff,
+                                    One(), backend...)
         end
     end
     return nothing