[region-isolation] Track operand info in a separate map rather than inline in a TransferringOperand data structure.

This is backing out an approach that I thought would be superior, but ended up
causing problems.

Originally, we mapped a region number to an immutable pointer set containing
Operand * where the region was tranferred. This worked great for a time... until
I began to need to propagate other information from the transferring code in the
analysis to the actual diagnostic emitter.

To be able to do that, my thought was to make a wrapper type around Operand
called TransferringOperand that contained the operand and the other information
I needed. This seemed to provide me what I wanted but I later found that since
the immutable pointer set was tracking TransferringOperands which were always
newly wrapped with an Operand *, we actually always created new pointer
sets. This is of course wasteful from a memory perspective, but also prevents me
from tracking transferring operand sets during the dataflow since we would never
converge.

In this commit, I fix that issue by again tracking just an Operand * in the
TransferringOperandSet and instead map each operand to a state structure which
we merge dataflow state into whenever we visit it. This provides us with
everything we need to in the next commit to including a region -> transferring
operand set equality check in our dataflow equations and always converge.

Author: gottesmm

include/swift/SILOptimizer/Analysis/RegionAnalysis.h

@@ -88,10 +88,13 @@ class BlockPartitionState {
 
   TransferringOperandSetFactory &ptrSetFactory;
 
+  TransferringOperandToStateMap &transferringOpToStateMap;
+
   BlockPartitionState(SILBasicBlock *basicBlock,
                       PartitionOpTranslator &translator,
                       TransferringOperandSetFactory &ptrSetFactory,
-                      IsolationHistory::Factory &isolationHistoryFactory);
+                      IsolationHistory::Factory &isolationHistoryFactory,
+                      TransferringOperandToStateMap &transferringOpToStateMap);
 
 public:
   bool getLiveness() const { return isLive; }
@@ -397,6 +400,8 @@ class RegionAnalysisFunctionInfo {
 
   IsolationHistory::Factory isolationHistoryFactory;
 
+  TransferringOperandToStateMap transferringOpToStateMap;
+
   // We make this optional to prevent an issue that we have seen on windows when
   // capturing a field in a closure that is used to initialize a different
   // field.
@@ -492,6 +497,16 @@ class RegionAnalysisFunctionInfo {
     return valueMap;
   }
 
+  IsolationHistory::Factory &getIsolationHistoryFactory() {
+    assert(supportedFunction && "Unsupported Function?!");
+    return isolationHistoryFactory;
+  }
+
+  TransferringOperandToStateMap &getTransferringOpToStateMap() {
+    assert(supportedFunction && "Unsupported Function?!");
+    return transferringOpToStateMap;
+  }
+
   bool isClosureCaptured(SILValue value, Operand *op);
 
   static SILValue getUnderlyingTrackedValue(SILValue value);

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -224,6 +224,7 @@ class SILIsolationInfo {
 };
 
 class Partition;
+class TransferringOperandToStateMap;
 
 /// A persistent data structure that is used to "rewind" partition history so
 /// that we can discover when values become part of the same region.
@@ -245,6 +246,7 @@ class IsolationHistory {
 
   // TODO: This shouldn't need to be a friend.
   friend class Partition;
+  friend TransferringOperandToStateMap;
 
   /// First node in the immutable linked list.
   Node *head = nullptr;
@@ -298,6 +300,11 @@ class IsolationHistory {
   /// Push that \p other should be merged into this region.
   void pushCFGHistoryJoin(Node *otherNode);
 
+  /// Push the top node of \p history as a CFG history join.
+  void pushCFGHistoryJoin(IsolationHistory history) {
+    return pushCFGHistoryJoin(history.getHead());
+  }
+
   Node *pop();
 };
 
@@ -457,10 +464,7 @@ class IsolationHistory::Factory {
   IsolationHistory get() { return IsolationHistory(this); }
 };
 
-class TransferringOperand {
-  using ValueType = llvm::PointerIntPair<Operand *, 1>;
-  ValueType value;
-
+struct TransferringOperandState {
   /// The dynamic isolation info of the region of value when we transferred.
   ///
   /// This will contain the isolated value if we found one.
@@ -469,65 +473,30 @@ class TransferringOperand {
   /// The dynamic isolation history at this point.
   IsolationHistory isolationHistory;
 
-  TransferringOperand(ValueType newValue, SILIsolationInfo isolationRegionInfo,
-                      IsolationHistory isolationHistory)
-      : value(newValue), isolationInfo(isolationRegionInfo),
-        isolationHistory(isolationHistory) {
-    assert(isolationInfo && "Should never see unknown isolation info");
-  }
-
-public:
-  TransferringOperand(Operand *op, bool isClosureCaptured,
-                      SILIsolationInfo isolationRegionInfo,
-                      IsolationHistory isolationHistory)
-      : TransferringOperand({op, isClosureCaptured}, isolationRegionInfo,
-                            isolationHistory) {}
-  explicit TransferringOperand(Operand *op,
-                               SILIsolationInfo isolationRegionInfo,
-                               IsolationHistory isolationHistory)
-      : TransferringOperand({op, false}, isolationRegionInfo,
-                            isolationHistory) {}
-
-  operator bool() const { return bool(value.getPointer()); }
-
-  Operand *getOperand() const { return value.getPointer(); }
-
-  SILValue get() const { return getOperand()->get(); }
-
-  bool isClosureCaptured() const { return value.getInt(); }
-
-  SILInstruction *getUser() const { return getOperand()->getUser(); }
+  /// Set to true if the element associated with the operand's vlaue is closure
+  /// captured by the user. In such a case, if our element is a sendable var of
+  /// a non-Sendable type, we cannot access it since we could race against an
+  /// assignment to the var in a closure.
+  bool isClosureCaptured;
 
-  SILIsolationInfo getIsolationInfo() const { return isolationInfo; }
-
-  IsolationHistory getIsolationHistory() const { return isolationHistory; }
-
-  unsigned getOperandNumber() const { return getOperand()->getOperandNumber(); }
-
-  void print(llvm::raw_ostream &os) const {
-    os << "Op Num: " << getOperand()->getOperandNumber() << ". "
-       << "Capture: " << (isClosureCaptured() ? "yes. " : "no.  ")
-       << "IsolationInfo: ";
-    isolationInfo.print(os);
-    os << "\nUser: " << *getUser();
-  }
+  TransferringOperandState(IsolationHistory history)
+      : isolationInfo(), isolationHistory(history), isClosureCaptured(false) {}
+};
 
-  static void Profile(llvm::FoldingSetNodeID &id, Operand *op,
-                      bool isClosureCaptured,
-                      SILIsolationInfo isolationRegionInfo,
-                      IsolationHistory isolationHistory) {
-    id.AddPointer(op);
-    id.AddBoolean(isClosureCaptured);
-    isolationRegionInfo.Profile(id);
-    id.AddPointer(isolationHistory.getHead());
-  }
+class TransferringOperandToStateMap {
+  llvm::SmallDenseMap<Operand *, TransferringOperandState> internalMap;
+  IsolationHistory::Factory &isolationHistoryFactory;
 
-  void Profile(llvm::FoldingSetNodeID &id) const {
-    Profile(id, getOperand(), isClosureCaptured(), isolationInfo,
-            isolationHistory);
+public:
+  TransferringOperandToStateMap(
+      IsolationHistory::Factory &isolationHistoryFactory)
+      : isolationHistoryFactory(isolationHistoryFactory) {}
+  TransferringOperandState &get(Operand *op) const {
+    auto *self = const_cast<TransferringOperandToStateMap *>(this);
+    auto history = IsolationHistory(&isolationHistoryFactory);
+    return self->internalMap.try_emplace(op, TransferringOperandState(history))
+        .first->getSecond();
   }
-
-  SWIFT_DEBUG_DUMP { print(llvm::dbgs()); }
 };
 
 } // namespace swift
@@ -676,9 +645,8 @@ class Partition {
 
   using Element = PartitionPrimitives::Element;
   using Region = PartitionPrimitives::Region;
-  using TransferringOperandSet = ImmutablePointerSet<TransferringOperand *>;
-  using TransferringOperandSetFactory =
-      ImmutablePointerSetFactory<TransferringOperand *>;
+  using TransferringOperandSet = ImmutablePointerSet<Operand *>;
+  using TransferringOperandSetFactory = ImmutablePointerSetFactory<Operand *>;
   using IsolationHistoryNode = IsolationHistory::Node;
 
 private:
@@ -1014,13 +982,16 @@ struct PartitionOpEvaluator {
 
 protected:
   TransferringOperandSetFactory &ptrSetFactory;
+  TransferringOperandToStateMap &operandToStateMap;
 
   Partition &p;
 
 public:
   PartitionOpEvaluator(Partition &p,
-                       TransferringOperandSetFactory &ptrSetFactory)
-      : ptrSetFactory(ptrSetFactory), p(p) {}
+                       TransferringOperandSetFactory &ptrSetFactory,
+                       TransferringOperandToStateMap &operandToStateMap)
+      : ptrSetFactory(ptrSetFactory), operandToStateMap(operandToStateMap),
+        p(p) {}
 
   /// Call shouldEmitVerboseLogging on our CRTP subclass.
   bool shouldEmitVerboseLogging() const {
@@ -1029,7 +1000,7 @@ struct PartitionOpEvaluator {
 
   /// Call handleLocalUseAfterTransfer on our CRTP subclass.
   void handleLocalUseAfterTransfer(const PartitionOp &op, Element elt,
-                                   TransferringOperand *transferringOp) const {
+                                   Operand *transferringOp) const {
     return asImpl().handleLocalUseAfterTransfer(op, elt, transferringOp);
   }
 
@@ -1194,9 +1165,13 @@ struct PartitionOpEvaluator {
       }
 
       // Mark op.getOpArgs()[0] as transferred.
-      auto *ptrSet = ptrSetFactory.emplace(
-          op.getSourceOp(), isClosureCapturedElt, transferredRegionIsolation,
-          p.getIsolationHistory());
+      TransferringOperandState &state = operandToStateMap.get(op.getSourceOp());
+      state.isClosureCaptured |= isClosureCapturedElt;
+      state.isolationInfo =
+          state.isolationInfo.merge(transferredRegionIsolation);
+      assert(state.isolationInfo && "Cannot have unknown");
+      state.isolationHistory.pushCFGHistoryJoin(p.getIsolationHistory());
+      auto *ptrSet = ptrSetFactory.get(op.getSourceOp());
       p.markTransferred(op.getOpArgs()[0], ptrSet);
       return;
     }
@@ -1266,9 +1241,8 @@ struct PartitionOpEvaluator {
 private:
   // Private helper that squelches the error if our transfer instruction and our
   // use have the same isolation.
-  void
-  handleLocalUseAfterTransferHelper(const PartitionOp &op, Element elt,
-                                    TransferringOperand *transferringOp) const {
+  void handleLocalUseAfterTransferHelper(const PartitionOp &op, Element elt,
+                                         Operand *transferringOp) const {
     if (shouldTryToSquelchErrors()) {
       if (auto isolationInfo = SILIsolationInfo::get(op.getSourceInst())) {
         if (isolationInfo.isActorIsolated() &&
@@ -1306,8 +1280,9 @@ struct PartitionOpEvaluatorBaseImpl : PartitionOpEvaluator<Subclass> {
   using Super = PartitionOpEvaluator<Subclass>;
 
   PartitionOpEvaluatorBaseImpl(Partition &workingPartition,
-                               TransferringOperandSetFactory &ptrSetFactory)
-      : Super(workingPartition, ptrSetFactory) {}
+                               TransferringOperandSetFactory &ptrSetFactory,
+                               TransferringOperandToStateMap &operandToStateMap)
+      : Super(workingPartition, ptrSetFactory, operandToStateMap) {}
 
   /// Should we emit extra verbose logging statements when evaluating
   /// PartitionOps.
@@ -1326,7 +1301,7 @@ struct PartitionOpEvaluatorBaseImpl : PartitionOpEvaluator<Subclass> {
   /// region. Can be used to get the immediate value transferred or the
   /// transferring instruction.
   void handleLocalUseAfterTransfer(const PartitionOp &op, Element elt,
-                                   TransferringOperand *transferringOp) const {}
+                                   Operand *transferringOp) const {}
 
   /// This is called if we detect a never transferred element that was passed to
   /// a transfer instruction.
@@ -1374,8 +1349,10 @@ struct PartitionOpEvaluatorBaseImpl : PartitionOpEvaluator<Subclass> {
 struct PartitionOpEvaluatorBasic final
     : PartitionOpEvaluatorBaseImpl<PartitionOpEvaluatorBasic> {
   PartitionOpEvaluatorBasic(Partition &workingPartition,
-                            TransferringOperandSetFactory &ptrSetFactory)
-      : PartitionOpEvaluatorBaseImpl(workingPartition, ptrSetFactory) {}
+                            TransferringOperandSetFactory &ptrSetFactory,
+                            TransferringOperandToStateMap &operandToStateMap)
+      : PartitionOpEvaluatorBaseImpl(workingPartition, ptrSetFactory,
+                                     operandToStateMap) {}
 };
 
 } // namespace swift

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -2954,10 +2954,12 @@ TranslationSemantics PartitionOpTranslator::visitCheckedCastAddrBranchInst(
 BlockPartitionState::BlockPartitionState(
     SILBasicBlock *basicBlock, PartitionOpTranslator &translator,
     TransferringOperandSetFactory &ptrSetFactory,
-    IsolationHistory::Factory &isolationHistoryFactory)
+    IsolationHistory::Factory &isolationHistoryFactory,
+    TransferringOperandToStateMap &transferringOpToStateMap)
     : entryPartition(isolationHistoryFactory.get()),
       exitPartition(isolationHistoryFactory.get()), basicBlock(basicBlock),
-      ptrSetFactory(ptrSetFactory) {
+      ptrSetFactory(ptrSetFactory),
+      transferringOpToStateMap(transferringOpToStateMap) {
   translator.translateSILBasicBlock(basicBlock, blockPartitionOps);
 }
 
@@ -2971,8 +2973,10 @@ bool BlockPartitionState::recomputeExitFromEntry(
 
     ComputeEvaluator(Partition &workingPartition,
                      TransferringOperandSetFactory &ptrSetFactory,
-                     PartitionOpTranslator &translator)
-        : PartitionOpEvaluatorBaseImpl(workingPartition, ptrSetFactory),
+                     PartitionOpTranslator &translator,
+                     TransferringOperandToStateMap &transferringOpToStateMap)
+        : PartitionOpEvaluatorBaseImpl(workingPartition, ptrSetFactory,
+                                       transferringOpToStateMap),
           translator(translator) {}
 
     SILIsolationInfo getIsolationRegionInfo(Element elt) const {
@@ -2989,7 +2993,8 @@ bool BlockPartitionState::recomputeExitFromEntry(
       return translator.isClosureCaptured(value, op->getUser());
     }
   };
-  ComputeEvaluator eval(workingPartition, ptrSetFactory, translator);
+  ComputeEvaluator eval(workingPartition, ptrSetFactory, translator,
+                        transferringOpToStateMap);
   for (const auto &partitionOp : blockPartitionOps) {
     // By calling apply without providing any error handling callbacks, errors
     // will be surpressed.  will be suppressed
@@ -3074,8 +3079,9 @@ static bool canComputeRegionsForFunction(SILFunction *fn) {
 RegionAnalysisFunctionInfo::RegionAnalysisFunctionInfo(
     SILFunction *fn, PostOrderFunctionInfo *pofi)
     : allocator(), fn(fn), valueMap(fn), translator(), ptrSetFactory(allocator),
-      isolationHistoryFactory(allocator), blockStates(), pofi(pofi),
-      solved(false), supportedFunction(true) {
+      isolationHistoryFactory(allocator),
+      transferringOpToStateMap(isolationHistoryFactory), blockStates(),
+      pofi(pofi), solved(false), supportedFunction(true) {
   // Before we do anything, make sure that we support processing this function.
   //
   // NOTE: See documentation on supportedFunction for criteria.
@@ -3088,7 +3094,8 @@ RegionAnalysisFunctionInfo::RegionAnalysisFunctionInfo(
       PartitionOpTranslator(fn, pofi, valueMap, isolationHistoryFactory);
   blockStates.emplace(fn, [this](SILBasicBlock *block) -> BlockPartitionState {
     return BlockPartitionState(block, *translator, ptrSetFactory,
-                               isolationHistoryFactory);
+                               isolationHistoryFactory,
+                               transferringOpToStateMap);
   });
   // Mark all blocks as needing to be updated.
   for (auto &block : *fn) {