[SME2/SVE2p1] Extend llvm.aarch64.sve.convert.to/from.svbool to accept target("aarch64.svcount")

[sdesmalen]

Summary

The convert intrinsics can be used to implement existing operations on svcount_t
when the actual bits/content of the predicate register doesn't matter (such
as PSEL, which copies the full contents of the first source register to the
destination register).

[david-arm]

LGTM!

[sdesmalen]

@hassnaa-arm made a good suggestion offline, that we could reuse the existing convert.to/from.svbool intrinsics.
I've updated the implementation.

[david-arm]

LGTM! The new version looks better now as it reuses the existing intrinsic.

llvm/include/llvm/IR/IntrinsicsAArch64.td

 def int_aarch64_sve_ptest_any   : AdvSIMD_SVE_PTEST_Intrinsic;
 def int_aarch64_sve_ptest_first : AdvSIMD_SVE_PTEST_Intrinsic;
 def int_aarch64_sve_ptest_last  : AdvSIMD_SVE_PTEST_Intrinsic;
 
 //
 // Reinterpreting data
 //
 
-def int_aarch64_sve_convert_from_svbool : DefaultAttrsIntrinsic<[llvm_anyvector_ty],
+def int_aarch64_sve_convert_from_svbool : DefaultAttrsIntrinsic<[llvm_any_ty],
                                                     [llvm_nxv16i1_ty],
                                                     [IntrNoMem]>;
 
 def int_aarch64_sve_convert_to_svbool : DefaultAttrsIntrinsic<[llvm_nxv16i1_ty],
-                                                  [llvm_anyvector_ty],
+                                                  [llvm_any_ty],
                                                   [IntrNoMem]>;
 
 //
 // Gather loads: scalar base + vector offsets
 //
 
 // 64 bit unscaled offsets
 def int_aarch64_sve_ld1_gather : AdvSIMD_GatherLoad_SV_64b_Offsets_Intrinsic;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

case Intrinsic::aarch64_sve_cnt: {
     if (Data.getValueType().isFloatingPoint())
       Data = DAG.getNode(ISD::BITCAST, dl, Op.getValueType(), Data);
     return DAG.getNode(AArch64ISD::CTPOP_MERGE_PASSTHRU, dl, Op.getValueType(),
                        Op.getOperand(2), Data, Op.getOperand(1));
   }
   case Intrinsic::aarch64_sve_dupq_lane:
     return LowerDUPQLane(Op, DAG);
   case Intrinsic::aarch64_sve_convert_from_svbool:
+    if (Op.getValueType() == MVT::aarch64svcount)
+      return DAG.getNode(ISD::BITCAST, dl, Op.getValueType(), Op.getOperand(1));
     return getSVEPredicateBitCast(Op.getValueType(), Op.getOperand(1), DAG);
   case Intrinsic::aarch64_sve_convert_to_svbool:
