[SYCL][NVPTX] Emit 'grid_constant' annotations for by-val kernel params

frasercrmck · frasercrmck · commit d2233ddcd095 · 2024-06-27T17:39:25.000+01:00
Also fix up the DeadArgumentElimination passes to correctly preserve the
annotations; when removing arguments from functions, dead parameters
need pruned and alive ones may need their values shifted down by the
number of dead arguments that came before them.
diff --git a/clang/lib/CodeGen/Targets/NVPTX.cpp b/clang/lib/CodeGen/Targets/NVPTX.cpp
@@ -8,6 +8,7 @@
 
 #include "ABIInfoImpl.h"
 #include "TargetInfo.h"
+#include "clang/Basic/Cuda.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
 
 using namespace clang;
@@ -80,6 +81,9 @@ class NVPTXTargetCodeGenInfo : public TargetCodeGenInfo {
   static void addNVVMMetadata(llvm::GlobalValue *GV, StringRef Name,
                               int Operand);
 
+  static void addNVVMMetadata(llvm::GlobalValue *GV, StringRef Name,
+                              const std::vector<int> &Operands);
+
 private:
   static void emitBuiltinSurfTexDeviceCopy(CodeGenFunction &CGF, LValue Dst,
                                            LValue Src) {
@@ -218,6 +222,98 @@ Address NVPTXABIInfo::EmitVAArg(CodeGenFunction &CGF, Address VAListAddr,
   llvm_unreachable("NVPTX does not support varargs");
 }
 
+// Get current CudaArch and ignore any unknown values
+// Copied from CGOpenMPRuntimeGPU
+static CudaArch getCudaArch(CodeGenModule &CGM) {
+  if (!CGM.getTarget().hasFeature("ptx"))
+    return CudaArch::UNKNOWN;
+  for (const auto &Feature : CGM.getTarget().getTargetOpts().FeatureMap) {
+    if (Feature.getValue()) {
+      CudaArch Arch = StringToCudaArch(Feature.getKey());
+      if (Arch != CudaArch::UNKNOWN)
+        return Arch;
+    }
+  }
+  return CudaArch::UNKNOWN;
+}
+
+static bool supportsGridConstant(CudaArch Arch) {
+  switch (Arch) {
+  case CudaArch::SM_70:
+  case CudaArch::SM_72:
+  case CudaArch::SM_75:
+  case CudaArch::SM_80:
+  case CudaArch::SM_86:
+  case CudaArch::SM_87:
+  case CudaArch::SM_89:
+  case CudaArch::SM_90:
+  case CudaArch::SM_90a:
+    return true;
+  case CudaArch::UNKNOWN:
+  case CudaArch::UNUSED:
+  case CudaArch::SM_20:
+  case CudaArch::SM_21:
+  case CudaArch::SM_30:
+  case CudaArch::SM_32_:
+  case CudaArch::SM_35:
+  case CudaArch::SM_37:
+  case CudaArch::SM_50:
+  case CudaArch::SM_52:
+  case CudaArch::SM_53:
+  case CudaArch::SM_60:
+  case CudaArch::SM_61:
+  case CudaArch::SM_62:
+    return false;
+  case CudaArch::GFX600:
+  case CudaArch::GFX601:
+  case CudaArch::GFX602:
+  case CudaArch::GFX700:
+  case CudaArch::GFX701:
+  case CudaArch::GFX702:
+  case CudaArch::GFX703:
+  case CudaArch::GFX704:
+  case CudaArch::GFX705:
+  case CudaArch::GFX801:
+  case CudaArch::GFX802:
+  case CudaArch::GFX803:
+  case CudaArch::GFX805:
+  case CudaArch::GFX810:
+  case CudaArch::GFX900:
+  case CudaArch::GFX902:
+  case CudaArch::GFX904:
+  case CudaArch::GFX906:
+  case CudaArch::GFX908:
+  case CudaArch::GFX909:
+  case CudaArch::GFX90a:
+  case CudaArch::GFX90c:
+  case CudaArch::GFX940:
+  case CudaArch::GFX941:
+  case CudaArch::GFX942:
+  case CudaArch::GFX1010:
+  case CudaArch::GFX1011:
+  case CudaArch::GFX1012:
+  case CudaArch::GFX1013:
+  case CudaArch::GFX1030:
+  case CudaArch::GFX1031:
+  case CudaArch::GFX1032:
+  case CudaArch::GFX1033:
+  case CudaArch::GFX1034:
+  case CudaArch::GFX1035:
+  case CudaArch::GFX1036:
+  case CudaArch::GFX1100:
+  case CudaArch::GFX1101:
+  case CudaArch::GFX1102:
+  case CudaArch::GFX1103:
+  case CudaArch::GFX1150:
+  case CudaArch::GFX1151:
+  case CudaArch::GFX1200:
+  case CudaArch::GFX1201:
+  case CudaArch::Generic:
+  case CudaArch::LAST:
+    llvm_unreachable("unhandled CudaArch");
+  }
+}
+
 void NVPTXTargetCodeGenInfo::setTargetAttributes(
     const Decl *D, llvm::GlobalValue *GV, CodeGen::CodeGenModule &M) const {
   if (GV->isDeclaration())
@@ -248,6 +344,21 @@ void NVPTXTargetCodeGenInfo::setTargetAttributes(
       addNVVMMetadata(F, "kernel", 1);
       // And kernel functions are not subject to inlining
       F->addFnAttr(llvm::Attribute::NoInline);
+
+      if (supportsGridConstant(getCudaArch(M))) {
+        // Add grid_constant annotations to all relevant kernel-function
+        // parameters. We can guarantee that in SYCL, all by-val kernel
+        // parameters are "grid_constant".
+        std::vector<int> GridConstantParamIdxs;
+        for (auto [Idx, Arg] : llvm::enumerate(F->args())) {
+          if (Arg.getType()->isPointerTy() && Arg.hasByValAttr()) {
+            // Note - the parameter indices are numbered from 1.
+            GridConstantParamIdxs.push_back(Idx + 1);
+          }
+        }
+        if (!GridConstantParamIdxs.empty())
+          addNVVMMetadata(F, "grid_constant", GridConstantParamIdxs);
+      }
     }
     bool HasMaxWorkGroupSize = false;
     bool HasMinWorkGroupPerCU = false;
@@ -329,6 +440,28 @@ void NVPTXTargetCodeGenInfo::addNVVMMetadata(llvm::GlobalValue *GV,
   MD->addOperand(llvm::MDNode::get(Ctx, MDVals));
 }
 
+void NVPTXTargetCodeGenInfo::addNVVMMetadata(llvm::GlobalValue *GV,
+                                             StringRef Name,
+                                             const std::vector<int> &Operands) {
+  llvm::Module *M = GV->getParent();
+  llvm::LLVMContext &Ctx = M->getContext();
+
+  // Get "nvvm.annotations" metadata node
+  llvm::NamedMDNode *MD = M->getOrInsertNamedMetadata("nvvm.annotations");
+
+  llvm::SmallVector<llvm::Metadata *, 8> MDOps;
+  for (int Op : Operands) {
+    MDOps.push_back(llvm::ConstantAsMetadata::get(
+        llvm::ConstantInt::get(llvm::Type::getInt32Ty(Ctx), Op)));
+  }
+  auto *OpList = llvm::MDNode::get(Ctx, MDOps);
+
+  llvm::Metadata *MDVals[] = {llvm::ConstantAsMetadata::get(GV),
+                              llvm::MDString::get(Ctx, Name), OpList};
+  // Append metadata to nvvm.annotations
+  MD->addOperand(llvm::MDNode::get(Ctx, MDVals));
+}
+
 bool NVPTXTargetCodeGenInfo::shouldEmitStaticExternCAliases() const {
   return false;
 }
diff --git a/clang/test/CodeGenSYCL/nvvm-annotations.cpp b/clang/test/CodeGenSYCL/nvvm-annotations.cpp
@@ -0,0 +1,35 @@
+// RUN: %clang_cc1 -fno-sycl-force-inline-kernel-lambda -fsycl-is-device -internal-isystem %S/Inputs -triple nvptx-nvidia-cuda -target-cpu sm_70 -disable-llvm-passes -sycl-std=2020 -emit-llvm -o - %s | FileCheck %s --check-prefixes=CHECK,GRIDCONST
+// RUN: %clang_cc1 -fno-sycl-force-inline-kernel-lambda -fsycl-is-device -internal-isystem %S/Inputs -triple nvptx64-nvidia-cuda -target-cpu sm_70 -disable-llvm-passes -sycl-std=2020 -emit-llvm -o - %s | FileCheck %s --check-prefixes=CHECK,GRIDCONST
+
+// RUN: %clang_cc1 -fno-sycl-force-inline-kernel-lambda -fsycl-is-device -internal-isystem %S/Inputs -triple nvptx-nvidia-cuda -target-cpu sm_60 -disable-llvm-passes -sycl-std=2020 -emit-llvm -o - %s | FileCheck %s --check-prefixes=CHECK,NOGRIDCONST
+// RUN: %clang_cc1 -fno-sycl-force-inline-kernel-lambda -fsycl-is-device -internal-isystem %S/Inputs -triple nvptx64-nvidia-cuda -target-cpu sm_60 -disable-llvm-passes -sycl-std=2020 -emit-llvm -o - %s | FileCheck %s --check-prefixes=CHECK,NOGRIDCONST
+
+// Tests that work_group_size_hint and reqd_work_group_size generate the same
+// metadata nodes for the same arguments.
+
+#include "sycl.hpp"
+
+using namespace sycl;
+
+int main() {
+  queue q;
+
+  struct S {
+    int a;
+  } s;
+  q.submit([&](handler &h) {
+    // CHECK: define{{.*}} void @[[FUNC1:.*kernel_grid_const_params]](ptr noundef byval(%struct.S) align 4 %_arg_s)
+    h.single_task<class kernel_grid_const_params>([=]() { (void) s;});
+  });
+
+  return 0;
+}
+// Don't emit grid_constant annotations for older architectures.
+
+// NOGRIDCONST-NOT: "grid_constant"
+
+// This isn't stable in general, as it depends on the order of the captured
+// parameters, but in this case there's only one parameter so we know it's 1.
+// GRIDCONST-DAG: = !{ptr @[[FUNC1]], !"grid_constant", [[MD:\![0-9]+]]}
+// GRIDCONST-DAG: [[MD]] = !{i32 1} 
+
diff --git a/llvm/include/llvm/Transforms/IPO/DeadArgumentElimination.h b/llvm/include/llvm/Transforms/IPO/DeadArgumentElimination.h
@@ -145,7 +145,8 @@ class DeadArgumentEliminationPass
   bool removeDeadArgumentsFromCallers(Function &F);
   void propagateVirtMustcallLiveness(const Module &M);
 
-  void UpdateNVPTXMetadata(Module &M, Function *F, Function *NF);
+  void UpdateNVPTXMetadata(Module &M, Function *F, Function *NF,
+                           const SmallVectorImpl<bool> &ArgAlive);
   llvm::DenseSet<Function *> NVPTXKernelSet;
 
   bool IsNVPTXKernel(const Function *F) { return NVPTXKernelSet.contains(F); };
diff --git a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
@@ -1171,7 +1171,7 @@ bool DeadArgumentEliminationPass::removeDeadStuffFromFunction(Function *F) {
     NF->addMetadata(KindID, *Node);
 
   if (IsNVPTXKernel(F))
-    UpdateNVPTXMetadata(*(F->getParent()), F, NF);
+    UpdateNVPTXMetadata(*(F->getParent()), F, NF, ArgAlive);
 
   // If either the return value(s) or argument(s) are removed, then probably the
   // function does not follow standard calling conventions anymore. Hence, add
@@ -1249,8 +1249,9 @@ PreservedAnalyses DeadArgumentEliminationPass::run(Module &M,
   return PreservedAnalyses::none();
 }
 
-void DeadArgumentEliminationPass::UpdateNVPTXMetadata(Module &M, Function *F,
-                                                      Function *NF) {
+void DeadArgumentEliminationPass::UpdateNVPTXMetadata(
+    Module &M, Function *F, Function *NF,
+    const SmallVectorImpl<bool> &ArgAlive) {
 
   auto *NvvmMetadata = M.getNamedMetadata("nvvm.annotations");
   if (!NvvmMetadata)
@@ -1260,13 +1261,74 @@ void DeadArgumentEliminationPass::UpdateNVPTXMetadata(Module &M, Function *F,
     const auto &FuncOperand = MetadataNode->getOperand(0);
     if (!FuncOperand)
       continue;
-    auto FuncConstant = dyn_cast<ConstantAsMetadata>(FuncOperand);
+    auto *FuncConstant = dyn_cast<ConstantAsMetadata>(FuncOperand);
     if (!FuncConstant)
       continue;
     auto *Func = dyn_cast<Function>(FuncConstant->getValue());
     if (Func != F)
       continue;
     // Update the metadata with the new function
     MetadataNode->replaceOperandWith(0, llvm::ConstantAsMetadata::get(NF));
+
+    // Carefully update grid_constant annotations, since those are denoted
+    // parameter indices, which may have changed
+    if (auto *Type = dyn_cast<MDString>(MetadataNode->getOperand(1));
+        Type && Type->getString() == "grid_constant") {
+      LLVMContext &Ctx = NF->getContext();
+      LLVM_DEBUG(dbgs() << "DeadArgumentEliminationPass - updating nvvm "
+                           "grid_constant annotations for fn: "
+                        << NF->getName() << "\n");
+      // The 'value' operand is a list of integers denoting parameter indices
+      auto *OldGridConstParamIdxs = dyn_cast<MDNode>(MetadataNode->getOperand(2));
+      if (!OldGridConstParamIdxs)
+        continue;
+      // For each parameter that's identified as a grid_constant, count how
+      // many arguments before that position are dead, and shift the number
+      // down by that amount.
+      // Note that there's no guaranteed order to the parameter indices, so
+      // there's fewer 'smart' things like counting up incrementally as we go.
+      SmallVector<Metadata *, 8> NewGridConstParamOps;
+      for (const auto &Op : OldGridConstParamIdxs->operands()) {
+        auto *ParamIdx = mdconst::dyn_extract<ConstantInt>(Op);
+        // If the operand's not a constant, or its constant value is not
+        // within the range of the old function's parameter list (note - it
+        // counts from 1), it's not well-defined. Just strip it out for
+        // safety.
+        if (!ParamIdx || ParamIdx->isZero() ||
+            ParamIdx->getZExtValue() > F->arg_size())
+          continue;
+
+        size_t OldParamIdx = ParamIdx->getZExtValue() - 1;
+        // If the parameter is no longer alive, it's definitely not a
+        // grid_constant. Strip it out.
+        if (!ArgAlive[OldParamIdx])
+          continue;
+
+        unsigned ShiftDownAmt = 0;
+        for (unsigned i = 0; i < std::min(F->arg_size(), OldParamIdx); i++) {
+          if (!ArgAlive[i])
+            ShiftDownAmt++;
+        }
+        NewGridConstParamOps.push_back(ConstantAsMetadata::get(ConstantInt::get(
+            Type::getInt32Ty(Ctx), OldParamIdx - ShiftDownAmt + 1)));
+      }
+
+      // Update the metadata with the new grid_constant information
+      MDNode *NewGridConstParamIdxs = MDNode::get(Ctx, NewGridConstParamOps);
+
+      LLVM_DEBUG(dbgs() << "  * updating old annotation {";
+                 auto PrintList =
+                     [](const MDNode *MD) {
+                       for (const auto &O : MD->operands())
+                         if (const auto *ParamNo =
+                                 mdconst::dyn_extract<ConstantInt>(O))
+                           dbgs() << ParamNo->getZExtValue() << ",";
+                     };
+                 PrintList(OldGridConstParamIdxs);
+                 dbgs() << "} to new annotation {";
+                 PrintList(NewGridConstParamIdxs); dbgs() << "}\n";);
+
+      MetadataNode->replaceOperandWith(2, NewGridConstParamIdxs);
+    }
   }
 }
diff --git a/llvm/test/Transforms/DeadArgElim/nvvm-annotations.ll b/llvm/test/Transforms/DeadArgElim/nvvm-annotations.ll
