[AArch64][SME2/SVE2p1] Add predicate-as-counter intrinsics for pext (multi)

[sdesmalen]

Summary

These intrinsics are used to implement the pext intrinsics that extract
two predicates (mask) from a predicate-as-counter value, e.g.

__attribute__((arm_streaming))
 svboolx2_t svpext_lane_c8_x2(svcount_t pnn, uint64_t imm);

As described in https://github.com/ARM-software/acle/pull/217

[kmclaughlin]

LGTM

llvm/include/llvm/IR/IntrinsicsAArch64.td

let TargetPrefix = "aarch64" in {
   // Predicate-as-counter intrinsics
   //
 
   def int_aarch64_sve_pext
       : DefaultAttrsIntrinsic<[llvm_anyvector_ty],
                               [llvm_aarch64_svcount_ty, llvm_i32_ty],
                               [IntrNoMem, ImmArg<ArgIndex<1>>]>;
 
+  def int_aarch64_sve_pext_x2
+      : DefaultAttrsIntrinsic<[llvm_anyvector_ty, LLVMMatchType<0>],
+                              [llvm_aarch64_svcount_ty, llvm_i32_ty],
+                              [IntrNoMem, ImmArg<ArgIndex<1>>]>;
 
   def int_aarch64_sve_ptrue_c8
       : DefaultAttrsIntrinsic<[llvm_aarch64_svcount_ty], [], [IntrNoMem]>;
   def int_aarch64_sve_ptrue_c16
       : DefaultAttrsIntrinsic<[llvm_aarch64_svcount_ty], [], [IntrNoMem]>;
   def int_aarch64_sve_ptrue_c32
       : DefaultAttrsIntrinsic<[llvm_aarch64_svcount_ty], [], [IntrNoMem]>;
   def int_aarch64_sve_ptrue_c64

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

void SelectPostLoad(SDNode *N, unsigned NumVecs, unsigned Opc,
                          unsigned SubRegIdx);
   void SelectLoadLane(SDNode *N, unsigned NumVecs, unsigned Opc);
   void SelectPostLoadLane(SDNode *N, unsigned NumVecs, unsigned Opc);
   void SelectPredicatedLoad(SDNode *N, unsigned NumVecs, unsigned Scale,
                             unsigned Opc_rr, unsigned Opc_ri,
                             bool IsIntr = false);
   void SelectDestructiveMultiIntrinsic(SDNode *N, unsigned NumVecs,
                                        bool IsZmMulti, unsigned Opcode);
+  void SelectPExtPair(SDNode *N, unsigned Opc);
   void SelectWhilePair(SDNode *N, unsigned Opc);
   void SelectCVTIntrinsic(SDNode *N, unsigned NumVecs, unsigned Opcode);
   void SelectClamp(SDNode *N, unsigned NumVecs, unsigned Opcode);
   void SelectUnaryMultiIntrinsic(SDNode *N, unsigned NumOutVecs,
                                  bool IsTupleInput, unsigned Opc);
   void SelectFrintFromVT(SDNode *N, unsigned NumVecs, unsigned Opcode);
 
   template <unsigned MaxIdx, unsigned Scale>
case 2: // 64-bit
     break;
   default:
     return 0;
   }
 
   return (Opcodes.size() <= Offset) ? 0 : Opcodes[Offset];
 }
 
