Swift Optimizer: simplification for apply, try_apply, begin_apply and partial_apply

* move the apply of partial_apply transformation from simplify-apply to simplify-partial_apply
* delete dead partial_apply instructions
* devirtualize apply, try_apply and begin_apply

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/CMakeLists.txt

@@ -17,6 +17,7 @@ swift_compiler_sources(Optimizer
   SimplifyDestructure.swift
   SimplifyGlobalValue.swift
   SimplifyLoad.swift
+  SimplifyPartialApply.swift
   SimplifyStrongRetainRelease.swift
   SimplifyStructExtract.swift
   SimplifyTupleExtract.swift

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyApply.swift

@@ -14,48 +14,18 @@ import SIL
 
 extension ApplyInst : OnoneSimplifyable {
   func simplify(_ context: SimplifyContext) {
-    tryReplaceTrivialApplyOfPartialApply(context)
+    _ = context.tryDevirtualize(apply: self, isMandatory: false)
   }
 }
 
-private extension ApplyInst {
-  func tryReplaceTrivialApplyOfPartialApply(_ context: SimplifyContext) {
-    guard let pa = callee as? PartialApplyInst else {
-      return
-    }
-
-    if pa.referencedFunction == nil {
-      return
-    }
-
-    // Currently we don't handle generic closures. For Onone this is good enough.
-    // TODO: handle it once we replace the SILCombine simplification with this.
-    if !allArgumentsAreTrivial(arguments) {
-      return
-    }
-
-    if !allArgumentsAreTrivial(pa.arguments) {
-      return
-    }
-
-    if !substitutionMap.isEmpty {
-      return
-    }
-
-    let allArgs = Array<Value>(arguments) + Array<Value>(pa.arguments)
-    let builder = Builder(before: self, context)
-    let newApply = builder.createApply(function: pa.callee, pa.substitutionMap, arguments: allArgs,
-                                       isNonThrowing: isNonThrowing, isNonAsync: isNonAsync,
-                                       specializationInfo: specializationInfo)
-    uses.replaceAll(with: newApply, context)
-    context.erase(instruction: self)
-
-    if context.tryDeleteDeadClosure(closure: pa) {
-      context.notifyInvalidatedStackNesting()
-    }
+extension TryApplyInst : OnoneSimplifyable {
+  func simplify(_ context: SimplifyContext) {
+    _ = context.tryDevirtualize(apply: self, isMandatory: false)
   }
 }
 
-private func allArgumentsAreTrivial(_ args: LazyMapSequence<OperandArray, Value>) -> Bool {
-  return !args.contains { !$0.hasTrivialType }
+extension BeginApplyInst : OnoneSimplifyable {
+  func simplify(_ context: SimplifyContext) {
+    _ = context.tryDevirtualize(apply: self, isMandatory: false)
+  }
 }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyPartialApply.swift

@@ -0,0 +1,34 @@
+//===--- SimplifyPartialApply.swift ---------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 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
+
+extension PartialApplyInst : OnoneSimplifyable {
+  func simplify(_ context: SimplifyContext) {
+    let optimizedApplyOfPartialApply = context.tryOptimizeApplyOfPartialApply(closure: self)
+    if optimizedApplyOfPartialApply {
+      context.notifyInvalidatedStackNesting()
+    }
+
+    if context.preserveDebugInfo && uses.contains(where: { $0.instruction is DebugValueInst }) {
+      return
+    }
+
+    // Try to delete the partial_apply.
+    // In case it became dead because of tryOptimizeApplyOfPartialApply, we don't
+    // need to copy all arguments again (to extend their lifetimes), because it
+    // was already done in tryOptimizeApplyOfPartialApply.
+    if context.tryDeleteDeadClosure(closure: self, needKeepArgsAlive: !optimizedApplyOfPartialApply) {
+      context.notifyInvalidatedStackNesting()
+    }
+  }
+}

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -71,6 +71,44 @@ extension MutatingContext {
     erase(instruction: inst)
   }
 
+  func tryOptimizeApplyOfPartialApply(closure: PartialApplyInst) -> Bool {
+    if _bridged.tryOptimizeApplyOfPartialApply(closure.bridged) {
+      notifyInstructionsChanged()
+      notifyCallsChanged()
+
+      for use in closure.callee.uses {
+        if use.instruction is FullApplySite {
+          notifyInstructionChanged(use.instruction)
+        }
+      }
+      return true
+    }
+    return false
+  }
+
+  func tryDeleteDeadClosure(closure: SingleValueInstruction, needKeepArgsAlive: Bool = true) -> Bool {
+    if _bridged.tryDeleteDeadClosure(closure.bridged, needKeepArgsAlive) {
+      notifyInstructionsChanged()
+      return true
+    }
+    return false
+  }
+
+  func tryDevirtualize(apply: FullApplySite, isMandatory: Bool) -> ApplySite? {
+    let result = _bridged.tryDevirtualizeApply(apply.bridged, isMandatory)
+    if let newApply = result.newApply.instruction {
+      erase(instruction: apply)
+      notifyInstructionsChanged()
+      notifyCallsChanged()
+      if result.cfgChanged {
+        notifyBranchesChanged()
+      }
+      notifyInstructionChanged(newApply)
+      return newApply as! FullApplySite
+    }
+    return nil
+  }
+
   /// Copies all instructions of a static init value of a global to the insertion point of `builder`.
   func copyStaticInitializer(fromInitValue: Value, to builder: Builder) -> Value? {
     let range = _bridged.copyStaticInitializer(fromInitValue.bridged, builder.bridged)
@@ -98,10 +136,6 @@ extension MutatingContext {
     }
   }
 
-  func tryDeleteDeadClosure(closure: SingleValueInstruction) -> Bool {
-    _bridged.tryDeleteDeadClosure(closure.bridged)
-  }
-
   func getContextSubstitutionMap(for type: Type) -> SubstitutionMap {
     SubstitutionMap(_bridged.getContextSubstitutionMap(type.bridged))
   }

include/swift/SILOptimizer/OptimizerBridging.h

@@ -202,7 +202,17 @@ struct BridgedPassContext {
     block.getBlock()->eraseFromParent();
   }
 
-  bool tryDeleteDeadClosure(BridgedInstruction closure) const;
+  bool tryOptimizeApplyOfPartialApply(BridgedInstruction closure) const;
+
+  bool tryDeleteDeadClosure(BridgedInstruction closure, bool needKeepArgsAlive) const;
+
+  struct DevirtResult {
+    OptionalBridgedInstruction newApply;
+    bool cfgChanged;
+  };
+
+  SWIFT_IMPORT_UNSAFE
+  DevirtResult tryDevirtualizeApply(BridgedInstruction apply, bool isMandatory) const;
 
   SWIFT_IMPORT_UNSAFE
   OptionalBridgedValue constantFoldBuiltin(BridgedInstruction builtin) const;

include/swift/SILOptimizer/Utils/Devirtualize.h

@@ -69,7 +69,8 @@ SubstitutionMap getWitnessMethodSubstitutions(SILModule &Module, ApplySite AI,
 /// Return the new apply and true if the CFG was also modified.
 std::pair<ApplySite, bool>
 tryDevirtualizeApply(ApplySite AI, ClassHierarchyAnalysis *CHA,
-                     OptRemark::Emitter *ORE = nullptr);
+                     OptRemark::Emitter *ORE = nullptr,
+                     bool isMandatory = false);
 bool canDevirtualizeApply(FullApplySite AI, ClassHierarchyAnalysis *CHA);
 bool canDevirtualizeClassMethod(FullApplySite AI, ClassDecl *CD,
                                 OptRemark::Emitter *ORE = nullptr,
@@ -108,7 +109,9 @@ tryDevirtualizeClassMethod(FullApplySite AI, SILValue ClassInstance,
 /// the original apply site.
 ///
 /// Return the new apply and true if the CFG was also modified.
-std::pair<ApplySite, bool> tryDevirtualizeWitnessMethod(ApplySite AI, OptRemark::Emitter *ORE);
+std::pair<ApplySite, bool> tryDevirtualizeWitnessMethod(ApplySite AI,
+                                                        OptRemark::Emitter *ORE,
+                                                        bool isMandatory);
 
 /// Delete a successfully-devirtualized apply site.  This must always be
 /// called after devirtualizing an apply; not only is it not semantically

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -30,6 +30,7 @@
 #include "swift/SILOptimizer/Utils/BasicBlockOptUtils.h"
 #include "swift/SILOptimizer/Utils/ConstantFolding.h"
 #include "swift/SILOptimizer/Utils/CFGOptUtils.h"
+#include "swift/SILOptimizer/Utils/Devirtualize.h"
 #include "swift/SILOptimizer/Utils/OptimizerStatsUtils.h"
 #include "swift/SILOptimizer/Utils/StackNesting.h"
 #include "swift/SILOptimizer/Utils/InstOptUtils.h"
@@ -1427,8 +1428,25 @@ std::string BridgedPassContext::getModuleDescription() const {
   return str;
 }
 
-bool BridgedPassContext::tryDeleteDeadClosure(BridgedInstruction closure) const {
-  return ::tryDeleteDeadClosure(closure.getAs<SingleValueInstruction>(), InstModCallbacks());
+bool BridgedPassContext::tryOptimizeApplyOfPartialApply(BridgedInstruction closure) const {
+  auto *pa = closure.getAs<PartialApplyInst>();
+  SILBuilder builder(pa);
+  return ::tryOptimizeApplyOfPartialApply(pa, builder.getBuilderContext(), InstModCallbacks());
+}
+
+bool BridgedPassContext::tryDeleteDeadClosure(BridgedInstruction closure, bool needKeepArgsAlive) const {
+  return ::tryDeleteDeadClosure(closure.getAs<SingleValueInstruction>(), InstModCallbacks(), needKeepArgsAlive);
+}
+
+BridgedPassContext::DevirtResult BridgedPassContext::tryDevirtualizeApply(BridgedInstruction apply,
+                                                                          bool isMandatory) const {
+  auto cha = invocation->getPassManager()->getAnalysis<ClassHierarchyAnalysis>();
+  auto result = ::tryDevirtualizeApply(ApplySite(apply.getInst()), cha, nullptr, isMandatory);
+  if (result.first) {
+    OptionalBridgedInstruction newApply(result.first.getInstruction()->asSILNode());
+    return {newApply, result.second};
+  }
+  return {{nullptr}, false};
 }
 
 OptionalBridgedValue BridgedPassContext::constantFoldBuiltin(BridgedInstruction builtin) const {

lib/SILOptimizer/Utils/Devirtualize.cpp

@@ -1140,7 +1140,7 @@ static bool isNonGenericThunkOfGenericExternalFunction(SILFunction *thunk) {
   return false;
 }
 
-static bool canDevirtualizeWitnessMethod(ApplySite applySite) {
+static bool canDevirtualizeWitnessMethod(ApplySite applySite, bool isMandatory) {
   SILFunction *f;
   SILWitnessTable *wt;
 
@@ -1183,7 +1183,7 @@ static bool canDevirtualizeWitnessMethod(ApplySite applySite) {
   // ```
   // In the defining module, the generic conformance can be specialized (which is not
   // possible in the client module, because it's not inlinable).
-  if (isNonGenericThunkOfGenericExternalFunction(f)) {
+  if (!isMandatory && isNonGenericThunkOfGenericExternalFunction(f)) {
     return false;
   }
 
@@ -1237,8 +1237,9 @@ static bool canDevirtualizeWitnessMethod(ApplySite applySite) {
 /// of a function_ref, returning the new apply.
 std::pair<ApplySite, bool>
 swift::tryDevirtualizeWitnessMethod(ApplySite applySite,
-                                    OptRemark::Emitter *ore) {
-  if (!canDevirtualizeWitnessMethod(applySite))
+                                    OptRemark::Emitter *ore,
+                                    bool isMandatory) {
+  if (!canDevirtualizeWitnessMethod(applySite, isMandatory))
     return {ApplySite(), false};
 
   SILFunction *f;
@@ -1262,7 +1263,7 @@ swift::tryDevirtualizeWitnessMethod(ApplySite applySite,
 /// Return the new apply and true if the CFG was also modified.
 std::pair<ApplySite, bool>
 swift::tryDevirtualizeApply(ApplySite applySite, ClassHierarchyAnalysis *cha,
-                            OptRemark::Emitter *ore) {
+                            OptRemark::Emitter *ore, bool isMandatory) {
   LLVM_DEBUG(llvm::dbgs() << "    Trying to devirtualize: "
                           << *applySite.getInstruction());
 
@@ -1272,7 +1273,7 @@ swift::tryDevirtualizeApply(ApplySite applySite, ClassHierarchyAnalysis *cha,
   //   %9 = apply %8<Self = CodeUnit?>(%6#1) : ...
   //
   if (isa<WitnessMethodInst>(applySite.getCallee()))
-    return tryDevirtualizeWitnessMethod(applySite, ore);
+    return tryDevirtualizeWitnessMethod(applySite, ore, isMandatory);
 
   // TODO: check if we can also de-virtualize partial applies of class methods.
   FullApplySite fas = FullApplySite::isa(applySite.getInstruction());
@@ -1344,7 +1345,7 @@ bool swift::canDevirtualizeApply(FullApplySite applySite,
   //   %9 = apply %8<Self = CodeUnit?>(%6#1) : ...
   //
   if (isa<WitnessMethodInst>(applySite.getCallee()))
-    return canDevirtualizeWitnessMethod(applySite);
+    return canDevirtualizeWitnessMethod(applySite, /*isMandatory*/ false);
 
   /// Optimize a class_method and alloc_ref pair into a direct function
   /// reference:

test/AutoDiff/SILOptimizer/differentiation_subset_parameters_thunk.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-frontend -emit-sil %s | %FileCheck %s
+// RUN: %target-swift-frontend -emit-sil -Xllvm -sil-disable-pass=OnoneSimplification %s | %FileCheck %s
 
 import _Differentiation
 

test/ClangImporter/serialization-sil.swift

@@ -1,5 +1,5 @@
 // RUN: %empty-directory(%t)
-// RUN: %target-swift-frontend -enable-copy-propagation=requested-passes-only -enable-lexical-lifetimes=false -emit-module-path %t/Test.swiftmodule -emit-sil -o /dev/null -module-name Test %s -sdk "" -import-objc-header %S/Inputs/serialization-sil.h
+// RUN: %target-swift-frontend -enable-copy-propagation=requested-passes-only -enable-lexical-lifetimes=false -Xllvm -sil-disable-pass=Simplification -emit-module-path %t/Test.swiftmodule -emit-sil -o /dev/null -module-name Test %s -sdk "" -import-objc-header %S/Inputs/serialization-sil.h
 // RUN: %target-sil-func-extractor %t/Test.swiftmodule -sil-print-debuginfo  -func='$s4Test16testPartialApplyyySoAA_pF' -o - | %FileCheck %s
 
 // REQUIRES: objc_interop