[OpenMP] Compute a proper module slice for the CGSCCC pass

[jdoerfert]

Summary

The module slice describes which functions we can analyze and transform
while working on an SCC as part of the CGSCC OpenMPOpt pass. So far, we
simply restricted it to the SCC. In a follow up we will need to have a
bigger scope which is why this patch introduces a proper identification
of the module slice. In short, everything that has a transitive
reference to a function in the SCC or is transitively referenced by one
is fair game.

[jdoerfert]

Limit foreachUse to the SCC and not module slice

[sstefan1]

LGTM

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

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
 
+/// Apply \p CB to all uses of \p F. If \p LookThroughConstantExprUses is
+/// true, constant expression users are not given to \p CB but their uses are
+/// traversed transitively.
+template <typename CBTy>
+static void foreachUse(Function &F, CBTy CB,
+                       bool LookThroughConstantExprUses = true) {
+  SmallVector<Use *, 8> Worklist(make_pointer_range(F.uses()));
+
+  for (unsigned idx = 0; idx < Worklist.size(); ++idx) {
+    Use &U = *Worklist[idx];
+
+    // Allow use in constant bitcasts and simply look through them.
+    if (LookThroughConstantExprUses && isa<ConstantExpr>(U.getUser())) {
+      for (Use &CEU : cast<ConstantExpr>(U.getUser())->uses())
+        Worklist.push_back(&CEU);
+      continue;
+    }
+
+    CB(U);
+  }
+}
+
 /// Helper struct to store tracked ICV values at specif instructions.
 struct ICVValue {
   Instruction *Inst;
   Value *TrackedValue;
 
   ICVValue(Instruction *I, Value *Val) : Inst(I), TrackedValue(Val) {}
 };
 

 namespace {
 
 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,
+                      BumpPtrAllocator &Allocator, SetVector<Function *> &CGSCC,
                       SmallPtrSetImpl<Kernel> &Kernels)
-      : InformationCache(M, AG, Allocator, CGSCC), ModuleSlice(ModuleSlice),
-        OMPBuilder(M), Kernels(Kernels) {
+      : InformationCache(M, AG, Allocator, &CGSCC), OMPBuilder(M),
+        Kernels(Kernels) {
+    initializeModuleSlice(CGSCC);
+
     OMPBuilder.initialize();
     initializeRuntimeFunctions();
     initializeInternalControlVars();
   }
 
   /// Generic information that describes an internal control variable.
   struct InternalControlVarInfo {
     /// The kind, as described by InternalControlVar enum.
struct RuntimeFunctionInfo {
     }
 
   private:
     /// Map from functions to all uses of this runtime function contained in
     /// them.
     DenseMap<Function *, std::shared_ptr<UseVector>> UsesMap;
   };
 
+  /// Initialize the ModuleSlice member based on \p SCC. ModuleSlices contains
+  /// (a subset of) all functions that we can look at during this SCC traversal.
+  /// This includes functions (transitively) called from the SCC and the
+  /// (transitive) callers of SCC functions. We also can look at a function if
+  /// there is a "reference edge", i.a., if the function somehow uses (!=calls)
+  /// a function in the SCC or a caller of a function in the SCC.
+  void initializeModuleSlice(SetVector<Function *> &SCC) {
+    ModuleSlice.insert(SCC.begin(), SCC.end());
+
+    SmallPtrSet<Function *, 16> Seen;
+    SmallVector<Function *, 16> Worklist(SCC.begin(), SCC.end());
+    while (!Worklist.empty()) {
+      Function *F = Worklist.pop_back_val();
+      ModuleSlice.insert(F);
+
+      for (Instruction &I : instructions(*F))
+        if (auto *CB = dyn_cast<CallBase>(&I))
+          if (Function *Callee = CB->getCalledFunction())
+            if (Seen.insert(Callee).second)
+              Worklist.push_back(Callee);
+    }
+
+    Seen.clear();
+    Worklist.append(SCC.begin(), SCC.end());
+    while (!Worklist.empty()) {
+      Function *F = Worklist.pop_back_val();
+      ModuleSlice.insert(F);
+
+      // Traverse all transitive uses.
+      foreachUse(*F, [&](Use &U) {
+        if (auto *UsrI = dyn_cast<Instruction>(U.getUser()))
+          if (Seen.insert(UsrI->getFunction()).second)
+            Worklist.push_back(UsrI->getFunction());
+      });
+    }
+  }
+
   /// The slice of the module we are allowed to look at.
-  SmallPtrSetImpl<Function *> &ModuleSlice;
+  SmallPtrSet<Function *, 8> ModuleSlice;
 
   /// An OpenMP-IR-Builder instance
   OpenMPIRBuilder OMPBuilder;
 
   /// Map from runtime function kind to the runtime function description.
   EnumeratedArray<RuntimeFunctionInfo, RuntimeFunction,
                   RuntimeFunction::OMPRTL___last>
       RFIs;
PreservedAnalyses OpenMPOptPass::run(LazyCallGraph::SCC &C,
                                      CGSCCAnalysisManager &AM,
                                      LazyCallGraph &CG, CGSCCUpdateResult &UR) {
   if (!containsOpenMP(*C.begin()->getFunction().getParent(), OMPInModule))
     return PreservedAnalyses::all();
 
   if (DisableOpenMPOptimizations)
     return PreservedAnalyses::all();
 
-  SmallPtrSet<Function *, 16> ModuleSlice;
   SmallVector<Function *, 16> SCC;
-  for (LazyCallGraph::Node &N : C) {
+  for (LazyCallGraph::Node &N : C)
     SCC.push_back(&N.getFunction());
-    ModuleSlice.insert(SCC.back());
-  }
 
   if (SCC.empty())
     return PreservedAnalyses::all();
 
   FunctionAnalysisManager &FAM =
       AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
 
   AnalysisGetter AG(FAM);
 
   auto OREGetter = [&FAM](Function *F) -> OptimizationRemarkEmitter & {
     return FAM.getResult<OptimizationRemarkEmitterAnalysis>(*F);
   };
 
   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,
-                                OMPInModule.getKernels());
+                                /*CGSCC*/ Functions, 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();
struct OpenMPOptLegacyPass : public CallGraphSCCPass {
   }
 
   bool runOnSCC(CallGraphSCC &CGSCC) override {
     if (!containsOpenMP(CGSCC.getCallGraph().getModule(), OMPInModule))
       return false;
     if (DisableOpenMPOptimizations || skipSCC(CGSCC))
       return false;
 
-    SmallPtrSet<Function *, 16> ModuleSlice;
     SmallVector<Function *, 16> SCC;
     for (CallGraphNode *CGN : CGSCC)
       if (Function *Fn = CGN->getFunction())
-        if (!Fn->isDeclaration()) {
+        if (!Fn->isDeclaration())
           SCC.push_back(Fn);
-          ModuleSlice.insert(Fn);
-        }
 
     if (SCC.empty())
       return false;
 
     CallGraph &CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();
     CGUpdater.initialize(CG, CGSCC);
 
     // Maintain a map of functions to avoid rebuilding the ORE
     DenseMap<Function *, std::unique_ptr<OptimizationRemarkEmitter>> OREMap;
     auto OREGetter = [&OREMap](Function *F) -> OptimizationRemarkEmitter & {
       std::unique_ptr<OptimizationRemarkEmitter> &ORE = OREMap[F];
       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, OMPInModule.getKernels());
+        /*CGSCC*/ Functions, 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();
   }