lib/Target/AMDGPU/AMDGPUOpenCLEnqueuedBlockLowering.cpp

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