[DAGCombine] Generalize foldSelectCCToShiftAnd

[mgudim]

Summary

[DAGCombine] Generalize foldSelectCCToShiftAnd

There are many equivalent ways to express the pattern `select_cc setlt
X, 0, A, 0". In this PR we extend the code in foldSelectCCToShiftAnd
to work on all such patterns.

By the time we call foldSelectCCToShiftAnd, the non-strict integer
inequalities must have been canonicalized to strict ones. But we can't
expect that the zero constant can only appear as the "false" argument of
the select. In this patch we extend foldSelectCCToShiftAnd to handle
these "inverted" cases.

[mgudim]

llvm/test/CodeGen/RISCV/select-to-shift-and.ll

2

This is a target-independent optimization, so I am not sure how to properly test it. I included this file only to demonstrate how the optimization works.

Please let me know what you think a proper place / form for this test should be.

[craig.topper]

Do the SETLE and SETGE cases occur? I thought DAGCombine always changes those to SETLT and SETGE by changing the constant?

[mgudim]

In D137422#3923194, @craig.topper wrote:

Do the SETLE and SETGE cases occur? I thought DAGCombine always changes those to SETLT and SETGE by changing the constant?

@craig.topper
Thanks for taking a look! Yes, you're right. I worked on this a month ago and some cases didn't work. Today I ran the latest master and everything worked. The simplification happend in SimpifySetCC. I am not sure what changed exactly. Anyway, the approach with canonicalizing SETCC is much better than making all other combines handle all the posibilities.

I'll see why the optimization didn't happen in the AArch64 test case.

[mgudim]

@craig.topper

When I previously said that now optimization works in all the cases I was wrong. You're right that we don't need to consider the non-strict inequalities. But we still need to extend the logic to work on the "inverted" patterns where the zero is the true value in the select instruction.

[craig.topper]

llvm/test/CodeGen/RISCV/select-to-shift-and.ll

2

Why did you stop on MIR instead of assembly?

[mgudim]

llvm/test/CodeGen/RISCV/select-to-shift-and.ll

2

Since we are testing something in ISel I thought it would make sense to run only ISel.

I do see that most of other tests go all the way to assembly. I am fine with either way.

[mgudim]

@craig.topper For the test, do you want to go all the way to assembly or stopping after isel is OK? Since the optimization is target-independent, where should the test go? Do you have further comments?

[reames]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

24476

A style, and process suggestion.

These first patterns can be handled via the common idiom of reading the operands, and then swapping the operand values before the main transform if required.

This would stand along well as it's own change. This is the part which currently matches your change description anyways.

24484

This part relies on an additional observation that X == 0, the two sides of the select return the same value.

There's an interesting question about whether we want to canonicalize these towards the other form. 0 and -1 are likely cheaper constants on at least some architectures. Might be worth an experiment to see how that plays out.

Assuming we don't, this can be a separate change with a clarified comment to emphasize the additional fact stated at the top of this comment.

llvm/test/CodeGen/RISCV/select-to-shift-and.ll

2

For codegen tests, it is expected to show assembly. That's normal practice, and you should follow that. Please replace update_mir with update_llc_test_checks.py

[mgudim]

Implemented approach suggested by @reames : instead of considering all the possible cases, assume that the zero constant always appears as the "false" option of select and if not, invert the comparison, switch the true and false option and call foldSelectCCToShiftAnd recursively.

Updated test to run all the way to final assembly.

[mgudim]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

24468

I thought about doing this reversal all the way up in SimplifySelectCC. This may enable other transformations that previously assumed the canonicalized form. For example:

// Fold select_cc setgt X, -1, C, ~C -> xor (ashr X, BW-1), C
// Fold select_cc setlt X, 0, C, ~C -> xor (ashr X, BW-1), ~C

I haven't tried it yet, but do you think this is worth exploring?

24484

Ok, let's leave at as a separate change for later.

llvm/test/CodeGen/AArch64/arm64-fmax.ll

61

I need to update 59 tests for these change. In some tests the checks were not auto-generated. What is the usual procedure to update many tests? Is it ok to insert auto-generated checks in the tests that didn't have them before?

[mgudim]

Added necessary checks missing in previous version. This fixed the failing tests.

[mgudim]

Removed sve-select.s which was added by accident.

[mgudim]

@reames Friendly ping.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::foldSelectCCToShiftAnd(const SDLoc &DL, SDValue N0,
                                             SDValue N1, SDValue N2, SDValue N3,
                                             ISD::CondCode CC) {
   // If this is a select where the false operand is zero and the compare is a
   // check of the sign bit, see if we can perform the "gzip trick":
   // select_cc setlt X, 0, A, 0 -> and (sra X, size(X)-1), A
   // select_cc setgt X, 0, A, 0 -> and (not (sra X, size(X)-1)), A
   EVT XType = N0.getValueType();
   EVT AType = N2.getValueType();
-  if (!isNullConstant(N3) || !XType.bitsGE(AType))
+  if (!XType.bitsGE(AType))
     return SDValue();
 
+  if (!isNullConstant(N3) && isNullConstant(N2)) {
+    // try to swap N2 and N3.
+    if (!XType.isSimple() || !XType.isInteger() || XType.isVector())
+      return SDValue();
+    ConstantSDNode *N1C = dyn_cast<ConstantSDNode>(N1);
+    if (!N1C)
+      return SDValue();
+    // Canonicalize GE/LE comparisons to use GT/LT comparisons.
+    ISD::CondCode InvertedCC = ISD::getSetCCInverse(CC, XType);
+    if (LegalOperations &&
+        !TLI.isCondCodeLegal(InvertedCC, XType.getSimpleVT()))
+      return SDValue();
+    const APInt &RHSC = N1C->getAPIntValue();
+    APInt C;
+    ISD::CondCode NewCC;
+    if (InvertedCC == ISD::SETLE) {
+      C = RHSC + 1;
+      NewCC = ISD::SETLT;
+    } else if (InvertedCC == ISD::SETGE) {
+      C = RHSC - 1;
+      NewCC = ISD::SETGT;
+    } else
+      return SDValue();
+    if (LegalOperations &&
+        (C.getBitWidth() > 64 || !TLI.isLegalICmpImmediate(C.getSExtValue())))
+      return SDValue();
+    return foldSelectCCToShiftAnd(DL, N0, DAG.getConstant(C, DL, XType), N3, N2,
+                                  NewCC);
+  }
+
   // If the comparison is testing for a positive value, we have to invert
   // the sign bit mask, so only do that transform if the target has a bitwise
   // 'and not' instruction (the invert is free).
   if (CC == ISD::SETGT && TLI.hasAndNot(N2)) {
     // (X > -1) ? A : 0
     // (X >  0) ? X : 0 <-- This is canonical signed max.
     if (!(isAllOnesConstant(N1) || (isNullConstant(N1) && N0 == N2)))
       return SDValue();

llvm/test/CodeGen/AArch64/arm64-fmax.ll

; CHECK-NEXT:    ret
   %tst = fcmp nnan une float %lhs, %rhs
   %res = select i1 %tst, float %rhs, float %lhs
   ret float %res
 }
 
 ; Make sure the transformation isn't triggered for integers
 define i64 @test_integer(i64  %in) {
 ; CHECK-LABEL: test_integer:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    cmp x0, #0
-; CHECK-NEXT:    csel x0, xzr, x0, lt
+; CHECK-NEXT:    bic x0, x0, x0, asr #63 
 ; CHECK-NEXT:    ret
   %cmp = icmp slt i64 %in, 0
   %val = select i1 %cmp, i64 0, i64 %in
   ret i64 %val
 }
 
 ; FIXME: It'd be nice for this to create an fmin instruction!
 define float @test_f16(half %in) {

llvm/test/CodeGen/RISCV/select-to-shift-and.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+
+; RUN: llc -mtriple=riscv64 -mattr=+zbb -target-abi=lp64 < %s | FileCheck %s
+
+define i64 @test0(i64 %a, i64 %b) {
+; CHECK-LABEL: test0:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    slti a0, a0, 0
+; CHECK-NEXT:    xori a0, a0, 1
+; CHECK-NEXT:    addi a0, a0, -1
+; CHECK-NEXT:    and a0, a0, a1
+; CHECK-NEXT:    ret
+entry:
+  %cmp.inv = icmp sgt i64 %a, -1
+  %res = select i1 %cmp.inv, i64 0, i64 %b
+  ret i64 %res
+}
+
+define i64 @test1(i64 %a, i64 %b) {
+; CHECK-LABEL: test1:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    srai a0, a0, 63
+; CHECK-NEXT:    and a0, a0, a1
+; CHECK-NEXT:    ret
+entry:
+  %cmp.inv = icmp slt i64 %a, 0
+  %res = select i1 %cmp.inv, i64 %b, i64 0
+  ret i64 %res
+}
+
+define i64 @test2(i64 %a, i64 %b) {
+; CHECK-LABEL: test2:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    srai a0, a0, 63
+; CHECK-NEXT:    andn a0, a1, a0
+; CHECK-NEXT:    ret
+entry:
+  %cmp.inv = icmp sgt i64 %a, -1
+  %res = select i1 %cmp.inv, i64 %b, i64 0
+  ret i64 %res
+}
+
+define i64 @test3(i64 %a, i64 %b) {
+; CHECK-LABEL: test3:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    slti a0, a0, 0
+; CHECK-NEXT:    addi a0, a0, -1
+; CHECK-NEXT:    and a0, a0, a1
+; CHECK-NEXT:    ret
+entry:
+  %cmp.inv = icmp slt i64 %a, 0
+  %res = select i1 %cmp.inv, i64 0, i64 %b
+  ret i64 %res
+}
+
+define i64 @test4(i64 %x) {
+; CHECK-LABEL: test4:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    min a0, a0, zero
+; CHECK-NEXT:    ret
+entry:
+  %cmp.inv = icmp sgt i64 %x, 0
+  %res = select i1 %cmp.inv, i64 0, i64 %x
+  ret i64 %res
+}
+
+define i64 @test5(i64 %x) {
+; CHECK-LABEL: test5:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    srai a1, a0, 63
+; CHECK-NEXT:    and a0, a1, a0
+; CHECK-NEXT:    ret
+entry:
+  %cmp.inv = icmp slt i64 %x, 1
+  %res = select i1 %cmp.inv, i64 %x, i64 0
+  ret i64 %res
+}
+
+define i64 @test6(i64 %x) {
+; CHECK-LABEL: test6:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    max a0, a0, zero
+; CHECK-NEXT:    ret
+entry:
+  %cmp.inv = icmp sgt i64 %x, 0
+  %res = select i1 %cmp.inv, i64 %x, i64 0
+  ret i64 %res
+}
+
+define i64 @test7(i64 %x) {
+; CHECK-LABEL: test7:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    slti a1, a0, 1
+; CHECK-NEXT:    addi a1, a1, -1
+; CHECK-NEXT:    and a0, a1, a0
+; CHECK-NEXT:    ret
+entry:
+  %cmp.inv = icmp slt i64 %x, 1
+  %res = select i1 %cmp.inv, i64 0, i64 %x
+  ret i64 %res
+}