[AArch64][SME] Allow predicate-as-counter operands for psel

[sdesmalen]

Summary

The specification says:

For programmer convenience, an assembler must also accept
predicate-as-counter register names for the destination predicate
register and the first source predicate register

[CarolineConcatto]

Thank you Sander,
It would be nice if you could point in the commit message to the spec that introduces this.
I believe it is the 2022 release. Right?

llvm/lib/Target/AArch64/AArch64RegisterInfo.td

class PNRAsmOperand<string name, string RegClass, int Width>
   let ParserMethod   = "tryParseSVEPredicateVector<RegKind::SVEPredicateAsCounter>";
 }
 
 class PNRRegOp<string Suffix, AsmOperandClass C, int EltSize, RegisterClass RC>
     : PPRRegOp<Suffix, C, ElementSizeNone, RC> {
   let PrintMethod = "printPredicateAsCounter<" # EltSize # ">";
 }
 
+def PNRAsmOpAny: PNRAsmOperand<"PNPredicateAny", "PPR", 0>;
 def PNRAsmOp8  : PNRAsmOperand<"PNPredicateB", "PPR", 8>;
 def PNRAsmOp16 : PNRAsmOperand<"PNPredicateH", "PPR", 16>;
 def PNRAsmOp32 : PNRAsmOperand<"PNPredicateS", "PPR", 32>;
 def PNRAsmOp64 : PNRAsmOperand<"PNPredicateD", "PPR", 64>;
 
+def PNRAny : PNRRegOp<"", PNRAsmOpAny, 0,  PPR>;
 def PNR8   : PNRRegOp<"b", PNRAsmOp8,  8,  PPR>;
 def PNR16  : PNRRegOp<"h", PNRAsmOp16, 16, PPR>;
 def PNR32  : PNRRegOp<"s", PNRAsmOp32, 32, PPR>;
 def PNR64  : PNRRegOp<"d", PNRAsmOp64, 64, PPR>;
 
 class PNRP8to15RegOp<string Suffix, AsmOperandClass C, int EltSize, RegisterClass RC>
     : PPRRegOp<Suffix, C, ElementSizeNone, RC> {
   let PrintMethod   = "printPredicateAsCounter<" # EltSize # ">";

llvm/lib/Target/AArch64/SMEInstrFormats.td

multiclass sve2_int_perm_sel_p<string asm, SDPatternOperator op> {
   }
   def _D : sve2_int_perm_sel_p<asm, PPR64, sme_elm_idx0_1> {
     bits<1> imm;
     let Inst{23}    = imm;
     let Inst{22}    = 0b1;
     let Inst{20-18} = 0b000;
   }
 
+  def : InstAlias<asm # "\t$Pd, $Pn, $Pm[$Rv, $imm]",
+                  (!cast<Instruction>(NAME # _B) PNRAny:$Pd,
+                      PNRAny:$Pn, PPR8:$Pm, MatrixIndexGPR32Op12_15:$Rv, sme_elm_idx0_15:$imm), 0>;
+  def : InstAlias<asm # "\t$Pd, $Pn, $Pm[$Rv, $imm]",
+                  (!cast<Instruction>(NAME # _H) PNRAny:$Pd,
+                      PNRAny:$Pn, PPR16:$Pm, MatrixIndexGPR32Op12_15:$Rv, sme_elm_idx0_7:$imm), 0>;
+  def : InstAlias<asm # "\t$Pd, $Pn, $Pm[$Rv, $imm]",
+                  (!cast<Instruction>(NAME # _S) PNRAny:$Pd,
+                      PNRAny:$Pn, PPR32:$Pm, MatrixIndexGPR32Op12_15:$Rv, sme_elm_idx0_3:$imm), 0>;
+  def : InstAlias<asm # "\t$Pd, $Pn, $Pm[$Rv, $imm]",
+                  (!cast<Instruction>(NAME # _D) PNRAny:$Pd,
+                      PNRAny:$Pn, PPR64:$Pm, MatrixIndexGPR32Op12_15:$Rv, sme_elm_idx0_1:$imm), 0>;
+
   def : Pat<(nxv16i1 (op (nxv16i1 PPRAny:$Pn), (nxv16i1 PPRAny:$Pm),
              MatrixIndexGPR32Op12_15:$idx)),
             (!cast<Instruction>(NAME # _B) $Pn, $Pm, $idx, 0)>;
   def : Pat<(nxv8i1 (op (nxv8i1 PPRAny:$Pn), (nxv8i1 PPRAny:$Pm),
              MatrixIndexGPR32Op12_15:$idx)),
             (!cast<Instruction>(NAME # _H) $Pn, $Pm, $idx, 0)>;
   def : Pat<(nxv4i1 (op (nxv4i1 PPRAny:$Pn), (nxv4i1 PPRAny:$Pm),
              MatrixIndexGPR32Op12_15:$idx)),

llvm/test/MC/AArch64/SME/psel.s

 // CHECK-ERROR: instruction requires: sme
 // CHECK-UNKNOWN: 25606da7 <unknown>
 
 psel    p15, p15, p15.d[w15, 1]
 // CHECK-INST: psel    p15, p15, p15.d[w15, 1]
 // CHECK-ENCODING: [0xef,0x7d,0xe3,0x25]
 // CHECK-ERROR: instruction requires: sme
 // CHECK-UNKNOWN: 25e37def <unknown>
+
+// --------------------------------------------------------------------------//
+// Also support with predicate-as-counter
+
+psel    pn15, pn15, p15.b[w15, 15]
+// CHECK-INST: psel    p15, p15, p15.b[w15, 15]
+// CHECK-ENCODING: [0xef,0x7d,0xff,0x25]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 25ff7def <unknown>
+
+psel    pn15, pn15, p15.h[w15, 7]
+// CHECK-INST: psel    p15, p15, p15.h[w15, 7]
+// CHECK-ENCODING: [0xef,0x7d,0xfb,0x25]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 25fb7def <unknown>
+
+psel    pn15, pn15, p15.s[w15, 3]
+// CHECK-INST: psel    p15, p15, p15.s[w15, 3]
+// CHECK-ENCODING: [0xef,0x7d,0xf3,0x25]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 25f37def <unknown>
+
+psel    pn15, pn15, p15.d[w15, 1]
+// CHECK-INST: psel    p15, p15, p15.d[w15, 1]
+// CHECK-ENCODING: [0xef,0x7d,0xe3,0x25]
+// CHECK-ERROR: instruction requires: sme
+// CHECK-UNKNOWN: 25e37def <unknown>