[AArch64][SVE] Implement pfirst and pnext intrinsics

[cameron.mcinally]

LGTM

llvm/include/llvm/IR/IntrinsicsAArch64.td

let TargetPrefix = "aarch64" in {  // All intrinsics start with "llvm.aarch64.".
   class AdvSIMD_3VectorArgIndexed_Intrinsic
     : Intrinsic<[llvm_anyvector_ty],
                 [LLVMMatchType<0>,
                  LLVMMatchType<0>,
                  LLVMMatchType<0>,
                  llvm_i32_ty],
                 [IntrNoMem]>;
 
+  class AdvSIMD_Pred1VectorArg_Intrinsic
+    : Intrinsic<[llvm_anyvector_ty],
+                [LLVMScalarOrSameVectorWidth<0, llvm_i1_ty>,
+                 LLVMMatchType<0>],
+                [IntrNoMem]>;
+
   class AdvSIMD_Pred2VectorArg_Intrinsic
     : Intrinsic<[llvm_anyvector_ty],
                 [LLVMScalarOrSameVectorWidth<0, llvm_i1_ty>,
                  LLVMMatchType<0>,
                  LLVMMatchType<0>],
                 [IntrNoMem]>;
 
   class AdvSIMD_Pred3VectorArg_Intrinsic

 //
 
 def int_aarch64_sve_ptrue : AdvSIMD_SVE_PTRUE_Intrinsic;
 
 //
 // Predicate operations
 //
 
+def int_aarch64_sve_pfirst  : AdvSIMD_Pred1VectorArg_Intrinsic;
+def int_aarch64_sve_pnext   : AdvSIMD_Pred1VectorArg_Intrinsic;
 def int_aarch64_sve_punpkhi : AdvSIMD_SVE_PUNPKHI_Intrinsic;
 def int_aarch64_sve_punpklo : AdvSIMD_SVE_PUNPKHI_Intrinsic;
 
 //
 // Gather loads:
 //
 
 // scalar + vector, 64 bit unscaled offsets

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

let Predicates = [HasSVE] in {
   defm BRKA_PPmP  : sve_int_break_m<0b001, "brka">;
   defm BRKAS_PPzP : sve_int_break_z<0b010, "brkas">;
   defm BRKB_PPzP  : sve_int_break_z<0b100, "brkb">;
   defm BRKB_PPmP  : sve_int_break_m<0b101, "brkb">;
   defm BRKBS_PPzP : sve_int_break_z<0b110, "brkbs">;
 
   def PTEST_PP : sve_int_ptest<0b010000, "ptest">;
   def PFALSE   : sve_int_pfalse<0b000000, "pfalse">;
-  defm PFIRST  : sve_int_pfirst<0b00000, "pfirst">;
-  defm PNEXT   : sve_int_pnext<0b00110, "pnext">;
+  defm PFIRST  : sve_int_pfirst<0b00000, "pfirst", int_aarch64_sve_pfirst>;
+  defm PNEXT   : sve_int_pnext<0b00110, "pnext", int_aarch64_sve_pnext>;
 
   defm AND_PPzPP   : sve_int_pred_log<0b0000, "and", int_aarch64_sve_and>;
   defm BIC_PPzPP   : sve_int_pred_log<0b0001, "bic", int_aarch64_sve_bic>;
   defm EOR_PPzPP   : sve_int_pred_log<0b0010, "eor", int_aarch64_sve_eor>;
   defm SEL_PPPP    : sve_int_pred_log<0b0011, "sel", vselect>;
   defm ANDS_PPzPP  : sve_int_pred_log<0b0100, "ands", int_aarch64_sve_ands>;
   defm BICS_PPzPP  : sve_int_pred_log<0b0101, "bics", int_aarch64_sve_bics>;
   defm EORS_PPzPP  : sve_int_pred_log<0b0110, "eors", int_aarch64_sve_eors>;

llvm/lib/Target/AArch64/SVEInstrFormats.td

: I<(outs pprty:$Pdn), (ins PPRAny:$Pg, pprty:$_Pdn),
   let Inst{8-5}   = Pg;
   let Inst{4}     = 0;
   let Inst{3-0}   = Pdn;
 
   let Constraints = "$Pdn = $_Pdn";
   let Defs = [NZCV];
 }
 
-multiclass sve_int_pfirst<bits<5> opc, string asm> {
-  def : sve_int_pfirst_next<0b01, opc, asm, PPR8>;
+multiclass sve_int_pfirst<bits<5> opc, string asm, SDPatternOperator op> {
+  def _B : sve_int_pfirst_next<0b01, opc, asm, PPR8>;
+
+  def : SVE_2_Op_Pat<nxv16i1, op, nxv16i1, nxv16i1, !cast<Instruction>(NAME # _B)>;
 }
 
-multiclass sve_int_pnext<bits<5> opc, string asm> {
+multiclass sve_int_pnext<bits<5> opc, string asm, SDPatternOperator op> {
   def _B : sve_int_pfirst_next<0b00, opc, asm, PPR8>;
   def _H : sve_int_pfirst_next<0b01, opc, asm, PPR16>;
   def _S : sve_int_pfirst_next<0b10, opc, asm, PPR32>;
   def _D : sve_int_pfirst_next<0b11, opc, asm, PPR64>;
+
+  def : SVE_2_Op_Pat<nxv16i1, op, nxv16i1, nxv16i1, !cast<Instruction>(NAME # _B)>;
+  def : SVE_2_Op_Pat<nxv8i1, op, nxv8i1, nxv8i1, !cast<Instruction>(NAME # _H)>;
+  def : SVE_2_Op_Pat<nxv4i1, op, nxv4i1, nxv4i1, !cast<Instruction>(NAME # _S)>;
+  def : SVE_2_Op_Pat<nxv2i1, op, nxv2i1, nxv2i1, !cast<Instruction>(NAME # _D)>;
 }
 
 //===----------------------------------------------------------------------===//
 // SVE Predicate Count Group
 //===----------------------------------------------------------------------===//
 
 class sve_int_count_r<bits<2> sz8_64, bits<5> opc, string asm,
                       RegisterOperand dty, PPRRegOp pprty, RegisterOperand sty>

llvm/test/CodeGen/AArch64/sve-intrinsics-pred-operations.ll

 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
 
 ;
+; PFIRST
+;
+
+define <vscale x 16 x i1> @pfirst_b8(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a) {
+; CHECK-LABEL: pfirst_b8:
+; CHECK: pfirst p1.b, p0, p1.b
+; CHECK-NEXT: mov p0.b, p1.b
+; CHECK-NEXT: ret
+  %out = call <vscale x 16 x i1> @llvm.aarch64.sve.pfirst.nxv16i1(<vscale x 16 x i1> %pg,
+                                                                  <vscale x 16 x i1> %a)
+  ret <vscale x 16 x i1> %out
+}
+
+;
+; PNEXT
+;
+
+define <vscale x 16 x i1> @pnext_b8(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a) {
+; CHECK-LABEL: pnext_b8:
+; CHECK: pnext p1.b, p0, p1.b
+; CHECK-NEXT: mov p0.b, p1.b
+; CHECK-NEXT: ret
+  %out = call <vscale x 16 x i1> @llvm.aarch64.sve.pnext.nxv16i1(<vscale x 16 x i1> %pg,
+                                                                 <vscale x 16 x i1> %a)
+  ret <vscale x 16 x i1> %out
+}
+
+define <vscale x 8 x i1> @pnext_b16(<vscale x 8 x i1> %pg, <vscale x 8 x i1> %a) {
+; CHECK-LABEL: pnext_b16:
+; CHECK: pnext p1.h, p0, p1.h
+; CHECK-NEXT: mov p0.b, p1.b
+; CHECK-NEXT: ret
+  %out = call <vscale x 8 x i1> @llvm.aarch64.sve.pnext.nxv8i1(<vscale x 8 x i1> %pg,
+                                                               <vscale x 8 x i1> %a)
+  ret <vscale x 8 x i1> %out
+}
+
+define <vscale x 4 x i1> @pnext_b32(<vscale x 4 x i1> %pg, <vscale x 4 x i1> %a) {
+; CHECK-LABEL: pnext_b32:
+; CHECK: pnext p1.s, p0, p1.s
+; CHECK-NEXT: mov p0.b, p1.b
+; CHECK-NEXT: ret
+  %out = call <vscale x 4 x i1> @llvm.aarch64.sve.pnext.nxv4i1(<vscale x 4 x i1> %pg,
+                                                               <vscale x 4 x i1> %a)
+  ret <vscale x 4 x i1> %out
+}
+
+define <vscale x 2 x i1> @pnext_b64(<vscale x 2 x i1> %pg, <vscale x 2 x i1> %a) {
+; CHECK-LABEL: pnext_b64:
+; CHECK: pnext p1.d, p0, p1.d
+; CHECK-NEXT: mov p0.b, p1.b
+; CHECK-NEXT: ret
+  %out = call <vscale x 2 x i1> @llvm.aarch64.sve.pnext.nxv2i1(<vscale x 2 x i1> %pg,
+                                                               <vscale x 2 x i1> %a)
+  ret <vscale x 2 x i1> %out
+}
+
+;
 ; PUNPKHI
 ;
 
 define <vscale x 8 x i1> @punpkhi_b16(<vscale x 16 x i1> %a) {
 ; CHECK-LABEL: punpkhi_b16
 ; CHECK: punpkhi p0.h, p0.b
 ; CHECK-NEXT: ret
   %res = call <vscale x 8 x i1> @llvm.aarch64.sve.punpkhi.nxv8i1(<vscale x 16 x i1> %a)

 define <vscale x 2 x i1> @punpklo_b4(<vscale x 4 x i1> %a) {
 ; CHECK-LABEL: punpklo_b4
 ; CHECK: punpklo p0.h, p0.b
 ; CHECK-NEXT: ret
   %res = call <vscale x 2 x i1> @llvm.aarch64.sve.punpklo.nxv2i1(<vscale x 4 x i1> %a)
   ret <vscale x 2 x i1> %res
 }
 
+declare <vscale x 16 x i1> @llvm.aarch64.sve.pfirst.nxv16i1(<vscale x 16 x i1>, <vscale x 16 x i1>)
+
+declare <vscale x 16 x i1> @llvm.aarch64.sve.pnext.nxv16i1(<vscale x 16 x i1>, <vscale x 16 x i1>)
+declare <vscale x 8 x i1> @llvm.aarch64.sve.pnext.nxv8i1(<vscale x 8 x i1>, <vscale x 8 x i1>)
+declare <vscale x 4 x i1> @llvm.aarch64.sve.pnext.nxv4i1(<vscale x 4 x i1>, <vscale x 4 x i1>)
+declare <vscale x 2 x i1> @llvm.aarch64.sve.pnext.nxv2i1(<vscale x 2 x i1>, <vscale x 2 x i1>)
+
 declare <vscale x 8 x i1> @llvm.aarch64.sve.punpkhi.nxv8i1(<vscale x 16 x i1>)
 declare <vscale x 4 x i1> @llvm.aarch64.sve.punpkhi.nxv4i1(<vscale x 8 x i1>)
 declare <vscale x 2 x i1> @llvm.aarch64.sve.punpkhi.nxv2i1(<vscale x 4 x i1>)
 
 declare <vscale x 8 x i1> @llvm.aarch64.sve.punpklo.nxv8i1(<vscale x 16 x i1>)
 declare <vscale x 4 x i1> @llvm.aarch64.sve.punpklo.nxv4i1(<vscale x 8 x i1>)
 declare <vscale x 2 x i1> @llvm.aarch64.sve.punpklo.nxv2i1(<vscale x 4 x i1>)