[AArch64] Adjust operand sequence for Add+Sub to combine more inline shift

[bcl5980]

Summary

((X >> C) - Y) + Z --> (Z - Y) + (X >> C)

Fix AArch part: #55714

[bcl5980]

Ping

[dmgreen]

Have you tried testing this? I think it will get stuck in DAG combines.

Otherwise the idea sounds good. Have you considered doing it for GlobalISel, or in a way that GlobalISel would also benefit?

[bcl5980]

In D136158#3878602, @dmgreen wrote:

Have you tried testing this? I think it will get stuck in DAG combines.

Otherwise the idea sounds good. Have you considered doing it for GlobalISel, or in a way that GlobalISel would also benefit?

I haven't found any stuck in DAG combine local and it looks all of the premerge test is also passed:
https://reviews.llvm.org/harbormaster/build/291788/

Which combination do you think will cause the dead loop in DAG combine with this patch? I can try to create a test for that.

I can try to do it in GISel also in another patch later.

[dmgreen]

In D136158#3878631, @bcl5980 wrote:

In D136158#3878602, @dmgreen wrote:

Have you tried testing this? I think it will get stuck in DAG combines.

Otherwise the idea sounds good. Have you considered doing it for GlobalISel, or in a way that GlobalISel would also benefit?

I haven't found any stuck in DAG combine local and it looks all of the premerge test is also passed:
https://reviews.llvm.org/harbormaster/build/291788/

Which combination do you think will cause the dead loop in DAG combine with this patch? I can try to create a test for that.

Yes - the llvm regression tests are very much a subset of all the patterns that can come up in practice. I'm not sure what this is running into. Perhaps try and bootstrap or compiling the llvm-test-suite?

[bcl5980]

In D136158#3880624, @dmgreen wrote:

In D136158#3878631, @bcl5980 wrote:

In D136158#3878602, @dmgreen wrote:

Have you tried testing this? I think it will get stuck in DAG combines.

Otherwise the idea sounds good. Have you considered doing it for GlobalISel, or in a way that GlobalISel would also benefit?

I haven't found any stuck in DAG combine local and it looks all of the premerge test is also passed:
https://reviews.llvm.org/harbormaster/build/291788/

Which combination do you think will cause the dead loop in DAG combine with this patch? I can try to create a test for that.

Yes - the llvm regression tests are very much a subset of all the patterns that can come up in practice. I'm not sure what this is running into. Perhaps try and bootstrap or compiling the llvm-test-suite?

Or can we do it on Machine IR(AArch64MIPeepholeOpt) to avoid dead loop and also GISel can get benefit?

[bcl5980]

dead loop happen in this file

combine.c.ll1 MBDownload

[bcl5980]

Z can't be constant to avoid dead loop

[bcl5980]

test pass:
llvm self-build
llvm test-suite
spec2017
sqlite3
python3.11

[dmgreen]

Thanks for finding the problem. This LGTM now.

Or can we do it on Machine IR(AArch64MIPeepholeOpt) to avoid dead loop and also GISel can get benefit?

