[llvm] [AArch64] Do not split bfloat HFA args between regs and stack (PR #128909)

[Oliver Stannard via llvm-commits]

https://github.com/ostannard created https://github.com/llvm/llvm-project/pull/128909

In AAPCS64, __fp16 and __bf16 share the same machine type, so they should be treated the same way for argument passing. In particular, arrays of them need to be treated as homogeneous aggregates, and not split between registers and the stack.

>From b606bea78bb74a304381a8d28bb7f7d3f5e92c2e Mon Sep 17 00:00:00 2001
From: Oliver Stannard <oliver.stannard at arm.com>
Date: Wed, 26 Feb 2025 16:36:42 +0000
Subject: [PATCH] [AArch64] Do not split bfloat HFA args between regs and stack

In AAPCS64, __fp16 and __bf16 share the same machine type, so they
should be treated the same way for argument passing. In particular,
arrays of them need to be treated as homogeneous aggregates, and not
split between registers and the stack.
---
 llvm/lib/Target/AArch64/AArch64CallingConvention.cpp | 2 +-
 llvm/test/CodeGen/AArch64/argument-blocks.ll         | 7 +++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AArch64/AArch64CallingConvention.cpp b/llvm/lib/Target/AArch64/AArch64CallingConvention.cpp
index 991d710c979b9..787a1a83613c9 100644
--- a/llvm/lib/Target/AArch64/AArch64CallingConvention.cpp
+++ b/llvm/lib/Target/AArch64/AArch64CallingConvention.cpp
@@ -142,7 +142,7 @@ static bool CC_AArch64_Custom_Block(unsigned &ValNo, MVT &ValVT, MVT &LocVT,
   ArrayRef<MCPhysReg> RegList;
   if (LocVT.SimpleTy == MVT::i64 || (IsDarwinILP32 && LocVT.SimpleTy == MVT::i32))
     RegList = XRegList;
-  else if (LocVT.SimpleTy == MVT::f16)
+  else if (LocVT.SimpleTy == MVT::f16 || LocVT.SimpleTy == MVT::bf16)
     RegList = HRegList;
   else if (LocVT.SimpleTy == MVT::f32 || LocVT.is32BitVector())
     RegList = SRegList;
diff --git a/llvm/test/CodeGen/AArch64/argument-blocks.ll b/llvm/test/CodeGen/AArch64/argument-blocks.ll
index b5374ca8ced53..8cef28a38970d 100644
--- a/llvm/test/CodeGen/AArch64/argument-blocks.ll
+++ b/llvm/test/CodeGen/AArch64/argument-blocks.ll
@@ -195,3 +195,10 @@ define half @test_f16_blocked([7 x double], [2 x half] %in) {
   %val = extractvalue [2 x half] %in, 0
   ret half %val
 }
+
+define bfloat @test_bf16_blocked([7 x double], [2 x bfloat] %in) {
+; CHECK-LABEL: test_bf16_blocked:
+; CHECK: ldr h0, [sp]
+  %val = extractvalue [2 x bfloat] %in, 0
+  ret bfloat %val
+}