final reflect cleanup

Author: YonahGoldberg

llvm/lib/Target/NVPTX/NVPTX.h

@@ -78,12 +78,12 @@ struct NVVMIntrRangePass : PassInfoMixin<NVVMIntrRangePass> {
 };
 
 struct NVVMReflectPass : PassInfoMixin<NVVMReflectPass> {
-  NVVMReflectPass() : NVVMReflectPass(0) {}
-  NVVMReflectPass(unsigned SmVersion);
+  NVVMReflectPass() : SmVersion(0) {}
+  NVVMReflectPass(unsigned SmVersion) : SmVersion(SmVersion) {}
   PreservedAnalyses run(Module &F, ModuleAnalysisManager &AM);
 
 private:
-  StringMap<int> VarMap;
+  unsigned SmVersion;
 };
 
 struct GenericToNVVMPass : PassInfoMixin<GenericToNVVMPass> {

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

@@ -87,6 +87,27 @@ static cl::opt<bool> EarlyByValArgsCopy(
     cl::desc("Create a copy of byval function arguments early."),
     cl::init(false), cl::Hidden);
 
+namespace llvm {
+
+void initializeGenericToNVVMLegacyPassPass(PassRegistry &);
+void initializeNVPTXAllocaHoistingPass(PassRegistry &);
+void initializeNVPTXAssignValidGlobalNamesPass(PassRegistry &);
+void initializeNVPTXAtomicLowerPass(PassRegistry &);
+void initializeNVPTXCtorDtorLoweringLegacyPass(PassRegistry &);
+void initializeNVPTXLowerAggrCopiesPass(PassRegistry &);
+void initializeNVPTXLowerAllocaPass(PassRegistry &);
+void initializeNVPTXLowerUnreachablePass(PassRegistry &);
+void initializeNVPTXCtorDtorLoweringLegacyPass(PassRegistry &);
+void initializeNVPTXLowerArgsLegacyPassPass(PassRegistry &);
+void initializeNVPTXProxyRegErasurePass(PassRegistry &);
+void initializeNVPTXForwardParamsPassPass(PassRegistry &);
+void initializeNVVMIntrRangePass(PassRegistry &);
+void initializeNVVMReflectLegacyPassPass(PassRegistry &);
+void initializeNVPTXAAWrapperPassPass(PassRegistry &);
+void initializeNVPTXExternalAAWrapperPass(PassRegistry &);
+
+} // end namespace llvm
+
 extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeNVPTXTarget() {
   // Register the target.
   RegisterTargetMachine<NVPTXTargetMachine32> X(getTheNVPTXTarget32());
@@ -95,7 +116,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeNVPTXTarget() {
   PassRegistry &PR = *PassRegistry::getPassRegistry();
   // FIXME: This pass is really intended to be invoked during IR optimization,
   // but it's very NVPTX-specific.
-  initializeNVVMReflectPass(PR);
+  initializeNVVMReflectLegacyPassPass(PR);
   initializeNVVMIntrRangePass(PR);
   initializeGenericToNVVMLegacyPassPass(PR);
   initializeNVPTXAllocaHoistingPass(PR);

llvm/lib/Target/NVPTX/NVVMReflect.cpp

@@ -41,44 +41,60 @@
 #include "llvm/ADT/StringExtras.h"
 #define NVVM_REFLECT_FUNCTION "__nvvm_reflect"
 #define NVVM_REFLECT_OCL_FUNCTION "__nvvm_reflect_ocl"
+// Argument of reflect call to retrive arch number
+#define CUDA_ARCH_NAME "__CUDA_ARCH"
+// Argument of reflect call to retrive ftz mode
+#define CUDA_FTZ_NAME "__CUDA_FTZ"
+// Name of module metadata where ftz mode is stored
+#define CUDA_FTZ_MODULE_NAME "nvvm-reflect-ftz"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "nvvm-reflect"
 
+namespace llvm {
+void initializeNVVMReflectLegacyPassPass(PassRegistry &);
+}
+
 namespace {
-class NVVMReflect : public ModulePass {
+class NVVMReflect {
 private:
-  StringMap<int> VarMap;
-  void handleReflectFunction(Function *F);
-  void setVarMap(Module &M);
+  // Map from reflect function call arguments to the value to replace the call with.
+  // Should include __CUDA_FTZ and __CUDA_ARCH values.
+  StringMap<int> ReflectMap;
+  bool handleReflectFunction(Module &M, StringRef ReflectName);
+  void populateReflectMap(Module &M);
   void foldReflectCall(CallInst *Call, Constant *NewValue);
 public:
-  static char ID;
-  NVVMReflect() : NVVMReflect(0) {}
   // __CUDA_FTZ is assigned in `runOnModule` by checking nvvm-reflect-ftz module
   // metadata.
-  explicit NVVMReflect(unsigned SmVersion) : ModulePass(ID), VarMap({{"__CUDA_ARCH", SmVersion * 10}}) {
-    initializeNVVMReflectPass(*PassRegistry::getPassRegistry());
-  }
-  // This mapping will contain should include __CUDA_FTZ and __CUDA_ARCH values.
-  explicit NVVMReflect(const StringMap<int> &Mapping) : ModulePass(ID), VarMap(Mapping) {
-    initializeNVVMReflectPass(*PassRegistry::getPassRegistry());
+  explicit NVVMReflect(unsigned SmVersion) : ReflectMap({{CUDA_ARCH_NAME, SmVersion * 10}}) {}
+  bool runOnModule(Module &M);
+};
+} // namespace
+
+class NVVMReflectLegacyPass : public ModulePass {
+private:
+  NVVMReflect Impl;
+public:
+  static char ID;
+    NVVMReflectLegacyPass(unsigned SmVersion) : ModulePass(ID), Impl(SmVersion) {
+    initializeNVVMReflectLegacyPassPass(*PassRegistry::getPassRegistry());
   }
   bool runOnModule(Module &M) override;
 };
-} // namespace
 
 ModulePass *llvm::createNVVMReflectPass(unsigned int SmVersion) {
-  return new NVVMReflect(SmVersion);
+  LLVM_DEBUG(dbgs() << "Creating NVVMReflectPass with SM version " << SmVersion << "\n");
+  return new NVVMReflectLegacyPass(SmVersion);
 }
 
 static cl::opt<bool>
     NVVMReflectEnabled("nvvm-reflect-enable", cl::init(true), cl::Hidden,
                        cl::desc("NVVM reflection, enabled by default"));
 
-char NVVMReflect::ID = 0;
-INITIALIZE_PASS(NVVMReflect, "nvvm-reflect",
+char NVVMReflectLegacyPass::ID = 0;
+INITIALIZE_PASS(NVVMReflectLegacyPass, "nvvm-reflect",
                 "Replace occurrences of __nvvm_reflect() calls with 0/1", false,
                 false)
 
@@ -92,10 +108,10 @@ ReflectList("nvvm-reflect-add", cl::value_desc("name=<int>"), cl::Hidden,
 
 // Set the VarMap with, first, the value of __CUDA_FTZ from module metadata, and then
 // the key/value pairs from the command line.
-void NVVMReflect::setVarMap(Module &M) {
+void NVVMReflect::populateReflectMap(Module &M) {
   if (auto *Flag = mdconst::extract_or_null<ConstantInt>(
-      M.getModuleFlag("nvvm-reflect-ftz")))
-    VarMap["__CUDA_FTZ"] = Flag->getSExtValue();
+      M.getModuleFlag(CUDA_FTZ_MODULE_NAME)))
+    ReflectMap[CUDA_FTZ_NAME] = Flag->getSExtValue();
 
   for (StringRef Option : ReflectList) {
     LLVM_DEBUG(dbgs() << "ReflectOption : " << Option << "\n");
@@ -107,94 +123,52 @@ void NVVMReflect::setVarMap(Module &M) {
     int ValInt;
     if (!to_integer(Val.trim(), ValInt, 10))
       report_fatal_error("integer value expected in nvvm-reflect-list option '" + Option + "'");
-    VarMap[Name] = ValInt;
+    ReflectMap[Name] = ValInt;
   }
 }
 
-/// Process a reflect function by finding all its uses and replacing them with
+/// Process a reflect function by finding all its calls and replacing them with
 /// appropriate constant values. For __CUDA_FTZ, uses the module flag value.
 /// For __CUDA_ARCH, uses SmVersion * 10. For all other strings, uses 0.
-void NVVMReflect::handleReflectFunction(Function *F) {
-  // Validate _reflect function
+bool NVVMReflect::handleReflectFunction(Module &M, StringRef ReflectName) {
+  Function *F = M.getFunction(ReflectName);
+  if (!F)
+    return false;
   assert(F->isDeclaration() && "_reflect function should not have a body");
   assert(F->getReturnType()->isIntegerTy() && "_reflect's return type should be integer");
 
-
-  // Go through the uses of the reflect function. Each use should be a CallInst
-  // with a ConstantArray argument. Replace the uses with the appropriate constant values.
-
-  // The IR for __nvvm_reflect calls differs between CUDA versions.
-  //
-  // CUDA 6.5 and earlier uses this sequence:
-  //    %ptr = tail call i8* @llvm.nvvm.ptr.constant.to.gen.p0i8.p4i8
-  //        (i8 addrspace(4)* getelementptr inbounds
-  //           ([8 x i8], [8 x i8] addrspace(4)* @str, i32 0, i32 0))
-  //    %reflect = tail call i32 @__nvvm_reflect(i8* %ptr)
-  //
-  // The value returned by Sym->getOperand(0) is a Constant with a
-  // ConstantDataSequential operand which can be converted to string and used
-  // for lookup.
-  //
-  // CUDA 7.0 does it slightly differently:
-  //   %reflect = call i32 @__nvvm_reflect(i8* addrspacecast
-  //        (i8 addrspace(1)* getelementptr inbounds
-  //           ([8 x i8], [8 x i8] addrspace(1)* @str, i32 0, i32 0) to i8*))
-  //
-  // In this case, we get a Constant with a GlobalVariable operand and we need
-  // to dig deeper to find its initializer with the string we'll use for lookup.
-
+  bool Changed = F->getNumUses() > 0;
   for (User *U : make_early_inc_range(F->users())) {
+    // Reflect function calls look like:
+    // @arch = private unnamed_addr addrspace(1) constant [12 x i8] c"__CUDA_ARCH\00"
+    // call i32 @__nvvm_reflect(ptr addrspacecast (ptr addrspace(1) @arch to ptr))
+    // We need to extract the string argument from the call (i.e. "__CUDA_ARCH")
     if (!isa<CallInst>(U))
       report_fatal_error("__nvvm_reflect can only be used in a call instruction");
     CallInst *Call = cast<CallInst>(U);
-
     if (Call->getNumOperands() != 2)
       report_fatal_error("__nvvm_reflect requires exactly one argument");
 
-    // In cuda 6.5 and earlier, we will have an extra constant-to-generic
-    // conversion of the string.
-    const Value *Str = Call->getArgOperand(0);
-    if (const CallInst *ConvCall = dyn_cast<CallInst>(Str)) {
-      // Verify this is the constant-to-generic intrinsic
-      Function *Callee = ConvCall->getCalledFunction();
-      if (!Callee || !Callee->isIntrinsic() || 
-          !Callee->getName().starts_with("llvm.nvvm.ptr.constant.to.gen"))
-        report_fatal_error("Expected llvm.nvvm.ptr.constant.to.gen intrinsic");
-      if (ConvCall->getNumOperands() != 2)
-        report_fatal_error("Expected one argument for ptr conversion");
-      Str = ConvCall->getArgOperand(0);
-    }
-    // Pre opaque pointers we have a constant expression wrapping the constant
-    Str = Str->stripPointerCasts();
-    if (!isa<Constant>(Str))
+    const Value *GlobalStr = Call->getArgOperand(0)->stripPointerCasts();
+    if (!isa<Constant>(GlobalStr))
       report_fatal_error("__nvvm_reflect argument must be a constant string");
 
-    const Value *Operand = cast<Constant>(Str)->getOperand(0);
-    if (const GlobalVariable *GV = dyn_cast<GlobalVariable>(Operand)) {
-      // For CUDA-7.0 style __nvvm_reflect calls, we need to find the operand's
-      // initializer.
-      if (!GV->hasInitializer())
-        report_fatal_error("__nvvm_reflect string must have an initializer");
-      const Constant *Initializer = GV->getInitializer();
-      Operand = Initializer;
-    }
-
-    if (!isa<ConstantDataSequential>(Operand))
+    const Value *ConstantStr = cast<Constant>(GlobalStr)->getOperand(0);
+    if (!isa<ConstantDataSequential>(ConstantStr))
       report_fatal_error("__nvvm_reflect argument must be a string constant");
-    if (!cast<ConstantDataSequential>(Operand)->isCString())
+    if (!cast<ConstantDataSequential>(ConstantStr)->isCString())
       report_fatal_error("__nvvm_reflect argument must be a null-terminated string");
 
-    StringRef ReflectArg = cast<ConstantDataSequential>(Operand)->getAsString();
+    StringRef ReflectArg = cast<ConstantDataSequential>(ConstantStr)->getAsString();
     // Remove the null terminator from the string
     ReflectArg = ReflectArg.substr(0, ReflectArg.size() - 1);
-
     if (ReflectArg.empty())
       report_fatal_error("__nvvm_reflect argument cannot be empty");
-
+    // Now that we have extracted the string argument, we can look it up in the VarMap
     int ReflectVal = 0; // The default value is 0
-    if (VarMap.contains(ReflectArg)) {
-      ReflectVal = VarMap[ReflectArg];
-    }
+    if (ReflectMap.contains(ReflectArg))
+      ReflectVal = ReflectMap[ReflectArg];
+
     LLVM_DEBUG(dbgs() << "Replacing call of reflect function " << F->getName() << "(" << ReflectArg << ") with value " << ReflectVal << "\n");
     Constant *NewValue = ConstantInt::get(Call->getType(), ReflectVal);
     foldReflectCall(Call, NewValue);
@@ -203,29 +177,26 @@ void NVVMReflect::handleReflectFunction(Function *F) {
 
   // Remove the __nvvm_reflect function from the module
   F->eraseFromParent();
+  return Changed;
 }
 
 void NVVMReflect::foldReflectCall(CallInst *Call, Constant *NewValue) {
-  // Initialize worklist with all users of the call
   SmallVector<Instruction*, 8> Worklist;
-  for (User *U : Call->users()) {
-    if (Instruction *I = dyn_cast<Instruction>(U)) {
-      Worklist.push_back(I);
+  // Replace an instruction with a constant and add all users of the instruction to the worklist
+  auto ReplaceInstructionWithConst = [&](Instruction *I, Constant *C) {
+    for (User *U : I->users()) {
+      if (Instruction *UI = dyn_cast<Instruction>(U))
+        Worklist.push_back(UI);
     }
-  }
+    I->replaceAllUsesWith(C);
+  };
 
-  Call->replaceAllUsesWith(NewValue);
+  ReplaceInstructionWithConst(Call, NewValue);
 
   while (!Worklist.empty()) {
     Instruction *I = Worklist.pop_back_val();
     if (Constant *C = ConstantFoldInstruction(I, Call->getModule()->getDataLayout())) {
-      // Add all users of this instruction to the worklist, replace it with the constant
-      // then delete it if it's dead
-      for (User *U : I->users()) {
-        if (Instruction *UI = dyn_cast<Instruction>(U))
-          Worklist.push_back(UI);
-      }
-      I->replaceAllUsesWith(C);
+      ReplaceInstructionWithConst(I, C);
       if (isInstructionTriviallyDead(I))
         I->eraseFromParent();
     } else if (I->isTerminator()) {
@@ -237,37 +208,20 @@ void NVVMReflect::foldReflectCall(CallInst *Call, Constant *NewValue) {
 bool NVVMReflect::runOnModule(Module &M) {
   if (!NVVMReflectEnabled)
     return false;
-
-  setVarMap(M);
-
-  bool Changed = false;
-  // Names of reflect function to find and replace
-  SmallVector<StringRef, 5> ReflectNames = {
-      NVVM_REFLECT_FUNCTION,
-      NVVM_REFLECT_OCL_FUNCTION,
-      Intrinsic::getName(Intrinsic::nvvm_reflect),
-  };
-
-  // Process all reflect functions
-  for (StringRef Name : ReflectNames) {
-    if (Function *ReflectFunction = M.getFunction(Name)) {
-      // If the reflect functition is called, we need to replace the call
-      // with the appropriate constant, modifying the IR.
-      Changed |= ReflectFunction->getNumUses() > 0;
-      handleReflectFunction(ReflectFunction);
-    }
-  }
-
+  populateReflectMap(M);
+  bool Changed = true;
+  handleReflectFunction(M, NVVM_REFLECT_FUNCTION);
+  handleReflectFunction(M, NVVM_REFLECT_OCL_FUNCTION);
+  handleReflectFunction(M, Intrinsic::getName(Intrinsic::nvvm_reflect));
   return Changed;
 }
 
-// Implementations for the pass that works with the new pass manager.
-NVVMReflectPass::NVVMReflectPass(unsigned SmVersion) {
-  VarMap["__CUDA_ARCH"] = SmVersion * 10;
+bool NVVMReflectLegacyPass::runOnModule(Module &M) {
+  return Impl.runOnModule(M);
 }
 
 PreservedAnalyses NVVMReflectPass::run(Module &M,
                                     ModuleAnalysisManager &AM) {
-  return NVVMReflect(VarMap).runOnModule(M) ? PreservedAnalyses::none()
+  return NVVMReflect(SmVersion).runOnModule(M) ? PreservedAnalyses::none()
                                    : PreservedAnalyses::all();
 }