I was wondering what to do about optimizations that can be shared between GISel and SDAG. AArch64MIPeepholeOpt often feels like the wrong place (MachineCombiner might be better for some optimizations like this one, but that currently doesn't scale very well). We don't really want to write them in the wrong place if they would be better as combines during ISel. For now at least this patch looks good.

[glandium]

Building Firefox with LTO for any aarch64 target (linux, macos, windows) enters in some sort of infinite loop since this change (linking never ends).

[bcl5980]

In D136158#3890955, @glandium wrote:

Building Firefox with LTO for any aarch64 target (linux, macos, windows) enters in some sort of infinite loop since this change (linking never ends).

reverted by rGdfb16bd5526b

[glandium]

Reproducer: https://drive.google.com/file/d/1n7KPlKCydCgSdtnp6990Zn5_Yjbbss8L/view?usp=sharing (command in testcase/cmd)

[bcl5980]

In D136158#3891134, @glandium wrote:

Reproducer: https://drive.google.com/file/d/1n7KPlKCydCgSdtnp6990Zn5_Yjbbss8L/view?usp=sharing (command in testcase/cmd)

Is it link time infinite loop? Somehow this change doesn't have any effect on link stage.
Can this reproducer work after I revert the code?

[mstorsjo]

Thanks for reverting - I've also bisected down hangs to the same commit. For my case, it can be reproduced with https://martin.st/temp/aarch64-hang.c, with clang -target aarch64-linux-gnu -c -O2 aarch64-hang.c.

[mstorsjo]

In D136158#3891161, @bcl5980 wrote:

In D136158#3891134, @glandium wrote:

Reproducer: https://drive.google.com/file/d/1n7KPlKCydCgSdtnp6990Zn5_Yjbbss8L/view?usp=sharing (command in testcase/cmd)

Is it link time infinite loop? Somehow this change doesn't have any effect on link stage.

In his case, he's doing LTO builds, so the compilation/codegen happens at link time. But my testcase might be smaller and easier to run.

[bcl5980]

In D136158#3891177, @mstorsjo wrote:

In D136158#3891161, @bcl5980 wrote:

In D136158#3891134, @glandium wrote:

Reproducer: https://drive.google.com/file/d/1n7KPlKCydCgSdtnp6990Zn5_Yjbbss8L/view?usp=sharing (command in testcase/cmd)

Is it link time infinite loop? Somehow this change doesn't have any effect on link stage.

In his case, he's doing LTO builds, so the compilation/codegen happens at link time. But my testcase might be smaller and easier to run.

Thanks for the case.

[bcl5980]

Fix infinite loop caused by Z is one-use shift C also.
Add some more tests.

[bcl5980]

@dmgreen, should I send a new patch review for the fix, or still in this page?

[thakis]

We also bisected hangs due to this (arm64 on all of android fuchsia ios mac). It looks like you have a bunch of reproducers already, but there's one more at https://bugs.chromium.org/p/chromium/issues/detail?id=1379286#c7 if you're interested (it's just a file from boringssl – we saw this hang on many files. I'd guess, but haven't checked, that doing a bootstrap build of clang on an arm64 box would also trigger this.)

Thanks for the revert!

[bcl5980]

Thanks for all of your tests. And current patch can pass all the tests now.

[dmgreen]

I see. Thanks for the update. Providing you have ran a bootstrap, this LGTM.

[bcl5980]

In D136158#3895618, @dmgreen wrote:

I see. Thanks for the update. Providing you have ran a bootstrap, this LGTM.

I'm sorry but I don't know what is the bootstrap. Can you give me some suggestion for that? Where can I get a opensource bootstrap and build?

[dmgreen]

Ah - it just means compiler the compiler with itself (on AArch64) and make sure it still passes all the tests. There are some details, including some handy cmake options in https://llvm.org/docs/AdvancedBuilds.html

[bcl5980]

In D136158#3895933, @dmgreen wrote:

Ah - it just means compiler the compiler with itself (on AArch64) and make sure it still passes all the tests. There are some details, including some handy cmake options in https://llvm.org/docs/AdvancedBuilds.html

Thanks for the explaination. I have already build and test llvm itself on MacbookAir with M2 processor.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Elt0->getOpcode() == ISD::EXTRACT_VECTOR_ELT &&
     SDValue Ext = DAG.getNode(ISD::ANY_EXTEND, DL, ExtVT, VecToExtend);
     return DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, MVT::v2i32, Ext,
                        SubvectorIdx);
   }
 
   return SDValue();
 }
 
