Swift Optimizer: add Onone simplification of apply instructions

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/CMakeLists.txt

@@ -7,6 +7,7 @@
 # See http://swift.org/CONTRIBUTORS.txt for Swift project authors
 
 swift_compiler_sources(Optimizer
+  SimplifyApply.swift
   SimplifyBeginCOWMutation.swift
   SimplifyGlobalValue.swift
   SimplifyStrongRetainRelease.swift)

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyApply.swift

@@ -0,0 +1,61 @@
+//===--- SimplifyApply.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 ApplyInst : OnoneSimplifyable {
+  func simplify(_ context: SimplifyContext) {
+    tryReplaceTrivialApplyOfPartialApply(context)
+  }
+}
+
+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()
+    }
+  }
+}
+
+private func allArgumentsAreTrivial(_ args: LazyMapSequence<OperandArray, Value>) -> Bool {
+  return !args.contains { !$0.hasTrivialType }
+}

test/SILOptimizer/simplify_apply.sil

@@ -0,0 +1,78 @@
+// RUN: %target-sil-opt -enable-sil-verify-all %s -onone-simplification -simplify-instruction=apply | %FileCheck %s
+
+// REQUIRES: swift_in_compiler
+
+import Swift
+import Builtin
+
+sil @closure_with_args : $@convention(thin) (Int, Bool) -> ()
+sil @closure2_with_args : $@convention(thin) (Int, String) -> ()
+sil @closure3_with_args : $@convention(thin) (String, Bool) -> ()
+sil @generic_callee_inguaranteed : $@convention(thin) <T, U> (@in_guaranteed T, @in_guaranteed U) -> ()
+
+// CHECK-LABEL: sil @test_apply_of_partial_apply
+// CHECK:         [[F:%.*]] = function_ref @closure_with_args
+// CHECK-NOT:     partial_apply
+// CHECK:         apply [[F]](%0, %1)
+// CHECK:       } // end sil function 'test_apply_of_partial_apply'
+sil @test_apply_of_partial_apply : $@convention(thin) (Int, Bool) -> () {
+bb0(%0 : $Int, %1 : $Bool):
+  %2 = function_ref @closure_with_args : $@convention(thin) (Int, Bool) -> ()
+  %3 = partial_apply %2(%1) : $@convention(thin) (Int, Bool) -> ()
+  apply %3(%0) : $@callee_owned (Int) -> ()
+  %r = tuple()
+  return %r : $()
+}
+
+// Currently this is not optimized by the simplification passes.
+// TODO: change the check lines if we can handle generic closures
+
+// CHECK-LABEL: sil @test_generic_partial_apply_apply_inguaranteed
+// CHECK:         [[F:%.*]] = function_ref @generic_callee_inguaranteed
+// CHECK:         [[PA:%.*]] = partial_apply [[F]]
+// CHECK:         apply [[PA]]
+// CHECK:       } // end sil function 'test_generic_partial_apply_apply_inguaranteed'
+sil @test_generic_partial_apply_apply_inguaranteed : $@convention(thin) <T> (@in T, @in T) -> () {
+bb0(%0 : $*T, %1 : $*T):
+  %f1 = function_ref @generic_callee_inguaranteed : $@convention(thin) <T, U> (@in_guaranteed T, @in_guaranteed U) -> ()
+  %pa = partial_apply %f1<T, T>(%1) : $@convention(thin) <T, U> (@in_guaranteed T, @in_guaranteed U) -> ()
+  %a1 = apply %pa(%0) : $@callee_owned (@in_guaranteed T) -> ()
+  destroy_addr %0 : $*T
+  %r = tuple ()
+  return %r : $()
+}
+
+// Currently this is not optimized by the simplification passes.
+// TODO: change the check lines if we can handle non-trivial arguments
+
+// CHECK-LABEL: sil @dont_handle_non_trivial_pa_args
+// CHECK:         [[F:%.*]] = function_ref @closure2_with_args
+// CHECK:         [[PA:%.*]] = partial_apply [[F]]
+// CHECK:         apply [[PA]]
+// CHECK:       } // end sil function 'dont_handle_non_trivial_pa_args'
+sil @dont_handle_non_trivial_pa_args : $@convention(thin) (Int, String) -> () {
+bb0(%0 : $Int, %1 : $String):
+  %2 = function_ref @closure2_with_args : $@convention(thin) (Int, String) -> ()
+  %3 = partial_apply %2(%1) : $@convention(thin) (Int, String) -> ()
+  apply %3(%0) : $@callee_owned (Int) -> ()
+  %r = tuple()
+  return %r : $()
+}
+
+// Currently this is not optimized by the simplification passes.
+// TODO: change the check lines if we can handle non-trivial arguments
+
+// CHECK-LABEL: sil @dont_handle_non_trivial_apply_args
+// CHECK:         [[F:%.*]] = function_ref @closure3_with_args
+// CHECK:         [[PA:%.*]] = partial_apply [[F]]
+// CHECK:         apply [[PA]]
+// CHECK:       } // end sil function 'dont_handle_non_trivial_apply_args'
+sil @dont_handle_non_trivial_apply_args : $@convention(thin) (String, Bool) -> () {
+bb0(%0 : $String, %1 : $Bool):
+  %2 = function_ref @closure3_with_args : $@convention(thin) (String, Bool) -> ()
+  %3 = partial_apply %2(%1) : $@convention(thin) (String, Bool) -> ()
+  apply %3(%0) : $@callee_owned (String) -> ()
+  %r = tuple()
+  return %r : $()
+}
+