SROA: Defer migrating the debug info for new allocas until after all partitions are created.

[aprantl]

Summary

SROA: Defer migrating the debug info for new allocas until after all partitions are created.
As discussed on #llvm today.

[chandlerc]

lib/Transforms/Scalar/SROA.cpp

3593

Rather than a member, how about a local here...

3607–3612

And change rewritePartition so that if it makes changes it returns the changed alloca. We can compare against &AI to handle the case that we got back the same alloca that went in, and we'll never get back the same alloca when we end up with more than one.

[chandlerc]

Is it possible to craft a test case for this? I'm thinking an alloca which needs rewriting but not splitting...

lib/Transforms/Scalar/SROA.cpp

3600

Just call the parameters the same name as the members. This does the right thing in the initializer list. I would probably use AI, Offset, and Size.

3603

I think you at least want a smallvector with 1 piece in it.

3611–3612

Any reason not to actually handle the null case for now? I think the code would be pretty clean:

if (AllocaInst *NewAI = rewritePartition(AI, AS, P)) {
  Changed = true;
  if (NewAI != &AI)
    Pieces.push_back(...);
}
++NumPartitions;

3627

I would say:

bool IsSplit = ...;

[aprantl]

• Addressed all of Chandler's comments.
• Added a testcase were the alloca is rewritten but not repartitioned.

[aprantl]

Tiny update: Don't add the dbg.declare to DeadInsts. It's redundant because it will get removed automatically when the alloca is removed. It's also a bad idea to remove the same instruction more than once.

[aprantl]

When compiling really large programs the pointer values in the DbgDeclares map may actually get recycled.

• clear the pointer map when entering a function
• remove values after the dbg.declare has been deleted

tested via an assertion.

[chandlerc]

Just want to say this LGTM in whatever form you decide to commit it.

I know there was a revert in between, so feel free to just commit the new version, or re-apply the old one and then commit this on top. Whatever works.

[aprantl]

This was r226598.

lib/Transforms/Scalar/SROA.cpp

public:
 
   const char *getPassName() const override { return "SROA"; }
   static char ID;
 
 private:
   friend class PHIOrSelectSpeculator;
   friend class AllocaSliceRewriter;
 
-  bool rewritePartition(AllocaInst &AI, AllocaSlices &AS,
+  AllocaInst *rewritePartition(AllocaInst &AI, AllocaSlices &AS,
                                AllocaSlices::Partition &P);
   bool splitAlloca(AllocaInst &AI, AllocaSlices &AS);
   bool runOnAlloca(AllocaInst &AI);
   void clobberUse(Use &U);
   void deleteDeadInstructions(SmallPtrSetImpl<AllocaInst *> &DeletedAllocas);
   bool promoteAllocas(Function &F);
 };
 }
 

 /// This routine drives both of the rewriting goals of the SROA pass. It tries
 /// to rewrite uses of an alloca partition to be conducive for SSA value
 /// promotion. If the partition needs a new, more refined alloca, this will
 /// build that new alloca, preserving as much type information as possible, and
 /// rewrite the uses of the old alloca to point at the new one and have the
 /// appropriate new offsets. It also evaluates how successful the rewrite was
 /// at enabling promotion and if it was successful queues the alloca to be
 /// promoted.
