Two optimization improvements to fix dead-object-elimination with OSSA

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/CMakeLists.txt

@@ -22,6 +22,7 @@ swift_compiler_sources(Optimizer
   SimplifyDebugStep.swift
   SimplifyDestroyValue.swift
   SimplifyDestructure.swift
+  SimplifyFixLifetime.swift
   SimplifyGlobalValue.swift
   SimplifyInitEnumDataAddr.swift
   SimplifyKeyPath.swift

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyFixLifetime.swift

@@ -0,0 +1,55 @@
+//===--- SimplifyFixLifetime.swift ----------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SIL
+
+/// Canonicalize a `fix_lifetime` from an address to a `load` + `fix_lifetime`:
+/// ```
+///    %1 = alloc_stack $T
+///    ...
+///    fix_lifetime %1
+/// ```
+/// ->
+/// ```
+///    %1 = alloc_stack $T
+///    ...
+///    %2 = load %1
+///    fix_lifetime %2
+/// ```
+///
+/// This transformation is done for `alloc_stack` and `store_borrow` (which always has an `alloc_stack`
+/// operand).
+/// The benefit of this transformation is that it enables other optimizations, like mem2reg.
+///
+extension FixLifetimeInst : Simplifyable, SILCombineSimplifyable {
+  func simplify(_ context: SimplifyContext) {
+    let opValue = operand.value
+    guard opValue is AllocStackInst || opValue is StoreBorrowInst,
+          opValue.type.isLoadable(in: parentFunction)
+    else {
+      return
+    }
+
+    let builder = Builder(before: self, context)
+    let loadedValue: Value
+    if !parentFunction.hasOwnership {
+      loadedValue = builder.createLoad(fromAddress: opValue, ownership: .unqualified)
+    } else if opValue.type.isTrivial(in: parentFunction) {
+      loadedValue = builder.createLoad(fromAddress: opValue, ownership: .trivial)
+    } else {
+      loadedValue = builder.createLoadBorrow(fromAddress: opValue)
+      Builder(after: self, context).createEndBorrow(of: loadedValue)
+    }
+    operand.set(to: loadedValue, context)
+  }
+}
+

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -103,6 +103,7 @@ private func registerSwiftPasses() {
   // Instruction passes
   registerForSILCombine(BeginBorrowInst.self,      { run(BeginBorrowInst.self, $0) })
   registerForSILCombine(BeginCOWMutationInst.self, { run(BeginCOWMutationInst.self, $0) })
+  registerForSILCombine(FixLifetimeInst.self,      { run(FixLifetimeInst.self, $0) })
   registerForSILCombine(GlobalValueInst.self,      { run(GlobalValueInst.self, $0) })
   registerForSILCombine(StrongRetainInst.self,     { run(StrongRetainInst.self, $0) })
   registerForSILCombine(StrongReleaseInst.self,    { run(StrongReleaseInst.self, $0) })

include/swift/SILOptimizer/PassManager/Passes.def

@@ -521,6 +521,7 @@ PASS_RANGE(AllPasses, AliasInfoDumper, PruneVTables)
 SWIFT_SILCOMBINE_PASS(BeginBorrowInst)
 SWIFT_SILCOMBINE_PASS(BeginCOWMutationInst)
 SWIFT_SILCOMBINE_PASS(ClassifyBridgeObjectInst)
+SWIFT_SILCOMBINE_PASS(FixLifetimeInst)
 SWIFT_SILCOMBINE_PASS(GlobalValueInst)
 SWIFT_SILCOMBINE_PASS(StrongRetainInst)
 SWIFT_SILCOMBINE_PASS(StrongReleaseInst)

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -281,7 +281,6 @@ class SILCombiner :
   SILInstruction *visitThickToObjCMetatypeInst(ThickToObjCMetatypeInst *TTOCMI);
   SILInstruction *visitObjCToThickMetatypeInst(ObjCToThickMetatypeInst *OCTTMI);
   SILInstruction *visitTupleExtractInst(TupleExtractInst *TEI);
-  SILInstruction *visitFixLifetimeInst(FixLifetimeInst *FLI);
   SILInstruction *visitSwitchValueInst(SwitchValueInst *SVI);
   SILInstruction *
   visitCheckedCastAddrBranchInst(CheckedCastAddrBranchInst *CCABI);

lib/SILOptimizer/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -1760,22 +1760,6 @@ SILInstruction *SILCombiner::visitTupleExtractInst(TupleExtractInst *TEI) {
   return nullptr;
 }
 
-SILInstruction *SILCombiner::visitFixLifetimeInst(FixLifetimeInst *fli) {
-  // fix_lifetime(alloc_stack) -> fix_lifetime(load(alloc_stack))
-  Builder.setCurrentDebugScope(fli->getDebugScope());
-  if (auto *ai = dyn_cast<AllocStackInst>(fli->getOperand())) {
-    if (fli->getOperand()->getType().isLoadable(*fli->getFunction())) {
-      // load when ossa is disabled
-      auto load = Builder.emitLoadBorrowOperation(fli->getLoc(), ai);
-      Builder.createFixLifetime(fli->getLoc(), load);
-      // no-op when ossa is disabled
-      Builder.emitEndBorrowOperation(fli->getLoc(), load);
-      return eraseInstFromFunction(*fli);
-    }
-  }
-  return nullptr;
-}
-
 static std::optional<SILType>
 shouldReplaceCallByContiguousArrayStorageAnyObject(SILFunction &F,
                                                    CanType storageMetaTy) {

lib/SILOptimizer/Transforms/DeadObjectElimination.cpp

@@ -604,7 +604,8 @@ recursivelyCollectInteriorUses(ValueBase *DefInst,
 
     // Lifetime endpoints that don't allow the address to escape.
     if (isa<RefCountingInst>(User) || isa<DebugValueInst>(User) ||
-        isa<FixLifetimeInst>(User) || isa<DestroyValueInst>(User)) {
+        isa<FixLifetimeInst>(User) || isa<DestroyValueInst>(User) ||
+        isa<EndBorrowInst>(User)) {
       AllUsers.insert(User);
       continue;
     }
@@ -632,6 +633,13 @@ recursivelyCollectInteriorUses(ValueBase *DefInst,
       }
       continue;
     }
+    if (auto *bb = dyn_cast<BeginBorrowInst>(User)) {
+      if (!recursivelyCollectInteriorUses(bb, AddressNode,
+                                          IsInteriorAddress)) {
+        return false;
+      }
+      continue;
+    }
     if (auto PTAI = dyn_cast<PointerToAddressInst>(User)) {
       // Only one pointer-to-address is allowed for safety.
       if (SeenPtrToAddr)

test/SILOptimizer/dead_array_elim_ossa.sil

@@ -31,6 +31,7 @@ struct NonTrivialStruct {
 sil [_semantics "array.uninitialized_intrinsic"] @allocArray : $@convention(thin) <τ_0_0> (Builtin.Word) -> @owned (Array<τ_0_0>, Builtin.RawPointer)
 
 sil [_semantics "array.uninitialized"] @adoptStorageSpecializedForInt : $@convention(method) (@guaranteed _ContiguousArrayStorage<Int>, Builtin.Word, @thin Array<Int>.Type) -> (@owned Array<Int>, UnsafeMutablePointer<Int>)
+sil [_semantics "array.uninitialized"] @allocUninitialized : $@convention(thin) (Int) -> (@owned Array<Int>, UnsafeMutablePointer<Int>)
 
 sil [_semantics "array.finalize_intrinsic"] @finalize : $@convention(thin) (@owned Array<Int>) -> @owned Array<Int>
 
@@ -190,3 +191,23 @@ bb0(%0 : @owned $Klass):
   return %t : $()
 }
 
+
+// CHECK-LABEL: sil [ossa] @dead_array_with_borrow :
+// CHECK-NOT:     apply
+// CHECK:         %1 = tuple
+// CHECK-NEXT:    return
+// CHECK-LABEL: } // end sil function 'dead_array_with_borrow'
+sil [ossa] @dead_array_with_borrow : $@convention(thin) (Int) -> () {
+bb0(%0 : $Int):
+  %2 = function_ref @allocUninitialized : $@convention(thin) (Int) -> (@owned Array<Int>, UnsafeMutablePointer<Int>)
+  %3 = apply %2(%0) : $@convention(thin) (Int) -> (@owned Array<Int>, UnsafeMutablePointer<Int>)
+  (%4, %5) = destructure_tuple %3
+  debug_value %4, let, name "a"
+  %7 = begin_borrow [lexical] %4
+  fix_lifetime %7
+  end_borrow %7
+  destroy_value %4
+  %11 = tuple ()
+  return %11
+}
+

test/SILOptimizer/simplify_fix_lifetime.sil

@@ -0,0 +1,106 @@
+// RUN: %target-sil-opt %s -simplification -simplify-instruction=fix_lifetime | %FileCheck %s
+
+import Swift
+import Builtin
+
+class C {}
+
+// CHECK-LABEL: sil @load_from_alloc_stack :
+// CHECK:         %1 = alloc_stack
+// CHECK-NEXT:    store %0 to %1
+// CHECK-NEXT:    %3 = load %1
+// CHECK-NEXT:    fix_lifetime %3
+// CHECK-NEXT:    dealloc_stack %1
+// CHECK:       } // end sil function 'load_from_alloc_stack'
+sil @load_from_alloc_stack : $@convention(thin) (@owned C) -> () {
+bb0(%0: $C):
+  %1 = alloc_stack $C
+  store %0 to %1
+  fix_lifetime %1
+  dealloc_stack %1
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil @load_from_inout :
+// CHECK-NOT:     load
+// CHECK:         fix_lifetime %0
+// CHECK:       } // end sil function 'load_from_inout'
+sil @load_from_inout : $@convention(thin) (@inout C) -> () {
+bb0(%0: $*C):
+  fix_lifetime %0
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil @not_loadable :
+// CHECK-NOT:     load
+// CHECK:         fix_lifetime %1
+// CHECK:       } // end sil function 'not_loadable'
+sil @not_loadable : $@convention(thin) <T> (@in_guaranteed T) -> () {
+bb0(%0: $*T):
+  %1 = alloc_stack $T
+  copy_addr %0 to %1
+  fix_lifetime %1
+  dealloc_stack %1
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil [ossa] @load_from_alloc_stack_ossa :
+// CHECK:         %1 = alloc_stack
+// CHECK-NEXT:    store %0 to [init] %1
+// CHECK-NEXT:    %3 = load_borrow %1
+// CHECK-NEXT:    fix_lifetime %3
+// CHECK-NEXT:    end_borrow %3
+// CHECK-NEXT:    destroy_addr %1
+// CHECK-NEXT:    dealloc_stack %1
+// CHECK:       } // end sil function 'load_from_alloc_stack_ossa'
+sil [ossa] @load_from_alloc_stack_ossa : $@convention(thin) (@owned C) -> () {
+bb0(%0: @owned $C):
+  %1 = alloc_stack $C
+  store %0 to [init] %1
+  fix_lifetime %1
+  destroy_addr %1
+  dealloc_stack %1
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil [ossa] @load_from_store_borrow :
+// CHECK:         %1 = alloc_stack
+// CHECK-NEXT:    %2 = store_borrow %0 to %1
+// CHECK-NEXT:    %3 = load_borrow %2
+// CHECK-NEXT:    fix_lifetime %3
+// CHECK-NEXT:    end_borrow %3
+// CHECK-NEXT:    end_borrow %2
+// CHECK-NEXT:    dealloc_stack %1
+// CHECK:       } // end sil function 'load_from_store_borrow'
+sil [ossa] @load_from_store_borrow : $@convention(thin) (@guaranteed C) -> () {
+bb0(%0: @guaranteed $C):
+  %1 = alloc_stack $C
+  %2 = store_borrow %0 to %1
+  fix_lifetime %2
+  end_borrow %2
+  dealloc_stack %1
+  %r = tuple ()
+  return %r : $()
+}
+
+// CHECK-LABEL: sil [ossa] @load_trivial_from_alloc_stack :
+// CHECK:         %1 = alloc_stack
+// CHECK-NEXT:    store %0 to [trivial] %1
+// CHECK-NEXT:    %3 = load [trivial] %1
+// CHECK-NEXT:    fix_lifetime %3
+// CHECK-NEXT:    dealloc_stack %1
+// CHECK:       } // end sil function 'load_trivial_from_alloc_stack'
+sil [ossa] @load_trivial_from_alloc_stack : $@convention(thin) (Int) -> () {
+bb0(%0: $Int):
+  %1 = alloc_stack $Int
+  store %0 to [trivial] %1
+  fix_lifetime %1
+  dealloc_stack %1
+  %r = tuple ()
+  return %r : $()
+}
+