[OpenMP] Identify GPU kernels (aka. OpenMP target regions)

[jdoerfert]

Summary

We now identify GPU kernels, that is entry points into the GPU code.
These kernels (can) correspond to OpenMP target regions. With this patch
we identify and on request print them via remarks.

[tianshilei1992]

LGTM but I would like others to take a look.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1186

This line of change looks not related to this patch.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1186

It is. I needed to get rid of the return statements and I wanted to keep the "early exit" out of the if-cascade. Entrance: else.

[JonChesterfield]

I think there's slightly more code here than is necessary.

Specifically, I think identifyKernels should return SmallPtrSetImpl<Kernel> instead of populating a member variable which can later be accessed. With a rename, proposing:
SmallPtrSetImpl<Kernel> getKernels(Module &M){/*roughly contents of current identifyKernels */}

The cache then stores the set by value instead of by reference. Less state lying around, can't accidentally add multiple copies of the name to a single set. Depending on the control flow we might look up the metadata more than once, but that seems fine given it usually goes in a cache.

Thoughts?

[jdoerfert]

In D83269#2137745, @JonChesterfield wrote:

I think there's slightly more code here than is necessary.

Specifically, I think identifyKernels should return SmallPtrSetImpl<Kernel> instead of populating a member variable which can later be accessed. With a rename, proposing:
SmallPtrSetImpl<Kernel> getKernels(Module &M){/*roughly contents of current identifyKernels */}

The cache then stores the set by value instead of by reference. Less state lying around, can't accidentally add multiple copies of the name to a single set. Depending on the control flow we might look up the metadata more than once, but that seems fine given it usually goes in a cache.

Thoughts?

We will end up looking at it once per SCC in the program, per invocation of the pass. I would prefer to cache module wide information explicitly and this was the "smallest" solution for this for now.
I can do recompute but the nvvm.annotations has ~100 (non-kernel) entries from the device runtime we'll have to go through every time.

[JonChesterfield]

Fair enough, stateful it is then.

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

 #include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 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; }
 
+  /// Return the known kernels (=GPU entry points) in the module.
+  SmallPtrSetImpl<Kernel> &getKernels() { return Kernels; }
+
+  /// 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;
+
+  /// 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);
 
 } // namespace omp
 
 /// OpenMP optimizations pass.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 
 static cl::opt<bool> DisableOpenMPOptimizations(
     "openmp-opt-disable", cl::ZeroOrMore,
     cl::desc("Disable OpenMP specific optimizations."), cl::Hidden,
     cl::init(false));
 
 static cl::opt<bool> PrintICVValues("openmp-print-icv-values", cl::init(false),
                                     cl::Hidden);
+static cl::opt<bool> PrintOpenMPKernels("openmp-print-gpu-kernels",
+                                        cl::init(false), cl::Hidden);
 
 STATISTIC(NumOpenMPRuntimeCallsDeduplicated,
           "Number of OpenMP runtime calls deduplicated");
 STATISTIC(NumOpenMPParallelRegionsDeleted,
           "Number of OpenMP parallel regions deleted");
 STATISTIC(NumOpenMPRuntimeFunctionsIdentified,
           "Number of OpenMP runtime functions identified");
 STATISTIC(NumOpenMPRuntimeFunctionUsesIdentified,
           "Number of OpenMP runtime function uses identified");
+STATISTIC(NumOpenMPTargetRegionKernels,
+          "Number of OpenMP target region entry points (=kernels) identified");
 
 #if !defined(NDEBUG)
 static constexpr auto TAG = "[" DEBUG_TYPE "]";
 #endif
 
 /// Helper struct to store tracked ICV values at specif instructions.
 struct ICVValue {
   Instruction *Inst;

 
 struct AAICVTracker;
 
 /// OpenMP specific information. For now, stores RFIs and ICVs also needed for
 /// Attributor runs.
 struct OMPInformationCache : public InformationCache {
   OMPInformationCache(Module &M, AnalysisGetter &AG,
                       BumpPtrAllocator &Allocator, SetVector<Function *> *CGSCC,
-                      SmallPtrSetImpl<Function *> &ModuleSlice)
+                      SmallPtrSetImpl<Function *> &ModuleSlice,
+                      SmallPtrSetImpl<Kernel> &Kernels)
       : InformationCache(M, AG, Allocator, CGSCC), ModuleSlice(ModuleSlice),
-        OMPBuilder(M) {
+        OMPBuilder(M), Kernels(Kernels) {
     OMPBuilder.initialize();
     initializeRuntimeFunctions();
     initializeInternalControlVars();
   }
 
   /// Generic information that describes an internal control variable.
   struct InternalControlVarInfo {
     /// The kind, as described by InternalControlVar enum.
if (declMatchesRTFTypes(F, OMPBuilder._ReturnType, ArgsTypes)) {           \
                  << " different functions.\n";                                 \
       });                                                                      \
     }                                                                          \
   }
 #include "llvm/Frontend/OpenMP/OMPKinds.def"
 
     // TODO: We should attach the attributes defined in OMPKinds.def.
   }
