LargeTypesReg2Mem: Avoid building an explosion schema just to count the number of registers for very large types

include/swift/SIL/SILTypeProperties.h

@@ -128,6 +128,13 @@ enum MayHaveCustomDeinit_t : bool {
   MayHaveCustomDeinit = true,
 };
 
+/// Is this a very large type whose expansion into SSA values is likely
+/// detrimential i.e should be kept indirect.
+enum IsVeryLargeType_t : bool {
+  IsNotVeryLargeType = false,
+  IsVeryLargeType = true
+};
+
 class SILTypeProperties {
   // These are chosen so that bitwise-or merges the flags properly.
   //
@@ -145,6 +152,7 @@ class SILTypeProperties {
     AddressableForDependenciesFlag = 1 << 9,
     HasRawLayoutFlag               = 1 << 10,
     CustomDeinitFlag               = 1 << 11,
+    IsVeryLargeTypeFlag            = 1 << 12,
   };
   // clang-format on
 
@@ -164,7 +172,8 @@ class SILTypeProperties {
       IsLexical_t isLexical = IsNotLexical, HasPack_t hasPack = HasNoPack,
       IsAddressableForDependencies_t isAFD = IsNotAddressableForDependencies,
       HasRawLayout_t hasRawLayout = DoesNotHaveRawLayout,
-      MayHaveCustomDeinit_t customDeinit = HasOnlyDefaultDeinit)
+      MayHaveCustomDeinit_t customDeinit = HasOnlyDefaultDeinit,
+      IsVeryLargeType_t largeType = IsNotVeryLargeType )
       : Flags((isTrivial ? 0U : NonTrivialFlag) |
               (isFixedABI ? 0U : NonFixedABIFlag) |
               (isAddressOnly ? AddressOnlyFlag : 0U) |
@@ -175,7 +184,8 @@ class SILTypeProperties {
               (hasPack ? HasPackFlag : 0U) |
               (isAFD ? AddressableForDependenciesFlag : 0U) |
               (hasRawLayout ? HasRawLayoutFlag : 0U) |
-              (customDeinit ? CustomDeinitFlag : 0U)) {}
+              (customDeinit ? CustomDeinitFlag : 0U) |
+              (largeType ? IsVeryLargeTypeFlag : 0U)) {}
 
   constexpr bool operator==(SILTypeProperties p) const {
     return Flags == p.Flags;
@@ -261,6 +271,9 @@ class SILTypeProperties {
   MayHaveCustomDeinit_t mayHaveCustomDeinit() const {
     return MayHaveCustomDeinit_t((Flags & CustomDeinitFlag) != 0);
   }
+  IsVeryLargeType_t isVeryLargeType() const {
+    return IsVeryLargeType_t((Flags & IsVeryLargeTypeFlag) != 0);
+  }
 
   void setNonTrivial() { Flags |= NonTrivialFlag; }
   void setIsOrContainsRawPointer() { Flags |= HasRawPointerFlag; }
@@ -287,6 +300,7 @@ class SILTypeProperties {
     Flags = (Flags & ~CustomDeinitFlag)
       | (hasCustomDeinit ? CustomDeinitFlag : 0);
   }
+  void setVeryLargeType() { Flags |= IsVeryLargeTypeFlag; }
 };
 
 } // end namespace swift

lib/IRGen/FixedTypeInfo.h

