mem-sep.cpp

Go to the documentation of this file.

/*
 * Copyright 2021 David Metz <david.c.metz@ntnu.no>
 * See COPYING for terms of redistribution.
 */
 
#include <jlm/hls/backend/rvsdg2rhls/add-prints.hpp>
#include <jlm/hls/backend/rvsdg2rhls/hls-function-util.hpp>
#include <jlm/hls/backend/rvsdg2rhls/mem-conv.hpp>
#include <jlm/hls/backend/rvsdg2rhls/mem-sep.hpp>
#include <jlm/hls/ir/hls.hpp>
#include <jlm/llvm/ir/LambdaMemoryState.hpp>
#include <jlm/llvm/ir/operators/call.hpp>
#include <jlm/llvm/ir/operators/lambda.hpp>
#include <jlm/llvm/ir/operators/Load.hpp>
#include <jlm/llvm/ir/operators/MemoryStateOperations.hpp>
#include <jlm/llvm/ir/operators/Store.hpp>
#include <jlm/rvsdg/gamma.hpp>
#include <jlm/rvsdg/substitution.hpp>
#include <jlm/rvsdg/theta.hpp>
#include <jlm/rvsdg/traverser.hpp>
#include <jlm/rvsdg/view.hpp>
 
#include <algorithm>
 
