[llvm] [SVE2.1][AArch64] Fix crash for bfadd with scalable vector types (PR #107274)

[via llvm-commits]

https://github.com/vfdff created https://github.com/llvm/llvm-project/pull/107274

Fix https://github.com/llvm/llvm-project/issues/104416

>From 6bff2ee2b6b5c089fdd82b276c104607483dbdbb Mon Sep 17 00:00:00 2001
From: zhongyunde 00443407 <zhongyunde at huawei.com>
Date: Tue, 3 Sep 2024 20:54:07 -0400
Subject: [PATCH] [SVE2.1][AArch64] Fix crash for bfadd with scalable vector
 types

Fix https://github.com/llvm/llvm-project/issues/104416
---
 .../Target/AArch64/AArch64ISelLowering.cpp    |  3 +-
 llvm/lib/Target/AArch64/SVEInstrFormats.td    | 20 ++++++++----
 llvm/test/CodeGen/AArch64/sve2p1-bf16-bin.ll  | 32 +++++++++++++++++++
 3 files changed, 47 insertions(+), 8 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/sve2p1-bf16-bin.ll

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 215f30128e7038..18448c9b3d4cda 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -1585,7 +1585,8 @@ AArch64TargetLowering::AArch64TargetLowering(const TargetMachine &TM,
     setTruncStoreAction(MVT::v2i32, MVT::v2i16, Custom);
 
     for (auto VT : {MVT::nxv2f16, MVT::nxv4f16, MVT::nxv8f16, MVT::nxv2f32,
-                    MVT::nxv4f32, MVT::nxv2f64}) {
+                    MVT::nxv4f32, MVT::nxv2f64, MVT::nxv2bf16, MVT::nxv4bf16,
+                    MVT::nxv8f16}) {
       setOperationAction(ISD::BITCAST, VT, Custom);
       setOperationAction(ISD::CONCAT_VECTORS, VT, Custom);
       setOperationAction(ISD::INSERT_SUBVECTOR, VT, Custom);
diff --git a/llvm/lib/Target/AArch64/SVEInstrFormats.td b/llvm/lib/Target/AArch64/SVEInstrFormats.td
index 107bc79f70dbcb..8e47932c32e9f1 100644
--- a/llvm/lib/Target/AArch64/SVEInstrFormats.td
+++ b/llvm/lib/Target/AArch64/SVEInstrFormats.td
@@ -2117,23 +2117,29 @@ multiclass sve2p1_bf_2op_p_zds<bits<4> opc, string asm, string Ps,
                             SDPatternOperator op, DestructiveInstTypeEnum flags,
                             string revname="", bit isReverseInstr=0> {
 let DestructiveInstType = flags in {
-  def NAME : sve_fp_2op_p_zds<0b00, opc, asm, ZPR16>,
-           SVEPseudo2Instr<Ps, 1>, SVEInstr2Rev<NAME , revname , isReverseInstr>;
+  def _H: sve_fp_2op_p_zds<0b00, opc, asm, ZPR16>,
+          SVEPseudo2Instr<Ps # _H, 1>, SVEInstr2Rev<NAME # _H, revname # _H, isReverseInstr>;
   }
 
-  def : SVE_3_Op_Pat<nxv8bf16, op, nxv8i1, nxv8bf16, nxv8bf16, !cast<Instruction>(NAME)>;
+  def : SVE_3_Op_Pat<nxv8bf16, op, nxv8i1, nxv8bf16, nxv8bf16, !cast<Instruction>(NAME # _H)>;
+  def : SVE_3_Op_Pat<nxv4bf16, op, nxv4i1, nxv4bf16, nxv4bf16, !cast<Instruction>(NAME # _H)>;
+  def : SVE_3_Op_Pat<nxv2bf16, op, nxv2i1, nxv2bf16, nxv2bf16, !cast<Instruction>(NAME # _H)>;
 }
 
 multiclass sve2p1_bf_bin_pred_zds<SDPatternOperator op> {
-  def _UNDEF : PredTwoOpPseudo<NAME, ZPR16, FalseLanesUndef>;
+  def _H_UNDEF : PredTwoOpPseudo<NAME # _H, ZPR16, FalseLanesUndef>;
 
-  def : SVE_3_Op_Pat<nxv8bf16, op, nxv8i1,  nxv8bf16, nxv8bf16, !cast<Pseudo>(NAME # _UNDEF)>;
+  def : SVE_3_Op_Pat<nxv8bf16, op, nxv8i1,  nxv8bf16, nxv8bf16, !cast<Pseudo>(NAME # _H_UNDEF)>;
+  def : SVE_3_Op_Pat<nxv4bf16, op, nxv4i1,  nxv4bf16, nxv4bf16, !cast<Pseudo>(NAME # _H_UNDEF)>;
+  def : SVE_3_Op_Pat<nxv2bf16, op, nxv2i1,  nxv2bf16, nxv2bf16, !cast<Pseudo>(NAME # _H_UNDEF)>;
 }
 
 multiclass sve2p1_bf_2op_p_zds_zeroing<SDPatternOperator op> {
-  def _ZERO : PredTwoOpPseudo<NAME, ZPR16, FalseLanesZero>;
+  def _H_ZERO : PredTwoOpPseudo<NAME # _H, ZPR16, FalseLanesZero>;
 
-  def : SVE_3_Op_Pat_SelZero<nxv8bf16, op, nxv8i1, nxv8bf16, nxv8bf16, !cast<Pseudo>(NAME # _ZERO)>;
+  def : SVE_3_Op_Pat_SelZero<nxv8bf16, op, nxv8i1, nxv8bf16, nxv8bf16, !cast<Pseudo>(NAME # _H_ZERO)>;
+  def : SVE_3_Op_Pat_SelZero<nxv4bf16, op, nxv4i1, nxv4bf16, nxv4bf16, !cast<Pseudo>(NAME # _H_ZERO)>;
+  def : SVE_3_Op_Pat_SelZero<nxv2bf16, op, nxv2i1, nxv2bf16, nxv2bf16, !cast<Pseudo>(NAME # _H_ZERO)>;
 }
 
 multiclass sve_fp_2op_p_zds<bits<4> opc, string asm, string Ps,
diff --git a/llvm/test/CodeGen/AArch64/sve2p1-bf16-bin.ll b/llvm/test/CodeGen/AArch64/sve2p1-bf16-bin.ll
new file mode 100644
index 00000000000000..0d161c3159f5ef
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/sve2p1-bf16-bin.ll
@@ -0,0 +1,32 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve2p1 -mattr=+sve-b16b16 -mattr=+use-experimental-zeroing-pseudos -verify-machineinstrs < %s \
+; RUN: | FileCheck %s
+
+define <vscale x 8 x bfloat> @bfadd_8_bf(<vscale x 8 x bfloat> %a, <vscale x 8 x bfloat> %b){
+; CHECK-LABEL: bfadd_8_bf:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    bfadd z0.h, z0.h, z1.h
+; CHECK-NEXT:    ret
+  %res = fadd <vscale x 8 x bfloat> %a, %b
+  ret <vscale x 8 x bfloat> %res
+}
+
+define <vscale x 4 x bfloat> @bfadd_4_bf(<vscale x 4 x bfloat> %a, <vscale x 4 x bfloat> %b){
+; CHECK-LABEL: bfadd_4_bf:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    bfadd z0.h, p0/m, z0.h, z1.h
+; CHECK-NEXT:    ret
+  %res = fadd <vscale x 4 x bfloat> %a, %b
+  ret <vscale x 4 x bfloat> %res
+}
+
+define <vscale x 2 x bfloat> @bfadd_2_bf(<vscale x 2 x bfloat> %a, <vscale x 2 x bfloat> %b){
+; CHECK-LABEL: bfadd_2_bf:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    bfadd z0.h, p0/m, z0.h, z1.h
+; CHECK-NEXT:    ret
+  %res = fadd <vscale x 2 x bfloat> %a, %b
+  ret <vscale x 2 x bfloat> %res
+}