LargeTypesReg2Mem: Forward the address of large formally by-val arguments and return values of C functions

rdar://47978338

Author: aschwaighofer

lib/IRGen/GenCall.cpp

@@ -31,6 +31,7 @@
 #include "clang/CodeGen/CodeGenABITypes.h"
 #include "clang/CodeGen/ModuleBuilder.h"
 #include "clang/Sema/Sema.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/GlobalPtrAuthInfo.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/Support/Compiler.h"
@@ -3962,6 +3963,26 @@ void irgen::emitClangExpandedParameter(IRGenFunction &IGF,
   swiftTI.deallocateStack(IGF, tempAlloc, swiftType);
 }
 
+Address getForwardableAlloca(const TypeInfo &TI, bool isForwardableArgument,
+                             Explosion &in) {
+  if (!isForwardableArgument)
+    return Address();
+
+  auto *load = dyn_cast<llvm::LoadInst>(*in.begin());
+  if (!load)
+    return Address();
+
+  auto *gep = dyn_cast<llvm::GetElementPtrInst>(load->getPointerOperand());
+  if (!gep)
+    return Address();
+
+  auto *alloca = dyn_cast<llvm::AllocaInst>(getUnderlyingObject(gep));
+  if (!alloca)
+    return Address();
+
+  return TI.getAddressForPointer(alloca);
+}
+
 static void externalizeArguments(IRGenFunction &IGF, const Callee &callee,
                                  Explosion &in, Explosion &out,
                                  TemporarySet &temporaries,
@@ -4021,6 +4042,8 @@ static void externalizeArguments(IRGenFunction &IGF, const Callee &callee,
     SILType paramType = silConv.getSILType(
         params[i - firstParam], IGF.IGM.getMaximalTypeExpansionContext());
 
+    bool isForwardableArgument = IGF.isForwardableArgument(i - firstParam);
+
     // In Swift, values that are foreign references types will always be
     // pointers. Additionally, we only import functions which use foreign
     // reference types indirectly (as pointers), so we know in every case, if
@@ -4082,8 +4105,17 @@ static void externalizeArguments(IRGenFunction &IGF, const Callee &callee,
                          Alignment(ABIAlign.getQuantity()));
         }
       }
-
-      ti.initialize(IGF, in, addr, isOutlined);
+      Address forwardFromAddr = getForwardableAlloca(ti, isForwardableArgument,
+                                                     in);
+      // Try to forward the address from a `load` instruction "immediately"
+      // preceeding the apply.
+      if (isForwardableArgument && forwardFromAddr.isValid()) {
+        ti.initializeWithTake(IGF, addr, forwardFromAddr,
+                              paramType.getAddressType(), isOutlined);
+        (void)in.claim(ti.getSchema().size());
+      } else {
+        ti.initialize(IGF, in, addr, isOutlined);
+      }
 
       out.add(addr.getAddress());
       break;

lib/IRGen/IRGen.cpp

@@ -90,6 +90,7 @@
 #include "llvm/Transforms/Instrumentation/ThreadSanitizer.h"
 #include "llvm/Transforms/ObjCARC.h"
 #include "llvm/Transforms/Scalar.h"
+#include "llvm/Transforms/Scalar/DCE.h"
 
 #include <thread>
 
@@ -329,7 +330,14 @@ void swift::performLLVMOptimizations(const IRGenOptions &Opts,
           MPM.addPass(InstrProfiling(options, false));
         });
   }
