[AArch64][SME2/SVE2p1] Add predicate-as-counter intrinsics for ptrue/cntp

[sdesmalen]

Summary

These intrinsics are used to implement:

• svptrue_c8(), svptrue_c16(), etc.
• svcntp_c8(svcount_t pnn, uint64_t vl), svcntp_c16(...), etc.

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

[david-arm]

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

4

Do we also want another RUN line for SME2, i.e.

; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme2 < %s | FileCheck %s

same as sve2p1-intrinsics-cntp.ll?

[sdesmalen]

Added RUN line for +sme2 to sve2p1-intrinsics-predicate-as-counter.ll

[sdesmalen]

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

4

Good spot, fixed!

[david-arm]

LGTM! Beautiful!

llvm/include/llvm/IR/Intrinsics.h

enum IITDescriptorKind {
       VecOfAnyPtrsToElt,
       VecElementArgument,
       Subdivide2Argument,
       Subdivide4Argument,
       VecOfBitcastsToInt,
       AMX,
       PPCQuad,
       AnyPtrToElt,
+      AArch64Svcount,
     } Kind;
 
     union {
       unsigned Integer_Width;
       unsigned Float_Width;
       unsigned Pointer_AddressSpace;
       unsigned Struct_NumElements;
       unsigned Argument_Info;

llvm/include/llvm/IR/Intrinsics.td

 def IIT_AMX : IIT_VT<x86amx, 51>;
 def IIT_PPCF128 : IIT_VT<ppcf128, 52>;
 def IIT_V3 : IIT_Vec<3, 53>;
 def IIT_EXTERNREF : IIT_VT<externref, 54>;
 def IIT_FUNCREF : IIT_VT<funcref, 55>;
 def IIT_ANYPTR_TO_ELT : IIT_Base<56>;
 def IIT_I2 : IIT_Int<2, 57>;
 def IIT_I4 : IIT_Int<4, 58>;
+def IIT_AARCH64_SVCOUNT : IIT_VT<aarch64svcount, 59>;
 }
 
 defvar IIT_all_FixedTypes = !filter(iit, IIT_all,
   !or(!isa<IIT_VT>(iit), !isa<IIT_Int>(iit)));
 
 defvar IIT_all_VectorTypes = !filter(iit, IIT_all,
   !isa<IIT_Vec>(iit));
 

 def llvm_empty_ty      : LLVMType<OtherVT>;                       // { }
 def llvm_descriptor_ty : LLVMPointerType<llvm_empty_ty>;          // { }*
 def llvm_metadata_ty   : LLVMType<MetadataVT>;                    // !{...}
 def llvm_token_ty      : LLVMType<token>;                         // token
 
 def llvm_x86mmx_ty     : LLVMType<x86mmx>;
 def llvm_ptrx86mmx_ty  : LLVMPointerType<llvm_x86mmx_ty>;         // <1 x i64>*
 
+def llvm_aarch64_svcount_ty : LLVMType<aarch64svcount>;
+
 def llvm_x86amx_ty     : LLVMType<x86amx>;
 
 def llvm_v2i1_ty       : LLVMType<v2i1>;     //   2 x i1
 def llvm_v4i1_ty       : LLVMType<v4i1>;     //   4 x i1
 def llvm_v8i1_ty       : LLVMType<v8i1>;     //   8 x i1
 def llvm_v16i1_ty      : LLVMType<v16i1>;    //  16 x i1
 def llvm_v32i1_ty      : LLVMType<v32i1>;    //  32 x i1
 def llvm_v64i1_ty      : LLVMType<v64i1>;    //  64 x i1

llvm/include/llvm/IR/IntrinsicsAArch64.td

