[region-isolation] Add a SILLocation to SequenceBoundary in IsolationHistory.

We package all isolation history nodes from a single instruction by placing a
sequence boundary at the bottom. When ever we pop, we actually pop a PartitionOp
at a time meaning that we pop until we see a SequenceBoundary. Thus the sequence
boundary will always be the last element visited when popping meaning that it is
a convenient place to stick the SILLocation associated with the entire
PartitionOp. As a benefit, there was some unused space in IsolationHistory::Node
for that case since we were not using the std::variant field at all.

Author: gottesmm

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -249,7 +249,10 @@ class IsolationHistory {
   /// Push a node that signals the end of a new sequence of history nodes that
   /// should execute together. Must be explicitly ended by a push sequence
   /// end. Is non-rentrant, so one cannot have multiple sequence starts.
-  Node *pushHistorySequenceBoundary();
+  ///
+  /// \p loc the SILLocation that identifies the instruction that the "package"
+  /// of history nodes that this sequence boundary ends is associated with.
+  Node *pushHistorySequenceBoundary(SILLocation loc);
 
   /// Push onto the history list that \p value should be added into its own
   /// independent region.
@@ -325,9 +328,12 @@ class IsolationHistory::Node final
   /// Child node. Never set on construction.
   Node *child = nullptr;
 
-  // Either the element that we are adding/removing/modifying or the Node that
-  // we are merging into.
-  std::variant<Element, Node *> subject;
+  /// Contains:
+  ///
+  /// 1. Node * if we have a CFGHistoryJoin.
+  /// 2. A SILLocation if we have a SequenceBoundary.
+  /// 3. An element otherwise.
+  std::variant<Element, Node *, SILLocation> subject;
 
   /// Number of additional element arguments stored in the tail allocated array.
   unsigned numAdditionalElements;
@@ -339,6 +345,8 @@ class IsolationHistory::Node final
 
   Node(Kind kind, Node *parent)
       : kind(kind), parent(parent), subject(nullptr) {}
+  Node(Kind kind, Node *parent, SILLocation loc)
+      : kind(kind), parent(parent), subject(loc) {}
   Node(Kind kind, Node *parent, Element value)
       : kind(kind), parent(parent), subject(value), numAdditionalElements(0) {}
   Node(Kind kind, Node *parent, Element primaryElement,
@@ -405,6 +413,12 @@ class IsolationHistory::Node final
       return nullptr;
     return getFirstArgAsNode();
   }
+
+  std::optional<SILLocation> getHistoryBoundaryLoc() const {
+    if (kind != SequenceBoundary)
+      return {};
+    return std::get<SILLocation>(subject);
+  }
 };
 
 class IsolationHistory::Factory {
@@ -671,12 +685,12 @@ class Partition {
 
   /// Return a new Partition that has a single region containing the elements of
   /// \p indices.
-  static Partition singleRegion(ArrayRef<Element> indices,
+  static Partition singleRegion(SILLocation loc, ArrayRef<Element> indices,
                                 IsolationHistory inputHistory);
 
   /// Return a new Partition that has each element of \p indices in their own
   /// region.
-  static Partition separateRegions(ArrayRef<Element> indices,
+  static Partition separateRegions(SILLocation loc, ArrayRef<Element> indices,
                                    IsolationHistory inputHistory);
 
   /// Test two partititons for equality by first putting them in canonical form
@@ -829,8 +843,8 @@ class Partition {
   void printHistory(llvm::raw_ostream &os) const;
 
   /// See docs on \p history.pushHistorySequenceBoundary().
-  IsolationHistoryNode *pushHistorySequenceBoundary() {
-    return history.pushHistorySequenceBoundary();
+  IsolationHistoryNode *pushHistorySequenceBoundary(SILLocation loc) {
+    return history.pushHistorySequenceBoundary(loc);
   }
 
   bool isTransferred(Element val) const {
@@ -1043,6 +1057,11 @@ struct PartitionOpEvaluator {
     return asImpl().isClosureCaptured(elt, op);
   }
 
+  /// Some evaluators pass in mock instructions that one cannot call getLoc()
+  /// upon. So to allow for this, provide a routine that our impl can override
+  /// if they need to.
+  static SILLocation getLoc(SILInstruction *inst) { return Impl::getLoc(inst); }
+
   /// Apply \p op to the partition op.
   void apply(const PartitionOp &op) const {
     if (shouldEmitVerboseLogging()) {
@@ -1061,7 +1080,7 @@ struct PartitionOpEvaluator {
 
     // Set the boundary so that as we push, this shows when to stop processing
     // for this PartitionOp.
-    p.pushHistorySequenceBoundary();
+    p.pushHistorySequenceBoundary(getLoc(op.getSourceInst()));
 
     switch (op.getKind()) {
     case PartitionOpKind::Assign:
@@ -1299,6 +1318,8 @@ struct PartitionOpEvaluatorBaseImpl : PartitionOpEvaluator<Subclass> {
 
   /// By default squelch errors.
   bool shouldTryToSquelchErrors() const { return true; }
+
+  static SILLocation getLoc(SILInstruction *inst) { return inst->getLoc(); }
 };
 
 /// A subclass of PartitionOpEvaluatorBaseImpl that doesn't have any special

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -1460,7 +1460,8 @@ class PartitionOpTranslator {
     auto functionArguments = function->getArguments();
     if (functionArguments.empty()) {
       LLVM_DEBUG(llvm::dbgs() << "    None.\n");
-      functionArgPartition = Partition::singleRegion({}, historyFactory.get());
+      functionArgPartition = Partition::singleRegion(SILLocation::invalid(), {},
+                                                     historyFactory.get());
       return;
     }
 
@@ -1489,8 +1490,8 @@ class PartitionOpTranslator {
       }
     }
 
-    functionArgPartition =
-        Partition::singleRegion(nonSendableJoinedIndices, historyFactory.get());
+    functionArgPartition = Partition::singleRegion(
+        SILLocation::invalid(), nonSendableJoinedIndices, historyFactory.get());
     for (Element elt : nonSendableSeparateIndices) {
       functionArgPartition->trackNewElement(elt);
     }

lib/SILOptimizer/Utils/PartitionUtils.cpp

@@ -237,7 +237,7 @@ void PartitionOp::print(llvm::raw_ostream &os, bool extraSpace) const {
 //                              MARK: Partition
 //===----------------------------------------------------------------------===//
 
-Partition Partition::singleRegion(ArrayRef<Element> indices,
+Partition Partition::singleRegion(SILLocation loc, ArrayRef<Element> indices,
                                   IsolationHistory inputHistory) {
   Partition p(inputHistory);
   if (!indices.empty()) {
@@ -249,7 +249,7 @@ Partition Partition::singleRegion(ArrayRef<Element> indices,
 
     // Place all of the operations until end of scope into one history
     // sequence.
-    p.pushHistorySequenceBoundary();
+    p.pushHistorySequenceBoundary(loc);
 
     // First create a region for repElement. We are going to merge all other
     // regions into its region.
@@ -269,14 +269,14 @@ Partition Partition::singleRegion(ArrayRef<Element> indices,
   return p;
 }
 
-Partition Partition::separateRegions(ArrayRef<Element> indices,
+Partition Partition::separateRegions(SILLocation loc, ArrayRef<Element> indices,
                                      IsolationHistory inputHistory) {
   Partition p(inputHistory);
   if (indices.empty())
     return p;
 
   // Place all operations in one history sequence.
-  p.pushHistorySequenceBoundary();
+  p.pushHistorySequenceBoundary(loc);
 
   auto maxIndex = Element(0);
   for (Element index : indices) {
@@ -954,10 +954,11 @@ IsolationHistory::pushNewElementRegion(Element element) {
   return getHead();
 }
 
-IsolationHistory::Node *IsolationHistory::pushHistorySequenceBoundary() {
+IsolationHistory::Node *
+IsolationHistory::pushHistorySequenceBoundary(SILLocation loc) {
   unsigned size = Node::totalSizeToAlloc<Element>(0);
   void *mem = factory->allocator.Allocate(size, alignof(Node));
-  head = new (mem) Node(Node::SequenceBoundary, head);
+  head = new (mem) Node(Node::SequenceBoundary, head, loc);
   return getHead();
 }
 

unittests/SILOptimizer/PartitionUtilsTest.cpp

@@ -49,6 +49,44 @@ struct MockedPartitionOpEvaluator final
   }
 
   bool shouldTryToSquelchErrors() const { return false; }
+
+  static SILLocation getLoc(SILInstruction *inst) {
+    return SILLocation::invalid();
+  }
+};
+
+} // namespace
+
+namespace {
+
+struct MockedPartitionOpEvaluatorWithFailureCallback final
+    : PartitionOpEvaluatorBaseImpl<
+          MockedPartitionOpEvaluatorWithFailureCallback> {
+  using FailureCallbackTy =
+      std::function<void(const PartitionOp &, unsigned, TransferringOperand *)>;
+  FailureCallbackTy failureCallback;
+
+  MockedPartitionOpEvaluatorWithFailureCallback(
+      Partition &workingPartition, TransferringOperandSetFactory &ptrSetFactory,
+      FailureCallbackTy failureCallback)
+      : PartitionOpEvaluatorBaseImpl(workingPartition, ptrSetFactory),
+        failureCallback(failureCallback) {}
+
+  void handleLocalUseAfterTransfer(const PartitionOp &op, Element elt,
+                                   TransferringOperand *transferringOp) const {
+    failureCallback(op, elt, transferringOp);
+  }
+
+  // Just say that we always have a disconnected value.
+  SILIsolationInfo getIsolationRegionInfo(Element elt) const {
+    return SILIsolationInfo::getDisconnected();
+  }
+
+  bool shouldTryToSquelchErrors() const { return false; }
+
+  static SILLocation getLoc(SILInstruction *inst) {
+    return SILLocation::invalid();
+  }
 };
 
 } // namespace
@@ -71,6 +109,8 @@ SILInstruction *instSingletons[5] = {
     (SILInstruction *)0xBBDA0000,
 };
 
+SILLocation fakeLoc = SILLocation::invalid();
+
 // This test tests that if a series of merges is split between two partitions
 // p1 and p2, but also applied in its entirety to p3, then joining p1 and p2
 // yields p3.
@@ -198,8 +238,8 @@ TEST(PartitionUtilsTest, Join1) {
 
   Element data1[] = {Element(0), Element(1), Element(2),
                      Element(3), Element(4), Element(5)};
-  Partition p1 =
-      Partition::separateRegions(llvm::ArrayRef(data1), historyFactory.get());
+  Partition p1 = Partition::separateRegions(fakeLoc, llvm::ArrayRef(data1),
+                                            historyFactory.get());
 
   {
     MockedPartitionOpEvaluator eval(p1, factory);
@@ -211,8 +251,8 @@ TEST(PartitionUtilsTest, Join1) {
                 PartitionOp::Assign(Element(5), Element(2))});
   }
 
-  Partition p2 =
-      Partition::separateRegions(llvm::ArrayRef(data1), historyFactory.get());
+  Partition p2 = Partition::separateRegions(fakeLoc, llvm::ArrayRef(data1),
+                                            historyFactory.get());
   {
     MockedPartitionOpEvaluator eval(p2, factory);
     eval.apply({PartitionOp::Assign(Element(0), Element(0)),
@@ -240,8 +280,8 @@ TEST(PartitionUtilsTest, Join2) {
 
   Element data1[] = {Element(0), Element(1), Element(2),
                      Element(3), Element(4), Element(5)};
-  Partition p1 =
-      Partition::separateRegions(llvm::ArrayRef(data1), historyFactory.get());
+  Partition p1 = Partition::separateRegions(fakeLoc, llvm::ArrayRef(data1),
+                                            historyFactory.get());
 
   {
     MockedPartitionOpEvaluator eval(p1, factory);
@@ -255,8 +295,8 @@ TEST(PartitionUtilsTest, Join2) {
 
   Element data2[] = {Element(4), Element(5), Element(6),
                      Element(7), Element(8), Element(9)};
-  Partition p2 =
-      Partition::separateRegions(llvm::ArrayRef(data2), historyFactory.get());
+  Partition p2 = Partition::separateRegions(fakeLoc, llvm::ArrayRef(data2),
+                                            historyFactory.get());
   {
     MockedPartitionOpEvaluator eval(p2, factory);
     eval.apply({PartitionOp::Assign(Element(4), Element(4)),
@@ -288,8 +328,8 @@ TEST(PartitionUtilsTest, Join2Reversed) {
 
   Element data1[] = {Element(0), Element(1), Element(2),
                      Element(3), Element(4), Element(5)};
-  Partition p1 =
-      Partition::separateRegions(llvm::ArrayRef(data1), historyFactory.get());
+  Partition p1 = Partition::separateRegions(fakeLoc, llvm::ArrayRef(data1),
+                                            historyFactory.get());
 
   {
     MockedPartitionOpEvaluator eval(p1, factory);
@@ -303,8 +343,8 @@ TEST(PartitionUtilsTest, Join2Reversed) {
 
   Element data2[] = {Element(4), Element(5), Element(6),
                      Element(7), Element(8), Element(9)};
-  Partition p2 =
-      Partition::separateRegions(llvm::ArrayRef(data2), historyFactory.get());
+  Partition p2 = Partition::separateRegions(fakeLoc, llvm::ArrayRef(data2),
+                                            historyFactory.get());
   {
     MockedPartitionOpEvaluator eval(p2, factory);
     eval.apply({PartitionOp::Assign(Element(4), Element(4)),
@@ -341,8 +381,8 @@ TEST(PartitionUtilsTest, JoinLarge) {
       Element(15), Element(16), Element(17), Element(18), Element(19),
       Element(20), Element(21), Element(22), Element(23), Element(24),
       Element(25), Element(26), Element(27), Element(28), Element(29)};
-  Partition p1 =
-      Partition::separateRegions(llvm::ArrayRef(data1), historyFactory.get());
+  Partition p1 = Partition::separateRegions(fakeLoc, llvm::ArrayRef(data1),
+                                            historyFactory.get());
   {
     MockedPartitionOpEvaluator eval(p1, factory);
     eval.apply({PartitionOp::Assign(Element(0), Element(29)),
@@ -384,8 +424,8 @@ TEST(PartitionUtilsTest, JoinLarge) {
       Element(30), Element(31), Element(32), Element(33), Element(34),
       Element(35), Element(36), Element(37), Element(38), Element(39),
       Element(40), Element(41), Element(42), Element(43), Element(44)};
-  Partition p2 =
-      Partition::separateRegions(llvm::ArrayRef(data2), historyFactory.get());
+  Partition p2 = Partition::separateRegions(fakeLoc, llvm::ArrayRef(data2),
+                                            historyFactory.get());
   {
     MockedPartitionOpEvaluator eval(p2, factory);
     eval.apply({PartitionOp::Assign(Element(15), Element(31)),
@@ -558,36 +598,6 @@ TEST(PartitionUtilsTest, TestAssign) {
   EXPECT_TRUE(Partition::equals(p1, p3));
 }
 
-namespace {
-
-struct MockedPartitionOpEvaluatorWithFailureCallback final
-    : PartitionOpEvaluatorBaseImpl<
-          MockedPartitionOpEvaluatorWithFailureCallback> {
-  using FailureCallbackTy =
-      std::function<void(const PartitionOp &, unsigned, TransferringOperand *)>;
-  FailureCallbackTy failureCallback;
-
-  MockedPartitionOpEvaluatorWithFailureCallback(
-      Partition &workingPartition, TransferringOperandSetFactory &ptrSetFactory,
-      FailureCallbackTy failureCallback)
-      : PartitionOpEvaluatorBaseImpl(workingPartition, ptrSetFactory),
-        failureCallback(failureCallback) {}
-
-  void handleLocalUseAfterTransfer(const PartitionOp &op, Element elt,
-                                   TransferringOperand *transferringOp) const {
-    failureCallback(op, elt, transferringOp);
-  }
-
-  // Just say that we always have a disconnected value.
-  SILIsolationInfo getIsolationRegionInfo(Element elt) const {
-    return SILIsolationInfo::getDisconnected();
-  }
-
-  bool shouldTryToSquelchErrors() const { return false; }
-};
-
-} // namespace
-
 // This test tests that consumption consumes entire regions as expected
 TEST(PartitionUtilsTest, TestConsumeAndRequire) {
   llvm::BumpPtrAllocator allocator;