Merge pull request #77379 from aschwaighofer/enable_aggressive_reg2mem

Enable heuristic that tries to keep large values in memory

Author: aschwaighofer

include/swift/AST/SILOptions.h

@@ -324,7 +324,7 @@ class SILOptions {
 
   /// Block expanding and register promotion more aggressively throughout the
   /// optimizer.
-  bool UseAggressiveReg2MemForCodeSize = false;
+  bool UseAggressiveReg2MemForCodeSize = true;
 
   SILOptions() {}
 

include/swift/Option/FrontendOptions.td

@@ -1299,10 +1299,10 @@ def disable_lifetime_dependence_diagnostics :
 
 def enable_aggressive_reg2mem :
   Flag<["-"], "enable-aggressive-reg2mem">,
-  HelpText<"Enable a more aggresive reg2mem heuristic">;
+  HelpText<"Enable a more aggressive reg2mem heuristic">;
 def disable_aggressive_reg2mem :
-  Flag<["-"], "disable-aggresive-reg2mem">,
-  HelpText<"Disable a more aggresive reg2mem heuristic">;
+  Flag<["-"], "disable-aggressive-reg2mem">,
+  HelpText<"Disable a more aggressive reg2mem heuristic">;
 
 def enable_collocate_metadata_functions :
   Flag<["-"], "enable-collocate-metadata-functions">,

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/IRGen/LoadableByAddress.cpp

@@ -3495,6 +3495,7 @@ class LargeLoadableHeuristic {
       UseAggressiveHeuristic(UseAggressiveHeuristic) {}
 
    void visit(SILInstruction *i);
+   void visit(SILArgument *arg);
 
    bool isLargeLoadableType(SILType ty);
    bool isPotentiallyCArray(SILType ty);
@@ -3539,6 +3540,27 @@ class LargeLoadableHeuristic {
 };
 }
 
+void LargeLoadableHeuristic::visit(SILArgument *arg) {
+  auto objType = arg->getType().getObjectType();
+  if (numRegisters(objType) < NumRegistersLargeType)
+    return;
+
+  auto &entry = largeTypeProperties[objType];
+  for (auto *use : arg->getUses()) {
+    auto *usr = use->getUser();
+    switch (usr->getKind()) {
+    case SILInstructionKind::TupleExtractInst:
+    case SILInstructionKind::StructExtractInst: {
+      auto projectionTy = cast<SingleValueInstruction>(usr)->getType();
+      if (numRegisters(projectionTy) >= NumRegistersLargeType)
+        entry.addProjection();
+      break;
+    }
+    default:
+      continue;
+    }
+  }
+}
 void LargeLoadableHeuristic::visit(SILInstruction *i) {
   if (!UseAggressiveHeuristic)
     return;
@@ -4611,6 +4633,9 @@ static void runPeepholesAndReg2Mem(SILPassManager *pm, SILModule *silMod,
                                      UseAggressiveHeuristic);
     Peepholes opts(pm, silMod, irgenModule);
     for (SILBasicBlock &BB : currF) {
+      for (auto *arg : BB.getArguments()) {
+        heuristic.visit(arg);
+      }
       for (SILInstruction &I : BB) {
         heuristic.visit(&I);
         if (opts.ignore(&I))

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()) {
@@ -1412,3 +1417,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

test/IRGen/big_types.sil

@@ -80,16 +80,17 @@ unwind:
 // CHECK-LABEL: sil @use_yield_big : $@convention(thin) () -> () {
 // CHECK:       bb0:
 // CHECK-NEXT:    [[TEMP:%.*]] = alloc_stack $BigStruct
+// CHECK-NEXT:    [[TEMP2:%.*]] = alloc_stack $BigStruct
 // CHECK-NEXT:    // function_ref
 // CHECK-NEXT:    [[CORO:%.*]] = function_ref @test_yield_big : $@yield_once @convention(thin) () -> @yields @in_guaranteed BigStruct
 // CHECK-NEXT:    ([[ADDR:%.*]], [[TOKEN:%.*]]) = begin_apply [[CORO]]()
-//   TODO: this isn't very efficient
-// CHECK-NEXT:    [[T0:%.*]] = load [[ADDR]] : $*BigStruct
-// CHECK-NEXT:    store [[T0]] to [[TEMP]] : $*BigStruct
+// CHECK-NEXT:   copy_addr [take] [[ADDR]] to [init] [[TEMP]] : $*BigStruct
+// CHECK-NEXT:   copy_addr [take] [[TEMP]] to [init] [[TEMP2]] : $*BigStruct
 // CHECK-NEXT:    // function_ref
 // CHECK-NEXT:    [[USE:%.*]] = function_ref @use_big_struct : $@convention(thin) (@in_guaranteed BigStruct) -> ()
-// CHECK-NEXT:    apply [[USE]]([[TEMP]])
+// CHECK-NEXT:    apply [[USE]]([[TEMP2]])
 // CHECK-NEXT:    [[RET:%.*]] = end_apply [[TOKEN]] as $()
+// CHECK-NEXT:    dealloc_stack [[TEMP2]] : $*BigStruct
 // CHECK-NEXT:    dealloc_stack [[TEMP]] : $*BigStruct
 // CHECK-NEXT:    return [[RET]] : $()
 sil @use_yield_big : $@convention(thin) () -> () {