[AArc64][SVE] Fold constant multiply of element count

[c-rhodes]

Summary

E.g.

%0 = tail call i64 @llvm.aarch64.sve.cntw(i32 31)
%mul = mul i64 %0, <const>

Should emit:

cntw    x0, all, mul #<const>

For <const> in the range 1-16.

Patch by Kerry McLaughlin

[sdesmalen]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

168

nit: this is perhaps personal preference, but I always prefer code like this being in the form:

if (invalidating condition)
  return false;
if (invalidating condition2)
  return false;
  :
return true;

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9551

this is missing a default case.

[c-rhodes]

• Refactored SelectRDVLImm and SelectShiftImm into same function by adding a extra template argument for Shift.
• Added default to switch in IsSVECntIntrinsic

[c-rhodes]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

168

Done for the check Scale is a factor of imm. I've left the low/high range checks as I think it's more readable as it is.

[sdesmalen]

LGTM! Thanks @c-rhodes

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

163

nit: Low/High -> Min/Max ?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

9571

nit:

There's nothing to do if the immediate can be packed into an instruction

->

Allow the scaling to be folded into the `cnt` instruction by preventing
the scaling to be obscured here. That makes it easier to pattern match.

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

bool SelectDupZeroOrUndef(SDValue N) {
     default:
       break;
     }
 
     return false;
   }
 
   template<MVT::SimpleValueType VT>
   bool SelectSVEAddSubImm(SDValue N, SDValue &Imm, SDValue &Shift) {
     return SelectSVEAddSubImm(N, VT, Imm, Shift);
   }
 
   template<MVT::SimpleValueType VT>
   bool SelectSVELogicalImm(SDValue N, SDValue &Imm) {
     return SelectSVELogicalImm(N, VT, Imm);
   }
 
+  // Returns a suitable CNT/INC/DEC/RDVL multiplier to calculate VSCALE*N.
+  template<signed Min, signed Max, signed Scale, bool Shift>
+  bool SelectCntImm(SDValue N, SDValue &Imm) {
+    if (!isa<ConstantSDNode>(N))
+      return false;
+
+    int64_t MulImm = cast<ConstantSDNode>(N)->getSExtValue();
+    if (Shift)
+      MulImm = 1 << MulImm;
+
+    if ((MulImm % std::abs(Scale)) != 0)
+      return false;
+
+    MulImm /= Scale;
+    if ((MulImm >= Min) && (MulImm <= Max)) {
+      Imm = CurDAG->getTargetConstant(MulImm, SDLoc(N), MVT::i32);
+      return true;
+    }
+
+    return false;
+  }
+
   /// Form sequences of consecutive 64/128-bit registers for use in NEON
   /// instructions making use of a vector-list (e.g. ldN, tbl). Vecs must have
   /// between 1 and 4 elements. If it contains a single element that is returned
   /// unchanged; otherwise a REG_SEQUENCE value is returned.
   SDValue createDTuple(ArrayRef<SDValue> Vecs);
   SDValue createQTuple(ArrayRef<SDValue> Vecs);
 
   /// Generic helper for the createDTuple/createQTuple

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

AArch64TargetLowering::BuildSDIVPow2(SDNode *N, const APInt &Divisor,
   // negate the result.
   if (Divisor.isNonNegative())
     return SRA;
 
   Created.push_back(SRA.getNode());
   return DAG.getNode(ISD::SUB, DL, VT, DAG.getConstant(0, DL, VT), SRA);
 }
 
