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/SmallString.h"
37 #include "llvm/IR/Instructions.h"
38 #include "llvm/IR/Mangler.h"
39 #include "llvm/IR/Module.h"
40 #include "llvm/Pass.h"
41 #include "llvm/Support/Debug.h"
42 
43 #define DEBUG_TYPE "amdgpu-lower-enqueued-block"
44 
45 using namespace llvm;
46 
47 namespace {
48 
49 /// Lower enqueued blocks.
50 class AMDGPUOpenCLEnqueuedBlockLowering : public ModulePass {
51 public:
52   static char ID;
53 
AMDGPUOpenCLEnqueuedBlockLowering()54   explicit AMDGPUOpenCLEnqueuedBlockLowering() : ModulePass(ID) {}
55 
56 private:
57   bool runOnModule(Module &M) override;
58 };
59 
60 } // end anonymous namespace
61 
62 char AMDGPUOpenCLEnqueuedBlockLowering::ID = 0;
63 
64 char &llvm::AMDGPUOpenCLEnqueuedBlockLoweringID =
65     AMDGPUOpenCLEnqueuedBlockLowering::ID;
66 
67 INITIALIZE_PASS(AMDGPUOpenCLEnqueuedBlockLowering, DEBUG_TYPE,
68                 "Lower OpenCL enqueued blocks", false, false)
69 
createAMDGPUOpenCLEnqueuedBlockLoweringPass()70 ModulePass* llvm::createAMDGPUOpenCLEnqueuedBlockLoweringPass() {
71   return new AMDGPUOpenCLEnqueuedBlockLowering();
72 }
73 
74 /// Collect direct or indrect callers of \p F and save them
75 /// to \p Callers.
collectCallers(Function * F,DenseSet<Function * > & Callers)76 static void collectCallers(Function *F, DenseSet<Function *> &Callers) {
77   for (auto U : F->users()) {
78     if (auto *CI = dyn_cast<CallInst>(&*U)) {
79       auto *Caller = CI->getParent()->getParent();
80       if (Callers.insert(Caller).second)
81         collectCallers(Caller, Callers);
82     }
83   }
84 }
85 
86 /// If \p U is instruction or constant, collect functions which directly or
87 /// indirectly use it.
collectFunctionUsers(User * U,DenseSet<Function * > & Funcs)88 static void collectFunctionUsers(User *U, DenseSet<Function *> &Funcs) {
89   if (auto *I = dyn_cast<Instruction>(U)) {
90     auto *F = I->getParent()->getParent();
91     if (Funcs.insert(F).second)
92       collectCallers(F, Funcs);
93     return;
94   }
95   if (!isa<Constant>(U))
96     return;
97   for (auto UU : U->users())
98     collectFunctionUsers(&*UU, Funcs);
99 }
100 
runOnModule(Module & M)101 bool AMDGPUOpenCLEnqueuedBlockLowering::runOnModule(Module &M) {
102   DenseSet<Function *> Callers;
103   auto &C = M.getContext();
104   bool Changed = false;
105   for (auto &F : M.functions()) {
106     if (F.hasFnAttribute("enqueued-block")) {
107       if (!F.hasName()) {
108         SmallString<64> Name;
109         Mangler::getNameWithPrefix(Name, "__amdgpu_enqueued_kernel",
110                                    M.getDataLayout());
111         F.setName(Name);
112       }
113       LLVM_DEBUG(dbgs() << "found enqueued kernel: " << F.getName() << '\n');
114       auto RuntimeHandle = (F.getName() + ".runtime_handle").str();
115       auto T = ArrayType::get(Type::getInt64Ty(C), 2);
116       auto *GV = new GlobalVariable(
117           M, T,
118           /*isConstant=*/false, GlobalValue::ExternalLinkage,
119           /*Initializer=*/Constant::getNullValue(T), RuntimeHandle,
120           /*InsertBefore=*/nullptr, GlobalValue::NotThreadLocal,
121           AMDGPUAS::GLOBAL_ADDRESS,
122           /*isExternallyInitialized=*/false);
123       LLVM_DEBUG(dbgs() << "runtime handle created: " << *GV << '\n');
124 
125       for (auto U : F.users()) {
126         auto *UU = &*U;
127         if (!isa<ConstantExpr>(UU))
128           continue;
129         collectFunctionUsers(UU, Callers);
130         auto *BitCast = cast<ConstantExpr>(UU);
131         auto *NewPtr = ConstantExpr::getPointerCast(GV, BitCast->getType());
132         BitCast->replaceAllUsesWith(NewPtr);
133         F.addFnAttr("runtime-handle", RuntimeHandle);
134         F.setLinkage(GlobalValue::ExternalLinkage);
135         Changed = true;
136       }
137     }
138   }
139 
140   for (auto F : Callers) {
141     if (F->getCallingConv() != CallingConv::AMDGPU_KERNEL)
142       continue;
143     F->addFnAttr("calls-enqueue-kernel");
144     LLVM_DEBUG(dbgs() << "mark enqueue_kernel caller:" << F->getName() << '\n');
145   }
146   return Changed;
147 }
148