Reland [SROA] Maintain shadow/backing alloca when some slices are noncapturnig read-only calls to allow alloca partitioning/promotion

This is inspired by the original variant of D109749 by Graham Hunter,
but is a more general version.

Roughly, instead of promoting the alloca, we call it
a shadow/backing alloca, go through all it's slices,
clone(!) instructions that operated on it,
but make them operate on the cloned alloca,
and promote cloned alloca instead.

This keeps the shadow/backing alloca, and all the original instructions
around, which results in said shadow/backing alloca being
a perfect mirror/representation of the promoted alloca's content,
so calls that take the alloca as arguments (non-capturingly!)
can be supported.

For now, we require that the calls also don't modify the alloca's content,
but that is only to simplify the initial implementation,
and that will be supported in a follow-up.

Overall, this leads to *smaller* codesize:
https://llvm-compile-time-tracker.com/compare.php?from=a8b4f5bbab62091835205f3d648902432a4a5b58&to=aeae054055b125b011c1122f82c86457e159436f&stat=size-total
and is roughly neutral compile-time wise:
https://llvm-compile-time-tracker.com/compare.php?from=a8b4f5bbab62091835205f3d648902432a4a5b58&to=aeae054055b125b011c1122f82c86457e159436f&stat=instructions

This relands commit 703240c,
that was reverted by commit 7405581,
because the assertion `isa<LoadInst>(OrigInstr)` didn't hold in practice,
as the newly added test `@select_of_ptrs` shows:
If the pointers into alloca are used by select's/PHI's, then even if
we manage to fracture the alloca, some sub-alloca's will likely remain.
And if there are any non-capturing calls, then we will also decide to
keep the original backing alloca around, and we suddenly ~doubled
the alloca size, and the amount of memory traffic.
I'm not sure if this is a problem or we could live with it,
but let's leave that for later...

Reviewed By: djtodoro

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

Author: LebedevRI

llvm/lib/Transforms/Scalar/SROA.cpp

@@ -167,6 +167,7 @@ class Slice {
   void makeUnsplittable() { UseAndIsSplittable.setInt(false); }
 
   Use *getUse() const { return UseAndIsSplittable.getPointer(); }
+  void setUse(Use *U) { UseAndIsSplittable.setPointer(U); }
 
   bool isDead() const { return getUse() == nullptr; }
   void kill() { UseAndIsSplittable.setPointer(nullptr); }
@@ -218,7 +219,7 @@ class Slice {
 class llvm::sroa::AllocaSlices {
 public:
   /// Construct the slices of a particular alloca.
-  AllocaSlices(const DataLayout &DL, AllocaInst &AI);
+  AllocaSlices(const DataLayout &DL, AllocaInst &AI, bool &Changed);
 
   /// Test whether a pointer to the allocation escapes our analysis.
   ///
@@ -270,6 +271,12 @@ class llvm::sroa::AllocaSlices {
     return DeadUseIfPromotable;
   }
 
+  void forgetTheDead() {
+    DeadUsers.clear();
+    DeadUseIfPromotable.clear();
+    DeadOperands.clear();
+  };
+
   /// Access the dead operands referring to this alloca.
   ///
   /// These are operands which have cannot actually be used to refer to the
@@ -295,11 +302,21 @@ class llvm::sroa::AllocaSlices {
 
   friend class AllocaSlices::SliceBuilder;
 
+  void formBackingAlloca(AllocaInst *AI, bool &Changed);
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Handle to alloca instruction to simplify method interfaces.
   AllocaInst &AI;
 #endif
 
+  /// Certain escaping uses of an alloca (non-capturing-ones)
+  /// do not prevent promotion, but force retention of the alloca.
+  /// This records if there are any such uses.
+  bool NeedsBackingAlloca = false;
+
+  /// Track if there are any `select`s/PHI's involving the alloca pointers.
+  bool HavePHINodesOrSelectInstrs = false;
+
   /// The instruction responsible for this alloca not having a known set
   /// of slices.
   ///
@@ -1055,18 +1072,35 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
       return;
     }
 
+    AS.HavePHINodesOrSelectInstrs = true;
+    if (AS.NeedsBackingAlloca && AS.HavePHINodesOrSelectInstrs)
+      return PI.setAborted(&I);
+
     insertUse(I, Offset, Size);
   }
 
   void visitPHINode(PHINode &PN) { visitPHINodeOrSelectInst(PN); }
 
   void visitSelectInst(SelectInst &SI) { visitPHINodeOrSelectInst(SI); }
 
+  void visitCallBase(CallBase &CB) {
+    if (!IsOffsetKnown || !CB.doesNotCapture(U->getOperandNo()))
+      return PI.setAborted(&CB);
+    // If we know that the callee does not retain the pointer,
+    // then it does not prevent SROA, although we have to workaround this.
+    // However, for now, only allow uses, that, at most, read from said memory.
+    if (!CB.onlyReadsMemory() && !CB.onlyReadsMemory(U->getOperandNo()))
+      return PI.setAborted(&CB);
+    AS.NeedsBackingAlloca = true;
+    if (AS.NeedsBackingAlloca && AS.HavePHINodesOrSelectInstrs)
+      return PI.setAborted(&CB);
+  }
+
   /// Disable SROA entirely if there are unhandled users of the alloca.
   void visitInstruction(Instruction &I) { PI.setAborted(&I); }
 };
 
-AllocaSlices::AllocaSlices(const DataLayout &DL, AllocaInst &AI)
+AllocaSlices::AllocaSlices(const DataLayout &DL, AllocaInst &AI, bool &Changed)
     :
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
       AI(AI),
@@ -1083,6 +1117,10 @@ AllocaSlices::AllocaSlices(const DataLayout &DL, AllocaInst &AI)
     return;
   }
 
