[AArch64] Do not split bfloat HFA args between regs and stack

[ostannard]

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.

[llvmbot]

@llvm/pr-subscribers-backend-aarch64

Author: Oliver Stannard (ostannard)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/128909.diff

2 Files Affected:

• (modified) llvm/lib/Target/AArch64/AArch64CallingConvention.cpp (+1-1)
• (modified) llvm/test/CodeGen/AArch64/argument-blocks.ll (+7)

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
+}

[davemgreen]

LGTM