Simplify parallel code.

Author: jack-dunham

ext/ITensorNetworksNextDaggerExt/ITensorNetworksNextDaggerExt.jl

@@ -6,17 +6,24 @@ import ITensorNetworksNext.ITensorNetworksNextParallel as ITNNP
 using Dagger
 using ITensorNetworksNext.ITensorNetworksNextParallel:
     DaggerNestedAlgorithm, DaggerState, ITensorNetworksNextParallel
+using Dictionaries: set!
 
-function ITNNP.DaggerNestedAlgorithm(f::Function, iterable; workers = workers(), kwargs...)
-    return DaggerNestedAlgorithm(; algorithms = map(f, iterable), workers, kwargs...)
+function ITNNP.DaggerNestedAlgorithm(f, iterable; kwargs...)
+    return DaggerNestedAlgorithm(; algorithms = map(f, iterable), kwargs...)
 end
 
-function initialize_dagger_state(problem::AIE.Problem, algorithm::AIE.Algorithm; iterate)
+function ITNNP.dagger_algorithm(f::Base.Callable, iterable; kwargs...)
+    return DaggerNestedAlgorithm(f, iterable; kwargs...)
+end
+
+function ITNNP.initialize_dagger_state(
+        problem::AIE.Problem, algorithm::AIE.Algorithm; iterate
+    )
     stopping_criterion_state = AI.initialize_state(
         problem, algorithm, algorithm.stopping_criterion
     )
 
