Merge pull request #39097 from meg-gupta/rlefix

Fix an edge case in OSSA RLE for loops

Author: meg-gupta

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -675,15 +675,10 @@ SILBasicBlock::iterator replaceSingleUse(Operand *use, SILValue newValue,
 SILValue
 makeCopiedValueAvailable(SILValue value, SILBasicBlock *inBlock);
 
-/// Given a newly created @owned value \p value without any uses, this utility
+/// Given an existing @owned value \p value, this utility
 /// inserts control equivalent copy and destroy at leaking blocks to adjust
 /// ownership and make \p value available for use at \p inBlock.
-///
-/// inBlock must be the only point at which \p value will be consumed. If this
-/// consuming point is within a loop, this will create and return a copy of \p
-/// value inside \p inBlock.
-SILValue
-makeNewValueAvailable(SILValue value, SILBasicBlock *inBlock);
+SILValue makeValueAvailable(SILValue value, SILBasicBlock *inBlock);
 
 /// Given an ssa value \p value, create destroy_values at leaking blocks
 ///

lib/SILOptimizer/Transforms/RedundantLoadElimination.cpp

@@ -1363,7 +1363,7 @@ SILValue RLEContext::computePredecessorLocationValue(SILBasicBlock *BB,
     }
     endLifetimeAtLeakingBlocks(phi, userBBs);
   }
