Revert "Relocate and use llvm::omp::getDeviceKernels"

This reverts commit bb9d5c2.

This will facilitate merging main due to 07ed818 (PR llvm#122320), which
changes llvm::omp::getDeviceKernels.  Will rewrite and reapply after
merging main.

Author: jdenny-ornl

llvm/include/llvm/IR/Module.h

@@ -16,7 +16,6 @@
 
 #include "llvm-c/Types.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/iterator_range.h"
@@ -514,9 +513,6 @@ class LLVM_ABI Module {
   /// Remove the given NamedMDNode from this module and delete it.
   void eraseNamedMetadata(NamedMDNode *NMD);
 
-  /// Get device kernels in the module.
-  SetVector<Function *> getDeviceKernels();
-
 /// @}
 /// @name Comdat Accessors
 /// @{

llvm/lib/Analysis/KernelInfo.cpp

@@ -78,6 +78,11 @@ class KernelInfo {
 
 } // end anonymous namespace
 
+static bool isKernelFunction(Function &F) {
+  // TODO: Is this general enough?  Consider languages beyond OpenMP.
+  return F.hasFnAttribute("kernel");
+}
+
 static void identifyCallee(OptimizationRemark &R, const Module *M,
                            const Value *V, StringRef Kind = "") {
   SmallString<100> Name; // might be function name or asm expression
@@ -287,8 +292,7 @@ void KernelInfo::emitKernelInfo(Function &F, FunctionAnalysisManager &FAM,
   KI.FlatAddrspace = TheTTI.getFlatAddressSpace();
 
   // Record function properties.
-  KI.ExternalNotKernel =
-      F.hasExternalLinkage() && !F.getParent()->getDeviceKernels().contains(&F);
+  KI.ExternalNotKernel = F.hasExternalLinkage() && !isKernelFunction(F);
   for (StringRef Name : {"omp_target_num_teams", "omp_target_thread_limit"}) {
     if (auto Val = parseFnAttrAsInteger(F, Name))
       KI.LaunchBounds.push_back({Name, *Val});

llvm/lib/IR/Module.cpp

@@ -322,32 +322,6 @@ void Module::eraseNamedMetadata(NamedMDNode *NMD) {
   eraseNamedMDNode(NMD);
 }
 
-SetVector<Function *> Module::getDeviceKernels() {
-  // TODO: Create a more cross-platform way of determining device kernels.
-  NamedMDNode *MD = getNamedMetadata("nvvm.annotations");
-  SetVector<Function *> Kernels;
-
-  if (!MD)
-    return Kernels;
-
-  for (auto *Op : MD->operands()) {
-    if (Op->getNumOperands() < 2)
-      continue;
-    MDString *KindID = dyn_cast<MDString>(Op->getOperand(1));
-    if (!KindID || KindID->getString() != "kernel")
-      continue;
-
-    Function *KernelFn =
-        mdconst::dyn_extract_or_null<Function>(Op->getOperand(0));
-    if (!KernelFn)
-      continue;
-
-    Kernels.insert(KernelFn);
-  }
-
-  return Kernels;
-}
-
 bool Module::isValidModFlagBehavior(Metadata *MD, ModFlagBehavior &MFB) {
   if (ConstantInt *Behavior = mdconst::dyn_extract_or_null<ConstantInt>(MD)) {
     uint64_t Val = Behavior->getLimitedValue();

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

@@ -5910,16 +5910,34 @@ bool llvm::omp::isOpenMPKernel(Function &Fn) {
 }
 
 KernelSet llvm::omp::getDeviceKernels(Module &M) {
+  // TODO: Create a more cross-platform way of determining device kernels.
+  NamedMDNode *MD = M.getNamedMetadata("nvvm.annotations");
   KernelSet Kernels;
-  for (Kernel K : M.getDeviceKernels()) {
+
+  if (!MD)
+    return Kernels;
+
+  for (auto *Op : MD->operands()) {
+    if (Op->getNumOperands() < 2)
+      continue;
+    MDString *KindID = dyn_cast<MDString>(Op->getOperand(1));
+    if (!KindID || KindID->getString() != "kernel")
+      continue;
+
+    Function *KernelFn =
+        mdconst::dyn_extract_or_null<Function>(Op->getOperand(0));
+    if (!KernelFn)
+      continue;
+
     // We are only interested in OpenMP target regions. Others, such as kernels
     // generated by CUDA but linked together, are not interesting to this pass.
-    if (isOpenMPKernel(*K)) {
+    if (isOpenMPKernel(*KernelFn)) {
       ++NumOpenMPTargetRegionKernels;
-      Kernels.insert(K);
+      Kernels.insert(KernelFn);
     } else
       ++NumNonOpenMPTargetRegionKernels;
   }
+
   return Kernels;
 }
 

llvm/test/Analysis/KernelInfo/linkage.ll

@@ -7,46 +7,32 @@
 target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"
 target triple = "nvptx64-nvidia-cuda"
 
-; CHECK: remark: test.c:3:0: in function 'extNotKer', ExternalNotKernel = 1
-define external void @extNotKer() !dbg !10 {
+; CHECK: remark: test.c:3:0: in function 'f', ExternalNotKernel = 1
+define external void @f() !dbg !10 {
 entry:
   ret void
 }
 
-; CHECK: remark: test.c:13:0: in artificial function 'impNotKer', ExternalNotKernel = 1
-define void @impNotKer() !dbg !20 {
+; CHECK: remark: test.c:13:0: in artificial function 'g', ExternalNotKernel = 1
+define void @g() !dbg !20 {
 entry:
   ret void
 }
 
-; CHECK: remark: test.c:23:0: in artificial function 'weakNotKer', ExternalNotKernel = 0
-define weak void @weakNotKer() !dbg !30 {
+; CHECK: remark: test.c:23:0: in function 'h', ExternalNotKernel = 0
+define external void @h() #0 !dbg !30 {
 entry:
   ret void
 }
 
-; CHECK: remark: test.c:33:0: in function 'extKerAttr', ExternalNotKernel = 0
-define external void @extKerAttr() #0 !dbg !40 {
-entry:
-  ret void
-}
-
-; CHECK: remark: test.c:43:0: in function 'extKer', ExternalNotKernel = 0
-define external void @extKer() !dbg !50 {
-entry:
-  ret void
-}
-
-; CHECK: remark: test.c:53:0: in artificial function 'weakKer', ExternalNotKernel = 0
-define weak void @weakKer() !dbg !60 {
+; CHECK: remark: test.c:33:0: in artificial function 'i', ExternalNotKernel = 0
+define weak void @i() !dbg !40 {
 entry:
   ret void
 }
 
 attributes #0 = { "kernel" }
 
-!nvvm.annotations = !{!42, !52, !62}
-
 !llvm.module.flags = !{!0}
 !llvm.dbg.cu = !{!1}
 
@@ -55,24 +41,11 @@ attributes #0 = { "kernel" }
 !2 = !DIFile(filename: "test.c", directory: "/tmp")
 !3 = !{null}
 !4 = !{}
-
-!10 = distinct !DISubprogram(name: "extNotKer", scope: !2, file: !2, line: 3, type: !11, scopeLine: 3, spFlags: DISPFlagDefinition, unit: !1, retainedNodes: !4)
+!10 = distinct !DISubprogram(name: "f", scope: !2, file: !2, line: 3, type: !11, scopeLine: 3, spFlags: DISPFlagDefinition, unit: !1, retainedNodes: !4)
 !11 = !DISubroutineType(types: !3)
-
-!20 = distinct !DISubprogram(name: "impNotKer", scope: !2, file: !2, line: 13, type: !21, scopeLine: 13, flags: DIFlagArtificial | DIFlagPrototyped, spFlags: DISPFlagLocalToUnit | DISPFlagDefinition, unit: !1, retainedNodes: !4)
+!20 = distinct !DISubprogram(name: "g", scope: !2, file: !2, line: 13, type: !21, scopeLine: 13, flags: DIFlagArtificial | DIFlagPrototyped, spFlags: DISPFlagLocalToUnit | DISPFlagDefinition, unit: !1, retainedNodes: !4)
 !21 = distinct !DISubroutineType(types: !3)
-
-!30 = distinct !DISubprogram(name: "weakNotKer", scope: !2, file: !2, line: 23, type: !31, scopeLine: 23, flags: DIFlagArtificial | DIFlagPrototyped, spFlags: DISPFlagLocalToUnit | DISPFlagDefinition, unit: !1, retainedNodes: !4)
+!30 = distinct !DISubprogram(name: "h", scope: !2, file: !2, line: 23, type: !31, scopeLine: 23, spFlags: DISPFlagDefinition, unit: !1, retainedNodes: !4)
 !31 = distinct !DISubroutineType(types: !3)
-
-!40 = distinct !DISubprogram(name: "extKerAttr", scope: !2, file: !2, line: 33, type: !41, scopeLine: 33, spFlags: DISPFlagDefinition, unit: !1, retainedNodes: !4)
+!40 = distinct !DISubprogram(name: "i", scope: !2, file: !2, line: 33, type: !41, scopeLine: 33, flags: DIFlagArtificial | DIFlagPrototyped, spFlags: DISPFlagLocalToUnit | DISPFlagDefinition, unit: !1, retainedNodes: !4)
 !41 = distinct !DISubroutineType(types: !3)
-!42 = !{ptr @extKerAttr, !"kernel", i32 1}
-
-!50 = distinct !DISubprogram(name: "extKer", scope: !2, file: !2, line: 43, type: !51, scopeLine: 43, spFlags: DISPFlagDefinition, unit: !1, retainedNodes: !4)
-!51 = distinct !DISubroutineType(types: !3)
-!52 = !{ptr @extKer, !"kernel", i32 1}
-
-!60 = distinct !DISubprogram(name: "weakKer", scope: !2, file: !2, line: 53, type: !61, scopeLine: 53, flags: DIFlagArtificial | DIFlagPrototyped, spFlags: DISPFlagLocalToUnit | DISPFlagDefinition, unit: !1, retainedNodes: !4)
-!61 = distinct !DISubroutineType(types: !3)
-!62 = !{ptr @weakKer, !"kernel", i32 1}