reflect improvements

Author: YonahGoldberg

llvm/lib/Target/NVPTX/NVVMReflect.cpp

@@ -4,6 +4,12 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
+// NVIDIA_COPYRIGHT_BEGIN
+//
+// Copyright (c) 2023-2025, NVIDIA CORPORATION.  All rights reserved.
+//
+// NVIDIA_COPYRIGHT_END
+//
 //===----------------------------------------------------------------------===//
 //
 // This pass replaces occurrences of __nvvm_reflect("foo") and llvm.nvvm.reflect
@@ -38,8 +44,7 @@
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
-#include "llvm/Transforms/Utils/StripGCRelocates.h"
-#include <algorithm>
+#include "llvm/ADT/StringExtras.h"
 #define NVVM_REFLECT_FUNCTION "__nvvm_reflect"
 #define NVVM_REFLECT_OCL_FUNCTION "__nvvm_reflect_ocl"
 
@@ -54,16 +59,15 @@ class NVVMReflect : public ModulePass {
   /// Process a reflect function by finding all its uses 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.
-  bool handleReflectFunction(Function *F);
+  void handleReflectFunction(Function *F);
   void setVarMap(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 int Sm) : ModulePass(ID) {
-    VarMap["__CUDA_ARCH"] = Sm * 10;
+  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.
@@ -87,51 +91,58 @@ INITIALIZE_PASS(NVVMReflect, "nvvm-reflect",
                 "Replace occurrences of __nvvm_reflect() calls with 0/1", false,
                 false)
 
+// Allow users to specify additional key/value pairs to reflect. These key/value pairs
+// are the last to be added to the VarMap, and therefore will take precedence over initial
+// values (i.e. __CUDA_FTZ from module medadata and __CUDA_ARCH from SmVersion).
 static cl::list<std::string>
-    ReflectList("nvvm-reflect-list", cl::value_desc("name=<int>"), cl::Hidden,
-                cl::desc("A list of string=num assignments"),
-                cl::ValueRequired);
-
-/// The command line can look as follows :
-/// -nvvm-reflect-list a=1,b=2 -nvvm-reflect-list c=3,d=0 -R e=2
-/// The strings "a=1,b=2", "c=3,d=0", "e=2" are available in the
-/// ReflectList vector. First, each of ReflectList[i] is 'split'
-/// using "," as the delimiter. Then each of this part is split
-/// using "=" as the delimiter.
+ReflectList("nvvm-reflect-add", cl::value_desc("name=<int>"), cl::Hidden,
+            cl::desc("list of comma-separated key=value pairs"),
+            cl::ValueRequired);
+
+// 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) {
+  LLVM_DEBUG(dbgs() << "Reflect list values:\n");
+  for (StringRef Option : ReflectList) {
+    LLVM_DEBUG(dbgs() << "  " << Option << "\n");
+  }
   if (auto *Flag = mdconst::extract_or_null<ConstantInt>(
-          M.getModuleFlag("nvvm-reflect-ftz")))
+      M.getModuleFlag("nvvm-reflect-ftz")))
     VarMap["__CUDA_FTZ"] = Flag->getSExtValue();
 
