[SilOpt] Add new layout type _TrivialStride and add pre-specialization suppport for it (swiftlang#70308)

drexin · web-flow · commit 36dd2c94505f · 2023-12-09T08:13:50.000-08:00
rdar://119329771

This layout allows adding pre-specializations for trivial types that have a different size, but the same stride. This is especially useful for collections, where the stride is the important factor.
diff --git a/docs/ABI/Mangling.rst b/docs/ABI/Mangling.rst
@@ -948,6 +948,7 @@ now codified into the ABI; the index 0 is therefore reserved.
   LAYOUT-CONSTRAINT ::= 'm' LAYOUT-SIZE  // Trivial of size at most N bits
   LAYOUT-CONSTRAINT ::= 'U'  // Unknown layout
   LAYOUT-CONSTRAINT ::= 'B' // BridgeObject
+  LAYOUT-CONSTRAINT ::= 'S' // TrivialStride
 
   LAYOUT-SIZE ::= INDEX // Size only
   LAYOUT-SIZE-AND-ALIGNMENT ::= INDEX INDEX // Size followed by alignment
diff --git a/include/swift/AST/KnownIdentifiers.def b/include/swift/AST/KnownIdentifiers.def
@@ -178,6 +178,7 @@ IDENTIFIER_WITH_NAME(NativeRefCountedObjectLayout, "_NativeRefCountedObject")
 IDENTIFIER_WITH_NAME(ClassLayout, "_Class")
 IDENTIFIER_WITH_NAME(NativeClassLayout, "_NativeClass")
 IDENTIFIER_WITH_NAME(BridgeObjectLayout, "_BridgeObject")
+IDENTIFIER_WITH_NAME(TrivialStrideLayout, "_TrivialStride")
 
 // Operators
 IDENTIFIER_WITH_NAME(MatchOperator, "~=")
diff --git a/include/swift/AST/LayoutConstraint.h b/include/swift/AST/LayoutConstraint.h
@@ -114,6 +114,8 @@ class LayoutConstraintInfo
 
   bool isBridgeObject() const { return isBridgeObject(Kind); }
 
+  bool isTrivialStride() const { return isTrivialStride(Kind); }
+
   unsigned getTrivialSizeInBytes() const {
     assert(isKnownSizeTrivial());
     return (SizeInBits + 7) / 8;
@@ -155,6 +157,11 @@ class LayoutConstraintInfo
     return 8*8;
   }
 
+  unsigned getTrivialStrideInBits() const {
+    assert(isTrivialStride());
+    return SizeInBits;
+  }
+
   operator bool() const {
     return isKnownLayout();
   }
@@ -204,6 +211,8 @@ class LayoutConstraintInfo
 
   static bool isBridgeObject(LayoutConstraintKind Kind);
 
+  static bool isTrivialStride(LayoutConstraintKind Kind);
+
   /// Uniquing for the LayoutConstraintInfo.
   void Profile(llvm::FoldingSetNodeID &ID) {
     Profile(ID, Kind, SizeInBits, Alignment);
diff --git a/include/swift/AST/LayoutConstraintKind.h b/include/swift/AST/LayoutConstraintKind.h
@@ -41,7 +41,9 @@ enum class LayoutConstraintKind : uint8_t {
   NativeRefCountedObject,
   // It is a layout constraint representing a bridge object
   BridgeObject,
-  LastLayout = BridgeObject,
+  // It is a layout constraint representing a trivial type of a known stride.
+  TrivialStride,
+  LastLayout = TrivialStride,
 };
 
 #endif
diff --git a/include/swift/Demangling/TypeDecoder.h b/include/swift/Demangling/TypeDecoder.h
@@ -430,6 +430,7 @@ void decodeRequirement(NodePointer node,
               .Case("B", LayoutConstraintKind::BridgeObject)
               .Cases("E", "e", LayoutConstraintKind::TrivialOfExactSize)
               .Cases("M", "m", LayoutConstraintKind::TrivialOfAtMostSize)
+              .Case("S", LayoutConstraintKind::TrivialStride)
               .Default(llvm::None);
 
       if (!kind)
@@ -438,7 +439,8 @@ void decodeRequirement(NodePointer node,
       BuiltLayoutConstraint layout;
 
       if (kind != LayoutConstraintKind::TrivialOfExactSize &&
-          kind != LayoutConstraintKind::TrivialOfAtMostSize) {
+          kind != LayoutConstraintKind::TrivialOfAtMostSize &&
+          kind != LayoutConstraintKind::TrivialStride) {
         layout = Builder.getLayoutConstraint(*kind);
       } else {
         auto size = child->getChild(2)->getIndex();
diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
@@ -6024,7 +6024,8 @@ LayoutConstraint LayoutConstraint::getLayoutConstraint(LayoutConstraintKind Kind
                                                       unsigned SizeInBits,
                                                       unsigned Alignment,
                                                       ASTContext &C) {
-  if (!LayoutConstraintInfo::isKnownSizeTrivial(Kind)) {
+  if (!LayoutConstraintInfo::isKnownSizeTrivial(Kind) &&
+      !LayoutConstraintInfo::isTrivialStride(Kind)) {
     assert(SizeInBits == 0);
     assert(Alignment == 0);
     return getLayoutConstraint(Kind);
diff --git a/lib/AST/ASTMangler.cpp b/lib/AST/ASTMangler.cpp
@@ -3804,6 +3804,9 @@ void ASTMangler::appendOpParamForLayoutConstraint(LayoutConstraint layout) {
   case LayoutConstraintKind::BridgeObject:
     appendOperatorParam("B");
     break;
+  case LayoutConstraintKind::TrivialStride:
+    appendOperatorParam("S", Index(layout->getTrivialSizeInBits()));
+    break;
   }
 }
 
diff --git a/lib/AST/ASTPrinter.cpp b/lib/AST/ASTPrinter.cpp
@@ -7687,6 +7687,7 @@ void LayoutConstraintInfo::print(ASTPrinter &Printer,
     return; // non-parameterized cases
   case LayoutConstraintKind::TrivialOfAtMostSize:
   case LayoutConstraintKind::TrivialOfExactSize:
+  case LayoutConstraintKind::TrivialStride:
     Printer << "(";
     Printer << SizeInBits;
     if (Alignment)
diff --git a/lib/AST/GenericSignature.cpp b/lib/AST/GenericSignature.cpp
@@ -527,6 +527,11 @@ GenericSignature GenericSignature::typeErased(ArrayRef<Type> typeErasedParams) c
         requirementsErased.push_back(
             Requirement(RequirementKind::SameType, req.getFirstType(),
                         CanType(BuiltinIntegerType::get(bitWidth, C))));
+      } else if (layout->isTrivialStride()) {
+        unsigned bitWidth = layout->getTrivialStrideInBits();
+        requirementsErased.push_back(
+            Requirement(RequirementKind::SameType, req.getFirstType(),
+                        CanType(BuiltinIntegerType::get(bitWidth, C))));
       } else {
         requirementsErased.push_back(req);
       }
diff --git a/lib/AST/LayoutConstraint.cpp b/lib/AST/LayoutConstraint.cpp
@@ -51,6 +51,10 @@ LayoutConstraint getLayoutConstraint(Identifier ID, ASTContext &Ctx) {
     return LayoutConstraint::getLayoutConstraint(
         LayoutConstraintKind::BridgeObject, Ctx);
 
+  if (ID == Ctx.Id_TrivialStrideLayout)
+    return LayoutConstraint::getLayoutConstraint(
+        LayoutConstraintKind::TrivialStride, 0, 0, Ctx);
+
   return LayoutConstraint::getLayoutConstraint(
       LayoutConstraintKind::UnknownLayout, Ctx);
 }
@@ -79,6 +83,8 @@ StringRef LayoutConstraintInfo::getName(LayoutConstraintKind Kind, bool internal
     return "_Trivial";
   case LayoutConstraintKind::BridgeObject:
     return "_BridgeObject";
+  case LayoutConstraintKind::TrivialStride:
+    return "_TrivialStride";
   }
 
   llvm_unreachable("Unhandled LayoutConstraintKind in switch.");
@@ -112,8 +118,9 @@ bool LayoutConstraintInfo::isAddressOnlyTrivial(LayoutConstraintKind Kind) {
 }
 
 bool LayoutConstraintInfo::isTrivial(LayoutConstraintKind Kind) {
-  return Kind > LayoutConstraintKind::UnknownLayout &&
-         Kind <= LayoutConstraintKind::Trivial;
+  return (Kind > LayoutConstraintKind::UnknownLayout &&
+          Kind <= LayoutConstraintKind::Trivial) ||
+         Kind == LayoutConstraintKind::TrivialStride;
 }
 
 bool LayoutConstraintInfo::isRefCountedObject(LayoutConstraintKind Kind) {
@@ -150,6 +157,10 @@ bool LayoutConstraintInfo::isBridgeObject(LayoutConstraintKind Kind) {
   return Kind == LayoutConstraintKind::BridgeObject;
 }
 
+bool LayoutConstraintInfo::isTrivialStride(LayoutConstraintKind Kind) {
+  return Kind == LayoutConstraintKind::TrivialStride;
+}
+
 SourceRange LayoutConstraintLoc::getSourceRange() const { return getLoc(); }
 
 #define MERGE_LOOKUP(lhs, rhs)                                                 \
@@ -180,55 +191,64 @@ static LayoutConstraintKind mergeTable[unsigned(E(LastLayout)) +
      E(/* TrivialOfAtMostSize */ TrivialOfAtMostSize), E(/* Trivial */ Trivial),
      E(/* Class */ Class), E(/* NativeClass */ NativeClass),
      E(/* RefCountedObject*/ RefCountedObject),
-     E(/* NativeRefCountedObject */ NativeRefCountedObject), MERGE_CONFLICT},
+     E(/* NativeRefCountedObject */ NativeRefCountedObject), MERGE_CONFLICT,
+     MERGE_CONFLICT},
 
     // Initialize the row for TrivialOfExactSize.
     {E(/* UnknownLayout */ TrivialOfExactSize),
      E(/* TrivialOfExactSize */ TrivialOfExactSize), MERGE_CONFLICT,
      E(/* Trivial */ TrivialOfExactSize), MERGE_CONFLICT, MERGE_CONFLICT,
-     MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT},
+     MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT},
 
     // Initialize the row for TrivialOfAtMostSize.
     {E(/* UnknownLayout */ TrivialOfAtMostSize), MERGE_CONFLICT,
      E(/* TrivialOfAtMostSize */ TrivialOfAtMostSize),
      E(/* Trivial */ TrivialOfAtMostSize), MERGE_CONFLICT, MERGE_CONFLICT,
-     MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT},
+     MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT},
 
     // Initialize the row for Trivial.
     {E(/* UnknownLayout */ Trivial),
      E(/* TrivialOfExactSize */ TrivialOfExactSize),
      E(/* TrivialOfAtMostSize */ TrivialOfAtMostSize), E(/* Trivial */ Trivial),
      MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT,
-     MERGE_CONFLICT},
+     MERGE_CONFLICT, MERGE_CONFLICT},
 
     // Initialize the row for Class.
     {E(/* UnknownLayout*/ Class), MERGE_CONFLICT, MERGE_CONFLICT,
      MERGE_CONFLICT, E(/* Class */ Class), E(/* NativeClass */ NativeClass),
      E(/* RefCountedObject */ Class),
-     E(/* NativeRefCountedObject */ NativeClass), MERGE_CONFLICT},
+     E(/* NativeRefCountedObject */ NativeClass), MERGE_CONFLICT,
+     MERGE_CONFLICT},
 
     // Initialize the row for NativeClass.
     {E(/* UnknownLayout */ NativeClass), MERGE_CONFLICT, MERGE_CONFLICT,
      MERGE_CONFLICT, E(/* Class */ NativeClass),
      E(/* NativeClass */ NativeClass), E(/* RefCountedObject */ NativeClass),
-     E(/* NativeRefCountedObject */ NativeClass), MERGE_CONFLICT},
+     E(/* NativeRefCountedObject */ NativeClass), MERGE_CONFLICT,
+     MERGE_CONFLICT},
 
     // Initialize the row for RefCountedObject.
     {E(/* UnknownLayout */ RefCountedObject), MERGE_CONFLICT, MERGE_CONFLICT,
      MERGE_CONFLICT, E(/* Class */ Class), E(/* NativeClass */ NativeClass),
      E(/* RefCountedObject */ RefCountedObject),
-     E(/* NativeRefCountedObject */ NativeRefCountedObject), MERGE_CONFLICT},
+     E(/* NativeRefCountedObject */ NativeRefCountedObject), MERGE_CONFLICT,
+     MERGE_CONFLICT},
 
     // Initialize the row for NativeRefCountedObject.
     {E(/* UnknownLayout */ NativeRefCountedObject), MERGE_CONFLICT,
      MERGE_CONFLICT, MERGE_CONFLICT, E(/* Class */ NativeClass),
      E(/* NativeClass */ NativeClass),
      E(/* RefCountedObject */ NativeRefCountedObject),
-     E(/* NativeRefCountedObject*/ NativeRefCountedObject), MERGE_CONFLICT},
+     E(/* NativeRefCountedObject*/ NativeRefCountedObject), MERGE_CONFLICT,
+     MERGE_CONFLICT},
 
     {E(BridgeObject), MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT,
      MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT,
-     E(/*BridgeObject*/ BridgeObject)},
+     E(/*BridgeObject*/ BridgeObject), MERGE_CONFLICT},
+
+    {E(TrivialStride), MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT,
+     MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT, MERGE_CONFLICT,
+     MERGE_CONFLICT, E(/*TrivialStride*/ TrivialStride)},
 };
 
 #undef E
