[CodeGen][SVE] Add patterns for whole vector predicate select

[sdesmalen]

Summary

Added patterns to implement select i1 %p, <vty> %a, <vty> %b

[efriedma]

We use whilelo to lower VECTOR_SPLAT; I think that ends up being one fewer vector instruction. What's the tradeoff between that vs. dup+cmpne?

[sdesmalen]

In D79356#2018893, @efriedma wrote:

We use whilelo to lower VECTOR_SPLAT; I think that ends up being one fewer vector instruction. What's the tradeoff between that vs. dup+cmpne?

There's no specific reason to choose dup+cmpne, other than that in these patterns the GPR32 can be any non-zero value, where to use whilelo we'd need a sign-extended i1 value.
I'm happy to change it to use whilelo, but I think it then needs to be implemented in ISelLowering code where we can do analysis on the value of the predicate using computeKnownBits, rather than using patterns. Is that correct?

[efriedma]

It's possible to write a TableGen pattern that explicitly generates an SBFM. In general that's what we would generate anyway. But you'd want to do custom lowering to give DAGCombine a chance to combine it away, yes.

[sdesmalen]

• select now lowers to splat_vector for the predicate (which codegens to whilelo) plus vselect.
• Updated tests and rebased patch.

[efriedma]

LGTM

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

def : Pat<(nxv16i1 (and PPR:$Ps1, PPR:$Ps2)),
             (AND_PPzPP (PTRUE_B 31), PPR:$Ps1, PPR:$Ps2)>;
   def : Pat<(nxv8i1 (and PPR:$Ps1, PPR:$Ps2)),
             (AND_PPzPP (PTRUE_H 31), PPR:$Ps1, PPR:$Ps2)>;
   def : Pat<(nxv4i1 (and PPR:$Ps1, PPR:$Ps2)),
             (AND_PPzPP (PTRUE_S 31), PPR:$Ps1, PPR:$Ps2)>;
   def : Pat<(nxv2i1 (and PPR:$Ps1, PPR:$Ps2)),
             (AND_PPzPP (PTRUE_D 31), PPR:$Ps1, PPR:$Ps2)>;
 