-  return makeNewValueAvailable(Val, BB);
+  return makeValueAvailable(Val, BB);
 }
 
 bool RLEContext::collectLocationValues(SILBasicBlock *BB, LSLocation &L,

lib/SILOptimizer/Transforms/SimplifyCFG.cpp

@@ -2672,7 +2672,7 @@ bool SimplifyCFG::simplifyTryApplyBlock(TryApplyInst *TAI) {
     if (requiresOSSACleanup(newCallee)) {
       newCallee = SILBuilderWithScope(newCallee->getNextInstruction())
         .createCopyValue(calleeLoc, newCallee);
-      newCallee = makeNewValueAvailable(newCallee, TAI->getParent());
+      newCallee = makeValueAvailable(newCallee, TAI->getParent());
     }
 
     ApplyOptions Options = TAI->getApplyOptions();

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -2012,25 +2012,30 @@ SILValue swift::makeCopiedValueAvailable(SILValue value, SILBasicBlock *inBlock)
   auto *copy = builder.createCopyValue(
       RegularLocation::getAutoGeneratedLocation(), value);
 
-  return makeNewValueAvailable(copy, inBlock);
+  return makeValueAvailable(copy, inBlock);
 }
 
-SILValue swift::makeNewValueAvailable(SILValue value, SILBasicBlock *inBlock) {
+SILValue swift::makeValueAvailable(SILValue value, SILBasicBlock *inBlock) {
   if (!value->getFunction()->hasOwnership())
     return value;
 
   if (value.getOwnershipKind() == OwnershipKind::None)
     return value;
 
-  assert(value->getUses().empty() &&
-         value.getOwnershipKind() == OwnershipKind::Owned);
+  assert(value.getOwnershipKind() == OwnershipKind::Owned);
+
+  SmallVector<SILBasicBlock *, 4> userBBs;
+  for (auto use : value->getUses()) {
+    userBBs.push_back(use->getParentBlock());
+  }
+  userBBs.push_back(inBlock);
 
   // Use \p jointPostDomComputer to:
   // 1. Create a control equivalent copy at \p inBlock if needed
   // 2. Insert destroy_value at leaking blocks
   SILValue controlEqCopy;
   findJointPostDominatingSet(
-      value->getParentBlock(), inBlock,
+      value->getParentBlock(), userBBs,
       [&](SILBasicBlock *loopBlock) {
         assert(loopBlock == inBlock);
         auto front = loopBlock->begin();

lib/SILOptimizer/Utils/LoadStoreOptUtils.cpp

@@ -119,7 +119,7 @@ void LSValue::reduceInner(LSLocation &Base, SILModule *M,
           Builder, RegularLocation::getAutoGeneratedLocation(),
           Base.getType(M, context).getObjectType(), Vals);
 
-  auto AvailVal = makeNewValueAvailable(AI.get(), InsertPt->getParent());
+  auto AvailVal = makeValueAvailable(AI.get(), InsertPt->getParent());
 
   // This is the Value for the current base.
   ProjectionPath P(Base.getType(M, context));

test/SILOptimizer/redundant_load_elim_nontrivial_ossa.sil

@@ -78,11 +78,13 @@ final class NewHalfOpenRangeGenerator : NewRangeGenerator1 {
 sil_global @total_klass : $Klass
 sil_global @total_nontrivialstruct : $NonTrivialStruct
 
+sil @use : $@convention(thin) (Builtin.Int32) -> ()
 sil @use_Klass : $@convention(thin) (@owned Klass) -> ()
 sil @use_nontrivialstruct : $@convention(thin) (@owned NonTrivialStruct) -> ()
 sil @use_a : $@convention(thin) (@owned A) -> ()
 sil @use_twofield : $@convention(thin) (@owned TwoField) -> ()
 sil @init_twofield : $@convention(thin) (@thin TwoField.Type) -> @owned TwoField
+sil @get_nontrivialstruct : $@convention(thin) () -> @owned NonTrivialStruct
 
 // We have a bug in the old projection code which this test case exposes.
 // Make sure its handled properly in the new projection.
@@ -1051,7 +1053,6 @@ bb8:
   unreachable
 }
 
-
 // CHECK-LABEL: @infinite_loop_and_unreachable :
 // CHECK:   [[V:%[0-9]+]] = load [copy]
 // CHECK:   [[C1:%[0-9]+]] = copy_value [[V]]
@@ -1078,3 +1079,171 @@ bb3:
   unreachable
 }
 
+// CHECK-LABEL: @test_available_value1 :
+// CHECK:   load [trivial]
+// CHECK-NOT:   load [trivial]
+// CHECK: } // end sil function 'test_available_value1'
+sil [ossa] @test_available_value1 : $@convention(thin) (@in Builtin.Int32) -> () {
+
+bb0(%0 : $*Builtin.Int32):
+  %1 = load [trivial] %0 : $*Builtin.Int32
+
+  %2 = function_ref @use : $@convention(thin) (Builtin.Int32) -> ()
+  %3 = apply %2(%1) : $@convention(thin) (Builtin.Int32) -> ()
+  cond_br undef, bb2, bb3
+
+bb1:
+  cond_br undef, bb9, bb10
+
+bb2:
+  br bb4
+
+bb3:
+  br bb8
+
+bb4:
+  br bb5
+
+bb5:
+
+  %9 = function_ref @use : $@convention(thin) (Builtin.Int32) -> ()
+  %10 = apply %9(%1) : $@convention(thin) (Builtin.Int32) -> ()
+  cond_br undef, bb6, bb7
+
+bb6:
+  br bb4
+
+bb7:
+  br bb8
+
+bb8:
+  br bb1
+
+bb9:
+  br bb11
+
+bb10:
+  br bb11
+
+bb11:
+  %17 = tuple ()
+  return %17 : $()
+}
+
+sil [ossa] @test_available_value2 : $@convention(thin) (Builtin.Int32, Builtin.Int32) -> () {
+bb0(%0 : $Builtin.Int32, %1 : $Builtin.Int32):
+  %2 = alloc_stack $Builtin.Int32
+  cond_br undef, bb1, bb2
+
+bb1:
+  br bb3(%0 : $Builtin.Int32)
+
+bb2:
+  br bb3(%1 : $Builtin.Int32)
+
+
+bb3(%6 : $Builtin.Int32):
+  store %6 to [trivial] %2 : $*Builtin.Int32
+  cond_br undef, bb5, bb6
+
+bb4:
+  cond_br undef, bb12, bb13
+
+bb5:
+  br bb7
+
+bb6:
+  br bb11
+
+bb7:
+  br bb8
+
+bb8:
+
+  %13 = function_ref @use : $@convention(thin) (Builtin.Int32) -> ()
+  %14 = apply %13(%6) : $@convention(thin) (Builtin.Int32) -> ()
+  cond_br undef, bb9, bb10
+
+bb9:
+  br bb7
+
+bb10:
+  br bb11
+
+bb11:
+  br bb4
+
+bb12:
+  br bb14
+
+bb13:
+  br bb14
+
+bb14:
+  dealloc_stack %2 : $*Builtin.Int32
+  %22 = tuple ()
+  return %22 : $()
+}
+
+// CHECK-LABEL: @test_available_value3 :
+// CHECK-NOT:   load
+// CHECK: } // end sil function 'test_available_value3'
+sil [ossa] @test_available_value3 : $@convention(method) (@owned NonTrivialStruct) -> () {
+bb0(%0 : @owned $NonTrivialStruct):
+  %1 = alloc_stack $NonTrivialStruct
+  store %0 to [init] %1 : $*NonTrivialStruct
+  br bb1
+
+bb1:
+  cond_br undef, bb2, bb3
+
+bb2:
+  %5 = function_ref @get_nontrivialstruct : $@convention(thin) () -> @owned NonTrivialStruct
+  %6 = apply %5() : $@convention(thin) () -> @owned NonTrivialStruct
+  store %6 to [assign] %1 : $*NonTrivialStruct
+  cond_br undef, bb4, bb13
+
+bb3:
+  unreachable
+
+bb4:
+  cond_br undef, bb5, bb12
+
+bb5:
+  br bb6
+
+bb6:
+  br bb7
+
+bb7:
+  cond_br undef, bb8, bb11
+
+bb8:
+  cond_br undef, bb9, bb10
+
+bb9:
+  %15 = load [take] %1 : $*NonTrivialStruct
+  destroy_value %15 : $NonTrivialStruct
+  dealloc_stack %1 : $*NonTrivialStruct
+  br bb14
+
+bb10:
+  br bb6
+
+bb11:
+  br bb1
+
+bb12:
+  br bb1
+
+bb13:
+  %22 = load [take] %1 : $*NonTrivialStruct
+  destroy_value %22 : $NonTrivialStruct
+  dealloc_stack %1 : $*NonTrivialStruct
+  br bb14
+
+bb14:
+  %26 = tuple ()
+  return %26 : $()
+}
+