Parallel BP

Project.toml

@@ -30,16 +30,19 @@ WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44"
 
 [weakdeps]
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
+Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
 
 [extensions]
 ITensorNetworksNextTensorOperationsExt = "TensorOperations"
+ITensorNetworksNextDaggerExt = "Dagger"
 
 [compat]
 AbstractTrees = "0.4.5"
 Adapt = "4.3"
 AlgorithmsInterface = "0.1"
 BackendSelection = "0.1.6"
 Combinatorics = "1"
+Dagger = "0.19.0"
 DataGraphs = "0.4"
 DiagonalArrays = "0.3.31"
 Dictionaries = "0.4.5"

ext/ITensorNetworksNextDaggerExt/ITensorNetworksNextDaggerExt.jl

@@ -0,0 +1,84 @@
+module ITensorNetworksNextDaggerExt
+
+import AlgorithmsInterface as AI
+import ITensorNetworksNext.AlgorithmsInterfaceExtensions as AIE
+import ITensorNetworksNext.ITensorNetworksNextParallel as ITNNP
+using Dagger
+using Dictionaries: Dictionary, set!
+
+@kwdef mutable struct DaggerState{Chunk <: Dagger.Chunk} <: AI.State
+    chunk::Chunk # A chunk living on the host worker.
+    futures = Dictionary{Int, Dagger.DTask}() # the futures from the remote steps.
+end
+
+function Base.getproperty(state::DaggerState, name::Symbol)
+    if name in (:chunk, :futures)
+        return getfield(state, name)
+    end
+    return getproperty(fetch(state.chunk), name)
+end
+function Base.setproperty!(state::DaggerState, name::Symbol, val)
+    if name === (:chunk, :futures)
+        return setfield!(state, name, val)
+    end
+    fetch(Dagger.@spawn setproperty!(state.chunk, name, val))
+    return state
+end
+
+# ====================================== overloads ======================================= #
+
+function ITNNP.default_workers(::AI.Algorithm, ::ITNNP.AbstractDaggerStrategy)
+    return Dagger.Distributed.workers()
+end
+
+function ITNNP.initialize_parallel_state(
+        problem::AI.Problem, algorithm::AI.Algorithm,
+        _strategy::ITNNP.GenericDaggerStrategy; kwargs...
+    )
+    chunk = Dagger.@mutable AI.initialize_state(problem, algorithm; iterate, kwargs...)
+
+    return DaggerState(; chunk)
+end
+
+function AI.step!(problem::AI.Problem, algorithm::ITNNP.Parallelized, state::DaggerState)
+    worker_list = algorithm.workers
+    algorithm = algorithm.parent
+
+    function remote_solve(state)
+        subsolve = AIE.initialize_subsolve(problem, algorithm, state)
+
+        subproblem, subalgorithm, substate = subsolve
+        AI.solve!(subproblem, subalgorithm, substate)
+
+        return subsolve
+    end
+
+    worker = worker_list[mod(state.iteration, 1:length(worker_list))]
+
+    dtask = Dagger.@spawn scope = Dagger.scope(; worker) remote_solve(fetch(state.chunk))
+
+    # Spawns on the host only (as we do not fetch the chunk before hand.)
+    dtask = Dagger.spawn(dtask, state.chunk) do subsolve, state
+        subproblem, subalgorithm, substate = subsolve
+
+        AIE.finalize_substate!(subproblem, subalgorithm, substate, state)
+        return state
+    end
+
+    set!(state.futures, state.iteration, dtask)
+
+    return state
+end
+function AI.finalize_state!(::AI.Problem, ::AI.Algorithm, state::DaggerState)
+    foreach(fetch, state.futures)
+    return state.iterate
+end
+
+function AIE.finalize_substate!(
+        problem::AI.Problem, algorithm::AI.Algorithm, substate::DaggerState, state::AI.State
+    )
+    AIE.finalize_substate!(problem, algorithm, fetch(substate.chunk), state)
+    return state
+end
+
+end # ITensorNetworksNextDaggerExt

src/AlgorithmsInterfaceExtensions/AlgorithmsInterfaceExtensions.jl

@@ -2,6 +2,37 @@ module AlgorithmsInterfaceExtensions
 
 import AlgorithmsInterface as AI
 
+@kwdef mutable struct GenericAlgorithmState{Iterate, StoppingCriterionState} <: AI.State
+    iterate::Iterate
+    iteration::Int = 0
+    stopping_criterion_state::StoppingCriterionState
+end
+
+function AI.initialize_state(
+        problem::AI.Problem,
+        algorithm::AI.Algorithm;
+        iterate,
+        iteration = 0
+    )
+    stopping_criterion_state = AI.initialize_state(
+        problem, algorithm, algorithm.stopping_criterion; iterate
+    )
+    return GenericAlgorithmState(; iterate, iteration, stopping_criterion_state)
+end
+
+function AI.initialize_state!(
+        problem::AI.Problem,
+        algorithm::AI.Algorithm,
+        state::AI.State;
+        iteration = 0
+    )
+    state.iteration = iteration
+    AI.initialize_state!(
+        problem, algorithm, algorithm.stopping_criterion, state.stopping_criterion_state
+    )
+    return state
+end
+
 # ============================ NestedAlgorithm =============================================
 
 abstract type NestedAlgorithm <: AI.Algorithm end