+// ((X >> C) - Y) + Z --> (Z - Y) + (X >> C)
+static SDValue performAddCombineSubShift(SDNode *N, SDValue SUB, SDValue Z,
+                                         SelectionDAG &DAG) {
+  auto IsOneUseShiftC = [&](SDValue Shift) {
+    if (!Shift.hasOneUse())
+      return false;
+
+    // TODO: support SRL and SRA also
+    if (Shift.getOpcode() != ISD::SHL)
+      return false;
+
+    if (!isa<ConstantSDNode>(Shift.getOperand(1)))
+      return false;
+    return true;
+  };
+
+  // DAGCombiner will revert the combination when Z is constant cause
+  // dead loop. So don't enable the combination when Z is constant.
+  // If Z is one use shift C, we also can't do the optimization.
+  // It will falling to self infinite loop.
+  if (isa<ConstantSDNode>(Z) || IsOneUseShiftC(Z))
+    return SDValue();
+
+  if (SUB.getOpcode() != ISD::SUB || !SUB.hasOneUse())
+    return SDValue();
+
+  SDValue Shift = SUB.getOperand(0);
+  if (!IsOneUseShiftC(Shift))
+    return SDValue();
+
+  SDLoc DL(N);
+  EVT VT = N->getValueType(0);
+
+  SDValue Y = SUB.getOperand(1);
+  SDValue NewSub = DAG.getNode(ISD::SUB, DL, VT, Z, Y);
+  return DAG.getNode(ISD::ADD, DL, VT, NewSub, Shift);
+}
+
 static SDValue performAddCombineForShiftedOperands(SDNode *N,
                                                    SelectionDAG &DAG) {
   // NOTE: Swapping LHS and RHS is not done for SUB, since SUB is not
   // commutative.
   if (N->getOpcode() != ISD::ADD)
     return SDValue();
 
   // Bail out when value type is not one of {i32, i64}, since AArch64 ADD with
   // shifted register is only available for i32 and i64.
   EVT VT = N->getValueType(0);
   if (VT != MVT::i32 && VT != MVT::i64)
     return SDValue();
 
   SDLoc DL(N);
   SDValue LHS = N->getOperand(0);
   SDValue RHS = N->getOperand(1);
 
+  if (SDValue Val = performAddCombineSubShift(N, LHS, RHS, DAG))
+    return Val;
+  if (SDValue Val = performAddCombineSubShift(N, RHS, LHS, DAG))
+    return Val;
+
   uint64_t LHSImm = 0, RHSImm = 0;
   // If both operand are shifted by imm and shift amount is not greater than 4
   // for one operand, swap LHS and RHS to put operand with smaller shift amount
   // on RHS.
   //
   // On many AArch64 processors (Cortex A78, Neoverse N1/N2/V1, etc), ADD with
   // LSL shift (shift <= 4) has smaller latency and larger throughput than ADD
   // with LSL (shift > 4). For the rest of processors, this is no-op for

llvm/test/CodeGen/AArch64/addsub.ll

if.then:                                          ; preds = %lor.lhs.false, %entry
   %tobool.not = icmp eq i32 %1, 0
   %cond = select i1 %tobool.not, i32* null, i32* @input_pool
   %call = tail call i32 bitcast (i32 (...)* @crng_reseed to i32 (i32*)*)(i32* noundef %cond)
   br label %if.end
 if.end:                                           ; preds = %if.then, %lor.lhs.false
   ret i32 undef
 }
 
+; ((X >> C) - Y) + Z --> (Z - Y) + (X >> C)
 define i32 @commute_subop0(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: commute_subop0:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    lsl w8, w0, #3
-; CHECK-NEXT:    sub w8, w8, w1
-; CHECK-NEXT:    add w0, w8, w2
+; CHECK-NEXT:    sub w8, w2, w1
+; CHECK-NEXT:    add w0, w8, w0, lsl #3
 ; CHECK-NEXT:    ret
   %shl = shl i32 %x, 3
   %sub = sub i32 %shl, %y
   %add = add i32 %sub, %z
   ret i32 %add
 }
 
