CommonNodeElimination.cpp

Go to the documentation of this file.

/*
 * Copyright 2017 Nico Reißmann <nico.reissmann@gmail.com>
 * Copyright 2025 Håvard Krogstie <krogstie.havard@gmail.com>
 * See COPYING for terms of redistribution.
 */
 
#include <jlm/llvm/ir/operators/delta.hpp>
#include <jlm/llvm/ir/operators/lambda.hpp>
#include <jlm/llvm/ir/RvsdgModule.hpp>
#include <jlm/llvm/opt/CommonNodeElimination.hpp>
#include <jlm/rvsdg/gamma.hpp>
#include <jlm/rvsdg/MatchType.hpp>
#include <jlm/rvsdg/Phi.hpp>
#include <jlm/rvsdg/theta.hpp>
#include <jlm/rvsdg/traverser.hpp>
#include <jlm/util/common.hpp>
#include <jlm/util/Hash.hpp>
#include <jlm/util/Statistics.hpp>
#include <jlm/util/time.hpp>
 
#include <map>
#include <unordered_map>
 
namespace jlm::llvm
{
 
class CommonNodeElimination::Statistics final : public util::Statistics
{
  const char * MarkTimerLabel_ = "MarkTime";
  const char * DivertTimerLabel_ = "DivertTime";
  const char * PruneTimerLabel_ = "PruneTime";
 
public:
  ~Statistics() override = default;
 
  explicit Statistics(const util::FilePath & sourceFile)
      : util::Statistics(Statistics::Id::CommonNodeElimination, sourceFile)
  {}
 
  void
  startMarkStatistics(const rvsdg::Graph & graph) noexcept
  {
    AddMeasurement(Label::NumRvsdgInputsBefore, rvsdg::ninputs(&graph.GetRootRegion()));
    AddTimer(MarkTimerLabel_).start();
  }
 
  void
  endMarkStatistics() noexcept
  {
    GetTimer(MarkTimerLabel_).stop();
  }
 
  void
  startDivertStatistics() noexcept
  {
    AddTimer(DivertTimerLabel_).start();
  }
 
  void
  endDivertStatistics(const rvsdg::Graph & graph) noexcept
  {
    AddMeasurement(Label::NumRvsdgInputsAfter, rvsdg::ninputs(&graph.GetRootRegion()));
    GetTimer(DivertTimerLabel_).stop();
  }
 
  void
  startPruneStatistics() noexcept
  {
    AddTimer(PruneTimerLabel_).start();
  }
 
  void
  stopPruneStatistics() noexcept
  {
    GetTimer(PruneTimerLabel_).stop();
  }
 
  static std::unique_ptr<Statistics>
  Create(const util::FilePath & sourceFile)
  {
    return std::make_unique<Statistics>(sourceFile);
  }
};
 
class CommonNodeElimination::Context
{
public:
  struct CongruenceSet
  {
    explicit CongruenceSet(const rvsdg::Output & leader)
        : leader(&leader)
    {}
 
    // The set leader. Never changes.
    // Once an output has become a leader, it will never be a follower.
    const rvsdg::Output * leader;
    // The set of followers. Can both grow and shrink during the marking phase.
    // Does not include the leader.
    util::HashSet<const rvsdg::Output *> followers;
  };
 
  using CongruenceSetIndex = size_t;
  static constexpr auto NoCongruenceSetIndex = std::numeric_limits<CongruenceSetIndex>::max();
 
  [[nodiscard]] CongruenceSetIndex
  numCongruenceSets() const
  {
    return sets_.size();
  }
 
  [[nodiscard]] CongruenceSetIndex
  tryGetSetFor(const rvsdg::Output & output) const
  {
    if (const auto it = congruenceSetMapping_.find(&output); it != congruenceSetMapping_.end())
    {
      return it->second;
    }
    return NoCongruenceSetIndex;
  }
 
  [[nodiscard]] CongruenceSetIndex
  getSetFor(const rvsdg::Output & output) const
  {
    const auto index = tryGetSetFor(output);
    if (index == NoCongruenceSetIndex)
      throw std::logic_error("Output does not belong to a congruence set");
    return index;
  }
 
