IRGen: Add indirect peepholes for arguments and return value coercion

aschwaighofer · aschwaighofer · commit 5f396cd13c86 · 2024-10-31T13:02:07.000-07:00
Large tuples of values (e.g char[32]) can be passed directly at the abi
boundry but expand to a big explosion of values.

Peephole this explosion at argument passing and return value passing
points to avoid code size growth associated with the explosion.
diff --git a/lib/IRGen/GenCall.cpp b/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;
diff --git a/lib/IRGen/IRGenFunction.h b/lib/IRGen/IRGenFunction.h
@@ -102,7 +102,7 @@ class IRGenFunction {
   Explosion collectParameters();
   void emitScalarReturn(SILType returnResultType, SILType funcResultType,
                         Explosion &scalars, bool isSwiftCCReturn,
-                        bool isOutlined);
+                        bool isOutlined, bool mayPeepholeLoad = false);
   void emitScalarReturn(llvm::Type *resultTy, Explosion &scalars);
   
   void emitBBForReturn();
diff --git a/lib/IRGen/IRGenSIL.cpp b/lib/IRGen/IRGenSIL.cpp
@@ -4373,7 +4373,8 @@ void IRGenFunction::emitCoroutineOrAsyncExit(bool isUnwind) {
 static void emitReturnInst(IRGenSILFunction &IGF,
                            SILType resultTy,
                            Explosion &result,
-                           CanSILFunctionType fnType) {
+                           CanSILFunctionType fnType,
+                           bool mayPeepholeLoad) {
   SILFunctionConventions conv(IGF.CurSILFn->getLoweredFunctionType(),
                               IGF.getSILModule());
 
@@ -4448,10 +4449,37 @@ static void emitReturnInst(IRGenSILFunction &IGF,
     assert(swiftCCReturn ||
            funcLang == SILFunctionLanguage::C && "Need to handle all cases");
     IGF.emitScalarReturn(resultTy, funcResultType, result, swiftCCReturn,
-                         false);
+                         false, mayPeepholeLoad);
   }
 }
 
+static bool canPeepholeLoadToReturn(IRGenModule &IGM, swift::ReturnInst *r) {
+  auto *load = dyn_cast<LoadInst>(r->getOperand());
+  if (!load)
+    return false;
+
+  // Later code can't deal with projections.
+  if (!isa<AllocStackInst>(load->getOperand()))
+    return false;
+
+  if (load->getParent() != r->getParent())
+    return false;
+
+  for (auto it = ++load->getIterator(), e = r->getIterator(); it != e; ++it) {
+    if (it->mayHaveSideEffects()) {
+      if (auto *dealloc = dyn_cast<DeallocStackInst>(&*it)) {
+        auto &ti = IGM.getTypeInfo(
+          dealloc->getOperand()->getType().getObjectType());
+        if (!ti.isLoadable())
+          return false;
+        continue;
+      }
+      return false;
+    }
+  }
+  return true;
+}
+
 void IRGenSILFunction::visitReturnInst(swift::ReturnInst *i) {
   Explosion result = getLoweredExplosion(i->getOperand());
 
@@ -4466,8 +4494,11 @@ void IRGenSILFunction::visitReturnInst(swift::ReturnInst *i) {
     result = std::move(temp);
   }
 
+  bool mayPeepholeLoad = canPeepholeLoadToReturn(IGM, i);
+
   emitReturnInst(*this, i->getOperand()->getType(), result,
-                 i->getFunction()->getLoweredFunctionType());
+                 i->getFunction()->getLoweredFunctionType(),
+                 mayPeepholeLoad);
 }
 
 void IRGenSILFunction::visitThrowInst(swift::ThrowInst *i) {
@@ -5524,6 +5555,81 @@ void IRGenSILFunction::visitLoadInst(swift::LoadInst *i) {
   }
   setLoweredExplosion(i, lowered);
 }
