Enable mem2reg and temprvalueopt on lexical enums

include/swift/SIL/OSSALifetimeCompletion.h

@@ -53,12 +53,19 @@ class OSSALifetimeCompletion {
   /// Insert a lifetime-ending instruction on every path to complete the OSSA
   /// lifetime of \p value. Lifetime completion is only relevant for owned
   /// values or borrow introducers.
-  ///
+  /// For lexical values lifetime is completed at unreachable instructions.
+  /// For non-lexical values lifetime is completed at the lifetime boundary.
+  /// When \p forceBoundaryCompletion is true, the client is able to guarantee
+  /// that lifetime completion of lexical values at the lifetime boundary is
+  /// sufficient.
+  /// Currently \p forceBoundaryCompletion is used by mem2reg and temprvalueopt
+  /// to complete lexical enum values on trivial paths.
   /// Returns true if any new instructions were created to complete the
   /// lifetime.
   ///
   /// TODO: We also need to complete scoped addresses (e.g. store_borrow)!
-  LifetimeCompletion completeOSSALifetime(SILValue value) {
+  LifetimeCompletion
+  completeOSSALifetime(SILValue value, bool forceBoundaryCompletion = false) {
     if (value->getOwnershipKind() == OwnershipKind::None)
       return LifetimeCompletion::NoLifetime;
 
@@ -73,13 +80,13 @@ class OSSALifetimeCompletion {
     if (!completedValues.insert(value))
       return LifetimeCompletion::AlreadyComplete;
 
-    return analyzeAndUpdateLifetime(value)
-      ? LifetimeCompletion::WasCompleted
-      : LifetimeCompletion::AlreadyComplete;
+    return analyzeAndUpdateLifetime(value, forceBoundaryCompletion)
+               ? LifetimeCompletion::WasCompleted
+               : LifetimeCompletion::AlreadyComplete;
   }
 
 protected:
-  bool analyzeAndUpdateLifetime(SILValue value);
+  bool analyzeAndUpdateLifetime(SILValue value, bool forceBoundaryCompletion);
 };
 
 //===----------------------------------------------------------------------===//

lib/SIL/Utils/OSSALifetimeCompletion.cpp

@@ -142,7 +142,8 @@ static bool endLifetimeAtUnreachableBlocks(SILValue value,
 /// This is only meant to cleanup lifetimes that lead to dead-end blocks. After
 /// recursively completing all nested scopes, it then simply ends the lifetime
 /// at the Unreachable instruction.
-bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(SILValue value) {
+bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(
+    SILValue value, bool forceBoundaryCompletion) {
   // Called for inner borrows, inner adjacent reborrows, inner reborrows, and
   // scoped addresses.
   auto handleInnerScope = [this](SILValue innerBorrowedValue) {
@@ -152,7 +153,7 @@ bool OSSALifetimeCompletion::analyzeAndUpdateLifetime(SILValue value) {
   liveness.compute(domInfo, handleInnerScope);
 
   bool changed = false;
-  if (value->isLexical()) {
+  if (value->isLexical() && !forceBoundaryCompletion) {
     changed |= endLifetimeAtUnreachableBlocks(value, liveness.getLiveness());
   } else {
     changed |= endLifetimeAtBoundary(value, liveness.getLiveness());

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -195,6 +195,13 @@ class LiveValues {
     return LiveValues::forOwned({stored, move});
   }
 
+  static LiveValues forValues(SILValue stored, SILValue lexical) {
+    if (stored->getOwnershipKind() == OwnershipKind::Guaranteed) {
+      return LiveValues::forGuaranteed({stored, lexical});
+    }
+    return LiveValues::forOwned({stored, lexical});
+  }
+
   static LiveValues toReplace(AllocStackInst *asi, SILValue replacement) {
     if (replacement->getOwnershipKind() == OwnershipKind::Guaranteed) {
       return LiveValues::forGuaranteed(Guaranteed::toReplace(asi, replacement));
@@ -715,6 +722,27 @@ static bool canEndLexicalLifetime(LiveValues values) {
   return values.canEndLexicalLifetime();
 }
 
+static SILValue getLexicalValueForStore(SILInstruction *inst,
+                                        AllocStackInst *asi) {
+  assert(isa<StoreInst>(inst) || isa<StoreBorrowInst>(inst));
+
+  SILValue stored = inst->getOperand(CopyLikeInstruction::Src);
+  LLVM_DEBUG(llvm::dbgs() << "*** Found Store def " << stored);
+
+  if (!lexicalLifetimeEnsured(asi)) {
+    return SILValue();
+  }
+  if (isa<StoreBorrowInst>(inst)) {
+    if (isGuaranteedLexicalValue(stored)) {
+      return SILValue();
+    }
+    auto borrow = cast<BeginBorrowInst>(inst->getNextInstruction());
+    return borrow;
+  }
+  auto move = cast<MoveValueInst>(inst->getNextInstruction());
+  return move;
+}
+
 /// Begin a lexical borrow scope for the value stored into the provided
 /// StoreInst after that instruction.
 ///
@@ -1233,27 +1261,9 @@ StackAllocationPromoter::getLiveOutValues(BlockSetVector &phiBlocks,
     BlockToInstMap::iterator it = initializationPoints.find(domBlock);
     if (it != initializationPoints.end()) {
       auto *inst = it->second;
-      assert(isa<StoreInst>(inst) || isa<StoreBorrowInst>(inst));
-
-      SILValue stored = inst->getOperand(CopyLikeInstruction::Src);
-      LLVM_DEBUG(llvm::dbgs() << "*** Found Store def " << stored);
-
-      if (!lexicalLifetimeEnsured(asi)) {
-        auto values = LiveValues::forOwned(stored, {});
-        return values;
-      }
-      if (isa<StoreBorrowInst>(inst)) {
-        if (isGuaranteedLexicalValue(stored)) {
-          auto values = LiveValues::forGuaranteed(stored, {});
-          return values;
-        }
-        auto borrow = cast<BeginBorrowInst>(inst->getNextInstruction());
-        auto values = LiveValues::forGuaranteed(stored, borrow);
-        return values;
-      }
-      auto move = cast<MoveValueInst>(inst->getNextInstruction());
-      auto values = LiveValues::forOwned(stored, move);
-      return values;
+      auto stored = inst->getOperand(CopyLikeInstruction::Src);
+      auto lexical = getLexicalValueForStore(inst, asi);
+      return LiveValues::forValues(stored, lexical);
     }
 
     // If there is a Phi definition in this block:
@@ -1806,8 +1816,13 @@ void StackAllocationPromoter::run() {
   if (asi->getType().isOrHasEnum()) {
     for (auto it : initializationPoints) {
       auto *si = it.second;
-      auto src = si->getOperand(0);
-      valuesToComplete.push_back(src);
+      auto stored = si->getOperand(CopyLikeInstruction::Src);
+      valuesToComplete.push_back(stored);
+      if (lexicalLifetimeEnsured(asi)) {
+        if (auto lexical = getLexicalValueForStore(si, asi)) {
+          valuesToComplete.push_back(lexical);
+        }
+      }
     }
   }
 
@@ -1817,8 +1832,12 @@ void StackAllocationPromoter::run() {
   // Now, complete lifetimes!
   OSSALifetimeCompletion completion(function, domInfo);
 
+  // We may have incomplete lifetimes for enum locations on trivial paths.
+  // After promoting them, complete lifetime here.
   for (auto it : valuesToComplete) {
-    completion.completeOSSALifetime(it);
+    // Set forceBoundaryCompletion as true so that we complete at boundary for
+    // lexical values as well.
+    completion.completeOSSALifetime(it, /* forceBoundaryCompletion */ true);
   }
 }
 
@@ -2119,33 +2138,6 @@ bool MemoryToRegisters::promoteSingleAllocation(AllocStackInst *alloc) {
                             << *alloc);
     deleter.forceDeleteWithUsers(alloc);
     return true;
-  } else {
-    auto enableOptimizationForEnum = [](AllocStackInst *asi) {
-      if (asi->isLexical()) {
-        return false;
-      }
-      for (auto *use : asi->getUses()) {
-        auto *user = use->getUser();
-        if (!isa<StoreInst>(user) && !isa<StoreBorrowInst>(user)) {
-          continue;
-        }
-        auto stored = user->getOperand(CopyLikeInstruction::Src);
-        if (stored->isLexical()) {
-          return false;
-        }
-      }
-      return true;
-    };
-    // For stack locs of enum type that are lexical or with lexical stored
-    // values, we require that all uses are in the same block. This is because
-    // we can have incomplete lifetime of enum typed addresses, and on
-    // converting to value form this causes verification error. For all other
-    // stack locs of enum type, we use the lifetime completion utility to fix
-    // the lifetime. But when we have a lexical value, the utility can complete
-    // lifetimes on dead end blocks only.
-    if (f.hasOwnership() && alloc->getType().isOrHasEnum() &&
-        !enableOptimizationForEnum(alloc))
-      return false;
   }
 
   LLVM_DEBUG(llvm::dbgs() << "*** Need to insert BB arguments for " << *alloc);

lib/SILOptimizer/Transforms/TempRValueElimination.cpp

@@ -731,8 +731,6 @@ TempRValueOptPass::tryOptimizeStoreIntoTemp(StoreInst *si) {
     return std::next(si->getIterator());
   }
 
-  bool isOrHasEnum = tempObj->getType().isOrHasEnum();
-
   // Scan all uses of the temporary storage (tempObj) to verify they all refer
   // to the value initialized by this copy. It is sufficient to check that the
   // only users that modify memory are the copy_addr [initialization] and
@@ -743,15 +741,6 @@ TempRValueOptPass::tryOptimizeStoreIntoTemp(StoreInst *si) {
     if (user == si)
       continue;
 
-    // For lexical stored values that are enums, we require that all uses are in
-    // the same block. This is because we can have incomplete address lifetimes
-    // on none/trivial paths. and OSSALifetimeCompletion currently can complete
-    // lexical values only in the presence of dead end blocks.
-    if (isOrHasEnum && si->getSrc()->isLexical() &&
-        user->getParent() != si->getParent() && !isa<DeallocStackInst>(user)) {
-      return std::next(si->getIterator());
-    }
-
     // Bail if there is any kind of user which is not handled in the code below.
     switch (user->getKind()) {
       case SILInstructionKind::DestroyAddrInst:
@@ -955,7 +944,7 @@ void TempRValueOptPass::run() {
   // Call the utlity to complete ossa lifetime.
   OSSALifetimeCompletion completion(function, da->get(function));
   for (auto it : valuesToComplete) {
-    completion.completeOSSALifetime(it);
+    completion.completeOSSALifetime(it, /* forceBoundaryCompletion */ true);
   }
 }
 

test/SILOptimizer/mem2reg_lifetime_nontrivial.sil

@@ -1061,7 +1061,7 @@ enum KlassOptional {
 sil @return_optional_or_error : $@convention(thin) () -> (@owned KlassOptional, @error Error)
 
 // CHECK-LABEL: sil [ossa] @test_deadphi4 :
-// CHECK:         alloc_stack
+// CHECK-NOT:         alloc_stack
 // CHECK-LABEL: } // end sil function 'test_deadphi4'
 sil [ossa] @test_deadphi4 : $@convention(thin) (@owned KlassOptional) -> () {
 bb0(%0 : @owned $KlassOptional):

test/SILOptimizer/mem2reg_ossa.sil

@@ -516,7 +516,7 @@ bb7:
 }
 
 // CHECK-LABEL: sil [ossa] @test_optional_in_multiple_blocks_lexical :
-// CHECK:         alloc_stack
+// CHECK-NOT:         alloc_stack
 // CHECK:       } // end sil function 'test_optional_in_multiple_blocks_lexical'
 sil [ossa] @test_optional_in_multiple_blocks_lexical : $@convention(method) (@guaranteed Optional<Klass>) -> () {
 bb0(%0 : @guaranteed $Optional<Klass>):
@@ -542,7 +542,7 @@ bb7:
 }
 
 // CHECK-LABEL: sil [ossa] @test_optional_in_multiple_blocks_lexical_storedvalue :
-// CHECK:         alloc_stack
+// CHECK-NOT:         alloc_stack
 // CHECK:       } // end sil function 'test_optional_in_multiple_blocks_lexical_storedvalue'
 sil [ossa] @test_optional_in_multiple_blocks_lexical_storedvalue : $@convention(method) (@owned Optional<Klass>) -> () {
 bb0(%0 : @owned $Optional<Klass>):

test/SILOptimizer/temp_rvalue_opt_ossa.sil

@@ -1364,7 +1364,7 @@ bb7:
 }
 
 // CHECK-LABEL: sil [ossa] @test_optimize_store_of_enum2
-// CHECK:      alloc_stack
+// CHECK-NOT:      alloc_stack
 // CHECK:       } // end sil function 'test_optimize_store_of_enum2'
 sil [ossa] @test_optimize_store_of_enum2 : $@convention(method) <Element> (@owned Optional<Klass>) -> () {
 bb0(%0 : @owned $Optional<Klass>):