[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 {
 
     /// Return the number of arguments (or the minimal number for variadic
     /// functions).
     size_t getNumArgs() const { return ArgumentTypes.size(); }
 
     /// Run the callback \p CB on each use and forget the use if the result is
     /// true. The callback will be fed the function in which the use was
     /// encountered as second argument.
-    void foreachUse(function_ref<bool(Use &, Function &)> CB) {
-      for (auto &It : UsesMap)
-        foreachUse(CB, It.first, It.second.get());
+    void foreachUse(SmallVectorImpl<Function *> &SCC,
+                    function_ref<bool(Use &, Function &)> CB) {
+      for (Function *F : SCC)
+        foreachUse(CB, F);
     }
 
     /// Run the callback \p CB on each use within the function \p F and forget
     /// the use if the result is true.
-    void foreachUse(function_ref<bool(Use &, Function &)> CB, Function *F,
-                    UseVector *Uses = nullptr) {
+    void foreachUse(function_ref<bool(Use &, Function &)> CB, Function *F) {
       SmallVector<unsigned, 8> ToBeDeleted;
       ToBeDeleted.clear();
 
       unsigned Idx = 0;
-      UseVector &UV = Uses ? *Uses : getOrCreateUseVector(F);
+      UseVector &UV = getOrCreateUseVector(F);
 
       for (Use *U : UV) {
         if (CB(*U, *F))
           ToBeDeleted.push_back(Idx);
         ++Idx;
       }
 
       // Remove the to-be-deleted indices in reverse order as prior
       // modifications will not modify the smaller indices.
       while (!ToBeDeleted.empty()) {
         unsigned Idx = ToBeDeleted.pop_back_val();
         UV[Idx] = UV.back();
         UV.pop_back();
       }
     }
 
   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;
auto DeleteCallCB = [&](Use &U, Function &) {
 
       CGUpdater.removeCallSite(*CI);
       CI->eraseFromParent();
       Changed = true;
       ++NumOpenMPParallelRegionsDeleted;
       return true;
     };
 
-    RFI.foreachUse(DeleteCallCB);
+    RFI.foreachUse(SCC, DeleteCallCB);
 
     return Changed;
   }
 
   /// Try to eliminate runtime calls by reusing existing ones.
   bool deduplicateRuntimeCalls() {
     bool Changed = false;
 
getCombinedIdentFromCallUsesIn(OMPInformationCache::RuntimeFunctionInfo &RFI,
     auto CombineIdentStruct = [&](Use &U, Function &Caller) {
       CallInst *CI = getCallIfRegularCall(U, &RFI);
       if (!CI || &F != &Caller)
         return false;
       Ident = combinedIdentStruct(Ident, CI->getArgOperand(0),
                                   /* GlobalOnly */ true, SingleChoice);
       return false;
     };
-    RFI.foreachUse(CombineIdentStruct);
+    RFI.foreachUse(SCC, CombineIdentStruct);
 
     if (!Ident || !SingleChoice) {
       // The IRBuilder uses the insertion block to get to the module, this is
       // unfortunate but we work around it for now.
       if (!OMPInfoCache.OMPBuilder.getInsertionPoint().getBlock())
         OMPInfoCache.OMPBuilder.updateToLocation(OpenMPIRBuilder::InsertPointTy(
             &F.getEntryBlock(), F.getEntryBlock().begin()));
       // Create a fallback location if non was found.
auto ReplaceAndDeleteCB = [&](Use &U, Function &Caller) {
 
       CGUpdater.removeCallSite(*CI);
       CI->replaceAllUsesWith(ReplVal);
       CI->eraseFromParent();
       ++NumOpenMPRuntimeCallsDeduplicated;
       Changed = true;
       return true;
     };
-    RFI.foreachUse(ReplaceAndDeleteCB);
+    RFI.foreachUse(SCC, ReplaceAndDeleteCB);
 
     return Changed;
   }
 
   /// Collect arguments that represent the global thread id in \p GTIdArgs.
   void collectGlobalThreadIdArguments(SmallSetVector<Value *, 16> &GTIdArgs) {
     // TODO: Below we basically perform a fixpoint iteration with a pessimistic
     //       initialization. We could define an AbstractAttribute instead and
auto AddUserArgs = [&](Value &GTId) {
               if (CallArgOpIsGTId(*Callee, U.getOperandNo(), *CI))
                 GTIdArgs.insert(Callee->getArg(U.getOperandNo()));
     };
 
     // The argument users of __kmpc_global_thread_num calls are GTIds.
     OMPInformationCache::RuntimeFunctionInfo &GlobThreadNumRFI =
         OMPInfoCache.RFIs[OMPRTL___kmpc_global_thread_num];
 
-    GlobThreadNumRFI.foreachUse([&](Use &U, Function &F) {
+    GlobThreadNumRFI.foreachUse(SCC, [&](Use &U, Function &F) {
       if (CallInst *CI = getCallIfRegularCall(U, &GlobThreadNumRFI))
         AddUserArgs(*CI);
       return false;
     });
 
     // Transitively search for more arguments by looking at the users of the
     // ones we know already. During the search the GTIdArgs vector is extended
     // so we cannot cache the size nor can we use a range based for.
auto ReplaceAndDeleteCB = [&](Use &U, Function &Caller) {
 
       A.removeCallSite(CI);
       CI->replaceAllUsesWith(ReplVal);
       CI->eraseFromParent();
       Changed = true;
       return true;
     };
 
-    GetterRFI.foreachUse(ReplaceAndDeleteCB);
+    GetterRFI.foreachUse(ReplaceAndDeleteCB, getAnchorScope());
     return Changed;
   }
 
   // Map of ICV to their values at specific program point.
   EnumeratedArray<SmallSetVector<ICVValue, 4>, InternalControlVar,
                   InternalControlVar::ICV___last>
       ICVValuesMap;
 
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();
   }