+static Address isSafeForMemCpyPeephole(const TypeInfo &TI, SILArgument *arg,
+                                       Explosion &argSrc, AllocStackInst *dst,
+                                       Address storeDst,
+                                       StoreInst *store,
+                                       llvm::Instruction * &insertPt) {
+  if (!arg || !dst)
+    return Address();
+
+  // Store of function argument.
+  if (store->getParent() != store->getFunction()->getEntryBlock())
+    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();
+  }
+
+  auto *dstAlloca = dyn_cast<llvm::AllocaInst>(storeDst.getAddress());
+  if (!dstAlloca)
+    return Address();
+
+  // Move the lifetime.begin above the load instruction (where we eventually
+  // will insert the memcpy.
+  llvm::Instruction *lifetimeBegin = nullptr;
+  for (const auto &use : dstAlloca->uses()) {
+    auto *begin = dyn_cast<llvm::LifetimeIntrinsic>(use.getUser());
+    if (!begin)
+      continue;
+    if (begin->getParent() != alloca->getParent())
+      continue;
+    if (begin->getIntrinsicID() != llvm::Intrinsic::lifetime_start)
+      continue;
+
+    if (lifetimeBegin) {
+      // Seen a second lifetime.begin in the entry block.
+      lifetimeBegin = nullptr;
+      break;
+    }
+    lifetimeBegin = begin;
+  }
+
+  if (!lifetimeBegin) {
+    return Address();
+  }
+
+  lifetimeBegin->moveBefore(load);
+
+  // 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);
+}
 
 static Address canForwardIndirectResultAlloca(const TypeInfo &TI,
                                               StoreInst *store,
@@ -5576,7 +5682,6 @@ void IRGenSILFunction::visitStoreInst(swift::StoreInst *i) {
   SILType objType = i->getSrc()->getType().getObjectType();
   const auto &typeInfo = cast<LoadableTypeInfo>(getTypeInfo(objType));
 
-
   llvm::Instruction *insertPt = nullptr;
   auto forwardAddr = canForwardIndirectResultAlloca(typeInfo, i, source,
                                                     insertPt);
@@ -5585,12 +5690,33 @@ void IRGenSILFunction::visitStoreInst(swift::StoreInst *i) {
     // 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);
+    ArtificialLocation Loc(getDebugScope(), IGM.DebugInfo.get(), Builder);
     addrTI.initializeWithTake(*this, dest, forwardAddr, i->getDest()->getType(),
                               false, /*zeroizeIfSensitive=*/ true);
     (void)source.claimAll();
     return;
   }
 
+  // See if we can forward a load from an alloca we have created for the purpose
+  // of argument coercion.
+  auto argSrc = dyn_cast<SILArgument>(i->getSrc());
+  auto stackDst = dyn_cast<AllocStackInst>(i->getDest());
+  const auto &addrTI = getTypeInfo(i->getDest()->getType());
+  insertPt = nullptr;
+
+  auto srcAddr = isSafeForMemCpyPeephole(addrTI, argSrc, source, stackDst, dest,
+                                         i, insertPt);
+  if (srcAddr.isValid() &&
+      (i->getOwnershipQualifier() == StoreOwnershipQualifier::Trivial ||
+       i->getOwnershipQualifier() == StoreOwnershipQualifier::Unqualified)) {
+    IRBuilder::SavedInsertionPointRAII insertRAII(this->Builder, insertPt);
+    ArtificialLocation Loc(getDebugScope(), IGM.DebugInfo.get(), Builder);
+    addrTI.initializeWithTake(*this, dest, srcAddr, i->getDest()->getType(),
+                              false, /*zeroizeIfSensitive*/true);
+    (void)source.claimAll();
+    return;
+  }
+
   switch (i->getOwnershipQualifier()) {
   case StoreOwnershipQualifier::Unqualified:
   case StoreOwnershipQualifier::Init:
diff --git a/lib/IRGen/NativeConventionSchema.h b/lib/IRGen/NativeConventionSchema.h
@@ -57,7 +57,7 @@ class NativeConventionSchema {
   /// calling convention's schema.
   Explosion mapIntoNative(IRGenModule &IGM, IRGenFunction &IGF,
                           Explosion &fromNonNative, SILType type,
-                          bool isOutlined) const;
+                          bool isOutlined, bool mayPeepholeLoad = false) const;
 
   /// Map form a native explosion that follows the native calling convention's
   /// schema to a non-native explosion whose schema is described by
diff --git a/test/IRGen/Inputs/arg_and_result_peepholes.h b/test/IRGen/Inputs/arg_and_result_peepholes.h
@@ -0,0 +1,7 @@
+#pragma once
+
+struct BigStruct {
+  char a[32];
+};
+
+struct BigStruct useBigStruct(struct BigStruct x);
diff --git a/test/IRGen/arg_and_result_peepholes.swift b/test/IRGen/arg_and_result_peepholes.swift
@@ -0,0 +1,52 @@
+// RUN: %target-swift-frontend %s  -import-objc-header %S/Inputs/arg_and_result_peepholes.h -emit-ir 2>&1 | %FileCheck %s
+
+// REQUIRES: PTRSIZE=64
+
+// CHECK: define{{.*}} swiftcc { i64, i64, i64, i64 } @"$s24arg_and_result_peepholes05test_D0ySo9BigStructVADF"(i64 %0, i64 %1, i64 %2, i64 %3)
+// CHECK: entry:
+// CHECK:   [[TMP:%.*]] = alloca { i64, i64, i64, i64 }
+// CHECK:   [[RES_MEM:%.*]] = alloca %TSo9BigStructV
+// CHECK:   [[ARG_MEM:%.*]] = alloca %TSo9BigStructV
+// CHECK:   [[TMP2:%.*]] = alloca %TSo9BigStructV
+// CHECK:   [[CALL:%.*]] = alloca %TSo9BigStructV
+// CHECK:   call void @llvm.lifetime.start.p0(i64 256, ptr [[TMP]])
+// CHECK:   [[A1:%.*]] = getelementptr inbounds { i64, i64, i64, i64 }, ptr [[TMP]], i32 0, i32 0
+// CHECK:   store i64 %0, ptr [[A1]]
+// CHECK:   [[A2:%.*]] = getelementptr inbounds { i64, i64, i64, i64 }, ptr [[TMP]], i32 0, i32 1
+// CHECK:   store i64 %1, ptr [[A2]]
+// CHECK:   [[A3:%.*]] = getelementptr inbounds { i64, i64, i64, i64 }, ptr [[TMP]], i32 0, i32 2
+// CHECK:   store i64 %2, ptr [[A3]]
+// CHECK:   [[A4:%.*]] = getelementptr inbounds { i64, i64, i64, i64 }, ptr [[TMP]], i32 0, i32 3
+// CHECK:   store i64 %3, ptr [[A4]]
+
+// CHECK:  call void @llvm.lifetime.start.p0(i64 32, ptr [[ARG_MEM]])
+// CHECK:  call void @llvm.memcpy.p0.p0.i64(ptr{{.*}} [[ARG_MEM]], ptr{{.*}} [[TMP]], i64 32, i1 false)
+
+// CHECK:  call void @llvm.lifetime.start.p0(i64 32, ptr [[TMP2]])
+// CHECK:  call void @llvm.memcpy.p0.p0.i64(ptr{{.*}} [[TMP2]], ptr{{.*}} [[ARG_MEM]], i64 32, i1 false)
+// CHECK:  call void @llvm.lifetime.start.p0(i64 32, ptr [[CALL]])
+// CHECK:  call void @useBigStruct(ptr{{.*}} [[CALL]], ptr{{.*}} [[TMP2]])
+// CHECK:  call void @llvm.lifetime.end.p0(i64 32, ptr [[TMP2]])
+
+// CHECK: call void @llvm.memcpy.p0.p0.i64(ptr{{.*}} [[RES_MEM]], ptr{{.*}} [[CALL]], i64 32, i1 false)
+
+// CHECK:  [[A5:%.*]] = getelementptr inbounds { i64, i64, i64, i64 }, ptr [[RES_MEM]], i32 0, i32 0
+// CHECK:  [[R1:%.*]] = load i64, ptr [[A5]]
+// CHECK:  [[A6:%.*]] = getelementptr inbounds { i64, i64, i64, i64 }, ptr [[RES_MEM]], i32 0, i32 1
+// CHECK:  [[R2:%.*]] = load i64, ptr [[A6]]
+// CHECK:  [[A7:%.*]] = getelementptr inbounds { i64, i64, i64, i64 }, ptr [[RES_MEM]], i32 0, i32 2
+// CHECK:  [[R3:%.*]] = load i64, ptr [[A7]]
+// CHECK:  [[A8:%.*]] = getelementptr inbounds { i64, i64, i64, i64 }, ptr [[RES_MEM]], i32 0, i32 3
+// CHECK:  [[R4:%.*]] = load i64, ptr [[A8]]
+// CHECK:  call void @llvm.lifetime.end.p0(i64 32, ptr [[ARG_MEM]])
+// CHECK:  call void @llvm.lifetime.end.p0(i64 32, ptr [[RES_MEM]])
+
+// CHECK:  [[R5:%.*]] = insertvalue { i64, i64, i64, i64 } undef, i64 [[R1]], 0
+// CHECK:  [[R6:%.*]] = insertvalue { i64, i64, i64, i64 } [[R5]], i64 [[R2]], 1
+// CHECK:  [[R7:%.*]] = insertvalue { i64, i64, i64, i64 } [[R6]], i64 [[R3]], 2
+// CHECK:  [[R8:%.*]] = insertvalue { i64, i64, i64, i64 } [[R7]], i64 [[R4]], 3
+// CHECK:  ret { i64, i64, i64, i64 } [[R8]]
+
+public func test_peepholes(_ v: BigStruct) -> BigStruct {
+  return useBigStruct(v)
+}
