Swift Optimizer: add the MandatoryPerformanceOptimizations pass

As a replacement for the old MandatoryGenericSpecializer

The pass it not enabled yet in the pass pipeline

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/SimplificationPasses.swift

@@ -71,9 +71,9 @@ let lateOnoneSimplificationPass = FunctionPass(name: "late-onone-simplification"
 //===--------------------------------------------------------------------===//
 
 
-private func runSimplification(on function: Function, _ context: FunctionPassContext,
-                               preserveDebugInfo: Bool,
-                               _ simplify: (Instruction, SimplifyContext) -> ()) {
+func runSimplification(on function: Function, _ context: FunctionPassContext,
+                       preserveDebugInfo: Bool,
+                       _ simplify: (Instruction, SimplifyContext) -> ()) {
   var worklist = InstructionWorklist(context)
   defer { worklist.deinitialize() }
 

SwiftCompilerSources/Sources/Optimizer/ModulePasses/CMakeLists.txt

@@ -7,6 +7,7 @@
 # See http://swift.org/CONTRIBUTORS.txt for Swift project authors
 
 swift_compiler_sources(Optimizer
+  MandatoryPerformanceOptimizations.swift
   ReadOnlyGlobalVariables.swift
   StackProtection.swift
 )

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -0,0 +1,170 @@
+//===--- MandatoryPerformanceOptimizations.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
+
+/// Performs mandatory optimizations for performance-annotated functions.
+///
+/// Optimizations include:
+/// * de-virtualization
+/// * mandatory inlining
+/// * generic specialization
+/// * mandatory memory optimizations
+/// * dead alloc elimination
+/// * instruction simplification
+///
+/// The pass starts with performance-annotated functions and transitively handles
+/// called functions.
+///
+let mandatoryPerformanceOptimizations = ModulePass(name: "mandatory-performance-optimizations") {
+  (moduleContext: ModulePassContext) in
+
+  var worklist = FunctionWorklist()
+  worklist.addAllPerformanceAnnotatedFunctions(of: moduleContext)
+
+  optimizeFunctionsTopDown(using: &worklist, moduleContext)
+}
+
+private func optimizeFunctionsTopDown(using worklist: inout FunctionWorklist,
+                                      _ moduleContext: ModulePassContext) {
+  while let f = worklist.pop() {
+    moduleContext.transform(function: f) { context in
+      if !context.loadFunction(function: f, loadCalleesRecursively: true) {
+        return
+      }
+      optimize(function: f, context)
+      worklist.add(calleesOf: f)
+    }
+  }
+}
+
+private func optimize(function: Function, _ context: FunctionPassContext) {
+  runSimplification(on: function, context, preserveDebugInfo: true) { instruction, simplifyCtxt in
+    if let i = instruction as? OnoneSimplifyable {
+      i.simplify(simplifyCtxt)
+      if instruction.isDeleted {
+        return
+      }
+    }
+    switch instruction {
+    case let apply as FullApplySite:
+      inlineAndDevirtualize(apply: apply, context, simplifyCtxt)
+    case let mt as MetatypeInst:
+      if mt.isTriviallyDeadIgnoringDebugUses {
+        simplifyCtxt.erase(instructionIncludingDebugUses: mt)
+      }
+    default:
+      break
+    }
+  }
+
+  _ = context.specializeApplies(in: function, isMandatory: true)
+
+  // If this is a just specialized function, try to optimize copy_addr, etc.
+  if context.optimizeMemoryAccesses(in: function) {
+    _ = context.eliminateDeadAllocations(in: function)
+  }
+}
+
+private func inlineAndDevirtualize(apply: FullApplySite, _ context: FunctionPassContext, _ simplifyCtxt: SimplifyContext) {
+
+  if simplifyCtxt.tryDevirtualize(apply: apply, isMandatory: true) != nil {
+    return
+  }
+
+  guard let callee = apply.referencedFunction else {
+    return
+  }
+
+  if !context.loadFunction(function: callee, loadCalleesRecursively: true) {
+    // We don't have the funcion body of the callee.
+    return
+  }
+
+  if shouldInline(apply: apply, callee: callee) {
+    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  {
+      simplifyCtxt.notifyInvalidatedStackNesting()
+    }
+  }
+}
+
+private func shouldInline(apply: FullApplySite, callee: Function) -> Bool {
+  if callee.isTransparent {
+    return true
+  }
+  if apply is BeginApplyInst {
+    // Avoid co-routines because they might allocate (their context).
+    return true
+  }
+  if apply.parentFunction.isGlobalInitOnceFunction && callee.inlineStrategy == .always {
+    // Some arithmetic operations, like integer conversions, are not transparent but `inline(__always)`.
+    // Force inlining them in global initializers so that it's possible to statically initialize the global.
+    return true
+  }
+  return false
+}
+
+fileprivate struct FunctionWorklist {
+  private(set) var functions = Array<Function>()
+  private var pushedFunctions = Set<Function>()
+  private var currentIndex = 0
+
+  mutating func pop() -> Function? {
+    if currentIndex < functions.count {
+      let f = functions[currentIndex]
+      currentIndex += 1
+      return f
+    }
+    return nil
+  }
+
+  mutating func addAllPerformanceAnnotatedFunctions(of moduleContext: ModulePassContext) {
+    for f in moduleContext.functions where f.performanceConstraints != .none {
+      pushIfNotVisited(f)
+    }
+  }
+
+  mutating func add(calleesOf function: Function) {
+    for inst in function.instructions {
+      switch inst {
+      case let apply as ApplySite:
+        if let callee = apply.referencedFunction {
+          pushIfNotVisited(callee)
+        }
+      case let bi as BuiltinInst:
+        switch bi.id {
+        case .Once, .OnceWithContext:
+          if let fri = bi.operands[1].value as? FunctionRefInst {
+            pushIfNotVisited(fri.referencedFunction)
+          }
+          break;
+        default:
+          break
+        }
+      default:
+        break
+      }
+    }
+  }
+
+  mutating func pushIfNotVisited(_ element: Function) {
+    if pushedFunctions.insert(element).inserted {
+      functions.append(element)
+    }
+  }
+}

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -109,6 +109,24 @@ extension MutatingContext {
     return nil
   }
 
+  func inlineFunction(apply: FullApplySite, mandatoryInline: Bool) {
+    let instAfterInling: Instruction?
+    switch apply {
+    case is ApplyInst, is BeginApplyInst:
+      instAfterInling = apply.next
+    case is TryApplyInst:
+      instAfterInling = apply.parentBlock.next?.instructions.first
+    default:
+      instAfterInling = nil
+    }
+
+    _bridged.inlineFunction(apply.bridged, mandatoryInline)
+
+    if let instAfterInling = instAfterInling {
+      notifyNewInstructions(from: apply, to: instAfterInling)
+    }
+  }
+
   /// 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)
@@ -191,12 +209,21 @@ struct FunctionPassContext : MutatingContext {
     return PostDominatorTree(bridged: bridgedPDT)
   }
 
-  func loadFunction(name: StaticString) -> Function? {
+  func loadFunction(name: StaticString, loadCalleesRecursively: Bool) -> Function? {
     return name.withUTF8Buffer { (nameBuffer: UnsafeBufferPointer<UInt8>) in
-      _bridged.loadFunction(llvm.StringRef(nameBuffer.baseAddress, nameBuffer.count)).function
+      let nameStr = llvm.StringRef(nameBuffer.baseAddress, nameBuffer.count)
+      return _bridged.loadFunction(nameStr, loadCalleesRecursively).function
     }
   }
 
+  func loadFunction(function: Function, loadCalleesRecursively: Bool) -> Bool {
+    if function.isDefinition {
+      return true
+    }
+    _bridged.loadFunction(function.bridged, loadCalleesRecursively)
+    return function.isDefinition
+  }
+
   func erase(block: BasicBlock) {
     _bridged.eraseBlock(block.bridged)
   }
@@ -210,6 +237,31 @@ struct FunctionPassContext : MutatingContext {
     _bridged.asNotificationHandler().notifyChanges(.effectsChanged)
   }
 
+  func optimizeMemoryAccesses(in function: Function) -> Bool {
+    if swift.optimizeMemoryAccesses(function.bridged.getFunction()) {
+      notifyInstructionsChanged()
+      return true
+    }
+    return false
+  }
+
+  func eliminateDeadAllocations(in function: Function) -> Bool {
+    if swift.eliminateDeadAllocations(function.bridged.getFunction()) {
+      notifyInstructionsChanged()
+      return true
+    }
+    return false
+  }
+
+  func specializeApplies(in function: Function, isMandatory: Bool) -> Bool {
+    if _bridged.specializeAppliesInFunction(function.bridged, isMandatory) {
+      notifyInstructionsChanged()
+      notifyCallsChanged()
+      return true
+    }
+    return false
+  }
+
   /// Copies `initValue` (including all operand instructions, transitively) to the
   /// static init value of `global`.
   func createStaticInitializer(for global: GlobalVariable, initValue: SingleValueInstruction) {

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -61,6 +61,7 @@ private func registerForSILCombine<InstType: SILCombineSimplifyable>(
 
 private func registerSwiftPasses() {
   // Module passes
+  registerPass(mandatoryPerformanceOptimizations, { mandatoryPerformanceOptimizations.run($0) })
   registerPass(readOnlyGlobalVariablesPass, { readOnlyGlobalVariablesPass.run($0) })
   registerPass(stackProtection, { stackProtection.run($0) })
 

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -78,7 +78,7 @@ extension Value {
   }
 }
 
-private extension Instruction {
+extension Instruction {
   var isTriviallyDead: Bool {
     if results.contains(where: { !$0.uses.isEmpty }) {
       return false

include/swift/SILOptimizer/OptimizerBridging.h

@@ -19,6 +19,7 @@
 #include "swift/SILOptimizer/Analysis/BasicCalleeAnalysis.h"
 #include "swift/SILOptimizer/Analysis/DeadEndBlocksAnalysis.h"
 #include "swift/SILOptimizer/Analysis/DominanceAnalysis.h"
+#include "swift/SILOptimizer/Utils/InstOptUtils.h"
 
 SWIFT_BEGIN_NULLABILITY_ANNOTATIONS
 
@@ -217,6 +218,8 @@ struct BridgedPassContext {
   SWIFT_IMPORT_UNSAFE
   OptionalBridgedValue constantFoldBuiltin(BridgedInstruction builtin) const;
 
+  bool specializeAppliesInFunction(BridgedFunction function, bool isMandatory) const;
+
   void createStaticInitializer(BridgedGlobalVar global, BridgedInstruction initValue) const;
 
   struct StaticInitCloneResult {
@@ -227,6 +230,8 @@ struct BridgedPassContext {
   SWIFT_IMPORT_UNSAFE
   StaticInitCloneResult copyStaticInitializer(BridgedValue initValue, BridgedBuilder b) const;
 
+  void inlineFunction(BridgedInstruction apply, bool mandatoryInline) const;
+
   SWIFT_IMPORT_UNSAFE
   BridgedValue getSILUndef(swift::SILType type) const {
     return {swift::SILUndef::get(type, *invocation->getFunction())};
@@ -390,9 +395,18 @@ struct BridgedPassContext {
   }
 
   SWIFT_IMPORT_UNSAFE
-  OptionalBridgedFunction loadFunction(llvm::StringRef name) const {
+  OptionalBridgedFunction loadFunction(llvm::StringRef name, bool loadCalleesRecursively) const {
+    swift::SILModule *mod = invocation->getPassManager()->getModule();
+    return {mod->loadFunction(name, loadCalleesRecursively ? swift::SILModule::LinkingMode::LinkAll
+                                                           : swift::SILModule::LinkingMode::LinkNormal)};
+  }
+
+  SWIFT_IMPORT_UNSAFE
+  void loadFunction(BridgedFunction function, bool loadCalleesRecursively) const {
     swift::SILModule *mod = invocation->getPassManager()->getModule();
-    return {mod->loadFunction(name, swift::SILModule::LinkingMode::LinkNormal)};
+    mod->loadFunction(function.getFunction(),
+                      loadCalleesRecursively ? swift::SILModule::LinkingMode::LinkAll
+                                             : swift::SILModule::LinkingMode::LinkNormal);
   }
 
   SWIFT_IMPORT_UNSAFE

include/swift/SILOptimizer/PassManager/Passes.def

@@ -382,6 +382,8 @@ SWIFT_FUNCTION_PASS(SILPrinter, "sil-printer",
      "Test pass which prints the SIL of a function")
 SWIFT_MODULE_PASS(FunctionUsesDumper, "dump-function-uses",
      "Dump the results of FunctionUses")
+SWIFT_MODULE_PASS(MandatoryPerformanceOptimizations, "mandatory-performance-optimizations",
+     "Performs optimizations for performance-annotated functions")
 SWIFT_MODULE_PASS(ReadOnlyGlobalVariablesPass, "read-only-global-variables",
      "Converts read-only var-globals to let-globals")
 SWIFT_MODULE_PASS(StackProtection, "stack-protection",

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -569,13 +569,16 @@ IntegerLiteralInst *optimizeBuiltinCanBeObjCClass(BuiltinInst *bi,
 /// Performs "predictable" memory access optimizations.
 ///
 /// See the PredictableMemoryAccessOptimizations pass.
-bool optimizeMemoryAccesses(SILFunction &fn);
+bool optimizeMemoryAccesses(SILFunction *fn);
 
 /// Performs "predictable" dead allocation optimizations.
 ///
 /// See the PredictableDeadAllocationElimination pass.
-bool eliminateDeadAllocations(SILFunction &fn);
+bool eliminateDeadAllocations(SILFunction *fn);
 
+bool specializeAppliesInFunction(SILFunction &F,
+                                 SILTransform *transform,
+                                 bool isMandatory);
 } // end namespace swift
 
 #endif // SWIFT_SILOPTIMIZER_UTILS_INSTOPTUTILS_H