@@ -66,7 +66,8 @@ class FixedTypeInfo : public TypeInfo {
                 SpecialTypeInfoKind stik = SpecialTypeInfoKind::Fixed)
       : TypeInfo(type, align, pod, bt, copy, alwaysFixedSize, isABIAccessible, stik),
         SpareBits(std::move(spareBits)) {
-    assert(SpareBits.size() == size.getValueInBits());
+    // SpareBits implementation is limited to 32bits.
+    assert(SpareBits.size() == (size.getValueInBits() & 0xFFFFFFFF));
     assert(isFixedSize());
     Bits.FixedTypeInfo.Size = size.getValue();
     assert(Bits.FixedTypeInfo.Size == size.getValue() && "truncation");

lib/IRGen/LoadableByAddress.cpp

@@ -3442,8 +3442,8 @@ struct Properties {
   uint16_t numRegisters = 0;
 
 public:
-  static uint16_t MaxNumUses;
-  static uint16_t MaxNumRegisters;
+  static const uint16_t MaxNumUses = 65535;
+  static const uint16_t MaxNumRegisters = 65535;
 
   void addConstructor() {
     sawConstructor = true;
@@ -3504,26 +3504,31 @@ struct Properties {
 };
 }
 
-uint16_t Properties::MaxNumUses = 65535;
-uint16_t Properties::MaxNumRegisters = 65535;
-
 namespace {
 class LargeLoadableHeuristic {
   GenericEnvironment *genEnv;
   IRGenModule *irgenModule;
+  SILFunction &fun;
 
   llvm::DenseMap<SILType, Properties> largeTypeProperties;
 
+  /// Cutoff used to determine when to deem a type as large solely on the number
+  /// of registers.
   static const unsigned NumRegistersVeryLargeType = 16;
+  /// Cutoff when we start considering other properties to determine when we
+  /// deem a type as large.
   static const unsigned NumRegistersLargeType = 8;
-  static const unsigned numUsesThreshold = 8;
+  /// Cutoff when we start to deem a type as large based on the number of uses
+  /// (requires the number or registers to be at least 8).
+  static const unsigned NumUsesThreshold = 8;
 
   bool UseAggressiveHeuristic;
 
 public:
    LargeLoadableHeuristic(IRGenModule *irgenModule, GenericEnvironment *genEnv,
+                          SILFunction &f,
                           bool UseAggressiveHeuristic)
-      : genEnv(genEnv), irgenModule(irgenModule),
+      : genEnv(genEnv), irgenModule(irgenModule), fun(f),
       UseAggressiveHeuristic(UseAggressiveHeuristic) {}
 
    void visit(SILInstruction *i);
@@ -3558,6 +3563,14 @@ class LargeLoadableHeuristic {
       if (cached)
         return cached;
 
+      // Avoid computing an explosion schema for very large types as this can be
+      // very compile time intensive.
+      if (fun.getTypeProperties(ty).isVeryLargeType()) {
+        entry.setAsVeryLargeType();
+        largeTypeProperties[ty] = entry;
+        return entry.getNumRegisters();
+      }
+
       const TypeInfo &TI = irgenModule->getTypeInfoForLowered(canType);
       auto &nativeSchemaOrigParam = TI.nativeParameterValueSchema(*irgenModule);
       // Passed compactly in registers but kept in many explosions.
@@ -3603,6 +3616,9 @@ void LargeLoadableHeuristic::propagate(PostOrderFunctionInfo &po) {
 }
 
 void LargeLoadableHeuristic::visit(SILArgument *arg) {
+  if (arg->getType().isAddress())
+    return;
+
   auto objType = arg->getType().getObjectType();
   if (numRegisters(objType) < NumRegistersLargeType)
     return;
@@ -3651,6 +3667,8 @@ void LargeLoadableHeuristic::visit(SILInstruction *i) {
   for (const auto &opd : i->getAllOperands()) {
     auto opdTy = opd.get()->getType();
     auto objType = opdTy.getObjectType();
+    bool isAddress = opdTy.isAddress();
+
     auto registerCount = numRegisters(objType);
     if (registerCount < NumRegistersLargeType)
       continue;
@@ -3699,10 +3717,12 @@ void LargeLoadableHeuristic::visit(SILInstruction *i) {
     case SILInstructionKind::PartialApplyInst:
     case SILInstructionKind::TryApplyInst:
     case SILInstructionKind::BeginApplyInst:
-      entry.addFunctionUseOrReturn();
+      if (!isAddress)
+        entry.addFunctionUseOrReturn();
       break;
     default:
-      entry.addUse();
+      if (!isAddress)
+        entry.addUse();
       break;
     }
 
@@ -3778,7 +3798,7 @@ bool LargeLoadableHeuristic::isLargeLoadableType(SILType ty) {
   if (regs >= NumRegistersVeryLargeType)
     return true;
 
-  if (entry.numUses() > numUsesThreshold)
+  if (entry.numUses() > NumUsesThreshold)
    return true;
 
   if (entry.hasStore())
@@ -4735,7 +4755,7 @@ static void runPeepholesAndReg2Mem(SILPassManager *pm, SILModule *silMod,
     // copy_addr peepholes
     bool UseAggressiveHeuristic =
       silMod->getOptions().UseAggressiveReg2MemForCodeSize;
-    LargeLoadableHeuristic heuristic(irgenModule, genEnv,
+    LargeLoadableHeuristic heuristic(irgenModule, genEnv, currF,
                                      UseAggressiveHeuristic);
     Peepholes opts(pm, silMod, irgenModule);
     for (SILBasicBlock &BB : currF) {

lib/SIL/IR/TypeLowering.cpp

@@ -404,6 +404,7 @@ namespace {
           && (!fixedSize.has_value()
               || *fixedSize > BuiltinFixedArrayType::MaximumLoadableSize)) {
         props.setAddressOnly();
+        props.setVeryLargeType();
       }
       return props;
     }