Enable jump threading

Author: meg-gupta

lib/SILOptimizer/Transforms/SimplifyCFG.cpp

@@ -906,24 +906,20 @@ static bool hasInjectedEnumAtEndOfBlock(SILBasicBlock *block, SILValue enumAddr)
 /// tryJumpThreading - Check to see if it looks profitable to duplicate the
 /// destination of an unconditional jump into the bottom of this block.
 bool SimplifyCFG::tryJumpThreading(BranchInst *BI) {
-  if (!EnableOSSASimplifyCFG && Fn.hasOwnership())
-    return false;
-
   auto *DestBB = BI->getDestBB();
   auto *SrcBB = BI->getParent();
   TermInst *destTerminator = DestBB->getTerminator();
-  if (!EnableOSSARewriteTerminator && Fn.hasOwnership()) {
-    if (llvm::any_of(DestBB->getArguments(), [this](SILValue op) {
-          return !op->getType().isTrivial(Fn);
-        })) {
-      return false;
-   }
-  }
+
   // If the destination block ends with a return, we don't want to duplicate it.
   // We want to maintain the canonical form of a single return where possible.
   if (destTerminator->isFunctionExiting())
     return false;
 
+  // There is no benefit duplicating such a destination.
+  if (DestBB->getSinglePredecessorBlock() != nullptr) {
+    return false;
+  }
+
   // Jump threading only makes sense if there is an argument on the branch
   // (which is reacted on in the DestBB), or if this goes through a memory
   // location (switch_enum_addr is the only address-instruction which we
@@ -942,6 +938,7 @@ bool SimplifyCFG::tryJumpThreading(BranchInst *BI) {
   for (unsigned i : indices(BI->getArgs())) {
     SILValue Arg = BI->getArg(i);
 
+    // TODO: Verify if we need to jump thread to remove releases in OSSA.
     // If the value being substituted on is release there is a chance we could
     // remove the release after jump threading.
     if (!Arg->getType().isTrivial(*SrcBB->getParent()) &&

lib/SILOptimizer/UtilityPasses/UnitTestRunner.cpp

@@ -361,6 +361,19 @@ struct SimplifyCFGSimplifyTermWithIdenticalDestBlocks : UnitTest {
   }
 };
 
+struct SimplifyCFGTryJumpThreading : UnitTest {
+  SimplifyCFGTryJumpThreading(UnitTestRunner *pass)
+      : UnitTest(pass) {}
+  void invoke(Arguments &arguments) override {
+    auto *passToRun = cast<SILFunctionTransform>(createSimplifyCFG());
+    passToRun->injectPassManager(getPass()->getPassManager());
+    passToRun->injectFunction(getFunction());
+    SimplifyCFG(*getFunction(), *passToRun, /*VerifyAll=*/false,
+                /*EnableJumpThread=*/false)
+        .tryJumpThreading(cast<BranchInst>(arguments.takeInstruction()));
+  }
+};
+
 // Arguments:
 // - string: list of characters, each of which specifies subsequent arguments
 //           - A: (block) argument
@@ -500,6 +513,9 @@ void UnitTestRunner::withTest(StringRef name, Doit doit) {
     ADD_UNIT_TEST_SUBCLASS(
         "simplify-cfg-simplify-term-with-identical-dest-blocks",
         SimplifyCFGSimplifyTermWithIdenticalDestBlocks)
+    ADD_UNIT_TEST_SUBCLASS(
+        "simplify-cfg-try-jump-threading",
+        SimplifyCFGTryJumpThreading)
 
     ADD_UNIT_TEST_SUBCLASS("test-specification-parsing", TestSpecificationTest)
     ADD_UNIT_TEST_SUBCLASS("visit-adjacent-reborrows-of-phi", VisitAdjacentReborrowsOfPhiTest)

test/SILOptimizer/simplify_cfg_ossa_jump_threading.sil

@@ -0,0 +1,60 @@
+
+class Klass {
+}
+
+enum FakeOptional<T> {
+  case some(T)
+  case none
+}
+
+sil [ossa] @test_simplify_switch_enum_jump_threading1 : $@convention(thin) (@owned Klass) -> () {
+bb0(%0 : @owned $Klass):
+  test_specification "simplify-cfg-try-jump-threading @instruction[1]"
+  %1 = enum $FakeOptional<Klass>, #FakeOptional.some!enumelt, %0 : $Klass
+  br bb1(%1 : $FakeOptional<Klass>)
+
+bb1(%3 : @owned $FakeOptional<Klass>):
+  switch_enum %3 : $FakeOptional<Klass>, case #FakeOptional.some!enumelt: bb3, case #FakeOptional.none!enumelt: bb2
+
+bb2:
+  br bb4
+
+bb3(%6 : @owned $Klass):
+  destroy_value %6 : $Klass
+  br bb4
+
+bb4:
+  %t = tuple ()
+  return %t : $()
+}
+
+
+sil [ossa] @test_simplify_switch_enum_jump_threading2 : $@convention(thin) (@owned Klass) -> () {
+bb0(%0 : @owned $Klass):
+  cond_br undef, bb1, bb2
+
+bb1:
+  test_specification "simplify-cfg-try-jump-threading @instruction[2]"
+  %1 = enum $FakeOptional<Klass>, #FakeOptional.some!enumelt, %0 : $Klass
+  br bb3(%1 : $FakeOptional<Klass>)
+
+bb2:  
+  test_specification "simplify-cfg-try-jump-threading @instruction[5]"
+  destroy_value %0 : $Klass
+  %2 = enum $FakeOptional<Klass>, #FakeOptional.none!enumelt
+  br bb3(%2 : $FakeOptional<Klass>)
+
+bb3(%3 : @owned $FakeOptional<Klass>):
+  switch_enum %3 : $FakeOptional<Klass>, case #FakeOptional.some!enumelt: bb5, case #FakeOptional.none!enumelt: bb4
+
+bb4:
+  br bb6
+
+bb5(%6 : @owned $Klass):
+  destroy_value %6 : $Klass
+  br bb6
+
+bb6:
+  %t = tuple ()
+  return %t : $()
+}