[NVPTX] Handle __nvvm_reflect("__CUDA_ARCH").

[tra]

Summary

libdevice in recent CUDA versions relies on __nvvm_reflect() to select
GPU-specific bitcode. This patch addresses the requirement.

[jlebar]

Just to check, the notion is that it's OK if I report a sm version less than what I end up running on?

llvm/lib/Target/NVPTX/NVVMReflect.cpp

55

explicit

[tra]

In D50207#1186491, @jlebar wrote:

Just to check, the notion is that it's OK if I report a sm version less than what I end up running on?

Yes. We may lose some performance, but not correctness as we are expected to be forward-compatible.
Until now reflect was being replaced with 0, so we were picking the variant suitable for the oldest GPU.

[tra]

Added 'explicit' keyword to contructor.

llvm/lib/Target/NVPTX/NVPTX.h

 };
 }
 
 FunctionPass *createNVPTXISelDag(NVPTXTargetMachine &TM,
                                  llvm::CodeGenOpt::Level OptLevel);
 ModulePass *createNVPTXAssignValidGlobalNamesPass();
 ModulePass *createGenericToNVVMPass();
 FunctionPass *createNVVMIntrRangePass(unsigned int SmVersion);
-FunctionPass *createNVVMReflectPass();
+FunctionPass *createNVVMReflectPass(unsigned int SmVersion);
 MachineFunctionPass *createNVPTXPrologEpilogPass();
 MachineFunctionPass *createNVPTXReplaceImageHandlesPass();
 FunctionPass *createNVPTXImageOptimizerPass();
 FunctionPass *createNVPTXLowerArgsPass(const NVPTXTargetMachine *TM);
 BasicBlockPass *createNVPTXLowerAllocaPass();
 MachineFunctionPass *createNVPTXPeephole();
 
 Target &getTheNVPTXTarget32();

llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp

 TargetPassConfig *NVPTXTargetMachine::createPassConfig(PassManagerBase &PM) {
   return new NVPTXPassConfig(*this, PM);
 }
 
 void NVPTXTargetMachine::adjustPassManager(PassManagerBuilder &Builder) {
   Builder.addExtension(
     PassManagerBuilder::EP_EarlyAsPossible,
     [&](const PassManagerBuilder &, legacy::PassManagerBase &PM) {
-      PM.add(createNVVMReflectPass());
+      PM.add(createNVVMReflectPass(Subtarget.getSmVersion()));
       PM.add(createNVVMIntrRangePass(Subtarget.getSmVersion()));
     });
 }
 
 TargetTransformInfo
 NVPTXTargetMachine::getTargetTransformInfo(const Function &F) {
   return TargetTransformInfo(NVPTXTTIImpl(this, F));
 }
void NVPTXPassConfig::addIRPasses() {
   disablePass(&FuncletLayoutID);
   disablePass(&PatchableFunctionID);
   disablePass(&ShrinkWrapID);
 
   // NVVMReflectPass is added in addEarlyAsPossiblePasses, so hopefully running
   // it here does nothing.  But since we need it for correctness when lowering
   // to NVPTX, run it here too, in case whoever built our pass pipeline didn't
   // call addEarlyAsPossiblePasses.
-  addPass(createNVVMReflectPass());
+  const NVPTXSubtarget &ST = *getTM<NVPTXTargetMachine>().getSubtargetImpl();
+  addPass(createNVVMReflectPass(ST.getSmVersion()));
 
   if (getOptLevel() != CodeGenOpt::None)
     addPass(createNVPTXImageOptimizerPass());
   addPass(createNVPTXAssignValidGlobalNamesPass());
   addPass(createGenericToNVVMPass());
 
   // NVPTXLowerArgs is required for correctness and should be run right
   // before the address space inference passes.

llvm/lib/Target/NVPTX/NVVMReflect.cpp

 #define DEBUG_TYPE "nvptx-reflect"
 
 namespace llvm { void initializeNVVMReflectPass(PassRegistry &); }
 
 namespace {
 class NVVMReflect : public FunctionPass {
 public:
   static char ID;
-  NVVMReflect() : FunctionPass(ID) {
+  unsigned int SmVersion;
+  NVVMReflect() : NVVMReflect(0) {}
+  NVVMReflect(unsigned int Sm) : FunctionPass(ID), SmVersion(Sm) {
     initializeNVVMReflectPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnFunction(Function &) override;
 };
 }
 
-FunctionPass *llvm::createNVVMReflectPass() { return new NVVMReflect(); }
+FunctionPass *llvm::createNVVMReflectPass(unsigned int SmVersion) {
+  return new NVVMReflect(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",
                 "Replace occurrences of __nvvm_reflect() calls with 0/1", false,
for (Instruction &I : instructions(F)) {
     int ReflectVal = 0; // The default value is 0
     if (ReflectArg == "__CUDA_FTZ") {
       // Try to pull __CUDA_FTZ from the nvvm-reflect-ftz module flag.  Our
       // choice here must be kept in sync with AutoUpgrade, which uses the same
       // technique to detect whether ftz is enabled.
       if (auto *Flag = mdconst::extract_or_null<ConstantInt>(
               F.getParent()->getModuleFlag("nvvm-reflect-ftz")))
         ReflectVal = Flag->getSExtValue();
+    } else if (ReflectArg == "__CUDA_ARCH") {
+      ReflectVal = SmVersion * 10;
     }
     Call->replaceAllUsesWith(ConstantInt::get(Call->getType(), ReflectVal));
     ToRemove.push_back(Call);
   }
 
   for (Instruction *I : ToRemove)
     I->eraseFromParent();
 
   return ToRemove.size() > 0;
 }

llvm/test/CodeGen/NVPTX/nvvm-reflect-arch.ll

+; Libdevice in recent CUDA versions relies on __CUDA_ARCH reflecting GPU type.
+; Verify that __nvvm_reflect() is replaced with an appropriate value.
+;
+; RUN: opt %s -S -nvvm-reflect -O2 -mtriple=nvptx64 \
+; RUN:   | FileCheck %s --check-prefixes=COMMON,SM20
+; RUN: opt %s -S -nvvm-reflect -O2 -mtriple=nvptx64 -mcpu=sm_35 \
+; RUN:   | FileCheck %s --check-prefixes=COMMON,SM35
+
+@"$str" = private addrspace(1) constant [12 x i8] c"__CUDA_ARCH\00"
+
+declare i32 @__nvvm_reflect(i8*)
+
+; COMMON-LABEL: @foo
+define i32 @foo(float %a, float %b) {
+; COMMON-NOT: call i32 @__nvvm_reflect
+  %reflect = call i32 @__nvvm_reflect(i8* addrspacecast (i8 addrspace(1)* getelementptr inbounds ([12 x i8], [12 x i8] addrspace(1)* @"$str", i32 0, i32 0) to i8*))
+; SM20: ret i32 200  
+; SM35: ret i32 350  
+  ret i32 %reflect
+}
+