rework default scheduler settings

Author: lkdvos

Project.toml

@@ -9,6 +9,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f"
@@ -35,6 +36,7 @@ LinearAlgebra = "1"
 OhMyThreads = "0.7.0"
 PackageExtensionCompat = "1"
 Random = "1"
+ScopedValues = "1.3.0"
 SparseArrays = "1"
 Strided = "2"
 TensorKitSectors = "0.1"

src/TensorKit.jl

@@ -102,6 +102,7 @@ const TO = TensorOperations
 
 using LRUCache
 using OhMyThreads
+using ScopedValues
 
 using TensorKitSectors
 import TensorKitSectors: dim, BraidingStyle, FusionStyle, ⊠, ⊗
@@ -185,6 +186,7 @@ include("spaces/vectorspaces.jl")
 #-------------------------------------
 # general definitions
 include("tensors/abstracttensor.jl")
+include("tensors/backends.jl")
 include("tensors/blockiterator.jl")
 include("tensors/tensor.jl")
 include("tensors/adjoint.jl")

src/tensors/backends.jl

@@ -0,0 +1,27 @@
+# Scheduler implementation
+# ------------------------
+const blockscheduler = ScopedValue{Scheduler}(SerialScheduler())
+const subblockscheduler = ScopedValue{Scheduler}(SerialScheduler())
+
+# Backend implementation
+# ----------------------
+# TODO: figure out a name
+# TODO: what should be the default scheduler?
+@kwdef struct TensorKitBackend{B<:AbstractBackend,BS,SBS} <: AbstractBackend
+    arraybackend::B = TO.DefaultBackend()
+    blockscheduler::BS = blockscheduler[]
+    subblockscheduler::SBS = subblockscheduler[]
+end
+
+function TO.select_backend(::typeof(TO.tensoradd!), C::AbstractTensorMap,
+                           A::AbstractTensorMap)
+    return TensorKitBackend()
+end
+function TO.select_backend(::typeof(TO.tensortrace!), C::AbstractTensorMap,
+                           A::AbstractTensorMap)
+    return TensorKitBackend()
+end
+function TO.select_backend(::typeof(TO.tensorcontract!), C::AbstractTensorMap,
+                           A::AbstractTensorMap, B::AbstractTensorMap)
+    return TensorKitBackend()
+end

src/tensors/indexmanipulations.jl

@@ -519,7 +519,7 @@ function add_transform_kernel!(tdst::TensorMap,
     structure_src = transformer.structure_src.fusiontreestructure
 
     tforeach(transformer.rows, transformer.cols, transformer.vals;
-             backend.scheduler) do row, col, val
+             scheduler=backend.subblockscheduler) do row, col, val
         sz_dst, str_dst, offset_dst = structure_dst[col]
         subblock_dst = StridedView(tdst.data, sz_dst, str_dst, offset_dst)
 
@@ -546,7 +546,7 @@ function add_transform_kernel!(tdst::TensorMap,
     rows = rowvals(transformer.matrix)
     vals = nonzeros(transformer.matrix)
 
-    tforeach(axes(transformer.matrix, 2); backend.scheduler) do j
+    tforeach(axes(transformer.matrix, 2); scheduler=backend.subblockscheduler) do j
         sz_dst, str_dst, offset_dst = structure_dst[j]
         subblock_dst = StridedView(tdst.data, sz_dst, str_dst, offset_dst)
         nzrows = nzrange(transformer.matrix, j)
@@ -602,7 +602,7 @@ end
 
 function _add_abelian_kernel!(tdst, tsrc, p, fusiontreetransform, α, β,
                               backend::TensorKitBackend, allocator)
-    tforeach(fusiontrees(tsrc); backend.scheduler) do (f₁, f₂)
+    tforeach(fusiontrees(tsrc); scheduler=backend.subblockscheduler) do (f₁, f₂)
         return _add_abelian_block!(tdst, tsrc, p, fusiontreetransform,
                                    f₁, f₂, α, β, backend.arraybackend, allocator)
     end
@@ -624,16 +624,17 @@ function _add_general_kernel!(tdst, tsrc, p, fusiontreetransform, α, β, backen
         tdst = scale!(tdst, β)
     end
     β′ = One()
-    if backend.scheduler isa SerialScheduler
+    if backend.subblockscheduler isa SerialScheduler
         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.arraybackend, allocator)
             end
         end
     else
-        tforeach(Iterators.product(sectors(codomain(tsrc)), sectors(domain(tsrc)))) do (s₁,
-                                                                                        s₂)
+        tforeach(Iterators.product(sectors(codomain(tsrc)), sectors(domain(tsrc)));
+                 scheduler=backend.subblockscheduler) do (s₁,
+                                                          s₂)
             return _add_nonabelian_sector!(tdts, tsrc, p, fusiontreetransform, s₁, s₂, α,
                                            β′, backend.arraybackend, allocator)
         end

src/tensors/tensoroperations.jl

@@ -144,30 +144,6 @@ TO.tensorcost(t::AbstractTensorMap, i::Int) = dim(space(t, i))
 # IMPLEMENTATONS
 #----------------
 
-# Backend implementation
-# ----------------------
-# TODO: figure out a name
-# TODO: what should be the default scheduler?
-# TODO: should we allow a separate scheduler for "blocks" and "subblocks"
-@kwdef struct TensorKitBackend{B<:AbstractBackend,BS,SBS} <: AbstractBackend
-    arraybackend::B = TO.DefaultBackend()
-    blockscheduler::BS = SerialScheduler()
-    subblockscheduler::SBS = SerialScheduler()
-end
-
-function TO.select_backend(::typeof(TO.tensoradd!), C::AbstractTensorMap,
-                           A::AbstractTensorMap)
-    return TensorKitBackend()
-end
-function TO.select_backend(::typeof(TO.tensortrace!), C::AbstractTensorMap,
-                           A::AbstractTensorMap)
-    return TensorKitBackend()
-end
-function TO.select_backend(::typeof(TO.tensorcontract!), C::AbstractTensorMap,
-                           A::AbstractTensorMap, B::AbstractTensorMap)
-    return TensorKitBackend()
-end
-
 # Trace implementation
 #----------------------
 """

test/runtests.jl

@@ -54,8 +54,8 @@ sectorlist = (Z2Irrep, Z3Irrep, Z4Irrep, Z3Irrep ⊠ Z4Irrep,
               Z2Irrep ⊠ FibonacciAnyon ⊠ FibonacciAnyon)
 
 Ti = time()
-include("fusiontrees.jl")
-include("spaces.jl")
+# include("fusiontrees.jl")
+# include("spaces.jl")
 include("tensors.jl")
 include("diagonal.jl")
 include("planar.jl")