DFG2LLVM_OpenCL.cpp

//=== DFG2LLVM_OpenCL.cpp ===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//

#define ENABLE_ASSERTS
#define TARGET_PTX 64
#define GENERIC_ADDRSPACE 0
#define GLOBAL_ADDRSPACE 1
#define CONSTANT_ADDRSPACE 4
#define SHARED_ADDRSPACE 3

#define DEBUG_TYPE "DFG2LLVM_OpenCL"
#include "SupportHPVM/DFG2LLVM.h"
#include "SupportHPVM/HPVMTimer.h"
#include "SupportHPVM/HPVMUtils.h"
#include "llvm-c/Core.h"
#include "llvm/IR/Attributes.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Module.h"
#include "llvm/IRReader/IRReader.h"
#include "llvm/Linker/Linker.h"
#include "llvm/Pass.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/ValueMapper.h"

#include "llvm/IR/IRPrintingPasses.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/UseListOrder.h"
#include "llvm/Support/ToolOutputFile.h"

#include <sstream>

#ifndef LLVM_BUILD_DIR
#error LLVM_BUILD_DIR is not defined
#endif

#define STR_VALUE(X) #X
#define STRINGIFY(X) STR_VALUE(X)
#define LLVM_BUILD_DIR_STR STRINGIFY(LLVM_BUILD_DIR)

using namespace llvm;
using namespace builddfg;
using namespace dfg2llvm;
using namespace hpvmUtils;

// HPVM Command line option to use timer or not
static cl::opt<bool> HPVMTimer_OpenCL("hpvm-timers-ptx",
                                      cl::desc("Enable hpvm timers"));

namespace {
// Helper class declarations

// Class to maintain the tuple of host pointer, device pointer and size
// in bytes. Would have preferred to use tuple but support not yet available
class OutputPtr {
public:
  OutputPtr(Value *_h_ptr, Value *_d_ptr, Value *_bytes)
      : h_ptr(_h_ptr), d_ptr(_d_ptr), bytes(_bytes) {}

  Value *h_ptr;
  Value *d_ptr;
  Value *bytes;
};

// Class to maintain important kernel info required for generating runtime
// calls
class Kernel {
public:
  Kernel(
      Function *_KF, DFLeafNode *_KLeafNode,
      std::map<unsigned, unsigned> _inArgMap = std::map<unsigned, unsigned>(),
      std::map<unsigned, std::pair<Value *, unsigned>> _sharedInArgMap =
          std::map<unsigned, std::pair<Value *, unsigned>>(),
      std::vector<unsigned> _outArgMap = std::vector<unsigned>(),
      unsigned _gridDim = 0,
      std::vector<Value *> _globalWGSize = std::vector<Value *>(),
      unsigned _blockDim = 0,
      std::vector<Value *> _localWGSize = std::vector<Value *>())
      : KernelFunction(_KF), KernelLeafNode(_KLeafNode), inArgMap(_inArgMap),
        sharedInArgMap(_sharedInArgMap), outArgMap(_outArgMap),
        gridDim(_gridDim), globalWGSize(_globalWGSize), blockDim(_blockDim),
        localWGSize(_localWGSize) {

    assert(gridDim == globalWGSize.size() &&
           "gridDim should be same as the size of vector globalWGSize");
    assert(blockDim == localWGSize.size() &&
           "blockDim should be same as the size of vector localWGSize");
  }

  Function *KernelFunction;
  DFLeafNode *KernelLeafNode;
  std::map<unsigned, unsigned> inArgMap;
  // Map for shared memory arguments
  std::map<unsigned, std::pair<Value *, unsigned>> sharedInArgMap;
  // Fields for (potential) allocation node
  DFLeafNode *AllocationNode;
  Function *AllocationFunction;
  std::map<unsigned, unsigned> allocInArgMap;

  std::vector<unsigned> outArgMap;
  unsigned gridDim;
  std::vector<Value *> globalWGSize;
  unsigned blockDim;
  std::vector<Value *> localWGSize;
  std::vector<int> localDimMap;

  std::map<unsigned, unsigned> &getInArgMap() { return inArgMap; }
  void setInArgMap(std::map<unsigned, unsigned> map) { inArgMap = map; }

  std::map<unsigned, std::pair<Value *, unsigned>> &getSharedInArgMap() {
    return sharedInArgMap;
  }
  void setSharedInArgMap(std::map<unsigned, std::pair<Value *, unsigned>> map) {
    sharedInArgMap = map;
  }

  std::vector<unsigned> &getOutArgMap() { return outArgMap; }
  void setOutArgMap(std::vector<unsigned> map) { outArgMap = map; }

  void setLocalWGSize(std::vector<Value *> V) { localWGSize = V; }

