[NVPTX] Read __CUDA_FTZ from module flags in NVVMReflect.

[jlebar]

Summary

Previously the NVVMReflect pass would read its configuration from
command-line flags or a static configuration given to the pass at
instantiation time.

This doesn't quite work for clang's use-case. It needs to pass a value
for __CUDA_FTZ down on a per-module basis. We use a module flag for
this, so the NVVMReflect pass needs to be updated to read said flag.

[tra]

LGTM.

lib/Target/NVPTX/NVVMReflect.cpp

 //===- NVVMReflect.cpp - NVVM Emulate conditional compilation -------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
-// This pass replaces occurrences of __nvvm_reflect("string") and
-// llvm.nvvm.reflect with an integer based on the value of -nvvm-reflect-list
-// string=<int>.
+// This pass replaces occurrences of __nvvm_reflect("foo") and llvm.nvvm.reflect
+// with an integer.
 //
-// If we see a string not specified in our flags, we replace that call with 0.
+// We choose the value we use by looking, in this order, at:
+//
+//  * the -nvvm-reflect-list flag, which has the format "foo=1,bar=42",
+//  * the StringMap passed to the pass's constructor, and
+//  * metadata in the module itself.
+//
+// If we see an unknown string, we replace its call with 0.
 //
 //===----------------------------------------------------------------------===//
 
 #include "NVPTX.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"

 class NVVMReflect : public FunctionPass {
 private:
   StringMap<int> VarMap;
 
 public:
   static char ID;
   NVVMReflect() : NVVMReflect(StringMap<int>()) {}
 
-  NVVMReflect(const StringMap<int> &Mapping) : FunctionPass(ID) {
+  NVVMReflect(const StringMap<int> &Mapping)
+      : FunctionPass(ID), VarMap(Mapping) {
     initializeNVVMReflectPass(*PassRegistry::getPassRegistry());
-    for (const auto &KV : Mapping)
-      VarMap[KV.getKey()] = KV.getValue();
     setVarMap();
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
   }
   bool runOnFunction(Function &) override;
 
for (Instruction &I : instructions(F)) {
     StringRef ReflectArg = cast<ConstantDataSequential>(Operand)->getAsString();
     ReflectArg = ReflectArg.substr(0, ReflectArg.size() - 1);
     DEBUG(dbgs() << "Arg of _reflect : " << ReflectArg << "\n");
 
     int ReflectVal = 0; // The default value is 0
     auto Iter = VarMap.find(ReflectArg);
     if (Iter != VarMap.end())
       ReflectVal = Iter->second;
+    else if (ReflectArg == "__CUDA_FTZ") {
+      // Try to pull __CUDA_FTZ from the nvvm-reflect-ftz module flag.
+      if (auto *Flag = mdconst::extract_or_null<ConstantInt>(
+              F.getParent()->getModuleFlag("nvvm-reflect-ftz")))
+        ReflectVal = Flag->getSExtValue();
+    }
     Call->replaceAllUsesWith(ConstantInt::get(Call->getType(), ReflectVal));
     ToRemove.push_back(Call);
   }
 
   for (Instruction *I : ToRemove)
     I->eraseFromParent();
 
   return ToRemove.size() > 0;
 }

test/CodeGen/NVPTX/nvvm-reflect-module-flag.ll

+; RUN: opt < %s -S -nvvm-reflect | FileCheck %s
+
+declare i32 @__nvvm_reflect(i8*)
+@str = private unnamed_addr addrspace(1) constant [11 x i8] c"__CUDA_FTZ\00"
+
+define i32 @foo() {
+  %call = call i32 @__nvvm_reflect(i8* addrspacecast (i8 addrspace(1)* getelementptr inbounds ([11 x i8], [11 x i8] addrspace(1)* @str, i32 0, i32 0) to i8*))
+  ; CHECK: ret i32 42
+  ret i32 %call
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 4, !"nvvm-reflect-ftz", i32 42}