[region-isolation] Rather than tracking task isolated values via a separate non-transferrable array... just track it by using ValueIsolationRegionInfo on a value.

gottesmm · gottesmm · commit b2c85a8294ce · 2024-03-10T22:08:40.000-07:00
In a subsequent commit, this is going to let me begin handling parameters with
actor regions in a nice way (and standardize all of the errors).

This is meant to be a refactoring commit that uses the current tests in tree to
make sure I did it correctly, so no tests need to be updated.
diff --git a/include/swift/SILOptimizer/Analysis/RegionAnalysis.h b/include/swift/SILOptimizer/Analysis/RegionAnalysis.h
@@ -139,33 +139,34 @@ class ValueIsolationRegionInfo {
   enum Kind {
     Unknown,
     Disconnected,
+    Task,
     Actor,
   };
 
 private:
   Kind kind;
-
-  /// When this is set it corresponds to a specific ActorIsolation from the AST
-  /// that we found.
-  ///
-  /// NOTE: actorInstanceType and actorIsolation should never be both set!
-  std::optional<ActorIsolation> actorIsolation;
-
-  /// When this is set it corresponds to a specific Actor type that we
-  /// identified as being actor isolated from a SIL level type. We do not have
-  /// the ActorIsolation information, so this is artificial.
-  ///
-  /// NOTE: actorInstanceType and actorIsolation should never be both set!
-  NominalTypeDecl *actorInstanceType = nullptr;
+  // clang-format off
+  std::variant<
+    // Used for actor isolated when we have ActorIsolation info from the AST.
+    std::optional<ActorIsolation>,
+    // Used for actor isolation when we infer the actor at the SIL level.
+    NominalTypeDecl *,
+    // The task isolated parameter when we find a task isolated value.
+    SILValue
+  > data;
+  // clang-format on
 
   ValueIsolationRegionInfo(Kind kind,
-                           std::optional<ActorIsolation> actorIsolation,
-                           NominalTypeDecl *actorInstanceType = nullptr)
-      : kind(kind), actorIsolation(actorIsolation),
-        actorInstanceType(actorInstanceType) {}
+                           std::optional<ActorIsolation> actorIsolation)
+      : kind(kind), data(actorIsolation) {}
+  ValueIsolationRegionInfo(Kind kind, NominalTypeDecl *decl)
+      : kind(kind), data(decl) {}
+
+  ValueIsolationRegionInfo(Kind kind, SILValue value)
+      : kind(kind), data(value) {}
 
 public:
-  ValueIsolationRegionInfo() : kind(Kind::Unknown), actorIsolation() {}
+  ValueIsolationRegionInfo() : kind(Kind::Unknown), data() {}
 
   operator bool() const { return kind != Kind::Unknown; }
 
@@ -175,6 +176,7 @@ class ValueIsolationRegionInfo {
 
   bool isDisconnected() const { return kind == Kind::Disconnected; }
   bool isActorIsolated() const { return kind == Kind::Actor; }
+  bool isTaskIsolated() const { return kind == Kind::Task; }
 
   void print(llvm::raw_ostream &os) const {
     switch (Kind(*this)) {
@@ -187,6 +189,9 @@ class ValueIsolationRegionInfo {
     case Actor:
       os << "actor";
       return;
+    case Task:
+      os << "task";
+      return;
     }
   }
 
@@ -196,14 +201,24 @@ class ValueIsolationRegionInfo {
   }
 
   std::optional<ActorIsolation> getActorIsolation() const {
-    assert(!actorInstanceType &&
-           "Should never be set if getActorIsolation is called");
-    return actorIsolation;
+    assert(kind == Actor);
+    assert(std::holds_alternative<NominalTypeDecl *>(data) &&
+           "Doesn't have an actor isolation?!");
+    return std::get<std::optional<ActorIsolation>>(data);
   }
 
   NominalTypeDecl *getActorInstance() const {
-    assert(!actorIsolation.has_value());
-    return actorInstanceType;
+    assert(kind == Actor);
+    assert(std::holds_alternative<NominalTypeDecl *>(data) &&
+           "Doesn't have an actor instance?!");
+    return std::get<NominalTypeDecl *>(data);
+  }
+
+  SILValue getTaskIsolatedValue() const {
+    assert(kind == Task);
+    assert(std::holds_alternative<SILValue>(data) &&
+           "Doesn't have a task isolated value");
+    return std::get<SILValue>(data);
   }
 
   [[nodiscard]] ValueIsolationRegionInfo
@@ -240,9 +255,13 @@ class ValueIsolationRegionInfo {
     if (actorIsolation.isActorIsolated())
       return getActorIsolated(actorIsolation);
     if (nomDecl->isActor())
-      return {Kind::Actor, {}, nomDecl};
+      return {Kind::Actor, nomDecl};
     return {};
   }
+
+  static ValueIsolationRegionInfo getTaskIsolated(SILValue value) {
+    return {Kind::Task, value};
+  }
 };
 
 } // namespace regionanalysisimpl
