[AArch64][SelectionDAG] Optimize multiplication by constant

[Allen]

Summary

InstCombine will merge the x*5+x into x*6,
then decompose MUL By const into shift/add or shift/sub, eg:
Change the costmodel to lower a = b * C where C = (2^n + 1) * 2^m to

add     w0, w0, w0, lsl n
lsl     w0, w0, m

Fixes AArch64 part of https://github.com/llvm/llvm-project/issues/57255.

[craig.topper]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13735

Is this copied from RISC-V? That doesn't look like an AArch64 instruction name

[Allen]

update comment

[Allen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13735

yes, I updated comment, thanks @craig.topper

[hiraditya]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13732

What is the rationale behind checking Imm+1 etc?

[Allen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13732

I think it can use shl + sub to replace mul, such as
https://gcc.godbolt.org/z/szh7Eb9K8

BTW: some other target, such as RISCVTargetLowering::decomposeMulByConstant also has such logic

llvm/test/CodeGen/AArch64/mul_pow2.ll

217

will be fixed with D132325

[hiraditya]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13732

as per godbolt link, it is already happening in some cases?

[efriedma]

It's not obvious that the replacement sequences are consistently faster. At least on some cores, "add x8, x8, w0, sxtw #1" and "smull x0, w0, w8" have exactly the same throughput, so transforming from the smull to a two-instruction sequence involving the add isn't really profitable.

On a related note, many cores have optimizations for arithmetic with lsl #n, so we should prefer that over uxtw/sxtw.

llvm/test/CodeGen/AArch64/mul_pow2.ll

179

I think the suggestion misses a zero-extension. Should be able to ubfiz, though.

281

I think you can save an instruction here: instead of "-(x*4+x*2)", compute x*2-x*8.

500

This seems to be overriding our existing logic here to produce a worse result.

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

169 ↗

(On Diff #454269)

Probably need some logic to allow folding inch/dech, assuming there isn't some reason to avoid them.

[dmgreen]

Hi - I had been looking at mul vs add+shift recently, but had not had the time to get very far and had not got to the important part yet - exactly where and when should we be doing the transform, especially when you consider all the various cpus present.

AArch64 usually handles mul vs add/shift here: https://github.com/llvm/llvm-project/blob/6c6c4f6a9b3ef2d7db937cb78784245ea8a61418/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp#L14465
With a cost-model of sorts here: https://github.com/llvm/llvm-project/blob/6c6c4f6a9b3ef2d7db937cb78784245ea8a61418/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp#L14442

My understanding (but some of this may not be correct) is that:

• The existing cost model is conservative and could do with adjustment, but is there for a reason. It prevents the transform when there is an add/sub user (to create a madd) or if there is a zext/sext operand (to create a umull).
• A good cost model is hard to be precise about but add cost 1, shift costs 1, add+shift usually costs 2 but can be 1 for some cpus and constants. mul costs somewhere between 2 and 5, depending on the cpu, size (i32/i64) and the values in the registers. madd costs the same as mul (so the add is free). Newer cpus have a lower cost for mul, especially i64 mul.
• Replacing a mul with 2 instructions is probably good, a mul with 3 instructions is more iffy.
• I was using https://godbolt.org/z/cPTEMnP5x with different C to compare when we do the transform compared to gcc.
• The cost model is certainly worse where the operands lead into a load/store pointer operand, as the add can be folded into the addressing mode so it won't form a madd. Small shifts can sometimes be free again. I was probably planning to alter the existing profitability checks.

So whilst this may be an improvement on some targets over what we have already, a lot of the changes here either look like obvious regressions (cntd/decd/etc), are non-obvious whether they are regressions or not (madd vs add+lsl+add_lsl), or are not longer testing the point of the tests (machine-combiner-madd.ll).

[Allen]

Thanks for your information, and it seems not a easy work :)

llvm/test/CodeGen/AArch64/mul_pow2.ll

179

Yes, this has been addressed with D132325, thanks.

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

169 ↗

(On Diff #454269)

Thanks, I'll try this first

[hiraditya]

I agree with @eli.friedman that decomposing multiply into two separate instructions may not be profitable.
Even when latency of multiply is high, having two instructions in the pipeline adds complexity (dependencies, re-ordering, retiring etc).

Another thing to consider would be the number of ALUs and MACs(Multiply Accumulate units).
In case the ALU pressure is high MAC operations may be more efficient as there wont be stalls.

[Allen]

Add more restrict to disable these have negitate effort

[Allen]

rebase to fix new added case Transforms/SeparateConstOffsetFromGEP/AArch64/split-gep.ll

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13726

We can't just expect the first user of a const will be the mul we are interested in.

All the other handling of decomposing mul into add+shift is currently done in performMulCombine. Would it be better to just alter the code that is already there? It would make it easier to be more precise with the costmodel.

llvm/test/CodeGen/AArch64/mul_pow2.ll

147–148

I think that (considering the mov as free in terms of latency), in this case the madd would be worth 2-3, the lsl+add_lsl+add would cost 3-4. It would depend heavily on the exact cpu though.
For i64 muls the madd would have a higher cost (it was 4 on one cpu I tested, but newer cpus are better).

307

Is this the case you are interested in? Could we change the existing costmodel to be more precise with which sub operand it considers free?

// Conservatively do not lower to shift+add+shift if the mul might be
// folded into madd or msub.
if (N->hasOneUse() && (N->use_begin()->getOpcode() == ISD::ADD ||
                       N->use_begin()->getOpcode() == ISD::SUB))
  return SDValue();

[Allen]

refactor with costmodel

[Allen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

13726

Apply your commend, and refactor it with costmodel, thanks

llvm/test/CodeGen/AArch64/mul_pow2.ll

307

Thanks for your idea, I can try this, but why the sub operand can be considered free ?

307

Done, thanks

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14562

Do you want to check specifically for a constant that fits into the immediate operand of an add/sub? If it doesn't fit, it doesn't really matter if it's constant.

14634

Move computation of "ShiftedVal1" into the if statment?

llvm/test/CodeGen/AArch64/mul_pow2.ll

306

There is one other possible sequence here:

mov w8, #6
mov w9, #-1
madd w0, w0, w8, w9

This is obviously not that exciting in isolation, but it might make sense if we can hoist the "mov" instructions out of a loop.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14562

I don't think it matters it the AddSub has one use? (It might for some other optimizations, but that can be a separate patch).

It matters that for a sub the AddSub->operand(1) == N. If it is AddSub->operand(0) == N, then we cannot fold to a msub, so that shouldn't block the transform to add+shifts.
The isIntOrFPConstant check can check for isLegalAddImmediate, to be a little more precise where it will generate a mul+addri or subri.

[Allen]

llvm/test/CodeGen/AArch64/mul_pow2.ll

306

hi @efriedma, I try to generate the new sequence as you showed, but I don't know how to put a const into a register? as the operand of madd should not be const value.

For the const **AddSub->getOperand(1)**, even when I use a ISD::ADD , it still retrun a const value.

 SDValue Const = DAG.getNode(ISD::ADD, DL, VT, AddSub->getOperand(1), DAG.getConstant(0, DL, VT));

(gdb) p Const->dump()
t5: i32 = Constant<-1>
$10 = void
(gdb) p AddSub->getOperand(1)->dump()
 t5: i32 = Constant<-1>

[Allen]

address comment

[Allen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14562

Thanks @dmgreen and @efriedma for detail suggestion

14634

Done ,thanks

llvm/test/CodeGen/AArch64/mul_pow2.ll

147–148

done, block the change as we can make use of madd

[efriedma]

llvm/test/CodeGen/AArch64/mul_pow2.ll

306

You don't have to do anything special to "put a const in a register". For example, if you write unsigned a(unsigned x) { return x * 0x33333333 + 0xF0F0F0F0; }, you get code like I suggested. Getting the same code with smaller immediates is just a matter of making sure we pick the pattern for madd, and not the pattern for sub-with-imm.

[Allen]

llvm/test/CodeGen/AArch64/mul_pow2.ll

306

I try to address the commit with D134336

[Allen]

abandoned as the accepted in D134934 and D134336

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 }
 
 // Return false to prevent folding
 // (mul (add x, c1), c2) -> (add (mul x, c2), c2*c1) in DAGCombine,
 // if the folding leads to worse code.
 bool AArch64TargetLowering::isMulAddWithConstProfitable(
     SDValue AddNode, SDValue ConstNode) const {
   // Let the DAGCombiner decide for vector types and large types.
   const EVT VT = AddNode.getValueType();
   if (VT.isVector() || VT.getScalarSizeInBits() > 64)
     return true;
 
   // It is worse if c1 is legal add immediate, while c1*c2 is not
   // and has to be composed by at least two instructions.
   const ConstantSDNode *C1Node = cast<ConstantSDNode>(AddNode.getOperand(1));
   const ConstantSDNode *C2Node = cast<ConstantSDNode>(ConstNode);
   const int64_t C1 = C1Node->getSExtValue();
   const APInt C1C2 = C1Node->getAPIntValue() * C2Node->getAPIntValue();
   if (!isLegalAddImmediate(C1) || isLegalAddImmediate(C1C2.getSExtValue()))
     return true;
   SmallVector<AArch64_IMM::ImmInsnModel, 4> Insn;
   AArch64_IMM::expandMOVImm(C1C2.getZExtValue(), VT.getSizeInBits(), Insn);
   if (Insn.size() > 1)
     return false;
 
   // Default to true and let the DAGCombiner decide.
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)))
       return SDValue();
     // Conservatively do not lower to shift+add+shift if the mul might be
     // folded into madd or msub.
     if (N->hasOneUse() && (N->use_begin()->getOpcode() == ISD::ADD ||
-                           N->use_begin()->getOpcode() == ISD::SUB))
+                           N->use_begin()->getOpcode() == ISD::SUB)) {
+      SDNode *AddSub = *N->use_begin();
+      // Shouldn't block the transform to add+shifts as the operand of madd/msub
+      // should not be const value.
+      SDValue BinaryOP;
+      if (N == AddSub->getOperand(0).getNode())
+        BinaryOP = AddSub->getOperand(1);
+      else
+        BinaryOP = AddSub->getOperand(0);
+      const ConstantSDNode *OpC = dyn_cast<ConstantSDNode>(BinaryOP);
+      bool match;
+      // For sub, may transform to msub only when OpC is a reg.
+      if (AddSub->getOpcode() == ISD::SUB)
+        match = !OpC;
+      // For add, may transform to madd when OpC is a reg or not match the
+      // isLegalAddImmediate.
+      const TargetLowering &TLI = DAG.getTargetLoweringInfo();
+      if (AddSub->getOpcode() == ISD::ADD)
+        match = !OpC ||
+                !TLI.isLegalAddImmediate(OpC->getAPIntValue().getSExtValue());
+      if (match)
         return SDValue();
     }
