[SILOptimizer] Add prespecialization for arbitray reference types (#58846)

* [SILOptimizer] Add prespecialization for arbitray reference types

* Fix benchmark Package.swift

* Move SimpleArray to utils

* Fix multiple indirect result case

* Remove leftover code from previous attempt

* Fix test after rebase

* Move code to compute type replacements to SpecializedFunction

* Fix ownership when OSSA is enabled

* Fixes after rebase

* Changes after rebasing

* Add feature flag for layout pre-specialization

* Fix pre_specialize-macos.swift

* Add compiler flag to benchmark build

* Fix benchmark SwiftPM flags

benchmark/CMakeLists.txt

@@ -162,6 +162,7 @@ set(SWIFT_BENCH_MODULES
     single-source/RomanNumbers
     single-source/SIMDRandomMask
     single-source/SIMDReduceInteger
+    single-source/SimpleArraySpecialization
     single-source/SequenceAlgos
     single-source/SetTests
     single-source/SevenBoom
@@ -215,6 +216,7 @@ endif()
 
 set(BENCH_LIBRARY_MODULES
     utils/TestsUtils
+    utils/SimpleArray
 )
 
 set(BENCH_DRIVER_LIBRARY_MODULES

benchmark/Package.swift

@@ -78,6 +78,7 @@ var multiSourceLibraries: [(parentSubDir: String, name: String)] = multiSourceLi
 
 var products: [Product] = []
 products.append(.library(name: "TestsUtils", type: .static, targets: ["TestsUtils"]))
+products.append(.library(name: "SimpleArray", type: .static, targets: ["SimpleArray"]))
 products.append(.library(name: "DriverUtils", type: .static, targets: ["DriverUtils"]))
 #if os(macOS) || os(iOS) || os(watchOS) || os(tvOS)
 products.append(.library(name: "ObjectiveCTests", type: .static, targets: ["ObjectiveCTests"]))
@@ -96,6 +97,12 @@ products += multiSourceLibraries.map {
 
 var targets: [Target] = []
 targets.append(.target(name: "TestsUtils", path: "utils", sources: ["TestsUtils.swift"]))
+targets.append(.target(
+  name: "SimpleArray",
+  path: "utils",
+  sources: ["SimpleArray.swift"],
+  swiftSettings: [.unsafeFlags(["-Xfrontend",
+                                "-enable-experimental-layout-prespecialization"])]))
 targets.append(.systemLibrary(name: "LibProc", path: "utils/LibProc"))
 targets.append(
   .target(name: "DriverUtils",
@@ -129,7 +136,7 @@ targets.append(
     publicHeadersPath: "."))
 #endif
 
-var singleSourceDeps: [Target.Dependency] = [.target(name: "TestsUtils")]
+var singleSourceDeps: [Target.Dependency] = [.target(name: "TestsUtils"), .target(name: "SimpleArray")]
 #if os(macOS) || os(iOS) || os(watchOS) || os(tvOS)
 singleSourceDeps.append(.target(name: "ObjectiveCTests"))
 #endif
@@ -147,7 +154,9 @@ targets += singleSourceLibraries.map { name in
   return .target(name: name,
       dependencies: singleSourceDeps,
       path: "single-source",
-      sources: ["\(name).swift"])
+      sources: ["\(name).swift"],
+      swiftSettings: [.unsafeFlags(["-Xfrontend",
+                                    "-enable-experimental-layout-prespecialization"])])
 }
 
 targets += cxxSingleSourceLibraries.map { name in

benchmark/cmake/modules/AddSwiftBenchmarkSuite.cmake

@@ -354,7 +354,8 @@ function (swift_benchmark_compile_archopts)
   set(common_options
       "-c"
       "-target" "${target}"
-      "-${BENCH_COMPILE_ARCHOPTS_OPT}" ${PAGE_ALIGNMENT_OPTION})
+      "-${BENCH_COMPILE_ARCHOPTS_OPT}" ${PAGE_ALIGNMENT_OPTION}
+      "-Xfrontend -enable-experimental-layout-prespecialization")
 
   if(SWIFT_BENCHMARK_GENERATE_DEBUG_INFO)
     list(APPEND common_options "-g")

benchmark/single-source/SimpleArraySpecialization.swift

@@ -0,0 +1,120 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2022 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
+//
+//===----------------------------------------------------------------------===//
+
+// This benchmark tests prespecialization of a simplified array type
+
+import TestsUtils
+import SimpleArray
+
+public let benchmarks = [
+  BenchmarkInfo(
+    name: "SimpleArraySpecialization",
+    runFunction: run_SimpleArraySpecializationBenchmarks,
+    tags: [.abstraction, .runtime, .cpubench]
+  ),
+  BenchmarkInfo(
+    name: "SimpleArraySpecialization2",
+    runFunction: run_SimpleArraySpecializationBenchmarks2,
+    tags: [.abstraction, .runtime, .cpubench]
+  ),
+  BenchmarkInfo(
+    name: "SimpleArraySpecialization3",
+    runFunction: run_SimpleArraySpecializationBenchmarks3,
+    tags: [.abstraction, .runtime, .cpubench]
+  ),
+  BenchmarkInfo(
+    name: "SimpleArraySpecialization4",
+    runFunction: run_SimpleArraySpecializationBenchmarks4,
+    tags: [.abstraction, .runtime, .cpubench]
+  ),
+]
+
+let xs = SimpleArray<MyClass>(capacity: 100_000)
+
+@_silgen_name("_swift_stdlib_immortalize")
+func _stdlib_immortalize(_ obj: AnyObject)
+
+import Foundation
+
+
+public final class MyClass {
+  public var x: Int = 23
+}
+
+
+@inline(never)
+public func run_SimpleArraySpecializationBenchmarks(_ n: Int) {
+  let myObject = MyClass()
+
+  // prevent refcount overflow
+  _stdlib_immortalize(myObject)
+  
+  for _ in 0..<n {
+    for i in 0..<100_000 {
+      xs.append(myObject)
+    }
+    xs.clear()
+  }
+
+  blackHole(xs)
+}
+
+@inline(never)
+public func run_SimpleArraySpecializationBenchmarks2(_ n: Int) {
+  let myObject = MyClass()
+  
+  // prevent refcount overflow
+  _stdlib_immortalize(myObject)
+
+  for _ in 0..<n {
+    for i in 0..<100_000 {
+      xs.append2(myObject)
+    }
+    xs.clear()
+  }
+
+  blackHole(xs)
+}
+
+@inline(never)
+public func run_SimpleArraySpecializationBenchmarks3(_ n: Int) {
+  let myObject = MyClass()
+  
+  // prevent refcount overflow
+  _stdlib_immortalize(myObject)
+
+  for _ in 0..<n {
+    for i in 0..<100_000 {
+      xs.append3(myObject)
+    }
+    xs.clear()
+  }
+
+  blackHole(xs)
+}
+
+@inline(never)
+public func run_SimpleArraySpecializationBenchmarks4(_ n: Int) {
+  let myObject = MyClass()
+  
+  // prevent refcount overflow
+  _stdlib_immortalize(myObject)
+
+  for _ in 0..<n {
+    for i in 0..<100_000 {
+      xs.append4(myObject)
+    }
+    xs.clear()
+  }
+
+  blackHole(xs)
+}

benchmark/utils/SimpleArray.swift

@@ -0,0 +1,82 @@
+@usableFromInline
+@frozen
+struct Header {
+  @usableFromInline
+  let capacity: Int
+
+  @usableFromInline
+  var count: Int
+}
+
+public final class SimpleArray<T> {
+  @usableFromInline let storage: ManagedBuffer<Header, T>
+
+  @_alwaysEmitIntoClient
+  @inline(__always)
+  @inlinable
+  var count: Int {
+    get {
+      return self.storage.withUnsafeMutablePointerToHeader { return $0.pointee.count }
+    }
+    set {
+      return self.storage.withUnsafeMutablePointerToHeader { $0.pointee.count = newValue }
+    }
+  }
+
+  @_alwaysEmitIntoClient
+  @inline(__always)
+  @inlinable
+  var capacity: Int {
+    return self.storage.withUnsafeMutablePointerToHeader { return $0.pointee.capacity }
+  }
+
+  public init(capacity: Int) {
+    self.storage = .create(minimumCapacity: capacity) { _ in
+      return Header(capacity: capacity, count: 0)
+    }
+  }
+
+  @_alwaysEmitIntoClient
+  @inline(__always)
+  @inlinable
+  func append_internal(_ element: __owned T) {
+    guard count < capacity else {
+      fatalError("Array is full")
+    }
+    storage.withUnsafeMutablePointerToElements { ($0 + count).initialize(to: element) }
+    count += 1
+  }
+
+  @inline(never)
+  @_effects(notEscaping self.**)
+  @_specialize(exported: true, where @_noMetadata T : _Class)
+  public func append(_ element: __owned T) {
+    append_internal(element)
+  }
+
+  @inline(never)
+  @inlinable
+  @_effects(notEscaping self.**)
+  public func append2(_ element: __owned T) {
+    append_internal(element)
+  }
+
+  @inline(__always)
+  @inlinable
+  @_effects(notEscaping self.**)
+  public func append3(_ element: __owned T) {
+    append_internal(element)
+  }
+
+  @inline(never)
+  @_effects(notEscaping self.**)
+  public func append4(_ element: __owned T) {
+    append_internal(element)
+  }
+
+  public func clear() {
+    // only reset the count to avoid measuring deinitialization
+    // overhead in benchmark
+    self.count = 0
+  }
+}

benchmark/utils/main.swift

@@ -165,6 +165,7 @@ import SequenceAlgos
 import SetTests
 import SevenBoom
 import Sim2DArray
+import SimpleArraySpecialization
 import SortArrayInClass
 import SortIntPyramids
 import SortLargeExistentials
@@ -348,6 +349,7 @@ register(SequenceAlgos.benchmarks)
 register(SetTests.benchmarks)
 register(SevenBoom.benchmarks)
 register(Sim2DArray.benchmarks)
+register(SimpleArraySpecialization.benchmarks)
 register(SortArrayInClass.benchmarks)
 register(SortIntPyramids.benchmarks)
 register(SortLargeExistentials.benchmarks)