+  // Whole vector selects.
+  def : Pat<(nxv16i8 (select GPR32:$cond, ZPR:$Zs1, ZPR:$Zs2)),
+            (SEL_ZPZZ_S (CMPNE_PPzZI_S (PTRUE_S 31), (DUP_ZR_S $cond), 0), ZPR:$Zs1, ZPR:$Zs2)>;
+  def : Pat<(nxv8i16 (select GPR32:$cond, ZPR:$Zs1, ZPR:$Zs2)),
+            (SEL_ZPZZ_S (CMPNE_PPzZI_S (PTRUE_S 31), (DUP_ZR_S $cond), 0), ZPR:$Zs1, ZPR:$Zs2)>;
+  def : Pat<(nxv4i32 (select GPR32:$cond, ZPR:$Zs1, ZPR:$Zs2)),
+            (SEL_ZPZZ_S (CMPNE_PPzZI_S (PTRUE_S 31), (DUP_ZR_S $cond), 0), ZPR:$Zs1, ZPR:$Zs2)>;
+  def : Pat<(nxv2i64 (select GPR32:$cond, ZPR:$Zs1, ZPR:$Zs2)),
+            (SEL_ZPZZ_S (CMPNE_PPzZI_S (PTRUE_S 31), (DUP_ZR_S $cond), 0), ZPR:$Zs1, ZPR:$Zs2)>;
+  def : Pat<(nxv8f16 (select GPR32:$cond, ZPR:$Zs1, ZPR:$Zs2)),
+            (SEL_ZPZZ_S (CMPNE_PPzZI_S (PTRUE_S 31), (DUP_ZR_S $cond), 0), ZPR:$Zs1, ZPR:$Zs2)>;
+  def : Pat<(nxv2f32 (select GPR32:$cond, ZPR:$Zs1, ZPR:$Zs2)),
+            (SEL_ZPZZ_S (CMPNE_PPzZI_S (PTRUE_S 31), (DUP_ZR_S $cond), 0), ZPR:$Zs1, ZPR:$Zs2)>;
+  def : Pat<(nxv4f32 (select GPR32:$cond, ZPR:$Zs1, ZPR:$Zs2)),
+            (SEL_ZPZZ_S (CMPNE_PPzZI_S (PTRUE_S 31), (DUP_ZR_S $cond), 0), ZPR:$Zs1, ZPR:$Zs2)>;
+  def : Pat<(nxv2f64 (select GPR32:$cond, ZPR:$Zs1, ZPR:$Zs2)),
+            (SEL_ZPZZ_S (CMPNE_PPzZI_S (PTRUE_S 31), (DUP_ZR_S $cond), 0), ZPR:$Zs1, ZPR:$Zs2)>;
+  def : Pat<(nxv16i1 (select GPR32:$cond, PPR:$Ps1, PPR:$Ps2)),
+            (SEL_PPPP (CMPNE_PPzZI_B (PTRUE_B 31), (DUP_ZR_B $cond), 0), PPR:$Ps1, PPR:$Ps2)>;
+  def : Pat<(nxv8i1 (select GPR32:$cond, PPR:$Ps1, PPR:$Ps2)),
+            (SEL_PPPP (CMPNE_PPzZI_B (PTRUE_B 31), (DUP_ZR_B $cond), 0), PPR:$Ps1, PPR:$Ps2)>;
+  def : Pat<(nxv4i1 (select GPR32:$cond, PPR:$Ps1, PPR:$Ps2)),
+            (SEL_PPPP (CMPNE_PPzZI_B (PTRUE_B 31), (DUP_ZR_B $cond), 0), PPR:$Ps1, PPR:$Ps2)>;
+  def : Pat<(nxv2i1 (select GPR32:$cond, PPR:$Ps1, PPR:$Ps2)),
+            (SEL_PPPP (CMPNE_PPzZI_B (PTRUE_B 31), (DUP_ZR_B $cond), 0), PPR:$Ps1, PPR:$Ps2)>;
+
   // Add more complex addressing modes here as required
   multiclass pred_load<ValueType Ty, ValueType PredTy, SDPatternOperator Load,
                        Instruction RegRegInst, Instruction RegImmInst, ComplexPattern AddrCP> {
     // reg + reg
     let AddedComplexity = 1 in {
       def _reg_reg_z : Pat<(Ty (Load (AddrCP GPR64:$base, GPR64:$offset), (PredTy PPR:$gp), (SVEDup0Undef))),
                            (RegRegInst PPR:$gp, GPR64:$base, GPR64:$offset)>;
     }

llvm/test/CodeGen/AArch64/select-sve.ll

+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve -verify-machineinstrs < %s | FileCheck %s
+
+define <vscale x  16 x i8> @select_nxv16i8(i1 %cond, <vscale x  16 x i8> %a, <vscale x  16 x i8> %b) {
+; CHECK-LABEL: select_nxv16i8:
+; CHECK-DAG: ptrue [[PG:p[0-9]+]].s
+; CHECK-DAG: and [[COND:w[0-9]+]], w0, #0x1
+; CHECK-DAG: mov [[VCOND:z[0-9]+]].s, [[COND]]
+; CHECK-NEXT: cmpne [[PRED:p[0-9]+]].s, [[PG]]/z, [[VCOND]].s, #0
+; CHECK-NEXT: sel z0.s, [[PRED]], z0.s, z1.s
+; CHECK-NEXT: ret
+  %res = select i1 %cond, <vscale x  16 x i8> %a, <vscale x  16 x i8> %b
+  ret <vscale x  16 x i8> %res
+}
+
+define <vscale x  8 x i16> @select_nxv8i16(i1 %cond, <vscale x  8 x i16> %a, <vscale x  8 x i16> %b) {
+; CHECK-LABEL: select_nxv8i16:
+; CHECK-DAG: ptrue [[PG:p[0-9]+]].s
+; CHECK-DAG: and [[COND:w[0-9]+]], w0, #0x1
+; CHECK-DAG: mov [[VCOND:z[0-9]+]].s, [[COND]]
+; CHECK-NEXT: cmpne [[PRED:p[0-9]+]].s, [[PG]]/z, [[VCOND]].s, #0
+; CHECK-NEXT: sel z0.s, [[PRED]], z0.s, z1.s
+; CHECK-NEXT: ret
+  %res = select i1 %cond, <vscale x  8 x i16> %a, <vscale x  8 x i16> %b
+  ret <vscale x  8 x i16> %res
+}
+
+define <vscale x  4 x i32> @select_nxv4i32(i1 %cond, <vscale x  4 x i32> %a, <vscale x  4 x i32> %b) {
+; CHECK-LABEL: select_nxv4i32:
+; CHECK-DAG: ptrue [[PG:p[0-9]+]].s
+; CHECK-DAG: and [[COND:w[0-9]+]], w0, #0x1
+; CHECK-DAG: mov [[VCOND:z[0-9]+]].s, [[COND]]
+; CHECK-NEXT: cmpne [[PRED:p[0-9]+]].s, [[PG]]/z, [[VCOND]].s, #0
+; CHECK-NEXT: sel z0.s, [[PRED]], z0.s, z1.s
+; CHECK-NEXT: ret
+  %res = select i1 %cond, <vscale x  4 x i32> %a, <vscale x  4 x i32> %b
+  ret <vscale x  4 x i32> %res
+}
+
+define <vscale x  2 x i64> @select_nxv2i64(i1 %cond, <vscale x  2 x i64> %a, <vscale x  2 x i64> %b) {
+; CHECK-LABEL: select_nxv2i64:
+; CHECK-DAG: ptrue [[PG:p[0-9]+]].s
+; CHECK-DAG: and [[COND:w[0-9]+]], w0, #0x1
+; CHECK-DAG: mov [[VCOND:z[0-9]+]].s, [[COND]]
+; CHECK-NEXT: cmpne [[PRED:p[0-9]+]].s, [[PG]]/z, [[VCOND]].s, #0
+; CHECK-NEXT: sel z0.s, [[PRED]], z0.s, z1.s
+; CHECK-NEXT: ret
+  %res = select i1 %cond, <vscale x  2 x i64> %a, <vscale x  2 x i64> %b
+  ret <vscale x  2 x i64> %res
+}
+
+define <vscale x  8 x half> @select_nxv8f16(i1 %cond, <vscale x  8 x half> %a, <vscale x  8 x half> %b) {
+; CHECK-LABEL: select_nxv8f16:
+; CHECK-DAG: ptrue [[PG:p[0-9]+]].s
+; CHECK-DAG: and [[COND:w[0-9]+]], w0, #0x1
+; CHECK-DAG: mov [[VCOND:z[0-9]+]].s, [[COND]]
+; CHECK-NEXT: cmpne [[PRED:p[0-9]+]].s, [[PG]]/z, [[VCOND]].s, #0
+; CHECK-NEXT: sel z0.s, [[PRED]], z0.s, z1.s
+; CHECK-NEXT: ret
+  %res = select i1 %cond, <vscale x  8 x half> %a, <vscale x  8 x half> %b
+  ret <vscale x  8 x half> %res
+}
+
+define <vscale x  4 x float> @select_nxv4f32(i1 %cond, <vscale x  4 x float> %a, <vscale x  4 x float> %b) {
+; CHECK-LABEL: select_nxv4f32:
+; CHECK-DAG: ptrue [[PG:p[0-9]+]].s
+; CHECK-DAG: and [[COND:w[0-9]+]], w0, #0x1
+; CHECK-DAG: mov [[VCOND:z[0-9]+]].s, [[COND]]
+; CHECK-NEXT: cmpne [[PRED:p[0-9]+]].s, [[PG]]/z, [[VCOND]].s, #0
+; CHECK-NEXT: sel z0.s, [[PRED]], z0.s, z1.s
+; CHECK-NEXT: ret
+  %res = select i1 %cond, <vscale x  4 x float> %a, <vscale x  4 x float> %b
+  ret <vscale x  4 x float> %res
+}
+
+define <vscale x  2 x double> @select_nxv2f64(i1 %cond, <vscale x  2 x double> %a, <vscale x  2 x double> %b) {
+; CHECK-LABEL: select_nxv2f64:
+; CHECK-DAG: ptrue [[PG:p[0-9]+]].s
+; CHECK-DAG: and [[COND:w[0-9]+]], w0, #0x1
+; CHECK-DAG: mov [[VCOND:z[0-9]+]].s, [[COND]]
+; CHECK-NEXT: cmpne [[PRED:p[0-9]+]].s, [[PG]]/z, [[VCOND]].s, #0
+; CHECK-NEXT: sel z0.s, [[PRED]], z0.s, z1.s
+; CHECK-NEXT: ret
+  %res = select i1 %cond, <vscale x  2 x double> %a, <vscale x  2 x double> %b
+  ret <vscale x  2 x double> %res
+}
+
+define <vscale x  16 x i1> @select_nxv16i1(i1 %cond, <vscale x  16 x i1> %a, <vscale x  16 x i1> %b) {
+; CHECK-LABEL: select_nxv16i1:
+; CHECK-DAG: ptrue [[PG:p[0-9]+]].b
+; CHECK-DAG: and [[COND:w[0-9]+]], w0, #0x1
+; CHECK-DAG: mov [[VCOND:z[0-9]+]].b, [[COND]]
+; CHECK-NEXT: cmpne [[PRED:p[0-9]+]].b, [[PG]]/z, [[VCOND]].b, #0
+; CHECK-NEXT: sel p0.b, [[PRED]], p0.b, p1.b
+; CHECK-NEXT: ret
+  %res = select i1 %cond, <vscale x  16 x i1> %a, <vscale x  16 x i1> %b
+  ret <vscale x  16 x i1> %res
+}
+
+define <vscale x  8 x i1> @select_nxv8i1(i1 %cond, <vscale x  8 x i1> %a, <vscale x  8 x i1> %b) {
+; CHECK-LABEL: select_nxv8i1:
+; CHECK-DAG: ptrue [[PG:p[0-9]+]].b
+; CHECK-DAG: and [[COND:w[0-9]+]], w0, #0x1
+; CHECK-DAG: mov [[VCOND:z[0-9]+]].b, [[COND]]
+; CHECK-NEXT: cmpne [[PRED:p[0-9]+]].b, [[PG]]/z, [[VCOND]].b, #0
+; CHECK-NEXT: sel p0.b, [[PRED]], p0.b, p1.b
+; CHECK-NEXT: ret
+  %res = select i1 %cond, <vscale x  8 x i1> %a, <vscale x  8 x i1> %b
+  ret <vscale x  8 x i1> %res
+}
+
+define <vscale x  4 x i1> @select_nxv4i1(i1 %cond, <vscale x  4 x i1> %a, <vscale x  4 x i1> %b) {
+; CHECK-LABEL: select_nxv4i1:
+; CHECK-DAG: ptrue [[PG:p[0-9]+]].b
+; CHECK-DAG: and [[COND:w[0-9]+]], w0, #0x1
+; CHECK-DAG: mov [[VCOND:z[0-9]+]].b, [[COND]]
+; CHECK-NEXT: cmpne [[PRED:p[0-9]+]].b, [[PG]]/z, [[VCOND]].b, #0
+; CHECK-NEXT: sel p0.b, [[PRED]], p0.b, p1.b
+; CHECK-NEXT: ret
+  %res = select i1 %cond, <vscale x  4 x i1> %a, <vscale x  4 x i1> %b
+  ret <vscale x  4 x i1> %res
+}
+
+define <vscale x  2 x i1> @select_nxv2i1(i1 %cond, <vscale x  2 x i1> %a, <vscale x  2 x i1> %b) {
+; CHECK-LABEL: select_nxv2i1:
+; CHECK-DAG: ptrue [[PG:p[0-9]+]].b
+; CHECK-DAG: and [[COND:w[0-9]+]], w0, #0x1
+; CHECK-DAG: mov [[VCOND:z[0-9]+]].b, [[COND]]
+; CHECK-NEXT: cmpne [[PRED:p[0-9]+]].b, [[PG]]/z, [[VCOND]].b, #0
+; CHECK-NEXT: sel p0.b, [[PRED]], p0.b, p1.b
+; CHECK-NEXT: ret
+  %res = select i1 %cond, <vscale x  2 x i1> %a, <vscale x  2 x i1> %b
+  ret <vscale x  2 x i1> %res
+}