@@ -432,11 +451,6 @@ class RegionAnalysisValueMap {
       equivalenceClassValuesToState;
   llvm::DenseMap<unsigned, RepresentativeValue> stateIndexToEquivalenceClass;
 
-  /// A list of values that can never be transferred.
-  ///
-  /// This only includes function arguments.
-  std::vector<TrackableValueID> neverTransferredValueIDs;
-
   SILFunction *fn;
 
 public:
@@ -451,12 +465,7 @@ class RegionAnalysisValueMap {
   SILValue maybeGetRepresentative(Element trackableValueID) const;
 
   ValueIsolationRegionInfo getIsolationRegion(Element trackableValueID) const;
-
-  ArrayRef<Element> getNonTransferrableElements() const {
-    return neverTransferredValueIDs;
-  }
-
-  void sortUniqueNeverTransferredValues();
+  ValueIsolationRegionInfo getIsolationRegion(SILValue trackableValueID) const;
 
   void print(llvm::raw_ostream &os) const;
   SWIFT_DEBUG_DUMP { print(llvm::dbgs()); }
@@ -473,9 +482,6 @@ class RegionAnalysisValueMap {
                                     ValueIsolationRegionInfo isolation) const;
   bool mergeIsolationRegionInfo(SILValue value,
                                 ValueIsolationRegionInfo isolation);
-  void addNeverTransferredValueID(TrackableValueID valueID) {
-    neverTransferredValueIDs.push_back(valueID);
-  }
   bool valueHasID(SILValue value, bool dumpIfHasNoID = false);
   TrackableValueID lookupValueID(SILValue value);
 };
diff --git a/include/swift/SILOptimizer/Utils/PartitionUtils.h b/include/swift/SILOptimizer/Utils/PartitionUtils.h
@@ -1040,21 +1040,26 @@ struct PartitionOpEvaluator {
     return asImpl().handleTransferNonTransferrable(op, elt);
   }
 
+  /// Just call our CRTP subclass.
+  void handleTransferNonTransferrable(const PartitionOp &op, Element elt,
+                                      Element otherElement) const {
+    return asImpl().handleTransferNonTransferrable(op, elt, otherElement);
+  }
+
   /// Call isActorDerived on our CRTP subclass.
   bool isActorDerived(Element elt) const {
     return asImpl().isActorDerived(elt);
   }
 
+  bool isTaskIsolatedDerived(Element elt) const {
+    return asImpl().isTaskIsolatedDerived(elt);
+  }
+
   /// Call isClosureCaptured on our CRTP subclass.
   bool isClosureCaptured(Element elt, Operand *op) const {
     return asImpl().isClosureCaptured(elt, op);
   }
 
-  /// Call getNonTransferrableElements() on our CRTP subclass.
-  ArrayRef<Element> getNonTransferrableElements() const {
-    return asImpl().getNonTransferrableElements();
-  }
-
   /// Apply \p op to the partition op.
   void apply(const PartitionOp &op) const {
     if (shouldEmitVerboseLogging()) {
@@ -1098,23 +1103,8 @@ struct PartitionOpEvaluator {
       assert(p.isTrackingElement(op.getOpArgs()[0]) &&
              "Transfer PartitionOp's argument should already be tracked");
 
-      // check if any nontransferrables are transferred here, and handle the
-      // failure if so
-      for (Element nonTransferrable : getNonTransferrableElements()) {
-        assert(
-            p.isTrackingElement(nonTransferrable) &&
-            "nontransferrables should be function args and self, and therefore"
-            "always present in the label map because of initialization at "
-            "entry");
-        if (!p.isTransferred(nonTransferrable) &&
-            p.areElementsInSameRegion(nonTransferrable, op.getOpArgs()[0])) {
-          return handleTransferNonTransferrable(op, nonTransferrable);
-        }
-      }
-
-      // If this value is actor derived or if any elements in its region are
-      // actor derived, we need to treat as nontransferrable.
-      if (isActorDerived(op.getOpArgs()[0]))
+      if (isActorDerived(op.getOpArgs()[0]) ||
+          isTaskIsolatedDerived(op.getOpArgs()[0]))
         return handleTransferNonTransferrable(op, op.getOpArgs()[0]);
 
       // NOTE: We purposely do not check here if a transferred value is already
@@ -1131,8 +1121,10 @@ struct PartitionOpEvaluator {
 
       Region elementRegion = p.getRegion(op.getOpArgs()[0]);
       for (const auto &pair : p.range()) {
-        if (pair.second == elementRegion && isActorDerived(pair.first))
-          return handleTransferNonTransferrable(op, op.getOpArgs()[0]);
+        if (pair.second == elementRegion &&
+            (isActorDerived(pair.first) || isTaskIsolatedDerived(pair.first)))
+          return handleTransferNonTransferrable(op, op.getOpArgs()[0],
+                                                pair.first);
         isClosureCapturedElt |= isClosureCaptured(pair.first, op.getSourceOp());
       }
 
@@ -1229,22 +1221,26 @@ struct PartitionOpEvaluatorBaseImpl : PartitionOpEvaluator<Subclass> {
   void handleFailure(const PartitionOp &op, Element elt,
                      TransferringOperand transferringOp) const {}
 
-  /// A list of elements that cannot be transferred. Whenever we transfer, we
-  /// check this list to see if we are transferring the element and then call
-  /// transferNonTransferrableCallback. This should consist only of function
-  /// arguments.
-  ArrayRef<Element> getNonTransferrableElements() const { return {}; }
-
   /// This is called if we detect a never transferred element that was passed to
   /// a transfer instruction.
   void handleTransferNonTransferrable(const PartitionOp &op,
                                       Element elt) const {}
 
+  /// This is called if we detect a never transferred element that was passed to
+  /// a transfer instruction but the actual element that could not be
+  /// transferred is a different element in its region.
+  void handleTransferNonTransferrable(const PartitionOp &op, Element elt,
+                                      Element otherEltInRegion) const {}
+
   /// This is used to determine if an element is actor derived. If we determine
   /// that a region containing such an element is transferred, we emit an error
   /// since actor regions cannot be transferred.
   bool isActorDerived(Element elt) const { return false; }
 
+  /// This is used to determine if an element is in the same region as a task
+  /// isolated value.
+  bool isTaskIsolatedDerived(Element elt) const { return false; }
+
   /// Check if the representative value of \p elt is closure captured at \p
   /// op.
   ///
diff --git a/lib/SILOptimizer/Analysis/RegionAnalysis.cpp b/lib/SILOptimizer/Analysis/RegionAnalysis.cpp
@@ -544,7 +544,7 @@ getGlobalActorInitIsolation(SILFunction *fn) {
 /// in the body of our function.
 static bool isTransferrableFunctionArgument(SILFunctionArgument *arg) {
   // Indirect out parameters cannot be an input transferring parameter.
-  if (arg->getArgumentConvention().isIndirectOutParameter())
+  if (arg->isIndirectResult() || arg->isIndirectErrorResult())
     return false;
 
   // If we have a function argument that is closure captured by a Sendable
@@ -563,6 +563,10 @@ static bool isTransferrableFunctionArgument(SILFunctionArgument *arg) {
   // 2. If we have an async-let based Sendable closure, we want to allow
   // for the argument to be transferred in the async let's statement and
   // not emit an error.
+  //
+  // TODO: Once the async let refactoring change this will no longer be needed
+  // since closure captures will have transferring parameters and be
+  // non-Sendable.
   if (arg->isClosureCapture() &&
       arg->getFunction()->getLoweredFunctionType()->isSendable())
     return true;
@@ -1414,7 +1418,8 @@ class PartitionOpTranslator {
         // Otherwise, it is one of our merged parameters. Add it to the never
         // transfer list and to the region join list.
         LLVM_DEBUG(llvm::dbgs() << "    %%" << state->getID() << ": " << *arg);
-        valueMap.addNeverTransferredValueID(state->getID());
+        valueMap.mergeIsolationRegionInfo(
+            arg, ValueIsolationRegionInfo::getTaskIsolated(arg));
         nonSendableJoinedIndices.push_back(state->getID());
       }
     }
@@ -1466,10 +1471,6 @@ class PartitionOpTranslator {
     return valueMap.mergeIsolationRegionInfo(value, isolationRegion);
   }
 
-  void addNeverTransferredValueID(TrackableValueID valueID) {
-    valueMap.addNeverTransferredValueID(valueID);
-  }
-
   bool valueHasID(SILValue value, bool dumpIfHasNoID = false) {
     return valueMap.valueHasID(value, dumpIfHasNoID);
   }
@@ -1478,18 +1479,6 @@ class PartitionOpTranslator {
     return valueMap.lookupValueID(value);
   }
 
-  void sortUniqueNeverTransferredValues() {
-    valueMap.sortUniqueNeverTransferredValues();
-  }
-
-  /// Get the vector of IDs that cannot be legally transferred at any point in
-  /// this function.
-  ArrayRef<TrackableValueID> getNeverTransferredValues() const {
-    return valueMap.getNonTransferrableElements();
-  }
-
-  // ===========================================================================
-
 public:
   /// Return the partition consisting of all function arguments.
   ///
@@ -3098,10 +3087,6 @@ void RegionAnalysisFunctionInfo::runDataflow() {
       }
     }
   }
-
-  // Now that we have finished processing, sort/unique our non transferred
-  // array.
-  translator->getValueMap().sortUniqueNeverTransferredValues();
 }
 
 //===----------------------------------------------------------------------===//
@@ -3137,6 +3122,14 @@ RegionAnalysisValueMap::getIsolationRegion(Element trackableValueID) const {
   return iter->getValueState().getIsolationRegionInfo();
 }
 
+ValueIsolationRegionInfo
+RegionAnalysisValueMap::getIsolationRegion(SILValue value) const {
+  auto iter = equivalenceClassValuesToState.find(RepresentativeValue(value));
+  if (iter == equivalenceClassValuesToState.end())
+    return {};
+  return iter->getSecond().getIsolationRegionInfo();
+}
+
 /// If \p isAddressCapturedByPartialApply is set to true, then this value is
 /// an address that is captured by a partial_apply and we want to treat it as
 /// may alias.
@@ -3213,9 +3206,11 @@ TrackableValue RegionAnalysisValueMap::getTrackableValue(
     }
   }
 
-  // Otherwise refer to the oracle.
-  if (!isNonSendableType(value->getType(), fn))
+  // Otherwise refer to the oracle. If we have a Sendable value, just return.
+  if (!isNonSendableType(value->getType(), fn)) {
     iter.first->getSecond().addFlag(TrackableValueFlag::isSendable);
+    return {iter.first->first, iter.first->second};
+  }
 
   // Check if our base is a ref_element_addr from an actor. In such a case,
   // mark this value as actor derived.
@@ -3285,8 +3280,15 @@ TrackableValue RegionAnalysisValueMap::getTrackableValue(
     }
   }
 
-  // If our access storage is from a class, then see if we have an actor. In
-  // such a case, we need to add this id to the neverTransferred set.
+  // See if we have a non-transferring argument from a function. In such a case,
+  // mark the value as task isolated.
+  if (auto *fArg =
+          dyn_cast<SILFunctionArgument>(iter.first->first.getValue())) {
+    if (!isTransferrableFunctionArgument(fArg)) {
+      iter.first->getSecond().mergeIsolationRegionInfo(
+          ValueIsolationRegionInfo::getTaskIsolated(fArg));
+    }
+  }
 
   return {iter.first->first, iter.first->second};
 }
@@ -3348,10 +3350,6 @@ TrackableValueID RegionAnalysisValueMap::lookupValueID(SILValue value) {
   return state.getID();
 }
 
-void RegionAnalysisValueMap::sortUniqueNeverTransferredValues() {
-  sortUnique(neverTransferredValueIDs);
-}
-
 void RegionAnalysisValueMap::print(llvm::raw_ostream &os) const {
 #ifndef NDEBUG
   // Since this is just used for debug output, be inefficient to make nicer
diff --git a/lib/SILOptimizer/Mandatory/TransferNonSendable.cpp b/lib/SILOptimizer/Mandatory/TransferNonSendable.cpp