  [[nodiscard]] bool
  hasSet(const rvsdg::Output & output) const
  {
    return tryGetSetFor(output) != NoCongruenceSetIndex;
  }
 
  CongruenceSetIndex
  getOrCreateSetForLeader(const rvsdg::Output & leader)
  {
    // The index of the new set, if this operation actually creates one
    auto nextSet = sets_.size();
    auto [it, added] = congruenceSetMapping_.try_emplace(&leader, nextSet);
 
    if (!added)
    {
      // If the leader already has its own set, we are done
      if (sets_[it->second].leader == &leader)
      {
        return it->second;
      }
 
      // Remove the output from the congruence set it is following
      sets_[it->second].followers.Remove(&leader);
      it->second = nextSet;
    }
 
    // Create the new set, lead by \p leader
    sets_.emplace_back(leader);
    return nextSet;
  }
 
  const rvsdg::Output &
  getLeader(CongruenceSetIndex index) const
  {
    JLM_ASSERT(index < sets_.size());
    return *sets_[index].leader;
  }
 
  void
  addFollower(CongruenceSetIndex index, const rvsdg::Output & follower)
  {
    JLM_ASSERT(index < sets_.size());
 
    const bool newFollower = sets_[index].followers.insert(&follower);
 
    // If the follower is already following the correct set, do nothing
    if (!newFollower)
      return;
 
    const auto [it, added] = congruenceSetMapping_.try_emplace(&follower, index);
 
    // If the follower already belonged to a congruence set, remove it from the old set
    if (!added)
    {
      JLM_ASSERT(it->second != index);
 
      if (sets_[it->second].leader == &follower)
        throw std::logic_error("Cannot turn a leader into a follower");
 
      const bool removed = sets_[it->second].followers.Remove(&follower);
      JLM_ASSERT(removed);
 
      it->second = index;
    }
  }
 
