Merge pull request swiftlang#78054 from eeckstein/ossa-opt-improvements

Some OSSA related optimization improvements

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/CMakeLists.txt

@@ -9,7 +9,7 @@
 swift_compiler_sources(Optimizer
   SimplifyAllocRefDynamic.swift
   SimplifyApply.swift
-  SimplifyBeginBorrow.swift
+  SimplifyBeginAndLoadBorrow.swift
   SimplifyBeginCOWMutation.swift
   SimplifyBranch.swift
   SimplifyBuiltin.swift

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyBeginBorrow.swift renamed to SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyBeginAndLoadBorrow.swift

@@ -1,4 +1,4 @@
-//===--- SimplifyBeginBorrow.swift ----------------------------------------===//
+//===--- SimplifyBeginAndLoadBorrow.swift ---------------------------------===//
 //
 // This source file is part of the Swift.org open source project
 //
@@ -12,7 +12,7 @@
 
 import SIL
 
-extension BeginBorrowInst : OnoneSimplifyable {
+extension BeginBorrowInst : OnoneSimplifyable, SILCombineSimplifyable {
   func simplify(_ context: SimplifyContext) {
     if borrowedValue.ownership == .owned,
        // We need to keep lexical lifetimes in place.
@@ -21,10 +21,51 @@ extension BeginBorrowInst : OnoneSimplifyable {
        !findPointerEscapingUse(of: borrowedValue)
     {
       tryReplaceBorrowWithOwnedOperand(beginBorrow: self, context)
+    } else {
+      removeBorrowOfThinFunction(beginBorrow: self, context)
     }
   }
 }
 
+extension LoadBorrowInst : Simplifyable, SILCombineSimplifyable {
+  func simplify(_ context: SimplifyContext) {
+    if uses.ignoreDebugUses.ignoreUsers(ofType: EndBorrowInst.self).isEmpty {
+      context.erase(instructionIncludingAllUsers: self)
+      return
+    }
+
+    // If the load_borrow is followed by a copy_value, combine both into a `load [copy]`:
+    // ```
+    //   %1 = load_borrow %0
+    //   %2 = some_forwarding_instruction %1 // zero or more forwarding instructions
+    //   %3 = copy_value %2
+    //   end_borrow %1
+    // ```
+    // ->
+    // ```
+    //   %1 = load [copy] %0
+    //   %3 = some_forwarding_instruction %1 // zero or more forwarding instructions
+    // ```
+    //
+    tryCombineWithCopy(context)
+  }
+
+  private func tryCombineWithCopy(_ context: SimplifyContext) {
+    let forwardedValue = lookThroughSingleForwardingUses()
+    guard let singleUser = forwardedValue.uses.ignoreUsers(ofType: EndBorrowInst.self).singleUse?.instruction,
+          let copy = singleUser as? CopyValueInst,
+          copy.parentBlock == self.parentBlock else {
+      return
+    }
+    let builder = Builder(before: self, context)
+    let loadCopy = builder.createLoad(fromAddress: address, ownership: .copy)
+    let forwardedOwnedValue = replace(guaranteedValue: self, withOwnedValue: loadCopy, context)
+    copy.uses.replaceAll(with: forwardedOwnedValue, context)
+    context.erase(instruction: copy)
+    context.erase(instructionIncludingAllUsers: self)
+  }
+}
+
 private func tryReplaceBorrowWithOwnedOperand(beginBorrow: BeginBorrowInst, _ context: SimplifyContext) {
   // The last value of a (potentially empty) forwarding chain, beginning at the `begin_borrow`.
   let forwardedValue = beginBorrow.lookThroughSingleForwardingUses()
@@ -38,6 +79,19 @@ private func tryReplaceBorrowWithOwnedOperand(beginBorrow: BeginBorrowInst, _ co
   }
 }
 
+private func removeBorrowOfThinFunction(beginBorrow: BeginBorrowInst, _ context: SimplifyContext) {
+  guard let thin2thickFn = beginBorrow.borrowedValue as? ThinToThickFunctionInst,
+        // For simplicity don't go into the trouble of removing reborrow phi arguments.
+        beginBorrow.uses.filterUsers(ofType: BranchInst.self).isEmpty else
+  {
+    return
+  }
+  // `thin_to_thick_function` has "none" ownership and is compatible with guaranteed values.
+  // Therefore the `begin_borrow` is not needed.
+  beginBorrow.uses.ignoreUsers(ofType: EndBorrowInst.self).replaceAll(with: thin2thickFn, context)
+  context.erase(instructionIncludingAllUsers: beginBorrow)
+}
+
 /// Replace
 /// ```
 ///   %1 = begin_borrow %0
@@ -156,7 +210,7 @@ private extension ForwardingInstruction {
 }
 
 /// Replaces a guaranteed value with an owned value.
-/// 
+///
 /// If the `guaranteedValue`'s use is a ForwardingInstruction (or forwarding instruction chain),
 /// it is converted to an owned version of the forwarding instruction (or instruction chain).
 ///

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyDestructure.swift

@@ -15,6 +15,21 @@ import SIL
 extension DestructureTupleInst : OnoneSimplifyable, SILCombineSimplifyable {
   func simplify(_ context: SimplifyContext) {
 
+    // If the tuple is trivial, replace
+    // ```
+    //   (%1, %2) = destructure_tuple %t
+    // ```
+    // ->
+    // ```
+    //   %1 = tuple_extract %t, 0
+    //   %2 = tuple_extract %t, 1
+    // ```
+    // This canonicalization helps other optimizations to e.g. CSE tuple_extracts.
+    //
+    if replaceWithTupleExtract(context) {
+      return
+    }
+
     // Eliminate the redundant instruction pair
     // ```
     //   %t = tuple (%0, %1, %2)
@@ -26,11 +41,39 @@ extension DestructureTupleInst : OnoneSimplifyable, SILCombineSimplifyable {
       tryReplaceConstructDestructPair(construct: tuple, destruct: self, context)
     }
   }
+
+  private func replaceWithTupleExtract(_ context: SimplifyContext) -> Bool {
+    guard self.tuple.type.isTrivial(in: parentFunction) else {
+      return false
+    }
+    let builder = Builder(before: self, context)
+    for (elementIdx, result) in results.enumerated() {
+      let elementValue = builder.createTupleExtract(tuple: self.tuple, elementIndex: elementIdx)
+      result.uses.replaceAll(with: elementValue, context)
+    }
+    context.erase(instruction: self)
+    return true
+  }
 }
 
 extension DestructureStructInst : OnoneSimplifyable, SILCombineSimplifyable {
   func simplify(_ context: SimplifyContext) {
 
+    // If the struct is trivial, replace
+    // ```
+    //   (%1, %2) = destructure_struct %s
+    // ```
+    // ->
+    // ```
+    //   %1 = struct_extract %s, #S.field0
+    //   %2 = struct_extract %s, #S.field1
+    // ```
+    // This canonicalization helps other optimizations to e.g. CSE tuple_extracts.
+    //
+    if replaceWithStructExtract(context) {
+      return
+    }
+
     switch self.struct {
     case let str as StructInst:
       // Eliminate the redundant instruction pair
@@ -82,6 +125,19 @@ extension DestructureStructInst : OnoneSimplifyable, SILCombineSimplifyable {
       break
     }
   }
+
+  private func replaceWithStructExtract(_ context: SimplifyContext) -> Bool {
+    guard self.struct.type.isTrivial(in: parentFunction) else {
+      return false
+    }
+    let builder = Builder(before: self, context)
+    for (fieldIdx, result) in results.enumerated() {
+      let fieldValue = builder.createStructExtract(struct: self.struct, fieldIndex: fieldIdx)
+      result.uses.replaceAll(with: fieldValue, context)
+    }
+    context.erase(instruction: self)
+    return true
+  }
 }
 
 private func tryReplaceConstructDestructPair(construct: SingleValueInstruction,

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -101,13 +101,15 @@ private func registerSwiftPasses() {
   registerPass(autodiffClosureSpecialization, { autodiffClosureSpecialization.run($0) })
 
   // Instruction passes
+  registerForSILCombine(BeginBorrowInst.self,      { run(BeginBorrowInst.self, $0) })
   registerForSILCombine(BeginCOWMutationInst.self, { run(BeginCOWMutationInst.self, $0) })
   registerForSILCombine(GlobalValueInst.self,      { run(GlobalValueInst.self, $0) })
   registerForSILCombine(StrongRetainInst.self,     { run(StrongRetainInst.self, $0) })
   registerForSILCombine(StrongReleaseInst.self,    { run(StrongReleaseInst.self, $0) })
   registerForSILCombine(RetainValueInst.self,      { run(RetainValueInst.self, $0) })
   registerForSILCombine(ReleaseValueInst.self,     { run(ReleaseValueInst.self, $0) })
   registerForSILCombine(LoadInst.self,             { run(LoadInst.self, $0) })
+  registerForSILCombine(LoadBorrowInst.self,       { run(LoadBorrowInst.self, $0) })
   registerForSILCombine(CopyValueInst.self,        { run(CopyValueInst.self, $0) })
   registerForSILCombine(DestroyValueInst.self,     { run(DestroyValueInst.self, $0) })
   registerForSILCombine(DestructureStructInst.self, { run(DestructureStructInst.self, $0) })

include/swift/SILOptimizer/PassManager/Passes.def

@@ -518,6 +518,7 @@ PASS(PruneVTables, "prune-vtables",
      "Mark class methods that do not require vtable dispatch")
 PASS_RANGE(AllPasses, AliasInfoDumper, PruneVTables)
 
+SWIFT_SILCOMBINE_PASS(BeginBorrowInst)
 SWIFT_SILCOMBINE_PASS(BeginCOWMutationInst)
 SWIFT_SILCOMBINE_PASS(ClassifyBridgeObjectInst)
 SWIFT_SILCOMBINE_PASS(GlobalValueInst)
@@ -526,6 +527,7 @@ SWIFT_SILCOMBINE_PASS(StrongReleaseInst)
 SWIFT_SILCOMBINE_PASS(RetainValueInst)
 SWIFT_SILCOMBINE_PASS(ReleaseValueInst)
 SWIFT_SILCOMBINE_PASS(LoadInst)
+SWIFT_SILCOMBINE_PASS(LoadBorrowInst)
 SWIFT_SILCOMBINE_PASS(CopyValueInst)
 SWIFT_SILCOMBINE_PASS(DestroyValueInst)
 SWIFT_SILCOMBINE_PASS(DestructureStructInst)

lib/SILOptimizer/IPO/GlobalPropertyOpt.cpp

@@ -219,9 +219,11 @@ bool GlobalPropertyOpt::canAddressEscape(SILValue V, bool acceptStore) {
 
     // These instructions do not cause the address to escape.
     if (isa<LoadInst>(User) ||
+        isa<LoadBorrowInst>(User) ||
         isa<DebugValueInst>(User) ||
         isa<StrongReleaseInst>(User) ||
         isa<StrongRetainInst>(User) ||
+        isa<DestroyAddrInst>(User) ||
         isa<DeallocBoxInst>(User) ||
         isa<DeallocStackInst>(User))
       continue;
@@ -285,10 +287,11 @@ void GlobalPropertyOpt::scanInstruction(swift::SILInstruction *Inst) {
       default:
         break;
     }
-  } else if (auto *LI = dyn_cast<LoadInst>(Inst)) {
+  } else if (isa<LoadInst>(Inst) || isa<LoadBorrowInst>(Inst)) {
+    auto *LI = cast<SingleValueInstruction>(Inst);
     if (isArrayType(LI->getType())) {
       // Add a dependency from the value at the address to the loaded value.
-      SILValue loadAddr = LI->getOperand();
+      SILValue loadAddr = LI->getOperand(0);
       assert(loadAddr->getType().isAddress());
       addDependency(getAddrEntry(loadAddr), getValueEntry(LI));
       return;

lib/SILOptimizer/LoopTransforms/ArrayOpt.h

@@ -81,10 +81,10 @@ class StructUseCollector {
 
   UserList AggregateAddressUsers;
   UserList StructAddressUsers;
-  SmallVector<LoadInst*, 16> StructLoads;
+  SmallVector<SingleValueInstruction *, 16> StructLoads;
   UserList StructValueUsers;
   UserOperList ElementAddressUsers;
-  SmallVector<std::pair<LoadInst*, Operand*>, 16> ElementLoads;
+  SmallVector<std::pair<SingleValueInstruction *, Operand*>, 16> ElementLoads;
   UserOperList ElementValueUsers;
   VisitedSet Visited;
 
@@ -126,7 +126,7 @@ class StructUseCollector {
       return false;
     for (SILInstruction *user : StructAddressUsers) {
       // ignore load users
-      if (isa<LoadInst>(user))
+      if (isa<LoadInst>(user) || isa<LoadBorrowInst>(user))
         continue;
       if (user != use1 && user != use2)
         return false;
@@ -162,8 +162,8 @@ class StructUseCollector {
       if (StructVal) {
         // Found a use of an element.
         assert(AccessPathSuffix.empty() && "should have accessed struct");
-        if (auto *LoadI = dyn_cast<LoadInst>(UseInst)) {
-          ElementLoads.push_back(std::make_pair(LoadI, StructVal));
+        if (isa<LoadInst>(UseInst) || isa<LoadBorrowInst>(UseInst)) {
+          ElementLoads.push_back(std::make_pair(cast<SingleValueInstruction>(UseInst), StructVal));
           continue;
         }
 
@@ -185,9 +185,9 @@ class StructUseCollector {
 
       if (AccessPathSuffix.empty()) {
         // Found a use of the struct at the given access path.
-        if (auto *LoadI = dyn_cast<LoadInst>(UseInst)) {
-          StructLoads.push_back(LoadI);
-          StructAddressUsers.push_back(LoadI);
+        if (isa<LoadInst>(UseInst) || isa<LoadBorrowInst>(UseInst)) {
+          StructLoads.push_back(cast<SingleValueInstruction>(UseInst));
+          StructAddressUsers.push_back(UseInst);
           continue;
         }
 

lib/SILOptimizer/LoopTransforms/COWArrayOpt.cpp

@@ -55,10 +55,10 @@ bool areArraysEqual(RCIdentityFunctionInfo *RCIA, SILValue A, SILValue B,
     return true;
   // We have stripped off struct_extracts. Remove the load to look at the
   // address we are loading from.
-  if (auto *ALoad = dyn_cast<LoadInst>(A))
-    A = ALoad->getOperand();
-  if (auto *BLoad = dyn_cast<LoadInst>(B))
-    B = BLoad->getOperand();
+  if (isa<LoadInst>(A) || isa<LoadBorrowInst>(A))
+    A = cast<SingleValueInstruction>(A)->getOperand(0);
+  if (isa<LoadInst>(B) || isa<LoadBorrowInst>(B))
+    B = cast<SingleValueInstruction>(B)->getOperand(0);
   // Strip off struct_extract_refs until we hit array address.
   if (ArrayAddress) {
     StructElementAddrInst *SEAI = nullptr;
@@ -463,6 +463,10 @@ bool COWArrayOpt::checkSafeArrayAddressUses(UserList &AddressUsers) {
       return false;
     }
 
+    if (isa<LoadBorrowInst>(UseInst)) {
+      continue;
+    }
+
     if (isa<DeallocStackInst>(UseInst)) {
       // Handle destruction of a local array.
       continue;
@@ -569,6 +573,9 @@ bool COWArrayOpt::checkSafeArrayValueUses(UserList &ArrayValueUsers) {
     if (isa<MarkDependenceInst>(UseInst))
       continue;
 
+    if (isa<EndBorrowInst>(UseInst))
+      continue;
+
     if (UseInst->isDebugInstruction())
       continue;