- Add lowering of cntsd back into LowerINTRINSIC_WO_CHAIN

- Remove patterns for cntsd and add for AArch64rdsvl
- Add nsw/nuw flags to mul of cntsd in EmitSMELd1St1

Author: kmclaughlin-arm

clang/lib/CodeGen/TargetBuiltins/ARM.cpp

@@ -4307,7 +4307,8 @@ Value *CodeGenFunction::EmitSMELd1St1(const SVETypeFlags &TypeFlags,
         CGM.getIntrinsic(Intrinsic::aarch64_sme_cntsd);
     llvm::Value *StreamingVectorLengthCall =
         Builder.CreateMul(Builder.CreateCall(StreamingVectorLength),
-                          llvm::ConstantInt::get(Int64Ty, 8), "svl");
+                          llvm::ConstantInt::get(Int64Ty, 8), "svl",
+                          /* HasNUW */ true, /* HasNSW */ true);
     llvm::Value *Mulvl =
         Builder.CreateMul(StreamingVectorLengthCall, Ops[4], "mulvl");
     // The type of the ptr parameter is void *, so use Int8Ty here.

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

@@ -940,20 +940,16 @@ bool AArch64DAGToDAGISel::SelectRDVLImm(SDValue N, SDValue &Imm) {
   return false;
 }
 
