[sil-generic-specializer] Add @_semantics("optimize.sil.specialize.generic.partial.never") to disable partial specialization on specific functions

lib/SILOptimizer/Utils/Generics.cpp

@@ -170,7 +170,8 @@ static bool isTypeTooComplex(Type t) {
 // ReabstractionInfo
 // =============================================================================
 
-static bool shouldNotSpecializeCallee(SILFunction *Callee) {
+static bool shouldNotSpecializeCallee(SILFunction *Callee,
+                                      SubstitutionList Subs = {}) {
   if (!Callee->shouldOptimize()) {
     DEBUG(llvm::dbgs() << "    Cannot specialize function " << Callee->getName()
           << " marked to be excluded from optimizations.\n");
@@ -180,6 +181,10 @@ static bool shouldNotSpecializeCallee(SILFunction *Callee) {
   if (Callee->hasSemanticsAttr("optimize.sil.specialize.generic.never"))
     return true;
 
+  if (!Subs.empty() &&
+      Callee->hasSemanticsAttr("optimize.sil.specialize.generic.partial.never"))
+    return true;
+
   return false;
 }
 
@@ -291,6 +296,10 @@ bool ReabstractionInfo::prepareAndCheck(ApplySite Apply, SILFunction *Callee,
     // We need a generic environment for the partial specialization.
     if (!CalleeGenericEnv)
       return false;
+
+    // Bail if the callee should not be partially specialized.
+    if (shouldNotSpecializeCallee(Callee, ParamSubs))
+      return false;
   }
 
   return true;

lib/SILOptimizer/Utils/PerformanceInlinerUtils.cpp

@@ -591,6 +591,31 @@ static bool calleeHasPartialApplyWithOpenedExistentials(FullApplySite AI) {
   return false;
 }
 
+// Returns true if a given apply site should be skipped during the
+// early inlining pass.
+//
+// NOTE: Add here the checks for any specific @_semantics/@_effects
+// attributes causing a given callee to be excluded from the inlining
+// during the early inlining pass.
+static bool shouldSkipApplyDuringEarlyInlining(FullApplySite AI) {
+  // Add here the checks for any specific @_semantics attributes that need
+  // to be skipped during the early inlining pass.
+  ArraySemanticsCall ASC(AI.getInstruction());
+  if (ASC && !ASC.canInlineEarly())
+    return true;
+
+  SILFunction *Callee = AI.getReferencedFunction();
+  if (!Callee)
+    return false;
+
+  // Add here the checks for any specific @_effects attributes that need
+  // to be skipped during the early inlining pass.
+  if (Callee->hasEffectsKind())
+    return true;
+
+  return false;
+}
+
 // Returns the callee of an apply_inst if it is basically inlineable.
 SILFunction *swift::getEligibleFunction(FullApplySite AI,
                                         InlineSelection WhatToInline) {
@@ -606,8 +631,7 @@ SILFunction *swift::getEligibleFunction(FullApplySite AI,
   // attribute if the inliner is asked not to inline them.
   if (Callee->hasSemanticsAttrs() || Callee->hasEffectsKind()) {
     if (WhatToInline == InlineSelection::NoSemanticsAndGlobalInit) {
-      ArraySemanticsCall ASC(AI.getInstruction());
-      if (!ASC.canInlineEarly())
+      if (shouldSkipApplyDuringEarlyInlining(AI))
         return nullptr;
     }
     // The "availability" semantics attribute is treated like global-init.

stdlib/public/core/Integers.swift.gyb

@@ -2271,6 +2271,7 @@ ${operatorComment(x.nonMaskingOperator, True)}
 }
 
 extension FixedWidthInteger {
+  @_semantics("optimize.sil.specialize.generic.partial.never")
   public init<Other: BinaryInteger>(clamping source: Other) {
     if _slowPath(source < Self.min) {
       self = Self.min
@@ -2416,6 +2417,7 @@ extension UnsignedInteger {
 }
 
 extension UnsignedInteger where Self : FixedWidthInteger {
+  @_semantics("optimize.sil.specialize.generic.partial.never")
   @inline(__always)
   public init<T : BinaryInteger>(_ source: T) {
     // This check is potentially removable by the optimizer
@@ -2430,6 +2432,7 @@ extension UnsignedInteger where Self : FixedWidthInteger {
     self.init(extendingOrTruncating: source)
   }
 
+  @_semantics("optimize.sil.specialize.generic.partial.never")
   @inline(__always)
   public init?<T : BinaryInteger>(exactly source: T) {
     // This check is potentially removable by the optimizer
@@ -2494,6 +2497,7 @@ extension SignedInteger {
 }
 
 extension SignedInteger where Self : FixedWidthInteger {
+  @_semantics("optimize.sil.specialize.generic.partial.never")
   @inline(__always)
   public init<T : BinaryInteger>(_ source: T) {
     // This check is potentially removable by the optimizer
@@ -2510,6 +2514,7 @@ extension SignedInteger where Self : FixedWidthInteger {
     self.init(extendingOrTruncating: source)
   }
 
+  @_semantics("optimize.sil.specialize.generic.partial.never")
   @inline(__always)
   public init?<T : BinaryInteger>(exactly source: T) {
     // This check is potentially removable by the optimizer

stdlib/public/core/String.swift

@@ -211,6 +211,7 @@ extension _StringCore {
     }
   }
 
+  @_semantics("optimize.sil.specialize.generic.partial.never")
   internal func _withCSubstringAndLength<
     Result, TargetEncoding: Unicode.Encoding
   >(

test/SILOptimizer/inline_semantics.sil

@@ -12,17 +12,43 @@ bb0:
   return %1 : $Int32                               // id: %2
 }
 
+//Not every @_semantics should be skipped during the early inlining pass, but
+//only those ones which are explicitly listed in shouldSkipApplyDuringEarlyInlining.
+
 //CHECK-LABEL: caller_func
-//CHECK: function_ref
-//CHECK: apply
-//CHECK-NEXT: ret
+//CHECK-NOT: function_ref
+//CHECK-NOT: apply
+//CHECK: end sil function 'caller_func'
 sil @caller_func : $@convention(thin) () -> Int32 {
 bb0:
   %0 = function_ref @callee_func : $@convention(thin) () -> Int32 // user: %1
   %1 = apply %0() : $@convention(thin) () -> Int32                // user: %2
   return %1 : $Int32                                  // id: %2
 }
 
+sil [_semantics "array.make_mutable"] @callee_func_with_to_be_skipped_during_inlining_semantics : $@convention(method) (@inout Int32) -> Int32 {
+bb0(%self : $*Int32):
+  %0 = integer_literal $Builtin.Int32, 3           // user: %1
+  %1 = struct $Int32 (%0 : $Builtin.Int32)         // user: %2
+  return %1 : $Int32                               // id: %2
+}
+
+//Not every @_semantics should be skipped during the early inlining pass, but
+//only those ones which are explicitly listed in shouldSkipApplyDuringEarlyInlining.
+
+//CHECK-LABEL: caller_func2
+//CHECK: function_ref
+//CHECK: apply
+//CHECK: end sil function 'caller_func2'
+sil @caller_func2 : $@convention(thin) () -> Int32 {
+bb0:
+  %self = alloc_stack $Int32
+  %0 = function_ref @callee_func_with_to_be_skipped_during_inlining_semantics  : $@convention(method) (@inout Int32) -> Int32 // user: %1
+  %1 = apply %0(%self) : $@convention(method) (@inout Int32) -> Int32                 // user: %2
+  dealloc_stack %self : $*Int32
+  return %1 : $Int32                                  // id: %2
+}
+
 // A function annotated with the @effects(readonly) attribute.
 sil [readonly] @callee_func2 : $@convention(thin) () -> Int32 {
 bb0: