[AArch64] Generate AND in place of CSEL for Table Based CTTZ lowering in -O3

[rahular-rrlogic]

Summary

This patch implements a for a target specific optimization on top of https://reviews.llvm.org/D113291 that replaces the cmp and csel with an and

Related submission and comments: https://reviews.llvm.org/D120462

Original issue: https://github.com/llvm/llvm-project/issues/45779

[djtodoro]

Thanks for this! Please add test.

[dmgreen]

Can you add some tests?

As far as I understand this should be testing for select (icmp eq X, 0), 0, cttz X and converting it to and(cttz X, #bw-1). If so there are more conditions that need to be checked.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17685

As far as I can tell this is checking that the condition code is 0? It would be better to check that it is equal to AArch64CC::EQ.

17687

Variables in llvm start with capital letters. We should make sure that i64 work too, it needs a different constant (there should be tests too).

[rahular-rrlogic]

In D123782#3451501, @dmgreen wrote:

Can you add some tests?

As far as I understand this should be testing for select (icmp eq X, 0), 0, cttz X and converting it to and(cttz X, #bw-1). If so there are more conditions that need to be checked.

What other conditions do you think should be added other than checking for 0 and cttz?

In D123782#3451500, @djtodoro wrote:

Thanks for this! Please add test.

I will add test with the changes

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17685

No, this is a mistake. I was intending to check if the operands are 0 and cttz. I will change that. Is a check for the condition being EQ really required, though?

[rahular-rrlogic]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17687

I will do the capitalization and add support for i64

[craig.topper]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17685

There's a function called isNullConstant that can be used here.

[rahular-rrlogic]

The condition for performing this optimization was modified to checking if the RHS of SUBS is 0 and if one of the values for the CSEL is a CTTZ.

Support for i64 was added.

Tests were added.

[djtodoro]

llvm/test/CodeGen/AArch64/table-based-cttz.ll

1 ↗

(On Diff #423342)

Can we add a top-level comment here, describing what we are testing?

[rahular-rrlogic]

Updated test with description and made formatting changes.

[rahular-rrlogic]

Modified test to add details and made formatting changes

[gsocshubham]

@momchil.velikov - Review request.

[dmgreen]

As far as I understand this should be testing for select (icmp eq X, 0), 0, cttz X and converting it to and(cttz X, #bw-1). If so there are more conditions that need to be checked.

What other conditions do you think should be added other than checking for 0 and cttz?

It looks like this checks for the select/CSEL, and the icmp/SUBS with a 0, but not the "eq" yet, and not the "0" in the select/csel. It also doesn't check that the X in the icmp and the X in the cttz are the same?

It would be worth adding test cases for the cases it should not fold.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17685

It should probably be something like this, if the value is a condition code.

AArch64CC CC = N->getConstantOperandVal(2);
if (CC == AArch64CC::EQ)

It sounds good to use it for the zero of the sub though.

17686–17687

Why is it hard-coding 31, as opposed to checking the size of the type? Why can't we get here with an i64?

llvm/test/CodeGen/AArch64/table-based-cttz.ll

27 ↗

(On Diff #423552)

Can you add an i64 case without the truncate down to an i32?

[rahular-rrlogic]

Removed hard coded constants and replaced that with bitwidth, added full support for i64 and added more conditions to match against the intended pattern. Modified test to include cases in which the optimization does not take place.

[rahular-rrlogic]

In D123782#3467076, @dmgreen wrote:

As far as I understand this should be testing for select (icmp eq X, 0), 0, cttz X and converting it to and(cttz X, #bw-1). If so there are more conditions that need to be checked.

What other conditions do you think should be added other than checking for 0 and cttz?

It looks like this checks for the select/CSEL, and the icmp/SUBS with a 0, but not the "eq" yet, and not the "0" in the select/csel. It also doesn't check that the X in the icmp and the X in the cttz are the same?

Is the last check really needed? Both icmp and cttz use the value from the same register in the IR itself.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17686–17687

I had misunderstood this part. Fixed now, thank you.

llvm/test/CodeGen/AArch64/table-based-cttz.ll

27 ↗

(On Diff #423552)

Added.

[craig.topper]

In D123782#3468601, @rahular-rrlogic wrote:

In D123782#3467076, @dmgreen wrote:

As far as I understand this should be testing for select (icmp eq X, 0), 0, cttz X and converting it to and(cttz X, #bw-1). If so there are more conditions that need to be checked.

What other conditions do you think should be added other than checking for 0 and cttz?

It looks like this checks for the select/CSEL, and the icmp/SUBS with a 0, but not the "eq" yet, and not the "0" in the select/csel. It also doesn't check that the X in the icmp and the X in the cttz are the same?

Is the last check really needed? Both icmp and cttz use the value from the same register in the IR itself.

Yes it's needed. You have to make you don't optimize select (icmp eq X, 0), 0, cttz Y if X and Y aren't the same.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17684

You can get the bitwidth a lot quicker with N1.getValueSizeInBits()

[craig.topper]

You should probably also handle select (icmp ne X, 0), (cttz X), 0.

[rahular-rrlogic]

Added condition to check if X in icmp and cttz are the same and added code to handle select (icmp ne X, 0), (cttz X), 0. Updated test for the same.

[craig.topper]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17690

else should be on the same line as the previous closing brace

17702

This line is longer than 80 columns.

17703

Space after if

[craig.topper]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17699

Technically, if you look through a truncate you need to know the truncate didn't drop any bits of the CTTZ result. But maybe AArch64ISD::SUBS and CSEL can only be created after type legalization so the only possible types are i32 and i64? I'm not an AArch64 expert so I don't know.

[dmgreen]

Can you pull the cttz code into a new function, so that it is a little more separate from the other code in performCSELCombine.

Running clang-format on the patch can also be good to remove all the formatting issues. There is a script to help in clang/tools/clang-format/clang-format-diff.py.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17699

Yeah I believe SUBS and CSEL will only be generated for legal types. Perhaps it is worth adding an assert just to be safe.

[rahular-rrlogic]

Created a new function to remove the folding itself from performCSELCombine(). Added an assert to check for legal types.

[rahular-rrlogic]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17699

The type legalization does occur in this case which I checked via the debug output. I have added an assert anyway as @dmgreen suggested.

[craig.topper]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17681–17687

You can get rid of the IsEQOrNE variable by adding

else
  return SDValue();

17681–17687

Remove the IsEQOrNE using the suggestion above, then invert this to make an early out.

17696

if (SDValue Folded = foldCTTZ(N, DAG))
  return Folded;

[rahular-rrlogic]

Modified the conditions to add an early out in the case of a non match.

[dmgreen]

Can you make sure the diff is against current main branch, not against the last version of the patch.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17632–17674

foldCTTZ -> foldCSELofCTTZ

17642

LLVM tends to prefer early exists, to reduce the indent level.

Instead of doing

if (X) {
  if (Y) {
    DoThing()
  }
}

it is preferred to do

if (!X)
  return SDValue();
if (!Y)
  return SDValue();
DoThing()

In this case, the AND variable can be removed, for example.

[rahular-rrlogic]

Renamed the function that performs the folding of CSEL into AND for CTTZ. Added more early outs. Fixed diff.

[dmgreen]

Thanks for the updates, I think this is looking good now.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17636

I think AND is now unused.

llvm/test/CodeGen/AArch64/table-based-cttz.ll

1 ↗

(On Diff #426350)

Use -mtriple=aarch64. Otherwise I think this will run differently on non-aarch64 native machines.
Can you also run update_llc_test_checks on the file. It fills in all the check lines in a consistent way.
And maybe name the test file cttz-and.ll, or something like it. It isn't directly related with the table based cttz's.

[rahular-rrlogic]

Changed test file name and generated assertions using update_llc_test_checks.py.

[craig.topper]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17655

SUB.getValue(1).getOperand(1) can be shortened to SUBS.getOperand(1). The Result number in isn't used by SDValue::getOperand.

[dmgreen]

Other than the point @craig.topper mentioned, this Looks OK to me.

[rahular-rrlogic]

Removed the unneeded call to a getValue as pointed out.

[rahular-rrlogic]

In D123782#3493238, @dmgreen wrote:

Other than the point @craig.topper mentioned, this Looks OK to me.

I have marked this as a child revision of https://reviews.llvm.org/D113291 but this patch should work independently of that. Should I remove that as a parent revision? Please let me know what you think (@djtodoro too). Additionally, how should I go about committing this since I don't have commit access?

[dmgreen]

I can commit this for you if you don't have commit access yet. I would just need a "name <[e-mail]" for the attribution.

[rahular-rrlogic]

In D123782#3498708, @dmgreen wrote:

I can commit this for you if you don't have commit access yet. I would just need a "name <[e-mail]" for the attribution.

You can use "Rahul Anand R <[e-mail]>" for the attribution.

[chapuni]

This causes miscompilation in llvm::lowertypetests::BitSetBuilder::build() in llvm/lib/Transforms/IPO/LowerTypeTests.cpp.
LLVM-Unit::IPOTests catches this.
BSI.AlignLog2 is masked by "and #0x1F".

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17665

CTTZ may be ISD::TRUNCATE.

[dmgreen]

Oh yeah, I missed that testcase. I'll revert for now. Thanks for the report.

[rahular-rrlogic]

@chapuni @dmgreen How do I trigger this unit test failure? When I run check-llvm-unit all the tests pass for me. I don't understand what I am missing here.

[chapuni]

In D123782#3504135, @rahular-rrlogic wrote:

@chapuni @dmgreen How do I trigger this unit test failure? When I run check-llvm-unit all the tests pass for me. I don't understand what I am missing here.

I did; Bootstrap stage2 (on aws c6g) on an internal builder at work.

$ ninja -C build/1 install
$ CC=/path/to/install/1/bin/clang CXX=/path/to/install/1/bin/clang++ cmake ... -DCMAKE_BUILD_TYPE=Release -B build/2
$ ninja -C build/2 IPOTests && path/to/IPOTests

[dmgreen]

Yeah I think the issue is that @cttztrunc should be doing and 0x3f, not and 0x1f.

We should really have caught that, but I was relying too much on a verification script that didn't catch it due to the way rbit is specified. Which reminds me that we should really be supporting ctlz too for the same transform, but perhaps that's best left to a followup.

[rahular-rrlogic]

Made bitwidth take its value depending on whether there is a truncation present

[dmgreen]

Thanks for the update. Have you tried a bootstrap to make sure it passes now?

[rahular-rrlogic]

In D123782#3508832, @dmgreen wrote:

Thanks for the update. Have you tried a bootstrap to make sure it passes now?

I never had any test failures even in the previous revision. How do I include all tests?

[dmgreen]

In D123782#3514938, @rahular-rrlogic wrote:

In D123782#3508832, @dmgreen wrote:

Thanks for the update. Have you tried a bootstrap to make sure it passes now?

I never had any test failures even in the previous revision. How do I include all tests?

Do you mean just ninja check-all? Those tests will only include the tests people have deemed worth adding in the past - they cannot cover all the possible cases and combinations of things that can come up. To run more thorough testing you need to start compiling code with the new compiler and making sure the results are correct. The issue that came up I think was when the compiler compiled itself - a bootstrap. We should make sure that doesn't still happen in the same way.

[rahular-rrlogic]

In D123782#3515031, @dmgreen wrote:

In D123782#3514938, @rahular-rrlogic wrote:

In D123782#3508832, @dmgreen wrote:

Thanks for the update. Have you tried a bootstrap to make sure it passes now?

I never had any test failures even in the previous revision. How do I include all tests?

Do you mean just ninja check-all? Those tests will only include the tests people have deemed worth adding in the past - they cannot cover all the possible cases and combinations of things that can come up. To run more thorough testing you need to start compiling code with the new compiler and making sure the results are correct. The issue that came up I think was when the compiler compiled itself - a bootstrap. We should make sure that doesn't still happen in the same way.

I tried a bootstrap and it seems to be building fine

[dmgreen]

Thanks for checking. Lets give this another go then. LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

static SDValue performBRCONDCombine(SDNode *N,
   else
     BR = DAG.getNode(AArch64ISD::CBNZ, SDLoc(N), MVT::Other, Chain, LHS, Dest);
 
   // Do not add new nodes to DAG combiner worklist.
   DCI.CombineTo(N, BR, false);
 
   return SDValue();
 }
 
+static SDValue foldCSELofCTTZ(SDNode *N, SelectionDAG &DAG) {
+  unsigned CC = N->getConstantOperandVal(2);
+  SDValue SUBS = N->getOperand(3);
+  SDValue Zero, CTTZ;
+
+  if (CC == AArch64CC::EQ && SUBS.getOpcode() == AArch64ISD::SUBS) {
+    Zero = N->getOperand(0);
+    CTTZ = N->getOperand(1);
+  } else if (CC == AArch64CC::NE && SUBS.getOpcode() == AArch64ISD::SUBS) {
+    Zero = N->getOperand(1);
+    CTTZ = N->getOperand(0);
+  } else
+    return SDValue();
+
+  if ((CTTZ.getOpcode() != ISD::CTTZ && CTTZ.getOpcode() != ISD::TRUNCATE) ||
+      (CTTZ.getOpcode() == ISD::TRUNCATE &&
+       CTTZ.getOperand(0).getOpcode() != ISD::CTTZ))
+    return SDValue();
+
+  assert((CTTZ.getValueType() == MVT::i32 || CTTZ.getValueType() == MVT::i64) &&
+         "Illegal type in CTTZ folding");
+
+  if (!isNullConstant(Zero) || !isNullConstant(SUBS.getOperand(1)))
+    return SDValue();
+
+  SDValue X = CTTZ.getOpcode() == ISD::TRUNCATE
+                  ? CTTZ.getOperand(0).getOperand(0)
+                  : CTTZ.getOperand(0);
+
+  if (X != SUBS.getOperand(0))
+    return SDValue();
+
+  unsigned BitWidth = CTTZ.getOpcode() == ISD::TRUNCATE
+                          ? CTTZ.getOperand(0).getValueSizeInBits()
+                          : CTTZ.getValueSizeInBits();
+  SDValue BitWidthMinusOne =
+      DAG.getConstant(BitWidth - 1, SDLoc(N), CTTZ.getValueType());
+  return DAG.getNode(ISD::AND, SDLoc(N), CTTZ.getValueType(), CTTZ,
+                     BitWidthMinusOne);
+}
+
 // Optimize CSEL instructions
 static SDValue performCSELCombine(SDNode *N,
                                   TargetLowering::DAGCombinerInfo &DCI,
                                   SelectionDAG &DAG) {
   // CSEL x, x, cc -> x
   if (N->getOperand(0) == N->getOperand(1))
     return N->getOperand(0);
 
+  // CSEL 0, cttz(X), eq(X, 0) -> AND cttz bitwidth-1
+  // CSEL cttz(X), 0, ne(X, 0) -> AND cttz bitwidth-1
+  if (SDValue Folded = foldCSELofCTTZ(N, DAG))
+    return Folded;
+
   return performCONDCombine(N, DCI, DAG, 2, 3);
 }
 
 static SDValue performSETCCCombine(SDNode *N, SelectionDAG &DAG) {
   assert(N->getOpcode() == ISD::SETCC && "Unexpected opcode!");
   SDValue LHS = N->getOperand(0);
   SDValue RHS = N->getOperand(1);
   ISD::CondCode Cond = cast<CondCodeSDNode>(N->getOperand(2))->get();
   SDLoc DL(N);
   EVT VT = N->getValueType(0);
 
   // setcc (csel 0, 1, cond, X), 1, ne ==> csel 0, 1, !cond, X
   if (Cond == ISD::SETNE && isOneConstant(RHS) &&
       LHS->getOpcode() == AArch64ISD::CSEL &&
       isNullConstant(LHS->getOperand(0)) && isOneConstant(LHS->getOperand(1)) &&
       LHS->hasOneUse()) {
     // Invert CSEL's condition.
     auto *OpCC = cast<ConstantSDNode>(LHS.getOperand(2));
     auto OldCond = static_cast<AArch64CC::CondCode>(OpCC->getZExtValue());
     auto NewCond = getInvertedCondCode(OldCond);
 
     // csel 0, 1, !cond, X
     SDValue CSEL =
         DAG.getNode(AArch64ISD::CSEL, DL, LHS.getValueType(), LHS.getOperand(0),
                     LHS.getOperand(1), DAG.getConstant(NewCond, DL, MVT::i32),

llvm/test/CodeGen/AArch64/fold-csel-cttz-and.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64 < %s | FileCheck %s
+
+;; Check the transformation
+;; CSEL 0, cttz, cc -> AND cttz numbits-1
+;; for cttz in the case of i32 and i64 respectively
+
+;; Cases for which the optimzation takes place
+define i32 @cttzi32(i32 %x) {
+; CHECK-LABEL: cttzi32:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    rbit w8, w0
+; CHECK-NEXT:    clz w8, w8
+; CHECK-NEXT:    and w0, w8, #0x1f
+; CHECK-NEXT:    ret
+entry:
+  %0 = call i32 @llvm.cttz.i32(i32 %x, i1 true)
+  %1 = icmp eq i32 %x, 0
+  %2 = select i1 %1, i32 0, i32 %0
+  ret i32 %2
+}
+
+define i64 @cttzi64(i64 %x) {
+; CHECK-LABEL: cttzi64:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    rbit x8, x0
+; CHECK-NEXT:    clz x8, x8
+; CHECK-NEXT:    and x0, x8, #0x3f
+; CHECK-NEXT:    ret
+entry:
+  %0 = call i64 @llvm.cttz.i64(i64 %x, i1 true)
+  %1 = icmp eq i64 %x, 0
+  %2 = select i1 %1, i64 0, i64 %0
+  ret i64 %2
+}
+
+define i32 @cttzi32ne(i32 %x) {
+; CHECK-LABEL: cttzi32ne:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    rbit w8, w0
+; CHECK-NEXT:    clz w8, w8
+; CHECK-NEXT:    and w0, w8, #0x1f
+; CHECK-NEXT:    ret
+entry:
+  %0 = call i32 @llvm.cttz.i32(i32 %x, i1 true)
+  %1 = icmp ne i32 %x, 0
+  %2 = select i1 %1, i32 %0, i32 0
+  ret i32 %2
+}
+
+define i64 @cttzi64ne(i64 %x) {
+; CHECK-LABEL: cttzi64ne:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    rbit x8, x0
+; CHECK-NEXT:    clz x8, x8
+; CHECK-NEXT:    and x0, x8, #0x3f
+; CHECK-NEXT:    ret
+entry:
+  %0 = call i64 @llvm.cttz.i64(i64 %x, i1 true)
+  %1 = icmp ne i64 %x, 0
+  %2 = select i1 %1, i64 %0, i64 0
+  ret i64 %2
+}
+
+define i32 @cttztrunc(i64 %x) {
+; CHECK-LABEL: cttztrunc:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    rbit x8, x0
+; CHECK-NEXT:    clz x8, x8
+; CHECK-NEXT:    and w0, w8, #0x3f
+; CHECK-NEXT:    ret
+entry:
+  %0 = call i64 @llvm.cttz.i64(i64 %x, i1 true)
+  %1 = icmp eq i64 %x, 0
+  %2 = select i1 %1, i64 0, i64 %0
+  %3 = trunc i64 %2 to i32
+  ret i32 %3
+}
+
+;; Cases for which the optimization does not take place
+define i32 @cttzne(i32 %x) {
+; CHECK-LABEL: cttzne:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    rbit w8, w0
+; CHECK-NEXT:    cmp w0, #0
+; CHECK-NEXT:    clz w8, w8
+; CHECK-NEXT:    csel w0, wzr, w8, ne
+; CHECK-NEXT:    ret
+entry:
+  %0 = call i32 @llvm.cttz.i32(i32 %x, i1 true)
+  %1 = icmp ne i32 %x, 0
+  %2 = select i1 %1, i32 0, i32 %0
+  ret i32 %2
+}
+
+define i32 @cttzxnot0(i32 %x) {
+; CHECK-LABEL: cttzxnot0:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    rbit w8, w0
+; CHECK-NEXT:    cmp w0, #10
+; CHECK-NEXT:    clz w8, w8
+; CHECK-NEXT:    csel w0, wzr, w8, eq
+; CHECK-NEXT:    ret
+entry:
+  %0 = call i32 @llvm.cttz.i32(i32 %x, i1 true)
+  %1 = icmp eq i32 %x, 10
+  %2 = select i1 %1, i32 0, i32 %0
+  ret i32 %2
+}
+
+define i32 @cttzlhsnot0(i32 %x) {
+; CHECK-LABEL: cttzlhsnot0:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    rbit w9, w0
+; CHECK-NEXT:    mov w8, #10
+; CHECK-NEXT:    clz w9, w9
+; CHECK-NEXT:    cmp w0, #0
+; CHECK-NEXT:    csel w0, w8, w9, eq
+; CHECK-NEXT:    ret
+entry:
+  %0 = call i32 @llvm.cttz.i32(i32 %x, i1 true)
+  %1 = icmp eq i32 %x, 0
+  %2 = select i1 %1, i32 10, i32 %0
+  ret i32 %2
+}
+
+define i32 @notcttz(i32 %x) {
+; CHECK-LABEL: notcttz:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    clz w8, w0
+; CHECK-NEXT:    cmp w0, #0
+; CHECK-NEXT:    csel w0, wzr, w8, eq
+; CHECK-NEXT:    ret
+entry:
+  %0 = call i32 @llvm.ctlz.i32(i32 %x, i1 true)
+  %1 = icmp eq i32 %x, 0
+  %2 = select i1 %1, i32 0, i32 %0
+  ret i32 %2
+}
+
+define i32 @cttzlhsnotx(i32 %x, i32 %y) {
+; CHECK-LABEL: cttzlhsnotx:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    rbit w8, w0
+; CHECK-NEXT:    cmp w1, #0
+; CHECK-NEXT:    clz w8, w8
+; CHECK-NEXT:    csel w0, wzr, w8, eq
+; CHECK-NEXT:    ret
+entry:
+  %0 = call i32 @llvm.cttz.i32(i32 %x, i1 true)
+  %1 = icmp eq i32 %y, 0
+  %2 = select i1 %1, i32 0, i32 %0
+  ret i32 %2
+}
+
+declare i32 @llvm.cttz.i32(i32, i1)
+
+declare i64 @llvm.cttz.i64(i64, i1)
+
+declare i32 @llvm.ctlz.i32(i32, i1)