  const util::HashSet<const rvsdg::Output *> &
  getFollowers(CongruenceSetIndex index) const
  {
    JLM_ASSERT(index < sets_.size());
    return sets_[index].followers;
  }
 
private:
  // The list of congruence sets
  std::vector<CongruenceSet> sets_;
  // A mapping from each output to the congruence set it belongs to, either as leader or follower
  std::unordered_map<const rvsdg::Output *, CongruenceSetIndex> congruenceSetMapping_;
};
 
static bool
areOutputsCongruent(
    const rvsdg::Output & o1,
    const rvsdg::Output & o2,
    CommonNodeElimination::Context & context)
{
  if (*o1.Type() != *o2.Type())
    return false;
 
  const auto o1Set = context.getSetFor(o1);
  const auto o2Set = context.getSetFor(o2);
 
  return o1Set == o2Set;
}
 
static bool
checkNodesCongruent(
    const rvsdg::Node & node1,
    const rvsdg::Node & node2,
    CommonNodeElimination::Context & context)
{
  const auto simpleNode1 = dynamic_cast<const rvsdg::SimpleNode *>(&node1);
  const auto simpleNode2 = dynamic_cast<const rvsdg::SimpleNode *>(&node2);
  if (!simpleNode1 || !simpleNode2)
    return false;
 
  if (simpleNode1->ninputs() != simpleNode2->ninputs())
    return false;
 
  if (simpleNode1->GetOperation() != simpleNode2->GetOperation())
    return false;
 
  // For each pair of corresponding inputs of simpleNode1 and simpleNode2,
  // they must have origins that are congruent
  for (auto & input : simpleNode1->Inputs())
  {
    const auto origin1 = input.origin();
    const auto origin2 = simpleNode2->input(input.index())->origin();
    if (!areOutputsCongruent(*origin1, *origin2, context))
      return false;
  }
 
  return true;
}
 
void
markNodeAsLeader(const rvsdg::Node & leader, CommonNodeElimination::Context & context)
{
  for (auto & output : leader.Outputs())
  {
    context.getOrCreateSetForLeader(output);
  }
}
 
void
markNodesAsCongruent(
    const rvsdg::Node & leader,
    const rvsdg::Node & follower,
    CommonNodeElimination::Context & context)
{
  JLM_ASSERT(leader.noutputs() == follower.noutputs());
  JLM_ASSERT(leader.region() == follower.region());
 
  for (size_t i = 0; i < leader.noutputs(); i++)
  {
    const auto & leaderOutput = *leader.output(i);
    const auto & followerOutput = *follower.output(i);
    const auto leaderSet = context.getSetFor(leaderOutput);
    context.addFollower(leaderSet, followerOutput);
  }
}
 
[[nodiscard]] const rvsdg::SimpleNode *
tryGetLeaderNode(const rvsdg::SimpleNode & node, CommonNodeElimination::Context & context)
{
  // Nodes with 0 outputs are never congruent with anything, so let the node be its own leader.
  if (node.noutputs() == 0)
    return &node;
 
  // Check the congruence set of the first output
  const auto output0Set = context.tryGetSetFor(*node.output(0));
  // If the output has not gotten a congruence set yet, the node has yet to be marked
  if (output0Set == CommonNodeElimination::Context::NoCongruenceSetIndex)
    return nullptr;
 
  // A simple node can only be congruent with other simple nodes
  const auto & output0Leader = context.getLeader(output0Set);
  const auto & leaderNode = rvsdg::AssertGetOwnerNode<rvsdg::SimpleNode>(output0Leader);
  JLM_ASSERT(leaderNode.noutputs() == node.noutputs());
  return &leaderNode;
}
 
using TopNodeLeaderList = std::vector<const rvsdg::Node *>;
 
static void
markSimpleTopNode(
    const rvsdg::SimpleNode & node,
    TopNodeLeaderList & leaders,
    CommonNodeElimination::Context & context)
{
  // If the node has already been marked, check if it is still congruent with its leader
  if (const auto existingLeaderNode = tryGetLeaderNode(node, context))
  {
    // We are our own leader, nothing to check
    if (existingLeaderNode == &node)
    {
      leaders.push_back(&node);
      return;
    }
 
    // Check if we are still congruent with the leader
    if (checkNodesCongruent(*existingLeaderNode, node, context))
    {
      return;
    }
 
    // This node is no longer congruent with its leader, continue looking for a new leader
  }
 
  // TODO: Use some sort of hashing to make this not O(n * m)
  // where n is the number of outputs and m is the number of congruence sets
 
  // Check if the node is congruent with any existing leader in the TopNodeLeaderList
  for (auto leader : leaders)
  {
    if (checkNodesCongruent(node, *leader, context))
    {
      markNodesAsCongruent(*leader, node, context);
      return;
    }
  }
 
  // No existing leader top node found, create new congruence sets for each of node's outputs
  markNodeAsLeader(node, context);
  leaders.push_back(&node);
}
 
static void
markSimpleNode(const rvsdg::SimpleNode & node, CommonNodeElimination::Context & context)
{
  // This function should never be called for TopNodes
  JLM_ASSERT(node.ninputs() > 0);
 
  // If node is already in a congruence set, check that it actually belongs there
  if (const auto leaderNode = tryGetLeaderNode(node, context))
  {
    // If node is its own leader, it definitely belongs, and we are done
    if (leaderNode == &node)
      return;
 
    // Double-check that we are congruent with our leader
    if (checkNodesCongruent(node, *leaderNode, context))
      return;
 
    // Otherwise we need to continue looking for a new leader
  }
 
  // This function looks at all nodes that take the given output as the origin for its first input.
  // If the other node is its own leader, and the current node is congruent with it,
  // the current node becomes a follower.
  const auto tryFindCongruentUserOf = [&](const rvsdg::Output & output) -> bool
  {
    // TODO: It would be possible to maintain a list of only users that are leader nodes,
    // to avoid needing to check every user
    for (auto & user : output.Users())
    {
      if (user.index() != 0)
        continue;
 
      const auto otherNode = rvsdg::TryGetOwnerNode<rvsdg::SimpleNode>(user);
      if (!otherNode)
        continue;
 
      // Do not compare against ourselves
      if (otherNode == &node)
        continue;
 
      // Only consider other nodes that are leaders
      auto otherNodeLeader = tryGetLeaderNode(*otherNode, context);
      if (otherNode != otherNodeLeader)
        continue;
 
      if (checkNodesCongruent(node, *otherNode, context))
      {
        // When nodes are congruent, they should always have the same amount of outputs
        JLM_ASSERT(node.noutputs() == otherNode->noutputs());
 
        markNodesAsCongruent(*otherNode, node, context);
        return true;
      }
    }
    return false;
  };
 
  // Use the origin of the first input to find other potential candidates
  const auto origin = node.input(0)->origin();
  const auto origin0Set = context.getSetFor(*origin);
  const auto & origin0Leader = context.getLeader(origin0Set);
  const auto & origin0Followers = context.getFollowers(origin0Set);
  if (tryFindCongruentUserOf(origin0Leader))
    return;
  for (auto follower : origin0Followers.Items())
  {
    if (tryFindCongruentUserOf(*follower))
      return;
  }
 
  // If we got here, no congruent node was found, so make the node its own leader
  markNodeAsLeader(node, context);
}
 
static void
markRegion(const rvsdg::Region &, CommonNodeElimination::Context & context);
 
static void
markGraphImports(const rvsdg::Region & region, CommonNodeElimination::Context & context)
{
  // FIXME: Multiple imports with identical names could in theory be aliases.
  // For now this function ignores that and makes all imports distinct. This is sound
  for (auto argument : region.Arguments())
  {
    context.getOrCreateSetForLeader(*argument);
  }
}
 
static bool
partitionArguments(
    const rvsdg::Region & region,
    const std::vector<CommonNodeElimination::Context::CongruenceSetIndex> & partitions,
    CommonNodeElimination::Context & context)
{
  JLM_ASSERT(region.narguments() == partitions.size());
 
  const auto numCongruenceSets = context.numCongruenceSets();
 
  // Keys in the map are (old congruence set index, provided partition key)
  // Values in the map are the new congruence set indices
  std::map<
      std::pair<
          CommonNodeElimination::Context::CongruenceSetIndex,
          CommonNodeElimination::Context::CongruenceSetIndex>,
      CommonNodeElimination::Context::CongruenceSetIndex>
      newSets;
  for (auto argument : region.Arguments())
  {
    const auto currentPartition = context.tryGetSetFor(*argument);
    const auto key = std::make_pair(currentPartition, partitions[argument->index()]);
 
    // If this argument is the first with the given key, it should be a leader
    // otherwise it should be a follower
 
    if (const auto it = newSets.find(key); it != newSets.end())
    {
      // This argument should be a follower of the given congruence set
      const auto toFollow = it->second;
 
      // If we are already a follower, we are done
      if (currentPartition == toFollow)
        continue;
 
      // Start following our leader
      context.addFollower(toFollow, *argument);
    }
    else
    {
      // This argument should be the leader of its congruence set
      newSets.emplace(key, context.getOrCreateSetForLeader(*argument));
    }
  }
 
  // Return true iff any new congruence sets were created
  return context.numCongruenceSets() != numCongruenceSets;
}
 
static bool
markSubregionsFromInputs(
    const rvsdg::StructuralNode & node,
    CommonNodeElimination::Context & context)
{
  bool anyChanges = false;
 
  for (auto & subregion : node.Subregions())
  {
    if (subregion.narguments() == 0)
    {
      markRegion(subregion, context);
    }
    else
    {
      // create a partitioning of the region arguments
      std::vector<size_t> partitions(subregion.narguments());
      // Arguments that do not belong to any input are given partition keys higher than any real
      // index
      size_t nextUniquePartitionKey = context.numCongruenceSets();
 
      for (const auto argument : subregion.Arguments())
      {
        if (const auto input = argument->input())
        {
          // If the argument corresponds to an input, use the partition key of the input
          partitions[argument->index()] = context.getSetFor(*input->origin());
        }
        else
        {
          // Otherwise make sure the argument is not partitioned with any other argument
          partitions[argument->index()] = nextUniquePartitionKey++;
        }
      }
 
      if (partitionArguments(subregion, partitions, context))
      {
        anyChanges = true;
        markRegion(subregion, context);
      }
    }
  }
 
  return anyChanges;
}
 
static std::optional<CommonNodeElimination::Context::CongruenceSetIndex>
tryGetGammaExitVarCongruenceSet(
    rvsdg::GammaNode::ExitVar & exitVar,
    CommonNodeElimination::Context & context)
{
  std::optional<CommonNodeElimination::Context::CongruenceSetIndex> sharedCongruenceSet;
 
  for (auto result : exitVar.branchResult)
  {
    if (const auto argument = dynamic_cast<rvsdg::RegionArgument *>(result->origin()))
    {
      const auto inputCongruenceSet = context.getSetFor(*argument->input()->origin());
      if (!sharedCongruenceSet.has_value())
      {
        sharedCongruenceSet = inputCongruenceSet;
      }
      else if (*sharedCongruenceSet != inputCongruenceSet)
      {
        // We have multiple different non-congruent origins
        return std::nullopt;
      }
    }
    else
    {
      // The branch result was not invariant
      return std::nullopt;
    }
  }
 
  return sharedCongruenceSet;
}
 
[[nodiscard]] static size_t
getGammaExitVariableHash(
    const rvsdg::GammaNode::ExitVar & exitVar,
    CommonNodeElimination::Context & context)
{
  size_t hash = 1;
  for (auto branchResult : exitVar.branchResult)
  {
    const auto set = context.getSetFor(*branchResult->origin());
    hash = util::CombineHashes(hash, set);
  }
  return hash;
}
 
[[nodiscard]] static bool
areGammaExitVariablesCongruent(
    const rvsdg::GammaNode::ExitVar & first,
    const rvsdg::GammaNode::ExitVar & second,
    CommonNodeElimination::Context & context)
{
  JLM_ASSERT(
      &rvsdg::AssertGetOwnerNode<rvsdg::GammaNode>(*first.output)
      == &rvsdg::AssertGetOwnerNode<rvsdg::GammaNode>(*second.output));
  const auto numResults = first.branchResult.size();
  for (size_t i = 0; i < numResults; i++)
  {
    const auto firstSet = context.getSetFor(*first.branchResult[i]->origin());
    const auto secondSet = context.getSetFor(*second.branchResult[i]->origin());
    if (firstSet != secondSet)
      return false;
  }
  return true;
}
 
static void
insertGammaExitVarInHashmap(
    rvsdg::GammaNode::ExitVar & exitVar,
    CommonNodeElimination::Context::CongruenceSetIndex congruenceSet,
    std::unordered_map<size_t, CommonNodeElimination::Context::CongruenceSetIndex> & leaderHashes,
    CommonNodeElimination::Context & context)
{
  JLM_ASSERT(&context.getLeader(congruenceSet) == exitVar.output);
 
  size_t hash = getGammaExitVariableHash(exitVar, context);
  do
  {
    const auto [_, inserted] = leaderHashes.emplace(hash, congruenceSet);
    if (inserted)
      return;
 
    // Open addressing, try next hash
    hash++;
  } while (true);
}
 
static void
lookupOrInsertGammaExitVarInHashmap(
    const rvsdg::GammaNode::ExitVar & exitVar,
    const rvsdg::GammaNode & gamma,
    std::unordered_map<size_t, CommonNodeElimination::Context::CongruenceSetIndex> & leaderHashes,
    CommonNodeElimination::Context & context)
{
  size_t hash = getGammaExitVariableHash(exitVar, context);
  do
  {
    // If no leader has the same hash already, insert a new leader exit variable
    const auto [it, inserted] = leaderHashes.emplace(hash, 0);
    if (inserted)
    {
      it->second = context.getOrCreateSetForLeader(*exitVar.output);
      return;
    }
 
    // There is already an exit variable with the same hash, check if it is a match
    auto & otherLeader = context.getLeader(it->second);
    auto otherExitVar = gamma.MapOutputExitVar(otherLeader);
    if (areGammaExitVariablesCongruent(exitVar, otherExitVar, context))
    {
      context.addFollower(it->second, *exitVar.output);
      return;
    }
 
    // Open addressing, try next hash
    hash++;
  } while (true);
}
 
static void
markGamma(const rvsdg::GammaNode & gamma, CommonNodeElimination::Context & context)
{
  markSubregionsFromInputs(gamma, context);
 
  // Mapping from exit variable hash to the leader of its congruence set.
  // The hash function is not collision free, and the table uses open addressing.
  std::unordered_map<size_t, CommonNodeElimination::Context::CongruenceSetIndex> leaderHashes;
 
  // Go through the outputs of the gamma node and create congruence sets for them.
  for (auto exitVar : gamma.GetExitVars())
  {
    // If the exit variable has previously been found to not be invariant, we skip trying again
    bool skipInvarianceCheck = false;
 
    // First check if the exit variable already belongs to a congruence set
    const auto existingSet = context.tryGetSetFor(*exitVar.output);
    if (existingSet != CommonNodeElimination::Context::NoCongruenceSetIndex)
    {
      // The output already belongs to a congruence set
 
      // Check if the output is the leader
      const auto & exisitingLeader = context.getLeader(existingSet);
      if (&exisitingLeader == exitVar.output)
      {
        // Add the exit var to the hashmap of exit variable leaders
        insertGammaExitVarInHashmap(exitVar, existingSet, leaderHashes, context);
 
        // A leader can never become a follower, so we are done
        continue;
      }
 
      // Check if the leader is another exit variable of the same gamma
      if (rvsdg::TryGetOwnerNode<rvsdg::GammaNode>(exisitingLeader) == &gamma)
      {
        const auto otherExitVar = gamma.MapOutputExitVar(exisitingLeader);
        // Double check that the output still belongs to the same congurence set
        if (areGammaExitVariablesCongruent(exitVar, otherExitVar, context))
        {
          // The exit variable should continue to be a follower
          continue;
        }
 
        // exitVar should find a new leader, but we can skip checking if it is invariant
        skipInvarianceCheck = true;
      }
    }
 
    if (!skipInvarianceCheck)
    {
      // First check if the exit variable is gamma invariant
      if (const auto entryVarCongruenceSet = tryGetGammaExitVarCongruenceSet(exitVar, context);
          entryVarCongruenceSet.has_value())
      {
        context.addFollower(*entryVarCongruenceSet, *exitVar.output);
        continue;
      }
    }
 
    // This function either finds a new congruent leader for the exitVar,
    // or makes exitVar itself the leader of its own congruence set.
    lookupOrInsertGammaExitVarInHashmap(exitVar, gamma, leaderHashes, context);
  }
}
 
static void
markTheta(const rvsdg::ThetaNode & theta, CommonNodeElimination::Context & context)
{
  bool anyChanges = markSubregionsFromInputs(theta, context);
  const auto loopVars = theta.GetLoopVars();
 
  // If anyChanges is false, it means no congruence sets changed within the theta node.
  // However, the outputs of invariant loop variables belong to the congruence set of their input,
  // which may have changed since the last time this function was called.
  bool onlyUpdateLoopInvariants = !anyChanges;
 
  // Use the loop variable post results to refine partitioning of loop variable arguments
  while (anyChanges)
  {
    // Create partition keys for each loop variable
    std::vector<CommonNodeElimination::Context::CongruenceSetIndex> partitions;
    for (const auto & loopVar : loopVars)
    {
      partitions.push_back(context.getSetFor(*loopVar.post->origin()));
    }
 
    anyChanges = partitionArguments(*theta.subregion(), partitions, context);
    if (anyChanges)
    {
      // Propagate refinement of argument congruence sets into the region
      markRegion(*theta.subregion(), context);
    }
  }
 
  // Check if the given loop variable's output is loop invariant.
  // If it is, the origin of the invariant value is returned.
  // If it is not, nullptr is returned.
  const auto isOutputInvariant = [&](const auto & loopVar) -> rvsdg::Output *
  {
    // First perform a trivial check
    if (rvsdg::ThetaLoopVarIsInvariant(loopVar))
    {
      return loopVar.input->origin();
    }
 
    const auto resultSet = context.getSetFor(*loopVar.post->origin());
    const auto & resultSetLeader = context.getLeader(resultSet);
    // If the origin's leader is the loop variable's pre, the result is also invariant
    if (&resultSetLeader == loopVar.pre)
    {
      return loopVar.input->origin();
    }
 
    // Finally check if the loop variable post takes its value from another loop variable,
    // and if that other loop variable is invariant
    if (rvsdg::TryGetRegionParentNode<rvsdg::ThetaNode>(resultSetLeader) == &theta)
    {
      const auto otherLoopVar = theta.MapPreLoopVar(resultSetLeader);
      // Check if the other loop variable is trivially invariant
      if (rvsdg::ThetaLoopVarIsInvariant(otherLoopVar))
      {
        return otherLoopVar.input->origin();
      }
 
      // Use the CNE context to check if the other loop variable is invariant
      const auto otherLoopVarPostSet = context.getSetFor(*otherLoopVar.post->origin());
      const auto & otherLoopVarPostLeader = context.getLeader(otherLoopVarPostSet);
      if (&otherLoopVarPostLeader == otherLoopVar.pre)
      {
        return otherLoopVar.input->origin();
      }
    }
 
    return nullptr;
  };
 
  // Partition theta outputs
  std::unordered_map<
      CommonNodeElimination::Context::CongruenceSetIndex,
      CommonNodeElimination::Context::CongruenceSetIndex>
      resultToOutputSetMapping;
  for (auto & loopVar : loopVars)
  {
    // If a loop variable is invariant, or its post result is congruent with
    // another loop variable that is invariant, the output becomes a follower of the input origin.
    if (const auto origin = isOutputInvariant(loopVar))
    {
      auto inputCongruenceSet = context.getSetFor(*origin);
      context.addFollower(inputCongruenceSet, *loopVar.output);
      continue;
    }
 
    if (onlyUpdateLoopInvariants)
      continue;
 
    // Other loop variable outputs are partitioned based on the origins of the post results.
    const auto resultSet = context.getSetFor(*loopVar.post->origin());
    const auto it = resultToOutputSetMapping.find(resultSet);
    if (it != resultToOutputSetMapping.end())
    {
      context.addFollower(it->second, *loopVar.output);
    }
    else
    {
      const auto outputSet = context.getOrCreateSetForLeader(*loopVar.output);
      resultToOutputSetMapping.emplace(resultSet, outputSet);
    }
  }
}
 
static void
markStructuralNode(const rvsdg::StructuralNode & node, CommonNodeElimination::Context & context)
{
  rvsdg::MatchTypeOrFail(
      node,
      [&](const rvsdg::GammaNode & gamma)
      {
        markGamma(gamma, context);
      },
      [&](const rvsdg::ThetaNode & theta)
      {
        markTheta(theta, context);
      },
      [&](const rvsdg::LambdaNode & lambda)
      {
        // Context variables are congruent if their origins are congruent.
        // All other arguments are given distinct congruence sets.
        markSubregionsFromInputs(lambda, context);
 
        // A lambda output is always unique
        markNodeAsLeader(lambda, context);
      },
      [&](const rvsdg::PhiNode & phi)
      {
        // Context variables are congruent if their origins are congruent.
        // Recursion variables are given distinct congruence sets.
        markSubregionsFromInputs(phi, context);
 
        // A phi node is always unique
        markNodeAsLeader(phi, context);
      },
      [&](const rvsdg::DeltaNode & delta)
      {
        // We skip doing CNE inside delta nodes
 
        // A delta node is always unique
        markNodeAsLeader(delta, context);
      });
}
 
static void
markRegion(const rvsdg::Region & region, CommonNodeElimination::Context & context)
{
  TopNodeLeaderList leaders;
 
  for (const auto node : rvsdg::TopDownConstTraverser(&region))
  {
    rvsdg::MatchTypeOrFail(
        *node,
        [&](const rvsdg::SimpleNode & simple)
        {
          // Handle top nodes as a special case
          if (node->ninputs() == 0)
          {
            markSimpleTopNode(simple, leaders, context);
          }
          else
          {
            markSimpleNode(simple, context);
          }
        },
        [&](const rvsdg::StructuralNode & structural)
        {
          markStructuralNode(structural, context);
        });
  }
}
 
/* divert phase */
 
static void
divertOutput(rvsdg::Output & output, CommonNodeElimination::Context & context)
{
  const auto outputSet = context.getSetFor(output);
 
  auto & leader = context.getLeader(outputSet);
  if (&leader == &output)
    return;
 
  output.divert_users(const_cast<rvsdg::Output *>(&leader));
}
 
static void
divertInRegion(rvsdg::Region &, CommonNodeElimination::Context &);
 
static void
divertInStructuralNode(rvsdg::StructuralNode & node, CommonNodeElimination::Context & context)
{
  bool divertInSubregions = false;
  rvsdg::MatchTypeOrFail(
      node,
      [&]([[maybe_unused]] rvsdg::GammaNode & gamma)
      {
        divertInSubregions = true;
      },
      [&]([[maybe_unused]] rvsdg::ThetaNode & theta)
      {
        divertInSubregions = true;
      },
      [&]([[maybe_unused]] rvsdg::LambdaNode & lambda)
      {
        divertInSubregions = true;
      },
      [&]([[maybe_unused]] rvsdg::PhiNode & phi)
      {
        divertInSubregions = true;
      },
      [&]([[maybe_unused]] rvsdg::DeltaNode & delta)
      {
        // Inside a delta node we can not perform diverting,
        // since we never marked the outputs there
      });
 
  if (divertInSubregions)
  {
    for (auto & subregion : node.Subregions())
    {
      divertInRegion(subregion, context);
    }
  }
}
 
static void
divertInRegion(rvsdg::Region & region, CommonNodeElimination::Context & context)
{
  // First divert all region arguments
  for (auto argument : region.Arguments())
  {
    divertOutput(*argument, context);
  }
 
  // Divert all nodes
  for (const auto node : rvsdg::TopDownTraverser(&region))
  {
    // When diverting structural nodes, also recurse into their subregions
    rvsdg::MatchType(
        *node,
        [&](rvsdg::StructuralNode & structural)
        {
          divertInStructuralNode(structural, context);
        });
 
    // Divert all outputs of the node
    for (auto & output : node->Outputs())
    {
      divertOutput(output, context);
    }
  }
 
  region.prune(false);
}
 
CommonNodeElimination::~CommonNodeElimination() noexcept = default;
 
void
CommonNodeElimination::Run(
    rvsdg::RvsdgModule & module,
    util::StatisticsCollector & statisticsCollector)
{
  auto & rvsdg = module.Rvsdg();
  auto & rootRegion = rvsdg.GetRootRegion();
 
  Context context;
  auto statistics = Statistics::Create(module.SourceFilePath().value());
 
  statistics->startMarkStatistics(rvsdg);
  markGraphImports(rootRegion, context);
  markRegion(rootRegion, context);
  statistics->endMarkStatistics();
 
  statistics->startDivertStatistics();
  divertInRegion(rootRegion, context);
  statistics->endDivertStatistics(rvsdg);
 
  // Pruning nodes as an extra stage instead of at the end of divertInRegion() was a deliberate
  // choice. The divertInRegion() method traverses the region tree to the innermost regions first
  // before handling the outermost region, which means we would prune in the innermost regions
  // before pruning in the outermost regions. However, it could be that CNE renders a structural
  // node dead. Thus, this innermost-first strategy would mean that we prune in the dead structural
  // node's subregion before removing the structural node itself, which is not desirable. In
  // contrast, the Graph::PruneNodes() method prunes the outermost regions first before traversing
  // into the subregions of the (leftover) structural nodes.
  statistics->startPruneStatistics();
  rvsdg.PruneNodes();
  statistics->stopPruneStatistics();
 
  statisticsCollector.CollectDemandedStatistics(std::move(statistics));
}
}