Parallel BP

Project.toml

@@ -29,25 +29,27 @@ 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"
-DataGraphs = "0.2.7"
+DataGraphs = "0.3.1"
 DiagonalArrays = "0.3.31"
 Dictionaries = "0.4.5"
-FunctionImplementations = "0.4"
+FunctionImplementations = "0.4.1"
 Graphs = "1.13.1"
 LinearAlgebra = "1.10"
 MacroTools = "0.5.16"
-NamedDimsArrays = "0.14.2"
-NamedGraphs = "0.6.9, 0.7, 0.8"
+NamedDimsArrays = "0.14.3"
+NamedGraphs = "0.10"
 SimpleTraits = "0.9.5"
 SplitApplyCombine = "1.2.3"
 TensorOperations = "5.3.1"

ext/ITensorNetworksNextDaggerExt/ITensorNetworksNextDaggerExt.jl

@@ -0,0 +1,63 @@
+module ITensorNetworksNextDaggerExt
+
+import AlgorithmsInterface as AI
+import ITensorNetworksNext.AlgorithmsInterfaceExtensions as AIE
+import ITensorNetworksNext.ITensorNetworksNextParallel as ITNNP
+using Dagger
+using ITensorNetworksNext.ITensorNetworksNextParallel:
+    DaggerNestedAlgorithm, DaggerState, ITensorNetworksNextParallel
+using Dictionaries: set!
+
+function ITNNP.DaggerNestedAlgorithm(f, iterable; kwargs...)
+    return DaggerNestedAlgorithm(; algorithms = map(f, iterable), kwargs...)
+end
+
+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 = Dictionary{Int, Dagger.DTask}()
+
+    return ITNNP.DaggerState(;
+        iterate,
+        remote_results,
+        stopping_criterion_state
+    )
+end
+
+function AI.initialize_state(
+        problem::AIE.Problem,
+        algorithm::ITNNP.DaggerNestedAlgorithm;
+        kwargs...
+    )
+    return ITNNP.initialize_dagger_state(problem, algorithm; kwargs...)
+end
+
+function AI.step!(
+        problem::AIE.Problem,
+        algorithm::ITNNP.DaggerNestedAlgorithm,
+        state::ITNNP.DaggerState;
+        kwargs...
+    )
+    subproblem = problem
+    subalgorithm = algorithm.algorithms[state.iteration]
+
+    iterate = ITNNP.get_subiterate(subproblem, subalgorithm, state)
+
+    dtask = Dagger.@spawn AI.solve(subproblem, subalgorithm; iterate)
+
+    set!(state.remote_results, state.iteration, dtask)
+
+    return state
+end
+
+include("daggerbeliefpropagation.jl")
+
+end # ITensorNetworksNextDaggerExt

ext/ITensorNetworksNextDaggerExt/daggerbeliefpropagation.jl

