PMO: Don't block pmo for large types - rather only block expansion of tuples

Author: aschwaighofer

include/swift/SIL/InstructionUtils.h

@@ -236,6 +236,12 @@ bool visitExplodedTupleValue(SILValue value,
 std::pair<SILFunction *, SILWitnessTable *>
 lookUpFunctionInWitnessTable(WitnessMethodInst *wmi, SILModule::LinkingMode linkingMode);
 
+/// True if a type can be expanded without a significant increase to code size.
+///
+/// False if expanding a type is invalid. For example, expanding a
+/// struct-with-deinit drops the deinit.
+bool shouldExpand(SILModule &module, SILType ty);
+
 } // end namespace swift
 
 #endif

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -417,12 +417,6 @@ ignore_expect_uses(ValueBase *value) {
 /// operations from it. These can be simplified and removed.
 bool simplifyUsers(SingleValueInstruction *inst);
 
-/// True if a type can be expanded without a significant increase to code size.
-///
-/// False if expanding a type is invalid. For example, expanding a
-/// struct-with-deinit drops the deinit.
-bool shouldExpand(SILModule &module, SILType ty);
-
 /// Check if the value of value is computed by means of a simple initialization.
 /// Store the actual SILValue into \p Val and the reversed list of instructions
 /// initializing it in \p Insns.

lib/SIL/Utils/InstructionUtils.cpp

@@ -27,9 +27,14 @@
 #include "swift/SIL/SILVisitor.h"
 
 #include "clang/AST/DeclObjC.h"
+#include "llvm/Support/CommandLine.h"
 
 using namespace swift;
 
+static llvm::cl::opt<bool> EnableExpandAll("enable-expand-all",
+                                           llvm::cl::init(false));
+
+
 SILValue swift::lookThroughOwnershipInsts(SILValue v) {
   while (true) {
     switch (v->getKind()) {
@@ -1411,3 +1416,29 @@ swift::lookUpFunctionInWitnessTable(WitnessMethodInst *wmi,
   return mod.lookUpFunctionInWitnessTable(wmi->getConformance(), wmi->getMember(),
                                           wmi->isSpecialized(), linkingMode);
 }
+
+// True if a type can be expanded without a significant increase to code size.
+//
+// False if expanding a type is invalid. For example, expanding a
+// struct-with-deinit drops the deinit.
+bool swift::shouldExpand(SILModule &module, SILType ty) {
+  // FIXME: Expansion
+  auto expansion = TypeExpansionContext::minimal();
+
+  if (module.Types.getTypeLowering(ty, expansion).isAddressOnly()) {
+    return false;
+  }
+  // A move-only-with-deinit type cannot be SROA.
+  //
+  // TODO: we could loosen this requirement if all paths lead to a drop_deinit.
+  if (auto *nominalTy = ty.getNominalOrBoundGenericNominal()) {
+    if (nominalTy->getValueTypeDestructor())
+      return false;
+  }
+  if (EnableExpandAll) {
+    return true;
+  }
+
+  unsigned numFields = module.Types.countNumberOfFields(ty, expansion);
+  return (numFields <= 6);
+}

lib/SILOptimizer/Mandatory/PMOMemoryUseCollector.cpp

@@ -286,9 +286,14 @@ bool ElementUseCollector::collectUses(SILValue Pointer) {
       continue;
     }
 
+    auto &mod = User->getFunction()->getModule();
+    bool shouldScalarizeTuple =
+      !mod.getOptions().UseAggressiveReg2MemForCodeSize ||
+      shouldExpand(mod, PointeeType);
+
     // Loads are a use of the value.
     if (isa<LoadInst>(User) || isa<LoadBorrowInst>(User)) {
-      if (PointeeType.is<TupleType>())
+      if (PointeeType.is<TupleType>() && shouldScalarizeTuple)
         UsesToScalarize.push_back(User);
       else
         Uses.emplace_back(User, PMOUseKind::Load);
@@ -300,7 +305,8 @@ bool ElementUseCollector::collectUses(SILValue Pointer) {
       if (UI->getOperandNumber() == StoreInst::Dest) {
         if (auto tupleType = PointeeType.getAs<TupleType>()) {
           if (!tupleType->isEqual(Module.getASTContext().TheEmptyTupleType) &&
-              !tupleType->containsPackExpansionType()) {
+              !tupleType->containsPackExpansionType() &&
+              shouldScalarizeTuple) {
             UsesToScalarize.push_back(User);
             continue;
           }
@@ -337,7 +343,8 @@ bool ElementUseCollector::collectUses(SILValue Pointer) {
       // have an access that crosses elements.
       if (auto tupleType = PointeeType.getAs<TupleType>()) {
         if (!tupleType->isEqual(Module.getASTContext().TheEmptyTupleType) &&
-            !tupleType->containsPackExpansionType()) {
+            !tupleType->containsPackExpansionType() &&
+            shouldScalarizeTuple) {
           UsesToScalarize.push_back(CAI);
           continue;
         }

lib/SILOptimizer/Mandatory/PredictableMemOpt.cpp

@@ -38,6 +38,9 @@
 
 using namespace swift;
 
+static llvm::cl::opt<bool> EnableAggressiveExpansionBlocking(
+  "enable-aggressive-expansion-blocking", llvm::cl::init(false));
+
 STATISTIC(NumLoadPromoted, "Number of loads promoted");
 STATISTIC(NumLoadTakePromoted, "Number of load takes promoted");
 STATISTIC(NumDestroyAddrPromoted, "Number of destroy_addrs promoted");
@@ -3156,7 +3159,8 @@ static AllocationInst *getOptimizableAllocation(SILInstruction *i) {
 
   // Don't promote large types.
   auto &mod = alloc->getFunction()->getModule();
-  if (mod.getOptions().UseAggressiveReg2MemForCodeSize &&
+  if (EnableAggressiveExpansionBlocking &&
+      mod.getOptions().UseAggressiveReg2MemForCodeSize &&
       !shouldExpand(mod, alloc->getType().getObjectType()))
     return nullptr;
 

lib/SILOptimizer/Utils/InstOptUtils.cpp

@@ -49,9 +49,6 @@
 
 using namespace swift;
 
-static llvm::cl::opt<bool> EnableExpandAll("enable-expand-all",
-                                           llvm::cl::init(false));
-
 static llvm::cl::opt<bool> KeepWillThrowCall(
     "keep-will-throw-call", llvm::cl::init(false),
     llvm::cl::desc(
@@ -1269,32 +1266,6 @@ bool swift::simplifyUsers(SingleValueInstruction *inst) {
   return changed;
 }
 
-// True if a type can be expanded without a significant increase to code size.
-//
-// False if expanding a type is invalid. For example, expanding a
-// struct-with-deinit drops the deinit.
-bool swift::shouldExpand(SILModule &module, SILType ty) {
-  // FIXME: Expansion
-  auto expansion = TypeExpansionContext::minimal();
-
-  if (module.Types.getTypeLowering(ty, expansion).isAddressOnly()) {
-    return false;
-  }
-  // A move-only-with-deinit type cannot be SROA.
-  //
-  // TODO: we could loosen this requirement if all paths lead to a drop_deinit.
-  if (auto *nominalTy = ty.getNominalOrBoundGenericNominal()) {
-    if (nominalTy->getValueTypeDestructor())
-      return false;
-  }
-  if (EnableExpandAll) {
-    return true;
-  }
-
-  unsigned numFields = module.Types.countNumberOfFields(ty, expansion);
-  return (numFields <= 6);
-}
-
 /// Some support functions for the global-opt and let-properties-opts
 
 // Encapsulate the state used for recursive analysis of a static