Merge pull request #77331 from aschwaighofer/large_types_reg2mem_aggressive

LargeTypesReg2Mem: Add a new heuristic that trys harder to keep large     values on the stack

Author: aschwaighofer

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/DeadStoreElimination.swift

@@ -75,7 +75,7 @@ let deadStoreElimination = FunctionPass(name: "dead-store-elimination") {
 }
 
 private func tryEliminate(store: StoreInst, complexityBudget: inout Int, _ context: FunctionPassContext) {
-  if !store.hasValidOwnershipForDeadStoreElimination {
+  if !store.hasValidOwnershipForDeadStoreElimination || !store.source.type.shouldExpand(context) {
     return
   }
 

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/RedundantLoadElimination.swift

@@ -97,6 +97,9 @@ private func eliminateRedundantLoads(in function: Function, ignoreArrays: Bool,
         {
           continue
         }
+        if !load.type.shouldExpand(context) {
+           continue
+        }
         tryEliminate(load: load, complexityBudget: &complexityBudget, context)
       }
     }

SwiftCompilerSources/Sources/Optimizer/PassManager/Options.swift

@@ -20,6 +20,10 @@ struct Options {
     _bridged.enableStackProtection()
   }
 
+  var useAggressiveReg2MemForCodeSize : Bool {
+    _bridged.useAggressiveReg2MemForCodeSize()
+  }
+
   var enableMoveInoutStackProtection: Bool {
     _bridged.enableMoveInoutStackProtection()
   }

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -809,3 +809,12 @@ extension CheckedCastAddrBranchInst {
     }
   }
 }
+
+extension Type {
+  func shouldExpand(_ context: some Context) -> Bool {
+    if !context.options.useAggressiveReg2MemForCodeSize {
+      return true
+    }
+    return context._bridged.shouldExpand(self.bridged)
+  }
+}

include/swift/AST/SILOptions.h