+  }
   // Use ShiftedConstValue instead of ConstValue to support both shift+add/sub
   // and shift+add+shift.
   APInt ShiftedConstValue = ConstValue.ashr(TrailingZeroes);
 
   unsigned ShiftAmt;
   // Is the shifted value the LHS operand of the add/sub?
   bool ShiftValUseIsN0 = true;
   // Do we need to negate the result?
   bool NegateResult = false;
+  // Is the sub has 2 shifted value operands?
+  bool Sub2Shift = false;
 
   if (ConstValue.isNonNegative()) {
     // (mul x, 2^N + 1) => (add (shl x, N), x)
     // (mul x, 2^N - 1) => (sub (shl x, N), x)
     // (mul x, (2^N + 1) * 2^M) => (shl (add (shl x, N), x), M)
+    // (mul x, (2^(N-M) - 1) * 2^M) => (sub (shl x, N), (shl x, M))
     APInt SCVMinus1 = ShiftedConstValue - 1;
+    APInt SCVPlus1 = ShiftedConstValue + 1;
     APInt CVPlus1 = ConstValue + 1;
     if (SCVMinus1.isPowerOf2()) {
       ShiftAmt = SCVMinus1.logBase2();
       AddSubOpc = ISD::ADD;
     } else if (CVPlus1.isPowerOf2()) {
       ShiftAmt = CVPlus1.logBase2();
       AddSubOpc = ISD::SUB;
+    } else if (SCVPlus1.isPowerOf2()) {
+      ShiftAmt = SCVPlus1.logBase2() + TrailingZeroes;
+      AddSubOpc = ISD::SUB;
+      Sub2Shift = true;
     } else
       return SDValue();
   } else {
     // (mul x, -(2^N - 1)) => (sub x, (shl x, N))
     // (mul x, -(2^N + 1)) => - (add (shl x, N), x)
     APInt CVNegPlus1 = -ConstValue + 1;
     APInt CVNegMinus1 = -ConstValue - 1;
     if (CVNegPlus1.isPowerOf2()) {
       ShiftAmt = CVNegPlus1.logBase2();
       AddSubOpc = ISD::SUB;
       ShiftValUseIsN0 = false;
     } else if (CVNegMinus1.isPowerOf2()) {
       ShiftAmt = CVNegMinus1.logBase2();
       AddSubOpc = ISD::ADD;
       NegateResult = true;
     } else
       return SDValue();
   }
 