let TargetPrefix = "aarch64" in {
   //
   foreach cmp = ["ge", "gt", "hi", "hs", "le", "lo", "ls", "lt"] in {
     def int_aarch64_sve_while # cmp # _x2
         : DefaultAttrsIntrinsic<[llvm_anyvector_ty, LLVMMatchType<0>],
                                 [llvm_i64_ty, llvm_i64_ty], [IntrNoMem]>;
   }
 
   //
+  // Predicate-as-counter intrinsics
+  //
+
+
+  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
+      : DefaultAttrsIntrinsic<[llvm_aarch64_svcount_ty], [], [IntrNoMem]>;
+
+  def int_aarch64_sve_cntp_c8
+      : DefaultAttrsIntrinsic<[llvm_i64_ty], [llvm_aarch64_svcount_ty, llvm_i32_ty],
+                              [IntrNoMem, ImmArg<ArgIndex<1>>]>;
+  def int_aarch64_sve_cntp_c16
+      : DefaultAttrsIntrinsic<[llvm_i64_ty], [llvm_aarch64_svcount_ty, llvm_i32_ty],
+                              [IntrNoMem, ImmArg<ArgIndex<1>>]>;
+  def int_aarch64_sve_cntp_c32
+      : DefaultAttrsIntrinsic<[llvm_i64_ty], [llvm_aarch64_svcount_ty, llvm_i32_ty],
+                              [IntrNoMem, ImmArg<ArgIndex<1>>]>;
+  def int_aarch64_sve_cntp_c64
+      : DefaultAttrsIntrinsic<[llvm_i64_ty], [llvm_aarch64_svcount_ty, llvm_i32_ty],
+                              [IntrNoMem, ImmArg<ArgIndex<1>>]>;
+
+  //
   // SME2 Intrinsics
   //
 
   class SME2_Matrix_ArrayVector_Single_Single_Intrinsic
     : DefaultAttrsIntrinsic<[],
                 [llvm_i32_ty,
                  llvm_anyvector_ty, LLVMMatchType<0>],
                 []>;

llvm/lib/IR/Function.cpp

case IIT_I1:
     OutputTable.push_back(IITDescriptor::get(IITDescriptor::Integer, 1));
     return;
   case IIT_I2:
     OutputTable.push_back(IITDescriptor::get(IITDescriptor::Integer, 2));
     return;
   case IIT_I4:
     OutputTable.push_back(IITDescriptor::get(IITDescriptor::Integer, 4));
     return;
+  case IIT_AARCH64_SVCOUNT:
+    OutputTable.push_back(IITDescriptor::get(IITDescriptor::AArch64Svcount, 0));
+    return;
   case IIT_I8:
     OutputTable.push_back(IITDescriptor::get(IITDescriptor::Integer, 8));
     return;
   case IIT_I16:
     OutputTable.push_back(IITDescriptor::get(IITDescriptor::Integer,16));
     return;
   case IIT_I32:
     OutputTable.push_back(IITDescriptor::get(IITDescriptor::Integer, 32));
static Type *DecodeFixedType(ArrayRef<Intrinsic::IITDescriptor> &Infos,
   case IITDescriptor::Token: return Type::getTokenTy(Context);
   case IITDescriptor::Metadata: return Type::getMetadataTy(Context);
   case IITDescriptor::Half: return Type::getHalfTy(Context);
   case IITDescriptor::BFloat: return Type::getBFloatTy(Context);
   case IITDescriptor::Float: return Type::getFloatTy(Context);
   case IITDescriptor::Double: return Type::getDoubleTy(Context);
   case IITDescriptor::Quad: return Type::getFP128Ty(Context);
   case IITDescriptor::PPCQuad: return Type::getPPC_FP128Ty(Context);
+  case IITDescriptor::AArch64Svcount:
+    return TargetExtType::get(Context, "aarch64.svcount");
 
   case IITDescriptor::Integer:
     return IntegerType::get(Context, D.Integer_Width);
   case IITDescriptor::Vector:
     return VectorType::get(DecodeFixedType(Infos, Tys, Context),
                            D.Vector_Width);
   case IITDescriptor::Pointer:
     return PointerType::get(DecodeFixedType(Infos, Tys, Context),
switch (D.Kind) {
     case IITDescriptor::Metadata: return !Ty->isMetadataTy();
     case IITDescriptor::Half: return !Ty->isHalfTy();
     case IITDescriptor::BFloat: return !Ty->isBFloatTy();
     case IITDescriptor::Float: return !Ty->isFloatTy();
     case IITDescriptor::Double: return !Ty->isDoubleTy();
     case IITDescriptor::Quad: return !Ty->isFP128Ty();
     case IITDescriptor::PPCQuad: return !Ty->isPPC_FP128Ty();
     case IITDescriptor::Integer: return !Ty->isIntegerTy(D.Integer_Width);
+    case IITDescriptor::AArch64Svcount:
+      return !isa<TargetExtType>(Ty) ||
+             cast<TargetExtType>(Ty)->getName() != "aarch64.svcount";
     case IITDescriptor::Vector: {
       VectorType *VT = dyn_cast<VectorType>(Ty);
       return !VT || VT->getElementCount() != D.Vector_Width ||
              matchIntrinsicType(VT->getElementType(), Infos, ArgTys,
                                 DeferredChecks, IsDeferredCheck);
     }
     case IITDescriptor::Pointer: {
       PointerType *PT = dyn_cast<PointerType>(Ty);

llvm/lib/Target/AArch64/SVEInstrFormats.td

 : Pat<(vtd (op vt1:$Op1, vt2:$Op2, vt3:$Op3, vt4:$Op4)),
       (inst $Op1, $Op2, $Op3, $Op4)>;
 
 class SVE_2_Op_Imm_Pat<ValueType vtd, SDPatternOperator op, ValueType vt1,
                        ValueType vt2, Operand ImmTy, Instruction inst>
 : Pat<(vtd (op vt1:$Op1, (vt2 ImmTy:$Op2))),
       (inst $Op1, ImmTy:$Op2)>;
 
+multiclass SVE2p1_Cntp_Pat<ValueType vtd, SDPatternOperator op, ValueType vt1,
+                           Instruction inst> {
+  def : Pat<(vtd (op vt1:$Op1, (i32 2))), (inst $Op1, 0)>;
+  def : Pat<(vtd (op vt1:$Op1, (i32 4))), (inst $Op1, 1)>;
+}
+
 class SVE_3_Op_Imm_Pat<ValueType vtd, SDPatternOperator op, ValueType vt1,
                        ValueType vt2, ValueType vt3, Operand ImmTy,
                        Instruction inst>
 : Pat<(vtd (op vt1:$Op1, vt2:$Op2, (vt3 ImmTy:$Op3))),
       (inst $Op1, $Op2, ImmTy:$Op3)>;
 
 class SVE_4_Op_Imm_Pat<ValueType vtd, SDPatternOperator op, ValueType vt1,
                        ValueType vt2, ValueType vt3, ValueType vt4,

 }
 
 multiclass sve2p1_two_way_dot_vvi<string mnemonic, bit u, SDPatternOperator intrinsic> {
   def NAME : sve2p1_two_way_dot_vvi<mnemonic, u>;
 
   def : SVE_4_Op_Imm_Pat<nxv4i32, intrinsic, nxv4i32, nxv8i16, nxv8i16, i32, VectorIndexS32b_timm, !cast<Instruction>(NAME)>;
 }
 
-class sve2p1_ptrue_pn<string mnemonic, bits<2> sz, PNRP8to15RegOp pnrty>
+class sve2p1_ptrue_pn<string mnemonic, bits<2> sz, PNRP8to15RegOp pnrty, SDPatternOperator op>
     : I<(outs pnrty:$PNd), (ins ), mnemonic, "\t$PNd",
-        "", []>, Sched<[]> {
+        "", [(set pnrty:$PNd, (op))]>, Sched<[]> {
   bits<3> PNd;
   let Inst{31-24}  = 0b00100101;
   let Inst{23-22} = sz;
   let Inst{21-3}  = 0b1000000111100000010;
   let Inst{2-0}   = PNd;
 
   let hasSideEffects = 0;
 }
 
 
 multiclass sve2p1_ptrue_pn<string mnemonic> {
- def _B : sve2p1_ptrue_pn<mnemonic, 0b00, PNR8_p8to15>;
- def _H : sve2p1_ptrue_pn<mnemonic, 0b01, PNR16_p8to15>;
- def _S : sve2p1_ptrue_pn<mnemonic, 0b10, PNR32_p8to15>;
- def _D : sve2p1_ptrue_pn<mnemonic, 0b11, PNR64_p8to15>;
+ def _B : sve2p1_ptrue_pn<mnemonic, 0b00, PNR8_p8to15, int_aarch64_sve_ptrue_c8>;
+ def _H : sve2p1_ptrue_pn<mnemonic, 0b01, PNR16_p8to15, int_aarch64_sve_ptrue_c16>;
+ def _S : sve2p1_ptrue_pn<mnemonic, 0b10, PNR32_p8to15, int_aarch64_sve_ptrue_c32>;
+ def _D : sve2p1_ptrue_pn<mnemonic, 0b11, PNR64_p8to15, int_aarch64_sve_ptrue_c64>;
 }
 
 
 // SVE extract mask predicate from predicate-as-counter
 class sve2p1_pred_as_ctr_to_mask_base<string mnemonic, bits<2> sz, bits<3> opc,
                                       RegisterOperand pprty, Operand idxty>
     : I<(outs pprty:$Pd), (ins PNRAny_p8to15:$PNn, idxty:$index),
         mnemonic, "\t$Pd, $PNn$index",
class sve2p1_pcount_pn<string mnemonic, bits<3> opc, bits<2> sz, PNRRegOp pnrty>
   let hasSideEffects = 0;
 }
 
 multiclass sve2p1_pcount_pn<string mnemonic, bits<3> opc> {
   def _B : sve2p1_pcount_pn<mnemonic, opc, 0b00, PNR8>;
   def _H : sve2p1_pcount_pn<mnemonic, opc, 0b01, PNR16>;
   def _S : sve2p1_pcount_pn<mnemonic, opc, 0b10, PNR32>;
   def _D : sve2p1_pcount_pn<mnemonic, opc, 0b11, PNR64>;
+
+  defm : SVE2p1_Cntp_Pat<i64, int_aarch64_sve_cntp_c8,  aarch64svcount, !cast<Instruction>(NAME # _B)>;
+  defm : SVE2p1_Cntp_Pat<i64, int_aarch64_sve_cntp_c16, aarch64svcount, !cast<Instruction>(NAME # _H)>;
+  defm : SVE2p1_Cntp_Pat<i64, int_aarch64_sve_cntp_c32, aarch64svcount, !cast<Instruction>(NAME # _S)>;
+  defm : SVE2p1_Cntp_Pat<i64, int_aarch64_sve_cntp_c64, aarch64svcount, !cast<Instruction>(NAME # _D)>;
 }
 
 
 // SVE integer compare scalar count and limit (predicate-as-counter)
 class sve2p1_int_while_rr_pn<string mnemonic, bits<2> sz, bits<3> opc,
                              PNRP8to15RegOp pnrty>
     : I<(outs pnrty:$PNd), (ins GPR64:$Rn, GPR64:$Rm, sve_vec_len_specifier_enum:$vl),
         mnemonic, "\t$PNd, $Rn, $Rm, $vl",

llvm/test/CodeGen/AArch64/sve2p1-intrinsics-cntp.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve2p1 < %s | FileCheck %s
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sme2 < %s | FileCheck %s
+
+define i64 @test_svcntp_c8_vlx2(target("aarch64.svcount") %pn) nounwind {
+; CHECK-LABEL: test_svcntp_c8_vlx2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cntp x0, pn0.b, vlx2
+; CHECK-NEXT:    ret
+  %res = call i64 @llvm.aarch64.sve.cntp.c8(target("aarch64.svcount") %pn, i32 2)
+  ret i64 %res
+}
+
+define i64 @test_svcntp_c8_vlx4(target("aarch64.svcount") %pn) nounwind {
+; CHECK-LABEL: test_svcntp_c8_vlx4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cntp x0, pn0.b, vlx4
+; CHECK-NEXT:    ret
+  %res = call i64 @llvm.aarch64.sve.cntp.c8(target("aarch64.svcount") %pn, i32 4)
+  ret i64 %res
+}
+
+define i64 @test_svcntp_c16_vlx2(target("aarch64.svcount") %pn) nounwind {
+; CHECK-LABEL: test_svcntp_c16_vlx2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cntp x0, pn0.h, vlx2
+; CHECK-NEXT:    ret
+  %res = call i64 @llvm.aarch64.sve.cntp.c16(target("aarch64.svcount") %pn, i32 2)
+  ret i64 %res
+}
+
+define i64 @test_svcntp_c16_vlx4(target("aarch64.svcount") %pn) nounwind {
+; CHECK-LABEL: test_svcntp_c16_vlx4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cntp x0, pn0.h, vlx4
+; CHECK-NEXT:    ret
+  %res = call i64 @llvm.aarch64.sve.cntp.c16(target("aarch64.svcount") %pn, i32 4)
+  ret i64 %res
+}
+
+define i64 @test_svcntp_c32_vlx2(target("aarch64.svcount") %pn) nounwind {
+; CHECK-LABEL: test_svcntp_c32_vlx2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cntp x0, pn0.s, vlx2
+; CHECK-NEXT:    ret
+  %res = call i64 @llvm.aarch64.sve.cntp.c32(target("aarch64.svcount") %pn, i32 2)
+  ret i64 %res
+}
+
+define i64 @test_svcntp_c32_vlx4(target("aarch64.svcount") %pn) nounwind {
+; CHECK-LABEL: test_svcntp_c32_vlx4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cntp x0, pn0.s, vlx4
+; CHECK-NEXT:    ret
+  %res = call i64 @llvm.aarch64.sve.cntp.c32(target("aarch64.svcount") %pn, i32 4)
+  ret i64 %res
+}
+
+define i64 @test_svcntp_c64_vlx2(target("aarch64.svcount") %pn) nounwind {
+; CHECK-LABEL: test_svcntp_c64_vlx2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cntp x0, pn0.d, vlx2
+; CHECK-NEXT:    ret
+  %res = call i64 @llvm.aarch64.sve.cntp.c64(target("aarch64.svcount") %pn, i32 2)
+  ret i64 %res
+}
+
+define i64 @test_svcntp_c64_vlx4(target("aarch64.svcount") %pn) nounwind {
+; CHECK-LABEL: test_svcntp_c64_vlx4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cntp x0, pn0.d, vlx4
+; CHECK-NEXT:    ret
+  %res = call i64 @llvm.aarch64.sve.cntp.c64(target("aarch64.svcount") %pn, i32 4)
+  ret i64 %res
+}
+
+
+declare i64 @llvm.aarch64.sve.cntp.c8(target("aarch64.svcount"), i32)
+declare i64 @llvm.aarch64.sve.cntp.c16(target("aarch64.svcount"), i32)
+declare i64 @llvm.aarch64.sve.cntp.c32(target("aarch64.svcount"), i32)
+declare i64 @llvm.aarch64.sve.cntp.c64(target("aarch64.svcount"), i32)

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

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64 -mattr=+sve2p1 < %s | FileCheck %s
+; RUN: llc -mtriple=aarch64 -mattr=+sme2 < %s | FileCheck %s
+
+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
+}
+
+define target("aarch64.svcount") @ptrue_h() nounwind {
+; CHECK-LABEL: ptrue_h:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue pn8.h
+; CHECK-NEXT:    mov p0.b, p8.b
+; CHECK-NEXT:    ret
+  %res = call target("aarch64.svcount") @llvm.aarch64.sve.ptrue.c16()
+  ret target("aarch64.svcount") %res
+}
+
+define target("aarch64.svcount") @ptrue_s() nounwind {
+; CHECK-LABEL: ptrue_s:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue pn8.s
+; CHECK-NEXT:    mov p0.b, p8.b
+; CHECK-NEXT:    ret
+  %res = call target("aarch64.svcount") @llvm.aarch64.sve.ptrue.c32()
+  ret target("aarch64.svcount") %res
+}
+
+define target("aarch64.svcount") @ptrue_d() nounwind {
+; CHECK-LABEL: ptrue_d:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue pn8.d
+; CHECK-NEXT:    mov p0.b, p8.b
+; CHECK-NEXT:    ret
+  %res = call target("aarch64.svcount") @llvm.aarch64.sve.ptrue.c64()
+  ret target("aarch64.svcount") %res
+}
+
+declare target("aarch64.svcount") @llvm.aarch64.sve.ptrue.c8()
+declare target("aarch64.svcount") @llvm.aarch64.sve.ptrue.c16()
+declare target("aarch64.svcount") @llvm.aarch64.sve.ptrue.c32()
+declare target("aarch64.svcount") @llvm.aarch64.sve.ptrue.c64()