+; ((X << C) - Y) + Z --> (Z - Y) + (X << C)
+define i32 @commute_subop0_lshr(i32 %x, i32 %y, i32 %z) {
+; CHECK-LABEL: commute_subop0_lshr:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsr w8, w0, #3
+; CHECK-NEXT:    sub w8, w8, w1
+; CHECK-NEXT:    add w0, w8, w2
+; CHECK-NEXT:    ret
+  %lshr = lshr i32 %x, 3
+  %sub = sub i32 %lshr, %y
+  %add = add i32 %sub, %z
+  ret i32 %add
+}
+
+; ((X << C) - Y) + Z --> (Z - Y) + (X << C)
+define i32 @commute_subop0_ashr(i32 %x, i32 %y, i32 %z) {
+; CHECK-LABEL: commute_subop0_ashr:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    asr w8, w0, #3
+; CHECK-NEXT:    sub w8, w8, w1
+; CHECK-NEXT:    add w0, w8, w2
+; CHECK-NEXT:    ret
+  %ashr = ashr i32 %x, 3
+  %sub = sub i32 %ashr, %y
+  %add = add i32 %sub, %z
+  ret i32 %add
+}
+
+; Z + ((X >> C) - Y) --> (Z - Y) + (X >> C)
 define i32 @commute_subop0_cadd(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: commute_subop0_cadd:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    lsl w8, w0, #3
-; CHECK-NEXT:    sub w8, w8, w1
-; CHECK-NEXT:    add w0, w2, w8
+; CHECK-NEXT:    sub w8, w2, w1
+; CHECK-NEXT:    add w0, w8, w0, lsl #3
 ; CHECK-NEXT:    ret
   %shl = shl i32 %x, 3
   %sub = sub i32 %shl, %y
   %add = add i32 %z, %sub
   ret i32 %add
 }
 
+; Y + ((X >> C) - X) --> (Y - X) + (X >> C)
 define i32 @commute_subop0_mul(i32 %x, i32 %y) {
 ; CHECK-LABEL: commute_subop0_mul:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    lsl w8, w0, #3
-; CHECK-NEXT:    sub w8, w8, w0
-; CHECK-NEXT:    add w0, w8, w1
+; CHECK-NEXT:    sub w8, w1, w0
+; CHECK-NEXT:    add w0, w8, w0, lsl #3
 ; CHECK-NEXT:    ret
   %mul = mul i32 %x, 7
   %add = add i32 %mul, %y
   ret i32 %add
 }
+
+; negative case for ((X >> C) - Y) + Z --> (Z - Y) + (X >> C)
+; Y can't be constant to avoid dead loop
+define i32 @commute_subop0_zconst(i32 %x, i32 %y) {
+; CHECK-LABEL: commute_subop0_zconst:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsl w8, w0, #3
+; CHECK-NEXT:    sub w8, w8, w1
+; CHECK-NEXT:    add w0, w8, #1
+; CHECK-NEXT:    ret
+  %shl = shl i32 %x, 3
+  %sub = sub i32 %shl, %y
+  %add = add i32 %sub, 1
+  ret i32 %add
+}
+
+; negative case for ((X >> C) - Y) + Z --> (Z - Y) + (X >> C)
+; Y can't be shift C also to avoid dead loop
+define i32 @commute_subop0_zshiftc_oneuse(i32 %x, i32 %y, i32 %z) {
+; CHECK-LABEL: commute_subop0_zshiftc_oneuse:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsl w8, w0, #3
+; CHECK-NEXT:    sub w8, w8, w1
+; CHECK-NEXT:    add w0, w8, w2, lsl #2
+; CHECK-NEXT:    ret
+  %xshl = shl i32 %x, 3
+  %sub = sub i32 %xshl, %y
+  %zshl = shl i32 %z, 2
+  %add = add i32 %sub, %zshl
+  ret i32 %add
+}
+
+define i32 @commute_subop0_zshiftc(i32 %x, i32 %y, i32 %z) {
+; CHECK-LABEL: commute_subop0_zshiftc:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsl w8, w2, #2
+; CHECK-NEXT:    sub w9, w8, w1
+; CHECK-NEXT:    add w9, w9, w0, lsl #3
+; CHECK-NEXT:    eor w0, w8, w9
+; CHECK-NEXT:    ret
+  %xshl = shl i32 %x, 3
+  %sub = sub i32 %xshl, %y
+  %zshl = shl i32 %z, 2
+  %add = add i32 %sub, %zshl
+  %r = xor i32 %zshl, %add
+  ret i32 %r
+}