[AArch64][SVE] Add patterns for VSELECT of immediates.

[efriedma]

Summary

This covers forms involving "CPY (immediate, zeroing)".

This doesn't handle the case where the operands are reversed, and the condition is freely invertible. Not sure how to handle that. Maybe a DAGCombine.

[fpetrogalli]

Hi @efriedma ,

thanks you for working on this.

This doesn't handle the case where the operands are reversed, and the condition is freely invertible.

Do you mean cases like select <vscale x 2 x i1> %p, <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i64> %three? If that the case, the "merging" version of CPY , where %three is allocated to the dest register, should do it? CPY <Zd>.<T>, <Pg>/M, <R><n|SP>.

Kind regards,

Francesco

llvm/test/CodeGen/AArch64/sve-vselect-imm.ll

9

The immediate is a signed immediate - could you please test negative values too?

Nit: the name %three doesn't seem to be related to anything other than the value 3 of some of the tests... could you change it to something more generic, like %vec, or even leave it as %0?

74–111

Can you do the same "invalid values" tests for when the scaling shift is 0? For example -130 and 130.

Nit: adding the token _invalid_imm_values in the names of the functions would help a lot identifying what you are testing here.

[efriedma]

Do you mean cases like select <vscale x 2 x i1> %p, <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i64> %three? If that the case, the "merging" version of CPY , where %three is allocated to the dest register, should do it? CPY <Zd>.<T>, <Pg>/M, <R><n|SP>.

I was more specifically thinking of the case where %p is something like an icmp, where we can potentially rewrite the operation to produce an inverted result for free. Otherwise, yes, we can use merging CPY.

[efriedma]

Update test

[fpetrogalli]

the case where %p is something like an icmp, where we can potentially rewrite the operation to produce an inverted result for free.

I am not sure I understand. Can you give an explicit example?

Regards,

Francesco

llvm/test/CodeGen/AArch64/sve-vselect-imm.ll

115

I am probably missing something trivial, but why would this be the case only if the odd bits if the predicate are zero?

Moreover, why not add the patterns needed for this directly in this patch? IIUC a higher AddedComplexity is what you need to prioritize the current pattern over the sel one.

[efriedma]

Say you have something like x < 0 ? 0 : 1. That doesn't match. But if you rewrite it to x >= 0 ? 1 : 0, it does.

llvm/test/CodeGen/AArch64/sve-vselect-imm.ll

115

Using ACLE intrinsics, the suggestion is essentially doing something like svreinterpret_s16(svsel(pred, svdup_s8(-128), svdup_s8(0))). Note that this is using the predicate as a <vscale x 16 x i1>, not a <vscale x 8 x i1>.

The pattern required for that is probably non-trivial

[sdesmalen]

LGTM

This doesn't handle the case where the operands are reversed, and the condition is freely invertible. Not sure how to handle that. Maybe a DAGCombine.

A combine is probably easiest as it avoids having to specify the sequence to negate the predicate.

llvm/lib/Target/AArch64/SVEInstrFormats.td

179–187

Can I assume this was already handled by the PredicateMethod of SVECpyImmOperand8?

[efriedma]

llvm/lib/Target/AArch64/SVEInstrFormats.td

179–187

ImmLeaf is only used for SelectionDAG lowering. We can't do SelectionDAG lowering for shifted immediates using ImmLeaf because it needs to produce two MachineOperands. Therefore, an ImmLeaf predicate is completely useless here.

For asm parsing, the necessary predicate is handled by the PredicateMethod, yes.

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

 //=- AArch64SVEInstrInfo.td -  AArch64 SVE Instructions -*- tablegen -*-----=//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // AArch64 Scalable Vector Extension (SVE) Instruction definitions.
 //
 //===----------------------------------------------------------------------===//
 
-def SVE8BitLslImm : ComplexPattern<i32, 2, "SelectSVE8BitLslImm", [imm]>;
 def SVELShiftImm64 : ComplexPattern<i32, 1, "SelectSVEShiftImm64<0, 64>", []>;
 
 // Contiguous loads - node definitions
 //
 def SDT_AArch64_LD1 : SDTypeProfile<1, 3, [
   SDTCisVec<0>, SDTCisVec<1>, SDTCisPtrTy<2>,
   SDTCVecEltisVT<1,i1>, SDTCisSameNumEltsAs<0,1>
 ]>;

