[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

Author: drexin

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)

include/swift/AST/Attr.h

@@ -1396,7 +1396,7 @@ class SynthesizedProtocolAttr : public DeclAttribute {
 class SpecializeAttr final
     : public DeclAttribute,
       private llvm::TrailingObjects<SpecializeAttr, Identifier,
-                                    AvailableAttr *> {
+                                    AvailableAttr *, Type> {
   friend class SpecializeAttrTargetDeclRequest;
   friend TrailingObjects;
 
@@ -1416,19 +1416,23 @@ class SpecializeAttr final
   uint64_t resolverContextData;
   size_t numSPIGroups;
   size_t numAvailableAttrs;
+  size_t numTypeErasedParams;
+  bool typeErasedParamsInitialized;
 
   SpecializeAttr(SourceLoc atLoc, SourceRange Range,
                  TrailingWhereClause *clause, bool exported,
                  SpecializationKind kind, GenericSignature specializedSignature,
                  DeclNameRef targetFunctionName, ArrayRef<Identifier> spiGroups,
-                 ArrayRef<AvailableAttr *> availabilityAttrs);
+                 ArrayRef<AvailableAttr *> availabilityAttrs,
+                 size_t typeErasedParamsCount);
 
 public:
   static SpecializeAttr *
   create(ASTContext &Ctx, SourceLoc atLoc, SourceRange Range,
          TrailingWhereClause *clause, bool exported, SpecializationKind kind,
          DeclNameRef targetFunctionName, ArrayRef<Identifier> spiGroups,
          ArrayRef<AvailableAttr *> availabilityAttrs,
+         size_t typeErasedParamsCount,
          GenericSignature specializedSignature = nullptr);
 
   static SpecializeAttr *create(ASTContext &ctx, bool exported,
@@ -1442,6 +1446,7 @@ class SpecializeAttr final
                                 SpecializationKind kind,
                                 ArrayRef<Identifier> spiGroups,
                                 ArrayRef<AvailableAttr *> availabilityAttrs,
+                                ArrayRef<Type> typeErasedParams,
                                 GenericSignature specializedSignature,
                                 DeclNameRef replacedFunction,
                                 LazyMemberLoader *resolver, uint64_t data);
@@ -1467,6 +1472,22 @@ class SpecializeAttr final
             numAvailableAttrs};
   }
 
+  ArrayRef<Type> getTypeErasedParams() const {
+    if (!typeErasedParamsInitialized)
+      return {};
+
+    return {this->template getTrailingObjects<Type>(),
+            numTypeErasedParams};
+  }
+
+  void setTypeErasedParams(const ArrayRef<Type> typeErasedParams) {
+    assert(typeErasedParams.size() == numTypeErasedParams);
+    if (!typeErasedParamsInitialized) {
+      std::uninitialized_copy(typeErasedParams.begin(), typeErasedParams.end(), getTrailingObjects<Type>());
+      typeErasedParamsInitialized = true;
+    }
+  }
+
   TrailingWhereClause *getTrailingWhereClause() const;
 
   GenericSignature getSpecializedSignature() const {
@@ -2183,6 +2204,19 @@ class TypeSequenceAttr : public DeclAttribute {
   }
 };
 
+// /// The @_noMetadata attribute, which marks a generic param as not providing metadata
+// class NoMetadataAttr : public DeclAttribute {
+//   NoMetadataAttr(SourceLoc atLoc, SourceRange Range);
+
+// public:
+//   static NoMetadataAttr *create(ASTContext &Ctx, SourceLoc atLoc,
+//                                 SourceRange Range);
+
+//   static bool classof(const DeclAttribute *DA) {
+//     return DA->getKind() == DAK_NoMetadata;
+//   }
+// };
+
 /// The @_unavailableFromAsync attribute, used to make function declarations
 /// unavailable from async contexts.
 class UnavailableFromAsyncAttr : public DeclAttribute {

include/swift/AST/DiagnosticsSema.def

@@ -6601,5 +6601,18 @@ ERROR(has_symbol_invalid_expr,none,
 ERROR(has_symbol_unsupported_in_closures,none,
       "#_hasSymbol is not supported in closures", ())
 
+//------------------------------------------------------------------------------
+// MARK: Type erasure
+//------------------------------------------------------------------------------
+
+ERROR(no_metadata_on_non_generic_param, none,
+"'@_noMetadata' must appear on a generic parameter",
+())
+ERROR(experimental_no_metadata_feature_can_only_be_used_when_enabled,
+      none, "Can not use feature when experimental layout based"
+      " pre-specializations are disabled! Pass the frontend flag"
+      " -enable-experimental-layout-prespecialization to swift to enable "
+      " the usage of this language feature", ())
+
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"