-    remote_results = Dict{Int, Dagger.DTask}()
+    remote_results = Dictionary{Int, Dagger.DTask}()
 
     return ITNNP.DaggerState(;
         iterate,
@@ -30,37 +37,23 @@ function AI.initialize_state(
         algorithm::ITNNP.DaggerNestedAlgorithm;
         kwargs...
     )
-    return initialize_dagger_state(problem, algorithm; kwargs...)
+    return ITNNP.initialize_dagger_state(problem, algorithm; kwargs...)
 end
 
-function AIE.get_subproblem(
+function AI.step!(
         problem::AIE.Problem,
         algorithm::ITNNP.DaggerNestedAlgorithm,
-        state::ITNNP.DaggerState
+        state::ITNNP.DaggerState;
+        kwargs...
     )
     subproblem = problem
     subalgorithm = algorithm.algorithms[state.iteration]
 
-    # This might be a Dagger.chun object.
     iterate = ITNNP.get_subiterate(subproblem, subalgorithm, state)
 
-    substate = Dagger.@spawn AI.initialize_state(subproblem, subalgorithm; iterate)
-
-    return subproblem, subalgorithm, substate
-end
-
-function AI.step!(
-        problem::AI.Problem,
-        algorithm::ITNNP.DaggerNestedAlgorithm,
-        state::ITNNP.DaggerState;
-        kwargs...
-    )
-    subproblem, subalgorithm, substate_future =
-        AIE.get_subproblem(problem, algorithm, state)
-
-    dtask = Dagger.@spawn AI.solve(subproblem, subalgorithm, substate_future)
+    dtask = Dagger.@spawn AI.solve(subproblem, subalgorithm; iterate)
 
-    state.remote_results[state.iteration] = dtask
+    set!(state.remote_results, state.iteration, dtask)
 
     return state
 end

ext/ITensorNetworksNextDaggerExt/daggerbeliefpropagation.jl

@@ -1,33 +1,49 @@
 import ITensorNetworksNext.ITensorNetworksNextParallel as ITNNP
 using Dagger
-using DataGraphs: DataGraphs, get_edge_data, get_vertex_data, is_edge_assigned,
+using DataGraphs: DataGraphs, edge_data, get_edge_data, get_vertex_data, is_edge_assigned,
     is_vertex_assigned, set_edge_data!, set_vertex_data!, underlying_graph
-using Dictionaries: Indices
+using Dictionaries: Dictionary, Indices, getindices
 using Graphs: AbstractEdge, AbstractGraph, dst, edges, src, vertices
 using ITensorNetworksNext: ITensorNetworksNext, BeliefPropagation, BeliefPropagationCache,
     BeliefPropagationProblem, BeliefPropagationState, beliefpropagation,
-    forest_cover_edge_sequence, select_algorithm
+    forest_cover_edge_sequence, select_algorithm, subcache
 using NamedGraphs.GraphsExtensions: boundary_edges
 using NamedGraphs.PartitionedGraphs: QuotientVertex, quotientedges, quotientvertices
 using NamedGraphs: NamedGraphs
 
-function ITNNP.DaggerBeliefPropagationCache(network::AbstractGraph)
+const DaggerBeliefPropagation = BeliefPropagation{<:ITNNP.DaggerNestedAlgorithm};
+
+function ITNNP.DaggerBeliefPropagationCache(
+        network::AbstractGraph;
+        workers = nothing,
+        scopes = nothing
+    )
     underlying_cache = BeliefPropagationCache(network)
 
     keys = Indices(quotientvertices(underlying_cache))
 
-    workers = Iterators.cycle(Dagger.Distributed.workers())
-    worker_dict = similar(keys, Int)
+    if isnothing(scopes)
+        workers = isnothing(workers) ? Dagger.Distributed.workers() : workers
 
-    for quotient_vertex in keys
-        worker, workers = Iterators.peel(workers)
-        worker_dict[quotient_vertex] = worker
+        sorted_workers = Iterators.take(Iterators.cycle(workers), length(keys))
+
+        scopes = map(Dagger.ProcessScope, collect(sorted_workers))
+    else
+        if length(keys) != length(scopes)
+            throw(
+                ArgumentError(
+                    "Number of provided scopes must match the number of vertex partitions of underlying graph"
+                )
+            )
+        end
     end
 
+    scope_dict = Dictionary(keys, scopes)
+
     quotient_chunks = map(keys) do quotient_vertex
-        worker = worker_dict[quotient_vertex]
+        scope = scope_dict[quotient_vertex]
         iterate = subcache(underlying_cache, quotient_vertex)
-        chunk = Dagger.@mutable worker = worker BeliefPropagationState(; iterate)
+        chunk = Dagger.@mutable scope = scope BeliefPropagationState(; iterate)
         return chunk
     end
 
@@ -42,7 +58,7 @@ function DataGraphs.is_vertex_assigned(bpc::ITNNP.DaggerBeliefPropagationCache,
     return is_vertex_assigned(bpc.underlying_cache, vertex)
 end
 function DataGraphs.is_edge_assigned(bpc::ITNNP.DaggerBeliefPropagationCache, edge)
-    return is_edge_assigned(bpc.undelying_cache, edge)
+    return is_edge_assigned(bpc.underlying_cache, edge)
 end
 
 function DataGraphs.get_vertex_data(bpc::ITNNP.DaggerBeliefPropagationCache, vertex)
@@ -52,7 +68,7 @@ function DataGraphs.get_edge_data(
         bpc::ITNNP.DaggerBeliefPropagationCache,
         edge::AbstractEdge
     )
-    return get_edge_data(bpc.undelying_caches, edge)
+    return get_edge_data(bpc.underlying_cache, edge)
 end
 
 function DataGraphs.set_vertex_data!(bpc::ITNNP.DaggerBeliefPropagationCache, val, vertex)
@@ -73,13 +89,11 @@ end
 function ITensorNetworksNext.beliefpropagation_sweep(
         cache::ITNNP.DaggerBeliefPropagationCache;
         edges,
-        workers = workers(),
         kwargs...
     )
-    keys = collect(quotientvertices(cache))
-
-    return ITNNP.dagger_algorithm(keys; keys, workers) do quotient_vertex
-        subcache = fetch(cache[quotient_vertex]).iterate
+    return ITNNP.dagger_algorithm(quotientvertices(cache)) do quotient_vertex
+        substate = fetch(cache[quotient_vertex])
+        subcache = substate.iterate
 
         subcache_edges = forest_cover_edge_sequence(subcache) ∩ edges
         incoming_edges = boundary_edges(cache, vertices(cache, quotient_vertex); dir = :in)
@@ -120,44 +134,61 @@ function ITNNP.get_subiterate(
 end
 
 function AIE.set_substate!(
-        ::BeliefPropagationProblem,
-        ::AIE.NestedAlgorithm,
+        problem::BeliefPropagationProblem,
+        algorithm::AIE.NestedAlgorithm,
         state::AIE.State,
         substate::ITNNP.DaggerState
     )
     dst_cache = state.iterate.iterate
 
     state.iterate.maxdiff = 0.0
 
-    for remote_result in substate.remote_results
-        get_maxdiff = dtask -> dtask.iterate.maxdiff
-        src_maxdiff = fetch(Dagger.@spawn get_maxdiff(remote_result))
-
-        if src_maxdiff > state.iterate.maxdiff
-            state.iterate.maxdiff = src_maxdiff
-        end
+    maxdiff_dtasks = map(substate.remote_results) do remote_result
+        return Dagger.spawn(dtask -> dtask.iterate.maxdiff, remote_result)
     end
 
-    function transfer_edges!(dst_chunk, src_chunk, edges)
-        src_subcache = src_chunk.iterate
-        dst_subcache = dst_chunk.iterate
-        for edge in edges
-            dst_subcache[edge] = src_subcache[edge]
-        end
-        return
+    maxdiff = maximum(fetch, maxdiff_dtasks)
+
+    if maxdiff > state.iterate.maxdiff
+        state.iterate.maxdiff = maxdiff
     end
 
     transfer_dtasks = map(quotientedges(dst_cache)) do quotient_edge
         src_subcache = dst_cache[src(quotient_edge)]
         dst_subcache = dst_cache[dst(quotient_edge)]
-        return Dagger.@spawn transfer_edges!(
+
+        src_subcache = fetch(src_subcache)
+
+        return Dagger.spawn(
             dst_subcache,
             fetch(src_subcache),
             edges(dst_cache, quotient_edge)
-        )
+        ) do dst, src, edges
+            src_subcache = src.iterate
+            dst_subcache = dst.iterate
+            for edge in edges
+                dst_subcache[edge] = src_subcache[edge]
+            end
+        end
     end
 
-    wait.(transfer_dtasks)
+    foreach(wait, transfer_dtasks)
+
+    return state
+end
+
+function ITNNP.finalize_state!(
+        ::BeliefPropagationProblem,
+        ::BeliefPropagation,
+        state::ITNNP.DaggerState
+    )
+    dst_cache = state.iterate.iterate
+
+    for quotient_vertex in quotientvertices(dst_cache)
+        substate = fetch(dst_cache[quotient_vertex])
+        subcache = substate.iterate
+        edge_data(dst_cache) .= edge_data(subcache)
+    end
 
     return state
 end

src/ITensorNetworksNextParallel/ITensorNetworksNextParallel.jl

@@ -1,6 +1,10 @@
 module ITensorNetworksNextParallel
 
 import AlgorithmsInterface as AI
+import ITensorNetworksNext.AlgorithmsInterfaceExtensions as AIE
+
+abstract type ParallelAlgorithm{Child} <: AIE.NestedAlgorithm{Child} end
+const IterativeParallelAlgorithm{Child <: ParallelAlgorithm} = AIE.NestedAlgorithm{Child}
 
 """
     get_subiterate(subproblem::AI.Problem, subalgorithm::AI.Algorithm, state::AI.State)
@@ -11,6 +15,25 @@ The returned value of this function is then pass to a remote call of `initialize
 """
 get_subiterate(::AI.Problem, ::AI.Algorithm, state::AI.State) = state.iterate
 
+finalize_state!(::AI.Problem, ::AI.Algorithm, state::AI.State) = state
+
+function AI.is_finished!(
+        problem::AI.Problem,
+        algorithm::IterativeParallelAlgorithm,
+        state::AI.State
+    )
+    c = algorithm.stopping_criterion
+    st = state.stopping_criterion_state
+
+    isfinished = AI.is_finished!(problem, algorithm, state, c, st)
+
+    if isfinished
+        finalize_state!(problem, algorithm, state)
+    end
+
+    return isfinished
+end
+
 include("dagger.jl")
 
 end # ITensorNetworksNextParallel

src/ITensorNetworksNextParallel/dagger.jl

@@ -1,31 +1,40 @@
 import ..ITensorNetworksNext.AlgorithmsInterfaceExtensions as AIE
 import AlgorithmsInterface as AI
+using Dictionaries: Dictionary
 using ITensorNetworksNext: AbstractBeliefPropagationCache
 
 @kwdef mutable struct DaggerState{
-        Iterate, StoppingCriterionState <: AI.StoppingCriterionState, Chunk, DTask,
+        Iterate, StoppingCriterionState <: AI.StoppingCriterionState, DTask,
     } <: AIE.State
     iterate::Iterate # DaggerBeliefPropagationCache
     iteration::Int = 0
     stopping_criterion_state::StoppingCriterionState
-    # remote_subiterates::Dict{Int, Chunk} = Dict{Int, Any}()
-    remote_results::Dict{Int, DTask} = Dict{Int, Any}()
+    remote_results::Dictionary{Int, DTask} = Dict{Int, Any}()
+end
+
+function initialize_dagger_state(problem, algorithm; kwargs...)
+    throw(
+        ErrorException(
+            "Package Dagger not loaded. Please install and load the Dagger package."
+        )
+    )
 end
 
 @kwdef struct DaggerNestedAlgorithm{
         ChildAlgorithm <: AIE.Algorithm,
         Algorithms <: AbstractVector{ChildAlgorithm},
         StoppingCriterion <: AI.StoppingCriterion,
-        KeyType,
-    } <: AIE.NestedAlgorithm
+    } <: ParallelAlgorithm{ChildAlgorithm}
     algorithms::Algorithms
     stopping_criterion::StoppingCriterion = AI.StopAfterIteration(length(algorithms))
-    workers::Vector{Int}
-    keys::Vector{KeyType} = collect(1:length(algorithms))
 end
 
-function dagger_algorithm(f::Function, iterable; kwargs...)
-    return DaggerNestedAlgorithm(f, iterable; kwargs...)
+function dagger_algorithm(f, iterable; kwargs...)
+    throw(
+        ErrorException(
+            "Package Dagger not loaded. Please install and load the Dagger package."
+        )
+    )
 end
 
 # ================================== belief propagation ================================== #

src/beliefpropagation/beliefpropagationproblem.jl

@@ -30,8 +30,8 @@ function AI.initialize_state!(
 end
 
 function AI.is_finished!(
-        ::AIE.Problem,
-        ::AIE.Algorithm,
+        problem::AIE.Problem,
+        algorithm::AIE.Algorithm,
         state::AIE.State,
         c::StopWhenConverged,
         st::StopWhenConvergedState
@@ -42,6 +42,16 @@ function AI.is_finished!(
         st.delta = state.iterate.maxdiff
     end
 
+    return AI.is_finished(problem, algorithm, state, c, st)
+end
+
+function AI.is_finished(
+        ::AIE.Problem,
+        ::AIE.Algorithm,
+        ::AIE.State,
+        c::StopWhenConverged,
+        st::StopWhenConvergedState
+    )
     return st.delta < c.tol
 end
 
@@ -60,7 +70,7 @@ end
         ChildAlgorithm <: AIE.Algorithm,
         Algorithms <: AbstractVector{ChildAlgorithm},
         StoppingCriterion <: AI.StoppingCriterion,
-    } <: AIE.NestedAlgorithm
+    } <: AIE.NestedAlgorithm{ChildAlgorithm}
     algorithms::Algorithms
     stopping_criterion::StoppingCriterion = AI.StopAfterIteration(length(algorithms))
 end
@@ -100,7 +110,7 @@ end
 struct BeliefPropagationSweep{
         ChildAlgorithm <: AIE.Algorithm,
         Algorithms <: AbstractVector{ChildAlgorithm},
-    } <: AIE.NestedAlgorithm
+    } <: AIE.NestedAlgorithm{ChildAlgorithm}
     algorithms::Algorithms
     stopping_criterion::AI.StopAfterIteration
     function BeliefPropagationSweep(; algorithms)