[AArch64] Enable using SVE2 bit-sel instructions with Neon types. (#146906)

This affects EOR3/BCAX/BSL/NBSL/BSL1N/BSL2N.

Author: rj-jesus

llvm/lib/Target/AArch64/SVEInstrFormats.td

@@ -5368,6 +5368,19 @@ multiclass sve2_int_bitwise_ternary_op<bits<3> opc, string asm,
   def : SVE_3_Op_Pat<nxv8i16, op, nxv8i16, nxv8i16, nxv8i16, !cast<Instruction>(NAME)>;
   def : SVE_3_Op_Pat<nxv4i32, op, nxv4i32, nxv4i32, nxv4i32, !cast<Instruction>(NAME)>;
   def : SVE_3_Op_Pat<nxv2i64, op, nxv2i64, nxv2i64, nxv2i64, !cast<Instruction>(NAME)>;
+
+  // Allow selecting SVE2 ternary ops with Neon types.
+  foreach VT = [nxv16i8, nxv8i16, nxv4i32, nxv2i64] in {
+    def : Pat<(SVEType<VT>.DSub (op V64:$op1, V64:$op2, V64:$op3)),
+              (EXTRACT_SUBREG (!cast<Instruction>(NAME) (INSERT_SUBREG (IMPLICIT_DEF), $op1, dsub),
+                                                        (INSERT_SUBREG (IMPLICIT_DEF), $op2, dsub),
+                                                        (INSERT_SUBREG (IMPLICIT_DEF), $op3, dsub)), dsub)>;
+
+    def : Pat<(SVEType<VT>.ZSub (op V128:$op1, V128:$op2, V128:$op3)),
+              (EXTRACT_SUBREG (!cast<Instruction>(NAME) (INSERT_SUBREG (IMPLICIT_DEF), $op1, zsub),
+                                                        (INSERT_SUBREG (IMPLICIT_DEF), $op2, zsub),
+                                                        (INSERT_SUBREG (IMPLICIT_DEF), $op3, zsub)), zsub)>;
+  }
 }
 
 class sve2_int_rotate_right_imm<bits<4> tsz8_64, string asm,

llvm/test/CodeGen/AArch64/bcax.ll

@@ -1,6 +1,8 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --extra_scrub
 ; RUN: llc -mtriple=aarch64 -mattr=+sha3 < %s | FileCheck --check-prefix=SHA3 %s
 ; RUN: llc -mtriple=aarch64 -mattr=-sha3 < %s | FileCheck --check-prefix=NOSHA3 %s
+; RUN: llc -mtriple=aarch64 -mattr=+sve2 < %s | FileCheck --check-prefix=SVE2 %s
+; RUN: llc -mtriple=aarch64 -mattr=+sha3,+sve2 < %s | FileCheck --check-prefix=SHA3 %s
 
 define <2 x i64> @bcax_64x2(<2 x i64> %0, <2 x i64> %1, <2 x i64> %2) {
 ; SHA3-LABEL: bcax_64x2:
@@ -13,6 +15,15 @@ define <2 x i64> @bcax_64x2(<2 x i64> %0, <2 x i64> %1, <2 x i64> %2) {
 ; NOSHA3-NEXT:    bic v0.16b, v0.16b, v1.16b
 ; NOSHA3-NEXT:    eor v0.16b, v0.16b, v2.16b
 ; NOSHA3-NEXT:    ret
+;
+; SVE2-LABEL: bcax_64x2:
+; SVE2:       // %bb.0:
+; SVE2-NEXT:    // kill: def $q2 killed $q2 def $z2
+; SVE2-NEXT:    // kill: def $q1 killed $q1 def $z1
+; SVE2-NEXT:    // kill: def $q0 killed $q0 def $z0
+; SVE2-NEXT:    bcax z2.d, z2.d, z0.d, z1.d
+; SVE2-NEXT:    mov v0.16b, v2.16b
+; SVE2-NEXT:    ret
   %4 = xor <2 x i64> %1, <i64 -1, i64 -1>
   %5 = and <2 x i64> %4, %0
   %6 = xor <2 x i64> %5, %2
@@ -30,6 +41,15 @@ define <4 x i32> @bcax_32x4(<4 x i32> %0, <4 x i32> %1, <4 x i32> %2) {
 ; NOSHA3-NEXT:    bic v0.16b, v0.16b, v1.16b
 ; NOSHA3-NEXT:    eor v0.16b, v0.16b, v2.16b
 ; NOSHA3-NEXT:    ret
+;
+; SVE2-LABEL: bcax_32x4:
+; SVE2:       // %bb.0:
+; SVE2-NEXT:    // kill: def $q2 killed $q2 def $z2
+; SVE2-NEXT:    // kill: def $q1 killed $q1 def $z1
+; SVE2-NEXT:    // kill: def $q0 killed $q0 def $z0
+; SVE2-NEXT:    bcax z2.d, z2.d, z0.d, z1.d
+; SVE2-NEXT:    mov v0.16b, v2.16b
+; SVE2-NEXT:    ret
   %4 = xor <4 x i32> %1, <i32 -1, i32 -1, i32 -1, i32 -1>
   %5 = and <4 x i32> %4, %0
   %6 = xor <4 x i32> %5, %2
@@ -47,6 +67,15 @@ define <8 x i16> @bcax_16x8(<8 x i16> %0, <8 x i16> %1, <8 x i16> %2) {
 ; NOSHA3-NEXT:    bic v0.16b, v0.16b, v1.16b
 ; NOSHA3-NEXT:    eor v0.16b, v0.16b, v2.16b
 ; NOSHA3-NEXT:    ret
+;
+; SVE2-LABEL: bcax_16x8:
+; SVE2:       // %bb.0:
+; SVE2-NEXT:    // kill: def $q2 killed $q2 def $z2
+; SVE2-NEXT:    // kill: def $q1 killed $q1 def $z1
+; SVE2-NEXT:    // kill: def $q0 killed $q0 def $z0
+; SVE2-NEXT:    bcax z2.d, z2.d, z0.d, z1.d
+; SVE2-NEXT:    mov v0.16b, v2.16b
+; SVE2-NEXT:    ret
   %4 = xor <8 x i16> %1, <i16 -1, i16 -1, i16 -1, i16 -1, i16 -1, i16 -1, i16 -1, i16 -1>
   %5 = and <8 x i16> %4, %0
   %6 = xor <8 x i16> %5, %2
@@ -64,6 +93,15 @@ define <16 x i8> @bcax_8x16(<16 x i8> %0, <16 x i8> %1, <16 x i8> %2) {
 ; NOSHA3-NEXT:    bic v0.16b, v0.16b, v1.16b
 ; NOSHA3-NEXT:    eor v0.16b, v0.16b, v2.16b
 ; NOSHA3-NEXT:    ret
+;
+; SVE2-LABEL: bcax_8x16:
+; SVE2:       // %bb.0:
+; SVE2-NEXT:    // kill: def $q2 killed $q2 def $z2
+; SVE2-NEXT:    // kill: def $q1 killed $q1 def $z1
+; SVE2-NEXT:    // kill: def $q0 killed $q0 def $z0
+; SVE2-NEXT:    bcax z2.d, z2.d, z0.d, z1.d
+; SVE2-NEXT:    mov v0.16b, v2.16b
+; SVE2-NEXT:    ret
   %4 = xor <16 x i8> %1, <i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1, i8 -1>
   %5 = and <16 x i8> %4, %0
   %6 = xor <16 x i8> %5, %2