instrument-ref.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/rhls2firrtl/base-hls.hpp>
#include <jlm/hls/backend/rvsdg2rhls/add-prints.hpp>
#include <jlm/hls/backend/rvsdg2rhls/instrument-ref.hpp>
#include <jlm/llvm/ir/CallingConvention.hpp>
#include <jlm/llvm/ir/operators/alloca.hpp>
#include <jlm/llvm/ir/operators/call.hpp>
#include <jlm/llvm/ir/operators/IntegerOperations.hpp>
#include <jlm/llvm/ir/operators/Load.hpp>
#include <jlm/rvsdg/gamma.hpp>
#include <jlm/rvsdg/traverser.hpp>
 
#include <cmath>
 
namespace jlm::hls
{
 
rvsdg::LambdaNode *
change_function_name(rvsdg::LambdaNode * ln, const std::string & name)
{
  const auto & op = dynamic_cast<llvm::LlvmLambdaOperation &>(ln->GetOperation());
  auto lambda = rvsdg::LambdaNode::Create(
      *ln->region(),
      llvm::LlvmLambdaOperation::Create(
          op.Type(),
          name,
          op.linkage(),
          op.callingConvention(),
          op.attributes()));
 
  /* add context variables */
  rvsdg::SubstitutionMap subregionmap;
  for (const auto & cv : ln->GetContextVars())
  {
    auto origin = cv.input->origin();
    auto newcv = lambda->AddContextVar(*origin);
    subregionmap.insert(cv.inner, newcv.inner);
  }
  /* collect function arguments */
  auto args = ln->GetFunctionArguments();
  auto newArgs = lambda->GetFunctionArguments();
  JLM_ASSERT(args.size() == newArgs.size());
  for (std::size_t n = 0; n < args.size(); ++n)
  {
    subregionmap.insert(args[n], newArgs[n]);
  }
 
  /* copy subregion */
  ln->subregion()->copy(lambda->subregion(), subregionmap);
 
  /* collect function results */
  std::vector<jlm::rvsdg::Output *> results;
  for (auto result : ln->GetFunctionResults())
    results.push_back(&subregionmap.lookup(*result->origin()));
 
  /* finalize lambda */
  lambda->finalize(results);
 
  divert_users(ln, outputs(lambda));
  jlm::rvsdg::remove(ln);
 
  return lambda;
}
 
void
instrument_ref(llvm::LlvmRvsdgModule & rm)
{
  auto & graph = rm.Rvsdg();
  auto root = &graph.GetRootRegion();
  auto lambda = dynamic_cast<rvsdg::LambdaNode *>(root->Nodes().begin().ptr());
 
  auto newLambda = change_function_name(lambda, "instrumented_ref");
 
  auto functionType = newLambda->GetOperation().type();
  auto numArguments = functionType.NumArguments();
  if (numArguments == 0)
  {
    // The lambda has no arguments so it shouldn't have any memory operations
    return;
  }
 
  auto memStateArgumentIndex = numArguments - 1;
  if (!rvsdg::is<llvm::MemoryStateType>(functionType.ArgumentType(memStateArgumentIndex)))
  {
    // The lambda has no memory state so it shouldn't have any memory operations
    return;
  }
  // The function should always have an IO state if it has a memory state
  auto ioStateArgumentIndex = numArguments - 2;
  JLM_ASSERT(rvsdg::is<llvm::IOStateType>(functionType.ArgumentType(ioStateArgumentIndex)));
 
  //  addr, width, memstate
  auto loadFunctionType = jlm::rvsdg::FunctionType::Create(
      { jlm::llvm::PointerType::Create(),
        jlm::rvsdg::BitType::Create(64),
        llvm::IOStateType::Create(),
        llvm::MemoryStateType::Create() },
      { llvm::IOStateType::Create(), llvm::MemoryStateType::Create() });
  auto & reference_load = llvm::LlvmGraphImport::createFunctionImport(
      graph,
      loadFunctionType,
      "reference_load",
      llvm::Linkage::externalLinkage,
      llvm::CallingConvention::Default);
  auto & reference_store = llvm::LlvmGraphImport::createFunctionImport(
      graph,
      loadFunctionType,
      "reference_store",
      llvm::Linkage::externalLinkage,
      llvm::CallingConvention::Default);
  // addr, size, memstate
  auto allocaFunctionType = jlm::rvsdg::FunctionType::Create(
      { jlm::llvm::PointerType::Create(),
        jlm::rvsdg::BitType::Create(64),
        llvm::IOStateType::Create(),
        jlm::llvm::MemoryStateType::Create() },
      { llvm::IOStateType::Create(), jlm::llvm::MemoryStateType::Create() });
  auto & reference_alloca = llvm::LlvmGraphImport::createFunctionImport(
      graph,
      allocaFunctionType,
      "reference_alloca",
      llvm::Linkage::externalLinkage,
      llvm::CallingConvention::Default);
 
  instrument_ref(
      root,
      newLambda->subregion()->argument(ioStateArgumentIndex),
      &reference_load,
      loadFunctionType,
      &reference_store,
      loadFunctionType,
      &reference_alloca,
      allocaFunctionType);
}
 
void
instrument_ref(
    rvsdg::Region * region,
    jlm::rvsdg::Output * ioState,
    jlm::rvsdg::Output * load_func,
    const std::shared_ptr<const jlm::rvsdg::FunctionType> & loadFunctionType,
    jlm::rvsdg::Output * store_func,
    const std::shared_ptr<const jlm::rvsdg::FunctionType> & storeFunctionType,
    jlm::rvsdg::Output * alloca_func,
    const std::shared_ptr<const jlm::rvsdg::FunctionType> & allocaFunctionType)
{
  load_func = &rvsdg::RouteToRegion(*load_func, *region);
  store_func = &rvsdg::RouteToRegion(*store_func, *region);
  alloca_func = &rvsdg::RouteToRegion(*alloca_func, *region);
  auto void_ptr = jlm::llvm::PointerType::Create();
  for (auto & node : rvsdg::TopDownTraverser(region))
  {
    if (auto structnode = dynamic_cast<rvsdg::StructuralNode *>(node))
    {
      for (size_t n = 0; n < structnode->nsubregions(); n++)
      {
        auto subregion = structnode->subregion(n);
        auto ioStateRouted = &rvsdg::RouteToRegion(*ioState, *subregion);
        instrument_ref(
            subregion,
            ioStateRouted,
            load_func,
            loadFunctionType,
            store_func,
            storeFunctionType,
            alloca_func,
            allocaFunctionType);
      }
    }
    else if (
        auto loadOp =
            dynamic_cast<const jlm::llvm::LoadNonVolatileOperation *>(&(node->GetOperation())))
    {
      auto addr = node->input(0)->origin();
      JLM_ASSERT(rvsdg::is<jlm::llvm::PointerType>(addr->Type()));
      size_t bitWidth = BaseHLS::JlmSize(&*loadOp->GetLoadedType());
      int log2Bytes = log2(bitWidth / 8);
      auto & widthNode = llvm::IntegerConstantOperation::Create(*region, 64, log2Bytes);
 
      // Does this IF make sense now when the void_ptr doesn't have a type?
      if (*addr->Type() != *void_ptr)
      {
        addr = llvm::BitCastOperation::create(addr, void_ptr);
      }
      auto memstate = node->input(1)->origin();
      auto callOp = jlm::llvm::CallOperation::Create(
          load_func,
          loadFunctionType,
          { addr, widthNode.output(0), ioState, memstate });
      // Divert the memory state of the load to the new memstate from the call operation
      node->input(1)->divert_to(callOp[1]);
    }
    else if (auto ao = dynamic_cast<const jlm::llvm::AllocaOperation *>(&(node->GetOperation())))
    {
      // ensure that the size is one
      JLM_ASSERT(node->ninputs() == 1);
      auto constant_output = dynamic_cast<rvsdg::NodeOutput *>(node->input(0)->origin());
      JLM_ASSERT(constant_output);
      auto constant_operation = dynamic_cast<const llvm::IntegerConstantOperation *>(
          &constant_output->node()->GetOperation());
      JLM_ASSERT(constant_operation);
      JLM_ASSERT(constant_operation->Representation().to_uint() == 1);
      jlm::rvsdg::Output * addr = node->output(0);
      // ensure that the alloca is an array type
      JLM_ASSERT(jlm::rvsdg::is<llvm::PointerType>(addr->Type()));
      auto at = dynamic_cast<const llvm::ArrayType *>(ao->allocatedType().get());
      JLM_ASSERT(at);
      auto & sizeNode =
          llvm::IntegerConstantOperation::Create(*region, 64, BaseHLS::JlmSize(at) / 8);
 
      // Does this IF make sense now when the void_ptr doesn't have a type?
      if (*addr->Type() != *void_ptr)
      {
        addr = llvm::BitCastOperation::create(addr, void_ptr);
      }
      std::vector<jlm::rvsdg::Input *> old_users;
      for (auto & user : node->output(1)->Users())
        old_users.push_back(&user);
      auto memstate = node->output(1);
      auto callOp = jlm::llvm::CallOperation::Create(
          alloca_func,
          allocaFunctionType,
          { addr, sizeNode.output(0), ioState, memstate });
      for (auto ou : old_users)
      {
        // Divert the memory state of the load to the new memstate from the call operation
        ou->divert_to(callOp[1]);
      }
    }
    else if (
        auto so =
            dynamic_cast<const jlm::llvm::StoreNonVolatileOperation *>(&(node->GetOperation())))
    {
      auto addr = node->input(0)->origin();
      JLM_ASSERT(rvsdg::is<jlm::llvm::PointerType>(addr->Type()));
      auto bitWidth = JlmSize(&so->GetStoredType());
      int log2Bytes = log2(bitWidth / 8);
      auto & widthNode = llvm::IntegerConstantOperation::Create(*region, 64, log2Bytes);
 
      // Does this IF make sense now when the void_ptr doesn't have a type?
      if (*addr->Type() != *void_ptr)
      {
        addr = llvm::BitCastOperation::create(addr, void_ptr);
      }
      auto memstate = node->output(0);
      std::vector<jlm::rvsdg::Input *> oldUsers;
      for (auto & user : memstate->Users())
        oldUsers.push_back(&user);
      auto callOp = jlm::llvm::CallOperation::Create(
          store_func,
          storeFunctionType,
          { addr, widthNode.output(0), ioState, memstate });
      // Divert the memory state after the store to the new memstate from the call operation
      for (auto user : oldUsers)
      {
        user->divert_to(callOp[1]);
      }
    }
  }
}
 
} // namespace jlm::hls