metagraphoptimization.jl/src/operations/clean.jl

# functions for "cleaning" nodes, i.e. regenerating the possible operations for a node

# function to find node fusions involving the given node if it's a data node
# pushes the found fusion everywhere it needs to be and returns nothing
function find_fusions!(graph::DAG, node::DataTaskNode)
   if length(node.parents) != 1 || length(node.children) != 1
      return nothing
   end

   child_node = first(node.children)
   parent_node = first(node.parents)

   #=if !(child_node in graph) || !(parent_node in graph)
      error("Parents/Children that are not in the graph!!!")
   end=#

   if length(child_node.parents) != 1
      return nothing
   end

   nf = NodeFusion((child_node, node, parent_node))
   push!(graph.possibleOperations.nodeFusions, nf)
   push!(child_node.operations, nf)
   push!(node.operations, nf)
   push!(parent_node.operations, nf)

   return nothing
end

# function to find node fusions involving the given node if it's a compute node
# pushes the found fusion(s) everywhere it needs to be and returns nothing
function find_fusions!(graph::DAG, node::ComputeTaskNode)
   # for loop that always runs once for a scoped block we can break out of
   for _ in 1:1
      # assume this node as child of the chain
      if length(node.parents) != 1
         break
      end
      node2 = first(node.parents)
      if length(node2.parents) != 1 || length(node2.children) != 1
         break
      end
      node3 = first(node2.parents)

      #=if !(node2 in graph) || !(node3 in graph)
         error("Parents/Children that are not in the graph!!!")
      end=#

      nf = NodeFusion((node, node2, node3))
      push!(graph.possibleOperations.nodeFusions, nf)
      push!(node.operations, nf)
      push!(node2.operations, nf)
      push!(node3.operations, nf)
   end

   for _ in 1:1
      # assume this node as parent of the chain
      if length(node.children) < 1
         break
      end
      node2 = first(node.children)
      if length(node2.parents) != 1 || length(node2.children) != 1
         break
      end
      node1 = first(node2.children)
      if (length(node1.parents) > 1)
         break
      end

      #=if !(node2 in graph) || !(node1 in graph)
         error("Parents/Children that are not in the graph!!!")
      end=#

      nf = NodeFusion((node1, node2, node))
      push!(graph.possibleOperations.nodeFusions, nf)
      push!(node1.operations, nf)
      push!(node2.operations, nf)
      push!(node.operations, nf)
   end

   return nothing
end

function find_reductions!(graph::DAG, node::Node)
   reductionVector = nothing
   # possible reductions are with nodes that are partners, i.e. parents of children
   partners_ = partners(node)
   delete!(partners_, node)
   for partner in partners_
      if can_reduce(node, partner)
         if reductionVector === nothing
            # only when there's at least one reduction partner, insert the vector
            reductionVector = Vector{Node}()
            push!(reductionVector, node)
         end

         push!(reductionVector, partner)
      end
   end

   if reductionVector !== nothing
      nr = NodeReduction(reductionVector)
      push!(graph.possibleOperations.nodeReductions, nr)
      for node in reductionVector
         push!(node.operations, nr)
      end
   end

   return nothing
end

function find_splits!(graph::DAG, node::Node)
   if (can_split(node))
      ns = NodeSplit(node)
      push!(graph.possibleOperations.nodeSplits, ns)
      push!(node.operations, ns)
   end

   return nothing
end

# "clean" the operations on a dirty node
function clean_node!(graph::DAG, node::Node)
   find_fusions!(graph, node)
   find_reductions!(graph, node)
   find_splits!(graph, node)
end