[CIR][ABI][AArch64] support for return struct types greater than 128 bits

clang/include/clang/CIR/ABIArgInfo.h

@@ -103,6 +103,9 @@ class ABIArgInfo {
   bool InReg : 1;          // isDirect() || isExtend() || isIndirect()
   bool CanBeFlattened : 1; // isDirect()
   bool SignExt : 1;        // isExtend()
+  bool IndirectByVal : 1;  // isIndirect()
+  bool IndirectRealign : 1; // isIndirect()
+  bool SRetAfterThis : 1;   // isIndirect()
 
   bool canHavePaddingType() const {
     return isDirect() || isExtend() || isIndirect() || isIndirectAliased() ||
@@ -195,6 +198,43 @@ class ABIArgInfo {
 
   static ABIArgInfo getIgnore() { return ABIArgInfo(Ignore); }
 
+  static ABIArgInfo getIndirect(unsigned Alignment, bool ByVal = true,
+                                bool Realign = false,
+                                mlir::Type Padding = nullptr) {
+    auto AI = ABIArgInfo(Indirect);
+    AI.setIndirectAlign(Alignment);
+    AI.setIndirectByVal(ByVal);
+    AI.setIndirectRealign(Realign);
+    AI.setSRetAfterThis(false);
+    AI.setPaddingType(Padding);
+    return AI;
+  }
+
+  void setIndirectAlign(unsigned align) {
+    assert((isIndirect() || isIndirectAliased()) && "Invalid kind!");
+    IndirectAttr.Align = align;
+  }
+
+  void setIndirectByVal(bool IBV) {
+    assert(isIndirect() && "Invalid kind!");
+    IndirectByVal = IBV;
+  }
+
+  void setIndirectRealign(bool IR) {
+    assert((isIndirect() || isIndirectAliased()) && "Invalid kind!");
+    IndirectRealign = IR;
+  }
+
+  void setSRetAfterThis(bool AfterThis) {
+    assert(isIndirect() && "Invalid kind!");
+    SRetAfterThis = AfterThis;
+  }
+
+  bool isSRetAfterThis() const {
+    assert(isIndirect() && "Invalid kind!");
+    return SRetAfterThis;
+  }
+
   Kind getKind() const { return TheKind; }
   bool isDirect() const { return TheKind == Direct; }
   bool isInAlloca() const { return TheKind == InAlloca; }

clang/include/clang/CIR/MissingFeatures.h

@@ -271,6 +271,7 @@ struct MissingFeatures {
   static bool ABIParameterCoercion() { return false; }
   static bool ABIPointerParameterAttrs() { return false; }
   static bool ABITransparentUnionHandling() { return false; }
+  static bool ABIPotentialArgAccess() { return false; }
 
   //-- Missing AST queries
 

clang/lib/CIR/Dialect/Transforms/TargetLowering/ABIInfo.cpp

@@ -42,5 +42,12 @@ bool ABIInfo::isPromotableIntegerTypeForABI(Type Ty) const {
   return false;
 }
 
+::cir::ABIArgInfo ABIInfo::getNaturalAlignIndirect(mlir::Type Ty, bool ByVal,
+                                                   bool Realign,
+                                                   mlir::Type Padding) const {
+  return ::cir::ABIArgInfo::getIndirect(getContext().getTypeAlign(Ty), ByVal,
+                                        Realign, Padding);
+}
+
 } // namespace cir
 } // namespace mlir

clang/lib/CIR/Dialect/Transforms/TargetLowering/ABIInfo.h

@@ -50,6 +50,10 @@ class ABIInfo {
   // Implement the Type::IsPromotableIntegerType for ABI specific needs. The
   // only difference is that this considers bit-precise integer types as well.
   bool isPromotableIntegerTypeForABI(Type Ty) const;
+
+  ::cir::ABIArgInfo getNaturalAlignIndirect(mlir::Type Ty, bool ByVal = true,
+                                            bool Realign = false,
+                                            mlir::Type Padding = {}) const;
 };
 
 } // namespace cir

clang/lib/CIR/Dialect/Transforms/TargetLowering/CIRToCIRArgMapping.h

@@ -29,6 +29,7 @@ namespace cir {
 /// LoweringFunctionInfo should be passed to actual CIR function.
 class CIRToCIRArgMapping {
   static const unsigned InvalidIndex = ~0U;
+  unsigned SRetArgNo;
   unsigned TotalIRArgs;
 
   /// Arguments of CIR function corresponding to single CIR argument.
@@ -51,7 +52,8 @@ class CIRToCIRArgMapping {
 public:
   CIRToCIRArgMapping(const CIRLowerContext &context,
                      const LowerFunctionInfo &FI, bool onlyRequiredArgs = false)
-      : ArgInfo(onlyRequiredArgs ? FI.getNumRequiredArgs() : FI.arg_size()) {
+      : SRetArgNo(InvalidIndex),
+        ArgInfo(onlyRequiredArgs ? FI.getNumRequiredArgs() : FI.arg_size()) {
     construct(context, FI, onlyRequiredArgs);
   };
 
@@ -69,7 +71,8 @@ class CIRToCIRArgMapping {
     const ::cir::ABIArgInfo &RetAI = FI.getReturnInfo();
 
     if (RetAI.getKind() == ::cir::ABIArgInfo::Indirect) {
-      cir_cconv_unreachable("NYI");
+      SwapThisWithSRet = RetAI.isSRetAfterThis();
+      SRetArgNo = SwapThisWithSRet ? 1 : IRArgNo++;
     }
 
     unsigned ArgNo = 0;
@@ -100,6 +103,11 @@ class CIRToCIRArgMapping {
         }
         break;
       }
+      case ::cir::ABIArgInfo::Indirect:
+      case ::cir::ABIArgInfo::IndirectAliased:
+        IRArgs.NumberOfArgs = 1;
+        break;
+
       default:
         cir_cconv_unreachable("Missing ABIArgInfo::Kind");
       }
@@ -130,6 +138,13 @@ class CIRToCIRArgMapping {
     return std::make_pair(ArgInfo[ArgNo].FirstArgIndex,
                           ArgInfo[ArgNo].NumberOfArgs);
   }
+
+  bool hasSRetArg() const { return SRetArgNo != InvalidIndex; }
+
+  unsigned getSRetArgNo() const {
+    assert(hasSRetArg());
+    return SRetArgNo;
+  }
 };
 
 } // namespace cir

clang/lib/CIR/Dialect/Transforms/TargetLowering/LowerCall.cpp

@@ -157,6 +157,8 @@ void LowerModule::constructAttributeList(StringRef Name,
     cir_cconv_assert(!::cir::MissingFeatures::noFPClass());
     break;
   case ABIArgInfo::Ignore:
+  case ABIArgInfo::Indirect:
+    cir_cconv_assert(!::cir::MissingFeatures::ABIPotentialArgAccess());
     break;
   default:
     cir_cconv_unreachable("Missing ABIArgInfo::Kind");

clang/lib/CIR/Dialect/Transforms/TargetLowering/LowerFunction.cpp

@@ -10,7 +10,6 @@
 // are adapted to operate on the CIR dialect, however.
 //
 //===----------------------------------------------------------------------===//
-
 #include "LowerFunction.h"
 #include "CIRToCIRArgMapping.h"
 #include "LowerCall.h"
@@ -433,6 +432,23 @@ LowerFunction::buildFunctionProlog(const LowerFunctionInfo &FI, FuncOp Fn,
   return success();
 }
 
+mlir::cir::AllocaOp findAlloca(Operation *op) {
+  if (!op)
+    return {};
+
+  if (auto al = dyn_cast<mlir::cir::AllocaOp>(op)) {
+    return al;
+  } else if (auto ret = dyn_cast<mlir::cir::ReturnOp>(op)) {
+    auto vals = ret.getInput();
+    if (vals.size() == 1)
+      return findAlloca(vals[0].getDefiningOp());
+  } else if (auto load = dyn_cast<mlir::cir::LoadOp>(op)) {
+    return findAlloca(load.getAddr().getDefiningOp());
+  }
+
+  return {};
+}
+
 LogicalResult LowerFunction::buildFunctionEpilog(const LowerFunctionInfo &FI) {
   // NOTE(cir): no-return, naked, and no result functions should be handled in
   // CIRGen.
@@ -446,6 +462,27 @@ LogicalResult LowerFunction::buildFunctionEpilog(const LowerFunctionInfo &FI) {
   case ABIArgInfo::Ignore:
     break;
 
+  case ABIArgInfo::Indirect: {
+    Value RVAddr = {};
+    CIRToCIRArgMapping IRFunctionArgs(LM.getContext(), FI, true);
+    if (IRFunctionArgs.hasSRetArg()) {
+      auto &entry = NewFn.getBody().front();
+      RVAddr = entry.getArgument(IRFunctionArgs.getSRetArgNo());
+    }
+
+    if (RVAddr) {
+      mlir::PatternRewriter::InsertionGuard guard(rewriter);
+      NewFn->walk([&](ReturnOp ret) {
+        if (auto al = findAlloca(ret)) {
+          rewriter.replaceAllUsesWith(al.getResult(), RVAddr);
+          rewriter.eraseOp(al);
+          rewriter.replaceOpWithNewOp<ReturnOp>(ret);
+        }
+      });
+    }
+    break;
+  }
+
   case ABIArgInfo::Extend:
   case ABIArgInfo::Direct:
     // FIXME(cir): Should we call ConvertType(RetTy) here?
@@ -517,6 +554,15 @@ LogicalResult LowerFunction::generateCode(FuncOp oldFn, FuncOp newFn,
   Block *srcBlock = &oldFn.getBody().front();
   Block *dstBlock = &newFn.getBody().front();
 
+  // Ensure both blocks have the same number of arguments in order to
+  // safely merge them.
+  CIRToCIRArgMapping IRFunctionArgs(LM.getContext(), FnInfo, true);
+  if (IRFunctionArgs.hasSRetArg()) {
+    auto dstIndex = IRFunctionArgs.getSRetArgNo();
+    auto retArg = dstBlock->getArguments()[dstIndex];
+    srcBlock->insertArgument(dstIndex, retArg.getType(), retArg.getLoc());
+  }
+
   // Migrate function body to new ABI-aware function.
   rewriter.inlineRegionBefore(oldFn.getBody(), newFn.getBody(),
                               newFn.getBody().end());

clang/lib/CIR/Dialect/Transforms/TargetLowering/LowerTypes.cpp

@@ -50,6 +50,7 @@ FuncType LowerTypes::getFunctionType(const LowerFunctionInfo &FI) {
     resultType = retAI.getCoerceToType();
     break;
   case ::cir::ABIArgInfo::Ignore:
+  case ::cir::ABIArgInfo::Indirect:
     resultType = VoidType::get(getMLIRContext());
     break;
   default:
@@ -60,7 +61,11 @@ FuncType LowerTypes::getFunctionType(const LowerFunctionInfo &FI) {
   SmallVector<Type, 8> ArgTypes(IRFunctionArgs.totalIRArgs());
 
   // Add type for sret argument.
-  cir_cconv_assert(!::cir::MissingFeatures::sretArgs());
+  if (IRFunctionArgs.hasSRetArg()) {
+    mlir::Type ret = FI.getReturnType();
+    ArgTypes[IRFunctionArgs.getSRetArgNo()] =
+        mlir::cir::PointerType::get(getMLIRContext(), ret);
+  }
 
   // Add type for inalloca argument.
   cir_cconv_assert(!::cir::MissingFeatures::inallocaArgs());

clang/lib/CIR/Dialect/Transforms/TargetLowering/Targets/AArch64.cpp

@@ -137,7 +137,7 @@ ABIArgInfo AArch64ABIInfo::classifyReturnType(Type RetTy,
     cir_cconv_unreachable("NYI");
   }
 
-  cir_cconv_unreachable("NYI");
+  return getNaturalAlignIndirect(RetTy);
 }
 
 ABIArgInfo

clang/test/CIR/CallConvLowering/AArch64/aarch64-cc-structs.c

@@ -21,6 +21,12 @@ typedef struct {
   int64_t b;
 } EQ_128;
 
+typedef struct {
+  int64_t a;
+  int64_t b;
+  int64_t c;
+} GT_128;
+
 // CHECK: cir.func {{.*@ret_lt_64}}() -> !u16i
 // CHECK:   %[[#V0:]] = cir.alloca !ty_LT_64_, !cir.ptr<!ty_LT_64_>, ["__retval"]
 // CHECK:   %[[#V1:]] = cir.cast(bitcast, %[[#V0]] : !cir.ptr<!ty_LT_64_>), !cir.ptr<!u16i>
@@ -60,3 +66,10 @@ EQ_128 ret_eq_128() {
   EQ_128 x;
   return x;
 }
+
+// CHECK:     cir.func {{.*@ret_gt_128}}(%arg0: !cir.ptr<!ty_GT_128_> 
+// CHECK-NOT:   cir.return {{%.*}}
+GT_128 ret_gt_128() {
+  GT_128 x;
+  return x;
+}