+
+  /// Collection of known kernels (\see Kernel) in the module.
+  SmallPtrSetImpl<Kernel> &Kernels;
 };
 
 struct OpenMPOpt {
 
   using OptimizationRemarkGetter =
       function_ref<OptimizationRemarkEmitter &(Function *)>;
 
   OpenMPOpt(SmallVectorImpl<Function *> &SCC, CallGraphUpdater &CGUpdater,
             OptimizationRemarkGetter OREGetter,
             OMPInformationCache &OMPInfoCache, Attributor &A)
       : M(*(*SCC.begin())->getParent()), SCC(SCC), CGUpdater(CGUpdater),
         OREGetter(OREGetter), OMPInfoCache(OMPInfoCache), A(A) {}
 
   /// Run all OpenMP optimizations on the underlying SCC/ModuleSlice.
   bool run() {
     if (SCC.empty())
       return false;
 
     bool Changed = false;
 
     LLVM_DEBUG(dbgs() << TAG << "Run on SCC with " << SCC.size()
                       << " functions in a slice with "
                       << OMPInfoCache.ModuleSlice.size() << " functions\n");
 
-    /// Print initial ICV values for testing.
-    /// FIXME: This should be done from the Attributor once it is added.
-    if (PrintICVValues) {
-      InternalControlVar ICVs[] = {ICV_nthreads, ICV_active_levels, ICV_cancel};
+    if (PrintICVValues)
+      printICVs();
+    if (PrintOpenMPKernels)
+      printKernels();
+
+    Changed |= runAttributor();
+
+    // Recollect uses, in case Attributor deleted any.
+    OMPInfoCache.recollectUses();
+
+    Changed |= deduplicateRuntimeCalls();
+    Changed |= deleteParallelRegions();
+
+    return Changed;
+  }
+
+  void printICVs() const {
 
-      for (Function *F : OMPInfoCache.ModuleSlice) {
-        for (auto ICV : ICVs) {
-          auto ICVInfo = OMPInfoCache.ICVs[ICV];
-          auto Remark = [&](OptimizationRemark OR) {
-            return OR << "OpenMP ICV " << ore::NV("OpenMPICV", ICVInfo.Name)
-                      << " Value: "
-                      << (ICVInfo.InitValue
-                              ? ICVInfo.InitValue->getValue().toString(10, true)
-                              : "IMPLEMENTATION_DEFINED");
-          };
 
-          emitRemarkOnFunction(F, "OpenMPICVTracker", Remark);
-        }
-      }
-    }
 
-    Changed |= runAttributor();
-
-    // Recollect uses, in case Attributor deleted any.
-    OMPInfoCache.recollectUses();
+  /// Print OpenMP GPU kernels for testing.
+  void printKernels() const {
+    for (Function *F : SCC) {
+      if (!OMPInfoCache.Kernels.count(F))
+        continue;
 
-    Changed |= deduplicateRuntimeCalls();
-    Changed |= deleteParallelRegions();
+      auto Remark = [&](OptimizationRemark OR) {
+        return OR << "OpenMP GPU kernel "
+                  << ore::NV("OpenMPGPUKernel", F->getName()) << "\n";
+      };
 
-    return Changed;
+      emitRemarkOnFunction(F, "OpenMPGPU", Remark);
+    }
   }
 
   /// Return the call if \p U is a callee use in a regular call. If \p RFI is
   /// given it has to be the callee or a nullptr is returned.
   static CallInst *getCallIfRegularCall(
       Use &U, OMPInformationCache::RuntimeFunctionInfo *RFI = nullptr) {
     CallInst *CI = dyn_cast<CallInst>(U.getUser());
     if (CI && CI->isCallee(&U) && !CI->hasOperandBundles() &&
private:
   ///   - OptimizationRemarkAnalysis to provide additional information about an
   ///     optimization attempt
   ///
   /// The remark is built using a callback function provided by the caller that
   /// takes a RemarkKind as input and returns a RemarkKind.
   template <typename RemarkKind,
             typename RemarkCallBack = function_ref<RemarkKind(RemarkKind &&)>>
   void emitRemark(Instruction *Inst, StringRef RemarkName,
-                  RemarkCallBack &&RemarkCB) {
+                  RemarkCallBack &&RemarkCB) const {
     Function *F = Inst->getParent()->getParent();
     auto &ORE = OREGetter(F);
 
     ORE.emit(
         [&]() { return RemarkCB(RemarkKind(DEBUG_TYPE, RemarkName, Inst)); });
   }
 
   /// Emit a remark on a function. Since only OptimizationRemark is supporting
   /// this, it can't be made generic.
-  void emitRemarkOnFunction(
-      Function *F, StringRef RemarkName,
-      function_ref<OptimizationRemark(OptimizationRemark &&)> &&RemarkCB) {
+  void
+  emitRemarkOnFunction(Function *F, StringRef RemarkName,
+                       function_ref<OptimizationRemark(OptimizationRemark &&)>
+                           &&RemarkCB) const {
     auto &ORE = OREGetter(F);
 
     ORE.emit([&]() {
       return RemarkCB(OptimizationRemark(DEBUG_TYPE, RemarkName, F));
     });
   }
 
   /// The underlying module.
PreservedAnalyses OpenMPOptPass::run(LazyCallGraph::SCC &C,
   };
 
   CallGraphUpdater CGUpdater;
   CGUpdater.initialize(CG, C, AM, UR);
 
   SetVector<Function *> Functions(SCC.begin(), SCC.end());
   BumpPtrAllocator Allocator;
   OMPInformationCache InfoCache(*(Functions.back()->getParent()), AG, Allocator,
-                                /*CGSCC*/ &Functions, ModuleSlice);
+                                /*CGSCC*/ &Functions, ModuleSlice,
+                                OMPInModule.getKernels());
 
   Attributor A(Functions, InfoCache, CGUpdater);
 
   // TODO: Compute the module slice we are allowed to look at.
   OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
   bool Changed = OMPOpt.run();
   (void)Changed;
   return PreservedAnalyses::all();
auto OREGetter = [&OREMap](Function *F) -> OptimizationRemarkEmitter & {
       if (!ORE)
         ORE = std::make_unique<OptimizationRemarkEmitter>(F);
       return *ORE;
     };
 
     AnalysisGetter AG;
     SetVector<Function *> Functions(SCC.begin(), SCC.end());
     BumpPtrAllocator Allocator;
-    OMPInformationCache InfoCache(*(Functions.back()->getParent()), AG,
-                                  Allocator,
-                                  /*CGSCC*/ &Functions, ModuleSlice);
+    OMPInformationCache InfoCache(
+        *(Functions.back()->getParent()), AG, Allocator,
+        /*CGSCC*/ &Functions, ModuleSlice, OMPInModule.getKernels());
 
     Attributor A(Functions, InfoCache, CGUpdater);
 
     // TODO: Compute the module slice we are allowed to look at.
     OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
     return OMPOpt.run();
   }
 
   bool doFinalization(CallGraph &CG) override { return CGUpdater.finalize(); }
 };
 
 } // end anonymous namespace
 
+void OpenMPInModule::identifyKernels(Module &M) {
+
+  NamedMDNode *MD = M.getOrInsertNamedMetadata("nvvm.annotations");
+  if (!MD)
+    return;
+
+  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;
+
+    ++NumOpenMPTargetRegionKernels;
+
+    Kernels.insert(KernelFn);
+  }
+}
+
 bool llvm::omp::containsOpenMP(Module &M, OpenMPInModule &OMPInModule) {
   if (OMPInModule.isKnown())
     return OMPInModule;
-
 #define OMP_RTL(_Enum, _Name, ...)                                             \
-  if (M.getFunction(_Name))                                                    \
-    return OMPInModule = true;
+  else if (M.getFunction(_Name)) OMPInModule = true;
 #include "llvm/Frontend/OpenMP/OMPKinds.def"
+
+  // 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;
+  }
+
   return OMPInModule = false;
 }
 
 char OpenMPOptLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(OpenMPOptLegacyPass, "openmpopt",
                       "OpenMP specific optimizations", false, false)
 INITIALIZE_PASS_DEPENDENCY(CallGraphWrapperPass)
 INITIALIZE_PASS_END(OpenMPOptLegacyPass, "openmpopt",
                     "OpenMP specific optimizations", false, false)
 
 Pass *llvm::createOpenMPOptLegacyPass() { return new OpenMPOptLegacyPass(); }

llvm/test/Transforms/OpenMP/gpu_kernel_detection_remarks.ll

+; RUN: opt -passes=openmpopt -pass-remarks=openmp-opt -openmp-print-gpu-kernels -disable-output < %s 2>&1 | FileCheck %s --implicit-check-not=non_kernel
+; RUN: opt        -openmpopt -pass-remarks=openmp-opt -openmp-print-gpu-kernels -disable-output < %s 2>&1 | FileCheck %s --implicit-check-not=non_kernel
+
+; CHECK-DAG: remark: <unknown>:0:0: OpenMP GPU kernel kernel1
+; CHECK-DAG: remark: <unknown>:0:0: OpenMP GPU kernel kernel2
+
+define void @kernel1() {
+  ret void
+}
+
+define void @kernel2() {
+  ret void
+}
+
+define void @non_kernel() {
+  ret void
+}
+
+; Needed to trigger the openmp-opt pass
+declare dso_local void @__kmpc_kernel_prepare_parallel(i8*)
+
+!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}