Add scheduler support

Author: lkdvos

Project.toml

@@ -7,8 +7,10 @@ version = "0.14.5"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
+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"
@@ -33,8 +35,10 @@ FiniteDifferences = "0.12"
 LRUCache = "1.0.2"
 LinearAlgebra = "1"
 MatrixAlgebraKit = "0.1.1"
+OhMyThreads = "0.8.0"
 PackageExtensionCompat = "1"
 Random = "1"
+ScopedValues = "1.3.0"
 SparseArrays = "1"
 Strided = "2"
 TensorKitSectors = "0.1"

src/TensorKit.jl

@@ -103,6 +103,8 @@ const TO = TensorOperations
 using MatrixAlgebraKit: MatrixAlgebraKit as MAK
 
 using LRUCache
+using OhMyThreads
+using ScopedValues
 
 using TensorKitSectors
 import TensorKitSectors: dim, BraidingStyle, FusionStyle, ⊠, ⊗
@@ -186,6 +188,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,28 @@
+# Scheduler implementation
+# ------------------------
+function select_scheduler(scheduler=OhMyThreads.Implementation.NotGiven(); kwargs...)
+    return if scheduler == OhMyThreads.Implementation.NotGiven() && isempty(kwargs)
+        Threads.nthreads() > 1 ? SerialScheduler() : DynamicScheduler()
+    else
+        OhMyThreads.Implementation._scheduler_from_userinput(scheduler; kwargs...)
+    end
+end
+
+"""
+    const blockscheduler = ScopedValue{Scheduler}(SerialScheduler())
+
+The default scheduler used when looping over different blocks in the matrix representation of a
+tensor.
+For controlling this value, see also [`set_blockscheduler`](@ref) and [`with_blockscheduler`](@ref).
+"""
+const blockscheduler = ScopedValue{Scheduler}(SerialScheduler())
+
+"""
+    with_blockscheduler(f, [scheduler]; kwargs...)
+
+Run `f` in a scope where the `blockscheduler` is determined by `scheduler' and `kwargs...`.
+"""
+@inline function with_blockscheduler(f, scheduler=OhMyThreads.Implementation.NotGiven();
+                                     kwargs...)
+    @with blockscheduler => select_scheduler(scheduler; kwargs...) f()
+end