-  for (unsigned I = 0, E = ReflectList.size(); I != E; ++I) {
-    LLVM_DEBUG(dbgs() << "Option : " << ReflectList[I] << "\n");
-    SmallVector<StringRef, 4> NameValList;
-    StringRef(ReflectList[I]).split(NameValList, ",");
-    for (unsigned J = 0, EJ = NameValList.size(); J != EJ; ++J) {
-      SmallVector<StringRef, 2> NameValPair;
-      NameValList[J].split(NameValPair, "=");
-      assert(NameValPair.size() == 2 && "name=val expected");
-      StringRef ValStr = NameValPair[1].trim();
+  /// The command line can look as follows :
+  /// -nvvm-reflect-add a=1,b=2 -nvvm-reflect-add c=3,d=0 -nvvm-reflect-add e=2
+  /// The strings "a=1,b=2", "c=3,d=0", "e=2" are available in the
+  /// ReflectList vector. First, each of ReflectList[i] is 'split'
+  /// using "," as the delimiter. Then each of this part is split
+  /// using "=" as the delimiter.
+  for (StringRef Option : ReflectList) {
+    LLVM_DEBUG(dbgs() << "ReflectOption : " << Option << "\n");
+    while (!Option.empty()) {
+      std::pair<StringRef, StringRef> Split = Option.split(',');
+      StringRef NameVal = Split.first;
+      Option = Split.second;
+
+      auto NameValPair = NameVal.split('=');
+      assert(!NameValPair.first.empty() && !NameValPair.second.empty() && 
+             "name=val expected");
+      
       int Val;
-      if (ValStr.getAsInteger(10, Val))
+      if (!to_integer(NameValPair.second.trim(), Val, 10))
         report_fatal_error("integer value expected");
-      VarMap[NameValPair[0]] = Val;
+      VarMap[NameValPair.first] = Val;
     }
   }
 }
 
-bool NVVMReflect::handleReflectFunction(Function *F) {
+void NVVMReflect::handleReflectFunction(Function *F) {
   // Validate _reflect function
   assert(F->isDeclaration() && "_reflect function should not have a body");
-  assert(F->getReturnType()->isIntegerTy() &&
-         "_reflect's return type should be integer");
+  assert(F->getReturnType()->isIntegerTy() && "_reflect's return type should be integer");
 
-  SmallVector<Instruction *, 4> ToRemove;
-  SmallVector<Instruction *, 4> ToSimplify;
 
   // 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.
+  // with a ConstantArray argument. Replace the uses with the appropriate constant values.
 
   // The IR for __nvvm_reflect calls differs between CUDA versions.
   //
@@ -153,7 +164,7 @@ bool NVVMReflect::handleReflectFunction(Function *F) {
   // 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.
 
-  for (User *U : F->users()) {
+  for (User *U : make_early_inc_range(F->users())) {
     assert(isa<CallInst>(U) && "Only a call instruction can use _reflect");
     CallInst *Call = cast<CallInst>(U);
 
@@ -165,21 +176,23 @@ bool NVVMReflect::handleReflectFunction(Function *F) {
     // conversion of the string.
     const Value *Str = Call->getArgOperand(0);
     if (const CallInst *ConvCall = dyn_cast<CallInst>(Str)) {
-      // FIXME: Add assertions about ConvCall.
+      // Verify this is the constant-to-generic intrinsic
+      Function *Callee = ConvCall->getCalledFunction();
+      assert(Callee && Callee->isIntrinsic() && 
+             Callee->getName().starts_with("llvm.nvvm.ptr.constant.to.gen") &&
+             "Expected llvm.nvvm.ptr.constant.to.gen intrinsic");
+      assert(ConvCall->getNumOperands() == 2 && "Expected one argument for ptr conversion");
       Str = ConvCall->getArgOperand(0);
     }
     // Pre opaque pointers we have a constant expression wrapping the constant
-    // string.
     Str = Str->stripPointerCasts();
-    assert(isa<Constant>(Str) &&
-           "Format of __nvvm_reflect function not recognized");
+    assert(isa<Constant>(Str) && "Format of __nvvm_reflect function not recognized");
 
     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.
-      assert(GV->hasInitializer() &&
-             "Format of _reflect function not recognized");
+      assert(GV->hasInitializer() && "Format of _reflect function not recognized");
       const Constant *Initializer = GV->getInitializer();
       Operand = Initializer;
     }
@@ -192,54 +205,48 @@ bool NVVMReflect::handleReflectFunction(Function *F) {
     StringRef ReflectArg = cast<ConstantDataSequential>(Operand)->getAsString();
     // Remove the null terminator from the string
     ReflectArg = ReflectArg.substr(0, ReflectArg.size() - 1);
-    LLVM_DEBUG(dbgs() << "Arg of _reflect : " << ReflectArg << "\n");
 
     int ReflectVal = 0; // The default value is 0
     if (VarMap.contains(ReflectArg)) {
       ReflectVal = VarMap[ReflectArg];
     }
-    LLVM_DEBUG(dbgs() << "ReflectVal: " << ReflectVal << "\n");
+    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);
+    Call->eraseFromParent();
+  }
 
-    // If the immediate user is a simple comparison we want to simplify it.
-    for (User *U : Call->users())
-      if (Instruction *I = dyn_cast<Instruction>(U))
-        ToSimplify.push_back(I);
+  // Remove the __nvvm_reflect function from the module
+  F->eraseFromParent();
+}
 
-    Call->replaceAllUsesWith(ConstantInt::get(Call->getType(), ReflectVal));
-    ToRemove.push_back(Call);
+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);
+    }
   }
 
-  // The code guarded by __nvvm_reflect may be invalid for the target machine.
-  // Traverse the use-def chain, continually simplifying constant expressions
-  // until we find a terminator that we can then remove.
-  while (!ToSimplify.empty()) {
-    Instruction *I = ToSimplify.pop_back_val();
-    if (Constant *C = ConstantFoldInstruction(I, F->getDataLayout())) {
-      for (User *U : I->users())
-        if (Instruction *I = dyn_cast<Instruction>(U))
-          ToSimplify.push_back(I);
+  Call->replaceAllUsesWith(NewValue);
 
-      I->replaceAllUsesWith(C);
-      if (isInstructionTriviallyDead(I)) {
-        ToRemove.push_back(I);
+  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);
+      if (isInstructionTriviallyDead(I))
+        I->eraseFromParent();
     } else if (I->isTerminator()) {
       ConstantFoldTerminator(I->getParent());
     }
   }
-
-  // Removing via isInstructionTriviallyDead may add duplicates to the ToRemove
-  // array. Filter out the duplicates before starting to erase from parent.
-  std::sort(ToRemove.begin(), ToRemove.end());
-  auto *NewLastIter = llvm::unique(ToRemove);
-  ToRemove.erase(NewLastIter, ToRemove.end());
-
-  for (Instruction *I : ToRemove)
-    I->eraseFromParent();
-
-  // Remove the __nvvm_reflect function from the module
-  F->eraseFromParent();
-  return ToRemove.size() > 0;
 }
 
 bool NVVMReflect::runOnModule(Module &M) {
@@ -250,15 +257,19 @@ bool NVVMReflect::runOnModule(Module &M) {
 
   bool Changed = false;
   // Names of reflect function to find and replace
-  SmallVector<std::string, 3> ReflectNames = {
-      NVVM_REFLECT_FUNCTION, NVVM_REFLECT_OCL_FUNCTION,
-      Intrinsic::getName(Intrinsic::nvvm_reflect).str()};
+  SmallVector<StringRef, 5> ReflectNames = {
+      NVVM_REFLECT_FUNCTION,
+      NVVM_REFLECT_OCL_FUNCTION,
+      Intrinsic::getName(Intrinsic::nvvm_reflect),
+  };
 
   // Process all reflect functions
-  for (const std::string &Name : ReflectNames) {
-    Function *ReflectFunction = M.getFunction(Name);
-    if (ReflectFunction) {
-      Changed |= handleReflectFunction(ReflectFunction);
+  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);
     }
   }
 
@@ -269,7 +280,9 @@ bool NVVMReflect::runOnModule(Module &M) {
 NVVMReflectPass::NVVMReflectPass(unsigned SmVersion) {
   VarMap["__CUDA_ARCH"] = SmVersion * 10;
 }
-PreservedAnalyses NVVMReflectPass::run(Module &M, ModuleAnalysisManager &AM) {
+
+PreservedAnalyses NVVMReflectPass::run(Module &M,
+                                    ModuleAnalysisManager &AM) {
   return NVVMReflect(VarMap).runOnModule(M) ? PreservedAnalyses::none()
-                                            : PreservedAnalyses::all();
-}
+                                   : PreservedAnalyses::all();
+}

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

@@ -0,0 +1,26 @@
+; Verify that when passing in command-line options to NVVMReflect, that reflect calls are replaced with
+; the appropriate command line values.
+
+declare i32 @__nvvm_reflect(ptr)
+@ftz = private unnamed_addr addrspace(1) constant [11 x i8] c"__CUDA_FTZ\00"
+@arch = private unnamed_addr addrspace(1) constant [12 x i8] c"__CUDA_ARCH\00"
+
+; RUN: opt -passes=nvvm-reflect -mtriple=nvptx-nvidia-cuda -nvvm-reflect-add=__CUDA_FTZ=1,__CUDA_ARCH=350 %s -S | FileCheck %s --check-prefix=CHECK-FTZ1-ARCH350
+; RUN: opt -passes=nvvm-reflect -mtriple=nvptx-nvidia-cuda -nvvm-reflect-add=__CUDA_FTZ=0 -nvvm-reflect-add=__CUDA_ARCH=520 %s -S | FileCheck %s --check-prefix=CHECK-FTZ0-ARCH520
+
+; Verify that if we have module metadata that sets __CUDA_FTZ=1, that gets overridden by the command line arguments
+
+; RUN: cat %s > %t.options
+; RUN: echo '!llvm.module.flags = !{!0}' >> %t.options
+; RUN: echo '!0 = !{i32 4, !"nvvm-reflect-ftz", i32 1}' >> %t.options
+; RUN: opt -passes=nvvm-reflect -mtriple=nvptx-nvidia-cuda -nvvm-reflect-add=__CUDA_FTZ=0,__CUDA_ARCH=520 %t.options -S | FileCheck %s --check-prefix=CHECK-FTZ0-ARCH520
+
+define i32 @options() {
+  %1 = call i32 @__nvvm_reflect(ptr addrspacecast (ptr addrspace(1) @ftz to ptr))
+  %2 = call i32 @__nvvm_reflect(ptr addrspacecast (ptr addrspace(1) @arch to ptr))
+  %3 = add i32 %1, %2
+  ret i32 %3
+}
+
+; CHECK-FTZ1-ARCH350: ret i32 351
+; CHECK-FTZ0-ARCH520: ret i32 520