1//===- NVVMReflect.cpp - NVVM Emulate conditional compilation -------------===//
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// This pass replaces occurrences of __nvvm_reflect("foo") and llvm.nvvm.reflect
10// with an integer.
11//
12// We choose the value we use by looking at metadata in the module itself. Note
13// that we intentionally only have one way to choose these values, because other
14// parts of LLVM (particularly, InstCombineCall) rely on being able to predict
15// the values chosen by this pass.
16//
17// If we see an unknown string, we replace its call with 0.
18//
19//===----------------------------------------------------------------------===//
20
21#include "NVPTX.h"
22#include "llvm/ADT/SmallVector.h"
23#include "llvm/ADT/StringExtras.h"
24#include "llvm/Analysis/ConstantFolding.h"
25#include "llvm/CodeGen/CommandFlags.h"
26#include "llvm/IR/Constants.h"
27#include "llvm/IR/DerivedTypes.h"
28#include "llvm/IR/Function.h"
29#include "llvm/IR/Instructions.h"
30#include "llvm/IR/Intrinsics.h"
31#include "llvm/IR/IntrinsicsNVPTX.h"
32#include "llvm/IR/Module.h"
33#include "llvm/IR/PassManager.h"
34#include "llvm/IR/Type.h"
35#include "llvm/Pass.h"
36#include "llvm/Support/CommandLine.h"
37#include "llvm/Support/Debug.h"
38#include "llvm/Support/raw_ostream.h"
39#include "llvm/Transforms/Scalar.h"
40#include "llvm/Transforms/Utils/BasicBlockUtils.h"
41#include "llvm/Transforms/Utils/Local.h"
42
43using namespace llvm;
44
45#define DEBUG_TYPE "nvvm-reflect"
46
47#define NVVM_REFLECT_FUNCTION "__nvvm_reflect"
48#define NVVM_REFLECT_OCL_FUNCTION "__nvvm_reflect_ocl"
49// Argument of reflect call to retrive arch number.
50#define CUDA_ARCH_NAME "__CUDA_ARCH"
51// Argument of reflect call to retrive ftz mode.
52#define CUDA_FTZ_NAME "__CUDA_FTZ"
53// Name of module metadata where ftz mode is stored.
54#define CUDA_FTZ_MODULE_NAME "nvvm-reflect-ftz"
55
56namespace {
57class NVVMReflect {
58 // Map from reflect function call arguments to the value to replace the call
59 // with. Should include __CUDA_FTZ and __CUDA_ARCH values.
60 StringMap<unsigned> ReflectMap;
61 bool handleReflectFunction(Module &M, StringRef ReflectName);
62 void populateReflectMap(Module &M);
63 void foldReflectCall(CallInst *Call, Constant *NewValue);
64
65public:
66 // __CUDA_FTZ is assigned in `runOnModule` by checking nvvm-reflect-ftz module
67 // metadata.
68 explicit NVVMReflect(unsigned SmVersion)
69 : ReflectMap({{CUDA_ARCH_NAME, SmVersion * 10}}) {}
70 bool runOnModule(Module &M);
71};
72
73class NVVMReflectLegacyPass : public ModulePass {
74 NVVMReflect Impl;
75
76public:
77 static char ID;
78 NVVMReflectLegacyPass(unsigned SmVersion) : ModulePass(ID), Impl(SmVersion) {}
79 bool runOnModule(Module &M) override;
80};
81} // namespace
82
83ModulePass *llvm::createNVVMReflectPass(unsigned SmVersion) {
84 return new NVVMReflectLegacyPass(SmVersion);
85}
86
87static cl::opt<bool>
88 NVVMReflectEnabled("nvvm-reflect-enable", cl::init(Val: true), cl::Hidden,
89 cl::desc("NVVM reflection, enabled by default"));
90
91char NVVMReflectLegacyPass::ID = 0;
92INITIALIZE_PASS(NVVMReflectLegacyPass, "nvvm-reflect",
93 "Replace occurrences of __nvvm_reflect() calls with 0/1", false,
94 false)
95
96// Allow users to specify additional key/value pairs to reflect. These key/value
97// pairs are the last to be added to the ReflectMap, and therefore will take
98// precedence over initial values (i.e. __CUDA_FTZ from module medadata and
99// __CUDA_ARCH from SmVersion).
100static cl::list<std::string> ReflectList(
101 "nvvm-reflect-add", cl::value_desc("name=<int>"), cl::Hidden,
102 cl::desc("A key=value pair. Replace __nvvm_reflect(name) with value."),
103 cl::ValueRequired);
104
105// Set the ReflectMap with, first, the value of __CUDA_FTZ from module metadata,
106// and then the key/value pairs from the command line.
107void NVVMReflect::populateReflectMap(Module &M) {
108 if (auto *Flag = mdconst::extract_or_null<ConstantInt>(
109 MD: M.getModuleFlag(CUDA_FTZ_MODULE_NAME)))
110 ReflectMap[CUDA_FTZ_NAME] = Flag->getSExtValue();
111
112 for (StringRef Option : ReflectList) {
113 LLVM_DEBUG(dbgs() << "ReflectOption : " << Option << "\n");
114 auto [Name, Val] = Option.split(Separator: '=');
115 if (Name.empty())
116 report_fatal_error(reason: "Empty name in nvvm-reflect-add option '" + Option +
117 "'");
118 if (Val.empty())
119 report_fatal_error(reason: "Missing value in nvvm-reflect-add option '" + Option +
120 "'");
121 unsigned ValInt;
122 if (!to_integer(S: Val.trim(), Num&: ValInt, Base: 10))
123 report_fatal_error(reason: "integer value expected in nvvm-reflect-add option '" +
124 Option + "'");
125 ReflectMap[Name] = ValInt;
126 }
127}
128
129/// Process a reflect function by finding all its calls and replacing them with
130/// appropriate constant values. For __CUDA_FTZ, uses the module flag value.
131/// For __CUDA_ARCH, uses SmVersion * 10. For all other strings, uses 0.
132bool NVVMReflect::handleReflectFunction(Module &M, StringRef ReflectName) {
133 Function *F = M.getFunction(Name: ReflectName);
134 if (!F)
135 return false;
136 assert(F->isDeclaration() && "_reflect function should not have a body");
137 assert(F->getReturnType()->isIntegerTy() &&
138 "_reflect's return type should be integer");
139
140 const bool Changed = !F->use_empty();
141 for (User *U : make_early_inc_range(Range: F->users())) {
142 // Reflect function calls look like:
143 // @arch = private unnamed_addr addrspace(1) constant [12 x i8]
144 // c"__CUDA_ARCH\00" call i32 @__nvvm_reflect(ptr addrspacecast (ptr
145 // addrspace(1) @arch to ptr)) We need to extract the string argument from
146 // the call (i.e. "__CUDA_ARCH")
147 auto *Call = dyn_cast<CallInst>(Val: U);
148 if (!Call)
149 report_fatal_error(
150 reason: "__nvvm_reflect can only be used in a call instruction");
151 if (Call->getNumOperands() != 2)
152 report_fatal_error(reason: "__nvvm_reflect requires exactly one argument");
153
154 auto *GlobalStr =
155 dyn_cast<Constant>(Val: Call->getArgOperand(i: 0)->stripPointerCasts());
156 if (!GlobalStr)
157 report_fatal_error(reason: "__nvvm_reflect argument must be a constant string");
158
159 auto *ConstantStr =
160 dyn_cast<ConstantDataSequential>(Val: GlobalStr->getOperand(i: 0));
161 if (!ConstantStr)
162 report_fatal_error(reason: "__nvvm_reflect argument must be a string constant");
163 if (!ConstantStr->isCString())
164 report_fatal_error(
165 reason: "__nvvm_reflect argument must be a null-terminated string");
166
167 StringRef ReflectArg = ConstantStr->getAsString().drop_back();
168 if (ReflectArg.empty())
169 report_fatal_error(reason: "__nvvm_reflect argument cannot be empty");
170 // Now that we have extracted the string argument, we can look it up in the
171 // ReflectMap
172 unsigned ReflectVal = 0; // The default value is 0
173 if (ReflectMap.contains(Key: ReflectArg))
174 ReflectVal = ReflectMap[ReflectArg];
175
176 LLVM_DEBUG(dbgs() << "Replacing call of reflect function " << F->getName()
177 << "(" << ReflectArg << ") with value " << ReflectVal
178 << "\n");
179 auto *NewValue = ConstantInt::get(Ty: Call->getType(), V: ReflectVal);
180 foldReflectCall(Call, NewValue);
181 Call->eraseFromParent();
182 }
183
184 // Remove the __nvvm_reflect function from the module
185 F->eraseFromParent();
186 return Changed;
187}
188
189void NVVMReflect::foldReflectCall(CallInst *Call, Constant *NewValue) {
190 SmallVector<Instruction *, 8> Worklist;
191 // Replace an instruction with a constant and add all users of the instruction
192 // to the worklist
193 auto ReplaceInstructionWithConst = [&](Instruction *I, Constant *C) {
194 for (User *U : I->users())
195 if (auto *UI = dyn_cast<Instruction>(Val: U))
196 Worklist.push_back(Elt: UI);
197 I->replaceAllUsesWith(V: C);
198 };
199
200 ReplaceInstructionWithConst(Call, NewValue);
201
202 auto &DL = Call->getModule()->getDataLayout();
203 while (!Worklist.empty()) {
204 auto *I = Worklist.pop_back_val();
205 if (Constant *C = ConstantFoldInstruction(I, DL)) {
206 ReplaceInstructionWithConst(I, C);
207 if (isInstructionTriviallyDead(I))
208 I->eraseFromParent();
209 } else if (I->isTerminator()) {
210 ConstantFoldTerminator(BB: I->getParent());
211 }
212 }
213}
214
215bool NVVMReflect::runOnModule(Module &M) {
216 if (!NVVMReflectEnabled)
217 return false;
218 populateReflectMap(M);
219 bool Changed = true;
220 Changed |= handleReflectFunction(M, NVVM_REFLECT_FUNCTION);
221 Changed |= handleReflectFunction(M, NVVM_REFLECT_OCL_FUNCTION);
222 Changed |=
223 handleReflectFunction(M, ReflectName: Intrinsic::getName(id: Intrinsic::nvvm_reflect));
224 return Changed;
225}
226
227bool NVVMReflectLegacyPass::runOnModule(Module &M) {
228 return Impl.runOnModule(M);
229}
230
231PreservedAnalyses NVVMReflectPass::run(Module &M, ModuleAnalysisManager &AM) {
232 return NVVMReflect(SmVersion).runOnModule(M) ? PreservedAnalyses::none()
233 : PreservedAnalyses::all();
234}
235