Add DaggerGraphs subpackage

Author: jpsamaroo

lib/DaggerGraphs/Project.toml

@@ -0,0 +1,13 @@
+name = "DaggerGraphs"
+uuid = "304567ff-242f-4479-af00-6fcdcd11a1dd"
+authors = ["Julian P Samaroo <jpsamaroo@jpsamaroo.me>"]
+version = "0.1.0"
+
+[deps]
+Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+
+[compat]
+Dagger = "0.17, 0.18"
+Graphs = "1"
+julia = "1"

lib/DaggerGraphs/src/DaggerGraphs.jl

@@ -0,0 +1,239 @@
+module DaggerGraphs
+
+using Dagger
+import Dagger: Chunk
+using Graphs
+
+export DGraph
+
+const ELTYPE = Union{Dagger.EagerThunk, Chunk}
+
+struct DGraphState{T}
+    # A set of locally-connected SimpleDiGraphs
+    parts::Vector{ELTYPE}
+    # The range of vertices within each of `parts`
+    parts_nv::Vector{UnitRange{T}}
+    # The number of edges in each of `parts`
+    parts_ne::Vector{T}
+    # The maximum number of nodes for each of `parts`
+    parts_v_max::Int
+
+    # A set of `AdjList` for each of `parts`
+    # An edge is present here if either src or dst (but not both) is in
+    # the respective `parts` graph
+    ext_adjs::Vector{ELTYPE}
+    # The number of edges in each of `ext_adjs`
+    ext_adjs_ne::Vector{T}
+    # The number of edges in each of `ext_adjs` where the source is this partition
+    ext_adjs_ne_src::Vector{T}
+end
+struct DGraph{T} <: Graphs.AbstractGraph{T}
+    state::Dagger.Chunk{DGraphState{T}}
+end
+function DGraph{T}(; chunksize=8) where {T}
+    state = DGraphState{T}(ELTYPE[],
+                           UnitRange{T}[],
+                           T[],
+                           chunksize,
+                           ELTYPE[],
+                           T[],
+                           T[])
+    return DGraph{T}(Dagger.tochunk(state))
+end
+DGraph(; kwargs...) = DGraph{Int}(; kwargs...)
+
+Base.eltype(::DGraph{T}) where T = T
+Graphs.edgetype(::DGraph{T}) where T = Tuple{T,T}
+Graphs.nv(g::DGraph) = fetch(Dagger.@spawn nv(g.state))
+function Graphs.nv(g::DGraphState)
+    if !isempty(g.parts_nv)
+        return last(g.parts_nv).stop
+    else
+        return 0
+    end
+end
+Graphs.ne(g::DGraph) = fetch(Dagger.@spawn ne(g.state))
+Graphs.ne(g::DGraphState) = sum(g.parts_ne; init=0) + sum(g.ext_adjs_ne_src; init=0)
+Graphs.has_vertex(g::DGraph, v::Int) = 1 <= v <= nv(g)
+Graphs.has_edge(g::DGraph, src::Int, dst::Int) =
+    fetch(Dagger.@spawn has_edge(g.state, src, dst))
+function Graphs.has_edge(g::DGraphState, src::Int, dst::Int)
+    src_part_idx = findfirst(span->src in span, g.parts_nv)
+    src_part_idx !== nothing || return false
+    dst_part_idx = findfirst(span->dst in span, g.parts_nv)
+    dst_part_idx !== nothing || return false
+
+    if src_part_idx == dst_part_idx
+        # The edge will be within a graph partition
+        part = g.parts[src_part_idx]
+        return fetch(Dagger.@spawn has_edge(part, src, dst))
+    else
+        # The edge will be in an AdjList
+        adj = g.ext_adjs[src_part_idx]
+        return fetch(Dagger.@spawn has_ext_adj(adj, src, dst))
+    end
+end
+Graphs.is_directed(::DGraph) = true
+Graphs.vertices(g::DGraph) = Base.OneTo(nv(g))
+Graphs.edges(g::DGraph) = DGraphEdgeIter(g)
+Graphs.zero(::Type{<:DGraph}) = DGraph()
+Graphs.add_vertex!(g::DGraph) = fetch(Dagger.@spawn add_vertex!(g.state))
+function Graphs.add_vertex!(g::DGraphState)
+    n = nv(g)
+    if fld(n, g.parts_v_max) == length(g.parts)
+        # We need to create a new partition for this vertex
+        push!(g.parts, Dagger.spawn() do
+            g = SimpleDiGraph()
+            add_vertex!(g)
+            g
+        end)
+        push!(g.parts_nv, (n+1):(n+1))
+        push!(g.parts_ne, 0)
+        push!(g.ext_adjs, Dagger.@spawn AdjList())
+        push!(g.ext_adjs_ne, 0)
+        push!(g.ext_adjs_ne_src, 0)
+    else
+        # We will add this vertex to the last partition
+        part = last(g.parts)
+        fetch(Dagger.@spawn add_vertex!(part))
+        span = g.parts_nv[end]
+        g.parts_nv[end] = UnitRange{Int}(span.start, span.stop+1)
+    end
+    return true
+end
+Graphs.add_edge!(g::DGraph, src::Int, dst::Int) =
+    fetch(Dagger.@spawn add_edge!(g.state, src, dst))
+function Graphs.add_edge!(g::DGraphState, src::Int, dst::Int)
+    src_part_idx = findfirst(span->src in span, g.parts_nv)
+    @assert src_part_idx !== nothing "Source vertex $src does not exist"
+
+    dst_part_idx = findfirst(span->dst in span, g.parts_nv)
+    @assert dst_part_idx !== nothing "Destination vertex $dst does not exist"
+
+    if src_part_idx == dst_part_idx
+        # Edge exists within a single partition
+        part = g.parts[src_part_idx]
+        if fetch(Dagger.@spawn add_edge!(part, src, dst))
+            g.parts_ne[src_part_idx] += 1
+        else
+            return false
+        end
+    else
+        # Edge spans two partitions
+        src_ext_adj = g.ext_adjs[src_part_idx]
+        dst_ext_adj = g.ext_adjs[dst_part_idx]
+        if !fetch(Dagger.@spawn add_ext_adj!(src_ext_adj, src, dst))
+            return false
+        end
+        fetch(Dagger.@spawn add_ext_adj!(dst_ext_adj, src, dst))
+        g.ext_adjs_ne_src[src_part_idx] += 1
+        g.ext_adjs_ne[src_part_idx] += 1
+        g.ext_adjs_ne[dst_part_idx] += 1
+    end
+
+    return true
+end
+
+struct AdjList{T}
+    adj::Vector{Tuple{T,T}}
+end
+AdjList() = AdjList{Int}(Tuple{Int,Int}[])
+function has_ext_adj(adj::AdjList, src::Int, dst::Int)
+    idx = findfirst(edge->edge[1] == src && edge[2] == dst, adj.adj)
+    if idx !== nothing
+        return true
+    end
+    return false
+end
+function add_ext_adj!(adj::AdjList, src::Int, dst::Int)
+    if has_ext_adj(adj, src, dst)
+        return false
+    end
+    push!(adj.adj, (src, dst))
+    return true
+end
+Graphs.edges(adj::AdjList) = adj.adj
+
+struct DGraphEdgeIter{T} <: Graphs.AbstractEdgeIter
+    graph::DGraphState{T}
+end
+DGraphEdgeIter(g::DGraph) = DGraphEdgeIter(fetch(g.state))
+struct DGraphEdgeIterState
+    adj::Bool
+    part::Int
+    idx::Int
+    cache::Any
+end
+Base.length(iter::DGraphEdgeIter) = ne(iter.graph)
+Base.eltype(iter::DGraphEdgeIter{T}) where T = Tuple{T,T}
+function Base.iterate(iter::DGraphEdgeIter)
+    g = iter.graph
+    if nv(g) == 0
+        return nothing
+    elseif sum(g.parts_ne; init=0) > 0
+        # Start with partitions
+        return iterate(iter, DGraphEdgeIterState(false, 1, 1, nothing))
+    elseif sum(g.ext_adjs_ne_src; init=0) > 0
+        # Start with external AdjLists
+        return iterate(iter, DGraphEdgeIterState(true, 1, 1, nothing))
+    else
+        return nothing
+    end
+end
+function Base.iterate(iter::DGraphEdgeIter{T}, state::DGraphEdgeIterState) where {T}
+    g = iter.graph
+    adj = state.adj
+    part = state.part
+    idx = state.idx
+    cache = state.cache
+
+    @label start
+    if !adj
+        if part > length(g.parts)
+            # Restart with external AdjLists
+            return iterate(iter, DGraphEdgeIterState(true, 1, 1, nothing))
+        end
+        if cache === nothing
+            cache = map(Tuple, fetch(Dagger.@spawn edges(g.parts[part])))
+        end
+    else
+        if part > length(g.ext_adjs)
+            # All done!
+            return nothing
+        end
+        if cache === nothing
+            cache = fetch(Dagger.@spawn edges(g.ext_adjs[part]))
+        end
+    end
+    cache::Vector{<:Tuple}
+
+    # Skip empty edge sets
+    if isempty(cache)
+        part += 1
+        idx = 1
+        cache = nothing
+        @goto start
+    end
+    cache::Vector{Tuple{T,T}}
+
+    # Get the current edge
+    value = cache[idx]
+    idx += 1
+    cur_part = part
+
+    # Reset if this partition/AdjList is exhausted
+    if idx > length(cache)
+        part += 1
+        idx = 1
+        cache = nothing
+    end
+
+    # Restart if this edge isn't "owned" by this AdjList
+    if adj && !(value[1] in g.parts_nv[cur_part])
+        @goto start
+    end
+
+    return (value, DGraphEdgeIterState(adj, part, idx, cache))
+end
+
+end # module DaggerGraphs

lib/DaggerGraphs/test/Project.toml

@@ -0,0 +1,3 @@
+[deps]
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

lib/DaggerGraphs/test/runtests.jl

@@ -0,0 +1,36 @@
+using DaggerGraphs, Graphs
+using Test
+
+function generate_random_dgraph(N)
+    g = DGraph()
+    n_ctr = 1
+    add_vertex!(g)
+    n_edges = 0
+    for i in 1:N
+        src = rand() < 0.1 ? n_ctr+1 : rand(1:n_ctr)
+        if src > n_ctr
+            add_vertex!(g)
+            n_ctr += 1
+        end
+        dst = rand() < 0.1 ? n_ctr+1 : rand(1:n_ctr)
+        if dst > n_ctr
+            add_vertex!(g)
+            n_ctr += 1
+        end
+        if add_edge!(g, src, dst)
+            n_edges += 1
+        end
+    end
+    @test nv(g) == n_ctr
+    @test ne(g) == n_edges
+    return g
+end
+
+@testset "Basics" begin
+    g = generate_random_dgraph(1000)
+    @test vertices(g) == Base.OneTo(nv(g))
+    @test length(edges(g)) == ne(g)
+    e = collect(edges(g))
+    @test e isa Vector{Tuple{Int,Int}}
+    @test all(edge->has_edge(g, edge[1], edge[2]), e)
+end