[llvm] [AArch64] Fix codegen for FEAT_F16F32DOT with SVE2/SME. (PR #192668)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: Sander de Smalen (sdesmalen-arm)

<details>
<summary>Changes</summary>

When compiling with +sve2/+sme, don't override to use Custom lowering for PARTIAL_REDUCE_FMLA when it previously determined the operation was legal due to +f16f32dot/+fp16fml.

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


2 Files Affected:

- (modified) llvm/lib/Target/AArch64/AArch64ISelLowering.cpp (+10-5) 
- (modified) llvm/test/CodeGen/AArch64/f16f32dot-fixed-length-fdot.ll (+14) 


``````````diff
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 9b34d9b385b4e..fbec7bb40a0da 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -2049,11 +2049,16 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
     if (Subtarget->hasSVE2() || Subtarget->hasSME()) {
       setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_FMLA, MVT::nxv4f32,
                                 MVT::nxv8f16, Legal);
-      // We can use SVE2p1 fdot to emulate the fixed-length variant.
-      setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_FMLA, MVT::v4f32,
-                                MVT::v8f16, Custom);
-      setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_FMLA, MVT::v2f32,
-                                MVT::v4f16, Custom);
+
+      // We can use SVE2p1 fdot or SVE2 fmlalb/t to emulate the fixed-length
+      // variant (unless NEON fdot is natively available).
+      if (!Subtarget->isNeonAvailable() ||
+          (!Subtarget->hasF16F32DOT() && !Subtarget->hasFP16FML())) {
+        setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_FMLA, MVT::v4f32,
+                                  MVT::v8f16, Custom);
+        setPartialReduceMLAAction(ISD::PARTIAL_REDUCE_FMLA, MVT::v2f32,
+                                  MVT::v4f16, Custom);
+      }
     }
 
     if (Subtarget->hasBF16())
diff --git a/llvm/test/CodeGen/AArch64/f16f32dot-fixed-length-fdot.ll b/llvm/test/CodeGen/AArch64/f16f32dot-fixed-length-fdot.ll
index 5b1fe56244fbe..ab72173755a61 100644
--- a/llvm/test/CodeGen/AArch64/f16f32dot-fixed-length-fdot.ll
+++ b/llvm/test/CodeGen/AArch64/f16f32dot-fixed-length-fdot.ll
@@ -3,6 +3,8 @@
 ; RUN: llc -mattr=+fp16fml < %s | FileCheck %s --check-prefixes=CHECK-FMLAL
 ; RUN: llc -mattr=+f16f32dot < %s | FileCheck %s --check-prefixes=CHECK-FDOT
 ; RUN: llc -mattr=+f16f32dot,+fp16fml < %s | FileCheck %s --check-prefixes=CHECK-FDOT
+; RUN: llc -mattr=+f16f32dot,+sve2 < %s | FileCheck %s --check-prefixes=CHECK-FDOT
+; RUN: llc -mattr=+f16f32dot,+sme -force-streaming < %s | FileCheck %s --check-prefixes=CHECK-SVE-FMLALBT
 
 target triple = "aarch64-linux-gnu"
 
@@ -29,6 +31,12 @@ define <4 x float> @fixed_fdot_wide(<4 x float> %acc, <8 x half> %a, <8 x half>
 ; CHECK-FDOT:       // %bb.0: // %entry
 ; CHECK-FDOT-NEXT:    fdot v0.4s, v1.8h, v2.8h
 ; CHECK-FDOT-NEXT:    ret
+;
+; CHECK-SVE-FMLALBT-LABEL: fixed_fdot_wide:
+; CHECK-SVE-FMLALBT:       // %bb.0: // %entry
+; CHECK-SVE-FMLALBT-NEXT:    fmlalb z0.s, z1.h, z2.h
+; CHECK-SVE-FMLALBT-NEXT:    fmlalt z0.s, z1.h, z2.h
+; CHECK-SVE-FMLALBT-NEXT:    ret
 entry:
   %a.wide = fpext <8 x half> %a to <8 x float>
   %b.wide = fpext <8 x half> %b to <8 x float>
@@ -58,6 +66,12 @@ define <2 x float> @fixed_fdot(<2 x float> %acc, <4 x half> %a, <4 x half> %b) {
 ; CHECK-FDOT:       // %bb.0: // %entry
 ; CHECK-FDOT-NEXT:    fdot v0.2s, v1.4h, v2.4h
 ; CHECK-FDOT-NEXT:    ret
+;
+; CHECK-SVE-FMLALBT-LABEL: fixed_fdot:
+; CHECK-SVE-FMLALBT:       // %bb.0: // %entry
+; CHECK-SVE-FMLALBT-NEXT:    fmlalb z0.s, z1.h, z2.h
+; CHECK-SVE-FMLALBT-NEXT:    fmlalt z0.s, z1.h, z2.h
+; CHECK-SVE-FMLALBT-NEXT:    ret
 entry:
   %a.wide = fpext <4 x half> %a to <4 x float>
   %b.wide = fpext <4 x half> %b to <4 x float>

``````````

</details>


https://github.com/llvm/llvm-project/pull/192668