[ARM64EC] Fix thunks for vector args (#96003)

The checks when building a thunk to decide if an arg needed to be cast
to/from an integer or redirected via a pointer didn't match how arg
types were changed in `canonicalizeThunkType`, this caused LLVM to ICE
when using vector types as args due to incorrect types in a call
instruction.

Instead of duplicating these checks, we should check if the arg type
differs between x64 and AArch64 and then cast or redirect as
appropriate.

Author: dpaoliello

llvm/lib/Target/AArch64/AArch64Arm64ECCallLowering.cpp

@@ -46,6 +46,18 @@ static cl::opt<bool> GenerateThunks("arm64ec-generate-thunks", cl::Hidden,
 
 namespace {
 
+enum ThunkArgTranslation : uint8_t {
+  Direct,
+  Bitcast,
+  PointerIndirection,
+};
+
+struct ThunkArgInfo {
+  Type *Arm64Ty;
+  Type *X64Ty;
+  ThunkArgTranslation Translation;
+};
+
 class AArch64Arm64ECCallLowering : public ModulePass {
 public:
   static char ID;
@@ -74,25 +86,30 @@ class AArch64Arm64ECCallLowering : public ModulePass {
 
   void getThunkType(FunctionType *FT, AttributeList AttrList,
                     Arm64ECThunkType TT, raw_ostream &Out,
-                    FunctionType *&Arm64Ty, FunctionType *&X64Ty);
+                    FunctionType *&Arm64Ty, FunctionType *&X64Ty,
+                    SmallVector<ThunkArgTranslation> &ArgTranslations);
   void getThunkRetType(FunctionType *FT, AttributeList AttrList,
                        raw_ostream &Out, Type *&Arm64RetTy, Type *&X64RetTy,
                        SmallVectorImpl<Type *> &Arm64ArgTypes,
-                       SmallVectorImpl<Type *> &X64ArgTypes, bool &HasSretPtr);
+                       SmallVectorImpl<Type *> &X64ArgTypes,
+                       SmallVector<ThunkArgTranslation> &ArgTranslations,
+                       bool &HasSretPtr);
   void getThunkArgTypes(FunctionType *FT, AttributeList AttrList,
                         Arm64ECThunkType TT, raw_ostream &Out,
                         SmallVectorImpl<Type *> &Arm64ArgTypes,
-                        SmallVectorImpl<Type *> &X64ArgTypes, bool HasSretPtr);
-  void canonicalizeThunkType(Type *T, Align Alignment, bool Ret,
-                             uint64_t ArgSizeBytes, raw_ostream &Out,
-                             Type *&Arm64Ty, Type *&X64Ty);
+                        SmallVectorImpl<Type *> &X64ArgTypes,
+                        SmallVectorImpl<ThunkArgTranslation> &ArgTranslations,
+                        bool HasSretPtr);
+  ThunkArgInfo canonicalizeThunkType(Type *T, Align Alignment, bool Ret,
+                                     uint64_t ArgSizeBytes, raw_ostream &Out);
 };
 
 } // end anonymous namespace
 
 void AArch64Arm64ECCallLowering::getThunkType(
     FunctionType *FT, AttributeList AttrList, Arm64ECThunkType TT,
-    raw_ostream &Out, FunctionType *&Arm64Ty, FunctionType *&X64Ty) {
+    raw_ostream &Out, FunctionType *&Arm64Ty, FunctionType *&X64Ty,
+    SmallVector<ThunkArgTranslation> &ArgTranslations) {
   Out << (TT == Arm64ECThunkType::Entry ? "$ientry_thunk$cdecl$"
                                         : "$iexit_thunk$cdecl$");
 
@@ -111,10 +128,10 @@ void AArch64Arm64ECCallLowering::getThunkType(
 
   bool HasSretPtr = false;
   getThunkRetType(FT, AttrList, Out, Arm64RetTy, X64RetTy, Arm64ArgTypes,
-                  X64ArgTypes, HasSretPtr);
+                  X64ArgTypes, ArgTranslations, HasSretPtr);
 
   getThunkArgTypes(FT, AttrList, TT, Out, Arm64ArgTypes, X64ArgTypes,
-                   HasSretPtr);
+                   ArgTranslations, HasSretPtr);
 
   Arm64Ty = FunctionType::get(Arm64RetTy, Arm64ArgTypes, false);
 
