[region-isolation] Implement MergeIsolationRegionInst.

I am adding this instruction to express artificially that two non-Sendable
values should be part of the same region. It is meant to be used in cases where
due to unsafe code using Sendable, we stop propagating a non-Sendable dependency
that needs to be made in the same region of a use of said Sendable value. I
included an example in ./docs/SIL.rst of where this comes up with @out results
of continuations.

Author: gottesmm

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -1612,3 +1612,6 @@ final public class CheckedCastAddrBranchInst : TermInst {
 
 final public class ThunkInst : Instruction {
 }
+
+final public class MergeIsolationRegionInst : Instruction {
+}

SwiftCompilerSources/Sources/SIL/Registration.swift

@@ -257,4 +257,5 @@ public func registerSILClasses() {
   register(CheckedCastBranchInst.self)
   register(CheckedCastAddrBranchInst.self)
   register(ThunkInst.self)
+  register(MergeIsolationRegionInst.self)
 }

docs/SIL.rst

@@ -4126,6 +4126,91 @@ pipeline.
 
 The operand is a guaranteed operand, i.e. not consumed.
 
+merge_isolation_region
+``````````````````````
+
+::
+
+  sil-instruction :: 'merge_isolation_region' (sil-operand ',')+ sil-operand
+
+  %2 = merge_isolation_region %first : $*T, %second : $U
+  %2 = merge_isolation_region %first : $*T, %second : $U, %third : $H
+
+Instruction that is only valid in Ownership SSA.
+
+This instruction informs region isolation that all of the operands should be
+considered to be artificially apart of the same region. It is intended to be
+used to express region dependency when due to unsafe codegen we have to traffic
+a non-Sendable value through computations with Sendable values (causing us to
+not track the non-Sendable value) but have to later express that a non-Sendable
+result of using the Sendable value needs to be in the same region as the
+original non-Sendable value. As an example of where this comes up, consider the
+following code::
+
+  // objc code
+  @interface CallbackData : NSObject
+  @end
+  
+  @interface Klass : NSObject
+  
+  - (void)loadDataWithCompletionHandler:(void (^)(CallbackData * _Nullable, NSError * _Nullable))completionHandler;
+  
+  @end
+
+  // swift code
+  extension Klass {
+    func loadCallbackData() async throws -> sending CallbackData {
+      try await loadData()
+    }
+  }
+
+This lowers to::
+
+  %5 = alloc_stack $CallbackData                  // users: %26, %25, %31, %16, %7
+  %6 = objc_method %0 : $Klass, #Klass.loadData!foreign : (Klass) -> () async throws -> CallbackData, $@convention(objc_method) (Optional<@convention(block) (Optional<CallbackData>, Optional<NSError>) -> ()>, Klass) -> () // user: %20
+  %7 = get_async_continuation_addr [throws] CallbackData, %5 : $*CallbackData // users: %23, %8
+  %8 = struct $UnsafeContinuation<CallbackData, any Error> (%7 : $Builtin.RawUnsafeContinuation) // user: %14
+  %9 = alloc_stack $@block_storage Any            // users: %22, %16, %10
+  %10 = project_block_storage %9 : $*@block_storage Any // user: %11
+  %11 = init_existential_addr %10 : $*Any, $CheckedContinuation<CallbackData, any Error> // user: %15
+  // function_ref _createCheckedThrowingContinuation<A>(_:)
+  %12 = function_ref @$ss34_createCheckedThrowingContinuationyScCyxs5Error_pGSccyxsAB_pGnlF : $@convention(thin) <τ_0_0> (UnsafeContinuation<τ_0_0, any Error>) -> @out CheckedContinuation<τ_0_0, any Error> // user: %14
+  %13 = alloc_stack $CheckedContinuation<CallbackData, any Error> // users: %21, %15, %14
+  %14 = apply %12<CallbackData>(%13, %8) : $@convention(thin) <τ_0_0> (UnsafeContinuation<τ_0_0, any Error>) -> @out CheckedContinuation<τ_0_0, any Error>
+  copy_addr [take] %13 to [init] %11 : $*CheckedContinuation<CallbackData, any Error> // id: %15
+  merge_isolation_region %9 : $*@block_storage Any, %5 : $*CallbackData // id: %16
+  // function_ref @objc completion handler block implementation for @escaping @callee_unowned @convention(block) (@unowned CallbackData?, @unowned NSError?) -> () with result type CallbackData
+  %17 = function_ref @$sSo12CallbackDataCSgSo7NSErrorCSgIeyByy_ABTz_ : $@convention(c) (@inout_aliasable @block_storage Any, Optional<CallbackData>, Optional<NSError>) -> () // user: %18
+  %18 = init_block_storage_header %9 : $*@block_storage Any, invoke %17 : $@convention(c) (@inout_aliasable @block_storage Any, Optional<CallbackData>, Optional<NSError>) -> (), type $@convention(block) (Optional<CallbackData>, Optional<NSError>) -> () // user: %19
+  %19 = enum $Optional<@convention(block) (Optional<CallbackData>, Optional<NSError>) -> ()>, #Optional.some!enumelt, %18 : $@convention(block) (Optional<CallbackData>, Optional<NSError>) -> () // user: %20
+  %20 = apply %6(%19, %0) : $@convention(objc_method) (Optional<@convention(block) (Optional<CallbackData>, Optional<NSError>) -> ()>, Klass) -> ()
+
+Notice how without the `merge_isolation_region`_ instruction (%16) there is no
+non-Sendable def-use chain from %5, the indirect return value of the block, to
+the actual non-Sendable block storage %9. This can result in region isolation
+not propagating restrictions on usage from %9 onto %5 risking the creation of
+races.
+
+Applying the previous discussion to this specific example, self (%0) is
+non-Sendable and is bound to the current task. If we did not have the
+`merge_isolation_region`_ instruction here, we would not tie the return value %5
+to %0 via %9. This would cause %5 to be treated as a disconnected value and thus
+be a valid sending return value potentially allowing for %5 in the caller of the
+function to be sent to another isolation domain and introduce a race.
+
+.. note::
+   This is effectively the same purpose that `mark_dependence`_ plays for memory
+   dependence (expressing memory dependence that the compiler cannot infer)
+   except in the world of region isolation. We purposely use a different
+   instruction since `mark_dependence`_ is often times used to create a
+   temporary dependence in between two values via the return value of
+   `mark_dependence`_. If `mark_dependence`_ had the semantics of acting like a
+   region merge we would in contrast have from that point on a region dependence
+   in between the base and value of the `mark_dependence`_ causing the
+   `mark_dependence`_ to have a less "local" effect since all paths through that
+   program point would have to maintain that region dependence until the end of
+   the function.
+
 dealloc_stack
 `````````````
 ::
@@ -9049,7 +9134,6 @@ NOTE: From the perspective of the address checker, a trivial `load`_ with a
 `moveonlywrapper_to_copyable_addr`_ operand is considered to be a use of a
 noncopyable type.
 
-
 Assertion configuration
 ~~~~~~~~~~~~~~~~~~~~~~~
 

include/swift/SIL/AddressWalker.h

@@ -215,7 +215,8 @@ TransitiveAddressWalker<Impl>::walk(SILValue projectedAddress) && {
         isa<RetainValueAddrInst>(user) || isa<ReleaseValueAddrInst>(user) ||
         isa<PackElementSetInst>(user) || isa<PackElementGetInst>(user) ||
         isa<DeinitExistentialAddrInst>(user) || isa<LoadBorrowInst>(user) ||
-        isa<TupleAddrConstructorInst>(user) || isa<DeallocPackInst>(user)) {
+        isa<TupleAddrConstructorInst>(user) || isa<DeallocPackInst>(user) ||
+        isa<MergeIsolationRegionInst>(user)) {
       callVisitUse(op);
       continue;
     }

include/swift/SIL/SILBuilder.h

@@ -2579,6 +2579,12 @@ class SILBuilder {
                     getSILDebugLocation(loc), fnValue, resultType));
   }
 
+  MergeIsolationRegionInst *
+  createMergeIsolationRegion(SILLocation loc, ArrayRef<SILValue> args) {
+    return insert(MergeIsolationRegionInst::create(getSILDebugLocation(loc),
+                                                   args, getModule()));
+  }
+
   //===--------------------------------------------------------------------===//
   // Terminator SILInstruction Creation Methods
   //===--------------------------------------------------------------------===//

include/swift/SIL/SILCloner.h

@@ -1168,6 +1168,15 @@ SILCloner<ImplClass>::visitBuiltinInst(BuiltinInst *Inst) {
                 getOpSubstitutionMap(Inst->getSubstitutions()), Args));
 }
 
+template <typename ImplClass>
+void SILCloner<ImplClass>::visitMergeIsolationRegionInst(
+    MergeIsolationRegionInst *Inst) {
+  auto Args = getOpValueArray<8>(Inst->getOperandValues());
+  getBuilder().setCurrentDebugScope(getOpScope(Inst->getDebugScope()));
+  recordClonedInstruction(Inst, getBuilder().createMergeIsolationRegion(
+                                    getOpLocation(Inst->getLoc()), Args));
+}
+
 template<typename ImplClass>
 void
 SILCloner<ImplClass>::visitApplyInst(ApplyInst *Inst) {

include/swift/SIL/SILInstruction.h

@@ -11666,6 +11666,25 @@ class TypeValueInst final
   }
 };
 
+class MergeIsolationRegionInst final
+    : public InstructionBaseWithTrailingOperands<
+          SILInstructionKind::MergeIsolationRegionInst,
+          MergeIsolationRegionInst, NonValueInstruction> {
+  friend SILBuilder;
+
+  MergeIsolationRegionInst(SILDebugLocation loc, ArrayRef<SILValue> operands)
+      : InstructionBaseWithTrailingOperands(operands, loc) {}
+
+  static MergeIsolationRegionInst *
+  create(SILDebugLocation loc, ArrayRef<SILValue> operands, SILModule &mod);
+
+public:
+  /// Return the SILValues for all operands of this instruction.
+  OperandValueArrayRef getArguments() const {
+    return OperandValueArrayRef(getAllOperands());
+  }
+};
+
 inline SILType *AllocRefInstBase::getTypeStorage() {
   // If the size of the subclasses are equal, then all of this compiles away.
   if (auto I = dyn_cast<AllocRefInst>(this))

include/swift/SIL/SILNode.h

@@ -317,6 +317,7 @@ class alignas(8) SILNode :
     SHARED_FIELD(SILFunctionArgument, uint32_t noImplicitCopy : 1,
                  lifetimeAnnotation : 2, closureCapture : 1,
                  parameterPack : 1);
+    SHARED_FIELD(MergeRegionIsolationInst, uint32_t numOperands);
 
     // Do not use `_sharedUInt32_private` outside of SILNode.
   } _sharedUInt32_private;

include/swift/SIL/SILNodes.def

@@ -894,6 +894,9 @@ NON_VALUE_INST(CondFailInst, cond_fail,
 NON_VALUE_INST(MarkUnresolvedMoveAddrInst, mark_unresolved_move_addr,
                SILInstruction, None, DoesNotRelease)
 
+NON_VALUE_INST(MergeIsolationRegionInst, merge_isolation_region,
+               SILInstruction, None, DoesNotRelease)
+
 NON_VALUE_INST(IncrementProfilerCounterInst, increment_profiler_counter,
                SILInstruction, MayReadWrite, DoesNotRelease)
 

lib/IRGen/IRGenSIL.cpp

@@ -1283,6 +1283,11 @@ class IRGenSILFunction :
     auto e = getLoweredExplosion(i->getOperand());
     setLoweredExplosion(i, e);
   }
+
+  void visitMergeIsolationRegionInst(MergeIsolationRegionInst *i) {
+    llvm_unreachable("Valid only when ownership is enabled");
+  }
+
   void visitReleaseValueInst(ReleaseValueInst *i);
   void visitReleaseValueAddrInst(ReleaseValueAddrInst *i);
   void visitDestroyValueInst(DestroyValueInst *i);

lib/SIL/IR/OperandOwnership.cpp

@@ -228,6 +228,7 @@ OPERAND_OWNERSHIP(InstantaneousUse, ClassifyBridgeObject)
 OPERAND_OWNERSHIP(InstantaneousUse, UnownedCopyValue)
 OPERAND_OWNERSHIP(InstantaneousUse, WeakCopyValue)
 OPERAND_OWNERSHIP(InstantaneousUse, ExtendLifetime)
+OPERAND_OWNERSHIP(InstantaneousUse, MergeIsolationRegion)
 #define REF_STORAGE(Name, ...)                                                 \
   OPERAND_OWNERSHIP(InstantaneousUse, StrongCopy##Name##Value)
 #include "swift/AST/ReferenceStorage.def"

lib/SIL/IR/SILInstructions.cpp

@@ -3552,3 +3552,11 @@ TypeValueInst *TypeValueInst::create(SILFunction &F, SILDebugLocation loc,
   return ::new (buffer)
       TypeValueInst(loc, typeDependentOperands, valueType, paramType);
 }
+
+MergeIsolationRegionInst *
+MergeIsolationRegionInst::create(SILDebugLocation loc, ArrayRef<SILValue> args,
+                                 SILModule &mod) {
+  auto size = totalSizeToAlloc<swift::Operand>(args.size());
+  auto buffer = mod.allocateInst(size, alignof(MergeIsolationRegionInst));
+  return ::new (buffer) MergeIsolationRegionInst(loc, args);
+}

lib/SIL/IR/SILPrinter.cpp

@@ -1658,7 +1658,13 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
     *this << ") : ";
     *this << BI->getType();
   }
-  
+
+  void visitMergeIsolationRegionInst(MergeIsolationRegionInst *mir) {
+    llvm::interleave(
+        mir->getArguments(), [&](SILValue v) { *this << getIDAndType(v); },
+        [&] { *this << ", "; });
+  }
+
   void visitAllocGlobalInst(AllocGlobalInst *AGI) {
     if (AGI->getReferencedGlobal()) {
       AGI->getReferencedGlobal()->printName(PrintState.OS);

lib/SIL/Parser/ParseSIL.cpp

@@ -3030,6 +3030,21 @@ bool SILParser::parseSpecificSILInstruction(SILBuilder &B,
     ResultVal = B.createBuiltin(InstLoc, Id, ResultTy, subMap, Args);
     break;
   }
+  case SILInstructionKind::MergeIsolationRegionInst: {
+    SmallVector<SILValue, 4> Args;
+    do {
+      SILValue Val;
+      if (parseTypedValueRef(Val, B))
+        return true;
+      Args.push_back(Val);
+    } while (P.consumeIf(tok::comma));
+
+    if (parseSILDebugLocation(InstLoc, B))
+      return true;
+
+    ResultVal = B.createMergeIsolationRegion(InstLoc, Args);
+    break;
+  }
   case SILInstructionKind::OpenExistentialAddrInst:
     if (parseOpenExistAddrKind() || parseTypedValueRef(Val, B) ||
         parseVerbatim("to") || parseSILType(Ty) ||

lib/SIL/Utils/InstructionUtils.cpp

@@ -517,6 +517,7 @@ RuntimeEffect swift::getRuntimeEffect(SILInstruction *inst, SILType &impactType)
   case SILInstructionKind::ClassifyBridgeObjectInst:
   case SILInstructionKind::ValueToBridgeObjectInst:
   case SILInstructionKind::MarkDependenceInst:
+  case SILInstructionKind::MergeIsolationRegionInst:
   case SILInstructionKind::MoveValueInst:
   case SILInstructionKind::DropDeinitInst:
   case SILInstructionKind::MarkUnresolvedNonCopyableValueInst:

lib/SIL/Verifier/SILVerifier.cpp

@@ -2187,6 +2187,11 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     }
   }
 
+  void checkMergeIsolationRegionInst(MergeIsolationRegionInst *mir) {
+    require(mir->getNumOperands() >= 2, "Must have at least two parameters");
+    require(F.hasOwnership(), "Only valid when OSSA is enabled");
+  }
+
   void checkMarkDependencInst(MarkDependenceInst *MDI) {
     require(isa<SILUndef>(MDI->getValue()) || MDI->getValue() != MDI->getBase(),
             "mark_dependence operands must be distinct");

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -2232,6 +2232,25 @@ class PartitionOpTranslator {
                              requireOperands);
   }
 
+  void translateSILMerge(MutableArrayRef<Operand> array,
+                         bool requireOperands = true) {
+    if (array.size() < 2)
+      return;
+
+    auto trackableDest = tryToTrackValue(array.front().get());
+    if (!trackableDest)
+      return;
+    for (Operand &op : array.drop_front()) {
+      if (auto trackableSrc = tryToTrackValue(op.get())) {
+        if (requireOperands) {
+          builder.addRequire(trackableSrc->getRepresentative().getValue());
+          builder.addRequire(trackableDest->getRepresentative().getValue());
+        }
+        builder.addMerge(trackableDest->getRepresentative().getValue(), &op);
+      }
+    }
+  }
+
   void translateSILAssignmentToTransferringParameter(TrackableValue destRoot,
                                                      Operand *destOperand,
                                                      TrackableValue srcRoot,
@@ -3176,6 +3195,13 @@ PartitionOpTranslator::visitMarkDependenceInst(MarkDependenceInst *mdi) {
   return TranslationSemantics::Special;
 }
 
+TranslationSemantics PartitionOpTranslator::visitMergeIsolationRegionInst(
+    MergeIsolationRegionInst *mir) {
+  // But we want to require and merge the base into the result.
+  translateSILMerge(mir->getAllOperands(), true /*require operands*/);
+  return TranslationSemantics::Special;
+}
+
 TranslationSemantics
 PartitionOpTranslator::visitPointerToAddressInst(PointerToAddressInst *ptai) {
   if (!isNonSendableType(ptai->getType())) {

lib/SILOptimizer/Mandatory/OwnershipModelEliminator.cpp

@@ -151,6 +151,11 @@ struct OwnershipModelEliminatorVisitor
     eraseInstruction(dbi);
     return true;
   }
+  bool visitMergeIsolationRegionInst(MergeIsolationRegionInst *mir) {
+    eraseInstruction(mir);
+    return true;
+  }
+
   bool visitLoadBorrowInst(LoadBorrowInst *lbi);
   bool visitMoveValueInst(MoveValueInst *mvi) {
     eraseInstructionAndRAUW(mvi, mvi->getOperand());

lib/SILOptimizer/UtilityPasses/SerializeSILPass.cpp

@@ -167,6 +167,7 @@ static bool hasOpaqueArchetype(TypeExpansionContext context,
   case SILInstructionKind::ClassifyBridgeObjectInst:
   case SILInstructionKind::ValueToBridgeObjectInst:
   case SILInstructionKind::MarkDependenceInst:
+  case SILInstructionKind::MergeIsolationRegionInst:
   case SILInstructionKind::CopyBlockInst:
   case SILInstructionKind::CopyBlockWithoutEscapingInst:
   case SILInstructionKind::CopyValueInst:

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -877,6 +877,7 @@ InlineCost swift::instructionInlineCost(SILInstruction &I) {
   case SILInstructionKind::BeginBorrowInst:
   case SILInstructionKind::BorrowedFromInst:
   case SILInstructionKind::MarkDependenceInst:
+  case SILInstructionKind::MergeIsolationRegionInst:
   case SILInstructionKind::PreviousDynamicFunctionRefInst:
   case SILInstructionKind::DynamicFunctionRefInst:
   case SILInstructionKind::FunctionRefInst: