Reapply [AArch64] fold subs ugt/ult to ands when the second operand is mask/pow2

[bcl5980]

Summary

Origianl patch made a mistake that ugt is reverse cc should be ule.
And ule < C will be generalize to ult < C + 1. So the new patch add support for ult < Pow2 case.

https://alive2.llvm.org/ce/z/naBw5A

[samtebbs]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19843

Do you know if CCIndex and CmpIndex always have consistent values? If so it would be cleaner and faster to not use a loop and instead initialise the Ops vector with the correct values in the correct places.

[bcl5980]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19843

Based on current code, CCIndex should be always 2, CmpIndex should be always 3. I'm not sure it will change or not in the future. But I can try to use the cleaner way and add a assert here.

[samtebbs]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19843

That's a good idea, thanks. I think accounting for how things are at the moment and tripping an assertion if it changes is a good approach.

llvm/test/CodeGen/AArch64/andcompare.ll

2406

Looks like the labels in the check statements are mixed up.

[bcl5980]

llvm/test/CodeGen/AArch64/andcompare.ll

2406

I'm sorry but I haven't understand what you said. What labels are mixed up? One of the result is SDAG and the other is GISel.

[samtebbs]

llvm/test/CodeGen/AArch64/andcompare.ll

2406

The IR function is called cmp_to_ands1 but the check statement is looking for cmp_to_subs1, that's why the test is failing: https://reviews.llvm.org/harbormaster/unit/view/5713977/

It's the same for the other IR functions below.

[bcl5980]

llvm/test/CodeGen/AArch64/andcompare.ll

2406

Ah, thanks. I will update it soon.

[bcl5980]

address comments.

[samtebbs]

Thanks for this, it's looking great, with one small suggestion.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19846

I'd reword the assertion message to be more like "Expected CCIndex to be 2 and CmpIndex to be 3", to make it more like existing assertions. The part about updating it would be better in a comment, in my opinion.

[bcl5980]

update comments.

[chapuni]

Seems it broke aarch64 bootstrap. Investigating.

[chapuni]

It miscompiled CodeGen/TargetLoweringObjectFileImpl.cpp

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19817

I think the latter is wrong.
eg. X=11 C=7 M=3

[bcl5980]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19817

Thanks for the case.
I will send a quick fix for that.

[chapuni]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19817

I reconfirmed the new expr is correct.
(not tested yet for bootstrap, though)

Could you enhance to add also them?

(X & C) >=u Pow2
(X & C) <=u Mask

[chapuni]

btw, was it not handled by instcombine?

[bcl5980]

In D141829#4061077, @chapuni wrote:

btw, was it not handled by instcombine?

Not handled by instcombine is because AArch64 has tst instruction that instcombine need two instruction to do that.

[bcl5980]

update two more tests for:
(X & C) >=u Pow2
(X & C) <=u Mask

[bcl5980]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19817

We needn't enhance the case in the code. DAGCombiner will generalize these to current two patterns. Check the test:
(X & C) >=u Pow2 in cmp_to_ands6
(X & C) <=u Mask in cmp_to_ands4

[chapuni]

Will this affect other tests? I met another.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19817

Thanks, I understood.

[bcl5980]

In D141829#4061257, @chapuni wrote:

Will this affect other tests? I met another.

I don't know actually. The motivation of this change is fix https://github.com/llvm/llvm-project/issues/59598.
Maybe @samtebbs can answer this question.

[bcl5980]

Fix tests failure in CodeGen/AArch64/and-mask-removal.ll

[chapuni]

Works fine from my side, thanks!

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

case AArch64CC::NV:
     return true;
   case AArch64CC::Invalid:
     break;
   }
 
   return false;
 }
 
+// (X & C) >u Mask --> (X & (C & (~Mask)) != 0
+// (X & C) <u Pow2 --> (X & (C & ~(Pow2-1)) == 0
+static SDValue performSubsToAndsCombine(SDNode *N, SDNode *SubsNode,
+                                        SDNode *AndNode, SelectionDAG &DAG,
+                                        unsigned CCIndex, unsigned CmpIndex,
+                                        unsigned CC) {
+  ConstantSDNode *SubsC = dyn_cast<ConstantSDNode>(SubsNode->getOperand(1));
+  if (!SubsC)
+    return SDValue();
+
+  APInt SubsAP = SubsC->getAPIntValue();
+  if (CC == AArch64CC::HI) {
+    if (!SubsAP.isMask())
+      return SDValue();
+  } else if (CC == AArch64CC::LO) {
+    if (!SubsAP.isPowerOf2())
+      return SDValue();
+  } else
+    return SDValue();
+
+  ConstantSDNode *AndC = dyn_cast<ConstantSDNode>(AndNode->getOperand(1));
+  if (!AndC)
+    return SDValue();
+
+  APInt MaskAP = CC == AArch64CC::HI ? SubsAP : (SubsAP - 1);
+
+  SDLoc DL(N);
+  APInt AndSMask = (~MaskAP) & AndC->getAPIntValue();
+  SDValue ANDS = DAG.getNode(
+      AArch64ISD::ANDS, DL, SubsNode->getVTList(), AndNode->getOperand(0),
+      DAG.getConstant(AndSMask, DL, SubsC->getValueType(0)));
+  SDValue AArch64_CC =
+      DAG.getConstant(CC == AArch64CC::HI ? AArch64CC::NE : AArch64CC::EQ, DL,
+                      N->getOperand(CCIndex)->getValueType(0));
+
+  // For now, only performCSELCombine and performBRCONDCombine call this
+  // function. And both of them pass 2 for CCIndex, 3 for CmpIndex with 4
+  // operands. So just init the ops direct to simplify the code. If we have some
+  // other case with different CCIndex, CmpIndex, we need to use for loop to
+  // rewrite the code here.
+  // TODO: Do we need to assert number of operand is 4 here?
+  assert((CCIndex == 2 && CmpIndex == 3) &&
+         "Expected CCIndex to be 2 and CmpIndex to be 3.");
+  SDValue Ops[] = {N->getOperand(0), N->getOperand(1), AArch64_CC,
+                   ANDS.getValue(1)};
+  return DAG.getNode(N->getOpcode(), N, N->getVTList(), Ops);
+}
+
 static
 SDValue performCONDCombine(SDNode *N,
                            TargetLowering::DAGCombinerInfo &DCI,
                            SelectionDAG &DAG, unsigned CCIndex,
                            unsigned CmpIndex) {
   unsigned CC = cast<ConstantSDNode>(N->getOperand(CCIndex))->getSExtValue();
   SDNode *SubsNode = N->getOperand(CmpIndex).getNode();
   unsigned CondOpcode = SubsNode->getOpcode();
 
   if (CondOpcode != AArch64ISD::SUBS || SubsNode->hasAnyUseOfValue(0))
     return SDValue();
 
   // There is a SUBS feeding this condition. Is it fed by a mask we can
   // use?
 
   SDNode *AndNode = SubsNode->getOperand(0).getNode();
   unsigned MaskBits = 0;
 
   if (AndNode->getOpcode() != ISD::AND)
     return SDValue();
 
+  if (SDValue Val = performSubsToAndsCombine(N, SubsNode, AndNode, DAG, CCIndex,
+                                             CmpIndex, CC))
+    return Val;
+
   if (ConstantSDNode *CN = dyn_cast<ConstantSDNode>(AndNode->getOperand(1))) {
     uint32_t CNV = CN->getZExtValue();
     if (CNV == 255)
       MaskBits = 8;
     else if (CNV == 65535)
       MaskBits = 16;
   }
 

llvm/test/CodeGen/AArch64/andcompare.ll

 ; GISEL-NEXT:    ret
 entry:
   %c0 = icmp sge i32 %s0, %s1
   %c1 = icmp sge i32 %s2, %s3
   %a = and i1 %c0, %c1
   %z = zext i1 %a to i32
   ret i32 %z
 }
+
+define i32 @cmp_to_ands1(i32 %num) {
+; SDISEL-LABEL: cmp_to_ands1:
+; SDISEL:       // %bb.0:
+; SDISEL-NEXT:    and w8, w0, #0xff
+; SDISEL-NEXT:    tst w0, #0xfe
+; SDISEL-NEXT:    csel w0, w8, wzr, ne
+; SDISEL-NEXT:    ret
+;
+; GISEL-LABEL: cmp_to_ands1:
+; GISEL:       // %bb.0:
+; GISEL-NEXT:    and w8, w0, #0xff
+; GISEL-NEXT:    cmp w8, #1
+; GISEL-NEXT:    csel w0, w8, wzr, hi
+; GISEL-NEXT:    ret
+  %and = and i32 %num, 255
+  %cmp = icmp ugt i32 %and, 1
+  %r = select i1 %cmp, i32 %and, i32 0
+  ret i32 %r
+}
+
+define i32 @cmp_to_ands2(i32 %num) {
+; SDISEL-LABEL: cmp_to_ands2:
+; SDISEL:       // %bb.0:
+; SDISEL-NEXT:    and w8, w0, #0xfe
+; SDISEL-NEXT:    tst w0, #0xc0
+; SDISEL-NEXT:    csel w0, w8, wzr, ne
+; SDISEL-NEXT:    ret
+;
+; GISEL-LABEL: cmp_to_ands2:
+; GISEL:       // %bb.0:
+; GISEL-NEXT:    and w8, w0, #0xfe
+; GISEL-NEXT:    cmp w8, #63
+; GISEL-NEXT:    csel w0, w8, wzr, hi
+; GISEL-NEXT:    ret
+  %and = and i32 %num, 254
+  %cmp = icmp ugt i32 %and, 63
+  %r = select i1 %cmp, i32 %and, i32 0
+  ret i32 %r
+}
+
+define i32 @cmp_to_ands3(i32 %num, i32 %a) {
+; SDISEL-LABEL: cmp_to_ands3:
+; SDISEL:       // %bb.0:
+; SDISEL-NEXT:    tst w0, #0x10
+; SDISEL-NEXT:    csel w0, w1, wzr, ne
+; SDISEL-NEXT:    ret
+;
+; GISEL-LABEL: cmp_to_ands3:
+; GISEL:       // %bb.0:
+; GISEL-NEXT:    mov w8, #23
+; GISEL-NEXT:    and w8, w0, w8
+; GISEL-NEXT:    cmp w8, #7
+; GISEL-NEXT:    csel w0, w1, wzr, hi
+; GISEL-NEXT:    ret
+  %and = and i32 %num, 23
+  %cmp = icmp ugt i32 %and, 7
+  %r = select i1 %cmp, i32 %a, i32 0
+  ret i32 %r
+}
+
+define i32 @cmp_to_ands4(i32 %num, i32 %a) {
+; SDISEL-LABEL: cmp_to_ands4:
+; SDISEL:       // %bb.0:
+; SDISEL-NEXT:    and w8, w0, #0x30
+; SDISEL-NEXT:    tst w0, #0x20
+; SDISEL-NEXT:    csel w0, w8, w1, eq
+; SDISEL-NEXT:    ret
+;
+; GISEL-LABEL: cmp_to_ands4:
+; GISEL:       // %bb.0:
+; GISEL-NEXT:    and w8, w0, #0x30
+; GISEL-NEXT:    cmp w8, #31
+; GISEL-NEXT:    csel w0, w8, w1, ls
+; GISEL-NEXT:    ret
+  %and = and i32 %num, 48
+  %cmp = icmp ule i32 %and, 31
+  %r = select i1 %cmp, i32 %and, i32 %a
+  ret i32 %r
+}
+
+define i32 @cmp_to_ands5(i32 %num, i32 %a) {
+; SDISEL-LABEL: cmp_to_ands5:
+; SDISEL:       // %bb.0:
+; SDISEL-NEXT:    and w8, w0, #0xf8
+; SDISEL-NEXT:    tst w0, #0xc0
+; SDISEL-NEXT:    csel w0, w8, w1, eq
+; SDISEL-NEXT:    ret
+;
+; GISEL-LABEL: cmp_to_ands5:
+; GISEL:       // %bb.0:
+; GISEL-NEXT:    and w8, w0, #0xf8
+; GISEL-NEXT:    cmp w8, #64
+; GISEL-NEXT:    csel w0, w8, w1, lo
+; GISEL-NEXT:    ret
+  %and = and i32 %num, 248
+  %cmp = icmp ult i32 %and, 64
+  %r = select i1 %cmp, i32 %and, i32 %a
+  ret i32 %r
+}
+
 ;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
 ; CHECK: {{.*}}