+  // We may have found that the pointer to the AI escapes, but isn't captured.
+  if (NeedsBackingAlloca)
+    formBackingAlloca(&AI, Changed);
+
   llvm::erase_if(Slices, [](const Slice &S) { return S.isDead(); });
 
   // Sort the uses. This arranges for the offsets to be in ascending order,
@@ -3587,6 +3625,122 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
 
 } // end anonymous namespace
 
+/// Apparently, we can promote the alloca, but some uses of the alloca
+/// are calls (that don't capture it's address), which require for the
+/// trace alloca to remain. To do so, we must form a new "backing" alloca,
+/// which will be kept as an up-to-date backup of the to-be-promoted-alloca's
+/// content, and used in it's place in these non-capturing calls.
+/// FIXME: support non-readonly non-capturing calls.
+void AllocaSlices::formBackingAlloca(AllocaInst *AllocaToPromote,
+                                     bool &Changed) {
+  assert(NeedsBackingAlloca &&
+         "Should not be called if there is no need to rewrite.");
+
+  // We are going to preserve all of the original instructions that were
+  // operating on the original alloca, so we must forget any instructions
+  // that were deemed as dead-to-be-deleted during normal promotion.
+  forgetTheDead();
+
+  Changed = true;
+
+  // Now, we want to retain all of the instructions operating on the original
+  // alloca, so to avoid much hassle, create a new alloca, and swap (RAUW) them.
+  AllocaInst *ShadowAlloca = cast<AllocaInst>(AllocaToPromote->clone());
+  ShadowAlloca->takeName(AllocaToPromote);
+  AllocaToPromote->setName(ShadowAlloca->getName() + ".prom");
+  ShadowAlloca->insertBefore(AllocaToPromote);
+  AllocaToPromote->replaceAllUsesWith(ShadowAlloca);
+
+  // Avoid recomputing the same pointer over and over again, cache it.
+  SmallDenseMap<std::pair<uint64_t, Type *>, Value *> RebasedPtrsCSE;
+
+  // Don't do anything fancy, just put new insts "right after" the alloca.
+  IRBuilderTy Builder(AllocaToPromote->getContext());
+  BasicBlock *AllocaToPromoteBB = AllocaToPromote->getParent();
+  Builder.SetInsertPoint(AllocaToPromoteBB,
+                         AllocaToPromoteBB->getFirstInsertionPt());
+
+  // Give a pointer `Offset` bytes into the `AllocaToPromote` with `PtrTy` type.
+  auto getRebasedPtr = [&RebasedPtrsCSE, &Builder, AllocaToPromote,
+                        DL = AllocaToPromote->getModule()->getDataLayout()](
+                           PointerType *PtrTy, const uint64_t Offset) {
+    // Look it up in a cache first.
+    auto It = RebasedPtrsCSE.find({Offset, PtrTy});
+    if (It != RebasedPtrsCSE.end())
+      return It->second;
+
+    // Otherwise, create a new pointer, and cache it for the future.
+    Value *NewPtr = getAdjustedPtr(
+        Builder, DL, AllocaToPromote,
+        APInt(DL.getIndexSizeInBits(PtrTy->getAddressSpace()), Offset), PtrTy,
+        "");
+    RebasedPtrsCSE[{Offset, PtrTy}] = NewPtr;
+
+    return NewPtr;
+  };
+
+  // Some instructions may have several uses of an alloca, and there's
+  // a separate slice for each use, so we must cache each instruction
+  // we clone, so that we only clone it once,
+  // not for each slice that references it.
+  SmallDenseMap<Instruction *, Instruction *> InstrCloneMap;
+
+  // Now, let's just deal with each slice. Roughly, we need to clone each
+  // instruction that is referenced by a slice (once per instruction!),
+  // and change the appropriate pointer from pointing at the shadow alloca
+  // into pointing into the alloca we are going to promote.
+  //
+  // NOTE: the original instruction is generally preserved,
+  // because we need to maintain the content parity between the two allocas!
+  for (Slice &S : Slices) {
+    // Just completely ignore dead slices.
+    if (S.isDead())
+      continue;
+
+    // Which instruction does this slice represent?
+    Use *OrigUse = S.getUse();
+    auto *OrigInstr = cast<Instruction>(OrigUse->getUser());
+
+    // Now, we need to make a clone of this instruction, but operating on
+    // the alloca-to-be-promoted instead.
+    Instruction *ClonedInstr;
+    // Only clone instruction once! See if we already did that for this instr.
+    auto It = InstrCloneMap.find(OrigInstr);
+    if (It != InstrCloneMap.end())
+      ClonedInstr = It->second;
+    else {
+      // This is the first time this instruction is seen.
+      // Clone it next to the original instruction, and cache it.
+      ClonedInstr = OrigInstr->clone();
+      ClonedInstr->insertBefore(OrigInstr);
+      InstrCloneMap.insert({OrigInstr, ClonedInstr});
+
+      // Also, if the instruction was returning anything, we do that instead.
+      if (!ClonedInstr->getType()->isVoidTy()) {
+        assert(isa<LoadInst>(OrigInstr) &&
+               "Not expecting to encounter here anything other than a `load`.");
+        ClonedInstr->setName(OrigInstr->getName() + ".prom");
+        OrigInstr->replaceAllUsesWith(ClonedInstr);
+      }
+
+      if (isa<LoadInst>(OrigInstr))
+        // We know that all the offending (non-capturing) calls do not modify
+        // the content of the shadow alloca, so we do not need to propagate
+        // the content of the shadow alloca to the alloca-to-be-promoted.
+        DeadUsers.push_back(OrigInstr);
+    }
+
+    // Final touch: the slice should refer to the
+    // use of the alloca-to-be-promoted, while it currently refers to
+    // use of the shadow alloca, so rectify that.
+    Value *NewPtr = getRebasedPtr(cast<PointerType>(OrigUse->get()->getType()),
+                                  S.beginOffset());
+    Use &ClonedUse = ClonedInstr->getOperandUse(OrigUse->getOperandNo());
+    ClonedUse.set(NewPtr);
+    S.setUse(&ClonedUse);
+  }
+}
+
 /// Strip aggregate type wrapping.
 ///
 /// This removes no-op aggregate types wrapping an underlying type. It will
@@ -4612,7 +4766,7 @@ bool SROAPass::runOnAlloca(AllocaInst &AI) {
   Changed |= AggRewriter.rewrite(AI);
 
   // Build the slices using a recursive instruction-visiting builder.
-  AllocaSlices AS(DL, AI);
+  AllocaSlices AS(DL, AI, Changed);
   LLVM_DEBUG(AS.print(dbgs()));
   if (AS.isEscaped())
     return Changed;