-// Given `cntsd = (rdsvl, #1) >> 3`, attempt to return a suitable multiplier
-// for RDSVL to calculate `cntsd << N`, i.e. `rdsvl, #(1 << (N - 3))`.
+// Returns a suitable RDSVL multiplier from a left shift.
 template <signed Low, signed High>
 bool AArch64DAGToDAGISel::SelectRDSVLShiftImm(SDValue N, SDValue &Imm) {
   if (!isa<ConstantSDNode>(N))
     return false;
 
-  int64_t ShlImm = cast<ConstantSDNode>(N)->getSExtValue();
-  if (ShlImm >= 3) {
-    int64_t MulImm = 1 << (ShlImm - 3);
-    if (MulImm >= Low && MulImm <= High) {
-      Imm = CurDAG->getSignedTargetConstant(MulImm, SDLoc(N), MVT::i32);
-      return true;
-    }
+  int64_t MulImm = 1 << cast<ConstantSDNode>(N)->getSExtValue();
+  if (MulImm >= Low && MulImm <= High) {
+    Imm = CurDAG->getSignedTargetConstant(MulImm, SDLoc(N), MVT::i32);
+    return true;
   }
 
   return false;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -6266,6 +6266,16 @@ SDValue AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
   case Intrinsic::aarch64_sve_clz:
     return DAG.getNode(AArch64ISD::CTLZ_MERGE_PASSTHRU, DL, Op.getValueType(),
                        Op.getOperand(2), Op.getOperand(3), Op.getOperand(1));
+  case Intrinsic::aarch64_sme_cntsd: {
+    auto Flags = SDNodeFlags();
+    Flags.setNoUnsignedWrap(true);
+    Flags.setNoSignedWrap(true);
+    Flags.setExact(true);
+    SDValue Bytes = DAG.getNode(AArch64ISD::RDSVL, DL, Op.getValueType(),
+                                DAG.getConstant(1, DL, MVT::i32));
+    return DAG.getNode(ISD::SRL, DL, Op.getValueType(), Bytes,
+                       DAG.getConstant(3, DL, MVT::i32), Flags);
+  }
   case Intrinsic::aarch64_sve_cnt: {
     SDValue Data = Op.getOperand(3);
     // CTPOP only supports integer operands.
@@ -19180,9 +19190,6 @@ static SDValue performMulCombine(SDNode *N, SelectionDAG &DAG,
        if (ConstValue.sge(1) && ConstValue.sle(16))
          return SDValue();
 
-  if (getIntrinsicID(N0.getNode()) == Intrinsic::aarch64_sme_cntsd)
-    return SDValue();
-
   // Multiplication of a power of two plus/minus one can be done more
   // cheaply as shift+add/sub. For now, this is true unilaterally. If
   // future CPUs have a cheaper MADD instruction, this may need to be

llvm/lib/Target/AArch64/AArch64SMEInstrInfo.td

@@ -139,25 +139,10 @@ def RDSVLI_XI  : sve_int_read_vl_a<0b0, 0b11111, "rdsvl", /*streaming_sve=*/0b1>
 def ADDSPL_XXI : sve_int_arith_vl<0b1, "addspl", /*streaming_sve=*/0b1>;
 def ADDSVL_XXI : sve_int_arith_vl<0b0, "addsvl", /*streaming_sve=*/0b1>;
 
-def : Pat<(AArch64rdsvl (i32 simm6_32b:$imm)), (RDSVLI_XI simm6_32b:$imm)>;
-
-// e.g. cntsb() * imm
-def : Pat<(i64 (mul (int_aarch64_sme_cntsd), (sme_cntsb_mul_imm i64:$imm))),
-          (RDSVLI_XI (!cast<SDNodeXForm>("trunc_imm") $imm))>;
-def : Pat<(i64 (mul (int_aarch64_sme_cntsd), (sme_cntsh_mul_imm i64:$imm))),
-          (UBFMXri (RDSVLI_XI (!cast<SDNodeXForm>("trunc_imm") $imm)), 1, 63)>;
-def : Pat<(i64 (mul (int_aarch64_sme_cntsd), (sme_cntsw_mul_imm i64:$imm))),
-          (UBFMXri (RDSVLI_XI (!cast<SDNodeXForm>("trunc_imm") $imm)), 2, 63)>;
-def : Pat<(i64 (mul (int_aarch64_sme_cntsd), (sme_cntsd_mul_imm i64:$imm))),
-          (UBFMXri (RDSVLI_XI (!cast<SDNodeXForm>("trunc_imm") $imm)), 3, 63)>;
-
-def : Pat<(i64 (shl (int_aarch64_sme_cntsd), (sme_cnts_shl_imm i64:$imm))),
+def : Pat<(i64 (shl (AArch64rdsvl (i32 1)), (sme_cnts_shl_imm i64:$imm))),
           (RDSVLI_XI (!cast<SDNodeXForm>("trunc_imm") $imm))>;
 
-// cntsh, cntsw, cntsd
-def : Pat<(i64 (shl (int_aarch64_sme_cntsd), (i64 2))), (UBFMXri (RDSVLI_XI 1), 1, 63)>;
-def : Pat<(i64 (shl (int_aarch64_sme_cntsd), (i64 1))), (UBFMXri (RDSVLI_XI 1), 2, 63)>;
-def : Pat<(i64 (int_aarch64_sme_cntsd)), (UBFMXri (RDSVLI_XI 1), 3, 63)>;
+def : Pat<(AArch64rdsvl (i32 simm6_32b:$imm)), (RDSVLI_XI simm6_32b:$imm)>;
 }
 
 let Predicates = [HasSME] in {

llvm/test/CodeGen/AArch64/sme-intrinsics-rdsvl.ll

@@ -46,46 +46,43 @@ define i64 @cntsd() {
 define i64 @sme_cntsb_mul() {
 ; CHECK-LABEL: sme_cntsb_mul:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    rdsvl x0, #2
+; CHECK-NEXT:    rdsvl x0, #4
 ; CHECK-NEXT:    ret
   %v = call i64 @llvm.aarch64.sme.cntsd()
   %shl = shl nuw nsw i64 %v, 3
-  %res = mul i64 %shl, 2
+  %res = mul nuw nsw i64 %shl, 4
   ret i64 %res
 }
 
 define i64 @sme_cntsh_mul() {
 ; CHECK-LABEL: sme_cntsh_mul:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    rdsvl x8, #5
-; CHECK-NEXT:    lsr x0, x8, #1
+; CHECK-NEXT:    rdsvl x0, #4
 ; CHECK-NEXT:    ret
   %v = call i64 @llvm.aarch64.sme.cntsd()
   %shl = shl nuw nsw i64 %v, 2
-  %res = mul i64 %shl, 5
+  %res = mul nuw nsw i64 %shl, 8
   ret i64 %res
 }
 
 define i64 @sme_cntsw_mul() {
 ; CHECK-LABEL: sme_cntsw_mul:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    rdsvl x8, #7
-; CHECK-NEXT:    lsr x0, x8, #2
+; CHECK-NEXT:    rdsvl x0, #4
 ; CHECK-NEXT:    ret
   %v = call i64 @llvm.aarch64.sme.cntsd()
   %shl = shl nuw nsw i64 %v, 1
-  %res = mul i64 %shl, 7
+  %res = mul nuw nsw i64 %shl, 16
   ret i64 %res
 }
 
 define i64 @sme_cntsd_mul() {
 ; CHECK-LABEL: sme_cntsd_mul:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    rdsvl x8, #3
-; CHECK-NEXT:    lsr x0, x8, #1
+; CHECK-NEXT:    rdsvl x0, #4
 ; CHECK-NEXT:    ret
   %v = call i64 @llvm.aarch64.sme.cntsd()
-  %res = mul i64 %v, 12
+  %res = mul nuw nsw i64 %v, 32
   ret i64 %res
 }