Migrate ConditionForwarding to OSSA (swiftlang#39734)

Author: meg-gupta

lib/SILOptimizer/Transforms/ConditionForwarding.cpp

@@ -11,14 +11,15 @@
 //===----------------------------------------------------------------------===//
 
 #define DEBUG_TYPE "condbranch-forwarding"
-#include "swift/SILOptimizer/PassManager/Passes.h"
-#include "swift/SILOptimizer/PassManager/Transforms.h"
-#include "swift/SIL/SILInstruction.h"
-#include "swift/SIL/SILArgument.h"
+#include "swift/SIL/BasicBlockBits.h"
 #include "swift/SIL/DebugUtils.h"
+#include "swift/SIL/OwnershipUtils.h"
+#include "swift/SIL/SILArgument.h"
 #include "swift/SIL/SILBuilder.h"
+#include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILUndef.h"
-#include "swift/SIL/BasicBlockBits.h"
+#include "swift/SILOptimizer/PassManager/Passes.h"
+#include "swift/SILOptimizer/PassManager/Transforms.h"
 
 using namespace swift;
 
@@ -101,15 +102,12 @@ class ConditionForwarding : public SILFunctionTransform {
 
   /// The entry point to the transformation.
   void run() override {
-    LLVM_DEBUG(llvm::dbgs() << "** StackPromotion **\n");
-
     bool Changed = false;
 
     SILFunction *F = getFunction();
 
-    // FIXME: Add ownership support.
-    if (F->hasOwnership())
-      return;
+    LLVM_DEBUG(llvm::dbgs()
+               << "** ConditionForwarding on " << F->getName() << " **\n");
 
     for (SILBasicBlock &BB : *F) {
       if (auto *SEI = dyn_cast<SwitchEnumInst>(BB.getTerminator())) {
@@ -143,6 +141,11 @@ static bool hasNoRelevantSideEffects(SILBasicBlock *BB) {
         return false;
       continue;
     }
+    if (isa<BeginBorrowInst>(&I) || isa<EndBorrowInst>(&I)) {
+      continue;
+    }
+    LLVM_DEBUG(llvm::dbgs() << "Bailing out, found inst with side-effects ");
+    LLVM_DEBUG(I.dump());
     return false;
   }
   return true;
@@ -226,6 +229,15 @@ bool ConditionForwarding::tryOptimize(SwitchEnumInst *SEI) {
   if (CommonBranchBlock->getSuccessors().size() != PredBlocks.size())
     return false;
 
+  if (getFunction()->hasOwnership()) {
+    // TODO: Currently disabled because this case may need lifetime extension
+    // Disabling this conservatively for now.
+    assert(Condition->getNumOperands() == 1);
+    BorrowedValue conditionOp(Condition->getOperand(0));
+    if (conditionOp && conditionOp.isLocalScope()) {
+      return false;
+    }
+  }
   // Now do the transformation!
   // First thing to do is to replace all uses of the Enum (= the merging block
   // argument), as this argument gets deleted.
@@ -247,8 +259,8 @@ bool ConditionForwarding::tryOptimize(SwitchEnumInst *SEI) {
       SILArgument *NewArg = nullptr;
       if (NeedEnumArg.insert(UseBlock)) {
         // The first Enum use in this UseBlock.
-        NewArg =
-            UseBlock->createPhiArgument(Arg->getType(), OwnershipKind::Owned);
+        NewArg = UseBlock->createPhiArgument(Arg->getType(),
+                                             Arg->getOwnershipKind());
       } else {
         // We already inserted the Enum argument for this UseBlock.
         assert(UseBlock->getNumArguments() >= 1);
@@ -281,10 +293,30 @@ bool ConditionForwarding::tryOptimize(SwitchEnumInst *SEI) {
     if (HasEnumArg) {
       // The successor block has a new argument (which we created above) where
       // we have to pass the Enum.
+      assert(!getFunction()->hasOwnership() ||
+             EI->getType().isTrivial(*getFunction()));
       BranchArgs.push_back(EI);
     }
     B.createBranch(BI->getLoc(), SEDest, BranchArgs);
     BI->eraseFromParent();
+    if (EI->use_empty()) {
+      assert(!HasEnumArg);
+      EI->eraseFromParent();
+    } else {
+      // If an @owned EI has uses remaining, ownership fixup is needed.
+      // 1. Create a copy_value of EI's operand and
+      // use it in the branch to avoid a double-consume.
+      // 2. Create a destroy_value of EI, to avoid a leak.
+      if (getFunction()->hasOwnership() && EI->hasOperand() &&
+          EI->getOwnershipKind() == OwnershipKind::Owned) {
+        auto *term = EI->getParent()->getTerminator();
+        assert(!HasEnumArg);
+        auto *copy = SILBuilderWithScope(EI).createCopyValue(EI->getLoc(),
+                                                             EI->getOperand());
+        term->getOperandRef(0).set(copy);
+        SILBuilderWithScope(term).createDestroyValue(EI->getLoc(), EI);
+      }
+    }
   }
 
   // Final step: replace the switch_enum by the condition.

test/SILOptimizer/conditionforwarding.sil

@@ -356,3 +356,35 @@ bb7:
   %r = tuple ()
   return %r : $()
 }
+
+// CHECK-LABEL: sil @forwarding_second_enum_use : 
+// CHECK-NOT: switch_enum
+// CHECK-LABEL:} // end sil function 'forwarding_second_enum_use'
+sil @forwarding_second_enum_use : $@convention(thin) (Builtin.Int1) -> () {
+bb0(%0 : $Builtin.Int1):
+  cond_br %0, bb1, bb2
+
+bb1:
+  %1 = enum $E, #E.A!enumelt
+  br bb3(%1 : $E)
+
+bb2:
+  %2 = enum $E, #E.B!enumelt
+  br bb3(%2 : $E)
+
+bb3(%14 : $E):
+  switch_enum %14 : $E, case #E.A!enumelt: bb4, case #E.B!enumelt: bb5
+
+bb4:
+  %15 = function_ref @use_enum : $@convention(thin) (E) -> ()
+  %16 = apply %15(%14) : $@convention(thin) (E) -> ()
+  br bb6
+
+bb5:
+  br bb6
+
+bb6:
+  %r = tuple ()
+  return %r : $()
+}
+