metagraphoptimization.jl/src/operations/find.jl

# functions that find operations on the inital graph

using Base.Threads

function insert_operation!(operations::PossibleOperations, nf::NodeFusion, locks::Dict{Node, SpinLock})
   n1 = nf.input[1]; n2 = nf.input[2]; n3 = nf.input[3]

   lock(locks[n1]) do; push!(nf.input[1].operations, nf); end
   lock(locks[n2]) do; push!(nf.input[2].operations, nf); end
   lock(locks[n3]) do; push!(nf.input[3].operations, nf); end
   return nothing
end

function insert_operation!(operations::PossibleOperations, nr::NodeReduction, locks::Dict{Node, SpinLock})
   # since node parents were sorted before, the NodeReductions contain elements in a known order
   # this, together with the locking, means that we can safely do the following without inserting duplicates
   first = true
   for n in nr.input
      skip_duplicate = false
      # careful here, this is a manual lock (because of the break)
      lock(locks[n])
      if first
         first = false
         for op in n.operations
            if typeof(op) <: NodeReduction
               skip_duplicate = true
               break
            end
         end
         if skip_duplicate
            unlock(locks[n])
            break
         end         
      end

      push!(n.operations, nr)
      unlock(locks[n])
   end
   return nothing
end

function insert_operation!(operations::PossibleOperations, ns::NodeSplit, locks::Dict{Node, SpinLock})
   lock(locks[ns.input]) do; push!(ns.input.operations, ns); end
   return nothing
end

function nr_insertion!(operations::PossibleOperations, nodeReductions::Vector{Vector{NodeReduction}}, locks::Dict{Node, SpinLock})
   total_len = 0
   for vec in nodeReductions
      total_len += length(vec)
   end
   sizehint!(operations.nodeReductions, total_len)

   t = @task for vec in nodeReductions
      union!(operations.nodeReductions, Set(vec))
   end
   schedule(t)

   @threads for vec in nodeReductions
      for op in vec
         insert_operation!(operations, op, locks)
      end
   end

   wait(t)

   return nothing
end

function nf_insertion!(operations::PossibleOperations, nodeFusions::Vector{Vector{NodeFusion}}, locks::Dict{Node, SpinLock})
   total_len = 0
   for vec in nodeFusions
      total_len += length(vec)
   end
   sizehint!(operations.nodeFusions, total_len)
   
   t = @task for vec in nodeFusions
      union!(operations.nodeFusions, Set(vec))
   end
   schedule(t)

   @threads for vec in nodeFusions 
      for op in vec
         insert_operation!(operations, op, locks)
      end
   end

   wait(t)

   return nothing
end

function ns_insertion!(operations::PossibleOperations, nodeSplits::Vector{Vector{NodeSplit}}, locks::Dict{Node, SpinLock})
   total_len = 0
   for vec in nodeSplits
      total_len += length(vec)
   end
   sizehint!(operations.nodeSplits, total_len)

   t = @task for vec in nodeSplits
      union!(operations.nodeSplits, Set(vec))
   end
   schedule(t)

   @threads for vec in nodeSplits
      for op in vec
         insert_operation!(operations, op, locks)
      end
   end

   wait(t)

   return nothing
end

# function to generate all possible operations on the graph
function generate_options(graph::DAG)
   locks = Dict{Node, SpinLock}()
   for n in graph.nodes
      locks[n] = SpinLock()
   end

   generatedFusions = [Vector{NodeFusion}() for _ in 1:nthreads()]
   generatedReductions = [Vector{NodeReduction}() for _ in 1:nthreads()]
   generatedSplits = [Vector{NodeSplit}() for _ in 1:nthreads()]

   # make sure the graph is fully generated through
   apply_all!(graph)

   nodeArray = collect(graph.nodes)

   # sort all nodes
   @threads for node in nodeArray
      sort_node!(node)
   end

   checkedNodes = Set{Node}()
   checkedNodesLock = SpinLock()
   # --- find possible node reductions ---
   @threads for node in nodeArray
      # we're looking for nodes with multiple parents, those parents can then potentially reduce with one another
      if (length(node.parents) <= 1)
         continue
      end

      candidates = node.parents

      # sort into equivalence classes
      trie = NodeTrie()

      for candidate in candidates
         # insert into trie
         insert!(trie, candidate)
      end

      nodeReductions = collect(trie)

      for nrVec in nodeReductions
         # parent sets are ordered and any node can only be part of one nodeReduction, so a NodeReduction is uniquely identifiable by its first element
         # this prevents duplicate nodeReductions being generated
         lock(checkedNodesLock)
         if (nrVec[1] in checkedNodes)
            unlock(checkedNodesLock)
            continue
         else
            push!(checkedNodes, nrVec[1])
         end
         unlock(checkedNodesLock)

         push!(generatedReductions[threadid()], NodeReduction(nrVec))
      end
   end

   
   # launch thread for node reduction insertion
   # remove duplicates
   nr_task = @task nr_insertion!(graph.possibleOperations, generatedReductions, locks)
   schedule(nr_task)

   # --- find possible node fusions ---
   @threads for node in nodeArray
      if (typeof(node) <: DataTaskNode)
         if length(node.parents) != 1
            # data node can only have a single parent
            continue
         end
         parent_node = first(node.parents)

         if length(node.children) != 1
            # this node is an entry node or has multiple children which should not be possible
            continue
         end
         child_node = first(node.children)
         if (length(child_node.parents) != 1)
            continue
         end

         push!(generatedFusions[threadid()], NodeFusion((child_node, node, parent_node)))
      end
   end

   # launch thread for node fusion insertion
   nf_task = @task nf_insertion!(graph.possibleOperations, generatedFusions, locks)
   schedule(nf_task)

   # find possible node splits
   @threads for node in nodeArray
      if (can_split(node))
         push!(generatedSplits[threadid()], NodeSplit(node))
      end
   end

   # launch thread for node split insertion
   ns_task = @task ns_insertion!(graph.possibleOperations, generatedSplits, locks)
   schedule(ns_task)

   empty!(graph.dirtyNodes)

   wait(nr_task)
   wait(nf_task)
   wait(ns_task)

   return nothing
end