[CIR][ABI][Lowering] covers return struct case with coercion through memory (#1059)

This PR covers one more case for return values of struct type, where
`memcpy` is emitted.

Author: gitoleg

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

@@ -257,6 +257,23 @@ Value emitAddressAtOffset(LowerFunction &LF, Value addr,
   return addr;
 }
 
+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 {};
+}
+
 /// After the calling convention is lowered, an ABI-agnostic type might have to
 /// be loaded back to its ABI-aware couterpart so it may be returned. If they
 /// differ, we have to do a coerced load. A coerced load, which means to load a
@@ -305,6 +322,31 @@ Value castReturnValue(Value Src, Type Ty, LowerFunction &LF) {
     return LF.getRewriter().create<LoadOp>(Src.getLoc(), Cast);
   }
 
+  // Otherwise do coercion through memory.
+  if (auto addr = findAlloca(Src.getDefiningOp())) {
+    auto &rewriter = LF.getRewriter();
+    auto *ctxt = LF.LM.getMLIRContext();
+    auto ptrTy = PointerType::get(ctxt, Ty);
+    auto voidPtr = PointerType::get(ctxt, mlir::cir::VoidType::get(ctxt));
+
+    // insert alloca near the previuos one
+    auto point = rewriter.saveInsertionPoint();
+    rewriter.setInsertionPointAfter(addr);
+    auto align = LF.LM.getDataLayout().getABITypeAlign(Ty);
+    auto alignAttr = rewriter.getI64IntegerAttr(align.value());
+    auto tmp =
+        rewriter.create<AllocaOp>(Src.getLoc(), ptrTy, Ty, "tmp", alignAttr);
+    rewriter.restoreInsertionPoint(point);
+
+    auto srcVoidPtr = createBitcast(addr, voidPtr, LF);
+    auto dstVoidPtr = createBitcast(tmp, voidPtr, LF);
+    auto i64Ty = IntType::get(ctxt, 64, false);
+    auto len = rewriter.create<ConstantOp>(
+        Src.getLoc(), IntAttr::get(i64Ty, SrcSize.getFixedValue()));
+    rewriter.create<MemCpyOp>(Src.getLoc(), dstVoidPtr, srcVoidPtr, len);
+    return rewriter.create<LoadOp>(Src.getLoc(), tmp.getResult());
+  }
+
   cir_cconv_unreachable("NYI");
 }
 
@@ -532,23 +574,6 @@ 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.

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

@@ -74,6 +74,32 @@ GT_128 ret_gt_128() {
   return x;
 }
 
+typedef struct {
+  int a;
+  int b;
+  int c;  
+} S;
+
+// CHECK: cir.func {{.*@retS}}() -> !cir.array<!u64i x 2>
+// CHECK:   %[[#V0:]] = cir.alloca !ty_S, !cir.ptr<!ty_S>, ["__retval"] {alignment = 4 : i64}
+// CHECK:   %[[#V1:]] = cir.alloca !cir.array<!u64i x 2>, !cir.ptr<!cir.array<!u64i x 2>>, ["tmp"] {alignment = 8 : i64}
+// CHECK:   %[[#V2:]] = cir.cast(bitcast, %[[#V0]] : !cir.ptr<!ty_S>), !cir.ptr<!void>
+// CHECK:   %[[#V3:]] = cir.cast(bitcast, %[[#V1]] : !cir.ptr<!cir.array<!u64i x 2>>), !cir.ptr<!void>
+// CHECK:   %[[#V4:]] = cir.const #cir.int<12> : !u64i
+// CHECK:   cir.libc.memcpy %[[#V4]] bytes from %[[#V2]] to %[[#V3]] : !u64i, !cir.ptr<!void> -> !cir.ptr<!void>
+// CHECK:   %[[#V5:]] = cir.load %[[#V1]] : !cir.ptr<!cir.array<!u64i x 2>>, !cir.array<!u64i x 2>
+// CHECK:   cir.return %[[#V5]] : !cir.array<!u64i x 2>
+
+// LLVM: [2 x i64] @retS() 
+// LLVM:   %[[#V1:]] = alloca %struct.S, i64 1, align 4
+// LLVM:   %[[#V2:]] = alloca [2 x i64], i64 1, align 8
+// LLVM:   call void @llvm.memcpy.p0.p0.i64(ptr %[[#V2]], ptr %[[#V1]], i64 12, i1 false)
+// LLVM:   %[[#V3:]] = load [2 x i64], ptr %[[#V2]], align 8
+// LLVM:   ret [2 x i64] %[[#V3]]
+S retS() {
+  S s;
+  return s;
+}
 // CHECK: cir.func {{.*@pass_lt_64}}(%arg0: !u64
 // CHECK:   %[[#V0:]] = cir.alloca !ty_LT_64_, !cir.ptr<!ty_LT_64_>
 // CHECK:   %[[#V1:]] = cir.cast(integral, %arg0 : !u64i), !u16i