@@ -18,7 +49,7 @@ function initialize_subsolve(
 end
 
 function finalize_substate!(
-        problem::AI.Problem, algorithm::AI.Algorithm, state::AI.State, substate::AI.State
+        ::AI.Problem, ::AI.Algorithm, substate::AI.State, state::AI.State
     )
     state.iterate = substate.iterate
     return state
@@ -27,7 +58,7 @@ end
 function AI.step!(problem::AI.Problem, algorithm::NestedAlgorithm, state::AI.State)
     subproblem, subalgorithm, substate = initialize_subsolve(problem, algorithm, state)
     AI.solve!(subproblem, subalgorithm, substate)
-    finalize_substate!(problem, algorithm, state, substate)
+    finalize_substate!(subproblem, subalgorithm, substate, state)
     return state
 end
 
@@ -38,6 +69,12 @@ end
 # state. Subtypes must store the inner state as a field named `substate`.
 abstract type NestedState <: AI.State end
 
+@kwdef mutable struct GenericNestedState{Substate, StoppingCriterionState} <: NestedState
+    substate::Substate
+    iteration::Int = 0
+    stopping_criterion_state::StoppingCriterionState
+end
+
 # Use `getfield` on the right-hand side so future edits to this forwarder
 # can't accidentally recurse through the overload.
 function Base.getproperty(state::NestedState, name::Symbol)

src/ITensorNetworksNext.jl

@@ -15,5 +15,9 @@ include("contract_network.jl")
 
 include("beliefpropagation/messagecache.jl")
 include("beliefpropagation/beliefpropagation.jl")
+include("beliefpropagation/partitioned.jl")
+
+# lib
+include("ITensorNetworksNextParallel/ITensorNetworksNextParallel.jl")
 
 end

src/ITensorNetworksNextParallel/ITensorNetworksNextParallel.jl

@@ -0,0 +1,60 @@
+module ITensorNetworksNextParallel
+
+import ..ITensorNetworksNext.AlgorithmsInterfaceExtensions as AIE
+import AlgorithmsInterface as AI
+
+abstract type AbstractParallelizationStrategy end
+
+function default_workers end
+function initialize_parallel_state end
+
+@kwdef struct Parallelized{Strategy, Workers, Algorithm <: AI.Algorithm} <: AI.Algorithm
+    parent::Algorithm
+    strategy::Strategy
+    workers::Workers = default_workers(parent, strategy)
+end
+
+function Base.getproperty(algorithm::Parallelized, name::Symbol)
+    if name in (:parent, :strategy, :workers)
+        return getfield(algorithm, name)
+    end
+    return getproperty(getfield(algorithm, :parent), name)
+end
+
+function AI.initialize_state(problem::AI.Problem, algorithm::Parallelized; kwargs...)
+    return initialize_parallel_state(
+        problem,
+        algorithm.parent,
+        algorithm.strategy;
+        kwargs...
+    )
+end
+
+# ====================================== Dagger.jl ======================================= #
+
+abstract type AbstractDaggerStrategy <: AbstractParallelizationStrategy end
+struct GenericDaggerStrategy <: AbstractDaggerStrategy end
+
+function initialize_parallel_state(
+        _problem,
+        _algorithm,
+        strategy::AbstractDaggerStrategy;
+        _kwargs...
+    )
+    throw(
+        ArgumentError(
+            "package Dagger.jl not loaded; please install and load Dagger.jl to use \
+            strategy of type $(typeof(strategy))."
+        )
+    )
+end
+
+function default_workers(algorithm, strategy::AbstractDaggerStrategy)
+    @warn(
+        "package Dagger.jl may not be loaded; please install and load Dagger.jl to use \
+        strategy of type `$(typeof(strategy))`"
+    )
+    throw(MethodError(default_workers, (algorithm, strategy)))
+end
+
+end # ITensorNetworksNextParallel

src/LazyNamedDimsArrays/symbolicnameddimsarray.jl

@@ -5,6 +5,9 @@ const SymbolicNamedDimsArray{T, N, Parent <: SymbolicArray{T, N}, DimNames} =
 function symnameddims(symname, dims)
     return lazy(nameddims(SymbolicArray(symname, denamed.(dims)), name.(dims)))
 end
+function symnameddims(name, ndarray::AbstractNamedDimsArray)
+    return symnameddims(name, Tuple(inds(ndarray)))
+end
 symnameddims(name) = symnameddims(name, ())
 using AbstractTrees: AbstractTrees
 function AbstractTrees.printnode(io::IO, a::SymbolicNamedDimsArray)

src/abstracttensornetwork.jl

@@ -34,6 +34,7 @@ Base.copy(::AbstractTensorNetwork) = not_implemented()
 
 # Iteration
 Base.iterate(tn::AbstractTensorNetwork, args...) = iterate(vertex_data(tn), args...)
+
 Base.keys(tn::AbstractTensorNetwork) = vertices(tn)
 
 # TODO: This contrasts with the `DataGraphs.AbstractDataGraph` definition,