1 files changed, 17 insertions, 7 deletions
diff --git a/llvm/lib/Target/NVPTX/NVVMReflect.cpp b/llvm/lib/Target/NVPTX/NVVMReflect.cpp
index 14f60761966..2f80791a433 100644
--- a/llvm/lib/Target/NVPTX/NVVMReflect.cpp
+++ b/llvm/lib/Target/NVPTX/NVVMReflect.cpp
@@ -7,11 +7,16 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// 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.
 //
 //===----------------------------------------------------------------------===//
 
@@ -55,10 +60,9 @@ 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();
   }
 
@@ -206,6 +210,12 @@ bool NVVMReflect::runOnFunction(Function &F) {
     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);
   }