[PIR] Consult the ParallelRegionInfo before promoting allocas

[jdoerfert]

Summary

The ParallelRegionInfo offers a safe approximation to determine if an
alloca can be promoted. With this patch SROA and Mem2Reg will consult
the ParallelRegionInfo before promotion.

TODO: Not all users of "isAllocaPromotable" are updated. The
DominatorTree and ParallelRegionInfo should not be optional arguments to
that function.

include/llvm/Transforms/Scalar/SROA.h

 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Support/Compiler.h"
 #include <vector>
 
 namespace llvm {
 
+class ParallelRegionInfo;
+
 /// A private "module" namespace for types and utilities used by SROA. These
 /// are implementation details and should not be used by clients.
 namespace sroa LLVM_LIBRARY_VISIBILITY {
 
 class AllocaSliceRewriter;
 class AllocaSlices;
 class Partition;
 class SROALegacyPass;

 /// 3) Finally, this will try to detect a pattern of accesses which map cleanly
 ///    onto insert and extract operations on a vector value, and convert them to
 ///    this form. By doing so, it will enable promotion of vector aggregates to
 ///    SSA vector values.
 class SROA : public PassInfoMixin<SROA> {
   LLVMContext *C = nullptr;
   DominatorTree *DT = nullptr;
   AssumptionCache *AC = nullptr;
+  ParallelRegionInfo *PRI = nullptr;
 
   /// \brief Worklist of alloca instructions to simplify.
   ///
   /// Each alloca in the function is added to this. Each new alloca formed gets
   /// added to it as well to recursively simplify unless that alloca can be
   /// directly promoted. Finally, each time we rewrite a use of an alloca other
   /// the one being actively rewritten, we add it back onto the list if not
   /// already present to ensure it is re-visited.
public:
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 
 private:
   friend class sroa::AllocaSliceRewriter;
   friend class sroa::SROALegacyPass;
 
   /// Helper used by both the public run method and by the legacy pass.
   PreservedAnalyses runImpl(Function &F, DominatorTree &RunDT,
-                            AssumptionCache &RunAC);
+                            AssumptionCache &RunAC, ParallelRegionInfo &RunPRI);
 
   bool presplitLoadsAndStores(AllocaInst &AI, sroa::AllocaSlices &AS);
   AllocaInst *rewritePartition(AllocaInst &AI, sroa::AllocaSlices &AS,
                                sroa::Partition &P);
   bool splitAlloca(AllocaInst &AI, sroa::AllocaSlices &AS);
   bool runOnAlloca(AllocaInst &AI);
   void clobberUse(Use &U);
   void deleteDeadInstructions(SmallPtrSetImpl<AllocaInst *> &DeletedAllocas);
   bool promoteAllocas(Function &F);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_SCALAR_SROA_H

include/llvm/Transforms/Utils/PromoteMemToReg.h

 
 namespace llvm {
 
 template <typename T> class ArrayRef;
 class AllocaInst;
 class DominatorTree;
 class AliasSetTracker;
 class AssumptionCache;
+class ParallelRegionInfo;
 
 /// \brief Return true if this alloca is legal for promotion.
 ///
 /// This is true if there are only loads, stores, and lifetime markers
 /// (transitively) using this alloca. This also enforces that there is only
 /// ever one layer of bitcasts or GEPs between the alloca and the lifetime
 /// markers.
-bool isAllocaPromotable(const AllocaInst *AI);
+///
+/// In case the dominance tree and the parallel region info are _both_ given
+/// we also verify that promoting the alloca does not break invariants of the
+/// parallel regions. See ParallelRegionInfo::isSafeToPromote(*) for details.
+bool isAllocaPromotable(const AllocaInst *AI, const DominatorTree *DT = nullptr,
+                        const ParallelRegionInfo *PRI = nullptr);
 
 /// \brief Promote the specified list of alloca instructions into scalar
 /// registers, inserting PHI nodes as appropriate.
 ///
 /// This function makes use of DominanceFrontier information.  This function
 /// does not modify the CFG of the function at all.  All allocas must be from
 /// the same function.
 ///

lib/Analysis/ParallelRegionInfo.cpp

 bool ParallelRegionInfo::maybeContainedInAny(const Loop *L,
                                              const DominatorTree &DT) const {
   // See ParallelTask::contains(const Loop *, DominatorTree &) for reasoning.
   return maybeContainedInAny(L->getHeader(), DT);
 }
 
 bool ParallelRegionInfo::isSafeToPromote(const AllocaInst &AI,
                                          const DominatorTree &DT) const {
+  if (TopLevelParallelRegions.empty())
+    return true;
 
   // First check if we know that AI is contained in a parallel region.
   ParallelRegion *AIPR = nullptr;
   if (AI.getParent() != &AI.getFunction()->getEntryBlock()) {
     for (ParallelRegion *PR : TopLevelParallelRegions) {
       // In order to not perform a potentially linear contains check we will
       // skip parallel regions that have multi-exit tasks. This is conservative
       // but sound.

lib/Transforms/Scalar/SROA.cpp

 
 #include "llvm/Transforms/Scalar/SROA.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/Loads.h"
+#include "llvm/Analysis/ParallelRegionInfo.h"
 #include "llvm/Analysis/PtrUseVisitor.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
bool SROA::runOnAlloca(AllocaInst &AI) {
   DEBUG(dbgs() << "SROA alloca: " << AI << "\n");
   ++NumAllocasAnalyzed;
 
   // Special case dead allocas, as they're trivial.
   if (AI.use_empty()) {
     AI.eraseFromParent();
     return true;
   }
+
+  // For now we do not perform SROA if it might break a parallel region.
+  // TODO Perform this check on the slices and promote the ones that are safe to
+  //      promote.
+  if (!PRI->isSafeToPromote(AI, *DT))
+    return false;
+
   const DataLayout &DL = AI.getModule()->getDataLayout();
 
   // Skip alloca forms that this analysis can't handle.
   if (AI.isArrayAllocation() || !AI.getAllocatedType()->isSized() ||
       DL.getTypeAllocSize(AI.getAllocatedType()) == 0)
     return false;
 
   bool Changed = false;
bool SROA::promoteAllocas(Function &F) {
 
   DEBUG(dbgs() << "Promoting allocas with mem2reg...\n");
   PromoteMemToReg(PromotableAllocas, *DT, nullptr, AC);
   PromotableAllocas.clear();
   return true;
 }
 
 PreservedAnalyses SROA::runImpl(Function &F, DominatorTree &RunDT,
-                                AssumptionCache &RunAC) {
+                                AssumptionCache &RunAC,
+                                ParallelRegionInfo &RunPRI) {
   DEBUG(dbgs() << "SROA function: " << F.getName() << "\n");
   C = &F.getContext();
   DT = &RunDT;
   AC = &RunAC;
+  PRI = &RunPRI;
 
   BasicBlock &EntryBB = F.getEntryBlock();
   for (BasicBlock::iterator I = EntryBB.begin(), E = std::prev(EntryBB.end());
        I != E; ++I) {
     if (AllocaInst *AI = dyn_cast<AllocaInst>(I))
       Worklist.insert(AI);
   }
 
PreservedAnalyses SROA::runImpl(Function &F, DominatorTree &RunDT,
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
   PA.preserve<GlobalsAA>();
   return PA;
 }
 
 PreservedAnalyses SROA::run(Function &F, FunctionAnalysisManager &AM) {
   return runImpl(F, AM.getResult<DominatorTreeAnalysis>(F),
-                 AM.getResult<AssumptionAnalysis>(F));
+                 AM.getResult<AssumptionAnalysis>(F),
+                 AM.getResult<ParallelRegionAnalysis>(F));
 }
 
 /// A legacy pass for the legacy pass manager that wraps the \c SROA pass.
 ///
 /// This is in the llvm namespace purely to allow it to be a friend of the \c
 /// SROA pass.
 class llvm::sroa::SROALegacyPass : public FunctionPass {
   /// The SROA implementation.
   SROA Impl;
 
 public:
   SROALegacyPass() : FunctionPass(ID) {
     initializeSROALegacyPassPass(*PassRegistry::getPassRegistry());
   }
   bool runOnFunction(Function &F) override {
     if (skipFunction(F))
       return false;
 
     auto PA = Impl.runImpl(
         F, getAnalysis<DominatorTreeWrapperPass>().getDomTree(),
-        getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F));
+        getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F),
+        getAnalysis<ParallelRegionInfoPass>().getParallelRegionInfo());
     return !PA.areAllPreserved();
   }
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<DominatorTreeWrapperPass>();
+    AU.addRequired<ParallelRegionInfoPass>();
     AU.addPreserved<GlobalsAAWrapperPass>();
     AU.setPreservesCFG();
   }
 
   StringRef getPassName() const override { return "SROA"; }
   static char ID;
 };
 
 char SROALegacyPass::ID = 0;
 
 FunctionPass *llvm::createSROAPass() { return new SROALegacyPass(); }
 
 INITIALIZE_PASS_BEGIN(SROALegacyPass, "sroa",
                       "Scalar Replacement Of Aggregates", false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(ParallelRegionInfoPass)
 INITIALIZE_PASS_END(SROALegacyPass, "sroa", "Scalar Replacement Of Aggregates",
                     false, false)

lib/Transforms/Utils/Mem2Reg.cpp

 // This pass is a simple pass wrapper around the PromoteMemToReg function call
 // exposed by the Utils library.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/Mem2Reg.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/ParallelRegionInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils/PromoteMemToReg.h"
 #include "llvm/Transforms/Utils/UnifyFunctionExitNodes.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "mem2reg"
 
 STATISTIC(NumPromoted, "Number of alloca's promoted");
 
 static bool promoteMemoryToRegister(Function &F, DominatorTree &DT,
-                                    AssumptionCache &AC) {
+                                    AssumptionCache &AC,
+                                    ParallelRegionInfo &PRI) {
   std::vector<AllocaInst *> Allocas;
   BasicBlock &BB = F.getEntryBlock(); // Get the entry node for the function
   bool Changed = false;
 
   while (1) {
     Allocas.clear();
 
     // Find allocas that are safe to promote, by looking at all instructions in
     // the entry node
     for (BasicBlock::iterator I = BB.begin(), E = --BB.end(); I != E; ++I)
       if (AllocaInst *AI = dyn_cast<AllocaInst>(I)) // Is it an alloca?
-        if (isAllocaPromotable(AI))
+        if (isAllocaPromotable(AI, &DT, &PRI))
           Allocas.push_back(AI);
 
     if (Allocas.empty())
       break;
 
     PromoteMemToReg(Allocas, DT, nullptr, &AC);
     NumPromoted += Allocas.size();
     Changed = true;
   }
   return Changed;
 }
 
 PreservedAnalyses PromotePass::run(Function &F, FunctionAnalysisManager &AM) {
   auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
   auto &AC = AM.getResult<AssumptionAnalysis>(F);
-  if (!promoteMemoryToRegister(F, DT, AC))
+  auto &PRI = AM.getResult<ParallelRegionAnalysis>(F);
+  if (!promoteMemoryToRegister(F, DT, AC, PRI))
     return PreservedAnalyses::all();
 
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
   return PA;
 }
 
 namespace {
 struct PromoteLegacyPass : public FunctionPass {
   static char ID; // Pass identification, replacement for typeid
   PromoteLegacyPass() : FunctionPass(ID) {
     initializePromoteLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   // runOnFunction - To run this pass, first we calculate the alloca
   // instructions that are safe for promotion, then we promote each one.
   //
   bool runOnFunction(Function &F) override {
     if (skipFunction(F))
       return false;
 
     DominatorTree &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     AssumptionCache &AC =
         getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
-    return promoteMemoryToRegister(F, DT, AC);
+    ParallelRegionInfo &PRI =
+        getAnalysis<ParallelRegionInfoPass>().getParallelRegionInfo();
+    return promoteMemoryToRegister(F, DT, AC, PRI);
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<DominatorTreeWrapperPass>();
+    AU.addRequired<ParallelRegionInfoPass>();
     AU.setPreservesCFG();
   }
   };
 }  // end of anonymous namespace
 
 char PromoteLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(PromoteLegacyPass, "mem2reg", "Promote Memory to "
                                                     "Register",
                       false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(ParallelRegionInfoPass)
 INITIALIZE_PASS_END(PromoteLegacyPass, "mem2reg", "Promote Memory to Register",
                     false, false)
 
 // createPromoteMemoryToRegister - Provide an entry point to create this pass.
 //
 FunctionPass *llvm::createPromoteMemoryToRegisterPass() {
   return new PromoteLegacyPass();
 }

lib/Transforms/Utils/PromoteMemoryToRegister.cpp

 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasSetTracker.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/IteratedDominanceFrontier.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/Analysis/ParallelRegionInfo.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include <algorithm>
 using namespace llvm;
 
 #define DEBUG_TYPE "mem2reg"
 
 STATISTIC(NumLocalPromoted, "Number of alloca's promoted within one block");
 STATISTIC(NumSingleStore,   "Number of alloca's promoted with a single store");
 STATISTIC(NumDeadAlloca,    "Number of dead alloca's removed");
 STATISTIC(NumPHIInsert,     "Number of PHI nodes inserted");
 
-bool llvm::isAllocaPromotable(const AllocaInst *AI) {
+bool llvm::isAllocaPromotable(const AllocaInst *AI, const DominatorTree *DT,
+                              const ParallelRegionInfo *PRI) {
+  // If the dominator tree and parallel region info were given we check if the
+  // promotion would break any parallel region requirements.
+  if (PRI && DT && !PRI->isSafeToPromote(*AI, *DT))
+    return false;
+
   // FIXME: If the memory unit is of pointer or integer type, we can permit
   // assignments to subsections of the memory unit.
   unsigned AS = AI->getType()->getAddressSpace();
 
   // Only allow direct and non-volatile loads and stores...
   for (const User *U : AI->users()) {
     if (const LoadInst *LI = dyn_cast<LoadInst>(U)) {
       // Note that atomic loads can be transformed; atomic semantics do

test/PIR/alloca_promotion.ll

+; RUN: opt -mem2reg -sroa -S < %s | FileCheck %s
+
+declare void @foo();
+declare void @bar(i32);
+
+; Verify we do not promote allocas that are used inside and outside a
+; parallel region.
+define i32 @alloca_used_in_seq_and_par_code() {
+entry:
+; CHECK: %local_alloca = alloca i32
+  %local_alloca = alloca i32
+  fork label %forked, %cont
+
+forked:                                       ; preds = %entry
+; CHECK: store i32 0, i32* %local_alloca
+  store i32 0, i32* %local_alloca
+  call void @foo()
+  halt label %cont
+
+cont:                                         ; preds = %entry, %forked
+  call void @foo()
+; CHECK: store i32 1, i32* %local_alloca
+  store i32 1, i32* %local_alloca
+  join label %join
+
+join:                                         ; preds = %cont
+; CHECK: %val = load i32, i32* %local_alloca
+  %val = load i32, i32* %local_alloca
+  ret i32 %val
+}
+
+; Verify we do not promote allocas even if they are used only inside a parallel
+; region but defined outside.
+define i32 @alloca_used_only_in_par_code() {
+entry:
+; CHECK: alloca i32
+  %local_alloca = alloca i32
+  fork label %forked, %cont
+
+forked:                                       ; preds = %entry
+  store i32 0, i32* %local_alloca
+  call void @foo()
+  halt label %cont
+
+cont:                                         ; preds = %entry, %forked
+  call void @foo()
+  store i32 1, i32* %local_alloca
+  %val = load i32, i32* %local_alloca
+  join label %join
+
+join:                                         ; preds = %cont
+; CHECK: ret i32 %val
+  ret i32 %val
+}
+
+; Verify we do promote allocas that are used only outside a parallel region.
+define i32 @alloca_used_only_in_seq_code() {
+entry:
+; CHECK-NOT: alloca i32
+  %local_alloca = alloca i32
+  store i32 0, i32* %local_alloca
+  fork label %forked, %cont
+
+forked:                                       ; preds = %entry
+  call void @foo()
+  halt label %cont
+
+cont:                                         ; preds = %entry, %forked
+  call void @foo()
+  %val = load i32, i32* %local_alloca
+  join label %join
+
+join:                                         ; preds = %cont
+  store i32 1, i32* %local_alloca
+; CHECK: ret i32 0
+  ret i32 %val
+}