Optimizer: add simplification for fix_lifetime

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.

This peephole optimization was already done in SILCombine, but it didn't handle store_borrow.
A good opportunity to make an instruction simplification out of it.

This is part of fixing regressions when enabling OSSA modules:
rdar://140229560

Author: eeckstein

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) {

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 : $()
+}
+