@@ -318,6 +318,10 @@ class SILOptions {
   /// for the debugger?
   bool ShouldFunctionsBePreservedToDebugger = true;
 
+  /// Block expanding and register promotion more aggressively throughout the
+  /// optimizer.
+  bool UseAggressiveReg2MemForCodeSize = false;
+
   SILOptions() {}
 
   /// Return a hash code of any components from these options that should

include/swift/Option/FrontendOptions.td

@@ -1294,6 +1294,13 @@ def disable_lifetime_dependence_diagnostics :
   Flag<["-"], "disable-lifetime-dependence-diagnostics">,
   HelpText<"Disable lifetime dependence diagnostics for Nonescapable types.">;
 
+def enable_aggressive_reg2mem :
+  Flag<["-"], "enable-aggressive-reg2mem">,
+  HelpText<"Enable a more aggresive reg2mem heuristic">;
+def disable_aggressive_reg2mem :
+  Flag<["-"], "disable-aggresive-reg2mem">,
+  HelpText<"Disable a more aggresive reg2mem heuristic">;
+
 def enable_collocate_metadata_functions :
   Flag<["-"], "enable-collocate-metadata-functions">,
   HelpText<"Enable collocate metadata functions">;

include/swift/SILOptimizer/OptimizerBridging.h

@@ -261,6 +261,8 @@ struct BridgedPassContext {
   BRIDGED_INLINE bool optimizeMemoryAccesses(BridgedFunction f) const;
   BRIDGED_INLINE bool eliminateDeadAllocations(BridgedFunction f) const;
 
+  BRIDGED_INLINE bool shouldExpand(BridgedType type) const;
+
   // IRGen
 
   SwiftInt getStaticSize(BridgedType type) const;
@@ -372,6 +374,7 @@ struct BridgedPassContext {
     Unchecked = 2
   };
 
+  BRIDGED_INLINE bool useAggressiveReg2MemForCodeSize() const;
   BRIDGED_INLINE bool enableStackProtection() const;
   BRIDGED_INLINE bool hasFeature(BridgedFeature feature) const;
   BRIDGED_INLINE bool enableMoveInoutStackProtection() const;

include/swift/SILOptimizer/OptimizerBridgingImpl.h

@@ -556,11 +556,21 @@ bool BridgedPassContext::enableMoveInoutStackProtection() const {
   return mod->getOptions().EnableMoveInoutStackProtection;
 }
 
+bool BridgedPassContext::useAggressiveReg2MemForCodeSize() const {
+  swift::SILModule *mod = invocation->getPassManager()->getModule();
+  return mod->getOptions().UseAggressiveReg2MemForCodeSize;
+}
+
 BridgedPassContext::AssertConfiguration BridgedPassContext::getAssertConfiguration() const {
   swift::SILModule *mod = invocation->getPassManager()->getModule();
   return (AssertConfiguration)mod->getOptions().AssertConfig;
 }
 
+bool BridgedPassContext::shouldExpand(BridgedType ty) const {
+  swift::SILModule &mod = *invocation->getPassManager()->getModule();
+  return swift::shouldExpand(mod, ty.unbridged());
+}
+
 static_assert((int)BridgedPassContext::SILStage::Raw == (int)swift::SILStage::Raw);
 static_assert((int)BridgedPassContext::SILStage::Canonical == (int)swift::SILStage::Canonical);
 static_assert((int)BridgedPassContext::SILStage::Lowered == (int)swift::SILStage::Lowered);

lib/Frontend/CompilerInvocation.cpp

@@ -3738,6 +3738,7 @@ bool CompilerInvocation::parseArgs(
     SILOpts.SkipFunctionBodies = FunctionBodySkipping::None;
     SILOpts.CMOMode = CrossModuleOptimizationMode::Everything;
     SILOpts.EmbeddedSwift = true;
+    SILOpts.UseAggressiveReg2MemForCodeSize = true;
     // OSSA modules are required for deinit de-virtualization.
     SILOpts.EnableOSSAModules = true;
     // -g is promoted to -gdwarf-types in embedded Swift
@@ -3751,6 +3752,11 @@ bool CompilerInvocation::parseArgs(
     }
   }
 
+  SILOpts.UseAggressiveReg2MemForCodeSize =
+    ParsedArgs.hasFlag(OPT_enable_aggressive_reg2mem,
+                       OPT_disable_aggressive_reg2mem,
+                       SILOpts.UseAggressiveReg2MemForCodeSize);
+
   // With Swift 6, enable @_spiOnly by default. This also enables proper error
   // reporting of ioi references from spi decls.
   if (LangOpts.EffectiveLanguageVersion.isVersionAtLeast(6)) {

lib/IRGen/GenCall.cpp

@@ -5346,7 +5346,8 @@ Explosion NativeConventionSchema::mapIntoNative(IRGenModule &IGM,
                                                 IRGenFunction &IGF,
                                                 Explosion &fromNonNative,
                                                 SILType type,
-                                                bool isOutlined) const {
+                                                bool isOutlined,
+                                                bool mayPeepholeLoad) const {
   if (fromNonNative.empty()) {
     assert(empty() && "Empty explosion must match the native convention");
     return Explosion();
@@ -5425,6 +5426,82 @@ Explosion NativeConventionSchema::mapIntoNative(IRGenModule &IGM,
                                          ? coercionTy
                                          : overlappedCoercionTy;
 
+  if (mayPeepholeLoad) {
+    auto succeeded = [&]() -> bool {
+      if (!overlappedCoercionTy->isEmptyTy())
+        return false;
+      auto load = dyn_cast<llvm::LoadInst>(*fromNonNative.begin());
+      if (!load)
+        return false;
+      auto *gep = dyn_cast<llvm::GetElementPtrInst>(load->getPointerOperand());
+      if (!gep)
+        return false;
+      auto *alloca = dyn_cast<llvm::AllocaInst>(getUnderlyingObject(gep));
+      if (!alloca)
+        return false;
+      auto numExplosions = fromNonNative.size();
+      if (numExplosions < 2)
+        return false;
+      for (unsigned i = 0, e = numExplosions; i < e; ++i) {
+        auto *otherLoad = dyn_cast<llvm::LoadInst>(*(fromNonNative.begin() + i));
+        if (!otherLoad)
+          return false;
+        auto otherAlloca = dyn_cast<llvm::AllocaInst>(
+          getUnderlyingObject(otherLoad->getPointerOperand()));
+        if (!otherAlloca || otherAlloca != alloca)
+          return false;
+        load = otherLoad;
+      }
+      auto allocaSize =
+        DataLayout.getTypeSizeInBits(alloca->getAllocatedType());
+
+      Address origAlloca(alloca, alloca->getAllocatedType(),
+                         Alignment(alloca->getAlign().value()));
+
+      IRBuilder Builder(*IGM.LLVMContext, false);
+      Builder.SetInsertPoint(load);
+
+      if (allocaSize < coercionSize) {
+        auto coerced = IGF.createAlloca(coercionTy, Alignment(alloca->getAlign().value()) , "tmp.coerce");
+        // Copy the defined bytes.
+        Builder.CreateMemCpy(coerced, origAlloca, Size(allocaSize/8));
+        origAlloca = coerced;
+      }
+
+      adjustAllocaAlignment(DataLayout, origAlloca, coercionTy);
+
+
+      unsigned expandedMapIdx = 0;
+      SmallVector<llvm::Value *, 8> expandedElts(expandedTys.size(), nullptr);
+      auto structAddr = Builder.CreateElementBitCast(origAlloca, coercionTy);
+      for (auto eltIndex : indices(coercionTy->elements())) {
+        auto layout = DataLayout.getStructLayout(coercionTy);
+        auto eltTy = coercionTy->getElementType(eltIndex);
+        // Skip padding fields.
+        if (eltTy->isArrayTy())
+          continue;
+        Address eltAddr = Builder.CreateStructGEP(structAddr, eltIndex, layout);
+        llvm::Value *elt = Builder.CreateLoad(eltAddr);
+        auto index = expandedTyIndicesMap[expandedMapIdx];
+        assert(expandedElts[index] == nullptr);
+        expandedElts[index] = elt;
+        ++expandedMapIdx;
+      }
+
+      // Add the values to the explosion.
+      for (auto *val : expandedElts)
+        nativeExplosion.add(val);
+      assert(expandedTys.size() == nativeExplosion.size());
+
+      return true;
+    }();
+
+    if (succeeded) {
+      (void)fromNonNative.claimAll();
+      return nativeExplosion;
+    }
+  }
+
   // Allocate a temporary for the coercion.
   Address temporary;
   Size tempSize;
@@ -5514,7 +5591,8 @@ Explosion IRGenFunction::coerceValueTo(SILType fromTy, Explosion &from,
 
 void IRGenFunction::emitScalarReturn(SILType returnResultType,
                                      SILType funcResultType, Explosion &result,
-                                     bool isSwiftCCReturn, bool isOutlined) {
+                                     bool isSwiftCCReturn, bool isOutlined,
+                                     bool mayPeepholeLoad) {
   if (result.empty()) {
     assert(IGM.getTypeInfo(returnResultType)
                .nativeReturnValueSchema(IGM)
@@ -5533,7 +5611,8 @@ void IRGenFunction::emitScalarReturn(SILType returnResultType,
     assert(!nativeSchema.requiresIndirect());
 
     Explosion native = nativeSchema.mapIntoNative(IGM, *this, result,
-                                                  funcResultType, isOutlined);
+                                                  funcResultType, isOutlined,
+                                                  mayPeepholeLoad);
     if (native.size() == 1) {
       Builder.CreateRet(native.claimNext());
       return;