-  SDValue ShiftedVal = DAG.getNode(ISD::SHL, DL, VT, N0,
+  SDValue ShiftedVal0 = DAG.getNode(ISD::SHL, DL, VT, N0,
                                     DAG.getConstant(ShiftAmt, DL, MVT::i64));
-
-  SDValue AddSubN0 = ShiftValUseIsN0 ? ShiftedVal : N0;
-  SDValue AddSubN1 = ShiftValUseIsN0 ? N0 : ShiftedVal;
+  if (Sub2Shift) {
+    SDValue ShiftedVal1 = DAG.getNode(
+        ISD::SHL, DL, VT, N0, DAG.getConstant(TrailingZeroes, DL, MVT::i64));
+    return DAG.getNode(ISD::SUB, DL, VT, ShiftedVal0, ShiftedVal1);
+  }
+  SDValue AddSubN0 = ShiftValUseIsN0 ? ShiftedVal0 : N0;
+  SDValue AddSubN1 = ShiftValUseIsN0 ? N0 : ShiftedVal0;
   SDValue Res = DAG.getNode(AddSubOpc, DL, VT, AddSubN0, AddSubN1);
   assert(!(NegateResult && TrailingZeroes) &&
          "NegateResult and TrailingZeroes cannot both be true for now.");
   // Negate the result.
   if (NegateResult)
     return DAG.getNode(ISD::SUB, DL, VT, DAG.getConstant(0, DL, VT), Res);
   // Shift the result.
   if (TrailingZeroes)

llvm/test/CodeGen/AArch64/mul_pow2.ll

; GISEL-NEXT:    ret
   %mul = mul nsw i64 %ext, 6
   ret i64 %mul
 }
 
 define i32 @test6_madd(i32 %x, i32 %y) {
 ; CHECK-LABEL: test6_madd:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov w8, #6
 ; CHECK-NEXT:    madd w0, w0, w8, w1
 ; CHECK-NEXT:    ret
 ;
 ; GISEL-LABEL: test6_madd:
 ; GISEL:       // %bb.0:
 ; GISEL-NEXT:    mov w8, #6
 ; GISEL-NEXT:    madd w0, w0, w8, w1
 ; GISEL-NEXT:    ret
 
   %mul = mul nsw i32 %x, 6

 ; GISEL-NEXT:    msub w0, w0, w8, w1
 ; GISEL-NEXT:    ret
 
   %mul = mul nsw i32 %x, 6
   %sub = sub i32 %y, %mul
   ret i32 %sub
 }
 
 define i64 @test6_umaddl(i32 %x, i64 %y) {
 ; CHECK-LABEL: test6_umaddl:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov w8, #6
 ; CHECK-NEXT:    umaddl x0, w0, w8, x1
 ; CHECK-NEXT:    ret
 ;
 ; GISEL-LABEL: test6_umaddl:
 ; GISEL:       // %bb.0:

 ; GISEL-NEXT:    ret
 
   %ext = sext i32 %x to i64
   %mul = mul nsw i64 %ext, 6
   %add = add i64 %mul, %y
   ret i64 %add
 }
 
 define i64 @test6_umsubl(i32 %x, i64 %y) {
 ; CHECK-LABEL: test6_umsubl:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov w8, #6
 ; CHECK-NEXT:    umsubl x0, w0, w8, x1
 ; CHECK-NEXT:    ret
 ;
 ; GISEL-LABEL: test6_umsubl:
 ; GISEL:       // %bb.0:

 
 define i64 @test6_smnegl(i32 %x) {
 ; CHECK-LABEL: test6_smnegl:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov w8, #6
 ; CHECK-NEXT:    smnegl x0, w0, w8
 ; CHECK-NEXT:    ret
 ;
 ; GISEL-LABEL: test6_smnegl:
 ; GISEL:       // %bb.0:
 ; GISEL-NEXT:    mov w8, #6
 ; GISEL-NEXT:    smnegl x0, w0, w8
 ; GISEL-NEXT:    ret
 
   %ext = sext i32 %x to i64
   %mul = mul nsw i64 %ext, 6
   %sub = sub i64 0, %mul
   ret i64 %sub
 }
 
+define i32 @mull6_sub(i32 %x) {
+; CHECK-LABEL: mull6_sub:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add w8, w0, w0, lsl #1
+; CHECK-NEXT:    lsl w8, w8, #1
+; CHECK-NEXT:    sub w0, w8, #1
+; CHECK-NEXT:    ret
+;
+; GISEL-LABEL: mull6_sub:
+; GISEL:       // %bb.0:
+; GISEL-NEXT:    mov w8, #6
+; GISEL-NEXT:    mul w8, w0, w8
+; GISEL-NEXT:    sub w0, w8, #1
+; GISEL-NEXT:    ret
+  %mul = mul nsw i32 %x, 6
+  %sub = add nsw i32 %mul, -1
+  ret i32 %sub
+}
+
 define i32 @test7(i32 %x) {
 ; CHECK-LABEL: test7:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    lsl w8, w0, #3
 ; CHECK-NEXT:    sub w0, w8, w0
 ; CHECK-NEXT:    ret
 ;
 ; GISEL-LABEL: test7:
; GISEL-NEXT:    ret
 
   %mul = mul nsw i32 %x, 13
   ret i32 %mul
 }
 
 define i32 @test14(i32 %x) {
 ; CHECK-LABEL: test14:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w8, #14
-; CHECK-NEXT:    mul w0, w0, w8
+; CHECK-NEXT:    lsl w8, w0, #4
+; CHECK-NEXT:    sub w0, w8, w0, lsl #1
 ; CHECK-NEXT:    ret
 ;
 ; GISEL-LABEL: test14:
 ; GISEL:       // %bb.0:
 ; GISEL-NEXT:    mov w8, #14
 ; GISEL-NEXT:    mul w0, w0, w8
 ; GISEL-NEXT:    ret
 
; GISEL-NEXT:    ret
   %mul = mul nsw i32 %x, -2
   ret i32 %mul
 }
 
 define i32 @ntest3(i32 %x) {
 ; CHECK-LABEL: ntest3:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    sub w0, w0, w0, lsl #2
 ; CHECK-NEXT:    ret
 ;
 ; GISEL-LABEL: ntest3:
 ; GISEL:       // %bb.0:
 ; GISEL-NEXT:    sub w0, w0, w0, lsl #2
 ; GISEL-NEXT:    ret
 
   %mul = mul nsw i32 %x, -3
   ret i32 %mul

 ; CHECK-NEXT:    movi v2.4s, #1, msl #16
 ; CHECK-NEXT:    shl v0.4s, v0.4s, #6
 ; CHECK-NEXT:    sub v0.4s, v1.4s, v0.4s
 ; CHECK-NEXT:    and v0.16b, v0.16b, v2.16b
 ; CHECK-NEXT:    ret
 ;
 ; GISEL-LABEL: muladd_demand_commute:
 ; GISEL:       // %bb.0:
-; GISEL-NEXT:    adrp x8, .LCPI42_1
-; GISEL-NEXT:    ldr q2, [x8, :lo12:.LCPI42_1]
-; GISEL-NEXT:    adrp x8, .LCPI42_0
+; GISEL-NEXT:    adrp x8, .LCPI43_1
+; GISEL-NEXT:    ldr q2, [x8, :lo12:.LCPI43_1]
+; GISEL-NEXT:    adrp x8, .LCPI43_0
 ; GISEL-NEXT:    mla v1.4s, v0.4s, v2.4s
-; GISEL-NEXT:    ldr q0, [x8, :lo12:.LCPI42_0]
+; GISEL-NEXT:    ldr q0, [x8, :lo12:.LCPI43_0]
 ; GISEL-NEXT:    and v0.16b, v1.16b, v0.16b
 ; GISEL-NEXT:    ret
   %m = mul <4 x i32> %x, <i32 131008, i32 131008, i32 131008, i32 131008>
   %a = add <4 x i32> %m, %y
   %r = and <4 x i32> %a, <i32 131071, i32 131071, i32 131071, i32 131071>
   ret <4 x i32> %r
 }