[llvm-branch-commits] [llvm] [AArch64][SVE] Support lowering masked loads of <4 x bf16> and <8 x bf16> (PR #208768)

Shanzhi Chen via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Fri Jul 10 09:19:09 PDT 2026


https://github.com/chenshanzhi created https://github.com/llvm/llvm-project/pull/208768

Add support for lowering masked loads of <4 x bf16> and <8 x bf16> when target features contain "+sve".


>From fe4bb07ec2e748f9bfd8aeb4838e0242e9ee5357 Mon Sep 17 00:00:00 2001
From: Shanzhi Chen <chenshanzhi at huawei.com>
Date: Fri, 10 Jul 2026 22:48:00 +0800
Subject: [PATCH 1/5] [AArch64][SVE] Support lowering masked loads of <4 x
 bf16> and <8 x bf16>

Add support for lowering masked loads of <4 x bf16> and <8 x bf16> when
target features contain "+sve".
---
 ...sve-fixed-length-masked-64-128bit-loads.ll | 37 +++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/llvm/test/CodeGen/AArch64/sve-fixed-length-masked-64-128bit-loads.ll b/llvm/test/CodeGen/AArch64/sve-fixed-length-masked-64-128bit-loads.ll
index 91af81c617c41..b30a3a9fc933d 100644
--- a/llvm/test/CodeGen/AArch64/sve-fixed-length-masked-64-128bit-loads.ll
+++ b/llvm/test/CodeGen/AArch64/sve-fixed-length-masked-64-128bit-loads.ll
@@ -35,6 +35,20 @@ define <8 x half> @masked_load_v8f16(ptr %src, <8 x i1> %mask) {
   ret <8 x half> %load
 }
 
+define <8 x bfloat> @masked_load_v8bf16(ptr %src, <8 x i1> %mask) #0 {
+; CHECK-LABEL: masked_load_v8bf16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ushll v0.8h, v0.8b, #0
+; CHECK-NEXT:    ptrue p0.h, vl8
+; CHECK-NEXT:    shl v0.8h, v0.8h, #15
+; CHECK-NEXT:    cmpne p1.h, p0/z, z0.h, #0
+; CHECK-NEXT:    ld1h { z0.h }, p1/z, [x0]
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
+; CHECK-NEXT:    ret
+  %load = call <8 x bfloat> @llvm.masked.load.v8bf16(ptr %src, i32 8, <8 x i1> %mask, <8 x bfloat> zeroinitializer)
+  ret <8 x bfloat> %load
+}
+
 define <4 x float> @masked_load_v4f32(ptr %src, <4 x i1> %mask) {
 ; CHECK-LABEL: masked_load_v4f32:
 ; CHECK:       // %bb.0:
@@ -134,3 +148,26 @@ define <4 x half> @masked_load_v4f16(ptr %ap, ptr %bp) {
   %load = call <4 x half> @llvm.masked.load.v4f16(ptr %ap, i32 2, <4 x i1> %mask, <4 x half> zeroinitializer)
   ret <4 x half> %load
 }
+
+define <4 x bfloat> @masked_load_v4bf16(ptr %ap, ptr %bp) #0 {
+; CHECK-LABEL: masked_load_v4bf16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ldr d0, [x0]
+; CHECK-NEXT:    ldr d1, [x1]
+; CHECK-NEXT:    ptrue p0.h, vl4
+; CHECK-NEXT:    shll v1.4s, v1.4h, #16
+; CHECK-NEXT:    shll v0.4s, v0.4h, #16
+; CHECK-NEXT:    fcmeq v0.4s, v0.4s, v1.4s
+; CHECK-NEXT:    xtn v0.4h, v0.4s
+; CHECK-NEXT:    cmpne p1.h, p0/z, z0.h, #0
+; CHECK-NEXT:    ld1h { z0.h }, p1/z, [x0]
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $z0
+; CHECK-NEXT:    ret
+  %a = load <4 x bfloat>, ptr %ap
+  %b = load <4 x bfloat>, ptr %bp
+  %mask = fcmp oeq <4 x bfloat> %a, %b
+  %load = call <4 x bfloat> @llvm.masked.load.v4bf16(ptr %ap, i32 2, <4 x i1> %mask, <4 x bfloat> zeroinitializer)
+  ret <4 x bfloat> %load
+}
+
+attributes #0 = { "target-features"="+sve,+bf16" }

>From 9a3cde782d98797769db233ca6ba4d40bffe593e Mon Sep 17 00:00:00 2001
From: Shanzhi Chen <chenshanzhi at huawei.com>
Date: Fri, 10 Jul 2026 23:14:17 +0800
Subject: [PATCH 2/5] fixup! [AArch64][SVE] Support lowering masked loads of <4
 x bf16> and <8 x bf16>

---
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 8935beb8b3472..7bb41e1d6dbe1 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -8855,6 +8855,7 @@ bool AArch64TargetLowering::useSVEForFixedLengthVectorVT(
   case MVT::f16:
   case MVT::f32:
   case MVT::f64:
+  case MVT::bf16:
     break;
   }
 