-bool SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS,
+AllocaInst *SROA::rewritePartition(AllocaInst &AI, AllocaSlices &AS,
                                    AllocaSlices::Partition &P) {
   // Try to compute a friendly type for this partition of the alloca. This
   // won't always succeed, in which case we fall back to a legal integer type
   // or an i8 array of an appropriate size.
   Type *SliceTy = nullptr;
   if (Type *CommonUseTy = findCommonType(P.begin(), P.end(), P.endOffset()))
     if (DL->getTypeAllocSize(CommonUseTy) >= P.size())
       SliceTy = CommonUseTy;
   if (!SliceTy)
if (SliceTy == AI.getAllocatedType()) {
     // If we will get at least this much alignment from the type alone, leave
     // the alloca's alignment unconstrained.
     if (Alignment <= DL->getABITypeAlignment(SliceTy))
       Alignment = 0;
     NewAI = new AllocaInst(
         SliceTy, nullptr, Alignment,
         AI.getName() + ".sroa." + Twine(P.begin() - AS.begin()), &AI);
     ++NumNewAllocas;
-
-    // Migrate debug information from the old alloca to the new alloca
-    // and the individial slices.
-    if (DbgDeclareInst *DbgDecl = DbgDeclares.lookup(&AI)) {
-      DIVariable Var(DbgDecl->getVariable());
-      DIExpression Piece;
-      DIBuilder DIB(*AI.getParent()->getParent()->getParent(),
-                    /*AllowUnresolved*/ false);
-      // Create a piece expression describing the slice, if the new slize is
-      // smaller than the old alloca or the old alloca already was described
-      // with a piece. It would be even better to just compare against the size
-      // of the type described in the debug info, but then we would need to
-      // build an expensive DIRefMap.
-      if (P.size() < DL->getTypeAllocSize(AI.getAllocatedType()) ||
-          DIExpression(DbgDecl->getExpression()).isVariablePiece())
-        Piece = DIB.createPieceExpression(P.beginOffset(), P.size());
-      Instruction *NewDDI = DIB.insertDeclare(NewAI, Var, Piece, &AI);
-      NewDDI->setDebugLoc(DbgDecl->getDebugLoc());
-      DbgDeclares.insert(std::make_pair(NewAI, cast<DbgDeclareInst>(NewDDI)));
-      DeadInsts.insert(DbgDecl);
-    }
   }
 
   DEBUG(dbgs() << "Rewriting alloca partition "
                << "[" << P.beginOffset() << "," << P.endOffset()
                << ") to: " << *NewAI << "\n");
 
   // Track the high watermark on the worklist as it is only relevant for
   // promoted allocas. We will reset it to this point if the alloca is not in
if (Promotable) {
     if (NewAI != &AI)
       Worklist.insert(NewAI);
 
     // Drop any post-promotion work items if promotion didn't happen.
     while (PostPromotionWorklist.size() > PPWOldSize)
       PostPromotionWorklist.pop_back();
   }
 
-  return true;
+  return NewAI;
 }
 
 /// \brief Walks the slices of an alloca and form partitions based on them,
 /// rewriting each of their uses.
 bool SROA::splitAlloca(AllocaInst &AI, AllocaSlices &AS) {
   if (AS.begin() == AS.end())
     return false;
 
   unsigned NumPartitions = 0;
   bool Changed = false;
 
+  /// \brief Describes the allocas introduced by rewritePartition
+  /// in order to migrate the debug info.
+  struct Piece {
+    AllocaInst *Alloca;
+    uint64_t Offset;
+    uint64_t Size;
+    Piece(AllocaInst *AI, uint64_t O, uint64_t S)
+      : Alloca(AI), Offset(O), Size(S) {}
+  };
+  std::vector<Piece> Pieces;
+
   // Rewrite each parttion.
   for (auto &P : AS.partitions()) {
-    Changed |= rewritePartition(AI, AS, P);
+    auto NewAI = rewritePartition(AI, AS, P);
+    if (NewAI != &AI)
+      Pieces.push_back(Piece(NewAI, P.beginOffset(), P.size()));
+
+    // FIXME: See related comment in rewritePartition.
+    Changed = true;
     ++NumPartitions;
   }
 
   NumAllocaPartitions += NumPartitions;
   MaxPartitionsPerAlloca =
       std::max<unsigned>(NumPartitions, MaxPartitionsPerAlloca);
 
+  // Migrate debug information from the old alloca to the new alloca(s)
+  // and the individial partitions.
+  if (DbgDeclareInst *DbgDecl = DbgDeclares.lookup(&AI)) {
+    DIVariable Var(DbgDecl->getVariable());
+    DIExpression Expr(DbgDecl->getExpression());
+    DIBuilder DIB(*AI.getParent()->getParent()->getParent(),
+                  /*AllowUnresolved*/ false);
+    bool didSplit = Pieces.size() > 1;
+    for (auto Piece : Pieces) {
+      // Create a piece expression describing the new partition or reuse AI's
+      // expression if there is only one partition.
+      if (didSplit)
+        Expr = DIB.createPieceExpression(Piece.Offset, Piece.Size);
+      Instruction *NewDDI = DIB.insertDeclare(Piece.Alloca, Var, Expr, &AI);
+      NewDDI->setDebugLoc(DbgDecl->getDebugLoc());
+      DbgDeclares.insert(std::make_pair(Piece.Alloca,
+                                        cast<DbgDeclareInst>(NewDDI)));
+      DeadInsts.insert(DbgDecl);
+    }
+  }
   return Changed;
 }
 
 /// \brief Clobber a use with undef, deleting the used value if it becomes dead.
 void SROA::clobberUse(Use &U) {
   Value *OldV = U;
   // Replace the use with an undef value.
   U = UndefValue::get(OldV->getType());