Merge pull request swiftlang#35601 from gottesmm/ossa-sil-combine-3

[sil-combine] More ownership updates

Author: gottesmm

include/swift/SILOptimizer/Utils/OwnershipOptUtils.h

@@ -121,7 +121,8 @@ class OwnershipRAUWHelper {
   /// "forwarding" transformation must be performed upon \p newValue at \p
   /// oldValue's insertion point so that we can then here RAUW the transformed
   /// \p newValue.
-  SILBasicBlock::iterator perform();
+  SILBasicBlock::iterator
+  perform(SingleValueInstruction *maybeTransformedNewValue = nullptr);
 
 private:
   SILBasicBlock::iterator replaceAddressUses(SingleValueInstruction *oldValue,

lib/SILOptimizer/SILCombiner/SILCombinerCastVisitors.cpp

@@ -739,9 +739,6 @@ visitUncheckedTrivialBitCastInst(UncheckedTrivialBitCastInst *UTBCI) {
 SILInstruction *
 SILCombiner::
 visitUncheckedBitwiseCastInst(UncheckedBitwiseCastInst *UBCI) {
-  if (UBCI->getFunction()->hasOwnership())
-    return nullptr;
-
   // (unchecked_bitwise_cast Y->Z (unchecked_bitwise_cast X->Y x))
   // OR (unchecked_trivial_cast Y->Z (unchecked_bitwise_cast X->Y x))
   //   ->
@@ -750,27 +747,59 @@ visitUncheckedBitwiseCastInst(UncheckedBitwiseCastInst *UBCI) {
   if (match(UBCI->getOperand(),
             m_CombineOr(m_UncheckedBitwiseCastInst(m_SILValue(Oper)),
                         m_UncheckedTrivialBitCastInst(m_SILValue(Oper))))) {
-    return Builder.createUncheckedBitwiseCast(UBCI->getLoc(), Oper,
-                                              UBCI->getType());
+    if (!Builder.hasOwnership()) {
+      return Builder.createUncheckedBitwiseCast(UBCI->getLoc(), Oper,
+                                                UBCI->getType());
+    }
+
+    OwnershipRAUWHelper helper(ownershipFixupContext, UBCI, Oper);
+    if (helper) {
+      auto *transformedOper = Builder.createUncheckedBitwiseCast(
+          UBCI->getLoc(), Oper, UBCI->getType());
+      helper.perform(transformedOper);
+      return nullptr;
+    }
+  }
+
+  if (UBCI->getType().isTrivial(*UBCI->getFunction())) {
+    // If our result is trivial, we can always just RAUW.
+    return Builder.createUncheckedTrivialBitCast(
+        UBCI->getLoc(), UBCI->getOperand(), UBCI->getType());
   }
-  if (UBCI->getType().isTrivial(*UBCI->getFunction()))
-    return Builder.createUncheckedTrivialBitCast(UBCI->getLoc(),
-                                                 UBCI->getOperand(),
-                                                 UBCI->getType());
 
   if (!SILType::canRefCast(UBCI->getOperand()->getType(), UBCI->getType(),
                            Builder.getModule()))
     return nullptr;
 
-  return Builder.createUncheckedRefCast(UBCI->getLoc(), UBCI->getOperand(),
-                                        UBCI->getType());
+  if (!Builder.hasOwnership()) {
+    return Builder.createUncheckedRefCast(UBCI->getLoc(), UBCI->getOperand(),
+                                          UBCI->getType());
+  }
+
+  {
+    OwnershipRAUWHelper helper(ownershipFixupContext, UBCI, UBCI->getOperand());
+    if (helper) {
+      auto *newInst = Builder.createUncheckedRefCast(UBCI->getLoc(), UBCI->getOperand(),
+                                                     UBCI->getType());
+      // If we have an operand with owned ownership, we change our
+      // unchecked_ref_cast to explicitly pass the owned operand as an unowned
+      // value. This is because otherwise, we would consume the owned value
+      // creating breaking OSSA. In contrast, if we have a guaranteed value, we
+      // are going to be replacing an UnownedInstantaneousUse with an
+      // InstantaneousUse which is always safe for a guaranteed value.
+      if (newInst->getOwnershipKind() == OwnershipKind::Owned) {
+        newInst->setOwnershipKind(OwnershipKind::Unowned);
+      }
+      helper.perform(newInst);
+      return nullptr;
+    }
+  }
+
+  return nullptr;
 }
 
 SILInstruction *
 SILCombiner::visitThickToObjCMetatypeInst(ThickToObjCMetatypeInst *TTOCMI) {
-  if (TTOCMI->getFunction()->hasOwnership())
-    return nullptr;
-
   if (auto *OCTTMI = dyn_cast<ObjCToThickMetatypeInst>(TTOCMI->getOperand())) {
     TTOCMI->replaceAllUsesWith(OCTTMI->getOperand());
     return eraseInstFromFunction(*TTOCMI);
@@ -793,9 +822,6 @@ SILCombiner::visitThickToObjCMetatypeInst(ThickToObjCMetatypeInst *TTOCMI) {
 
 SILInstruction *
 SILCombiner::visitObjCToThickMetatypeInst(ObjCToThickMetatypeInst *OCTTMI) {
-  if (OCTTMI->getFunction()->hasOwnership())
-    return nullptr;
-
   if (auto *TTOCMI = dyn_cast<ThickToObjCMetatypeInst>(OCTTMI->getOperand())) {
     OCTTMI->replaceAllUsesWith(TTOCMI->getOperand());
     return eraseInstFromFunction(*OCTTMI);

lib/SILOptimizer/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -147,20 +147,24 @@ static EnumElementDecl *getInjectEnumCaseTo(SILValue Addr) {
 }
 
 SILInstruction *SILCombiner::visitSwitchEnumAddrInst(SwitchEnumAddrInst *SEAI) {
-  if (SEAI->getFunction()->hasOwnership())
-    return nullptr;
+  SILValue Addr = SEAI->getOperand();
 
   // Convert switch_enum_addr -> br
-  // if the only thing which writes to the address is an inject_enum_addr.
-  SILValue Addr = SEAI->getOperand();
-  if (EnumElementDecl *EnumCase = getInjectEnumCaseTo(Addr)) {
-    SILBasicBlock *Dest = SEAI->getCaseDestination(EnumCase);
-    // If the only instruction which writes to Addr is an inject_enum_addr we
-    // know that there cannot be an enum payload.
-    assert(Dest->getNumArguments() == 0 &&
-           "didn't expect a payload argument");
-    Builder.createBranch(SEAI->getLoc(), Dest);
-    return eraseInstFromFunction(*SEAI);
+  //
+  // If the only thing which writes to the address is an inject_enum_addr. We
+  // only perform these optimizations when we are not in OSSA since this
+  // eliminates an edge from the CFG and we want SILCombine in OSSA to never do
+  // that, so in the future we can invalidate less.
+  if (!SEAI->getFunction()->hasOwnership()) {
+    if (EnumElementDecl *EnumCase = getInjectEnumCaseTo(Addr)) {
+      SILBasicBlock *Dest = SEAI->getCaseDestination(EnumCase);
+      // If the only instruction which writes to Addr is an inject_enum_addr we
+      // know that there cannot be an enum payload.
+      assert(Dest->getNumArguments() == 0 &&
+             "didn't expect a payload argument");
+      Builder.createBranch(SEAI->getLoc(), Dest);
+      return eraseInstFromFunction(*SEAI);
+    }
   }
 
   SILType Ty = Addr->getType();
@@ -172,15 +176,44 @@ SILInstruction *SILCombiner::visitSwitchEnumAddrInst(SwitchEnumAddrInst *SEAI) {
   //     ->
   //   %value = load %ptr
   //   switch_enum %value
-  SmallVector<std::pair<EnumElementDecl*, SILBasicBlock*>, 8> Cases;
-  for (int i = 0, e = SEAI->getNumCases(); i < e; ++i)
+  //
+  // If we are using ownership, we perform a load_borrow right before the new
+  // switch_enum and end the borrow scope right afterwards.
+  Builder.setCurrentDebugScope(SEAI->getDebugScope());
+  SmallVector<std::pair<EnumElementDecl *, SILBasicBlock *>, 8> Cases;
+  for (int i : range(SEAI->getNumCases())) {
     Cases.push_back(SEAI->getCase(i));
+  }
 
-  Builder.setCurrentDebugScope(SEAI->getDebugScope());
   SILBasicBlock *Default = SEAI->hasDefault() ? SEAI->getDefaultBB() : nullptr;
-  LoadInst *EnumVal = Builder.createLoad(SEAI->getLoc(), Addr,
-                                         LoadOwnershipQualifier::Unqualified);
-  Builder.createSwitchEnum(SEAI->getLoc(), EnumVal, Default, Cases);
+  SILValue EnumVal = Builder.emitLoadBorrowOperation(SEAI->getLoc(), Addr);
+  auto *sei = Builder.createSwitchEnum(SEAI->getLoc(), EnumVal, Default, Cases);
+
+  if (Builder.hasOwnership()) {
+    for (int i : range(sei->getNumCases())) {
+      auto c = sei->getCase(i);
+      if (c.first->hasAssociatedValues()) {
+        auto eltType = Addr->getType().getEnumElementType(
+            c.first, Builder.getModule(), Builder.getTypeExpansionContext());
+        eltType = eltType.getObjectType();
+        if (eltType.isTrivial(Builder.getFunction())) {
+          c.second->createPhiArgument(eltType, OwnershipKind::None);
+        } else {
+          c.second->createPhiArgument(eltType, OwnershipKind::Guaranteed);
+        }
+      }
+      Builder.setInsertionPoint(c.second->front().getIterator());
+      Builder.emitEndBorrowOperation(SEAI->getLoc(), EnumVal);
+    }
+
+    if (auto defaultBlock = sei->getDefaultBBOrNull()) {
+      defaultBlock.get()->createPhiArgument(EnumVal->getType(),
+                                            OwnershipKind::Guaranteed);
+      Builder.setInsertionPoint(defaultBlock.get()->front().getIterator());
+      Builder.emitEndBorrowOperation(SEAI->getLoc(), EnumVal);
+    }
+  }
+
   return eraseInstFromFunction(*SEAI);
 }
 
@@ -1581,9 +1614,6 @@ SILCombiner::visitInjectEnumAddrInst(InjectEnumAddrInst *IEAI) {
 SILInstruction *
 SILCombiner::
 visitUnreachableInst(UnreachableInst *UI) {
-  if (UI->getFunction()->hasOwnership())
-    return nullptr;
-
   // Make sure that this unreachable instruction
   // is the last instruction in the basic block.
   if (UI->getParent()->getTerminator() == UI)
@@ -1857,9 +1887,6 @@ SILInstruction *SILCombiner::visitCondBranchInst(CondBranchInst *CBI) {
 }
 
 SILInstruction *SILCombiner::visitSelectEnumInst(SelectEnumInst *SEI) {
-  if (SEI->getFunction()->hasOwnership())
-    return nullptr;
-
   // Canonicalize a select_enum: if the default refers to exactly one case, then
   // replace the default with that case.
   if (SEI->hasDefault()) {

lib/SILOptimizer/Utils/OwnershipOptUtils.cpp

@@ -608,14 +608,19 @@ SILBasicBlock::iterator OwnershipRAUWUtility::handleUnowned() {
     // value, we need to ensure that the guaranteed value is live at all use
     // points of the unowned value. If so, just replace and continue.
     //
-    // TODO: Implement this.
+    // TODO: Implement this for more interesting cases.
+    if (isa<SILFunctionArgument>(newValue))
+      return replaceAllUsesAndErase(oldValue, newValue, callbacks);
 
     // Otherwise, we need to lifetime extend the borrow over all of the use
-    // points. To do so, we copy the value, borrow it, insert an unchecked
-    // ownership conversion to unowned at oldValue and then RAUW.
+    // points. To do so, we copy the value, borrow it, and insert an unchecked
+    // ownership conversion to unowned at all uses that are terminator uses.
+    //
+    // We need to insert the conversion since if we have a non-argument
+    // guaranteed value since its scope will end before the terminator so we
+    // need to convert the value to unowned early.
     //
-    // We need to insert the conversion to ensure that we do not violate
-    // ownership propagation rules of forwarding insts.
+    // TODO: Do we need a separate array here?
     SmallVector<Operand *, 8> oldValueUses(oldValue->getUses());
     for (auto *use : oldValueUses) {
       if (auto *ti = dyn_cast<TermInst>(use->getUser())) {
@@ -628,14 +633,12 @@ SILBasicBlock::iterator OwnershipRAUWUtility::handleUnowned() {
         }
       }
     }
+
     auto extender = getLifetimeExtender();
     SILValue borrow =
         extender.createPlusZeroBorrow(newValue, oldValue->getUses());
     SILBuilderWithScope builder(oldValue);
-    auto *newInst = builder.createUncheckedOwnershipConversion(
-        oldValue->getLoc(), borrow, OwnershipKind::Unowned);
-    callbacks.createdNewInst(newInst);
-    return replaceAllUsesAndErase(oldValue, newInst, callbacks);
+    return replaceAllUsesAndErase(oldValue, borrow, callbacks);
   }
   case OwnershipKind::Owned: {
     // If we have an unowned value that we want to replace with an owned value,
@@ -646,13 +649,15 @@ SILBasicBlock::iterator OwnershipRAUWUtility::handleUnowned() {
 
     // Otherwise, insert a copy of the owned value and lifetime extend that over
     // all uses of the value and then RAUW.
+    //
+    // NOTE: For terminator uses, we funnel the use through an
+    // unchecked_ownership_conversion to ensure that we can end the lifetime of
+    // our owned/guaranteed value before the terminator.
     SmallVector<Operand *, 8> oldValueUses(oldValue->getUses());
     for (auto *use : oldValueUses) {
       if (auto *ti = dyn_cast<TermInst>(use->getUser())) {
         if (ti->isFunctionExiting()) {
           SILBuilderWithScope builder(ti);
-          // We insert this to ensure that we can extend our owned value's
-          // lifetime to before the function end point.
           auto *newInst = builder.createUncheckedOwnershipConversion(
               ti->getLoc(), use->get(), OwnershipKind::Unowned);
           callbacks.createdNewInst(newInst);
@@ -663,10 +668,7 @@ SILBasicBlock::iterator OwnershipRAUWUtility::handleUnowned() {
     auto extender = getLifetimeExtender();
     SILValue copy = extender.createPlusZeroCopy(newValue, oldValue->getUses());
     SILBuilderWithScope builder(oldValue);
-    auto *newInst = builder.createUncheckedOwnershipConversion(
-        oldValue->getLoc(), copy, OwnershipKind::Unowned);
-    callbacks.createdNewInst(newInst);
-    auto result = replaceAllUsesAndErase(oldValue, newInst, callbacks);
+    auto result = replaceAllUsesAndErase(oldValue, copy, callbacks);
     return result;
   }
   }
@@ -1014,15 +1016,19 @@ OwnershipRAUWHelper::OwnershipRAUWHelper(OwnershipFixupContext &inputCtx,
   }
 }
 
-SILBasicBlock::iterator OwnershipRAUWHelper::perform() {
+SILBasicBlock::iterator
+OwnershipRAUWHelper::perform(SingleValueInstruction *maybeTransformedNewValue) {
   assert(isValid() && "OwnershipRAUWHelper invalid?!");
 
   // Make sure to always clear our context after we transform.
   SWIFT_DEFER { ctx->clear(); };
+  SILValue actualNewValue = newValue;
+  if (maybeTransformedNewValue)
+    actualNewValue = maybeTransformedNewValue;
 
   if (oldValue->getType().isAddress())
-    return replaceAddressUses(oldValue, newValue);
+    return replaceAddressUses(oldValue, actualNewValue);
 
-  OwnershipRAUWUtility utility{oldValue, newValue, *ctx};
+  OwnershipRAUWUtility utility{oldValue, actualNewValue, *ctx};
   return utility.perform();
 }

test/SILOptimizer/ossa_rauw_tests.sil

@@ -318,6 +318,40 @@ bb0(%0 : @guaranteed $Klass):
   return %3 : $(Klass, Klass)
 }
 
+// CHECK-LABEL: sil [ossa] @unowned_to_guaranteed_rauw_2a : $@convention(thin) (@guaranteed Builtin.NativeObject) -> (Klass, Klass) {
+// CHECK: bb0(
+// CHECK-NEXT: unchecked_ref_cast
+// CHECK-NEXT: tuple
+// CHECK-NEXT: return
+// CHECK: } // end sil function 'unowned_to_guaranteed_rauw_2a'
+sil [ossa] @unowned_to_guaranteed_rauw_2a : $@convention(thin) (@guaranteed Builtin.NativeObject) -> (Klass, Klass) {
+bb0(%0 : @guaranteed $Builtin.NativeObject):
+  %0a = unchecked_ref_cast %0 : $Builtin.NativeObject to $Klass
+  %1 = unchecked_bitwise_cast %0a : $Klass to $SubKlass
+  %2 = unchecked_bitwise_cast %1 : $SubKlass to $Klass
+  %3 = tuple(%2 : $Klass, %2 : $Klass)
+  return %3 : $(Klass, Klass)
+}
+
+// We need the unchecked_ownership_conversion since our base value is
+// guaranteed, not a function argument, and our user is a function exiting
+// terminator.
+//
+// CHECK-LABEL: sil [ossa] @unowned_to_guaranteed_rauw_2b : $@convention(thin) (@guaranteed Builtin.NativeObject) -> Klass {
+// CHECK: bb0(
+// CHECK-NEXT: unchecked_ref_cast
+// CHECK-NEXT: unchecked_ownership_conversion
+// CHECK-NEXT: return
+// CHECK: } // end sil function 'unowned_to_guaranteed_rauw_2b'
+sil [ossa] @unowned_to_guaranteed_rauw_2b : $@convention(thin) (@guaranteed Builtin.NativeObject) -> Klass {
+bb0(%0 : @guaranteed $Builtin.NativeObject):
+  %0a = unchecked_ref_cast %0 : $Builtin.NativeObject to $Klass
+  %1 = unchecked_bitwise_cast %0a : $Klass to $SubKlass
+  %2 = unchecked_bitwise_cast %1 : $SubKlass to $Klass
+  return %2 : $Klass
+}
+
+
 // CHECK-LABEL: sil [ossa] @unowned_to_guaranteed_rauw_2_loop : $@convention(thin) (@guaranteed Klass) -> @owned FakeOptional<(Klass, Klass)> {
 // CHECK: bb0([[ARG:%.*]] : @guaranteed $Klass):
 // CHECK-NOT: unchecked_bitwise_cast
@@ -376,7 +410,6 @@ bbExitBlock(%result : @owned $FakeOptional<(Klass, Klass)>):
 
 // CHECK-LABEL: sil [ossa] @unowned_to_guaranteed_rauw_3 : $@convention(thin) (@guaranteed Klass) -> Klass {
 // CHECK: bb0(
-// CHECK-NEXT: unchecked_ownership_conversion
 // CHECK-NEXT: return
 // CHECK: } // end sil function 'unowned_to_guaranteed_rauw_3'
 sil [ossa] @unowned_to_guaranteed_rauw_3 : $@convention(thin) (@guaranteed Klass) -> Klass {