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
45using namespace llvm;
46
47namespace {
48
49/// Lower enqueued blocks.
50class AMDGPUOpenCLEnqueuedBlockLowering : public ModulePass {
51public:
52  static char ID;
53
54  explicit AMDGPUOpenCLEnqueuedBlockLowering() : ModulePass(ID) {}
55
56private:
57  bool runOnModule(Module &M) override;
58};
59
60} // end anonymous namespace
61
62char AMDGPUOpenCLEnqueuedBlockLowering::ID = 0;
63
64char &llvm::AMDGPUOpenCLEnqueuedBlockLoweringID =
65    AMDGPUOpenCLEnqueuedBlockLowering::ID;
66
67INITIALIZE_PASS(AMDGPUOpenCLEnqueuedBlockLowering, DEBUG_TYPE,
68                "Lower OpenCL enqueued blocks", false, false)
69
70ModulePass* llvm::createAMDGPUOpenCLEnqueuedBlockLoweringPass() {
71  return new AMDGPUOpenCLEnqueuedBlockLowering();
72}
73
74/// Collect direct or indrect callers of \p F and save them
75/// to \p Callers.
76static 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.
88static 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
101bool 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