+static bool IsSVECntIntrinsic(SDValue S) {
+  switch(getIntrinsicID(S.getNode())) {
+  default:
+    break;
+  case Intrinsic::aarch64_sve_cntb:
+  case Intrinsic::aarch64_sve_cnth:
+  case Intrinsic::aarch64_sve_cntw:
+  case Intrinsic::aarch64_sve_cntd:
+    return true;
+  }
+  return false;
+}
+
 static SDValue performMulCombine(SDNode *N, SelectionDAG &DAG,
                                  TargetLowering::DAGCombinerInfo &DCI,
                                  const AArch64Subtarget *Subtarget) {
   if (DCI.isBeforeLegalizeOps())
     return SDValue();
 
   // The below optimizations require a constant RHS.
   if (!isa<ConstantSDNode>(N->getOperand(1)))
     return SDValue();
 
+  SDValue N0 = N->getOperand(0);
   ConstantSDNode *C = cast<ConstantSDNode>(N->getOperand(1));
   const APInt &ConstValue = C->getAPIntValue();
 
+  // Allow the scaling to be folded into the `cnt` instruction by preventing
+  // the scaling to be obscured here. This makes it easier to pattern match.
+  if (IsSVECntIntrinsic(N0) ||
+     (N0->getOpcode() == ISD::TRUNCATE &&
+      (IsSVECntIntrinsic(N0->getOperand(0)))))
+       if (ConstValue.sge(1) && ConstValue.sle(16))
+         return SDValue();
+
   // Multiplication of a power of two plus/minus one can be done more
   // cheaply as as shift+add/sub. For now, this is true unilaterally. If
   // future CPUs have a cheaper MADD instruction, this may need to be
   // gated on a subtarget feature. For Cyclone, 32-bit MADD is 4 cycles and
   // 64-bit is 5 cycles, so this is always a win.
   // More aggressively, some multiplications N0 * C can be lowered to
   // shift+add+shift if the constant C = A * B where A = 2^N + 1 and B = 2^M,
   // e.g. 6=3*2=(2+1)*2.
   // TODO: consider lowering more cases, e.g. C = 14, -6, -14 or even 45
   // which equals to (1+2)*16-(1+2).
-  SDValue N0 = N->getOperand(0);
   // TrailingZeroes is used to test if the mul can be lowered to
   // shift+add+shift.
   unsigned TrailingZeroes = ConstValue.countTrailingZeros();
   if (TrailingZeroes) {
     // Conservatively do not lower to shift+add+shift if the mul might be
     // folded into smul or umul.
     if (N0->hasOneUse() && (isSignExtended(N0.getNode(), DAG) ||
                             isZeroExtended(N0.getNode(), DAG)))

llvm/lib/Target/AArch64/SVEInstrFormats.td

 def sve_incdec_imm : Operand<i32>, TImmLeaf<i32, [{
   return (((uint32_t)Imm) > 0) && (((uint32_t)Imm) < 17);
 }]> {
   let ParserMatchClass = Imm1_16Operand;
   let EncoderMethod = "getSVEIncDecImm";
   let DecoderMethod = "DecodeSVEIncDecImm";
 }
 
+// This allows i32 immediate extraction from i64 based arithmetic.
+def sve_cnt_mul_imm : ComplexPattern<i32, 1, "SelectCntImm<1, 16, 1, false>">;
+def sve_cnt_shl_imm : ComplexPattern<i32, 1, "SelectCntImm<1, 16, 1, true>">;
+
 //===----------------------------------------------------------------------===//
 // SVE PTrue - These are used extensively throughout the pattern matching so
 //             it's important we define them first.
 //===----------------------------------------------------------------------===//
 
 class sve_int_ptrue<bits<2> sz8_64, bits<3> opc, string asm, PPRRegOp pprty,
                     ValueType vt, SDPatternOperator op>
 : I<(outs pprty:$Pd), (ins sve_pred_enum:$pattern),

 multiclass sve_int_count<bits<3> opc, string asm, SDPatternOperator op> {
   def NAME : sve_int_count<opc, asm>;
 
   def : InstAlias<asm # "\t$Rd, $pattern",
                   (!cast<Instruction>(NAME) GPR64:$Rd, sve_pred_enum:$pattern, 1), 1>;
   def : InstAlias<asm # "\t$Rd",
                   (!cast<Instruction>(NAME) GPR64:$Rd, 0b11111, 1), 2>;
 
+  def : Pat<(i64 (mul (op sve_pred_enum:$pattern), (sve_cnt_mul_imm i32:$imm))),
+            (!cast<Instruction>(NAME) sve_pred_enum:$pattern, sve_incdec_imm:$imm)>;
+
+  def : Pat<(i64 (shl (op sve_pred_enum:$pattern), (i64 (sve_cnt_shl_imm i32:$imm)))),
+            (!cast<Instruction>(NAME) sve_pred_enum:$pattern, sve_incdec_imm:$imm)>;
+
   def : Pat<(i64 (op sve_pred_enum:$pattern)),
             (!cast<Instruction>(NAME) sve_pred_enum:$pattern, 1)>;
 }
 
 class sve_int_countvlv<bits<5> opc, string asm, ZPRRegOp zprty>
 : I<(outs zprty:$Zdn), (ins zprty:$_Zdn, sve_pred_enum:$pattern, sve_incdec_imm:$imm4),
   asm, "\t$Zdn, $pattern, mul $imm4",
   "",

llvm/test/CodeGen/AArch64/sve-intrinsics-counting-elems.ll

 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
 
 ;
 ; CNTB
 ;
 
 define i64 @cntb() {
 ; CHECK-LABEL: cntb:
 ; CHECK: cntb x0, vl2
 ; CHECK-NEXT: ret
   %out = call i64 @llvm.aarch64.sve.cntb(i32 2)
   ret i64 %out
 }
 
+define i64 @cntb_mul3() {
+; CHECK-LABEL: cntb_mul3:
+; CHECK: cntb x0, vl6, mul #3
+; CHECK-NEXT: ret
+  %cnt = call i64 @llvm.aarch64.sve.cntb(i32 6)
+  %out = mul i64 %cnt, 3
+  ret i64 %out
+}
+
+define i64 @cntb_mul4() {
+; CHECK-LABEL: cntb_mul4:
+; CHECK: cntb x0, vl8, mul #4
+; CHECK-NEXT: ret
+  %cnt = call i64 @llvm.aarch64.sve.cntb(i32 8)
+  %out = mul i64 %cnt, 4
+  ret i64 %out
+}
+
 ;
 ; CNTH
 ;
 
 define i64 @cnth() {
 ; CHECK-LABEL: cnth:
 ; CHECK: cnth x0, vl3
 ; CHECK-NEXT: ret
   %out = call i64 @llvm.aarch64.sve.cnth(i32 3)
   ret i64 %out
 }
 
+define i64 @cnth_mul5() {
+; CHECK-LABEL: cnth_mul5:
+; CHECK: cnth x0, vl7, mul #5
+; CHECK-NEXT: ret
+  %cnt = call i64 @llvm.aarch64.sve.cnth(i32 7)
+  %out = mul i64 %cnt, 5
+  ret i64 %out
+}
+
+define i64 @cnth_mul8() {
+; CHECK-LABEL: cnth_mul8:
+; CHECK: cnth x0, vl5, mul #8
+; CHECK-NEXT: ret
+  %cnt = call i64 @llvm.aarch64.sve.cnth(i32 5)
+  %out = mul i64 %cnt, 8
+  ret i64 %out
+}
+
 ;
 ; CNTW
 ;
 
 define i64 @cntw() {
 ; CHECK-LABEL: cntw:
 ; CHECK: cntw x0, vl4
 ; CHECK-NEXT: ret
   %out = call i64 @llvm.aarch64.sve.cntw(i32 4)
   ret i64 %out
 }
 
+define i64 @cntw_mul11() {
+; CHECK-LABEL: cntw_mul11:
+; CHECK: cntw x0, vl8, mul #11
+; CHECK-NEXT: ret
+  %cnt = call i64 @llvm.aarch64.sve.cntw(i32 8)
+  %out = mul i64 %cnt, 11
+  ret i64 %out
+}
+
+define i64 @cntw_mul2() {
+; CHECK-LABEL: cntw_mul2:
+; CHECK: cntw x0, vl6, mul #2
+; CHECK-NEXT: ret
+  %cnt = call i64 @llvm.aarch64.sve.cntw(i32 6)
+  %out = mul i64 %cnt, 2
+  ret i64 %out
+}
+
 ;
 ; CNTD
 ;
 
 define i64 @cntd() {
 ; CHECK-LABEL: cntd:
 ; CHECK: cntd x0, vl5
 ; CHECK-NEXT: ret
   %out = call i64 @llvm.aarch64.sve.cntd(i32 5)
   ret i64 %out
 }
 
+define i64 @cntd_mul15() {
+; CHECK-LABEL: cntd_mul15:
+; CHECK: cntd x0, vl16, mul #15
+; CHECK-NEXT: ret
+  %cnt = call i64 @llvm.aarch64.sve.cntd(i32 9)
+  %out = mul i64 %cnt, 15
+  ret i64 %out
+}
+
+define i64 @cntd_mul16() {
+; CHECK-LABEL: cntd_mul16:
+; CHECK: cntd x0, vl32, mul #16
+; CHECK-NEXT: ret
+  %cnt = call i64 @llvm.aarch64.sve.cntd(i32 10)
+  %out = mul i64 %cnt, 16
+  ret i64 %out
+}
+
 ;
 ; CNTP
 ;
 
 define i64 @cntp_b8(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a) {
 ; CHECK-LABEL: cntp_b8:
 ; CHECK: cntp x0, p0, p1.b
 ; CHECK-NEXT: ret