-
+  if (Opts.shouldOptimize()) {
+    PB.registerPipelineStartEPCallback(
+        [](ModulePassManager &MPM, OptimizationLevel level) {
+          // Run this before SROA to avoid un-neccessary expansion of dead
+          // loads.
+          MPM.addPass(createModuleToFunctionPassAdaptor(DCEPass()));
+        });
+  }
   bool isThinLTO = Opts.LLVMLTOKind  == IRGenLLVMLTOKind::Thin;
   bool isFullLTO = Opts.LLVMLTOKind  == IRGenLLVMLTOKind::Full;
   if (!Opts.shouldOptimize() || Opts.DisableLLVMOptzns) {

lib/IRGen/IRGenFunction.h

@@ -808,6 +808,21 @@ class IRGenFunction {
   LocalTypeDataCache const *getLocalTypeData() { return LocalTypeData; }
 #endif
 
+  /// A forwardable argument is a load that is immediately preceeds the apply it
+  /// is used as argument to. If there is no side-effecting instructions between
+  /// said load and the apply, we can memcpy the loads address to the apply's
+  /// indirect argument alloca.
+  void clearForwardableArguments() {
+    forwardableArguments.clear();
+  }
+
+  void setForwardableArgument(unsigned idx) {
+    forwardableArguments.insert(idx);
+  }
+
+  bool isForwardableArgument(unsigned idx) const {
+    return forwardableArguments.contains(idx);
+  }
 private:
   LocalTypeDataCache &getOrCreateLocalTypeData();
   void destroyLocalTypeData();
@@ -826,6 +841,8 @@ class IRGenFunction {
   CanType SelfType;
   bool SelfTypeIsExact = false;
   DynamicSelfKind SelfKind;
+
+  llvm::SmallSetVector<unsigned, 16> forwardableArguments;
 };
 
 using ConditionalDominanceScope = IRGenFunction::ConditionalDominanceScope;

lib/IRGen/IRGenSIL.cpp

@@ -16,10 +16,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "GenKeyPath.h"
-#include "swift/AST/ExtInfo.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTMangler.h"
 #include "swift/AST/DiagnosticsIRGen.h"
+#include "swift/AST/ExtInfo.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/IRGenOptions.h"
 #include "swift/AST/ParameterList.h"
@@ -53,6 +53,7 @@
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/TinyPtrVector.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"
@@ -3258,7 +3259,9 @@ void IRGenSILFunction::visitExistentialMetatypeInst(
 static void emitApplyArgument(IRGenSILFunction &IGF,
                               SILValue arg,
                               SILType paramType,
-                              Explosion &out) {
+                              Explosion &out,
+                              SILInstruction *apply = nullptr,
+                              unsigned idx = 0) {
   bool isSubstituted = (arg->getType() != paramType);
 
   // For indirect arguments, we just need to pass a pointer.
@@ -3282,7 +3285,24 @@ static void emitApplyArgument(IRGenSILFunction &IGF,
 
   // Fast path: avoid an unnecessary temporary explosion.
   if (!isSubstituted) {
+    bool canForwardLoadToIndirect = false;
+    auto *load = dyn_cast<LoadInst>(arg);
+    [&]() {
+      if (apply && load && apply->getParent() == load->getParent()) {
+        for (auto it = std::next(load->getIterator()), e = apply->getIterator();
+             it != e; ++it) {
+          if (isa<LoadInst>(&(*it))) {
+            continue;
+          }
+          return;
+        }
+        canForwardLoadToIndirect = true;
+      }
+    }();
     IGF.getLoweredExplosion(arg, out);
+    if (canForwardLoadToIndirect) {
+      IGF.setForwardableArgument(idx);
+    }
     return;
   }
 
@@ -3710,6 +3730,8 @@ void IRGenSILFunction::visitFullApplySite(FullApplySite site) {
   // Lower the SIL arguments to IR arguments.
 
   // Turn the formal SIL parameters into IR-gen things.
+  clearForwardableArguments();
+
   for (auto index : indices(args)) {
     if (origConv.hasIndirectSILErrorResults() &&
         index == origConv.getNumIndirectSILResults()) {
@@ -3718,7 +3740,7 @@ void IRGenSILFunction::visitFullApplySite(FullApplySite site) {
       continue;
     }
     emitApplyArgument(*this, args[index], emission->getParameterType(index),
-                      llArgs);
+                      llArgs, site.getInstruction(), index);
   }
 
   auto &calleeFP = emission->getCallee().getFunctionPointer();
@@ -3744,6 +3766,9 @@ void IRGenSILFunction::visitFullApplySite(FullApplySite site) {
   Explosion result;
   emission->emitToExplosion(result, false);
 
+  // We might have set forwardable arguments. Clear it for the next round.
+  clearForwardableArguments();
+
   // For a simple apply, just bind the apply result to the result of the call.
   if (auto apply = dyn_cast<ApplyInst>(i)) {
     setLoweredExplosion(apply, result);
@@ -5336,12 +5361,72 @@ void IRGenSILFunction::visitLoadInst(swift::LoadInst *i) {
   setLoweredExplosion(i, lowered);
 }
 
+static Address canForwardIndirectResultAlloca(const TypeInfo &TI,
+                                              StoreInst *store,
+                                              Explosion &argSrc,
+                                              llvm::Instruction * &insertPt) {
+  // Check that the store stores the result of and apply instruction immediately
+  // preceeding the store.
+  auto *apply = dyn_cast<ApplyInst>(store->getSrc());
+  auto *allocStack = dyn_cast<AllocStackInst>(store->getDest());
+  if (!apply || !allocStack || apply->getParent() != store->getParent() ||
+      std::next(apply->getIterator()) != store->getIterator())
+    return Address();
+
+  auto explosionSize = TI.getSchema().size();
+  if (argSrc.size() < 1 || explosionSize < 4)
+    return Address();
+
+  auto *load = dyn_cast<llvm::LoadInst>(*argSrc.begin());
+  if (!load)
+    return Address();
+  auto *gep = dyn_cast<llvm::GetElementPtrInst>(load->getPointerOperand());
+  if (!gep)
+    return Address();
+
+  auto *alloca = dyn_cast<llvm::AllocaInst>(getUnderlyingObject(gep));
+  if (!alloca)
+    return Address();
+
+  // Check all the other loads.
+  for (size_t i = 1, e = explosionSize; i != e; ++i) {
+    auto *load = dyn_cast<llvm::LoadInst>(*(argSrc.begin() + i));
+    if (!load)
+      return Address();
+    auto *alloca2 = dyn_cast<llvm::AllocaInst>(
+        getUnderlyingObject(load->getPointerOperand()));
+    if (!alloca2 || alloca2 != alloca)
+      return Address();
+  }
+
+  // Set insertPt to the first load such that we are within the lifetime of the
+  // alloca marked by the lifetime intrinsic.
+  insertPt = load;
+
+  return TI.getAddressForPointer(alloca);
+}
+
 void IRGenSILFunction::visitStoreInst(swift::StoreInst *i) {
   Explosion source = getLoweredExplosion(i->getSrc());
   Address dest = getLoweredAddress(i->getDest());
   SILType objType = i->getSrc()->getType().getObjectType();
-
   const auto &typeInfo = cast<LoadableTypeInfo>(getTypeInfo(objType));
+
+
+  llvm::Instruction *insertPt = nullptr;
+  auto forwardAddr = canForwardIndirectResultAlloca(typeInfo, i, source,
+                                                    insertPt);
+  if (forwardAddr.isValid()) {
+    const auto &addrTI = getTypeInfo(i->getDest()->getType());
+    // Set the insert point to the first load instruction. We need to be with
+    // the lifetime of the alloca.
+    IRBuilder::SavedInsertionPointRAII insertRAII(this->Builder, insertPt);
+    addrTI.initializeWithTake(*this, dest, forwardAddr, i->getDest()->getType(),
+                              false);
+    (void)source.claimAll();
+    return;
+  }
+
   switch (i->getOwnershipQualifier()) {
   case StoreOwnershipQualifier::Unqualified:
   case StoreOwnershipQualifier::Init:

test/IRGen/Inputs/large_argument_result_c.h

@@ -0,0 +1,7 @@
+#pragma once
+
+typedef struct {
+  unsigned long long arr[16];
+} large_thing;
+
+large_thing pass_and_return(large_thing a, large_thing b);

test/IRGen/large_argument_result_c.swift

@@ -0,0 +1,20 @@
+// RUN: %target-swift-frontend %s -Osize -import-objc-header %S/Inputs/large_argument_result_c.h -emit-ir -o - 2>&1 | %FileCheck %s
+
+// REQUIRES: PTRSIZE=64
+
+// CHECK: define swiftcc void @"$s23large_argument_result_c7runTestyySo0A6_thingaF"(ptr {{.*}} %0)
+// CHECK:  [[CALL_ALLOCA:%.*]] = alloca <{ %Ts6UInt64V, %Ts6UInt64V, %Ts6UInt64V
+// CHECK:  [[TMP_ALLOCA:%.*]] = alloca %TSo11large_thinga, align 8
+// CHECK:  [[TMP_ALLOCA2:%.*]] = alloca %TSo11large_thinga, align 8
+// CHECK:  call void @llvm.memcpy.p0.p0.i64(ptr {{.*}} [[TMP_ALLOCA2]], ptr {{.*}} %0, i64 128, i1 false)
+// CHECK:  call void @llvm.memcpy.p0.p0.i64(ptr {{.*}} [[TMP_ALLOCA]], ptr {{.*}} %0, i64 128, i1 false)
+// CHECK:  call void @pass_and_return(ptr {{.*}} [[CALL_ALLOCA]], ptr nonnull [[TMP_ALLOCA]], ptr nonnull [[TMP_ALLOCA2]])
+// CHECK:  call {{.*}} @swift_allocObject
+// CHECK:  [[BOX:%.*]] = call noalias ptr @swift_allocObject(
+// CHECK:  [[ADDR_IN_BOX:%.*]] = getelementptr inbounds <{ %swift.refcounted, [128 x i8] }>, ptr [[BOX]], i64 0, i32 1
+// CHECK:  call void @llvm.memcpy.p0.p0.i64(ptr {{.*}} [[ADDR_IN_BOX]], ptr {{.*}} [[CALL_ALLOCA]], i64 128, i1 false)
+// CHECK:  call void @llvm.lifetime.end.p0(i64 128, ptr nonnull [[CALL_ALLOCA]])
+public func runTest(_ l : large_thing) {
+  let r = pass_and_return(l, l)
+  print(r)
+}