@@ -0,0 +1,194 @@
+import ITensorNetworksNext.ITensorNetworksNextParallel as ITNNP
+using Dagger
+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: 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, subcache
+using NamedGraphs.GraphsExtensions: boundary_edges
+using NamedGraphs.PartitionedGraphs: QuotientVertex, quotientedges, quotientvertices
+using NamedGraphs: NamedGraphs
+
+const DaggerBeliefPropagation = BeliefPropagation{<:ITNNP.DaggerNestedAlgorithm};
+
+function ITNNP.DaggerBeliefPropagationCache(
+        network::AbstractGraph;
+        workers = nothing,
+        scopes = nothing
+    )
+    underlying_cache = BeliefPropagationCache(network)
+
+    keys = Indices(quotientvertices(underlying_cache))
+
+    if isnothing(scopes)
+        workers = isnothing(workers) ? Dagger.Distributed.workers() : workers
+
+        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
+        scope = scope_dict[quotient_vertex]
+        iterate = subcache(underlying_cache, quotient_vertex)
+        chunk = Dagger.@mutable scope = scope BeliefPropagationState(; iterate)
+        return chunk
+    end
+
+    return ITNNP.DaggerBeliefPropagationCache(underlying_cache, quotient_chunks)
+end
+
+function DataGraphs.underlying_graph(cache::ITNNP.DaggerBeliefPropagationCache)
+    return underlying_graph(cache.underlying_cache)
+end
+
+function DataGraphs.is_vertex_assigned(bpc::ITNNP.DaggerBeliefPropagationCache, vertex)
+    return is_vertex_assigned(bpc.underlying_cache, vertex)
+end
+function DataGraphs.is_edge_assigned(bpc::ITNNP.DaggerBeliefPropagationCache, edge)
+    return is_edge_assigned(bpc.underlying_cache, edge)
+end
+
+function DataGraphs.get_vertex_data(bpc::ITNNP.DaggerBeliefPropagationCache, vertex)
+    return get_vertex_data(bpc.underlying_cache, vertex)
+end
+function DataGraphs.get_edge_data(
+        bpc::ITNNP.DaggerBeliefPropagationCache,
+        edge::AbstractEdge
+    )
+    return get_edge_data(bpc.underlying_cache, edge)
+end
+
+function DataGraphs.set_vertex_data!(bpc::ITNNP.DaggerBeliefPropagationCache, val, vertex)
+    return set_vertex_data!(bpc.underlying_cache, val, vertex)
+end
+function DataGraphs.set_edge_data!(bpc::ITNNP.DaggerBeliefPropagationCache, val, edge)
+    return set_edge_data!(bpc.underlying_cache, val, edge)
+end
+
+NamedGraphs.to_graph_index(::ITNNP.DaggerBeliefPropagationCache, qv::QuotientVertex) = qv
+function DataGraphs.get_index_data(
+        cache::ITNNP.DaggerBeliefPropagationCache,
+        qv::QuotientVertex
+    )
+    return cache.quotient_chunks[qv]
+end
+
+function ITensorNetworksNext.beliefpropagation_sweep(
+        cache::ITNNP.DaggerBeliefPropagationCache;
+        edges,
+        kwargs...
+    )
+    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)
+
+        alg = select_algorithm(
+            beliefpropagation,
+            subcache;
+            # Don't update the incoming messages
+            edges = setdiff(subcache_edges, incoming_edges),
+            maxiter = 1,
+            kwargs...
+        )
+
+        return alg
+    end
+end
+
+function AI.initialize_state(
+        problem::AIE.Problem,
+        algorithm::BeliefPropagation{<:ITNNP.DaggerNestedAlgorithm};
+        kwargs...
+    )
+    return ITNNP.initialize_dagger_state(problem, algorithm; kwargs...)
+end
+
+function ITNNP.get_subiterate(
+        ::BeliefPropagationProblem,
+        ::BeliefPropagation, # Our parallel region runs a small BP
+        state::ITNNP.DaggerState
+    )
+    cache = state.iterate.iterate
+
+    quotient_vertex = collect(quotientvertices(cache))[state.iteration]
+
+    subiterate = cache[quotient_vertex]
+
+    return subiterate
+end
+
+function AIE.set_substate!(
+        problem::BeliefPropagationProblem,
+        algorithm::AIE.NestedAlgorithm,
+        state::AIE.State,
+        substate::ITNNP.DaggerState
+    )
+    dst_cache = state.iterate.iterate
+
+    state.iterate.maxdiff = 0.0
+
+    maxdiff_dtasks = map(substate.remote_results) do remote_result
+        return Dagger.spawn(dtask -> dtask.iterate.maxdiff, remote_result)
+    end
+
+    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)]
+
+        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
+
+    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/AlgorithmsInterfaceExtensions/AlgorithmsInterfaceExtensions.jl

@@ -149,10 +149,10 @@ end
 
 # ============================ NestedAlgorithm =============================================
 
-abstract type NestedAlgorithm <: Algorithm end
+abstract type NestedAlgorithm{Child} <: Algorithm end
 
-function nested_algorithm(f::Function, nalgorithms::Int; kwargs...)
-    return DefaultNestedAlgorithm(f, nalgorithms; kwargs...)
+function nested_algorithm(f::Function, iterable; kwargs...)
+    return DefaultNestedAlgorithm(f, iterable; kwargs...)
 end
 
 max_iterations(algorithm::NestedAlgorithm) = length(algorithm.algorithms)
@@ -202,7 +202,7 @@ from a list of stored algorithms.
         ChildAlgorithm <: Algorithm,
         Algorithms <: AbstractVector{ChildAlgorithm},
         StoppingCriterion <: AI.StoppingCriterion,
-    } <: NestedAlgorithm
+    } <: NestedAlgorithm{ChildAlgorithm}
     algorithms::Algorithms
     stopping_criterion::StoppingCriterion = AI.StopAfterIteration(length(algorithms))
 end

src/ITensorNetworksNext.jl

@@ -9,4 +9,10 @@ include("contract_network.jl")
 include("sweeping/utils.jl")
 include("sweeping/eigenproblem.jl")
 
+include("beliefpropagation/abstractbeliefpropagationcache.jl")
+include("beliefpropagation/beliefpropagationcache.jl")
+include("beliefpropagation/beliefpropagationproblem.jl")
+
+include("ITensorNetworksNextParallel/ITensorNetworksNextParallel.jl")
+
 end

src/ITensorNetworksNextParallel/ITensorNetworksNextParallel.jl

@@ -0,0 +1,39 @@
+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)
+
+For a given `subproblem` and `subalgorithm` of a parent nested algorithm,
+derive (from the parent state `state`) the iterate to be used in the associated sub state.
+The returned value of this function is then pass to a remote call of `initialize_state`.
+"""
+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