>From 23d4b850a27b01a4a7452bd1d37725cdae4cd48a Mon Sep 17 00:00:00 2001
From: Shanzhi Chen <chenshanzhi at huawei.com>
Date: Fri, 10 Jul 2026 23:41:08 +0800
Subject: [PATCH 3/5] fixup! [AArch64][SVE] Support lowering masked loads of <4
 x bf16> and <8 x bf16>

---
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 7bb41e1d6dbe1..043ee230be3ae 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -2522,7 +2522,7 @@ void AArch64TargetLowering::addTypeForFixedLengthSVE(MVT VT) {
 
   // Mark floating-point truncating stores/extending loads as having custom
   // lowering
-  if (VT.isFloatingPoint()) {
+  if (VT.isFloatingPoint() && !VT.isVectorOf(MVT::bf16)) {
     MVT InnerVT = VT.changeVectorElementType(MVT::f16);
     while (InnerVT != VT) {
       setTruncStoreAction(VT, InnerVT, Custom);

>From f3b83ffde5ad5f4024e6a267c27f0795f5bead7c Mon Sep 17 00:00:00 2001
From: Shanzhi Chen <chenshanzhi at huawei.com>
Date: Fri, 10 Jul 2026 23:52:27 +0800
Subject: [PATCH 4/5] fixup! [AArch64][SVE] Support lowering masked loads of <4
 x bf16> and <8 x bf16>

---
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 043ee230be3ae..115f8fb5d04ac 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -8863,6 +8863,10 @@ bool AArch64TargetLowering::useSVEForFixedLengthVectorVT(
   if (OverrideNEON && (VT.is128BitVector() || VT.is64BitVector()))
     return Subtarget->isSVEorStreamingSVEAvailable();
 
+  // Don't use SVE for bf16 vectors other than v8bf16 and v4bf16.
+  if (VT.isVectorOf(MVT::bf16))
+    return false;
+
   // Ensure NEON MVTs only belong to a single register class.
   if (VT.getFixedSizeInBits() <= 128)
     return false;

>From e924ff2b23a2a9753d8aae84ee464fc0c8e547c4 Mon Sep 17 00:00:00 2001
From: Shanzhi Chen <chenshanzhi at huawei.com>
Date: Sat, 11 Jul 2026 00:08:18 +0800
Subject: [PATCH 5/5] fixup! [AArch64][SVE] Support lowering masked loads of <4
 x bf16> and <8 x bf16>

---
 .../CodeGen/AArch64/sve2p1-vector-shuffles.ll | 28 ++-----------------
 1 file changed, 3 insertions(+), 25 deletions(-)

diff --git a/llvm/test/CodeGen/AArch64/sve2p1-vector-shuffles.ll b/llvm/test/CodeGen/AArch64/sve2p1-vector-shuffles.ll
index a2d4ababd919f..d13837f1c3da6 100644
--- a/llvm/test/CodeGen/AArch64/sve2p1-vector-shuffles.ll
+++ b/llvm/test/CodeGen/AArch64/sve2p1-vector-shuffles.ll
@@ -82,32 +82,10 @@ define void @dupq_bf16_256b(ptr %addr) #0 {
 ;
 ; SME-LABEL: dupq_bf16_256b:
 ; SME:       // %bb.0:
-; SME-NEXT:    ldp q1, q0, [x0]
-; SME-NEXT:    str q0, [sp, #-64]!
-; SME-NEXT:    .cfi_def_cfa_offset 64
-; SME-NEXT:    ldr h0, [sp, #4]
-; SME-NEXT:    str q1, [sp, #32]
-; SME-NEXT:    str h0, [sp, #30]
-; SME-NEXT:    str h0, [sp, #28]
-; SME-NEXT:    str h0, [sp, #26]
-; SME-NEXT:    str h0, [sp, #24]
-; SME-NEXT:    str h0, [sp, #22]
-; SME-NEXT:    str h0, [sp, #20]
-; SME-NEXT:    str h0, [sp, #18]
-; SME-NEXT:    str h0, [sp, #16]
-; SME-NEXT:    ldr h0, [sp, #36]
-; SME-NEXT:    ldr q1, [sp, #16]
-; SME-NEXT:    str h0, [sp, #62]
-; SME-NEXT:    str h0, [sp, #60]
-; SME-NEXT:    str h0, [sp, #58]
-; SME-NEXT:    str h0, [sp, #56]
-; SME-NEXT:    str h0, [sp, #54]
-; SME-NEXT:    str h0, [sp, #52]
-; SME-NEXT:    str h0, [sp, #50]
-; SME-NEXT:    str h0, [sp, #48]
-; SME-NEXT:    ldr q0, [sp, #48]
+; SME-NEXT:    ldp q0, q1, [x0]
+; SME-NEXT:    mov z0.h, z0.h[2]
+; SME-NEXT:    mov z1.h, z1.h[2]
 ; SME-NEXT:    stp q0, q1, [x0]
-; SME-NEXT:    add sp, sp, #64
 ; SME-NEXT:    ret
   %load = load <16 x bfloat>, ptr %addr
   %splat.lanes = shufflevector <16 x bfloat> %load, <16 x bfloat> poison, <16 x i32> <i32  2, i32  2, i32  2, i32  2, i32  2, i32  2, i32  2, i32  2,



More information about the llvm-branch-commits mailing list