Index: lib/Transforms/Scalar/SROA.cpp
===================================================================
--- lib/Transforms/Scalar/SROA.cpp
+++ lib/Transforms/Scalar/SROA.cpp
@@ -1240,8 +1240,8 @@
   friend class PHIOrSelectSpeculator;
   friend class AllocaSliceRewriter;
 
-  bool rewritePartition(AllocaInst &AI, AllocaSlices &AS,
-                        AllocaSlices::Partition &P);
+  AllocaInst *rewritePartition(AllocaInst &AI, AllocaSlices &AS,
+                               AllocaSlices::Partition &P);
   bool splitAlloca(AllocaInst &AI, AllocaSlices &AS);
   bool runOnAlloca(AllocaInst &AI);
   void clobberUse(Use &U);
@@ -3435,8 +3435,8 @@
 /// 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,
-                            AllocaSlices::Partition &P) {
+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.
@@ -3497,27 +3497,6 @@
         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 "
@@ -3600,7 +3579,7 @@
       PostPromotionWorklist.pop_back();
   }
 
-  return true;
+  return NewAI;
 }
 
 /// \brief Walks the slices of an alloca and form partitions based on them,
@@ -3612,9 +3591,25 @@
   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;
   }
 
@@ -3622,6 +3617,26 @@
   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;
 }