Optimizer: fix an ownership violation when duplicating loops

If a guaranteed value is used in a dead-end exit block and the enclosing value is _not_ destroyed in this block, we end up missing the enclosing value as phi-argument after duplicating the loop.
TODO: once we have complete lifetimes we can remove this check again.

rdar://159125605

Author: eeckstein

include/swift/SILOptimizer/Utils/LoopUtils.h

@@ -28,6 +28,7 @@ class SILInstruction;
 class SILLoop;
 class DominanceInfo;
 class SILLoopInfo;
+class DeadEndBlocks;
 
 /// Canonicalize the loop for rotation and downstream passes.
 ///
@@ -40,7 +41,7 @@ bool canonicalizeAllLoops(DominanceInfo *DT, SILLoopInfo *LI);
 
 /// Check whether it is safe to duplicate this instruction when duplicating
 /// this loop by unrolling or versioning.
-bool canDuplicateLoopInstruction(SILLoop *L, SILInstruction *Inst);
+bool canDuplicateLoopInstruction(SILLoop *L, SILInstruction *Inst, DeadEndBlocks *deb);
 
 /// A visitor that visits loops in a function in a bottom up order. It only
 /// performs the visit.

lib/SILOptimizer/LoopTransforms/ArrayPropertyOpt.cpp

@@ -61,6 +61,7 @@
 #include "swift/SIL/Projection.h"
 #include "swift/SIL/SILCloner.h"
 #include "swift/SILOptimizer/Analysis/ArraySemantic.h"
+#include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
 #include "swift/SILOptimizer/Analysis/LoopAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/BasicBlockOptUtils.h"