@@ -124,7 +141,8 @@ void AArch64Arm64ECCallLowering::getThunkType(
 void AArch64Arm64ECCallLowering::getThunkArgTypes(
     FunctionType *FT, AttributeList AttrList, Arm64ECThunkType TT,
     raw_ostream &Out, SmallVectorImpl<Type *> &Arm64ArgTypes,
-    SmallVectorImpl<Type *> &X64ArgTypes, bool HasSretPtr) {
+    SmallVectorImpl<Type *> &X64ArgTypes,
+    SmallVectorImpl<ThunkArgTranslation> &ArgTranslations, bool HasSretPtr) {
 
   Out << "$";
   if (FT->isVarArg()) {
@@ -153,17 +171,20 @@ void AArch64Arm64ECCallLowering::getThunkArgTypes(
     for (int i = HasSretPtr ? 1 : 0; i < 4; i++) {
       Arm64ArgTypes.push_back(I64Ty);
       X64ArgTypes.push_back(I64Ty);
+      ArgTranslations.push_back(ThunkArgTranslation::Direct);
     }
 
     // x4
     Arm64ArgTypes.push_back(PtrTy);
     X64ArgTypes.push_back(PtrTy);
+    ArgTranslations.push_back(ThunkArgTranslation::Direct);
     // x5
     Arm64ArgTypes.push_back(I64Ty);
     if (TT != Arm64ECThunkType::Entry) {
       // FIXME: x5 isn't actually used by the x64 side; revisit once we
       // have proper isel for varargs
       X64ArgTypes.push_back(I64Ty);
+      ArgTranslations.push_back(ThunkArgTranslation::Direct);
     }
     return;
   }
@@ -187,18 +208,20 @@ void AArch64Arm64ECCallLowering::getThunkArgTypes(
     uint64_t ArgSizeBytes = 0;
     Align ParamAlign = Align();
 #endif
-    Type *Arm64Ty, *X64Ty;
-    canonicalizeThunkType(FT->getParamType(I), ParamAlign,
-                          /*Ret*/ false, ArgSizeBytes, Out, Arm64Ty, X64Ty);
+    auto [Arm64Ty, X64Ty, ArgTranslation] =
+        canonicalizeThunkType(FT->getParamType(I), ParamAlign,
+                              /*Ret*/ false, ArgSizeBytes, Out);
     Arm64ArgTypes.push_back(Arm64Ty);
     X64ArgTypes.push_back(X64Ty);
+    ArgTranslations.push_back(ArgTranslation);
   }
 }
 
 void AArch64Arm64ECCallLowering::getThunkRetType(
     FunctionType *FT, AttributeList AttrList, raw_ostream &Out,
     Type *&Arm64RetTy, Type *&X64RetTy, SmallVectorImpl<Type *> &Arm64ArgTypes,
-    SmallVectorImpl<Type *> &X64ArgTypes, bool &HasSretPtr) {
+    SmallVectorImpl<Type *> &X64ArgTypes,
+    SmallVector<ThunkArgTranslation> &ArgTranslations, bool &HasSretPtr) {
   Type *T = FT->getReturnType();
 #if 0
   // FIXME: Need more information about argument size; see
@@ -240,13 +263,13 @@ void AArch64Arm64ECCallLowering::getThunkRetType(
         // that's a miscompile.)
         Type *SRetType = SRetAttr0.getValueAsType();
         Align SRetAlign = AttrList.getParamAlignment(0).valueOrOne();
-        Type *Arm64Ty, *X64Ty;
         canonicalizeThunkType(SRetType, SRetAlign, /*Ret*/ true, ArgSizeBytes,
-                              Out, Arm64Ty, X64Ty);
+                              Out);
         Arm64RetTy = VoidTy;
         X64RetTy = VoidTy;
         Arm64ArgTypes.push_back(FT->getParamType(0));
         X64ArgTypes.push_back(FT->getParamType(0));
+        ArgTranslations.push_back(ThunkArgTranslation::Direct);
         HasSretPtr = true;
         return;
       }
@@ -258,8 +281,10 @@ void AArch64Arm64ECCallLowering::getThunkRetType(
     return;
   }
 
-  canonicalizeThunkType(T, Align(), /*Ret*/ true, ArgSizeBytes, Out, Arm64RetTy,
-                        X64RetTy);
+  auto info =
+      canonicalizeThunkType(T, Align(), /*Ret*/ true, ArgSizeBytes, Out);
+  Arm64RetTy = info.Arm64Ty;
+  X64RetTy = info.X64Ty;
   if (X64RetTy->isPointerTy()) {
     // If the X64 type is canonicalized to a pointer, that means it's
     // passed/returned indirectly. For a return value, that means it's an
@@ -269,21 +294,33 @@ void AArch64Arm64ECCallLowering::getThunkRetType(
   }
 }
 
-void AArch64Arm64ECCallLowering::canonicalizeThunkType(
-    Type *T, Align Alignment, bool Ret, uint64_t ArgSizeBytes, raw_ostream &Out,
-    Type *&Arm64Ty, Type *&X64Ty) {
+ThunkArgInfo AArch64Arm64ECCallLowering::canonicalizeThunkType(
+    Type *T, Align Alignment, bool Ret, uint64_t ArgSizeBytes,
+    raw_ostream &Out) {
+
+  auto direct = [](Type *T) {
+    return ThunkArgInfo{T, T, ThunkArgTranslation::Direct};
+  };
+
+  auto bitcast = [this](Type *Arm64Ty, uint64_t SizeInBytes) {
+    return ThunkArgInfo{Arm64Ty,
+                        llvm::Type::getIntNTy(M->getContext(), SizeInBytes * 8),
+                        ThunkArgTranslation::Bitcast};
+  };
+
+  auto pointerIndirection = [this](Type *Arm64Ty) {
+    return ThunkArgInfo{Arm64Ty, PtrTy,
+                        ThunkArgTranslation::PointerIndirection};
+  };
+
   if (T->isFloatTy()) {
     Out << "f";
-    Arm64Ty = T;
-    X64Ty = T;
-    return;
+    return direct(T);
   }
 
   if (T->isDoubleTy()) {
     Out << "d";
-    Arm64Ty = T;
-    X64Ty = T;
-    return;
+    return direct(T);
   }
 
   if (T->isFloatingPointTy()) {
@@ -306,16 +343,14 @@ void AArch64Arm64ECCallLowering::canonicalizeThunkType(
       Out << (ElementTy->isFloatTy() ? "F" : "D") << TotalSizeBytes;
       if (Alignment.value() >= 16 && !Ret)
         Out << "a" << Alignment.value();
-      Arm64Ty = T;
       if (TotalSizeBytes <= 8) {
         // Arm64 returns small structs of float/double in float registers;
         // X64 uses RAX.
-        X64Ty = llvm::Type::getIntNTy(M->getContext(), TotalSizeBytes * 8);
+        return bitcast(T, TotalSizeBytes);
       } else {
         // Struct is passed directly on Arm64, but indirectly on X64.
-        X64Ty = PtrTy;
+        return pointerIndirection(T);
       }
-      return;
     } else if (T->isFloatingPointTy()) {
       report_fatal_error("Only 32 and 64 bit floating points are supported for "
                          "ARM64EC thunks");
@@ -324,9 +359,7 @@ void AArch64Arm64ECCallLowering::canonicalizeThunkType(
 
   if ((T->isIntegerTy() || T->isPointerTy()) && DL.getTypeSizeInBits(T) <= 64) {
     Out << "i8";
-    Arm64Ty = I64Ty;
-    X64Ty = I64Ty;
-    return;
+    return direct(I64Ty);
   }
 
   unsigned TypeSize = ArgSizeBytes;
@@ -338,13 +371,12 @@ void AArch64Arm64ECCallLowering::canonicalizeThunkType(
   if (Alignment.value() >= 16 && !Ret)
     Out << "a" << Alignment.value();
   // FIXME: Try to canonicalize Arm64Ty more thoroughly?
-  Arm64Ty = T;
   if (TypeSize == 1 || TypeSize == 2 || TypeSize == 4 || TypeSize == 8) {
     // Pass directly in an integer register
-    X64Ty = llvm::Type::getIntNTy(M->getContext(), TypeSize * 8);
+    return bitcast(T, TypeSize);
   } else {
     // Passed directly on Arm64, but indirectly on X64.
-    X64Ty = PtrTy;
+    return pointerIndirection(T);
   }
 }
 
@@ -355,8 +387,9 @@ Function *AArch64Arm64ECCallLowering::buildExitThunk(FunctionType *FT,
   SmallString<256> ExitThunkName;
   llvm::raw_svector_ostream ExitThunkStream(ExitThunkName);
   FunctionType *Arm64Ty, *X64Ty;
+  SmallVector<ThunkArgTranslation> ArgTranslations;
   getThunkType(FT, Attrs, Arm64ECThunkType::Exit, ExitThunkStream, Arm64Ty,
-               X64Ty);
+               X64Ty, ArgTranslations);
   if (Function *F = M->getFunction(ExitThunkName))
     return F;
 
@@ -387,6 +420,7 @@ Function *AArch64Arm64ECCallLowering::buildExitThunk(FunctionType *FT,
   SmallVector<Value *> Args;
 
   // Pass the called function in x9.
+  auto X64TyOffset = 1;
   Args.push_back(F->arg_begin());
 
   Type *RetTy = Arm64Ty->getReturnType();
@@ -396,10 +430,14 @@ Function *AArch64Arm64ECCallLowering::buildExitThunk(FunctionType *FT,
     // pointer.
     if (DL.getTypeStoreSize(RetTy) > 8) {
       Args.push_back(IRB.CreateAlloca(RetTy));
+      X64TyOffset++;
     }
   }
 
-  for (auto &Arg : make_range(F->arg_begin() + 1, F->arg_end())) {
+  for (auto [Arg, X64ArgType, ArgTranslation] : llvm::zip_equal(
+           make_range(F->arg_begin() + 1, F->arg_end()),
+           make_range(X64Ty->param_begin() + X64TyOffset, X64Ty->param_end()),
+           ArgTranslations)) {
     // Translate arguments from AArch64 calling convention to x86 calling
     // convention.
     //
@@ -414,18 +452,20 @@ Function *AArch64Arm64ECCallLowering::buildExitThunk(FunctionType *FT,
     // with an attribute.)
     //
     // The first argument is the called function, stored in x9.
-    if (Arg.getType()->isArrayTy() || Arg.getType()->isStructTy() ||
-        DL.getTypeStoreSize(Arg.getType()) > 8) {
+    if (ArgTranslation != ThunkArgTranslation::Direct) {
       Value *Mem = IRB.CreateAlloca(Arg.getType());
       IRB.CreateStore(&Arg, Mem);
-      if (DL.getTypeStoreSize(Arg.getType()) <= 8) {
+      if (ArgTranslation == ThunkArgTranslation::Bitcast) {
         Type *IntTy = IRB.getIntNTy(DL.getTypeStoreSizeInBits(Arg.getType()));
         Args.push_back(IRB.CreateLoad(IntTy, IRB.CreateBitCast(Mem, PtrTy)));
-      } else
+      } else {
+        assert(ArgTranslation == ThunkArgTranslation::PointerIndirection);
         Args.push_back(Mem);
+      }
     } else {
       Args.push_back(&Arg);
     }
+    assert(Args.back()->getType() == X64ArgType);
   }
   // FIXME: Transfer necessary attributes? sret? anything else?
 
@@ -459,8 +499,10 @@ Function *AArch64Arm64ECCallLowering::buildEntryThunk(Function *F) {
   SmallString<256> EntryThunkName;
   llvm::raw_svector_ostream EntryThunkStream(EntryThunkName);
   FunctionType *Arm64Ty, *X64Ty;
+  SmallVector<ThunkArgTranslation> ArgTranslations;
   getThunkType(F->getFunctionType(), F->getAttributes(),
-               Arm64ECThunkType::Entry, EntryThunkStream, Arm64Ty, X64Ty);
+               Arm64ECThunkType::Entry, EntryThunkStream, Arm64Ty, X64Ty,
+               ArgTranslations);
   if (Function *F = M->getFunction(EntryThunkName))
     return F;
 
@@ -472,7 +514,6 @@ Function *AArch64Arm64ECCallLowering::buildEntryThunk(Function *F) {
   // Copy MSVC, and always set up a frame pointer. (Maybe this isn't necessary.)
   Thunk->addFnAttr("frame-pointer", "all");
 
-  auto &DL = M->getDataLayout();
   BasicBlock *BB = BasicBlock::Create(M->getContext(), "", Thunk);
   IRBuilder<> IRB(BB);
 
@@ -481,24 +522,28 @@ Function *AArch64Arm64ECCallLowering::buildEntryThunk(Function *F) {
 
   bool TransformDirectToSRet = X64RetType->isVoidTy() && !RetTy->isVoidTy();
   unsigned ThunkArgOffset = TransformDirectToSRet ? 2 : 1;
-  unsigned PassthroughArgSize = F->isVarArg() ? 5 : Thunk->arg_size();
+  unsigned PassthroughArgSize =
+      (F->isVarArg() ? 5 : Thunk->arg_size()) - ThunkArgOffset;
+  assert(ArgTranslations.size() == F->isVarArg() ? 5 : PassthroughArgSize);
 
   // Translate arguments to call.
   SmallVector<Value *> Args;
-  for (unsigned i = ThunkArgOffset, e = PassthroughArgSize; i != e; ++i) {
-    Value *Arg = Thunk->getArg(i);
-    Type *ArgTy = Arm64Ty->getParamType(i - ThunkArgOffset);
-    if (ArgTy->isArrayTy() || ArgTy->isStructTy() ||
-        DL.getTypeStoreSize(ArgTy) > 8) {
+  for (unsigned i = 0; i != PassthroughArgSize; ++i) {
+    Value *Arg = Thunk->getArg(i + ThunkArgOffset);
+    Type *ArgTy = Arm64Ty->getParamType(i);
+    ThunkArgTranslation ArgTranslation = ArgTranslations[i];
+    if (ArgTranslation != ThunkArgTranslation::Direct) {
       // Translate array/struct arguments to the expected type.
-      if (DL.getTypeStoreSize(ArgTy) <= 8) {
+      if (ArgTranslation == ThunkArgTranslation::Bitcast) {
         Value *CastAlloca = IRB.CreateAlloca(ArgTy);
         IRB.CreateStore(Arg, IRB.CreateBitCast(CastAlloca, PtrTy));
         Arg = IRB.CreateLoad(ArgTy, CastAlloca);
       } else {
+        assert(ArgTranslation == ThunkArgTranslation::PointerIndirection);
         Arg = IRB.CreateLoad(ArgTy, IRB.CreateBitCast(Arg, PtrTy));
       }
     }
+    assert(Arg->getType() == ArgTy);
     Args.push_back(Arg);
   }
 
@@ -558,8 +603,10 @@ Function *AArch64Arm64ECCallLowering::buildEntryThunk(Function *F) {
 Function *AArch64Arm64ECCallLowering::buildGuestExitThunk(Function *F) {
   llvm::raw_null_ostream NullThunkName;
   FunctionType *Arm64Ty, *X64Ty;
+  SmallVector<ThunkArgTranslation> ArgTranslations;
   getThunkType(F->getFunctionType(), F->getAttributes(),
-               Arm64ECThunkType::GuestExit, NullThunkName, Arm64Ty, X64Ty);
+               Arm64ECThunkType::GuestExit, NullThunkName, Arm64Ty, X64Ty,
+               ArgTranslations);
   auto MangledName = getArm64ECMangledFunctionName(F->getName().str());
   assert(MangledName && "Can't guest exit to function that's already native");
   std::string ThunkName = *MangledName;