namespace jlm::hls
{
 
static rvsdg::RegionResult *
trace_edge(
    jlm::rvsdg::Output * common_edge,
    jlm::rvsdg::Output * new_edge,
    std::vector<rvsdg::Node *> & load_nodes,
    const std::vector<rvsdg::Node *> & store_nodes,
    std::vector<rvsdg::Node *> & decouple_nodes)
{
  // follows along common edge and routes new edge through the same regions
  // redirects the supplied loads, stores and decouples to the new edge
  // the new edge might be routed through unnecessary regions. This should be fixed by running DNE
  while (true)
  {
    // each iteration should update common_edge and/or new_edge
    JLM_ASSERT(common_edge->nusers() == 1);
    JLM_ASSERT(new_edge->nusers() == 1);
    auto & user = *common_edge->Users().begin();
    auto & new_next = *new_edge->Users().begin();
    if (auto res = dynamic_cast<rvsdg::RegionResult *>(&user))
    {
      // end of region reached
      return res;
    }
    else if (auto gammaNode = rvsdg::TryGetOwnerNode<rvsdg::GammaNode>(user))
    {
      auto ip = gammaNode->AddEntryVar(new_edge);
      std::vector<jlm::rvsdg::Output *> vec;
      new_edge = gammaNode->AddExitVar(ip.branchArgument).output;
      new_next.divert_to(new_edge);
 
      auto rolevar = gammaNode->MapInput(user);
 
      if (auto entryvar = std::get_if<rvsdg::GammaNode::EntryVar>(&rolevar))
      {
        for (size_t i = 0; i < gammaNode->nsubregions(); ++i)
        {
          auto subres = trace_edge(
              entryvar->branchArgument[i],
              ip.branchArgument[i],
              load_nodes,
              store_nodes,
              decouple_nodes);
          common_edge = subres->output();
        }
      }
    }
    else if (auto theta = rvsdg::TryGetOwnerNode<rvsdg::ThetaNode>(user))
    {
      auto olv = theta->MapInputLoopVar(user);
      auto lv = theta->AddLoopVar(new_edge);
      trace_edge(olv.pre, lv.pre, load_nodes, store_nodes, decouple_nodes);
      common_edge = olv.output;
      new_edge = lv.output;
      new_next.divert_to(new_edge);
    }
    else if (auto sn = rvsdg::TryGetOwnerNode<rvsdg::SimpleNode>(user))
    {
      auto op = &sn->GetOperation();
      if (dynamic_cast<const jlm::llvm::StoreNonVolatileOperation *>(op))
      {
        JLM_ASSERT(sn->noutputs() == 1);
        if (store_nodes.end() != std::find(store_nodes.begin(), store_nodes.end(), sn))
        {
          user.divert_to(new_edge);
          sn->output(0)->divert_users(common_edge);
          new_edge = sn->output(0);
          new_next.divert_to(new_edge);
        }
        else
        {
          common_edge = sn->output(0);
        }
      }
      else if (dynamic_cast<const jlm::llvm::LoadNonVolatileOperation *>(op))
      {
        JLM_ASSERT(sn->noutputs() == 2);
        if (load_nodes.end() != std::find(load_nodes.begin(), load_nodes.end(), sn))
        {
          auto & new_next = *new_edge->Users().begin();
          user.divert_to(new_edge);
          sn->output(1)->divert_users(common_edge);
          new_next.divert_to(sn->output(1));
          new_edge = sn->output(1);
          new_next.divert_to(new_edge);
        }
        else
        {
          common_edge = sn->output(1);
        }
      }
      else if (dynamic_cast<const jlm::llvm::CallOperation *>(op))
      {
        int oi = sn->noutputs() - sn->ninputs() + user.index();
        // TODO: verify this is the right type of function call
        if (decouple_nodes.end() != std::find(decouple_nodes.begin(), decouple_nodes.end(), sn))
        {
          auto & new_next = *new_edge->Users().begin();
          user.divert_to(new_edge);
          sn->output(oi)->divert_users(common_edge);
          new_next.divert_to(sn->output(oi));
          new_edge = new_next.origin();
        }
        else
        {
          common_edge = sn->output(oi);
        }
      }
      else
      {
        JLM_ASSERT(sn->noutputs() == 1);
        common_edge = sn->output(0);
      }
    }
    else
    {
      JLM_UNREACHABLE("THIS SHOULD NOT HAPPEN");
    }
  }
}
 
std::vector<rvsdg::Node *>
MemoryStateSeparation::gatherNonDecoupleCalls(rvsdg::Region & region)
{
  std::function<void(rvsdg::Region &, std::vector<rvsdg::Node *> &)> gatherCalls =
      [&gatherCalls](rvsdg::Region & region, std::vector<rvsdg::Node *> & calls)
  {
    for (auto node : rvsdg::TopDownTraverser(&region))
    {
      // Handle innermost regions first
      if (const auto structuralNode = dynamic_cast<rvsdg::StructuralNode *>(node))
      {
        for (auto & subregion : structuralNode->Subregions())
        {
          gatherCalls(subregion, calls);
        }
      }
 
      if (rvsdg::is<llvm::CallOperation>(node))
      {
        auto functionName = get_function_name(node->input(0));
        if (functionName.rfind("decouple") == functionName.npos)
        {
          calls.push_back(node);
        }
      }
    }
  };
 
  std::vector<rvsdg::Node *> calls;
  gatherCalls(region, calls);
  return calls;
}
 
void
MemoryStateSeparation::separateMemoryStates(const rvsdg::LambdaNode & lambdaNode)
{
  const auto lambdaSubregion = lambdaNode.subregion();
  auto & memoryStateArgument = llvm::GetMemoryStateRegionArgument(lambdaNode);
 
  auto tracedPointerNodesVector = TracePointerArguments(&lambdaNode);
  for (auto & tp : tracedPointerNodesVector)
  {
    auto & decouple_nodes = tp.decoupleNodes;
    auto decouple_requests_cnt = decouple_nodes.size();
    // place decouple responses along same state edge
    for (size_t i = 0; i < decouple_requests_cnt; ++i)
    {
      auto req = decouple_nodes[i];
      auto channel = req->input(1)->origin();
      auto channel_constant = jlm::hls::trace_constant(channel);
      auto decouple_response = find_decouple_response(&lambdaNode, channel_constant);
      decouple_nodes.push_back(decouple_response);
    }
  }
 
  // Create fake ports for non-decouple calls
  const auto nonDecoupleCalls = gatherNonDecoupleCalls(*lambdaSubregion);
  for (auto call : nonDecoupleCalls)
  {
    tracedPointerNodesVector.emplace_back();
    tracedPointerNodesVector.back().decoupleNodes.push_back(call);
  }
 
  const size_t numMemoryStates = tracedPointerNodesVector.size() + 1;
 
  // Assign memory node ids incrementally, used by both the split and merge
  std::vector<llvm::MemoryNodeId> memoryNodeIds;
  for (size_t i = 0; i < numMemoryStates; ++i)
  {
    memoryNodeIds.push_back(i);
  }
 
  auto & lambdaEntrySplitNode =
      llvm::LambdaEntryMemoryStateSplitOperation::CreateNode(memoryStateArgument, memoryNodeIds);
  auto memoryStates = outputs(&lambdaEntrySplitNode);
 
  // handle existing state edge - TODO: remove entirely?
  // The memory state chain that was previously between the state argument and state result,
  // should now be attached as a chain between the final state on the split and merge, respectively
  auto common_edge = memoryStates.back();
  // Divert everything except the lambdaEntrySplit itself, to the lambdaEntrySplit
  memoryStateArgument.divertUsersWhere(
      *common_edge,
      [&](const rvsdg::Input & input)
      {
        return &input != lambdaEntrySplitNode.input(0);
      });
 
  auto state_result = &llvm::GetMemoryStateRegionResult(lambdaNode);
  memoryStates.back() = state_result->origin();
 
  auto & lambdaExitMergeNode = llvm::LambdaExitMemoryStateMergeOperation::CreateNode(
      *lambdaSubregion,
      memoryStates,
      memoryNodeIds);
  memoryStates.pop_back();
  state_result->divert_to(lambdaExitMergeNode.output(0));
 
  for (auto tp : tracedPointerNodesVector)
  {
    auto new_edge = memoryStates.back();
    memoryStates.pop_back();
    trace_edge(common_edge, new_edge, tp.loadNodes, tp.storeNodes, tp.decoupleNodes);
  }
}
 
MemoryStateSeparation::~MemoryStateSeparation() noexcept = default;
 
MemoryStateSeparation::MemoryStateSeparation()
    : Transformation("MemoryStateSeparation")
{}
 
void
MemoryStateSeparation::Run(rvsdg::RvsdgModule & rvsdgModule, util::StatisticsCollector &)
{
  const auto & graph = rvsdgModule.Rvsdg();
  const auto rootRegion = &graph.GetRootRegion();
  if (rootRegion->numNodes() != 1)
  {
    throw std::logic_error("Root should have only one node now");
  }
 
  const auto lambdaNode =
      dynamic_cast<const rvsdg::LambdaNode *>(rootRegion->Nodes().begin().ptr());
  if (!lambdaNode)
  {
    throw std::logic_error("Node needs to be a lambda");
  }
 
  separateMemoryStates(*lambdaNode);
}
 
} // namespace jlm::hls