contrib/llvm/lib/Target/AMDGPU/AMDGPUOpenCLEnqueuedBlockLowering.cpp

   1 //===- AMDGPUOpenCLEnqueuedBlockLowering.cpp - Lower enqueued block -------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // \file
  11 // This post-linking pass replaces the function pointer of enqueued
  12 // block kernel with a global variable (runtime handle) and adds
  13 // "runtime-handle" attribute to the enqueued block kernel.
  14 //
  15 // In LLVM CodeGen the runtime-handle metadata will be translated to
  16 // RuntimeHandle metadata in code object. Runtime allocates a global buffer
  17 // for each kernel with RuntimeHandel metadata and saves the kernel address
  18 // required for the AQL packet into the buffer. __enqueue_kernel function
  19 // in device library knows that the invoke function pointer in the block
  20 // literal is actually runtime handle and loads the kernel address from it
  21 // and put it into AQL packet for dispatching.
  22 //
  23 // This cannot be done in FE since FE cannot create a unique global variable
  24 // with external linkage across LLVM modules. The global variable with internal
  25 // linkage does not work since optimization passes will try to replace loads
  26 // of the global variable with its initialization value.
  27 //
  28 // It also identifies the kernels directly or indirectly enqueues kernels
  29 // and adds "calls-enqueue-kernel" function attribute to them, which will
  30 // be used to determine whether to emit runtime metadata for the kernel
  31 // enqueue related hidden kernel arguments.
  32 //
  33 //===----------------------------------------------------------------------===//
  34
  35 #include "AMDGPU.h"
  36 #include "llvm/ADT/DenseSet.h"
  37 #include "llvm/ADT/StringRef.h"
  38 #include "llvm/IR/Constants.h"
  39 #include "llvm/IR/DerivedTypes.h"
  40 #include "llvm/IR/Instructions.h"
  41 #include "llvm/IR/Mangler.h"
  42 #include "llvm/IR/Module.h"
  43 #include "llvm/IR/User.h"
  44 #include "llvm/Pass.h"
  45 #include "llvm/Support/Debug.h"
  46 #include "llvm/Support/raw_ostream.h"
  47
  48 #define DEBUG_TYPE "amdgpu-lower-enqueued-block"
  49
  50 using namespace llvm;
  51
  52 namespace {
  53
  54 /// Lower enqueued blocks.
  55 class AMDGPUOpenCLEnqueuedBlockLowering : public ModulePass {
  56 public:
  57   static char ID;
  58
  59   explicit AMDGPUOpenCLEnqueuedBlockLowering() : ModulePass(ID) {}
  60
  61 private:
  62   bool runOnModule(Module &M) override;
  63 };
  64
  65 } // end anonymous namespace
  66
  67 char AMDGPUOpenCLEnqueuedBlockLowering::ID = 0;
  68
  69 char &llvm::AMDGPUOpenCLEnqueuedBlockLoweringID =
  70     AMDGPUOpenCLEnqueuedBlockLowering::ID;
  71
  72 INITIALIZE_PASS(AMDGPUOpenCLEnqueuedBlockLowering, DEBUG_TYPE,
  73                 "Lower OpenCL enqueued blocks", false, false)
  74
  75 ModulePass* llvm::createAMDGPUOpenCLEnqueuedBlockLoweringPass() {
  76   return new AMDGPUOpenCLEnqueuedBlockLowering();
  77 }
  78
  79 /// Collect direct or indrect callers of \p F and save them
  80 /// to \p Callers.
  81 static void collectCallers(Function *F, DenseSet<Function *> &Callers) {
  82   for (auto U : F->users()) {
  83     if (auto *CI = dyn_cast<CallInst>(&*U)) {
  84       auto *Caller = CI->getParent()->getParent();
  85       if (Callers.insert(Caller).second)
  86         collectCallers(Caller, Callers);
  87     }
  88   }
  89 }
  90
  91 /// If \p U is instruction or constant, collect functions which directly or
  92 /// indirectly use it.
  93 static void collectFunctionUsers(User *U, DenseSet<Function *> &Funcs) {
  94   if (auto *I = dyn_cast<Instruction>(U)) {
  95     auto *F = I->getParent()->getParent();
  96     if (Funcs.insert(F).second)
  97       collectCallers(F, Funcs);
  98     return;
  99   }
 100   if (!isa<Constant>(U))
 101     return;
 102   for (auto UU : U->users())
 103     collectFunctionUsers(&*UU, Funcs);
 104 }
 105
 106 bool AMDGPUOpenCLEnqueuedBlockLowering::runOnModule(Module &M) {
 107   DenseSet<Function *> Callers;
 108   auto &C = M.getContext();
 109   bool Changed = false;
 110   for (auto &F : M.functions()) {
 111     if (F.hasFnAttribute("enqueued-block")) {
 112       if (!F.hasName()) {
 113         SmallString<64> Name;
 114         Mangler::getNameWithPrefix(Name, "__amdgpu_enqueued_kernel",
 115                                    M.getDataLayout());
 116         F.setName(Name);
 117       }
 118       LLVM_DEBUG(dbgs() << "found enqueued kernel: " << F.getName() << '\n');
 119       auto RuntimeHandle = (F.getName() + ".runtime_handle").str();
 120       auto T = ArrayType::get(Type::getInt64Ty(C), 2);
 121       auto *GV = new GlobalVariable(
 122           M, T,
 123           /*IsConstant=*/false, GlobalValue::ExternalLinkage,
 124           /*Initializer=*/Constant::getNullValue(T), RuntimeHandle,
 125           /*InsertBefore=*/nullptr, GlobalValue::NotThreadLocal,
 126           AMDGPUAS::GLOBAL_ADDRESS,
 127           /*IsExternallyInitialized=*/false);
 128       LLVM_DEBUG(dbgs() << "runtime handle created: " << *GV << '\n');
 129
 130       for (auto U : F.users()) {
 131         auto *UU = &*U;
 132         if (!isa<ConstantExpr>(UU))
 133           continue;
 134         collectFunctionUsers(UU, Callers);
 135         auto *BitCast = cast<ConstantExpr>(UU);
 136         auto *NewPtr = ConstantExpr::getPointerCast(GV, BitCast->getType());
 137         BitCast->replaceAllUsesWith(NewPtr);
 138         F.addFnAttr("runtime-handle", RuntimeHandle);
 139         F.setLinkage(GlobalValue::ExternalLinkage);
 140         Changed = true;
 141       }
 142     }
 143   }
 144
 145   for (auto F : Callers) {
 146     if (F->getCallingConv() != CallingConv::AMDGPU_KERNEL)
 147       continue;
 148     F->addFnAttr("calls-enqueue-kernel");
 149     LLVM_DEBUG(dbgs() << "mark enqueue_kernel caller:" << F->getName() << '\n');
 150   }
 151   return Changed;
 152 }