+// This function is almost identical to SelectWhilePair, but has an
+// extra check on the range of the immediate operand.
+// TODO: Merge these two functions together at some point?
+void AArch64DAGToDAGISel::SelectPExtPair(SDNode *N, unsigned Opc) {
+  // Immediate can be either 0 or 1.
+  if (ConstantSDNode *Imm = dyn_cast<ConstantSDNode>(N->getOperand(2)))
+    if (Imm->getZExtValue() > 1)
+      return;
+
+  SDLoc DL(N);
+  EVT VT = N->getValueType(0);
+  SDValue Ops[] = {N->getOperand(1), N->getOperand(2)};
+  SDNode *WhilePair = CurDAG->getMachineNode(Opc, DL, MVT::Untyped, Ops);
+  SDValue SuperReg = SDValue(WhilePair, 0);
+
+  for (unsigned I = 0; I < 2; ++I)
+    ReplaceUses(SDValue(N, I), CurDAG->getTargetExtractSubreg(
+                                   AArch64::psub0 + I, DL, VT, SuperReg));
+
+  CurDAG->RemoveDeadNode(N);
+}
+
 void AArch64DAGToDAGISel::SelectWhilePair(SDNode *N, unsigned Opc) {
   SDLoc DL(N);
   EVT VT = N->getValueType(0);
 
   SDValue Ops[] = {N->getOperand(1), N->getOperand(2)};
 
   SDNode *WhilePair = CurDAG->getMachineNode(Opc, DL, MVT::Untyped, Ops);
   SDValue SuperReg = SDValue(WhilePair, 0);
case Intrinsic::aarch64_sve_sunpk_x4:
       return;
     case Intrinsic::aarch64_sve_uunpk_x4:
       if (auto Op = SelectOpcodeFromVT<SelectTypeKind::Int>(
               Node->getValueType(0),
               {0, AArch64::UUNPK_VG4_4Z2Z_H, AArch64::UUNPK_VG4_4Z2Z_S,
                AArch64::UUNPK_VG4_4Z2Z_D}))
         SelectUnaryMultiIntrinsic(Node, 4, /*IsTupleInput=*/true, Op);
       return;
+    case Intrinsic::aarch64_sve_pext_x2: {
+      if (auto Op = SelectOpcodeFromVT<SelectTypeKind::AnyType>(
+              Node->getValueType(0),
+              {AArch64::PEXT_2PCI_B, AArch64::PEXT_2PCI_H, AArch64::PEXT_2PCI_S,
+               AArch64::PEXT_2PCI_D}))
+        SelectPExtPair(Node, Op);
+      return;
+    }
     }
     break;
   }
   case ISD::INTRINSIC_VOID: {
     unsigned IntNo = cast<ConstantSDNode>(Node->getOperand(1))->getZExtValue();
     if (Node->getNumOperands() >= 3)
       VT = Node->getOperand(2)->getValueType(0);
     switch (IntNo) {

llvm/test/CodeGen/AArch64/sve2p1-intrinsics-predicate-as-counter.ll

; CHECK-NEXT:    ret
   ret <vscale x 2 x i1> %res
 }
 
 declare <vscale x 16 x i1> @llvm.aarch64.sve.pext.nxv16i1(target("aarch64.svcount"), i32)
 declare <vscale x 8 x i1> @llvm.aarch64.sve.pext.nxv8i1(target("aarch64.svcount"), i32)
 declare <vscale x 4 x i1> @llvm.aarch64.sve.pext.nxv4i1(target("aarch64.svcount"), i32)
 declare <vscale x 2 x i1> @llvm.aarch64.sve.pext.nxv2i1(target("aarch64.svcount"), i32)
 
+define  {<vscale x 16 x i1>,<vscale x 16 x i1>} @pext_x2_b(target("aarch64.svcount") %x) nounwind {
+; CHECK-LABEL: pext_x2_b:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    addvl sp, sp, #-1
+; CHECK-NEXT:    str p8, [sp, #7, mul vl] // 2-byte Folded Spill
+; CHECK-NEXT:    mov p8.b, p0.b
+; CHECK-NEXT:    pext { p0.b, p1.b }, pn8[1]
+; CHECK-NEXT:    ldr p8, [sp, #7, mul vl] // 2-byte Folded Reload
+; CHECK-NEXT:    addvl sp, sp, #1
+; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+  %res = call {<vscale x 16 x i1>,<vscale x 16 x i1>} @llvm.aarch64.sve.pext.x2.nxv16i1(target("aarch64.svcount") %x, i32 1)
+  ret  {<vscale x 16 x i1>,<vscale x 16 x i1>} %res
+}
+
+define {<vscale x 8 x i1>,<vscale x 8 x i1>} @pext_x2_h(target("aarch64.svcount") %x) nounwind {
+; CHECK-LABEL: pext_x2_h:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    addvl sp, sp, #-1
+; CHECK-NEXT:    str p8, [sp, #7, mul vl] // 2-byte Folded Spill
+; CHECK-NEXT:    mov p8.b, p0.b
+; CHECK-NEXT:    pext { p0.h, p1.h }, pn8[1]
+; CHECK-NEXT:    ldr p8, [sp, #7, mul vl] // 2-byte Folded Reload
+; CHECK-NEXT:    addvl sp, sp, #1
+; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+  %res = call {<vscale x 8 x i1>,<vscale x 8 x i1>} @llvm.aarch64.sve.pext.x2.nxv8i1(target("aarch64.svcount") %x, i32 1)
+  ret {<vscale x 8 x i1>,<vscale x 8 x i1>} %res
+}
+
+define {<vscale x 4 x i1>,<vscale x 4 x i1>} @pext_x2_s(target("aarch64.svcount") %x) nounwind {
+; CHECK-LABEL: pext_x2_s:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    addvl sp, sp, #-1
+; CHECK-NEXT:    str p8, [sp, #7, mul vl] // 2-byte Folded Spill
+; CHECK-NEXT:    mov p8.b, p0.b
+; CHECK-NEXT:    pext { p0.s, p1.s }, pn8[1]
+; CHECK-NEXT:    ldr p8, [sp, #7, mul vl] // 2-byte Folded Reload
+; CHECK-NEXT:    addvl sp, sp, #1
+; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+  %res = call {<vscale x 4 x i1>,<vscale x 4 x i1>} @llvm.aarch64.sve.pext.x2.nxv4i1(target("aarch64.svcount") %x, i32 1)
+  ret {<vscale x 4 x i1>,<vscale x 4 x i1>} %res
+}
+
+define {<vscale x 2 x i1>,<vscale x 2 x i1>} @pext_x2_d(target("aarch64.svcount") %x) nounwind {
+; CHECK-LABEL: pext_x2_d:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    addvl sp, sp, #-1
+; CHECK-NEXT:    str p8, [sp, #7, mul vl] // 2-byte Folded Spill
+; CHECK-NEXT:    mov p8.b, p0.b
+; CHECK-NEXT:    pext { p0.d, p1.d }, pn8[1]
+; CHECK-NEXT:    ldr p8, [sp, #7, mul vl] // 2-byte Folded Reload
+; CHECK-NEXT:    addvl sp, sp, #1
+; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+  %res = call {<vscale x 2 x i1>,<vscale x 2 x i1>} @llvm.aarch64.sve.pext.x2.nxv2i1(target("aarch64.svcount") %x, i32 1)
+  ret {<vscale x 2 x i1>,<vscale x 2 x i1>} %res
+}
+
+declare {<vscale x 16 x i1>,<vscale x 16 x i1>} @llvm.aarch64.sve.pext.x2.nxv16i1(target("aarch64.svcount"), i32)
+declare {<vscale x 8 x i1>,<vscale x 8 x i1>} @llvm.aarch64.sve.pext.x2.nxv8i1(target("aarch64.svcount"), i32)
+declare {<vscale x 4 x i1>,<vscale x 4 x i1>} @llvm.aarch64.sve.pext.x2.nxv4i1(target("aarch64.svcount"), i32)
+declare {<vscale x 2 x i1>,<vscale x 2 x i1>} @llvm.aarch64.sve.pext.x2.nxv2i1(target("aarch64.svcount"), i32)
+
 define target("aarch64.svcount") @ptrue_b() nounwind {
 ; CHECK-LABEL: ptrue_b:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue pn8.b
 ; CHECK-NEXT:    mov p0.b, p8.b
 ; CHECK-NEXT:    ret
   %res = call target("aarch64.svcount") @llvm.aarch64.sve.ptrue.c8()
   ret target("aarch64.svcount") %res