[SILOpt] Allow pre-specializations for _Trivial of known size (#70256)

* [SILOpt] Allow pre-specializations for _Trivial of known size

rdar://119224542

This allows pre-specializations to be generated and applied for trivial types of a shared size.

Author: drexin

include/swift/SILOptimizer/PassManager/PassManager.h

@@ -73,12 +73,6 @@ class SwiftPassInvocation {
 
   SILSSAUpdater *ssaUpdater = nullptr;
 
-  /// IRGen module for passes that request it (e.g. simplification pass)
-  irgen::IRGenModule *irgenModule = nullptr;
-
-  /// IRGenerator used by IRGenModule above
-  irgen::IRGenerator *irgen = nullptr;
-
   static constexpr int BlockSetCapacity = 8;
   char blockSetStorage[sizeof(BasicBlockSet) * BlockSetCapacity];
   bool aliveBlockSets[BlockSetCapacity];
@@ -184,6 +178,7 @@ class SILPassManager {
 
   /// An optional IRGenModule associated with this PassManager.
   irgen::IRGenModule *IRMod;
+  irgen::IRGenerator *irgen;
 
   /// The list of transformations to run.
   llvm::SmallVector<SILTransform *, 16> Transformations;
@@ -290,7 +285,7 @@ class SILPassManager {
 
   /// \returns the associated IGenModule or null if this is not an IRGen
   /// pass manager.
-  irgen::IRGenModule *getIRGenModule() { return IRMod; }
+  irgen::IRGenModule *getIRGenModule();
 
   SwiftPassInvocation *getSwiftPassInvocation() {
     return &swiftPassInvocation;

include/swift/SILOptimizer/Utils/SILOptFunctionBuilder.h

@@ -56,6 +56,9 @@ class SILOptFunctionBuilder {
   }
 
   SILModule &getModule() const { return *getPassManager().getModule(); }
+  irgen::IRGenModule *getIRGenModule() const {
+    return transform.getIRGenModule();
+  }
 
 private:
   SILPassManager &getPassManager() const {

lib/AST/GenericSignature.cpp

@@ -503,6 +503,7 @@ CanType GenericSignature::getReducedType(Type type) const {
 GenericSignature GenericSignature::typeErased(ArrayRef<Type> typeErasedParams) const {
   bool changedSignature = false;
   llvm::SmallVector<Requirement, 2> requirementsErased;
+  auto &C = Ptr->getASTContext();
 
   for (auto req : getRequirements()) {
     bool found = std::any_of(typeErasedParams.begin(),
@@ -512,14 +513,20 @@ GenericSignature GenericSignature::typeErased(ArrayRef<Type> typeErasedParams) c
       return t->isEqual(other);
     });
     if (found && req.getKind() == RequirementKind::Layout) {
-      if (req.getLayoutConstraint()->isClass()) {
+      auto layout = req.getLayoutConstraint();
+      if (layout->isClass()) {
+        requirementsErased.push_back(Requirement(RequirementKind::SameType,
+                                                 req.getFirstType(),
+                                                 C.getAnyObjectType()));
+      } else if (layout->isBridgeObject()) {
+        requirementsErased.push_back(Requirement(RequirementKind::SameType,
+                                                 req.getFirstType(),
+                                                 C.TheBridgeObjectType));
+      } else if (layout->isFixedSizeTrivial()) {
+        unsigned bitWidth = layout->getTrivialSizeInBits();
         requirementsErased.push_back(
             Requirement(RequirementKind::SameType, req.getFirstType(),
-                        Ptr->getASTContext().getAnyObjectType()));
-      } else if (req.getLayoutConstraint()->isBridgeObject()) {
-        requirementsErased.push_back(
-            Requirement(RequirementKind::SameType, req.getFirstType(),
-                        Ptr->getASTContext().TheBridgeObjectType));
+                        CanType(BuiltinIntegerType::get(bitWidth, C))));
       } else {
         requirementsErased.push_back(req);
       }

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -374,7 +374,7 @@ void swift::executePassPipelinePlan(SILModule *SM,
 
 SILPassManager::SILPassManager(SILModule *M, bool isMandatory,
                                irgen::IRGenModule *IRMod)
-    : Mod(M), IRMod(IRMod),
+    : Mod(M), IRMod(IRMod), irgen(nullptr),
       swiftPassInvocation(this),
       isMandatory(isMandatory), deserializationNotificationHandler(nullptr) {
 #define SIL_ANALYSIS(NAME) \
@@ -901,6 +901,27 @@ void SILPassManager::execute() {
   }
 }
 
+irgen::IRGenModule *SILPassManager::getIRGenModule() {
+  // We need an IRGenModule to get the actual sizes from type lowering.
+  // Creating an IRGenModule involves some effort, let's cache it for the
+  // whole pass.
+  if (IRMod == nullptr) {
+    SILModule *module = getModule();
+
+    auto *irgenOpts = module->getIRGenOptionsOrNull();
+    if (!irgenOpts)
+      return nullptr;
+
+    if (irgen == nullptr)
+      irgen = new irgen::IRGenerator(*irgenOpts, *module);
+    auto targetMachine = irgen->createTargetMachine();
+    assert(targetMachine && "failed to create target");
+    IRMod = new irgen::IRGenModule(*irgen, std::move(targetMachine));
+  }
+
+  return IRMod;
+}
+
 /// D'tor.
 SILPassManager::~SILPassManager() {
 
@@ -940,6 +961,16 @@ SILPassManager::~SILPassManager() {
            "Deleting a locked analysis. Did we forget to unlock ?");
     delete A;
   }
+
+  if (irgen) {
+    // If irgen is set, we also own the IRGenModule
+    if (IRMod) {
+      delete IRMod;
+      IRMod = nullptr;
+    }
+    delete irgen;
+    irgen = nullptr;
+  }
 }
 
 void SILPassManager::notifyOfNewFunction(SILFunction *F, SILTransform *T) {
@@ -1382,22 +1413,7 @@ void SwiftPassInvocation::finishedInstructionPassRun() {
 }
 
 irgen::IRGenModule *SwiftPassInvocation::getIRGenModule() {
-  // We need an IRGenModule to get the actual sizes from type lowering.
-  // Creating an IRGenModule involves some effort, let's cache it for the
-  // whole pass.
-  if (irgenModule == nullptr && irgen == nullptr) {
-    SILModule *module = getPassManager()->getModule();
-
-    auto *irgenOpts = module->getIRGenOptionsOrNull();
-    if (!irgenOpts)
-      return nullptr;
-
-    irgen = new irgen::IRGenerator(*irgenOpts, *module);
-    auto targetMachine = irgen->createTargetMachine();
-    assert(targetMachine && "failed to create target");
-    irgenModule = new irgen::IRGenModule(*irgen, std::move(targetMachine));
-  }
-  return irgenModule;
+  return passManager->getIRGenModule();
 }
 
 void SwiftPassInvocation::endPass() {
@@ -1430,16 +1446,7 @@ void SwiftPassInvocation::endTransformFunction() {
   assert(numNodeSetsAllocated == 0 && "Not all NodeSets deallocated");
 }
 
-SwiftPassInvocation::~SwiftPassInvocation() {
-  if (irgenModule) {
-    delete irgenModule;
-    irgenModule = nullptr;
-  }
-  if (irgen) {
-    delete irgen;
-    irgen = nullptr;
-  }
-}
+SwiftPassInvocation::~SwiftPassInvocation() {}
 
 //===----------------------------------------------------------------------===//
 //                           OptimizerBridging

lib/SILOptimizer/Utils/Generics.cpp

@@ -13,6 +13,7 @@
 #define DEBUG_TYPE "generic-specializer"
 
 #include "swift/SILOptimizer/Utils/Generics.h"
+#include "../../IRGen/IRGenModule.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/DiagnosticEngine.h"
 #include "swift/AST/DiagnosticsSIL.h"
@@ -3001,7 +3002,8 @@ bool usePrespecialized(
         }
 
         if (!erased || !layout ||
-            (!layout->isClass() && !layout->isBridgeObject())) {
+            (!layout->isClass() && !layout->isBridgeObject() &&
+             !layout->isFixedSizeTrivial())) {
           newSubs.push_back(entry.value());
           continue;
         }
@@ -3022,6 +3024,18 @@ bool usePrespecialized(
           } else {
             newSubs.push_back(genericParam->getASTContext().getAnyObjectType());
           }
+        } else if (layout->isFixedSizeTrivial() && lowered.isTrivial(refF)) {
+          auto *IGM = funcBuilder.getIRGenModule();
+          auto &ti = IGM->getTypeInfo(lowered);
+          auto fixedSize =
+              ti.buildTypeLayoutEntry(*IGM, lowered, false)->fixedSize(*IGM);
+
+          if (fixedSize &&
+              fixedSize->getValueInBits() == layout->getTrivialSizeInBits()) {
+            newSubs.push_back(CanType(
+                BuiltinIntegerType::get(layout->getTrivialSizeInBits(),
+                                        genericParam->getASTContext())));
+          }
         } else {
           // no match
           break;

test/SILOptimizer/Inputs/pre_specialized_module_layouts.swift

@@ -11,6 +11,7 @@ public class SomeClass {
 @_specialize(exported: true, where T == Double)
 @_specialize(exported: true, where @_noMetadata T : _Class)
 @_specialize(exported: true, where @_noMetadata T : _BridgeObject)
+@_specialize(exported: true, where @_noMetadata T : _Trivial(64))
 @_specialize(exported: true, availability: macOS 10.50, *; where T == SomeData)
 public func publicPrespecialized<T>(_ t: T) {
 }

test/SILOptimizer/pre_specialize_layouts.swift

@@ -31,6 +31,11 @@ public struct OveralignedReferenceWrapperStruct {
   let x: AnyObject
 }
 
+public struct TwoInt32 {
+  let x: Int32 = 0
+  let y: Int32 = 0
+}
+
 // Make sure we generate the public pre-specialized entry points.
 
 // OPT-DAG: sil @$s22pre_specialize_layouts10testPublic1tyx_tlFSf_Ts5 : $@convention(thin) (Float) -> () {
@@ -79,6 +84,11 @@ internal func testEmitIntoClient<T>(t: T) {
 // OPT:   apply [[F1]]
 // OPT:   [[F2:%.*]] = function_ref @$s30pre_specialized_module_layouts20publicPrespecializedyyxlFSd_Ts5 : $@convention(thin) (Double) -> ()
 // OPT:   apply [[F2]]
+// OPT:   [[F9:%.*]] = function_ref @$s30pre_specialized_module_layouts20publicPrespecializedyyxlFBi64__Ts5 : $@convention(thin) (Builtin.Int64) -> ()
+// OPT:   [[A5:%.*]] = unchecked_trivial_bit_cast {{%.*}} : $UInt64 to $Builtin.Int64
+// OPT:   apply [[F9]]([[A5]]) : $@convention(thin) (Builtin.Int64) -> ()
+// OPT:   [[A6:%.*]] = unchecked_trivial_bit_cast {{%.*}} : $TwoInt32 to $Builtin.Int64
+// OPT:   apply [[F9]]([[A6]]) : $@convention(thin) (Builtin.Int64) -> ()
 // OPT-macosx:   [[F6:%.*]] = function_ref @$s30pre_specialized_module_layouts20publicPrespecializedyyxlF : $@convention(thin) <τ_0_0> (@in_guaranteed τ_0_0) -> ()
 // OPT-macosx:    apply [[F6]]<SomeData>
 // OPT: [[F7:%.*]] = function_ref @$s30pre_specialized_module_layouts20publicPrespecializedyyxlFyXl_Ts5 : $@convention(thin) (@guaranteed AnyObject) -> ()
@@ -145,6 +155,8 @@ internal func testEmitIntoClient<T>(t: T) {
 public func usePrespecializedEntryPoints(wrapperStruct: ReferenceWrapperStruct, overaligned: OveralignedReferenceWrapperStruct, array: [Int]) {
   publicPrespecialized(1)
   publicPrespecialized(1.0)
+  publicPrespecialized(UInt64(1))
+  publicPrespecialized(TwoInt32())
   publicPrespecialized(SomeData())
   publicPrespecialized(SomeClass())
   publicPrespecialized(wrapperStruct)