Merge pull request #77009 from eeckstein/fix-inlining-in-mpo

MandatoryPerformanceOptimization: don't let not-inlinable functions to be inlined

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -208,15 +208,30 @@ private func removeUnusedMetatypeInstructions(in function: Function, _ context:
 }
 
 private func shouldInline(apply: FullApplySite, callee: Function, alreadyInlinedFunctions: inout Set<PathFunctionTuple>) -> Bool {
-  if callee.isTransparent {
-    precondition(callee.hasOwnership, "transparent functions should have ownership at this stage of the pipeline")
-    return true
+  if let beginApply = apply as? BeginApplyInst,
+     !beginApply.canInline
+  {
+    return false
+  }
+
+  if !callee.canBeInlinedIntoCaller(withSerializedKind: apply.parentFunction.serializedKind) &&
+     // Even if the serialization kind doesn't match, we need to make sure to inline witness method thunks
+     // in embedded swift.
+     callee.thunkKind != .thunk
+  {
+    return false
   }
 
-  if !apply.canInline {
+  // Cannot inline a non-ossa function into an ossa function
+  if apply.parentFunction.hasOwnership && !callee.hasOwnership {
     return false
   }
 
+  if callee.isTransparent {
+    precondition(callee.hasOwnership, "transparent functions should have ownership at this stage of the pipeline")
+    return true
+  }
+
   if apply is BeginApplyInst {
     // Avoid co-routines because they might allocate (their context).
     return true

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -650,27 +650,6 @@ 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 let calleeFunction = referencedFunction,
-       !calleeFunction.canBeInlinedIntoCaller(parentFunction.serializedKind) {
-      return false
-    }
-
-    // Cannot inline a non-ossa function into an ossa function
-    if parentFunction.hasOwnership,
-      let calleeFunction = referencedFunction,
-      !calleeFunction.hasOwnership {
-      return false
-    }
-
-    return true
-  }
-
   var inliningCanInvalidateStackNesting: Bool {
     guard let calleeFunction = referencedFunction else {
       return false
@@ -691,6 +670,10 @@ extension FullApplySite {
   }
 }
 
+extension BeginApplyInst {
+  var canInline: Bool { BeginApply_canInline(bridged) }
+}
+
 extension GlobalVariable {
   /// Removes all `begin_access` and `end_access` instructions from the initializer.
   ///

SwiftCompilerSources/Sources/SIL/Function.swift

@@ -162,8 +162,26 @@ final public class Function : CustomStringConvertible, HasShortDescription, Hash
     }
   }
 
-  public func canBeInlinedIntoCaller(_ kind: SerializedKind) -> Bool {
-    bridged.canBeInlinedIntoCaller(serializedKindBridged(kind))
+  public func canBeInlinedIntoCaller(withSerializedKind callerSerializedKind: SerializedKind) -> Bool {
+    switch serializedKind {
+    // If both callee and caller are not_serialized, the callee can be inlined into the caller
+    // during SIL inlining passes even if it (and the caller) might contain private symbols.
+    case .notSerialized:
+      return callerSerializedKind == .notSerialized;
+
+    // If Package-CMO is enabled, we serialize package, public, and @usableFromInline decls as
+    // [serialized_for_package].
+    // Their bodies must not, however, leak into @inlinable functons (that are [serialized])
+    // since they are inlined outside of their defining module.
+    //
+    // If this callee is [serialized_for_package], the caller must be either non-serialized
+    // or [serialized_for_package] for this callee's body to be inlined into the caller.
+    // It can however be referenced by [serialized] caller.
+    case .serializedForPackage:
+      return callerSerializedKind != .serialized;
+    case .serialized:
+      return true;
+    }
   }
 
   public func hasValidLinkageForFragileRef(_ kind: SerializedKind) -> Bool {

include/swift/SILOptimizer/OptimizerBridging.h

@@ -401,7 +401,7 @@ struct BridgedPassContext {
   bool completeLifetime(BridgedValue value) const;
 };
 
-bool FullApplySite_canInline(BridgedInstruction apply);
+bool BeginApply_canInline(BridgedInstruction beginApply);
 
 enum class BridgedDynamicCastResult {
   willSucceed,

include/swift/SILOptimizer/Utils/SILInliner.h

@@ -58,6 +58,8 @@ class SILInliner {
       : FuncBuilder(FuncBuilder), deleter(deleter), IKind(IKind),
         ApplySubs(ApplySubs) {}
 
+  static bool canInlineBeginApply(BeginApplyInst *BA);
+
   /// Returns true if we are able to inline \arg AI.
   ///
   /// *NOTE* This must be checked before attempting to inline \arg AI. If one

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -2062,9 +2062,8 @@ bool BridgedPassContext::completeLifetime(BridgedValue value) const {
   return result == LifetimeCompletion::WasCompleted;
 }
 
-bool FullApplySite_canInline(BridgedInstruction apply) {
-  return swift::SILInliner::canInlineApplySite(
-      swift::FullApplySite(apply.unbridged()));
+bool BeginApply_canInline(BridgedInstruction beginApply) {
+  return swift::SILInliner::canInlineBeginApply(beginApply.getAs<BeginApplyInst>());
 }
 
 BridgedDynamicCastResult classifyDynamicCastBridged(BridgedType sourceTy, BridgedType destTy,

lib/SILOptimizer/Utils/SILInliner.cpp

@@ -29,7 +29,7 @@
 
 using namespace swift;
 
-static bool canInlineBeginApply(BeginApplyInst *BA) {
+bool SILInliner::canInlineBeginApply(BeginApplyInst *BA) {
   // Don't inline if we have multiple resumption sites (i.e. end_apply or
   // abort_apply instructions).  The current implementation clones a single
   // copy of the end_apply and abort_apply paths, so it can't handle values

test/embedded/dictionary-runtime.swift

@@ -7,18 +7,33 @@
 // REQUIRES: optimized_stdlib
 // REQUIRES: OS=macosx
 
+var dict1: [Int: Int] = [1:10]
+var dict2: [Int: Int] = [1:20]
+
+func do_swap(n: Int) {
+  swap(&dict2[n], &dict1[n])
+}
+
+
 @main
 struct Main {
   static func main() {
     var dict: [Int: StaticString] = [:]
     dict[11] = "hello"
     dict[33] = "!"
     dict[22] = "world"
+
+    // CHECK: hello world !
     for key in dict.keys.sorted() {
       print(dict[key]!, terminator: " ")
     }
     print("")
+
+    do_swap(n: 1)
+    // CHECK: 20
+    print(dict1[1]!)
+    // CHECK: 10
+    print(dict2[1]!)
   }
 }
 
-// CHECK: hello world !