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

3607

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

3611

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

3615

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.

3615–3616

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;

3618

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

3626

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

class SROA : public FunctionPass {
   /// speculation and by being speculated will allow promoting allocas
   /// currently in the promotable queue.
   SetVector<SelectInst *, SmallVector<SelectInst *, 2>> SpeculatableSelects;
 
   /// Debug intrinsics do not show up as regular uses in the
   /// IR. This side-table holds the missing use edges.
   DenseMap<AllocaInst *, DbgDeclareInst *> DbgDeclares;
 
+  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) {}
+  };
+  /// \brief Describes the allocas introduced by rewritePartition in order to
+  /// migrate the debug info.
+  std::vector<Piece> Pieces;
+
 public:
   SROA(bool RequiresDomTree = true)
       : FunctionPass(ID), RequiresDomTree(RequiresDomTree), C(nullptr),
         DL(nullptr), DT(nullptr) {
     initializeSROAPass(*PassRegistry::getPassRegistry());
   }
   bool runOnFunction(Function &F) override;
   void getAnalysisUsage(AnalysisUsage &AU) const override;
if (SliceTy == AI.getAllocatedType()) {
     // 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);
-    }
+    // Remember the alloca so we can update the debug info for it.
+    Pieces.push_back( Piece(NewAI, P.beginOffset(), P.size()) );
   }
 
   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

 /// \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;
+  Pieces.clear();
 
   // Rewrite each parttion.
   for (auto &P : AS.partitions()) {
     Changed |= rewritePartition(AI, AS, P);
     ++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());