Use OSSALifetimeCompletion in PredictableMemOpt

The current algorithm to complete lifetimes is incorrect in a few cases.
Use OSSALifetimeCompletion instead.

Fixes rdar://119204768

Author: meg-gupta

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -286,15 +286,15 @@ struct FunctionPassContext : MutatingContext {
   }
 
   func optimizeMemoryAccesses(in function: Function) -> Bool {
-    if BridgedPassContext.optimizeMemoryAccesses(function.bridged) {
+    if _bridged.optimizeMemoryAccesses(function.bridged) {
       notifyInstructionsChanged()
       return true
     }
     return false
   }
 
   func eliminateDeadAllocations(in function: Function) -> Bool {
-    if BridgedPassContext.eliminateDeadAllocations(function.bridged) {
+    if _bridged.eliminateDeadAllocations(function.bridged) {
       notifyInstructionsChanged()
       return true
     }

include/swift/SILOptimizer/OptimizerBridging.h

@@ -216,8 +216,8 @@ struct BridgedPassContext {
                                                             bool isPrivate) const;
   void inlineFunction(BridgedInstruction apply, bool mandatoryInline) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedValue getSILUndef(BridgedType type) const;
-  BRIDGED_INLINE static bool optimizeMemoryAccesses(BridgedFunction f);
-  BRIDGED_INLINE static bool eliminateDeadAllocations(BridgedFunction f);
+  BRIDGED_INLINE bool optimizeMemoryAccesses(BridgedFunction f) const;
+  BRIDGED_INLINE bool eliminateDeadAllocations(BridgedFunction f) const;
 
   // IRGen
 

include/swift/SILOptimizer/OptimizerBridgingImpl.h

@@ -218,11 +218,12 @@ BridgedValue BridgedPassContext::getSILUndef(BridgedType type) const {
   return {swift::SILUndef::get(type.unbridged(), *invocation->getFunction())};
 }
 
-bool BridgedPassContext::optimizeMemoryAccesses(BridgedFunction f) {
-  return swift::optimizeMemoryAccesses(f.getFunction());
+bool BridgedPassContext::optimizeMemoryAccesses(BridgedFunction f) const {
+  return swift::optimizeMemoryAccesses(f.getFunction(), this->getDomTree().di);
 }
-bool BridgedPassContext::eliminateDeadAllocations(BridgedFunction f) {
-  return swift::eliminateDeadAllocations(f.getFunction());
+bool BridgedPassContext::eliminateDeadAllocations(BridgedFunction f) const {
+  return swift::eliminateDeadAllocations(f.getFunction(),
+                                         this->getDomTree().di);
 }
 
 BridgedBasicBlockSet BridgedPassContext::allocBasicBlockSet() const {

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -582,12 +582,12 @@ IntegerLiteralInst *optimizeBuiltinCanBeObjCClass(BuiltinInst *bi,
 /// Performs "predictable" memory access optimizations.
 ///
 /// See the PredictableMemoryAccessOptimizations pass.
-bool optimizeMemoryAccesses(SILFunction *fn);
+bool optimizeMemoryAccesses(SILFunction *fn, DominanceInfo *domInfo);
 
 /// Performs "predictable" dead allocation optimizations.
 ///
 /// See the PredictableDeadAllocationElimination pass.
-bool eliminateDeadAllocations(SILFunction *fn);
+bool eliminateDeadAllocations(SILFunction *fn, DominanceInfo *domInfo);
 
 SILVTable *specializeVTableForType(SILType type, SILModule &mod, SILTransform *transform);
 

lib/SILOptimizer/Mandatory/PredictableMemOpt.cpp

@@ -20,6 +20,7 @@
 #include "swift/SIL/BasicBlockBits.h"
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/LinearLifetimeChecker.h"
+#include "swift/SIL/OSSALifetimeCompletion.h"
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/SILOptimizer/PassManager/Passes.h"
@@ -1945,19 +1946,22 @@ class AllocOptimize {
 
   InstructionDeleter &deleter;
 
+  DominanceInfo *domInfo;
+
   /// A structure that we use to compute our available values.
   AvailableValueDataflowContext DataflowContext;
 
 public:
   AllocOptimize(AllocationInst *memory, SmallVectorImpl<PMOMemoryUse> &uses,
                 SmallVectorImpl<SILInstruction *> &releases,
-                DeadEndBlocks &deadEndBlocks, InstructionDeleter &deleter)
+                DeadEndBlocks &deadEndBlocks, InstructionDeleter &deleter,
+                DominanceInfo *domInfo)
       : Module(memory->getModule()), TheMemory(memory),
         MemoryType(getMemoryType(memory)),
         NumMemorySubElements(getNumSubElements(
             MemoryType, Module, TypeExpansionContext(*memory->getFunction()))),
         Uses(uses), Releases(releases), deadEndBlocks(deadEndBlocks),
-        deleter(deleter),
+        deleter(deleter), domInfo(domInfo),
         DataflowContext(TheMemory, NumMemorySubElements, uses, deleter) {}
 
   bool optimizeMemoryAccesses();
@@ -2655,145 +2659,20 @@ bool AllocOptimize::tryToRemoveDeadAllocation() {
   // post-dominating consuming use sets. This can happen if we have an enum that
   // is known dynamically none along a path. This is dynamically correct, but
   // can not be represented in OSSA so we insert these destroys along said path.
-  SmallVector<SILBasicBlock *, 32> consumingUseBlocks;
+  OSSALifetimeCompletion completion(TheMemory->getFunction(), domInfo);
+
   while (!valuesNeedingLifetimeCompletion.empty()) {
     auto optV = valuesNeedingLifetimeCompletion.pop_back_val();
     if (!optV)
       continue;
     SILValue v = *optV;
-    if (v->getOwnershipKind() != OwnershipKind::Owned)
-      continue;
-
-    // FIXME: This doesn't handle situations where non-consuming uses are in
-    // blocks in which the consuming use is already deleted. Replace the
-    // following code with an general OSSA utility for owned lifetime
-    // fixup. ValueLifetimeAnalysis does this easily, but we also need to handle
-    // reborrows by adding copies. An OwnershipOptUtils utility will handle
-    // that soon.
-
-    // First see if our value doesn't have any uses. In such a case, just
-    // insert a destroy_value at the next instruction and return.
-    if (v->use_empty()) {
-      auto *next = v->getNextInstruction();
-      auto loc = RegularLocation::getAutoGeneratedLocation();
-      SILBuilderWithScope localBuilder(next);
-      localBuilder.createDestroyValue(loc, v);
-      continue;
-    }
-
-    // Otherwise, we first see if we have any consuming uses at all. If we do,
-    // then we know that any such consuming uses since we have an owned value
-    // /must/ be strongly control equivalent to our value and unreachable from
-    // each other, so we can just use findJointPostDominatingSet to complete
-    // the set.
-    consumingUseBlocks.clear();
-    for (auto *use : v->getConsumingUses())
-      consumingUseBlocks.push_back(use->getParentBlock());
-
-    if (!consumingUseBlocks.empty()) {
-      findJointPostDominatingSet(
-          v->getParentBlock(), consumingUseBlocks, [](SILBasicBlock *) {},
-          [&](SILBasicBlock *result) {
-            auto loc = RegularLocation::getAutoGeneratedLocation();
-            SILBuilderWithScope builder(result);
-            builder.createDestroyValue(loc, v);
-          });
-      continue;
-    }
-
-    // If we do not have at least one consuming use, we need to do something
-    // different. This situation can occur given a non-trivial enum typed
-    // stack allocation that:
-    //
-    // 1. Had a destroy_addr eliminated along a path where we dynamically know
-    //    that the stack allocation is storing a trivial case.
-    //
-    // 2. Had some other paths where due to dead end blocks, no destroy_addr
-    //    is needed.
-    //
-    // To fix this, we just treat all uses as consuming blocks and insert
-    // destroys using the joint post dominance set computer and insert
-    // destroys at the end of all input blocks in the post dom set and at the
-    // beginning of any leaking blocks.
-    {
-      // TODO: Can we just pass this in to findJointPostDominatingSet instead
-      // of recomputing it there? Maybe an overload that lets us do this?
-      BasicBlockSet foundUseBlocks(v->getFunction());
-      for (auto *use : v->getUses()) {
-        auto *block = use->getParentBlock();
-        if (!foundUseBlocks.insert(block))
-          continue;
-        consumingUseBlocks.push_back(block);
-      }
-    }
-    findJointPostDominatingSet(
-        v->getParentBlock(), consumingUseBlocks,
-        [&](SILBasicBlock *foundInputBlock) {
-          // This is a block that is reachable from another use. We are not
-          // interested in these.
-        },
-        [&](SILBasicBlock *leakingBlock) {
-          auto loc = RegularLocation::getAutoGeneratedLocation();
-          SILBuilderWithScope builder(leakingBlock);
-          builder.createDestroyValue(loc, v);
-        },
-        [&](SILBasicBlock *inputBlockInPostDomSet) {
-          auto *termInst = inputBlockInPostDomSet->getTerminator();
-          switch (termInst->getTermKind()) {
-          case TermKind::UnreachableInst:
-            // We do not care about input blocks that end in unreachables. We
-            // are going to leak down them so do not insert a destroy_value
-            // there.
-            return;
-
-          // NOTE: Given that our input value is owned, our branch can only
-          // accept the use as a non-consuming use if the branch is forwarding
-          // unowned ownership. Luckily for use, we checked early if we had
-          // any such uses and bailed, so we know the branch can not use our
-          // value. This is just avoiding a corner case that we don't need to
-          // handle.
-          case TermKind::BranchInst:
-            LLVM_FALLTHROUGH;
-          // NOTE: We put cond_br here since in OSSA, cond_br can never have
-          // a non-trivial value operand, meaning we can insert before.
-          case TermKind::CondBranchInst:
-            LLVM_FALLTHROUGH;
-          case TermKind::ReturnInst:
-          case TermKind::ThrowInst:
-          case TermKind::ThrowAddrInst:
-          case TermKind::UnwindInst:
-          case TermKind::YieldInst: {
-            // These terminators can never be non-consuming uses of an owned
-            // value since we would be leaking the owned value no matter what
-            // we do. Given that, we can assume that what ever the
-            // non-consuming use actually was, must be before this
-            // instruction. So insert the destroy_value at the end of the
-            // block, before the terminator.
-            auto loc = RegularLocation::getAutoGeneratedLocation();
-            SILBuilderWithScope localBuilder(termInst);
-            localBuilder.createDestroyValue(loc, v);
-            return;
-          }
-          case TermKind::TryApplyInst:
-          case TermKind::SwitchValueInst:
-          case TermKind::SwitchEnumInst:
-          case TermKind::SwitchEnumAddrInst:
-          case TermKind::DynamicMethodBranchInst:
-          case TermKind::AwaitAsyncContinuationInst:
-          case TermKind::CheckedCastBranchInst:
-          case TermKind::CheckedCastAddrBranchInst: {
-            // Otherwise, we insert the destroy_addr /after/ the
-            // terminator. All of these are guaranteed to have each successor
-            // to have the block as its only predecessor block.
-            SILBuilderWithScope::insertAfter(termInst, [&](auto &b) {
-              auto loc = RegularLocation::getAutoGeneratedLocation();
-              b.createDestroyValue(loc, v);
-            });
-            return;
-          }
-          }
-          llvm_unreachable("Case that did not return in its body?!");
-        });
+    // Lexical enums can have incomplete lifetimes in non payload paths that
+    // don't end in unreachable. Force their lifetime to end immediately after
+    // the last use instead.
+    bool forceBoundaryCompletion = v->getType().isOrHasEnum();
+    LLVM_DEBUG(llvm::dbgs() << "Completing lifetime of: ");
+    LLVM_DEBUG(v->dump());
+    completion.completeOSSALifetime(v, forceBoundaryCompletion);
   }
 
   return true;
@@ -2863,7 +2742,7 @@ static AllocationInst *getOptimizableAllocation(SILInstruction *i) {
   return alloc;
 }
 
-bool swift::optimizeMemoryAccesses(SILFunction *fn) {
+bool swift::optimizeMemoryAccesses(SILFunction *fn, DominanceInfo *domInfo) {
   bool changed = false;
   DeadEndBlocks deadEndBlocks(fn);
 
@@ -2892,8 +2771,8 @@ bool swift::optimizeMemoryAccesses(SILFunction *fn) {
         // runs. It creates and deletes instructions other than alloc.
         continue;
       }
-      AllocOptimize allocOptimize(alloc, uses, destroys, deadEndBlocks,
-                                  deleter);
+      AllocOptimize allocOptimize(alloc, uses, destroys, deadEndBlocks, deleter,
+                                  domInfo);
       changed |= allocOptimize.optimizeMemoryAccesses();
 
       // Move onto the next instruction. We know this is safe since we do not
@@ -2904,7 +2783,7 @@ bool swift::optimizeMemoryAccesses(SILFunction *fn) {
   return changed;
 }
 
-bool swift::eliminateDeadAllocations(SILFunction *fn) {
+bool swift::eliminateDeadAllocations(SILFunction *fn, DominanceInfo *domInfo) {
   if (!fn->hasOwnership())
     return false;
 
@@ -2935,8 +2814,8 @@ bool swift::eliminateDeadAllocations(SILFunction *fn) {
       if (!collectPMOElementUsesFrom(memInfo, uses, destroys)) {
         continue;
       }
-      AllocOptimize allocOptimize(alloc, uses, destroys, deadEndBlocks,
-                                  deleter);
+      AllocOptimize allocOptimize(alloc, uses, destroys, deadEndBlocks, deleter,
+                                  domInfo);
       if (allocOptimize.tryToRemoveDeadAllocation()) {
         deleter.cleanupDeadInstructions();
         ++NumAllocRemoved;
@@ -2958,19 +2837,24 @@ class PredictableMemoryAccessOptimizations : public SILFunctionTransform {
   /// either indicates that this pass missing some opportunities the first time,
   /// or has a pass order dependency on other early passes.
   void run() override {
+    auto *func = getFunction();
+    LLVM_DEBUG(llvm::dbgs() << "Looking at: " << func->getName() << "\n");
+    auto *da = getAnalysis<DominanceAnalysis>();
     // TODO: Can we invalidate here just instructions?
-    if (optimizeMemoryAccesses(getFunction()))
+    if (optimizeMemoryAccesses(func, da->get(func)))
       invalidateAnalysis(SILAnalysis::InvalidationKind::FunctionBody);
   }
 };
 
 class PredictableDeadAllocationElimination : public SILFunctionTransform {
   void run() override {
+    auto *func = getFunction();
+    LLVM_DEBUG(llvm::dbgs() << "Looking at: " << func->getName() << "\n");
+    auto *da = getAnalysis<DominanceAnalysis>();
     // If we are already canonical or do not have ownership, just bail.
-    if (getFunction()->wasDeserializedCanonical() ||
-        !getFunction()->hasOwnership())
+    if (func->wasDeserializedCanonical() || !func->hasOwnership())
       return;
-    if (eliminateDeadAllocations(getFunction()))
+    if (eliminateDeadAllocations(func, da->get(func)))
       invalidateAnalysis(SILAnalysis::InvalidationKind::FunctionBody);
   }
 };

test/SILOptimizer/predictable_deadalloc_elim_ownership.sil

@@ -697,3 +697,55 @@ bb0(%arg : @owned $Builtin.NativeObject):
   destroy_value %1 : $Builtin.NativeObject
   unreachable
 }
+
+enum ErrorEnum : Error {
+  case errorCase(label: [any Error])
+  case other
+}
+
+sil [ossa] @test_complete_lifetime : $@convention(thin) (@owned ErrorEnum) -> () {
+
+bb0(%0 : @owned $ErrorEnum):
+  %1 = alloc_stack $any Error
+  %2 = alloc_existential_box $any Error, $ErrorEnum
+  %3 = project_existential_box $ErrorEnum in %2 : $any Error
+  store %2 to [init] %1 : $*any Error
+  store %0 to [init] %3 : $*ErrorEnum
+  %6 = load [take] %1 : $*any Error
+  dealloc_stack %1 : $*any Error
+  %8 = alloc_stack $any Error
+  %9 = copy_value %6 : $any Error
+  store %9 to [init] %8 : $*any Error
+  %11 = alloc_stack $ErrorEnum
+  checked_cast_addr_br copy_on_success any Error in %8 : $*any Error to ErrorEnum in %11 : $*ErrorEnum, bb2, bb3
+
+bb1:
+  %13 = tuple ()
+  return %13 : $()
+
+bb2:
+  %15 = load [take] %11 : $*ErrorEnum
+  destroy_value %15 : $ErrorEnum
+  dealloc_stack %11 : $*ErrorEnum
+  destroy_addr %8 : $*any Error
+  dealloc_stack %8 : $*any Error
+  destroy_value %6 : $any Error
+  br bb1
+
+bb3:
+  dealloc_stack %11 : $*ErrorEnum
+  destroy_addr %8 : $*any Error
+  dealloc_stack %8 : $*any Error
+  %25 = copy_value %6 : $any Error
+  cond_br undef, bb4, bb5
+
+bb4:
+  br bb6
+
+bb5:
+  br bb6
+
+bb6:
+  unreachable
+}
+