6.1: [OSSACanonicalizeOwned] Fix liveness passed to completion.

lib/SIL/Utils/PrettyStackTrace.cpp

@@ -96,7 +96,7 @@ void PrettyStackTraceSILFunction::printFunctionInfo(llvm::raw_ostream &out) cons
   if (SILPrintOnError)
     func->print(out);
   if (SILPrintModuleOnError)
-    func->getModule().print(out);
+    func->getModule().print(out, func->getModule().getSwiftModule());
 }
 
 void PrettyStackTraceSILNode::print(llvm::raw_ostream &out) const {

lib/SIL/Utils/PrunedLiveness.cpp

@@ -860,7 +860,7 @@ void PrunedLiveRange<LivenessWithDefs>::computeBoundary(
   // Visit each post-dominating block as the starting point for a
   // backward CFG traversal.
   for (auto *block : postDomBlocks) {
-    blockWorklist.push(block);
+    blockWorklist.pushIfNotVisited(block);
   }
   while (auto *block = blockWorklist.pop()) {
     // Process each block that has not been visited and is not LiveOut.

lib/SILOptimizer/SILCombiner/SILCombine.cpp

@@ -120,6 +120,8 @@ void SILCombiner::addReachableCodeToWorklist(SILBasicBlock *BB) {
 //                               Implementation
 //===----------------------------------------------------------------------===//
 
+namespace swift {
+
 // Define a CanonicalizeInstruction subclass for use in SILCombine.
 class SILCombineCanonicalize final : CanonicalizeInstruction {
   SmallSILInstructionWorklist<256> &Worklist;
@@ -162,6 +164,8 @@ class SILCombineCanonicalize final : CanonicalizeInstruction {
   }
 };
 
+} // end namespace swift
+
 SILCombiner::SILCombiner(SILFunctionTransform *trans,
                          bool removeCondFails, bool enableCopyPropagation) :
   parentTransform(trans),
@@ -414,99 +418,136 @@ bool SILCombiner::doOneIteration(SILFunction &F, unsigned Iteration) {
     if (!parentTransform->continueWithNextSubpassRun(I))
       return false;
 
-    // Check to see if we can DCE the instruction.
-    if (isInstructionTriviallyDead(I)) {
-      LLVM_DEBUG(llvm::dbgs() << "SC: DCE: " << *I << '\n');
-      eraseInstFromFunction(*I);
-      ++NumDeadInst;
-      MadeChange = true;
-      continue;
-    }
+    processInstruction(I, scCanonicalize, MadeChange);
+  }
+
+  Worklist.resetChecked();
+  return MadeChange;
+}
+
+void SILCombiner::processInstruction(SILInstruction *I,
+                                     SILCombineCanonicalize &scCanonicalize,
+                                     bool &MadeChange) {
+  // Check to see if we can DCE the instruction.
+  if (isInstructionTriviallyDead(I)) {
+    LLVM_DEBUG(llvm::dbgs() << "SC: DCE: " << *I << '\n');
+    eraseInstFromFunction(*I);
+    ++NumDeadInst;
+    MadeChange = true;
+    return;
+  }
 #ifndef NDEBUG
-    std::string OrigIStr;
+  std::string OrigIStr;
 #endif
-    LLVM_DEBUG(llvm::raw_string_ostream SS(OrigIStr); I->print(SS);
-               OrigIStr = SS.str(););
-    LLVM_DEBUG(llvm::dbgs() << "SC: Visiting: " << OrigIStr << '\n');
+  LLVM_DEBUG(llvm::raw_string_ostream SS(OrigIStr); I->print(SS);
+             OrigIStr = SS.str(););
+  LLVM_DEBUG(llvm::dbgs() << "SC: Visiting: " << OrigIStr << '\n');
 
-    // Canonicalize the instruction.
-    if (scCanonicalize.tryCanonicalize(I)) {
-      MadeChange = true;
-      continue;
-    }
+  // Canonicalize the instruction.
+  if (scCanonicalize.tryCanonicalize(I)) {
+    MadeChange = true;
+    return;
+  }
 
-    // If we have reached this point, all attempts to do simple simplifications
-    // have failed. First if we have an owned forwarding value, we try to
-    // sink. Otherwise, we perform the actual SILCombine operation.
-    if (EnableSinkingOwnedForwardingInstToUses) {
-      // If we have an ownership forwarding single value inst that forwards
-      // through its first argument and it is trivially duplicatable, see if it
-      // only has consuming uses. If so, we can duplicate the instruction into
-      // the consuming use blocks and destroy any destroy_value uses of it that
-      // we see. This makes it easier for SILCombine to fold instructions with
-      // owned parameters since chains of these values will be in the same
-      // block.
-      if (auto *svi = dyn_cast<SingleValueInstruction>(I)) {
-        if (auto fwdOp = ForwardingOperation(svi)) {
-          if (fwdOp.getSingleForwardingOperand() &&
-              SILValue(svi)->getOwnershipKind() == OwnershipKind::Owned) {
-            // Try to sink the value. If we sank the value and deleted it,
-            // continue. If we didn't optimize or sank but we are still able to
-            // optimize further, we fall through to SILCombine below.
-            if (trySinkOwnedForwardingInst(svi)) {
-              continue;
-            }
+  // If we have reached this point, all attempts to do simple simplifications
+  // have failed. First if we have an owned forwarding value, we try to
+  // sink. Otherwise, we perform the actual SILCombine operation.
+  if (EnableSinkingOwnedForwardingInstToUses) {
+    // If we have an ownership forwarding single value inst that forwards
+    // through its first argument and it is trivially duplicatable, see if it
+    // only has consuming uses. If so, we can duplicate the instruction into
+    // the consuming use blocks and destroy any destroy_value uses of it that
+    // we see. This makes it easier for SILCombine to fold instructions with
+    // owned parameters since chains of these values will be in the same
+    // block.
+    if (auto *svi = dyn_cast<SingleValueInstruction>(I)) {
+      if (auto fwdOp = ForwardingOperation(svi)) {
+        if (fwdOp.getSingleForwardingOperand() &&
+            SILValue(svi)->getOwnershipKind() == OwnershipKind::Owned) {
+          // Try to sink the value. If we sank the value and deleted it,
+          // return. If we didn't optimize or sank but we are still able to
+          // optimize further, we fall through to SILCombine below.
+          if (trySinkOwnedForwardingInst(svi)) {
+            return;
           }
         }
       }
     }
+  }
 
-    // Then begin... SILCombine.
-    Builder.setInsertionPoint(I);
+  // Then begin... SILCombine.
+  Builder.setInsertionPoint(I);
 
-    SILInstruction *currentInst = I;
-    if (SILInstruction *Result = visit(I)) {
-      ++NumCombined;
-      // Should we replace the old instruction with a new one?
-      Worklist.replaceInstructionWithInstruction(I, Result
+  SILInstruction *currentInst = I;
+  if (SILInstruction *Result = visit(I)) {
+    ++NumCombined;
+    // Should we replace the old instruction with a new one?
+    Worklist.replaceInstructionWithInstruction(I, Result
 #ifndef NDEBUG
-                                                 ,
-                                                 OrigIStr
+                                               ,
+                                               OrigIStr
 #endif
-      );
-      currentInst = Result;
-      MadeChange = true;
-    }
+    );
+    currentInst = Result;
+    MadeChange = true;
+  }
 
-    // Eliminate copies created that this SILCombine iteration may have
-    // introduced during OSSA-RAUW.
-    canonicalizeOSSALifetimes(currentInst->isDeleted() ? nullptr : currentInst);
-
-    // Builder's tracking list has been accumulating instructions created by the
-    // during this SILCombine iteration. To finish this iteration, go through
-    // the tracking list and add its contents to the worklist and then clear
-    // said list in preparation for the next iteration.
-    for (SILInstruction *I : *Builder.getTrackingList()) {
-      if (!I->isDeleted()) {
-        LLVM_DEBUG(llvm::dbgs()
-                   << "SC: add " << *I << " from tracking list to worklist\n");
-        Worklist.add(I);
-      }
+  // Eliminate copies created that this SILCombine iteration may have
+  // introduced during OSSA-RAUW.
+  canonicalizeOSSALifetimes(currentInst->isDeleted() ? nullptr : currentInst);
+
+  // Builder's tracking list has been accumulating instructions created by the
+  // during this SILCombine iteration. To finish this iteration, go through
+  // the tracking list and add its contents to the worklist and then clear
+  // said list in preparation for the next iteration.
+  for (SILInstruction *I : *Builder.getTrackingList()) {
+    if (!I->isDeleted()) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << "SC: add " << *I << " from tracking list to worklist\n");
+      Worklist.add(I);
     }
-    Builder.getTrackingList()->clear();
   }
-
-  Worklist.resetChecked();
-  return MadeChange;
+  Builder.getTrackingList()->clear();
 }
 
+namespace swift::test {
+struct SILCombinerProcessInstruction {
+  void operator()(SILCombiner &combiner, SILInstruction *inst,
+                  SILCombineCanonicalize &canonicalizer, bool &madeChange) {
+    combiner.processInstruction(inst, canonicalizer, madeChange);
+  }
+};
+// Arguments:
+// - instruction: the instruction to be processed
+// - bool: remove cond_fails
+// - bool: enable lifetime canonicalization
+// Dumps:
+// - the function after the processing is attempted
+static FunctionTest SILCombineProcessInstruction(
+    "sil_combine_process_instruction",
+    [](auto &function, auto &arguments, auto &test) {
+      auto inst = arguments.takeInstruction();
+      auto removeCondFails = arguments.takeBool();
+      auto enableCopyPropagation = arguments.takeBool();
+      SILCombiner combiner(test.getPass(), removeCondFails,
+                           enableCopyPropagation);
+      SILCombineCanonicalize canonicalizer(combiner.Worklist,
+                                           *test.getDeadEndBlocks());
+      bool madeChange = false;
+      SILCombinerProcessInstruction()(combiner, inst, canonicalizer,
+                                      madeChange);
+      function.dump();
+    });
+} // end namespace swift::test
+
 namespace swift::test {
 // Arguments:
-// - instruction: the instruction to be canonicalized
+// - instruction: the instruction to be visited
 // Dumps:
-// - the function after the canonicalization is attempted
-static FunctionTest SILCombineCanonicalizeInstruction(
-    "sil_combine_instruction", [](auto &function, auto &arguments, auto &test) {
+// - the function after the visitation is attempted
+static FunctionTest SILCombineVisitInstruction(
+    "sil_combine_visit_instruction",
+    [](auto &function, auto &arguments, auto &test) {
       SILCombiner combiner(test.getPass(), false, false);
       auto inst = arguments.takeInstruction();
       combiner.Builder.setInsertionPoint(inst);

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -47,6 +47,10 @@
 namespace swift {
 
 class AliasAnalysis;
+class SILCombineCanonicalize;
+namespace test {
+struct SILCombinerProcessInstruction;
+}
 
 /// This is a class which maintains the state of the combiner and simplifies
 /// many operations such as removing/adding instructions and syncing them with
@@ -417,6 +421,11 @@ class SILCombiner :
   /// Perform one SILCombine iteration.
   bool doOneIteration(SILFunction &F, unsigned Iteration);
 
+  void processInstruction(SILInstruction *instruction,
+                          SILCombineCanonicalize &scCanonicalize,
+                          bool &MadeChange);
+  friend test::SILCombinerProcessInstruction;
+
   /// Add reachable code to the worklist. Meant to be used when starting to
   /// process a new function.
   void addReachableCodeToWorklist(SILBasicBlock *BB);

lib/SILOptimizer/Utils/CanonicalizeOSSALifetime.cpp

@@ -285,6 +285,43 @@ void CanonicalizeOSSALifetime::extendLivenessToDeadEnds() {
     completeLiveness.updateForUse(destroy, /*lifetimeEnding*/ true);
   }
 
+  // Demote consuming uses within complete liveness to non-consuming uses.
+  //
+  // OSSALifetimeCompletion considers the lifetime of a single value.  Such
+  // lifetimes never continue beyond consumes.
+  std::optional<llvm::SmallPtrSet<SILInstruction *, 8>> lastUsers;
+  auto isConsumeOnBoundary = [&](SILInstruction *instruction) -> bool {
+    if (!lastUsers) {
+      // Avoid computing lastUsers if possible.
+      auto *function = getCurrentDef()->getFunction();
+      auto *deadEnds = deadEndBlocksAnalysis->get(function);
+      llvm::SmallVector<SILBasicBlock *, 8> completeConsumingBlocks(
+          consumingBlocks.getArrayRef());
+      for (auto &block : *function) {
+        if (!deadEnds->isDeadEnd(&block))
+          continue;
+        completeConsumingBlocks.push_back(&block);
+      }
+      PrunedLivenessBoundary boundary;
+      liveness->computeBoundary(boundary, completeConsumingBlocks);
+
+      lastUsers.emplace();
+      for (auto *lastUser : boundary.lastUsers) {
+        lastUsers->insert(lastUser);
+      }
+    }
+    return lastUsers->contains(instruction);
+  };
+  for (auto pair : liveness->getAllUsers()) {
+    if (!pair.second.isEnding())
+      continue;
+    auto *instruction = pair.first;
+    if (isConsumeOnBoundary(instruction))
+      continue;
+    // Demote instruction's lifetime-ending-ness to non-lifetime-ending.
+    completeLiveness.updateForUse(pair.first, /*lifetimeEnding=*/false);
+  }
+
   OSSALifetimeCompletion::visitAvailabilityBoundary(
       getCurrentDef(), completeLiveness, [&](auto *unreachable, auto end) {
         if (end == OSSALifetimeCompletion::LifetimeEnd::Boundary) {

test/SILOptimizer/canonicalize_ossa_lifetime_unit.sil

@@ -162,9 +162,9 @@ exit(%phi : @owned $C, %typhi : $S):
 sil @empty : $@convention(thin) () -> () {
 [global: ]
 bb0:
-  %0 = tuple ()                                   
-  return %0 : $()                                 
-} 
+  %0 = tuple ()
+  return %0 : $()
+}
 
 // Even though the apply of %empty is not a deinit barrier, verify that the
 // destroy is not hoisted, because MoS is move-only.
@@ -568,16 +568,16 @@ entry(%c1 : @owned $C):
 // CHECK:       bb0([[C1:%[^,]+]] : @owned $C):
 // CHECK:         [[TAKE_C:%[^,]+]] = function_ref @takeC
 // CHECK:         [[BARRIER:%[^,]+]] = function_ref @barrier
-// CHECK:         cond_br undef, [[LEFT]], [[RIGHT]]                         
-// CHECK:       [[LEFT]]:                                              
+// CHECK:         cond_br undef, [[LEFT]], [[RIGHT]]
+// CHECK:       [[LEFT]]:
 // CHECK:         [[M:%[^,]+]] = move_value [[C1]]
 // CHECK:         apply [[TAKE_C]]([[M]])
-// CHECK:         br [[EXIT]]                                          
-// CHECK:       [[RIGHT]]:                                              
+// CHECK:         br [[EXIT]]
+// CHECK:       [[RIGHT]]:
 // CHECK:         apply [[BARRIER]]()
 // CHECK:         destroy_value [[C1]]
-// CHECK:         br [[EXIT]]                                          
-// CHECK:       [[EXIT]]:                                              
+// CHECK:         br [[EXIT]]
+// CHECK:       [[EXIT]]:
 // CHECK:         apply [[BARRIER]]()
 // CHECK-LABEL: } // end sil function 'lexical_end_at_end_2'
 // CHECK-LABEL: end running test {{.*}} on lexical_end_at_end_2: canonicalize_ossa_lifetime
@@ -857,3 +857,31 @@ exit:
   %retval = tuple ()
   return %retval : $()
 }
+
+
+// CHECK-LABEL: begin running test {{.*}} on consume_copy_before_use_in_dead_end
+// CHECK-LABEL: sil [ossa] @consume_copy_before_use_in_dead_end : {{.*}} {
+// CHECK-NOT:     destroy_value [dead_end]
+// CHECK-LABEL: } // end sil function 'consume_copy_before_use_in_dead_end'
+// CHECK-LABEL: end running test {{.*}} on consume_copy_before_use_in_dead_end
+sil [ossa] @consume_copy_before_use_in_dead_end : $@convention(thin) (@owned C) -> () {
+entry(%c : @owned $C):
+  cond_br undef, exit, die
+
+exit:
+  destroy_value %c : $C
+  %retval = tuple ()
+  return %retval : $()
+
+die:
+  %move = move_value [lexical] %c : $C
+  %copy = copy_value %move : $C
+  specify_test "canonicalize_ossa_lifetime true false true %move"
+  apply undef(%move) : $@convention(thin) (@owned C) -> ()
+  %addr = alloc_stack $C
+  %token = store_borrow %copy to %addr : $*C
+  apply undef() : $@convention(thin) () -> ()
+  %reload = load_borrow %token : $*C
+  apply undef(%reload) : $@convention(thin) (@guaranteed C) -> ()
+  unreachable
+}