[OpenMP] Change OpenMPOpt to check openmp metadata

The metadata added in D102361 introduces a module flag that we can check
to determine if the module was compiled with `-fopenmp` enables. We can
now check for the precense of this instead of scanning the call graph
for OpenMP runtime functions.

Depends on D102361

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D102423

Author: jhuber6

llvm/include/llvm/Transforms/IPO/OpenMPOpt.h

@@ -20,60 +20,27 @@ namespace omp {
 /// Summary of a kernel (=entry point for target offloading).
 using Kernel = Function *;
 
-/// Helper to remember if the module contains OpenMP (runtime calls), to be used
-/// foremost with containsOpenMP.
-struct OpenMPInModule {
-  OpenMPInModule &operator=(bool Found) {
-    if (Found)
-      Value = OpenMPInModule::OpenMP::FOUND;
-    else
-      Value = OpenMPInModule::OpenMP::NOT_FOUND;
-    return *this;
-  }
-  bool isKnown() { return Value != OpenMP::UNKNOWN; }
-  operator bool() { return Value != OpenMP::NOT_FOUND; }
+/// Set of kernels in the module
+using KernelSet = SmallPtrSet<Kernel, 4>;
 
-  /// Does this function \p F contain any OpenMP runtime calls?
-  bool containsOMPRuntimeCalls(Function *F) const {
-    return FuncsWithOMPRuntimeCalls.contains(F);
-  }
+/// Helper to determine if \p M contains OpenMP.
+bool containsOpenMP(Module &M);
 
-  /// Return the known kernels (=GPU entry points) in the module.
-  SmallPtrSetImpl<Kernel> &getKernels() { return Kernels; }
+/// Helper to determine if \p M is a OpenMP target offloading device module.
+bool isOpenMPDevice(Module &M);
 
-  /// Identify kernels in the module and populate the Kernels set.
-  void identifyKernels(Module &M);
-
-private:
-  enum class OpenMP { FOUND, NOT_FOUND, UNKNOWN } Value = OpenMP::UNKNOWN;
-
-  friend bool containsOpenMP(Module &M, OpenMPInModule &OMPInModule);
-
-  /// In which functions are OpenMP runtime calls present?
-  SmallPtrSet<Function *, 32> FuncsWithOMPRuntimeCalls;
-
-  /// Collection of known kernels (=GPU entry points) in the module.
-  SmallPtrSet<Kernel, 8> Kernels;
-};
-
-/// Helper to determine if \p M contains OpenMP (runtime calls).
-bool containsOpenMP(Module &M, OpenMPInModule &OMPInModule);
+/// Get OpenMP device kernels in \p M.
+KernelSet getDeviceKernels(Module &M);
 
 } // namespace omp
 
 /// OpenMP optimizations pass.
 class OpenMPOptPass : public PassInfoMixin<OpenMPOptPass> {
-  /// Helper to remember if the module contains OpenMP (runtime calls).
-  omp::OpenMPInModule OMPInModule;
-
 public:
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
 };
 
 class OpenMPOptCGSCCPass : public PassInfoMixin<OpenMPOptCGSCCPass> {
-  /// Helper to remember if the module contains OpenMP (runtime calls).
-  omp::OpenMPInModule OMPInModule;
-
 public:
   PreservedAnalyses run(LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
                         LazyCallGraph &CG, CGSCCUpdateResult &UR);

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

@@ -1629,7 +1629,7 @@ struct OpenMPOpt {
     for (auto *F : SCC) {
       if (!F->isDeclaration())
         A.getOrCreateAAFor<AAExecutionDomain>(IRPosition::function(*F));
-      if (!OMPInfoCache.Kernels.empty())
+      if (isOpenMPDevice(M))
         A.getOrCreateAAFor<AAHeapToStack>(IRPosition::function(*F));
     }
   }
@@ -2629,17 +2629,18 @@ AAHeapToShared &AAHeapToShared::createForPosition(const IRPosition &IRP,
 }
 
 PreservedAnalyses OpenMPOptPass::run(Module &M, ModuleAnalysisManager &AM) {
-  if (!containsOpenMP(M, OMPInModule))
+  if (!containsOpenMP(M))
     return PreservedAnalyses::all();
-
   if (DisableOpenMPOptimizations)
     return PreservedAnalyses::all();
 
+  KernelSet Kernels = getDeviceKernels(M);
+
   // Create internal copies of each function if this is a kernel Module.
   DenseSet<const Function *> InternalizedFuncs;
-  if (!OMPInModule.getKernels().empty())
+  if (isOpenMPDevice(M))
     for (Function &F : M)
-      if (!F.isDeclaration() && !OMPInModule.getKernels().contains(&F))
+      if (!F.isDeclaration() && !Kernels.contains(&F))
         if (Attributor::internalizeFunction(F, /* Force */ true))
           InternalizedFuncs.insert(&F);
 
@@ -2665,10 +2666,9 @@ PreservedAnalyses OpenMPOptPass::run(Module &M, ModuleAnalysisManager &AM) {
   CallGraphUpdater CGUpdater;
 
   SetVector<Function *> Functions(SCC.begin(), SCC.end());
-  OMPInformationCache InfoCache(M, AG, Allocator, /*CGSCC*/ Functions,
-                                OMPInModule.getKernels());
+  OMPInformationCache InfoCache(M, AG, Allocator, /*CGSCC*/ Functions, Kernels);
 
-  unsigned MaxFixponitIterations = (!OMPInModule.getKernels().empty()) ? 64 : 32;
+  unsigned MaxFixponitIterations = (Kernels.empty()) ? 64 : 32;
   Attributor A(Functions, InfoCache, CGUpdater, nullptr, true, false, MaxFixponitIterations, OREGetter,
                DEBUG_TYPE);
 
@@ -2684,30 +2684,25 @@ PreservedAnalyses OpenMPOptCGSCCPass::run(LazyCallGraph::SCC &C,
                                           CGSCCAnalysisManager &AM,
                                           LazyCallGraph &CG,
                                           CGSCCUpdateResult &UR) {
-  if (!containsOpenMP(*C.begin()->getFunction().getParent(), OMPInModule))
+  if (!containsOpenMP(*C.begin()->getFunction().getParent()))
     return PreservedAnalyses::all();
-
   if (DisableOpenMPOptimizations)
     return PreservedAnalyses::all();
 
   SmallVector<Function *, 16> SCC;
   // If there are kernels in the module, we have to run on all SCC's.
-  bool SCCIsInteresting = !OMPInModule.getKernels().empty();
   for (LazyCallGraph::Node &N : C) {
     Function *Fn = &N.getFunction();
     SCC.push_back(Fn);
-
-    // Do we already know that the SCC contains kernels,
-    // or that OpenMP functions are called from this SCC?
-    if (SCCIsInteresting)
-      continue;
-    // If not, let's check that.
-    SCCIsInteresting |= OMPInModule.containsOMPRuntimeCalls(Fn);
   }
 
-  if (!SCCIsInteresting || SCC.empty())
+  if (SCC.empty())
     return PreservedAnalyses::all();
 
+  Module &M = *C.begin()->getFunction().getParent();
+
+  KernelSet Kernels = getDeviceKernels(M);
+
   FunctionAnalysisManager &FAM =
       AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
 
@@ -2723,9 +2718,9 @@ PreservedAnalyses OpenMPOptCGSCCPass::run(LazyCallGraph::SCC &C,
 
   SetVector<Function *> Functions(SCC.begin(), SCC.end());
   OMPInformationCache InfoCache(*(Functions.back()->getParent()), AG, Allocator,
-                                /*CGSCC*/ Functions, OMPInModule.getKernels());
+                                /*CGSCC*/ Functions, Kernels);
 
-  unsigned MaxFixponitIterations = (!OMPInModule.getKernels().empty()) ? 64 : 32;
+  unsigned MaxFixponitIterations = (isOpenMPDevice(M)) ? 64 : 32;
   Attributor A(Functions, InfoCache, CGUpdater, nullptr, false, true, MaxFixponitIterations, OREGetter,
                DEBUG_TYPE);
 
@@ -2741,7 +2736,6 @@ namespace {
 
 struct OpenMPOptCGSCCLegacyPass : public CallGraphSCCPass {
   CallGraphUpdater CGUpdater;
-  OpenMPInModule OMPInModule;
   static char ID;
 
   OpenMPOptCGSCCLegacyPass() : CallGraphSCCPass(ID) {
@@ -2752,38 +2746,27 @@ struct OpenMPOptCGSCCLegacyPass : public CallGraphSCCPass {
     CallGraphSCCPass::getAnalysisUsage(AU);
   }
 
-  bool doInitialization(CallGraph &CG) override {
-    // Disable the pass if there is no OpenMP (runtime call) in the module.
-    containsOpenMP(CG.getModule(), OMPInModule);
-    return false;
-  }
-
   bool runOnSCC(CallGraphSCC &CGSCC) override {
-    if (!containsOpenMP(CGSCC.getCallGraph().getModule(), OMPInModule))
+    if (!containsOpenMP(CGSCC.getCallGraph().getModule()))
       return false;
     if (DisableOpenMPOptimizations || skipSCC(CGSCC))
       return false;
 
     SmallVector<Function *, 16> SCC;
     // If there are kernels in the module, we have to run on all SCC's.
-    bool SCCIsInteresting = !OMPInModule.getKernels().empty();
     for (CallGraphNode *CGN : CGSCC) {
       Function *Fn = CGN->getFunction();
       if (!Fn || Fn->isDeclaration())
         continue;
       SCC.push_back(Fn);
-
-      // Do we already know that the SCC contains kernels,
-      // or that OpenMP functions are called from this SCC?
-      if (SCCIsInteresting)
-        continue;
-      // If not, let's check that.
-      SCCIsInteresting |= OMPInModule.containsOMPRuntimeCalls(Fn);
     }
 
-    if (!SCCIsInteresting || SCC.empty())
+    if (SCC.empty())
       return false;
 
+    Module &M = CGSCC.getCallGraph().getModule();
+    KernelSet Kernels = getDeviceKernels(M);
+
     CallGraph &CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();
     CGUpdater.initialize(CG, CGSCC);
 
@@ -2799,11 +2782,11 @@ struct OpenMPOptCGSCCLegacyPass : public CallGraphSCCPass {
     AnalysisGetter AG;
     SetVector<Function *> Functions(SCC.begin(), SCC.end());
     BumpPtrAllocator Allocator;
-    OMPInformationCache InfoCache(
-        *(Functions.back()->getParent()), AG, Allocator,
-        /*CGSCC*/ Functions, OMPInModule.getKernels());
+    OMPInformationCache InfoCache(*(Functions.back()->getParent()), AG,
+                                  Allocator,
+                                  /*CGSCC*/ Functions, Kernels);
 
-    unsigned MaxFixponitIterations = (!OMPInModule.getKernels().empty()) ? 64 : 32;
+    unsigned MaxFixponitIterations = (isOpenMPDevice(M)) ? 64 : 32;
     Attributor A(Functions, InfoCache, CGUpdater, nullptr, false, true,
                  MaxFixponitIterations, OREGetter, DEBUG_TYPE);
 
@@ -2816,11 +2799,13 @@ struct OpenMPOptCGSCCLegacyPass : public CallGraphSCCPass {
 
 } // end anonymous namespace
 
-void OpenMPInModule::identifyKernels(Module &M) {
-
+KernelSet llvm::omp::getDeviceKernels(Module &M) {
+  // TODO: Create a more cross-platform way of determining device kernels.
   NamedMDNode *MD = M.getOrInsertNamedMetadata("nvvm.annotations");
+  KernelSet Kernels;
+
   if (!MD)
-    return;
+    return Kernels;
 
   for (auto *Op : MD->operands()) {
     if (Op->getNumOperands() < 2)
@@ -2838,38 +2823,24 @@ void OpenMPInModule::identifyKernels(Module &M) {
 
     Kernels.insert(KernelFn);
   }
-}
 
-bool llvm::omp::containsOpenMP(Module &M, OpenMPInModule &OMPInModule) {
-  if (OMPInModule.isKnown())
-    return OMPInModule;
+  return Kernels;
+}
 
-  auto RecordFunctionsContainingUsesOf = [&](Function *F) {
-    for (User *U : F->users())
-      if (auto *I = dyn_cast<Instruction>(U))
-        OMPInModule.FuncsWithOMPRuntimeCalls.insert(I->getFunction());
-  };
+bool llvm::omp::containsOpenMP(Module &M) {
+  Metadata *MD = M.getModuleFlag("openmp");
+  if (!MD)
+    return false;
 
-  // MSVC doesn't like long if-else chains for some reason and instead just
-  // issues an error. Work around it..
-  do {
-#define OMP_RTL(_Enum, _Name, ...)                                             \
-  if (Function *F = M.getFunction(_Name)) {                                    \
-    RecordFunctionsContainingUsesOf(F);                                        \
-    OMPInModule = true;                                                        \
-  }
-#include "llvm/Frontend/OpenMP/OMPKinds.def"
-  } while (false);
+  return true;
+}
 
-  // Identify kernels once. TODO: We should split the OMPInformationCache into a
-  // module and an SCC part. The kernel information, among other things, could
-  // go into the module part.
-  if (OMPInModule.isKnown() && OMPInModule) {
-    OMPInModule.identifyKernels(M);
-    return true;
-  }
+bool llvm::omp::isOpenMPDevice(Module &M) {
+  Metadata *MD = M.getModuleFlag("openmp-device");
+  if (!MD)
+    return false;
 
-  return OMPInModule = false;
+  return true;
 }
 
 char OpenMPOptCGSCCLegacyPass::ID = 0;

llvm/test/Transforms/OpenMP/add_attributes.ll

@@ -1739,3 +1739,6 @@ attributes #0 = { noinline }
 ; OPTIMISTIC: ; Function Attrs: convergent noinline nounwind 
 ; OPTIMISTIC-NEXT: declare void @__kmpc_barrier_simple_spmd(%struct.ident_t* nocapture nofree readonly, i32)
 
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 7, !"openmp", i32 50}

llvm/test/Transforms/OpenMP/add_attributes_amdgcn.ll

@@ -26,3 +26,7 @@ declare void @__kmpc_syncwarp(i64)
 
 ; OPTIMISTIC: ; Function Attrs: convergent nounwind
 ; OPTIMISTIC-NEXT: declare void @__kmpc_syncwarp(i64)
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 7, !"openmp", i32 50}

llvm/test/Transforms/OpenMP/deduplication.ll

@@ -221,3 +221,7 @@ entry:
   call void @useI32(i32 %tid5)
   ret void
 }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 7, !"openmp", i32 50}

llvm/test/Transforms/OpenMP/deduplication_remarks.ll

@@ -30,7 +30,7 @@ declare !dbg !4 void @useI32(i32) local_unnamed_addr
 declare void @llvm.dbg.value(metadata, metadata, metadata)
 
 !llvm.dbg.cu = !{!0}
-!llvm.module.flags = !{!8, !9, !10, !11, !12}
+!llvm.module.flags = !{!8, !9, !10, !11, !12, !29}
 !llvm.ident = !{!13}
 
 !0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, producer: "clang version 10.0.0 ", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, retainedTypes: !3, splitDebugInlining: false, nameTableKind: None)
@@ -62,3 +62,4 @@ declare void @llvm.dbg.value(metadata, metadata, metadata)
 !26 = !DILocation(line: 9, column: 10, scope: !14)
 !27 = !DILocation(line: 10, column: 2, scope: !14)
 !28 = !DILocation(line: 13, column: 1, scope: !14)
+!29 = !{i32 7, !"openmp", i32 50}

llvm/test/Transforms/OpenMP/globalization_remarks.ll

@@ -4,13 +4,14 @@ source_filename = "declare_target_codegen_globalization.cpp"
 target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"
 target triple = "nvptx64"
 
+; CHECK: remark: globalization_remarks.c:5:7: Could not move globalized variable to the stack. Variable is potentially captured.
 ; CHECK: remark: globalization_remarks.c:5:7: Found thread data sharing on the GPU. Expect degraded performance due to data globalization.
 
 @S = external local_unnamed_addr global i8*
 
 define void @foo() {
 entry:
-  %0 = call i8* @__kmpc_alloc_shared(i64 4), !dbg !8
+  %0 = call i8* @__kmpc_alloc_shared(i64 4), !dbg !10
   %x_on_stack = bitcast i8* %0 to i32*
   %1 = bitcast i32* %x_on_stack to i8*
   call void @share(i8* %1)
@@ -30,13 +31,17 @@ declare void @__kmpc_free_shared(i8*)
 
 
 !llvm.dbg.cu = !{!0}
-!llvm.module.flags = !{!3, !4}
+!llvm.module.flags = !{!3, !4, !5, !6}
+!nvvm.annotations = !{!7, !8}
 
 !0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, producer: "clang version 12.0.0", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, splitDebugInlining: false, nameTableKind: None)
 !1 = !DIFile(filename: "globalization_remarks.c", directory: "/tmp/globalization_remarks.c")
 !2 = !{}
 !3 = !{i32 2, !"Debug Info Version", i32 3}
 !4 = !{i32 1, !"wchar_size", i32 4}
-!6 = distinct !DISubprogram(name: "foo", scope: !1, file: !1, line: 1, type: !7, scopeLine: 1, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !2)
-!7 = !DISubroutineType(types: !2)
-!8 = !DILocation(line: 5, column: 7, scope: !6)
+!5 = !{i32 7, !"openmp", i32 50}
+!6 = !{i32 7, !"openmp-device", i32 50}
+!7 = !{void ()* @foo, !"kernel", i32 1}
+!8 = distinct !DISubprogram(name: "foo", scope: !1, file: !1, line: 1, type: !9, scopeLine: 1, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !2)
+!9 = !DISubroutineType(types: !2)
+!10 = !DILocation(line: 5, column: 7, scope: !8)

llvm/test/Transforms/OpenMP/gpu_kernel_detection_remarks.ll

@@ -19,9 +19,11 @@ define void @non_kernel() {
 ; Needed to trigger the openmp-opt pass
 declare dso_local void @__kmpc_kernel_prepare_parallel(i8*)
 
+!llvm.module.flags = !{!4}
 !nvvm.annotations = !{!2, !0, !1, !3, !1, !2}
 
 !0 = !{void ()* @kernel1, !"kernel", i32 1}
 !1 = !{void ()* @non_kernel, !"non_kernel", i32 1}
 !2 = !{null, !"align", i32 1}
 !3 = !{void ()* @kernel2, !"kernel", i32 1}
+!4 = !{i32 7, !"openmp", i32 50}

llvm/test/Transforms/OpenMP/gpu_state_machine_function_ptr_replacement.ll

@@ -275,5 +275,8 @@ declare void @__kmpc_kernel_end_parallel()
 
 
 !nvvm.annotations = !{!1}
+!llvm.module.flags = !{!2, !3}
 
 !1 = !{void ()* @__omp_offloading_50_6dfa0f01_foo_l6, !"kernel", i32 1}
+!2 = !{i32 7, !"openmp", i32 50}
+!3 = !{i32 7, !"openmp-device", i32 50}