@@ -87,6 +88,7 @@ class ArrayPropertiesAnalysis {
   SILLoop *Loop;
   SILBasicBlock *Preheader;
   DominanceInfo *DomTree;
+  DeadEndBlocks *deb;
 
   SinkAddressProjections sinkProj;
 
@@ -104,9 +106,9 @@ class ArrayPropertiesAnalysis {
   bool reachingBlocksComputed = false;
 
 public:
-  ArrayPropertiesAnalysis(SILLoop *L, DominanceAnalysis *DA)
+  ArrayPropertiesAnalysis(SILLoop *L, DominanceAnalysis *DA, DeadEndBlocks *deb)
       : Fun(L->getHeader()->getParent()), Loop(L), Preheader(nullptr),
-        DomTree(DA->get(Fun)), ReachingBlocks(Fun) {}
+        DomTree(DA->get(Fun)), deb(deb), ReachingBlocks(Fun) {}
 
   /// Check if it is profitable to specialize a loop when you see an apply
   /// instruction. We consider it is not profitable to specialize the loop when:
@@ -178,7 +180,7 @@ class ArrayPropertiesAnalysis {
       for (auto &Inst : *BB) {
         // Can't clone alloc_stack instructions whose dealloc_stack is outside
         // the loop.
-        if (!canDuplicateLoopInstruction(Loop, &Inst))
+        if (!canDuplicateLoopInstruction(Loop, &Inst, deb))
           return false;
 
         if (!sinkProj.analyzeAddressProjections(&Inst)) {
@@ -796,6 +798,7 @@ class SwiftArrayPropertyOptPass : public SILFunctionTransform {
       return;
 
     DominanceAnalysis *DA = PM->getAnalysis<DominanceAnalysis>();
+    DeadEndBlocks *deb = PM->getAnalysis<DeadEndBlocksAnalysis>()->get(Fn);
     SILLoopAnalysis *LA = PM->getAnalysis<SILLoopAnalysis>();
     SILLoopInfo *LI = LA->get(Fn);
 
@@ -808,7 +811,7 @@ class SwiftArrayPropertyOptPass : public SILFunctionTransform {
     // possible loop-nest.
     SmallVector<SILBasicBlock *, 16> HoistableLoopNests;
     std::function<void(SILLoop *)> processChildren = [&](SILLoop *L) {
-      ArrayPropertiesAnalysis Analysis(L, DA);
+      ArrayPropertiesAnalysis Analysis(L, DA, deb);
       if (Analysis.run()) {
         // Hoist in the current loop nest.
         HasChanged = true;

lib/SILOptimizer/LoopTransforms/LoopUnroll.cpp

@@ -19,6 +19,7 @@
 #include "swift/SIL/SILCloner.h"
 #include "swift/SILOptimizer/Analysis/LoopAnalysis.h"
 #include "swift/SILOptimizer/Analysis/IsSelfRecursiveAnalysis.h"
+#include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
 #include "swift/SILOptimizer/Utils/BasicBlockOptUtils.h"
@@ -211,7 +212,8 @@ static std::optional<uint64_t> getMaxLoopTripCount(SILLoop *Loop,
 /// heuristic that looks at the trip count and the cost of the instructions in
 /// the loop to determine whether we should unroll this loop.
 static bool canAndShouldUnrollLoop(SILLoop *Loop, uint64_t TripCount,
-                                   IsSelfRecursiveAnalysis *SRA) {
+                                   IsSelfRecursiveAnalysis *SRA,
+                                   DeadEndBlocks *deb) {
   assert(Loop->getSubLoops().empty() && "Expect innermost loops");
   if (TripCount > 32)
     return false;
@@ -227,7 +229,7 @@ static bool canAndShouldUnrollLoop(SILLoop *Loop, uint64_t TripCount,
     (Loop->getBlocks())[0]->getParent()->getModule().getOptions().UnrollThreshold;
   for (auto *BB : Loop->getBlocks()) {
     for (auto &Inst : *BB) {
-      if (!canDuplicateLoopInstruction(Loop, &Inst))
+      if (!canDuplicateLoopInstruction(Loop, &Inst, deb))
         return false;
       if (instructionInlineCost(Inst) != InlineCost::Free)
         ++Cost;
@@ -388,7 +390,7 @@ updateSSA(SILFunction *Fn, SILLoop *Loop,
 
 /// Try to fully unroll the loop if we can determine the trip count and the trip
 /// count is below a threshold.
-static bool tryToUnrollLoop(SILLoop *Loop, IsSelfRecursiveAnalysis *SRA) {
+static bool tryToUnrollLoop(SILLoop *Loop, IsSelfRecursiveAnalysis *SRA, DeadEndBlocks *deb) {
   assert(Loop->getSubLoops().empty() && "Expecting innermost loops");
 
   LLVM_DEBUG(llvm::dbgs() << "Trying to unroll loop : \n" << *Loop);
@@ -409,7 +411,7 @@ static bool tryToUnrollLoop(SILLoop *Loop, IsSelfRecursiveAnalysis *SRA) {
     return false;
   }
 
-  if (!canAndShouldUnrollLoop(Loop, MaxTripCount.value(), SRA)) {
+  if (!canAndShouldUnrollLoop(Loop, MaxTripCount.value(), SRA, deb)) {
     LLVM_DEBUG(llvm::dbgs() << "Not unrolling, exceeds cost threshold\n");
     return false;
   }
@@ -490,6 +492,7 @@ class LoopUnrolling : public SILFunctionTransform {
     auto *Fun = getFunction();
     SILLoopInfo *LoopInfo = PM->getAnalysis<SILLoopAnalysis>()->get(Fun);
     IsSelfRecursiveAnalysis *SRA = PM->getAnalysis<IsSelfRecursiveAnalysis>();
+    DeadEndBlocks *deb = PM->getAnalysis<DeadEndBlocksAnalysis>()->get(Fun);
 
     LLVM_DEBUG(llvm::dbgs() << "Loop Unroll running on function : "
                             << Fun->getName() << "\n");
@@ -517,7 +520,7 @@ class LoopUnrolling : public SILFunctionTransform {
 
     // Try to unroll innermost loops.
     for (auto *Loop : InnermostLoops)
-      Changed |= tryToUnrollLoop(Loop, SRA);
+      Changed |= tryToUnrollLoop(Loop, SRA, deb);
 
     if (Changed) {
       invalidateAnalysis(SILAnalysis::InvalidationKind::FunctionBody);

lib/SILOptimizer/Utils/LoopUtils.cpp

@@ -218,9 +218,21 @@ bool swift::canonicalizeAllLoops(DominanceInfo *DT, SILLoopInfo *LI) {
   return MadeChange;
 }
 
-bool swift::canDuplicateLoopInstruction(SILLoop *L, SILInstruction *I) {
+bool swift::canDuplicateLoopInstruction(SILLoop *L, SILInstruction *I, DeadEndBlocks *deb) {
   SinkAddressProjections sinkProj;
   for (auto res : I->getResults()) {
+    // If a guaranteed value is used in a dead-end exit block and the enclosing value
+    // is _not_ destroyed in this block, we end up missing the enclosing value as
+    // phi-argument after duplicating the loop.
+    // TODO: once we have complete lifetimes we can remove this check
+    if (res->getOwnershipKind() == OwnershipKind::Guaranteed) {
+      for (Operand *use : res->getUses()) {
+        SILBasicBlock *useBlock = use->getUser()->getParent();
+        if (!L->contains(useBlock) && deb->isDeadEnd(useBlock))
+          return false;
+      }
+    }
+
     if (!res->getType().isAddress()) {
       continue;
     }

test/SILOptimizer/array_property_opt_ossa_guaranteed.sil

@@ -76,6 +76,38 @@ bb3:
   return %4 : $MyBool
 }
 
+// Check that we don't crash with an ownership violation
+// CHECK-LABEL: sil [ossa] @borrow_use_in_dead_end : 
+// CHECK-LABEL: } // end sil function 'borrow_use_in_dead_end'
+sil [ossa] @borrow_use_in_dead_end : $@convention(thin) (@guaranteed MyArray<MyClass>) -> () {
+bb0(%0 : @guaranteed $MyArray<MyClass>):
+  br bb1
+
+bb1:
+  %2 = function_ref @arrayPropertyIsNative : $@convention(method) (@guaranteed MyArray<MyClass>) -> Bool
+  %5 = apply %2(%0) : $@convention(method) (@guaranteed MyArray<MyClass>) -> Bool
+  %10 = alloc_ref $MyClass
+  %11 = begin_borrow %10
+  cond_br undef, bb2, bb3
+
+bb2:
+  fix_lifetime %11
+  unreachable
+
+bb3:
+  end_borrow %11
+  destroy_value %10
+  cond_br undef, bb4, bb5
+
+bb4:
+  br bb1
+
+bb5:
+  %r = tuple ()
+  return %r
+}
+
+
 // CHECK-LABEL: sil [ossa] @load_and_copy_within_loop : 
 // CHECK: bb1:
 // CHECK: [[FUNC1:%.*]] = function_ref @arrayPropertyIsNative