Add scheduler support in mul!

Author: lkdvos

src/tensors/blockiterator.jl

@@ -13,3 +13,5 @@ Base.IteratorEltype(::BlockIterator) = Base.HasEltype()
 Base.eltype(::Type{<:BlockIterator{T}}) where {T} = blocktype(T)
 Base.length(iter::BlockIterator) = length(iter.structure)
 Base.isdone(iter::BlockIterator, state...) = Base.isdone(iter.structure, state...)
+
+Base.haskey(iter::BlockIterator, c) = haskey(iter.structure, c)

src/tensors/linalg.jl

@@ -309,12 +309,47 @@ function LinearAlgebra.tr(t::AbstractTensorMap)
 end
 
 # TensorMap multiplication
-function LinearAlgebra.mul!(tC::AbstractTensorMap,
-                            tA::AbstractTensorMap,
-                            tB::AbstractTensorMap, α=true, β=false)
+function LinearAlgebra.mul!(tC::AbstractTensorMap, tA::AbstractTensorMap,
+                            tB::AbstractTensorMap,
+                            α::Number, β::Number,
+                            backend::AbstractBackend=TO.DefaultBackend())
+    if backend isa TO.DefaultBackend
+        newbackend = TO.select_backend(mul!, tC, tA, tB)
+        return mul!(tC, tA, tB, α, β, newbackend)
+    elseif backend isa TO.NoBackend # error for missing backend
+        TC = typeof(tC)
+        TA = typeof(tA)
+        TB = typeof(tB)
+        throw(ArgumentError("No suitable backend found for `mul!` and tensor types $TC, $TA and $TB"))
+    else # error for unknown backend
+        TC = typeof(tC)
+        TA = typeof(tA)
+        TB = typeof(tB)
+        throw(ArgumentError("Unknown backend for `mul!` and tensor types $TC, $TA and $TB"))
+    end
+end
+
+function TO.select_backend(::typeof(mul!), C::AbstractTensorMap, A::AbstractTensorMap,
+                           B::AbstractTensorMap)
+    return TensorKitBackend()
+end
+
+function LinearAlgebra.mul!(tC::AbstractTensorMap, tA::AbstractTensorMap,
+                            tB::AbstractTensorMap, α::Number, β::Number,
+                            backend::TensorKitBackend)
     compose(space(tA), space(tB)) == space(tC) ||
         throw(SpaceMismatch(lazy"$(space(tC)) ≠ $(space(tA)) * $(space(tB))"))
 
+    scheduler = backend.blockscheduler
+    if isnothing(scheduler)
+        return sequential_mul!(tC, tA, tB, α, β)
+    else
+        return threaded_mul!(tC, tA, tB, α, β, scheduler)
+    end
+end
+
+function sequential_mul!(tC::AbstractTensorMap, tA::AbstractTensorMap,
+                         tB::AbstractTensorMap, α::Number, β::Number)
     iterC = blocks(tC)
     iterA = blocks(tA)
     iterB = blocks(tB)
@@ -336,13 +371,13 @@ function LinearAlgebra.mul!(tC::AbstractTensorMap,
             elseif cB < cC
                 nextB = iterate(iterB, stateB)
             else
-                if β != one(β)
+                if !isone(β)
                     rmul!(C, β)
                 end
                 nextC = iterate(iterC, stateC)
             end
         else
-            if β != one(β)
+            if !isone(β)
                 rmul!(C, β)
             end
             nextC = iterate(iterC, stateC)
@@ -351,7 +386,21 @@ function LinearAlgebra.mul!(tC::AbstractTensorMap,
     return tC
 end
 
-# TODO: consider spawning threads for different blocks, support backends
+function threaded_mul!(tC::AbstractTensorMap, tA::AbstractTensorMap, tB::AbstractTensorMap,
+                       α::Number, β::Number, scheduler::Scheduler)
+    # obtain cached data before multithreading
+    bCs, bAs, bBs = blocks(tC), blocks(tA), blocks(tB)
+
+    tforeach(blocksectors(tC); scheduler) do c
+        if haskey(bAs, c) # then also bBs should have it
+            mul!(bCs[c], bAs[c], bBs[c], α, β)
+        elseif !isone(β)
+            scale!(bCs[c], β)
+        end
+    end
+
+    return tC
+end
 
 # TensorMap inverse
 function Base.inv(t::AbstractTensorMap)

src/tensors/tensoroperations.jl

@@ -149,9 +149,10 @@ TO.tensorcost(t::AbstractTensorMap, i::Int) = dim(space(t, i))
 # 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,S<:Scheduler} <: AbstractBackend
+@kwdef struct TensorKitBackend{B<:AbstractBackend,BS,SBS} <: AbstractBackend
     arraybackend::B = TO.DefaultBackend()
-    scheduler::S = SerialScheduler()
+    blockscheduler::BS = SerialScheduler()
+    subblockscheduler::SBS = SerialScheduler()
 end
 
 function TO.select_backend(::typeof(TO.tensoradd!), C::AbstractTensorMap,