  bool hasLocalWG() const { return blockDim != 0; }
};

// Helper function declarations
static bool canBePromoted(Argument *arg, Function *F);
static void getExecuteNodeParams(Module &M, Value *&, Value *&, Value *&,
                                 Kernel *, ValueToValueMapTy &, Instruction *);
static Value *genWorkGroupPtr(Module &M, std::vector<Value *>,
                              ValueToValueMapTy &, Instruction *,
                              const Twine &WGName = "WGSize");
static std::string getPTXFilename(const Module &);
static std::string getFilenameFromModule(const Module &M);
static void changeDataLayout(Module &);
static void changeTargetTriple(Module &);
static void findReturnInst(Function *, std::vector<ReturnInst *> &);
static void findIntrinsicInst(Function *, Intrinsic::ID,
                              std::vector<IntrinsicInst *> &);
static AtomicRMWInst::BinOp getAtomicOp(Intrinsic::ID);
static std::string getAtomicOpName(Intrinsic::ID);

// DFG2LLVM_OpenCL - The first implementation.
struct DFG2LLVM_OpenCL : public DFG2LLVM {
  static char ID; // Pass identification, replacement for typeid
  DFG2LLVM_OpenCL() : DFG2LLVM(ID) {}

private:
public:
  bool runOnModule(Module &M);
};

// Visitor for Code generation traversal (tree traversal for now)
class CGT_OpenCL : public CodeGenTraversal {

private:
  // Member variables
  std::unique_ptr<Module> KernelM;
  DFNode *KernelLaunchNode = NULL;
  Kernel *kernel;

  // HPVM Runtime API
  FunctionCallee llvm_hpvm_ocl_launch;
  FunctionCallee llvm_hpvm_ocl_wait;
  FunctionCallee llvm_hpvm_ocl_initContext;
  FunctionCallee llvm_hpvm_ocl_clearContext;
  FunctionCallee llvm_hpvm_ocl_argument_shared;
  FunctionCallee llvm_hpvm_ocl_argument_scalar;
  FunctionCallee llvm_hpvm_ocl_argument_ptr;
  FunctionCallee llvm_hpvm_ocl_output_ptr;
  FunctionCallee llvm_hpvm_ocl_free;
  FunctionCallee llvm_hpvm_ocl_getOutput;
  FunctionCallee llvm_hpvm_ocl_executeNode;

  // Functions
  std::string getKernelsModuleName(Module &M);
  void fixValueAddrspace(Value *V, unsigned addrspace);
  std::vector<unsigned> globalToConstantMemoryOpt(std::vector<unsigned> *,
                                                  Function *);
  Function *changeArgAddrspace(Function *F, std::vector<unsigned> &Ags,
                               unsigned i);
  void addCLMetadata(Function *F);
  Function *transformFunctionToVoid(Function *F);
  void insertRuntimeCalls(DFInternalNode *N, Kernel *K, const Twine &FileName);

  // Virtual Functions
  void init() {
    HPVMTimer = HPVMTimer_OpenCL;
    TargetName = "OpenCL";
  }
  void initRuntimeAPI();
  void codeGen(DFInternalNode *N);
  void codeGen(DFLeafNode *N);

public:
  // Constructor
  CGT_OpenCL(Module &_M, BuildDFG &_DFG)
      : CodeGenTraversal(_M, _DFG), KernelM(CloneModule(_M)) {
    init();
    initRuntimeAPI();
    DEBUG(errs() << "Old module pointer: " << &_M << "\n");
    DEBUG(errs() << "New module pointer: " << KernelM.get() << "\n");

    // Copying instead of creating new, in order to preserve required info
    // (metadata) Remove functions, global variables and aliases
    std::vector<GlobalVariable *> GVVect;
    for (Module::global_iterator mi = KernelM->global_begin(),
                                 me = KernelM->global_end();
         (mi != me); ++mi) {
      GlobalVariable *GV = &*mi;
      GVVect.push_back(GV);
    }
    for (auto *GV : GVVect) {
      GV->replaceAllUsesWith(UndefValue::get(GV->getType()));
      GV->eraseFromParent();
    }

    std::vector<Function *> FuncVect;
    for (Module::iterator mi = KernelM->begin(), me = KernelM->end();
         (mi != me); ++mi) {
      Function *F = &*mi;
      FuncVect.push_back(F);
    }
    for (auto *F : FuncVect) {
      F->replaceAllUsesWith(UndefValue::get(F->getType()));
      F->eraseFromParent();
    }

    std::vector<GlobalAlias *> GAVect;
    for (Module::alias_iterator mi = KernelM->alias_begin(),
                                me = KernelM->alias_end();
         (mi != me); ++mi) {
      GlobalAlias *GA = &*mi;
      GAVect.push_back(GA);
    }
    for (auto *GA : GAVect) {
      GA->replaceAllUsesWith(UndefValue::get(GA->getType()));
      GA->eraseFromParent();
    }

    changeDataLayout(*KernelM);
    changeTargetTriple(*KernelM);

    DEBUG(errs() << *KernelM);
  }