+    if (Op.getOperand(1).getValueType() == MVT::aarch64svcount)
+      return DAG.getNode(ISD::BITCAST, dl, MVT::nxv16i1, Op.getOperand(1));
     return getSVEPredicateBitCast(MVT::nxv16i1, Op.getOperand(1), DAG);
   case Intrinsic::aarch64_sve_fneg:
     return DAG.getNode(AArch64ISD::FNEG_MERGE_PASSTHRU, dl, Op.getValueType(),
                        Op.getOperand(2), Op.getOperand(3), Op.getOperand(1));
   case Intrinsic::aarch64_sve_frintp:
     return DAG.getNode(AArch64ISD::FCEIL_MERGE_PASSTHRU, dl, Op.getValueType(),
                        Op.getOperand(2), Op.getOperand(3), Op.getOperand(1));
   case Intrinsic::aarch64_sve_frintm:

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

 instCombineConvertFromSVBool(InstCombiner &IC, IntrinsicInst &II) {
   // If the reinterpret instruction operand is a PHI Node
   if (isa<PHINode>(II.getArgOperand(0)))
     return processPhiNode(IC, II);
 
   if (auto BinOpCombine = tryCombineFromSVBoolBinOp(IC, II))
     return BinOpCombine;
 
+  // Ignore converts to/from svcount_t.
+  if (isa<TargetExtType>(II.getArgOperand(0)->getType()) ||
+      isa<TargetExtType>(II.getType()))
+    return std::nullopt;
+
   SmallVector<Instruction *, 32> CandidatesForRemoval;
   Value *Cursor = II.getOperand(0), *EarliestReplacement = nullptr;
 
   const auto *IVTy = cast<VectorType>(II.getType());
 
   // Walk the chain of conversions.
   while (Cursor) {
     // If the type of the cursor has fewer lanes than the final result, zeroing

llvm/test/CodeGen/AArch64/sve-intrinsics-reinterpret.ll

 ; CHECK-NEXT:    ptrue p1.d
 ; CHECK-NEXT:    punpklo p1.h, p1.b
 ; CHECK-NEXT:    and p0.b, p0/z, p0.b, p1.b
 ; CHECK-NEXT:    ret
   %res = call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv1i1(<vscale x 1 x i1> %arg)
   ret <vscale x 16 x i1> %res
 }
 
+define <vscale x 16 x i1> @reinterpret_bool_from_svcount(target("aarch64.svcount") %pg) "target-features"="+sme2" {
+; CHECK-LABEL: reinterpret_bool_from_svcount:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ret
+  %out = call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.taarch64.svcountt(target("aarch64.svcount") %pg)
+  ret <vscale x 16 x i1> %out
+}
+
 ;
 ; Converting from svbool_t
 ;
 
 define <vscale x 16 x i1> @reinterpret_bool_to_b(<vscale x 16 x i1> %pg) {
 ; CHECK-LABEL: reinterpret_bool_to_b:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ret

 define <vscale x 1 x i1> @reinterpret_bool_to_q(<vscale x 16 x i1> %pg) {
 ; CHECK-LABEL: reinterpret_bool_to_q:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ret
   %out = call <vscale x 1 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv1i1(<vscale x 16 x i1> %pg)
   ret <vscale x 1 x i1> %out
 }
 
+define target("aarch64.svcount") @reinterpret_bool_to_svcount(<vscale x 16 x i1> %pg) "target-features"="+sme2" {
+; CHECK-LABEL: reinterpret_bool_to_svcount:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ret
+  %out = call target("aarch64.svcount") @llvm.aarch64.sve.convert.from.svbool.taarch64.svcountt(<vscale x 16 x i1> %pg)
+  ret target("aarch64.svcount") %out
+}
+
+
 ; Reinterpreting a ptrue should not introduce an `and` instruction.
 define <vscale x 16 x i1> @reinterpret_ptrue() {
 ; CHECK-LABEL: reinterpret_ptrue:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.h
 ; CHECK-NEXT:    ret
   %in = tail call <vscale x 8 x i1> @llvm.aarch64.sve.ptrue.nxv8i1(i32 31)
   %out = tail call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv8i1(<vscale x 8 x i1> %in)

 declare <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 immarg)
 declare <vscale x 8 x i1> @llvm.aarch64.sve.cmpgt.nxv8i16(<vscale x 8 x i1>, <vscale x 8 x i16>, <vscale x 8 x i16>)
 
 declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv16i1(<vscale x 16 x i1>)
 declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv8i1(<vscale x 8 x i1>)
 declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv4i1(<vscale x 4 x i1>)
 declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv2i1(<vscale x 2 x i1>)
 declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv1i1(<vscale x 1 x i1>)
+declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.taarch64.svcountt(target("aarch64.svcount"))
 
 declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv16i1(<vscale x 16 x i1>)
 declare <vscale x 8 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv8i1(<vscale x 16 x i1>)
 declare <vscale x 4 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv4i1(<vscale x 16 x i1>)
 declare <vscale x 2 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv2i1(<vscale x 16 x i1>)
 declare <vscale x 1 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv1i1(<vscale x 16 x i1>)
+declare target("aarch64.svcount") @llvm.aarch64.sve.convert.from.svbool.taarch64.svcountt(<vscale x 16 x i1>)