llvm/lib/Target/AArch64/SVEInstrFormats.td

 def SVECpyImmOperand64 : SVEShiftedImmOperand<64, "Cpy", "isSVECpyImm<int64_t>">;
 
 def SVEAddSubImmOperand8  : SVEShiftedImmOperand<8,  "AddSub", "isSVEAddSubImm<int8_t>">;
 def SVEAddSubImmOperand16 : SVEShiftedImmOperand<16, "AddSub", "isSVEAddSubImm<int16_t>">;
 def SVEAddSubImmOperand32 : SVEShiftedImmOperand<32, "AddSub", "isSVEAddSubImm<int32_t>">;
 def SVEAddSubImmOperand64 : SVEShiftedImmOperand<64, "AddSub", "isSVEAddSubImm<int64_t>">;
 
 class imm8_opt_lsl<int ElementWidth, string printType,
-                   AsmOperandClass OpndClass, code Predicate>
-    : Operand<i32>, ImmLeaf<i32, Predicate> {
+                   AsmOperandClass OpndClass>
+    : Operand<i32> {
   let EncoderMethod = "getImm8OptLsl";
   let DecoderMethod = "DecodeImm8OptLsl<" # ElementWidth # ">";
   let PrintMethod = "printImm8OptLsl<" # printType # ">";
   let ParserMatchClass = OpndClass;
   let MIOperandInfo = (ops i32imm, i32imm);
 }
 
-def cpy_imm8_opt_lsl_i8  : imm8_opt_lsl<8,  "int8_t",  SVECpyImmOperand8,  [{
-  return AArch64_AM::isSVECpyImm<int8_t>(Imm);
-}]>;
-def cpy_imm8_opt_lsl_i16 : imm8_opt_lsl<16, "int16_t", SVECpyImmOperand16, [{
-  return AArch64_AM::isSVECpyImm<int16_t>(Imm);
-}]>;
-def cpy_imm8_opt_lsl_i32 : imm8_opt_lsl<32, "int32_t", SVECpyImmOperand32, [{
-  return AArch64_AM::isSVECpyImm<int32_t>(Imm);
-}]>;
-def cpy_imm8_opt_lsl_i64 : imm8_opt_lsl<64, "int64_t", SVECpyImmOperand64, [{
-  return AArch64_AM::isSVECpyImm<int64_t>(Imm);
-}]>;
-
-def addsub_imm8_opt_lsl_i8  : imm8_opt_lsl<8,  "uint8_t",  SVEAddSubImmOperand8,  [{
-  return AArch64_AM::isSVEAddSubImm<int8_t>(Imm);
-}]>;
-def addsub_imm8_opt_lsl_i16 : imm8_opt_lsl<16, "uint16_t", SVEAddSubImmOperand16, [{
-  return AArch64_AM::isSVEAddSubImm<int16_t>(Imm);
-}]>;
-def addsub_imm8_opt_lsl_i32 : imm8_opt_lsl<32, "uint32_t", SVEAddSubImmOperand32, [{
-  return AArch64_AM::isSVEAddSubImm<int32_t>(Imm);
-}]>;
-def addsub_imm8_opt_lsl_i64 : imm8_opt_lsl<64, "uint64_t", SVEAddSubImmOperand64, [{
-  return AArch64_AM::isSVEAddSubImm<int64_t>(Imm);
-}]>;
+def cpy_imm8_opt_lsl_i8  : imm8_opt_lsl<8,  "int8_t",  SVECpyImmOperand8>;
+def cpy_imm8_opt_lsl_i16 : imm8_opt_lsl<16, "int16_t", SVECpyImmOperand16>;
+def cpy_imm8_opt_lsl_i32 : imm8_opt_lsl<32, "int32_t", SVECpyImmOperand32>;
+def cpy_imm8_opt_lsl_i64 : imm8_opt_lsl<64, "int64_t", SVECpyImmOperand64>;
+
+def addsub_imm8_opt_lsl_i8  : imm8_opt_lsl<8,  "uint8_t",  SVEAddSubImmOperand8>;
+def addsub_imm8_opt_lsl_i16 : imm8_opt_lsl<16, "uint16_t", SVEAddSubImmOperand16>;
+def addsub_imm8_opt_lsl_i32 : imm8_opt_lsl<32, "uint32_t", SVEAddSubImmOperand32>;
+def addsub_imm8_opt_lsl_i64 : imm8_opt_lsl<64, "uint64_t", SVEAddSubImmOperand64>;
 
 def SVEAddSubImm8Pat  : ComplexPattern<i32, 2, "SelectSVEAddSubImm<MVT::i8>", []>;
 def SVEAddSubImm16Pat : ComplexPattern<i32, 2, "SelectSVEAddSubImm<MVT::i16>", []>;
 def SVEAddSubImm32Pat : ComplexPattern<i32, 2, "SelectSVEAddSubImm<MVT::i32>", []>;
 def SVEAddSubImm64Pat : ComplexPattern<i32, 2, "SelectSVEAddSubImm<MVT::i64>", []>;
 
 def SVELogicalImm8Pat  : ComplexPattern<i64, 1, "SelectSVELogicalImm<MVT::i8>", []>;
 def SVELogicalImm16Pat : ComplexPattern<i64, 1, "SelectSVELogicalImm<MVT::i16>", []>;
 def SVELogicalImm32Pat : ComplexPattern<i64, 1, "SelectSVELogicalImm<MVT::i32>", []>;
 def SVELogicalImm64Pat : ComplexPattern<i64, 1, "SelectSVELogicalImm<MVT::i64>", []>;
 
+def SVE8BitLslImm : ComplexPattern<i32, 2, "SelectSVE8BitLslImm", [imm]>;
+
 def SVEArithUImmPat  : ComplexPattern<i32, 1, "SelectSVEArithImm", []>;
 def SVEArithSImmPat  : ComplexPattern<i32, 1, "SelectSVESignedArithImm", []>;
 
 class SVEExactFPImm<string Suffix, string ValA, string ValB> : AsmOperandClass {
   let Name = "SVEExactFPImmOperand" # Suffix;
   let DiagnosticType = "Invalid" # Name;
   let ParserMethod = "tryParseFPImm<false>";
   let PredicateMethod = "isExactFPImm<" # ValA # ", " # ValB # ">";
multiclass sve_int_dup_imm_pred_merge<string asm> {
   def : InstAlias<"fmov $Zd, $Pg/m, #0.0",
                   (!cast<Instruction>(NAME # _S) ZPR32:$Zd, PPRAny:$Pg, 0, 0), 0>;
   def : InstAlias<"fmov $Zd, $Pg/m, #0.0",
                   (!cast<Instruction>(NAME # _D) ZPR64:$Zd, PPRAny:$Pg, 0, 0), 0>;
 }
 
 multiclass sve_int_dup_imm_pred_zero_inst<
     bits<2> sz8_64, string asm, ZPRRegOp zprty, ValueType intty,
-    ValueType predty, imm8_opt_lsl cpyimm> {
-  def NAME : sve_int_dup_imm_pred<sz8_64, 0, asm, zprty, "/z", (ins PPRAny:$Pg, cpyimm:$imm)>;
+    ValueType predty, ValueType scalarty, imm8_opt_lsl cpyimm> {
+  def NAME : sve_int_dup_imm_pred<sz8_64, 0, asm, zprty, "/z",
+                                  (ins PPRAny:$Pg, cpyimm:$imm)>;
   def : InstAlias<"mov $Zd, $Pg/z, $imm",
                   (!cast<Instruction>(NAME) zprty:$Zd, PPRAny:$Pg, cpyimm:$imm), 1>;
   def : Pat<(intty (zext (predty PPRAny:$Ps1))),
             (!cast<Instruction>(NAME) PPRAny:$Ps1, 1, 0)>;
   def : Pat<(intty (sext (predty PPRAny:$Ps1))),
             (!cast<Instruction>(NAME) PPRAny:$Ps1, -1, 0)>;
   def : Pat<(intty (anyext (predty PPRAny:$Ps1))),
             (!cast<Instruction>(NAME) PPRAny:$Ps1, 1, 0)>;
+  def : Pat<(intty
+              (vselect predty:$Pg,
+                (intty (AArch64dup (scalarty (SVE8BitLslImm i32:$imm, i32:$shift)))),
+                (intty (AArch64dup (scalarty 0))))),
+            (!cast<Instruction>(NAME) $Pg, i32:$imm, i32:$shift)>;
 }
 
 multiclass sve_int_dup_imm_pred_zero<string asm> {
-  defm _B : sve_int_dup_imm_pred_zero_inst<0b00, asm, ZPR8,  nxv16i8, nxv16i1, cpy_imm8_opt_lsl_i8>;
-  defm _H : sve_int_dup_imm_pred_zero_inst<0b01, asm, ZPR16, nxv8i16, nxv8i1,  cpy_imm8_opt_lsl_i16>;
-  defm _S : sve_int_dup_imm_pred_zero_inst<0b10, asm, ZPR32, nxv4i32, nxv4i1,  cpy_imm8_opt_lsl_i32>;
-  defm _D : sve_int_dup_imm_pred_zero_inst<0b11, asm, ZPR64, nxv2i64, nxv2i1,  cpy_imm8_opt_lsl_i64>;
+  defm _B : sve_int_dup_imm_pred_zero_inst<0b00, asm, ZPR8,  nxv16i8, nxv16i1,
+                                           i32, cpy_imm8_opt_lsl_i8>;
+  defm _H : sve_int_dup_imm_pred_zero_inst<0b01, asm, ZPR16, nxv8i16, nxv8i1,
+                                           i32, cpy_imm8_opt_lsl_i16>;
+  defm _S : sve_int_dup_imm_pred_zero_inst<0b10, asm, ZPR32, nxv4i32, nxv4i1,
+                                           i32, cpy_imm8_opt_lsl_i32>;
+  defm _D : sve_int_dup_imm_pred_zero_inst<0b11, asm, ZPR64, nxv2i64, nxv2i1,
+                                           i64, cpy_imm8_opt_lsl_i64>;
 }
 
 //===----------------------------------------------------------------------===//
 // SVE Integer Compare - Vectors Group
 //===----------------------------------------------------------------------===//
 
 class sve_int_cmp<bit cmp_1, bits<2> sz8_64, bits<3> opc, string asm,
                   PPRRegOp pprty, ZPRRegOp zprty1, ZPRRegOp zprty2>

llvm/test/CodeGen/AArch64/sve-vselect-imm.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64--linux-gnu -mattr=+sve < %s | FileCheck %s
+
+define <vscale x 16 x i8> @sel_8(<vscale x 16 x i1> %p) {
+; CHECK-LABEL: sel_8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.b, p0/z, #3 // =0x3
+; CHECK-NEXT:    ret
+%three = shufflevector <vscale x 16 x i8> insertelement (<vscale x 16 x i8> undef, i8 3, i32 0), <vscale x 16 x i8> zeroinitializer, <vscale x 16 x i32> zeroinitializer
+%sel = select <vscale x 16 x i1> %p, <vscale x 16 x i8> %three, <vscale x 16 x i8> zeroinitializer
+ret <vscale x 16 x i8> %sel
+}
+
+define <vscale x 8 x i16> @sel_16(<vscale x 8 x i1> %p) {
+; CHECK-LABEL: sel_16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.h, p0/z, #3 // =0x3
+; CHECK-NEXT:    ret
+%three = shufflevector <vscale x 8 x i16> insertelement (<vscale x 8 x i16> undef, i16 3, i32 0), <vscale x 8 x i16> zeroinitializer, <vscale x 8 x i32> zeroinitializer
+%sel = select <vscale x 8 x i1> %p, <vscale x 8 x i16> %three, <vscale x 8 x i16> zeroinitializer
+ret <vscale x 8 x i16> %sel
+}
+
+define <vscale x 4 x i32> @sel_32(<vscale x 4 x i1> %p) {
+; CHECK-LABEL: sel_32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.s, p0/z, #3 // =0x3
+; CHECK-NEXT:    ret
+%three = shufflevector <vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 3, i32 0), <vscale x 4 x i32> zeroinitializer, <vscale x 4 x i32> zeroinitializer
+%sel = select <vscale x 4 x i1> %p, <vscale x 4 x i32> %three, <vscale x 4 x i32> zeroinitializer
+ret <vscale x 4 x i32> %sel
+}
+
+define <vscale x 2 x i64> @sel_64(<vscale x 2 x i1> %p) {
+; CHECK-LABEL: sel_64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.d, p0/z, #3 // =0x3
+; CHECK-NEXT:    ret
+%three = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 3, i32 0), <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i32> zeroinitializer
+%sel = select <vscale x 2 x i1> %p, <vscale x 2 x i64> %three, <vscale x 2 x i64> zeroinitializer
+ret <vscale x 2 x i64> %sel
+}
+
+define <vscale x 8 x i16> @sel2_16(<vscale x 8 x i1> %p) {
+; CHECK-LABEL: sel2_16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.h, p0/z, #512 // =0x200
+; CHECK-NEXT:    ret
+%three = shufflevector <vscale x 8 x i16> insertelement (<vscale x 8 x i16> undef, i16 512, i32 0), <vscale x 8 x i16> zeroinitializer, <vscale x 8 x i32> zeroinitializer
+%sel = select <vscale x 8 x i1> %p, <vscale x 8 x i16> %three, <vscale x 8 x i16> zeroinitializer
+ret <vscale x 8 x i16> %sel
+}
+
+define <vscale x 4 x i32> @sel2_32(<vscale x 4 x i1> %p) {
+; CHECK-LABEL: sel2_32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.s, p0/z, #512 // =0x200
+; CHECK-NEXT:    ret
+%three = shufflevector <vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 512, i32 0), <vscale x 4 x i32> zeroinitializer, <vscale x 4 x i32> zeroinitializer
+%sel = select <vscale x 4 x i1> %p, <vscale x 4 x i32> %three, <vscale x 4 x i32> zeroinitializer
+ret <vscale x 4 x i32> %sel
+}
+
+define <vscale x 2 x i64> @sel2_64(<vscale x 2 x i1> %p) {
+; CHECK-LABEL: sel2_64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.d, p0/z, #512 // =0x200
+; CHECK-NEXT:    ret
+%three = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 512, i32 0), <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i32> zeroinitializer
+%sel = select <vscale x 2 x i1> %p, <vscale x 2 x i64> %three, <vscale x 2 x i64> zeroinitializer
+ret <vscale x 2 x i64> %sel
+}
+
+define <vscale x 8 x i16> @sel3_16(<vscale x 8 x i1> %p) {
+; CHECK-LABEL: sel3_16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #513
+; CHECK-NEXT:    mov z0.h, w8
+; CHECK-NEXT:    mov z1.h, #0 // =0x0
+; CHECK-NEXT:    sel z0.h, p0, z0.h, z1.h
+; CHECK-NEXT:    ret
+%three = shufflevector <vscale x 8 x i16> insertelement (<vscale x 8 x i16> undef, i16 513, i32 0), <vscale x 8 x i16> zeroinitializer, <vscale x 8 x i32> zeroinitializer
+%sel = select <vscale x 8 x i1> %p, <vscale x 8 x i16> %three, <vscale x 8 x i16> zeroinitializer
+ret <vscale x 8 x i16> %sel
+}
+
+define <vscale x 4 x i32> @sel3_32(<vscale x 4 x i1> %p) {
+; CHECK-LABEL: sel3_32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #513
+; CHECK-NEXT:    mov z0.s, w8
+; CHECK-NEXT:    mov z1.s, #0 // =0x0
+; CHECK-NEXT:    sel z0.s, p0, z0.s, z1.s
+; CHECK-NEXT:    ret
+%three = shufflevector <vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 513, i32 0), <vscale x 4 x i32> zeroinitializer, <vscale x 4 x i32> zeroinitializer
+%sel = select <vscale x 4 x i1> %p, <vscale x 4 x i32> %three, <vscale x 4 x i32> zeroinitializer
+ret <vscale x 4 x i32> %sel
+}
+
+define <vscale x 2 x i64> @sel3_64(<vscale x 2 x i1> %p) {
+; CHECK-LABEL: sel3_64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w8, #513
+; CHECK-NEXT:    mov z0.d, x8
+; CHECK-NEXT:    mov z1.d, #0 // =0x0
+; CHECK-NEXT:    sel z0.d, p0, z0.d, z1.d
+; CHECK-NEXT:    ret
+%three = shufflevector <vscale x 2 x i64> insertelement (<vscale x 2 x i64> undef, i64 513, i32 0), <vscale x 2 x i64> zeroinitializer, <vscale x 2 x i32> zeroinitializer
+%sel = select <vscale x 2 x i1> %p, <vscale x 2 x i64> %three, <vscale x 2 x i64> zeroinitializer
+ret <vscale x 2 x i64> %sel
+}