  void writeKernelsModule();
};

// Initialize the HPVM runtime API. This makes it easier to insert these calls
void CGT_OpenCL::initRuntimeAPI() {

  // Load Runtime API Module
  SMDiagnostic Err;

  std::string runtimeAPI = std::string(LLVM_BUILD_DIR_STR) +
                           "/tools/hpvm/projects/hpvm-rt/hpvm-rt.bc";

  runtimeModule = parseIRFile(runtimeAPI, Err, M.getContext());
  if (runtimeModule == nullptr) {
    DEBUG(errs() << Err.getMessage() << " " << runtimeAPI << "\n");
    assert(false && "couldn't parse runtime");
  } else
    DEBUG(errs() << "Successfully loaded hpvm-rt API module\n");

  // Get or insert the global declarations for launch/wait functions
  DECLARE(llvm_hpvm_ocl_launch);
  DECLARE(llvm_hpvm_ocl_wait);
  DECLARE(llvm_hpvm_ocl_initContext);
  DECLARE(llvm_hpvm_ocl_clearContext);
  DECLARE(llvm_hpvm_ocl_argument_shared);
  DECLARE(llvm_hpvm_ocl_argument_scalar);
  DECLARE(llvm_hpvm_ocl_argument_ptr);
  DECLARE(llvm_hpvm_ocl_output_ptr);
  DECLARE(llvm_hpvm_ocl_free);
  DECLARE(llvm_hpvm_ocl_getOutput);
  DECLARE(llvm_hpvm_ocl_executeNode);

  // Get or insert timerAPI functions as well if you plan to use timers
  initTimerAPI();

  // Insert init context in main
  DEBUG(errs() << "Gen Code to initialize OpenCL Timer\n");
  Function *VI = M.getFunction("llvm.hpvm.init");
  assert(VI->getNumUses() == 1 && "__hpvm__init should only be used once");

  InitCall = cast<Instruction>(*VI->user_begin());
  initializeTimerSet(InitCall);
  switchToTimer(hpvm_TimerID_INIT_CTX, InitCall);
  CallInst::Create(llvm_hpvm_ocl_initContext,
                   ArrayRef<Value *>(getTargetID(M, hpvm::GPU_TARGET)), "",
                   InitCall);
  switchToTimer(hpvm_TimerID_NONE, InitCall);

  // Insert print instruction at hpvm exit
  DEBUG(errs() << "Gen Code to print OpenCL Timer\n");
  Function *VC = M.getFunction("llvm.hpvm.cleanup");
  DEBUG(errs() << *VC << "\n");
  assert(VC->getNumUses() == 1 && "__hpvm__clear should only be used once");

  CleanupCall = cast<Instruction>(*VC->user_begin());
  printTimerSet(CleanupCall);
}

// Generate Code to call the kernel
// The plan is to replace the internal node with a leaf node. This method is
// used to generate a function to associate with this leaf node. The function
// is responsible for all the memory allocation/transfer and invoking the
// kernel call on the device
void CGT_OpenCL::insertRuntimeCalls(DFInternalNode *N, Kernel *K,
                                    const Twine &FileName) {
  // Check if clone already exists. If it does, it means we have visited this
  // function before.
  //  assert(N->getGenFunc() == NULL && "Code already generated for this node");

  assert(N->getGenFuncForTarget(hpvm::GPU_TARGET) == NULL &&
         "Code already generated for this node");

  // Useful values
  Value *True = ConstantInt::get(Type::getInt1Ty(M.getContext()), 1);
  Value *False = ConstantInt::get(Type::getInt1Ty(M.getContext()), 0);

  // If kernel struct has not been initialized with kernel function, then fail
  assert(K != NULL && "No kernel found!!");

  DEBUG(errs() << "Generating kernel call code\n");

  Function *F = N->getFuncPointer();

  // Create of clone of F with no instructions. Only the type is the same as F
  // without the extra arguments.
  Function *F_CPU;

  // Clone the function, if we are seeing this function for the first time. We
  // only need a clone in terms of type.
  ValueToValueMapTy VMap;

  // Create new function with the same type
  F_CPU =
      Function::Create(F->getFunctionType(), F->getLinkage(), F->getName(), &M);

  // Loop over the arguments, copying the names of arguments over.
  Function::arg_iterator dest_iterator = F_CPU->arg_begin();
  for (Function::const_arg_iterator i = F->arg_begin(), e = F->arg_end();
       i != e; ++i) {
    dest_iterator->setName(i->getName()); // Copy the name over...
    // Increment dest iterator
    ++dest_iterator;
  }

  // Add a basic block to this empty function
  BasicBlock *BB = BasicBlock::Create(M.getContext(), "entry", F_CPU);
  ReturnInst *RI = ReturnInst::Create(
      M.getContext(), UndefValue::get(F_CPU->getReturnType()), BB);

  // FIXME: Adding Index and Dim arguments are probably not required except
  // for consistency purpose (DFG2LLVM_CPU does assume that all leaf nodes do
  // have those arguments)

  // Add Index and Dim arguments except for the root node
  if (!N->isRoot() && !N->getParent()->isChildGraphStreaming())
    F_CPU = addIdxDimArgs(F_CPU);

  BB = &*F_CPU->begin();
  RI = cast<ReturnInst>(BB->getTerminator());

  // Add the generated function info to DFNode
  //  N->setGenFunc(F_CPU, hpvm::CPU_TARGET);
  N->addGenFunc(F_CPU, hpvm::GPU_TARGET, true);
  DEBUG(errs() << "Added GPUGenFunc: " << F_CPU->getName() << " for node "
               << N->getFuncPointer()->getName() << "\n");

  // Loop over the arguments, to create the VMap
  dest_iterator = F_CPU->arg_begin();
  for (Function::const_arg_iterator i = F->arg_begin(), e = F->arg_end();
       i != e; ++i) {
    // Add mapping to VMap and increment dest iterator
    VMap[&*i] = &*dest_iterator;
    ++dest_iterator;
  }

  /* TODO: Use this code to verufy if this is a good pattern for PTX kernel

  // Sort children in topological order before code generation for kernel call
  N->getChildGraph()->sortChildren();

  // The DFNode N has the property that it has only one child (leaving Entry
  // and Exit dummy nodes). This child is the PTX kernel. This simplifies code
  // generation for kernel calls significantly. All the inputs to this child
  // node would either be constants or from the parent node N.

  assert(N->getChildGraph()->size() == 3
         && "Node expected to have just one non-dummy node!");

  DFNode* C;
  for(DFGraph::children_iterator ci = N->getChildGraph()->begin(),
      ce = N->getChildGraph()->end(); ci != ce; ++ci) {
    C = *ci;
    // Skip dummy node call
    if (!C->isDummyNode())
      break;
  }

  assert(C->isDummyNode() == false && "Internal Node only contains dummy
  nodes!");

  Function* CF = C->getFuncPointer();
  */
  Function *KF = K->KernelLeafNode->getFuncPointer();
  // Initialize context
  // DEBUG(errs() << "Initializing context" << "\n");
  // CallInst::Create(llvm_hpvm_ocl_initContext, None, "", RI);

  DEBUG(errs() << "Initializing commandQ"
               << "\n");
  // Initialize command queue
  switchToTimer(hpvm_TimerID_SETUP, InitCall);
  Value *fileStr = getStringPointer(FileName, InitCall, "Filename");
  DEBUG(errs() << "Kernel Filename constant: " << *fileStr << "\n");
  DEBUG(errs() << "Generating code for kernel - "
               << K->KernelFunction->getName() << "\n");
  Value *kernelStr =
      getStringPointer(K->KernelFunction->getName(), InitCall, "KernelName");

  Value *LaunchInstArgs[] = {fileStr, kernelStr};

  DEBUG(errs() << "Inserting launch call"
               << "\n");
  CallInst *OpenCL_Ctx = CallInst::Create(llvm_hpvm_ocl_launch,
                                          ArrayRef<Value *>(LaunchInstArgs, 2),
                                          "graph" + KF->getName(), InitCall);
  DEBUG(errs() << *OpenCL_Ctx << "\n");
  GraphIDAddr = new GlobalVariable(
      M, OpenCL_Ctx->getType(), false, GlobalValue::CommonLinkage,
      Constant::getNullValue(OpenCL_Ctx->getType()),
      "graph" + KF->getName() + ".addr");
  DEBUG(errs() << "Store at: " << *GraphIDAddr << "\n");
  StoreInst *SI = new StoreInst(OpenCL_Ctx, GraphIDAddr, InitCall);
  DEBUG(errs() << *SI << "\n");
  switchToTimer(hpvm_TimerID_NONE, InitCall);
  switchToTimer(hpvm_TimerID_SETUP, RI);
  Value *GraphID = new LoadInst(GraphIDAddr, "graph." + KF->getName(), RI);

  // Iterate over the required input edges of the node and use the hpvm-rt API
  // to set inputs
  DEBUG(errs() << "Iterate over input edges of node and insert hpvm api\n");
  std::vector<OutputPtr> OutputPointers;
  // Vector to hold the device memory object that need to be cleared before we
  // release context
  std::vector<Value *> DevicePointers;

  std::map<unsigned, unsigned> &kernelInArgMap = K->getInArgMap();
  /*
    for(unsigned i=0; i<KF->getFunctionType()->getNumParams(); i++) {

      // The kernel object gives us the mapping of arguments from kernel launch
      // node function (F_CPU) to kernel (kernel->KF)
      Value* inputVal = getArgumentAt(F_CPU, K->getInArgMap()[i]);

  */

  for (auto &InArgMapPair : kernelInArgMap) {
    unsigned i = InArgMapPair.first;
    Value *inputVal = getArgumentAt(F_CPU, InArgMapPair.second);
    DEBUG(errs() << "\tArgument " << i << " = " << *inputVal << "\n");

    // input value has been obtained.
    // Check if input is a scalar value or a pointer operand
    // For scalar values such as int, float, etc. the size is simply the size of
    // type on target machine, but for pointers, the size of data would be the
    // next integer argument
    if (inputVal->getType()->isPointerTy()) {

      switchToTimer(hpvm_TimerID_COPY_PTR, RI);
      // Pointer Input
      // CheckAttribute
      Value *isOutput = (hasAttribute(KF, i, Attribute::Out)) ? True : False;
      Value *isInput = ((hasAttribute(KF, i, Attribute::Out)) &&
                        !(hasAttribute(KF, i, Attribute::In)))
                           ? False
                           : True;

      Argument *A = getArgumentAt(KF, i);
      if (isOutput == True) {
        DEBUG(errs() << *A << " is an OUTPUT argument\n");
      }
      if (isInput == True) {
        DEBUG(errs() << *A << " is an INPUT argument\n");
      }

      Value *inputValI8Ptr = CastInst::CreatePointerCast(
          inputVal, Type::getInt8PtrTy(M.getContext()),
          inputVal->getName() + ".i8ptr", RI);

      // Assert that the pointer argument size (next argument) is in the map
      assert(kernelInArgMap.find(i + 1) != kernelInArgMap.end());

      Value *inputSize = getArgumentAt(F_CPU, kernelInArgMap[i + 1]);
      assert(
          inputSize->getType() == Type::getInt64Ty(M.getContext()) &&
          "Pointer type input must always be followed by size (integer type)");
      Value *setInputArgs[] = {
          GraphID,
          inputValI8Ptr,
          ConstantInt::get(Type::getInt32Ty(M.getContext()), i),
          inputSize,
          isInput,
          isOutput};
      Value *d_ptr =
          CallInst::Create(llvm_hpvm_ocl_argument_ptr,
                           ArrayRef<Value *>(setInputArgs, 6), "", RI);
      DevicePointers.push_back(d_ptr);
      // If this has out attribute, store the returned device pointer in
      // memory to read device memory later
      if (isOutput == True)
        OutputPointers.push_back(OutputPtr(inputValI8Ptr, d_ptr, inputSize));
    } else {
      switchToTimer(hpvm_TimerID_COPY_SCALAR, RI);
      // Scalar Input
      // Store the scalar value on stack and then pass the pointer to its
      // location
      AllocaInst *inputValPtr = new AllocaInst(
          inputVal->getType(), 0, inputVal->getName() + ".ptr", RI);
      StoreInst *SI = new StoreInst(inputVal, inputValPtr, RI);

      Value *inputValI8Ptr = CastInst::CreatePointerCast(
          inputValPtr, Type::getInt8PtrTy(M.getContext()),
          inputVal->getName() + ".i8ptr", RI);

      Value *setInputArgs[] = {
          GraphID, inputValI8Ptr,
          ConstantInt::get(Type::getInt32Ty(M.getContext()), i),
          ConstantExpr::getSizeOf(inputVal->getType())};
      CallInst::Create(llvm_hpvm_ocl_argument_scalar,
                       ArrayRef<Value *>(setInputArgs, 4), "", RI);
    }
  }

  DEBUG(
      errs() << "Setup shared memory arguments of node and insert hpvm api\n");

  // Check to see if all the allocation sizes are constant (determined
  // statically)
  bool constSizes = true;
  for (auto &e : K->getSharedInArgMap()) {
    constSizes &= isa<Constant>(e.second.first);
  }

  // If the sizes are all constant
  if (constSizes) {
    for (auto &e : K->getSharedInArgMap()) {
      unsigned argNum = e.first;
      Value *allocSize = e.second.first;

      DEBUG(errs() << "\tLocal Memory at " << argNum
                   << ", size = " << *allocSize << "\n");

      if (KF->getFunctionType()->getParamType(argNum)->isPointerTy()) {
        // Shared memory ptr argument - scalar at size position
        switchToTimer(hpvm_TimerID_COPY_SCALAR, RI);

        assert(isa<Constant>(allocSize) &&
               "Constant shared memory size is expected");

        Value *setInputArgs[] = {
            GraphID, ConstantInt::get(Type::getInt32Ty(M.getContext()), argNum),
            allocSize};
        CallInst::Create(llvm_hpvm_ocl_argument_shared,
                         ArrayRef<Value *>(setInputArgs, 3), "", RI);
      } else {
        // Sharem memory size argument - scalar at address position
        switchToTimer(hpvm_TimerID_COPY_SCALAR, RI);
        // Store the scalar value on stack and then pass the pointer to its
        // location
        AllocaInst *allocSizePtr =
            new AllocaInst(allocSize->getType(), 0,
                           allocSize->getName() + ".sharedMem.ptr", RI);
        StoreInst *SI = new StoreInst(allocSize, allocSizePtr, RI);

        Value *allocSizeI8Ptr = CastInst::CreatePointerCast(
            allocSizePtr, Type::getInt8PtrTy(M.getContext()),
            allocSize->getName() + ".sharedMem.i8ptr", RI);

        Value *setInputArgs[] = {
            GraphID, allocSizeI8Ptr,
            ConstantInt::get(Type::getInt32Ty(M.getContext()), argNum),
            ConstantExpr::getSizeOf(allocSize->getType())};
        CallInst::Create(llvm_hpvm_ocl_argument_scalar,
                         ArrayRef<Value *>(setInputArgs, 4), "", RI);
      }
    }
  } else {

    Function *F_alloc = K->AllocationFunction;
    StructType *FAllocRetTy = dyn_cast<StructType>(F_alloc->getReturnType());
    assert(FAllocRetTy && "Allocation node with no struct return type");

    std::vector<Value *> AllocInputArgs;
    for (unsigned i = 0; i < K->allocInArgMap.size(); i++) {
      AllocInputArgs.push_back(getArgumentAt(F_CPU, K->allocInArgMap.at(i)));
    }

    CallInst *CI = CallInst::Create(F_alloc, AllocInputArgs, "", RI);
    std::vector<ExtractValueInst *> ExtractValueInstVec;
    for (unsigned i = 1; i < FAllocRetTy->getNumElements(); i += 2) {
      ExtractValueInst *EI = ExtractValueInst::Create(CI, i, "", RI);
      ExtractValueInstVec.push_back(EI);
    }

    for (auto &e : K->getSharedInArgMap()) {
      unsigned argNum = e.first;
      Value *allocSize = ExtractValueInstVec[e.second.second / 2];

      DEBUG(errs() << "\tLocal Memory at " << argNum
                   << ", size = " << *allocSize << "\n");

      if (KF->getFunctionType()->getParamType(argNum)->isPointerTy()) {
        // Shared memory ptr argument - scalar at size position
        switchToTimer(hpvm_TimerID_COPY_SCALAR, RI);

        Value *setInputArgs[] = {
            GraphID, ConstantInt::get(Type::getInt32Ty(M.getContext()), argNum),
            allocSize};
        CallInst::Create(llvm_hpvm_ocl_argument_shared,
                         ArrayRef<Value *>(setInputArgs, 3), "", RI);
      } else {
        // Sharem memory size argument - scalar at address position
        switchToTimer(hpvm_TimerID_COPY_SCALAR, RI);
        // Store the scalar value on stack and then pass the pointer to its
        // location
        AllocaInst *allocSizePtr =
            new AllocaInst(allocSize->getType(), 0,
                           allocSize->getName() + ".sharedMem.ptr", RI);
        StoreInst *SI = new StoreInst(allocSize, allocSizePtr, RI);

        Value *allocSizeI8Ptr = CastInst::CreatePointerCast(
            allocSizePtr, Type::getInt8PtrTy(M.getContext()),
            allocSize->getName() + ".sharedMem.i8ptr", RI);

        Value *setInputArgs[] = {
            GraphID, allocSizeI8Ptr,
            ConstantInt::get(Type::getInt32Ty(M.getContext()), argNum),
            ConstantExpr::getSizeOf(allocSize->getType())};
        CallInst::Create(llvm_hpvm_ocl_argument_scalar,
                         ArrayRef<Value *>(setInputArgs, 4), "", RI);
      }
    }
  }

  DEBUG(errs() << "Setup output edges of node and insert hpvm api\n");
  // Set output if struct is not an empty struct
  StructType *OutputTy = K->KernelLeafNode->getOutputType();
  std::vector<Value *> d_Outputs;
  if (!OutputTy->isEmptyTy()) {
    switchToTimer(hpvm_TimerID_COPY_PTR, RI);
    // Not an empty struct
    // Iterate over all elements of the struct and put them in
    for (unsigned i = 0; i < OutputTy->getNumElements(); i++) {
      unsigned outputIndex = KF->getFunctionType()->getNumParams() + i;
      Value *setOutputArgs[] = {
          GraphID,
          ConstantInt::get(Type::getInt32Ty(M.getContext()), outputIndex),
          ConstantExpr::getSizeOf(OutputTy->getElementType(i))};

      CallInst *d_Output = CallInst::Create(llvm_hpvm_ocl_output_ptr,
                                            ArrayRef<Value *>(setOutputArgs, 3),
                                            "d_output." + KF->getName(), RI);
      d_Outputs.push_back(d_Output);
    }
  }

  // Enqueue kernel
  // Need work dim, localworksize, globalworksize
  // Allocate size_t[numDims] space on stack. Store the work group sizes and
  // pass it as an argument to ExecNode

  switchToTimer(hpvm_TimerID_MISC, RI);
  Value *workDim, *LocalWGPtr, *GlobalWGPtr;
  getExecuteNodeParams(M, workDim, LocalWGPtr, GlobalWGPtr, K, VMap, RI);
  switchToTimer(hpvm_TimerID_KERNEL, RI);
  Value *ExecNodeArgs[] = {GraphID, workDim, LocalWGPtr, GlobalWGPtr};
  CallInst *Event = CallInst::Create(llvm_hpvm_ocl_executeNode,
                                     ArrayRef<Value *>(ExecNodeArgs, 4),
                                     "event." + KF->getName(), RI);
  DEBUG(errs() << "Execute Node Call: " << *Event << "\n");

  // Wait for Kernel to Finish
  CallInst::Create(llvm_hpvm_ocl_wait, ArrayRef<Value *>(GraphID), "", RI);

  switchToTimer(hpvm_TimerID_READ_OUTPUT, RI);
  // Read Output Struct if not empty
  if (!OutputTy->isEmptyTy()) {
    std::vector<Value *> h_Outputs;
    Value *KernelOutput = UndefValue::get(OutputTy);
    for (unsigned i = 0; i < OutputTy->getNumElements(); i++) {
      Value *GetOutputArgs[] = {
          GraphID, Constant::getNullValue(Type::getInt8PtrTy(M.getContext())),
          d_Outputs[i], ConstantExpr::getSizeOf(OutputTy->getElementType(i))};
      CallInst *h_Output = CallInst::Create(
          llvm_hpvm_ocl_getOutput, ArrayRef<Value *>(GetOutputArgs, 4),
          "h_output." + KF->getName() + ".addr", RI);
      // Read each device pointer listed in output struct
      // Load the output struct
      CastInst *BI = BitCastInst::CreatePointerCast(
          h_Output, OutputTy->getElementType(i)->getPointerTo(), "output.ptr",
          RI);

      Value *OutputElement = new LoadInst(BI, "output." + KF->getName(), RI);
      KernelOutput = InsertValueInst::Create(KernelOutput, OutputElement,
                                             ArrayRef<unsigned>(i),
                                             KF->getName() + "output", RI);
    }
    OutputMap[K->KernelLeafNode] = KernelOutput;
  }

  // Read all the pointer arguments which had side effects i.e., had out
  // attribute
  DEBUG(errs() << "Output Pointers : " << OutputPointers.size() << "\n");
  // FIXME: Not reading output pointers anymore as we read them when data is
  // actually requested
  /*for(auto output: OutputPointers) {
    DEBUG(errs() << "Read: " << *output.d_ptr << "\n");
    DEBUG(errs() << "\tTo: " << *output.h_ptr << "\n");
    DEBUG(errs() << "\t#bytes: " << *output.bytes << "\n");

    Value* GetOutputArgs[] = {GraphID, output.h_ptr, output.d_ptr,
  output.bytes}; CallInst* CI = CallInst::Create(llvm_hpvm_ocl_getOutput,
                                    ArrayRef<Value*>(GetOutputArgs, 4),
                                    "", RI);
  }*/
  switchToTimer(hpvm_TimerID_MEM_FREE, RI);
  // Clear Context and free device memory
  DEBUG(errs() << "Clearing context"
               << "\n");
  // Free Device Memory
  for (auto d_ptr : DevicePointers) {
    CallInst::Create(llvm_hpvm_ocl_free, ArrayRef<Value *>(d_ptr), "", RI);
  }
  switchToTimer(hpvm_TimerID_CLEAR_CTX, CleanupCall);
  // Clear Context
  LoadInst *LI = new LoadInst(GraphIDAddr, "", CleanupCall);
  CallInst::Create(llvm_hpvm_ocl_clearContext, ArrayRef<Value *>(LI), "",
                   CleanupCall);
  switchToTimer(hpvm_TimerID_NONE, CleanupCall);

  switchToTimer(hpvm_TimerID_MISC, RI);
  DEBUG(errs() << "*** Generating epilogue code for the function****\n");
  // Generate code for output bindings
  // Get Exit node
  DFNode *C = N->getChildGraph()->getExit();
  // Get OutputType of this node
  StructType *OutTy = N->getOutputType();
  Value *retVal = UndefValue::get(F_CPU->getReturnType());
  // Find the kernel's output arg map, to use instead of the bindings
  std::vector<unsigned> outArgMap = kernel->getOutArgMap();
  // Find all the input edges to exit node
  for (unsigned i = 0; i < OutTy->getNumElements(); i++) {
    DEBUG(errs() << "Output Edge " << i << "\n");
    // Find the incoming edge at the requested input port
    DFEdge *E = C->getInDFEdgeAt(i);

    assert(E && "No Binding for output element!");
    // Find the Source DFNode associated with the incoming edge
    DFNode *SrcDF = E->getSourceDF();

    DEBUG(errs() << "Edge source -- " << SrcDF->getFuncPointer()->getName()
                 << "\n");

    // If Source DFNode is a dummyNode, edge is from parent. Get the
    // argument from argument list of this internal node
    Value *inputVal;
    if (SrcDF->isEntryNode()) {
      inputVal = getArgumentAt(F_CPU, i);
      DEBUG(errs() << "Argument " << i << " = " << *inputVal << "\n");
    } else {
      // edge is from a internal node
      // Check - code should already be generated for this source dfnode
      // FIXME: Since the 2-level kernel code gen has aspecific structure, we
      // can assume the SrcDF is same as Kernel Leaf node.
      // Use outArgMap to get correct mapping
      SrcDF = K->KernelLeafNode;
      assert(OutputMap.count(SrcDF) &&
             "Source node call not found. Dependency violation!");

      // Find Output Value associated with the Source DFNode using OutputMap
      Value *CI = OutputMap[SrcDF];

      // Extract element at source position from this call instruction
      std::vector<unsigned> IndexList;
      // i is the destination of DFEdge E
      // Use the mapping instead of the bindings
      //      IndexList.push_back(E->getSourcePosition());
      IndexList.push_back(outArgMap[i]);
      DEBUG(errs() << "Going to generate ExtarctVal inst from " << *CI << "\n");
      ExtractValueInst *EI = ExtractValueInst::Create(CI, IndexList, "", RI);
      inputVal = EI;
    }
    std::vector<unsigned> IdxList;
    IdxList.push_back(i);
    retVal = InsertValueInst::Create(retVal, inputVal, IdxList, "", RI);
  }

  DEBUG(errs() << "Extracted all\n");
  switchToTimer(hpvm_TimerID_NONE, RI);
  retVal->setName("output");
  ReturnInst *newRI = ReturnInst::Create(F_CPU->getContext(), retVal);
  ReplaceInstWithInst(RI, newRI);
}

// Right now, only targeting the one level case. In general, device functions
// can return values so we don't need to change them
void CGT_OpenCL::codeGen(DFInternalNode *N) {
  DEBUG(errs() << "Inside internal node: " << N->getFuncPointer()->getName()
               << "\n");
  if (KernelLaunchNode == NULL)
    DEBUG(errs() << "No kernel launch node\n");
  else {
    DEBUG(errs() << "KernelLaunchNode: "
                 << KernelLaunchNode->getFuncPointer()->getName() << "\n");
  }

  if (!KernelLaunchNode) {
    DEBUG(errs()
          << "No code generated (host code for kernel launch complete).\n");
    return;
  }

  if (N == KernelLaunchNode) {
    DEBUG(errs() << "Found kernel launch node. Generating host code.\n");
    // TODO

    // Now the remaining nodes to be visited should be ignored
    KernelLaunchNode = NULL;
    DEBUG(errs() << "Insert Runtime calls\n");
    insertRuntimeCalls(N, kernel, getPTXFilename(M));

  } else {
    DEBUG(errs() << "Found intermediate node. Getting size parameters.\n");
    // Keep track of the arguments order.
    std::map<unsigned, unsigned> inmap1 = N->getInArgMap();
    std::map<unsigned, unsigned> inmap2 = kernel->getInArgMap();
    // TODO: Structure assumed: one thread node, one allocation node (at most),
    // TB node
    std::map<unsigned, unsigned> inmapFinal;