Merge pull request #67860 from eeckstein/fix-mandatory-optmization-inlining

MandatoryPerformanceOptimizations: fix some problems with inilning

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -99,16 +99,14 @@ private func inlineAndDevirtualize(apply: FullApplySite, alreadyInlinedFunctions
     return
   }
 
-  if shouldInline(apply: apply, callee: callee, alreadyInlinedFunctions: &alreadyInlinedFunctions) {
-    simplifyCtxt.inlineFunction(apply: apply, mandatoryInline: true)
-
-    // In OSSA `partial_apply [on_stack]`s are represented as owned values rather than stack locations.
-    // It is possible for their destroys to violate stack discipline.
-    // When inlining into non-OSSA, those destroys are lowered to dealloc_stacks.
-    // This can result in invalid stack nesting.
-    if callee.hasOwnership && !apply.parentFunction.hasOwnership  {
+  if apply.canInline &&
+     shouldInline(apply: apply, callee: callee, alreadyInlinedFunctions: &alreadyInlinedFunctions)
+  {
+    if apply.inliningCanInvalidateStackNesting  {
       simplifyCtxt.notifyInvalidatedStackNesting()
     }
+
+    simplifyCtxt.inlineFunction(apply: apply, mandatoryInline: true)
   }
 }
 

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -110,10 +110,15 @@ extension MutatingContext {
   }
 
   func inlineFunction(apply: FullApplySite, mandatoryInline: Bool) {
+    // This is only a best-effort attempt to notity the new cloned instructions as changed.
+    // TODO: get a list of cloned instructions from the `inlineFunction`
     let instAfterInling: Instruction?
     switch apply {
-    case is ApplyInst, is BeginApplyInst:
+    case is ApplyInst:
       instAfterInling = apply.next
+    case let beginApply as BeginApplyInst:
+      let next = beginApply.next!
+      instAfterInling = (next is EndApplyInst ? nil : next)
     case is TryApplyInst:
       instAfterInling = apply.parentBlock.next?.instructions.first
     default:

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 import SIL
+import OptimizerBridging
 
 extension Value {
   var nonDebugUses: LazyFilterSequence<UseList> {
@@ -396,3 +397,37 @@ extension Function {
   }
 }
 
+extension FullApplySite {
+  var canInline: Bool {
+    // Some checks which are implemented in C++
+    if !FullApplySite_canInline(bridged) {
+      return false
+    }
+    // Cannot inline a non-inlinable function it an inlinable function.
+    if parentFunction.isSerialized,
+       let calleeFunction = referencedFunction,
+       !calleeFunction.isSerialized {
+      return false
+    }
+    return true
+  }
+
+  var inliningCanInvalidateStackNesting: Bool {
+    guard let calleeFunction = referencedFunction else {
+      return false
+    }
+
+    // In OSSA `partial_apply [on_stack]`s are represented as owned values rather than stack locations.
+    // It is possible for their destroys to violate stack discipline.
+    // When inlining into non-OSSA, those destroys are lowered to dealloc_stacks.
+    // This can result in invalid stack nesting.
+    if calleeFunction.hasOwnership && !parentFunction.hasOwnership {
+      return true
+    }
+    // Inlining of coroutines can result in improperly nested stack allocations.
+    if self is BeginApplyInst {
+      return true
+    }
+    return false
+  }
+}

include/swift/SILOptimizer/OptimizerBridging.h

@@ -513,6 +513,8 @@ struct BridgedPassContext {
   bool enableSimplificationFor(BridgedInstruction inst) const;
 };
 
+bool FullApplySite_canInline(BridgedInstruction apply);
+
 //===----------------------------------------------------------------------===//
 //                          Pass registration
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -1636,6 +1636,10 @@ bool BridgedPassContext::enableSimplificationFor(BridgedInstruction inst) const
   return false;
 }
 
+bool FullApplySite_canInline(BridgedInstruction apply) {
+  return swift::SILInliner::canInlineApplySite(swift::FullApplySite(apply.getInst()));
+}
+
 // TODO: can't be inlined to work around https://github.com/apple/swift/issues/64502
 CalleeList BridgedCalleeAnalysis::getCallees(BridgedValue callee) const {
   return ca->getCalleeListOfValue(callee.getSILValue());

test/SILOptimizer/mandatory_performance_optimizations.sil

@@ -8,6 +8,8 @@ import Builtin
 import Swift
 import SwiftShims
 
+sil_global [let] @g1 : $Int32
+sil_global [let] @g2 : $Int32
 
 sil @paable : $@convention(thin) (Builtin.Int64) -> ()
 sil @moved_pai_callee : $@convention(thin) (@inout_aliasable Builtin.Int64) -> ()
@@ -130,3 +132,79 @@ sil_witness_table public_external [serialized] Int: Comparable module Swift {
   method #Comparable."<": <Self where Self : Comparable> (Self.Type) -> (Self, Self) -> Bool : @$sSiSLsSL1loiySbx_xtFZTW
 }
 
+sil @get_int_value : $@convention(thin) () -> Int32 {
+bb0:
+  %0 = integer_literal $Builtin.Int32, 10
+  %1 = struct $Int32 (%0 : $Builtin.Int32)
+  return %1 : $Int32
+}
+
+// CHECK-LABEL: sil [global_init_once_fn] [no_locks] @globalinit_inline_into_init :
+// CHECK-NOT:     apply
+// CHECK:       } // end sil function 'globalinit_inline_into_init'
+sil [global_init_once_fn] [no_locks] @globalinit_inline_into_init : $@convention(c) () -> () {
+bb0:
+  alloc_global @g1
+  %1 = global_addr @g1 : $*Int32
+  %2 = function_ref @get_int_value : $@convention(thin) () -> Int32
+  %3 = apply %2() : $@convention(thin) () -> Int32
+  store %3 to %1 : $*Int32
+  %6 = tuple ()
+  return %6 : $()
+}
+
+// CHECK-LABEL: sil [serialized] [global_init_once_fn] [no_locks] @globalinit_dont_inline_non_inlinable_into_inlinable :
+// CHECK:         apply
+// CHECK:       } // end sil function 'globalinit_dont_inline_non_inlinable_into_inlinable'
+sil [serialized] [global_init_once_fn] [no_locks] @globalinit_dont_inline_non_inlinable_into_inlinable : $@convention(c) () -> () {
+bb0:
+  alloc_global @g2
+  %1 = global_addr @g2 : $*Int32
+  %2 = function_ref @get_int_value : $@convention(thin) () -> Int32
+  %3 = apply %2() : $@convention(thin) () -> Int32
+  store %3 to %1 : $*Int32
+  %6 = tuple ()
+  return %6 : $()
+}
+
+sil @yield_int_value : $@convention(thin) @yield_once () -> (@yields Int32) {
+bb0:
+  %0 = integer_literal $Builtin.Int32, 10
+  %1 = struct $Int32 (%0 : $Builtin.Int32)
+  yield %1 : $Int32, resume bb1, unwind bb2
+bb1:
+  %3 = tuple ()
+  return %3 : $()
+bb2:
+  unwind
+}
+
+// CHECK-LABEL: sil [no_locks] @inline_begin_apply :
+// CHECK-NOT:     begin_apply
+// CHECK:       } // end sil function 'inline_begin_apply'
+sil [no_locks] @inline_begin_apply : $@convention(thin) () -> Int32 {
+bb0:
+  %0 = function_ref @yield_int_value : $@convention(thin) @yield_once () -> (@yields Int32)
+  (%1, %2) = begin_apply %0() : $@convention(thin) @yield_once () -> (@yields Int32)
+  end_apply %2
+  return %1 : $Int32
+}
+
+// CHECK-LABEL: sil [no_locks] @dont_inline_begin_apply :
+// CHECK:         begin_apply
+// CHECK:       } // end sil function 'dont_inline_begin_apply'
+sil [no_locks] @dont_inline_begin_apply : $@convention(thin) () -> Int32 {
+bb0:
+  %0 = function_ref @yield_int_value : $@convention(thin) @yield_once () -> (@yields Int32)
+  (%1, %2) = begin_apply %0() : $@convention(thin) @yield_once () -> (@yields Int32)
+  cond_br undef, bb1, bb2
+bb1:
+  end_apply %2
+  br bb3
+bb2:
+  end_apply %2
+  br bb3
+bb3:
+  return %1 : $Int32
+}
+