@@ -324,6 +344,7 @@ LayoutConstraint::merge(LayoutConstraint Other) {
 LayoutConstraint
 LayoutConstraint::getLayoutConstraint(LayoutConstraintKind Kind) {
   assert(!LayoutConstraintInfo::isKnownSizeTrivial(Kind));
+  assert(!LayoutConstraintInfo::isTrivialStride(Kind));
   switch(Kind) {
   case LayoutConstraintKind::Trivial:
     return LayoutConstraint(&LayoutConstraintInfo::TrivialConstraintInfo);
@@ -343,6 +364,7 @@ LayoutConstraint::getLayoutConstraint(LayoutConstraintKind Kind) {
     return LayoutConstraint(&LayoutConstraintInfo::BridgeObjectConstraintInfo);
   case LayoutConstraintKind::TrivialOfAtMostSize:
   case LayoutConstraintKind::TrivialOfExactSize:
+  case LayoutConstraintKind::TrivialStride:
     llvm_unreachable("Wrong layout constraint kind");
   }
   llvm_unreachable("unhandled kind");
diff --git a/lib/Demangling/Demangler.cpp b/lib/Demangling/Demangler.cpp
@@ -4052,6 +4052,11 @@ NodePointer Demangler::demangleGenericRequirement() {
       if (!size)
         return nullptr;
       name = "m";
+    } else if (c == 'S') {
+      size = demangleIndexAsNode();
+      if (!size)
+        return nullptr;
+      name = "S";
     } else {
       // Unknown layout constraint.
       return nullptr;
diff --git a/lib/Demangling/OldDemangler.cpp b/lib/Demangling/OldDemangler.cpp
@@ -1664,6 +1664,11 @@ class OldDemangler {
         if (!demangleNatural(size, depth + 1))
           return nullptr;
         name = "m";
+      } else if (Mangled.nextIf('S')) {
+        kind = Node::Kind::Identifier;
+        if (!demangleNatural(size, depth + 1))
+          return nullptr;
+        name = "S";
       } else {
         return nullptr;
       }
diff --git a/lib/SIL/IR/AbstractionPattern.cpp b/lib/SIL/IR/AbstractionPattern.cpp
@@ -2239,6 +2239,7 @@ class SubstFunctionTypePatternVisitor
       case LayoutConstraintKind::RefCountedObject:
       case LayoutConstraintKind::TrivialOfAtMostSize:
       case LayoutConstraintKind::BridgeObject:
+      case LayoutConstraintKind::TrivialStride:
         break;
       
       case LayoutConstraintKind::UnknownLayout:
diff --git a/lib/SILOptimizer/Utils/Generics.cpp b/lib/SILOptimizer/Utils/Generics.cpp
@@ -3003,7 +3003,7 @@ bool usePrespecialized(
 
         if (!erased || !layout ||
             (!layout->isClass() && !layout->isBridgeObject() &&
-             !layout->isFixedSizeTrivial())) {
+             !layout->isFixedSizeTrivial() && !layout->isTrivialStride())) {
           newSubs.push_back(entry.value());
           continue;
         }
@@ -3036,6 +3036,23 @@ bool usePrespecialized(
                 BuiltinIntegerType::get(layout->getTrivialSizeInBits(),
                                         genericParam->getASTContext())));
           }
+        } else if (layout->isTrivialStride() && lowered.isTrivial(refF)) {
+          auto *IGM = funcBuilder.getIRGenModule();
+          auto &ti = IGM->getTypeInfo(lowered);
+          auto *typeLayout = ti.buildTypeLayoutEntry(*IGM, lowered, false);
+          auto fixedSize = typeLayout->fixedSize(*IGM);
+          if (fixedSize) {
+            auto stride = fixedSize->roundUpToAlignment(
+                *typeLayout->fixedAlignment(*IGM));
+            if (stride.isZero())
+              stride = irgen::Size(1);
+
+            if (stride.getValueInBits() == layout->getTrivialStrideInBits()) {
+              newSubs.push_back(CanType(
+                  BuiltinIntegerType::get(layout->getTrivialStrideInBits(),
+                                          genericParam->getASTContext())));
+            }
+          }
         } else {
           // no match
           break;
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp
@@ -1185,6 +1185,7 @@ getActualLayoutConstraintKind(uint64_t rawKind) {
   CASE(NativeClass)
   CASE(UnknownLayout)
   CASE(BridgeObject)
+  CASE(TrivialStride)
   }
 #undef CASE
 
@@ -1244,7 +1245,8 @@ llvm::Error ModuleFile::deserializeGenericRequirementsChecked(
       ASTContext &ctx = getContext();
       LayoutConstraint layout;
       if (kind != LayoutConstraintKind::TrivialOfAtMostSize &&
-          kind != LayoutConstraintKind::TrivialOfExactSize)
+          kind != LayoutConstraintKind::TrivialOfExactSize &&
+          kind != LayoutConstraintKind::TrivialStride)
         layout = LayoutConstraint::getLayoutConstraint(kind, ctx);
       else
         layout =
diff --git a/lib/Serialization/ModuleFormat.h b/lib/Serialization/ModuleFormat.h
@@ -58,7 +58,7 @@ const uint16_t SWIFTMODULE_VERSION_MAJOR = 0;
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
 /// Don't worry about adhering to the 80-column limit for this line.
-const uint16_t SWIFTMODULE_VERSION_MINOR = 824; // LayoutRequirementKindField
+const uint16_t SWIFTMODULE_VERSION_MINOR = 825; // TrivialStride
 
 /// A standard hash seed used for all string hashes in a serialized module.
 ///
@@ -484,6 +484,7 @@ enum LayoutRequirementKind : uint8_t {
   Class = 6,
   NativeClass = 7,
   BridgeObject = 8,
+  TrivialStride = 9,
 };
 using LayoutRequirementKindField = BCFixed<4>;
 
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
@@ -1480,6 +1480,8 @@ void Serializer::serializeGenericRequirements(
       if (layout->isKnownSizeTrivial()) {
         size = layout->getTrivialSizeInBits();
         alignment = layout->getAlignmentInBits();
+      } else if (layout->isTrivialStride()) {
+        size = layout->getTrivialStrideInBits();
       }
       LayoutRequirementKind rawKind = LayoutRequirementKind::UnknownLayout;
       switch (layout->getKind()) {
@@ -1510,6 +1512,9 @@ void Serializer::serializeGenericRequirements(
       case LayoutConstraintKind::BridgeObject:
         rawKind = LayoutRequirementKind::BridgeObject;
         break;
+      case LayoutConstraintKind::TrivialStride:
+        rawKind = LayoutRequirementKind::TrivialStride;
+        break;
       }
       scratch.push_back(rawKind);
       scratch.push_back(addTypeRef(req.getFirstType()));
diff --git a/test/SILOptimizer/Inputs/pre_specialized_module_layouts.swift b/test/SILOptimizer/Inputs/pre_specialized_module_layouts.swift
@@ -12,6 +12,7 @@ public class SomeClass {
 @_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, where @_noMetadata T : _TrivialStride(128))
 @_specialize(exported: true, availability: macOS 10.50, *; where T == SomeData)
 public func publicPrespecialized<T>(_ t: T) {
 }
diff --git a/test/SILOptimizer/pre_specialize_layouts.swift b/test/SILOptimizer/pre_specialize_layouts.swift

Original file line number	Diff line number	Diff line change
`@@ -3804,6 +3804,9 @@ void ASTMangler::appendOpParamForLayoutConstraint(LayoutConstraint layout) {`
`3804`	`3804`	`case LayoutConstraintKind::BridgeObject:`
`3805`	`3805`	`appendOperatorParam("B");`
`3806`	`3806`	`break;`
	`3807`	`+ case LayoutConstraintKind::TrivialStride:`
	`3808`	`+ appendOperatorParam("S", Index(layout->getTrivialSizeInBits()));`
	`3809`	`+ break;`
`3807`	`3810`	`}`
`3808`	`3811`	`}`
`3809`	`3812`