[region-isolation] Cache getUnderlyingTrackedValue.

TLDR: Was looking at some performance traces and saw that we need to cache the
result of this value.

----

Specifically, I noticed that we were spending a lot of time computing this
operation. When I looked at the code I saw that we already had a cache along the
relevant code paths... but the cache was from equivalence class representative
-> state. Before we hit that cache, we were performing the work to map the value
to the equivalence class representative... so the work to perform the relevant
lookup from value -> state (which goes through the equivalence class
representative) was not just a hash table lookup. This operation makes it
cheaper by making it two cache lookups.

It may be possible to make this cheaper by redoing the actual mapping of
information so that we can go straight from value to state. I think it would be
slightly different since we would probably need to represent the state in a
separate array and map with indices... which is really just a more efficient
hash table. We could also use malloc/etc but lets not even talk about that.

rdar://139520959

Author: gottesmm

include/swift/SILOptimizer/Analysis/RegionAnalysis.h

@@ -295,13 +295,52 @@ class RegionAnalysisValueMap {
   /// into this map
   llvm::DenseMap<RepresentativeValue, TrackableValueState>
       equivalenceClassValuesToState;
+
+  /// The inverse map of equivalenceClassValuesToState.
   llvm::DenseMap<unsigned, RepresentativeValue> stateIndexToEquivalenceClass;
 
+  /// State that the value -> representative computation yields to us.
+  struct UnderlyingTrackedValueInfo {
+    SILValue value;
+
+    /// Only used for addresses.
+    std::optional<ActorIsolation> actorIsolation;
+
+    explicit UnderlyingTrackedValueInfo(SILValue value) : value(value) {}
+
+    UnderlyingTrackedValueInfo() : value(), actorIsolation() {}
+
+    UnderlyingTrackedValueInfo(const UnderlyingTrackedValueInfo &newVal)
+        : value(newVal.value), actorIsolation(newVal.actorIsolation) {}
+
+    UnderlyingTrackedValueInfo &
+    operator=(const UnderlyingTrackedValueInfo &newVal) {
+      value = newVal.value;
+      actorIsolation = newVal.actorIsolation;
+      return *this;
+    }
+
+    UnderlyingTrackedValueInfo(SILValue value,
+                               std::optional<ActorIsolation> actorIsolation)
+        : value(value), actorIsolation(actorIsolation) {}
+
+    operator bool() const { return value; }
+  };
+
+  /// A map from a SILValue to its equivalence class representative.
+  llvm::DenseMap<SILValue, UnderlyingTrackedValueInfo> valueToEquivalenceClass;
+
   SILFunction *fn;
 
 public:
   RegionAnalysisValueMap(SILFunction *fn) : fn(fn) { }
 
+  /// Maps a value to its representative value if one exists. Return an empty
+  /// representative value if we do not find one.
+  SILValue getRepresentative(SILValue value) const {
+    return getUnderlyingTrackedValue(value).value;
+  }
+
   /// Returns the value for this instruction if it isn't a fake "represenative
   /// value" to inject actor isolatedness. Asserts in such a case.
   SILValue getRepresentative(Element trackableValueID) const;
@@ -354,6 +393,30 @@ class RegionAnalysisValueMap {
   /// {TrackableValue(), false}.
   std::pair<TrackableValue, bool>
   initializeTrackableValue(SILValue value, SILIsolationInfo info) const;
+
+  /// A helper function that performs the actual getUnderlyingTrackedValue
+  /// computation that is cached in getUnderlyingTrackedValue(). Please never
+  /// call this directly! Only call it from getUnderlyingTrackedValue.
+  UnderlyingTrackedValueInfo
+  getUnderlyingTrackedValueHelper(SILValue value) const;
+
+  UnderlyingTrackedValueInfo getUnderlyingTrackedValue(SILValue value) const {
+    // Use try_emplace so we only construct underlying tracked value info on
+    // success and only lookup once in the hash table.
+    auto *self = const_cast<RegionAnalysisValueMap *>(this);
+    auto iter = self->valueToEquivalenceClass.try_emplace(
+        value, UnderlyingTrackedValueInfo());
+
+    // Didn't insert... we have a value!
+    if (!iter.second)
+      return iter.first->getSecond();
+
+    // Otherwise, update with the actual tracked value info.
+    iter.first->getSecond() = getUnderlyingTrackedValueHelper(value);
+
+    // And return the value.
+    return iter.first->getSecond();
+  }
 };
 
 class RegionAnalysisFunctionInfo {
@@ -485,7 +548,9 @@ class RegionAnalysisFunctionInfo {
 
   bool isClosureCaptured(SILValue value, Operand *op);
 
-  static SILValue getUnderlyingTrackedValue(SILValue value);
+  SILValue getUnderlyingTrackedValue(SILValue value) {
+    return getValueMap().getRepresentative(value);
+  }
 
 private:
   void runDataflow();

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -93,19 +93,6 @@ regionanalysisimpl::getApplyIsolationCrossing(SILInstruction *inst) {
 
 namespace {
 
-struct UnderlyingTrackedValueInfo {
-  SILValue value;
-
-  /// Only used for addresses.
-  std::optional<ActorIsolation> actorIsolation;
-
-  explicit UnderlyingTrackedValueInfo(SILValue value) : value(value) {}
-
-  UnderlyingTrackedValueInfo(SILValue value,
-                             std::optional<ActorIsolation> actorIsolation)
-      : value(value), actorIsolation(actorIsolation) {}
-};
-
 struct UseDefChainVisitor
     : public AccessUseDefChainVisitor<UseDefChainVisitor, SILValue> {
   bool isMerge = false;
@@ -337,95 +324,6 @@ static bool isLookThroughIfOperandAndResultNonSendable(SILInstruction *inst) {
   }
 }
 
-static SILValue getUnderlyingTrackedObjectValue(SILValue value) {
-  auto *fn = value->getFunction();
-  SILValue result = value;
-  while (true) {
-    SILValue temp = result;
-
-    if (auto *svi = dyn_cast<SingleValueInstruction>(temp)) {
-      if (isStaticallyLookThroughInst(svi)) {
-        temp = svi->getOperand(0);
-      }
-
-      // If we have a cast and our operand and result are non-Sendable, treat it
-      // as a look through.
-      if (isLookThroughIfOperandAndResultNonSendable(svi)) {
-        if (SILIsolationInfo::isNonSendableType(svi->getType(), fn) &&
-            SILIsolationInfo::isNonSendableType(svi->getOperand(0)->getType(),
-                                                fn)) {
-          temp = svi->getOperand(0);
-        }
-      }
-
-      if (isLookThroughIfResultNonSendable(svi)) {
-        if (SILIsolationInfo::isNonSendableType(svi->getType(), fn)) {
-          temp = svi->getOperand(0);
-        }
-      }
-
-      if (isLookThroughIfOperandNonSendable(svi)) {
-        // If our operand is a non-Sendable type, look through this instruction.
-        if (SILIsolationInfo::isNonSendableType(svi->getOperand(0)->getType(),
-                                                fn)) {
-          temp = svi->getOperand(0);
-        }
-      }
-    }
-
-    if (auto *inst = temp->getDefiningInstruction()) {
-      if (isStaticallyLookThroughInst(inst)) {
-        temp = inst->getOperand(0);
-      }
-    }
-
-    if (temp != result) {
-      result = temp;
-      continue;
-    }
-
-    return result;
-  }
-}
-
-static UnderlyingTrackedValueInfo getUnderlyingTrackedValue(SILValue value) {
-  // Before a check if the value we are attempting to access is Sendable. In
-  // such a case, just return early.
-  if (!SILIsolationInfo::isNonSendableType(value))
-    return UnderlyingTrackedValueInfo(value);
-
-  // Look through a project_box, so that we process it like its operand object.
-  if (auto *pbi = dyn_cast<ProjectBoxInst>(value)) {
-    value = pbi->getOperand();
-  }
-
-  if (!value->getType().isAddress()) {
-    SILValue underlyingValue = getUnderlyingTrackedObjectValue(value);
-
-    // If we do not have a load inst, just return the value.
-    if (!isa<LoadInst, LoadBorrowInst>(underlyingValue)) {
-      return UnderlyingTrackedValueInfo(underlyingValue);
-    }
-
-    // If we got an address, lets see if we can do even better by looking at the
-    // address.
-    value = cast<SingleValueInstruction>(underlyingValue)->getOperand(0);
-  }
-  assert(value->getType().isAddress());
-
-  UseDefChainVisitor visitor;
-  SILValue base = visitor.visitAll(value);
-  assert(base);
-  if (base->getType().isObject())
-    return {getUnderlyingTrackedValue(base).value, visitor.actorIsolation};
-
-  return {base, visitor.actorIsolation};
-}
-
-SILValue RegionAnalysisFunctionInfo::getUnderlyingTrackedValue(SILValue value) {
-  return ::getUnderlyingTrackedValue(value).value;
-}
-
 namespace {
 
 struct TermArgSources {
@@ -915,6 +813,97 @@ void RegionAnalysisValueMap::print(llvm::raw_ostream &os) const {
 #endif
 }
 
+static SILValue getUnderlyingTrackedObjectValue(SILValue value) {
+  auto *fn = value->getFunction();
+  SILValue result = value;
+  while (true) {
+    SILValue temp = result;
+
+    if (auto *svi = dyn_cast<SingleValueInstruction>(temp)) {
+      if (isStaticallyLookThroughInst(svi)) {
+        temp = svi->getOperand(0);
+      }
+
+      // If we have a cast and our operand and result are non-Sendable, treat it
+      // as a look through.
+      if (isLookThroughIfOperandAndResultNonSendable(svi)) {
+        if (SILIsolationInfo::isNonSendableType(svi->getType(), fn) &&
+            SILIsolationInfo::isNonSendableType(svi->getOperand(0)->getType(),
+                                                fn)) {
+          temp = svi->getOperand(0);
+        }
+      }
+
+      if (isLookThroughIfResultNonSendable(svi)) {
+        if (SILIsolationInfo::isNonSendableType(svi->getType(), fn)) {
+          temp = svi->getOperand(0);
+        }
+      }
+
+      if (isLookThroughIfOperandNonSendable(svi)) {
+        // If our operand is a non-Sendable type, look through this instruction.
+        if (SILIsolationInfo::isNonSendableType(svi->getOperand(0)->getType(),
+                                                fn)) {
+          temp = svi->getOperand(0);
+        }
+      }
+    }
+
+    if (auto *inst = temp->getDefiningInstruction()) {
+      if (isStaticallyLookThroughInst(inst)) {
+        temp = inst->getOperand(0);
+      }
+    }
+
+    if (temp != result) {
+      result = temp;
+      continue;
+    }
+
+    return result;
+  }
+}
+
+RegionAnalysisValueMap::UnderlyingTrackedValueInfo
+RegionAnalysisValueMap::getUnderlyingTrackedValueHelper(SILValue value) const {
+  // Before a check if the value we are attempting to access is Sendable. In
+  // such a case, just return early.
+  if (!SILIsolationInfo::isNonSendableType(value))
+    return UnderlyingTrackedValueInfo(value);
+
+  // Look through a project_box, so that we process it like its operand object.
+  if (auto *pbi = dyn_cast<ProjectBoxInst>(value)) {
+    value = pbi->getOperand();
+  }
+
+  if (!value->getType().isAddress()) {
+    SILValue underlyingValue = getUnderlyingTrackedObjectValue(value);
+
+    // If we do not have a load inst, just return the value.
+    if (!isa<LoadInst, LoadBorrowInst>(underlyingValue)) {
+      return UnderlyingTrackedValueInfo(underlyingValue);
+    }
+
+    // If we got an address, lets see if we can do even better by looking at the
+    // address.
+    value = cast<SingleValueInstruction>(underlyingValue)->getOperand(0);
+  }
+  assert(value->getType().isAddress());
+
+  UseDefChainVisitor visitor;
+  SILValue base = visitor.visitAll(value);
+  assert(base);
+  if (base->getType().isObject()) {
+    // NOTE: We purposely recurse into the cached version of our computation
+    // rather than recurse into getUnderlyingTrackedObjectValueHelper. This is
+    // safe since we know that value was previously an address so if our base is
+    // an object, it cannot be the same object.
+    return {getUnderlyingTrackedValue(base).value, visitor.actorIsolation};
+  }
+
+  return {base, visitor.actorIsolation};
+}
+
 //===----------------------------------------------------------------------===//
 //                            MARK: TrackableValue
 //===----------------------------------------------------------------------===//
@@ -988,6 +977,7 @@ namespace {
 ///
 /// 2. Just computing reachability early is a very easy way to do this.
 struct PartialApplyReachabilityDataflow {
+  RegionAnalysisValueMap &valueMap;
   PostOrderFunctionInfo *pofi;
   llvm::DenseMap<SILValue, unsigned> valueToBit;
   std::vector<std::pair<SILValue, SILInstruction *>> valueToGenInsts;
@@ -1002,8 +992,10 @@ struct PartialApplyReachabilityDataflow {
   BasicBlockData<BlockState> blockData;
   bool propagatedReachability = false;
 
-  PartialApplyReachabilityDataflow(SILFunction *fn, PostOrderFunctionInfo *pofi)
-      : pofi(pofi), blockData(fn) {}
+  PartialApplyReachabilityDataflow(SILFunction *fn,
+                                   RegionAnalysisValueMap &valueMap,
+                                   PostOrderFunctionInfo *pofi)
+      : valueMap(valueMap), pofi(pofi), blockData(fn) {}
 
   /// Begin tracking an operand of a partial apply.
   void add(Operand *op);
@@ -1035,7 +1027,7 @@ struct PartialApplyReachabilityDataflow {
 
 private:
   SILValue getRootValue(SILValue value) const {
-    return getUnderlyingTrackedValue(value).value;
+    return valueMap.getRepresentative(value);
   }
 
   unsigned getBitForValue(SILValue value) const {
@@ -1767,7 +1759,8 @@ class PartitionOpTranslator {
                         RegionAnalysisValueMap &valueMap,
                         IsolationHistory::Factory &historyFactory)
       : function(function), functionArgPartition(), builder(),
-        partialApplyReachabilityDataflow(function, pofi), valueMap(valueMap) {
+        partialApplyReachabilityDataflow(function, valueMap, pofi),
+        valueMap(valueMap) {
     builder.translator = this;
 
     REGIONBASEDISOLATION_LOG(