[WIP][SelectionDAG] Fold shift constants into cmp

[Allen]

Summary

Perform the following transform for shift constants if legal

((X+1) >> ShiftBits) <= NewC -> X <  (NewC << ShiftBits)
((X+1) >> ShiftBits) >  NewC -> X >= (NewC << ShiftBits)

Fixes https://github.com/llvm/llvm-project/issues/54283

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

4299

This comment doesn't provide much information to the reader outside of the context of this patch. Can you elaborate for future readers?

4322

Same

4547

The first two lines of this have the same thing on both sides. Is that saying we're re-creating the code that is already there?

6061

What changed here?

6319

And here

[Allen]

update comment for more readable

[Allen]

rebase to top commit to fix the code premerge conflict

[Allen]

Fixes clang-format: please reformat the code

[Allen]

@craig.topper please let me know if you have some more idea, thanks

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

4299

Add more detail reason in comment

4322

Done, thanks

4547

Thanks for your quick review, updated.

6061

the origin '-' is not UTF8 code, and it show garbled messages when show with vim

6319

same above

[Allen]

ping ? @craig.topper

[Allen]

this is similar to D121449, so I'll abandon it that commit land

[bcl5980]

In D121355#3384564, @Allen wrote:

this is similar to D121449, so I'll abandon it that commit land

It looks this code works also on arm backend, what I do is only on aarch64.
And D121449 optimize more cases (arm64-xaluo.ll)
So I think both of them are useful.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

4300

it's not the type that is opaque, it's the constant value.

4545–4563

Why does the new transform only happen if shouldAvoidTransformToShift returns false? You aren't creating a shift instruction.

4547

This comment makes it look like the pattern (X >> ShiftBits) >= NewC is being matched, but there is no SRL instruction using X is there?

[craig.topper]

ARMTargetLowering::getARMCmp and AArch64TargetLowering::getAArch64Cmp both have code that tries to adjust the condition code to get the immediate to fit. So it looks like the primary problem is that introduction of the ISD::SRL that blocks the target from seeing a constant they can fix it.

[Allen]

update according comment advise

[Allen]

In D121355#3384841, @craig.topper wrote:

ARMTargetLowering::getARMCmp and AArch64TargetLowering::getAArch64Cmp both have code that tries to adjust the condition code to get the immediate to fit. So it looks like the primary problem is that introduction of the ISD::SRL that blocks the target from seeing a constant they can fix it.

hi, @craig.topper
Yes, ISD::SRL already support the 2^N (isLegalAddImmediate will return false), but SelectionDAG still
fail to match it as it convert the 2^N to (2^N +1)/(2^N -1) , which is not a legal add immediate.
This is because we always do the transform：Canonicalize GE/LE comparisons to use GT/LT comparisons

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

4300

thanks, update comment

4545–4563

good catch, thanks.

4547

thanks, update comment

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

4548

The first two cases never happen do they?

In those cases C1 == NewC << ShiftBits and we already checked isLegalICmpImmediate on line 4500.

[Allen]

ping? @craig.topper

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

4548

Yes, thanks very much!

I think AndRHS->getAPIntValue().isPowerOf2() will be matched, so I update the comment.

4548

hi, @craig.topper:

I had update the comment, do you have other suggestion ?

[craig.topper]

I'm not sure I like this patch as is.

The fundamental problem isn't directly the shift code. We have conflict between canonicalization and target lowering. Canonicalization prefers to create SET(U)GT/SET(U)LT. But that may create an unsupported constant. Is it the target's responsibility to detect reverse the canonicalization during lowering? We do that already for setgt X, -1 on many targets. ARM and AArch64 do it for even more constant values. Or should canonicalization be consulting isLegalICmpImmediate?

If it's the target's responsibility during, then we should just block the shift transform and leave the unsupported immediate so the target can fix it.
If canonicalization in this code should deal with it, then we should be converting SETGT/SETUGT to SETGE/SETUGE based on isLegalICmpImmediate around line 4363 in the original code. The same idea as the changes in this patch to the SETGE/SETUGE canonicalization.

What we have a right now is somewhere in the middle of these two ideas. We're only considering very specific constants by creating the constant the shift code. But we've made the SETGE/SETUGE canonicalization block transforming a potentially larger set of the constants.

[craig.topper]

@spatel or @RKSimon what are your thoughts?

[RKSimon]

I'm a little worried that this is likely to lead to infinite loops with constant/condcode combines fighting - this feels like it should be be backend specific tbh, or maybe even ISelDAGToDAG.

[Allen]

In D121355#3413414, @RKSimon wrote:

I'm a little worried that this is likely to lead to infinite loops with constant/condcode combines fighting - this feels like it should be be backend specific tbh, or maybe even ISelDAGToDAG.

I added condition isLegalICmpImmediate on line 4335, which will guard the infinite loops.

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

if (auto *N1C = isConstOrConstSplat(N1)) {
     // Canonicalize GE/LE comparisons to use GT/LT comparisons.
     if (Cond == ISD::SETGE || Cond == ISD::SETUGE) {
       // X >= MIN --> true
       if (C1 == MinVal)
         return DAG.getBoolConstant(true, dl, VT, OpVT);
 
       if (!VT.isVector()) { // TODO: Support this for vectors.
         // X >= C0 --> X > (C0 - 1)
+        // Add new extra condition to avoid infinitely transform.
         APInt C = C1 - 1;
         ISD::CondCode NewCC = (Cond == ISD::SETGE) ? ISD::SETGT : ISD::SETUGT;
         if ((DCI.isBeforeLegalizeOps() ||
              isCondCodeLegal(NewCC, VT.getSimpleVT())) &&
-            (!N1C->isOpaque() || (C.getBitWidth() <= 64 &&
-                                  isLegalICmpImmediate(C.getSExtValue())))) {
+            ((!N1C->isOpaque() && (C.getBitWidth() > 64 ||
+                                   !isLegalICmpImmediate(C1.getSExtValue()) ||
+                                   isLegalICmpImmediate(C.getSExtValue()))) ||
+             (N1C->isOpaque() && C.getBitWidth() <= 64 &&
           return DAG.getSetCC(dl, VT, N0,
                               DAG.getConstant(C, dl, N1.getValueType()),
                               NewCC);
         }
       }
     }
 
     if (Cond == ISD::SETLE || Cond == ISD::SETULE) {
       // X <= MAX --> true
       if (C1 == MaxVal)
         return DAG.getBoolConstant(true, dl, VT, OpVT);
 
       // X <= C0 --> X < (C0 + 1)
+      // Add new extra condition to avoid infinitely transform.
       if (!VT.isVector()) { // TODO: Support this for vectors.
         APInt C = C1 + 1;
         ISD::CondCode NewCC = (Cond == ISD::SETLE) ? ISD::SETLT : ISD::SETULT;
         if ((DCI.isBeforeLegalizeOps() ||
              isCondCodeLegal(NewCC, VT.getSimpleVT())) &&
-            (!N1C->isOpaque() || (C.getBitWidth() <= 64 &&
-                                  isLegalICmpImmediate(C.getSExtValue())))) {
+            ((!N1C->isOpaque() && (C.getBitWidth() > 64 ||
+                                   !isLegalICmpImmediate(C1.getSExtValue()) ||
+                                   isLegalICmpImmediate(C.getSExtValue()))) ||
+             (N1C->isOpaque() && C.getBitWidth() <= 64 &&
           return DAG.getSetCC(dl, VT, N0,
                               DAG.getConstant(C, dl, N1.getValueType()),
                               NewCC);
         }
       }
     }
 
     if (Cond == ISD::SETLT || Cond == ISD::SETULT) {
if (C1.getMinSignedBits() <= 64 &&
               SDValue CmpRHS = DAG.getConstant(C1.lshr(ShiftBits), dl, ShValTy);
               return DAG.getSetCC(dl, VT, Shift, CmpRHS, Cond);
             }
           }
         }
       } else if (Cond == ISD::SETULT || Cond == ISD::SETUGE ||
                  Cond == ISD::SETULE || Cond == ISD::SETUGT) {
         bool AdjOne = (Cond == ISD::SETULE || Cond == ISD::SETUGT);
-        // X <  0x100000000 -> (X >> 32) <  1
-        // X >= 0x100000000 -> (X >> 32) >= 1
-        // X <= 0x0ffffffff -> (X >> 32) <  1
-        // X >  0x0ffffffff -> (X >> 32) >= 1
         unsigned ShiftBits;
         APInt NewC = C1;
         ISD::CondCode NewCond = Cond;
         if (AdjOne) {
           ShiftBits = C1.countTrailingOnes();
           NewC = NewC + 1;
           NewCond = (Cond == ISD::SETULE) ? ISD::SETULT : ISD::SETUGE;
         } else {
           ShiftBits = C1.countTrailingZeros();
         }
         NewC.lshrInPlace(ShiftBits);
         if (ShiftBits && NewC.getMinSignedBits() <= 64 &&
             isLegalICmpImmediate(NewC.getSExtValue()) &&
             !TLI.shouldAvoidTransformToShift(ShValTy, ShiftBits)) {
+          if (isLegalICmpImmediate(NewC.getSExtValue() << ShiftBits)) {
+            // Single instruction can describe as (NewC << ShiftBits) is legal
+            // X <  (NewC << ShiftBits)    -> X <  (NewC << ShiftBits)
+            // X >= (NewC << ShiftBits)    -> X >= (NewC << ShiftBits)
+            // X <= (NewC << ShiftBits) -1 -> X <  (NewC << ShiftBits)
+            // X >  (NewC << ShiftBits) -1 -> X >= (NewC << ShiftBits)
+            SDValue Shift = DAG.getConstant(NewC.getSExtValue() << ShiftBits,
+                                            dl, ShValTy);
+            return DAG.getSetCC(dl, VT, N0, Shift, NewCond);
+          } else {
+            // X <  0x100000000 -> (X >> 32) <  1
+            // X >= 0x100000000 -> (X >> 32) >= 1
+            // X <= 0x0ffffffff -> (X >> 32) <  1
+            // X >  0x0ffffffff -> (X >> 32) >= 1
             SDValue Shift = DAG.getNode(ISD::SRL, dl, ShValTy, N0,
-                                      DAG.getConstant(ShiftBits, dl, ShiftTy));
             SDValue CmpRHS = DAG.getConstant(NewC, dl, ShValTy);
             return DAG.getSetCC(dl, VT, Shift, CmpRHS, NewCond);
           }
         }
       }
     }
+  }
 
   if (!isa<ConstantFPSDNode>(N0) && isa<ConstantFPSDNode>(N1)) {
     auto *CFP = cast<ConstantFPSDNode>(N1);
     assert(!CFP->getValueAPF().isNaN() && "Unexpected NaN value");
 
     // Otherwise, we know the RHS is not a NaN.  Simplify the node to drop the
     // constant if knowing that the operand is non-nan is enough.  We prefer to
     // have SETO(x,x) instead of SETO(x, 0.0) because this avoids having to
TargetLowering::prepareUREMEqFold(EVT SETCCVT, SDValue REMNode,
   // if C2 is not less than C1, the comparison is always false.
   // But we have produced the comparison that will give the
   // opposive tautological answer. So these lanes would need to be fixed up.
   SDValue TautologicalInvertedChannels =
       DAG.getSetCC(DL, SETCCVT, D, CompTargetNode, ISD::SETULE);
   Created.push_back(TautologicalInvertedChannels.getNode());
 
   // NOTE: we avoid letting illegal types through even if we're before legalize
-  // ops – legalization has a hard time producing good code for this.
+  // ops - legalization has a hard time producing good code for this.
   if (isOperationLegalOrCustom(ISD::VSELECT, SETCCVT)) {
     // If we have a vector select, let's replace the comparison results in the
     // affected lanes with the correct tautological result.
     SDValue Replacement = DAG.getBoolConstant(Cond == ISD::SETEQ ? false : true,
                                               DL, SETCCVT, SETCCVT);
     return DAG.getNode(ISD::VSELECT, DL, SETCCVT, TautologicalInvertedChannels,
                        Replacement, NewCC);
   }
if (!HadIntMinDivisor)
     return Fold;
 
   // That fold is only valid for positive divisors. Which effectively means,
   // it is invalid for INT_MIN divisors. So if we have such a lane,
   // we must fix-up results for said lanes.
   assert(VT.isVector() && "Can/should only get here for vectors.");
 
   // NOTE: we avoid letting illegal types through even if we're before legalize
-  // ops – legalization has a hard time producing good code for the code that
+  // ops - legalization has a hard time producing good code for the code that
   // follows.
   if (!isOperationLegalOrCustom(ISD::SETEQ, VT) ||
       !isOperationLegalOrCustom(ISD::AND, VT) ||
       !isOperationLegalOrCustom(Cond, VT) ||
       !isOperationLegalOrCustom(ISD::VSELECT, SETCCVT))
     return SDValue();
 
   Created.push_back(Fold.getNode());

llvm/test/CodeGen/AArch64/icmp-shift-opt.ll

 ; CHECK-NEXT:    cmp x8, #0
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %shl = shl i256 %a, 17
   %cmp = icmp ne i256 %shl, 0
   ret i1 %cmp
 }
 
+; bool high_mask1 (unsigned int x) { return (x >> 13) == 0; }
+define i1 @high_mask1(i32 noundef %x) {
+; CHECK-LABEL: high_mask1:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    cmp w0, #2, lsl #12 // =8192
+; CHECK-NEXT:    cset w0, hs
+; CHECK-NEXT:    ret
+entry:
+  %cmp = icmp ugt i32 %x, 8191
+  ret i1 %cmp
+}
+
+; bool high_mask2 (unsigned int x) { return (x >> 13) != 0; }
+define i1 @high_mask2(i32 noundef %x) {
+; CHECK-LABEL: high_mask2:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    cmp w0, #2, lsl #12 // =8192
+; CHECK-NEXT:    cset w0, lo
+; CHECK-NEXT:    ret
+entry:
+  %cmp = icmp ult i32 %x, 8192
+  ret i1 %cmp
+}
+
 declare void @use(i128 %a)

llvm/test/CodeGen/ARM/and-cmpz.ll

false:
   ret i32 2
 }
 
 ; CHECK-LABEL: i16_cmpz:
 ; T1:      uxth    r0, r0
 ; T1-NEXT: lsrs    r0, r0, #9
 ; T1-NEXT: bne
 ; T2:      uxth    r0, r0
-; T2-NEXT: movs    r2, #0
-; T2-NEXT: cmp.w   r2, r0, lsr #9
+; T2-NEXT: cmp.w  r0, #512
 define void @i16_cmpz(i16 %x, void (i32)* %foo) {
 entry:
   %cmp = icmp ult i16 %x, 512
   br i1 %cmp, label %if.then, label %if.end
 
 if.then:                                          ; preds = %entry
   tail call void %foo(i32 0) #1
   br label %if.end
 
 if.end:                                           ; preds = %if.then, %entry
   ret void
 }

llvm/test/CodeGen/ARM/consthoist-icmpimm.ll

 ; CHECKV7M-NEXT:    .save {r7, lr}
 ; CHECKV7M-NEXT:    push {r7, lr}
 ; CHECKV7M-NEXT:    mov r12, r0
 ; CHECKV7M-NEXT:    ldr r0, [sp, #16]
 ; CHECKV7M-NEXT:    lsls r0, r0, #31
 ; CHECKV7M-NEXT:    ldrd lr, r0, [sp, #8]
 ; CHECKV7M-NEXT:    beq .LBB6_2
 ; CHECKV7M-NEXT:  @ %bb.1: @ %then
-; CHECKV7M-NEXT:    orrs r2, r3
-; CHECKV7M-NEXT:    lsrs r2, r2, #17
-; CHECKV7M-NEXT:    orr.w r2, r2, r3, lsl #15
-; CHECKV7M-NEXT:    orr.w r3, r12, r1
-; CHECKV7M-NEXT:    cmp r2, #0
+; CHECKV7M-NEXT:    subs.w r2, r2, #131072
+; CHECKV7M-NEXT:    sbcs r2, r3, #0
 ; CHECKV7M-NEXT:    mov r2, r0
-; CHECKV7M-NEXT:    lsr.w r3, r3, #17
-; CHECKV7M-NEXT:    orr.w r1, r3, r1, lsl #15
-; CHECKV7M-NEXT:    it ne
-; CHECKV7M-NEXT:    movne r2, lr
-; CHECKV7M-NEXT:    cmp r1, #0
-; CHECKV7M-NEXT:    it ne
-; CHECKV7M-NEXT:    movne r0, lr
+; CHECKV7M-NEXT:    it hs
+; CHECKV7M-NEXT:    movhs r2, lr
+; CHECKV7M-NEXT:    subs.w r3, r12, #131072
+; CHECKV7M-NEXT:    sbcs r1, r1, #0
+; CHECKV7M-NEXT:    it hs
+; CHECKV7M-NEXT:    movhs r0, lr
 ; CHECKV7M-NEXT:    add r0, r2
 ; CHECKV7M-NEXT:    pop {r7, pc}
 ; CHECKV7M-NEXT:  .LBB6_2: @ %else
-; CHECKV7M-NEXT:    orr.w r1, r2, r3
-; CHECKV7M-NEXT:    lsrs r1, r1, #17
-; CHECKV7M-NEXT:    orr.w r1, r1, r3, lsl #15
-; CHECKV7M-NEXT:    cmp r1, #0
-; CHECKV7M-NEXT:    it ne
-; CHECKV7M-NEXT:    movne r0, lr
+; CHECKV7M-NEXT:    subs.w r1, r2, #131072
+; CHECKV7M-NEXT:    sbcs r1, r3, #0
+; CHECKV7M-NEXT:    it hs
+; CHECKV7M-NEXT:    movhs r0, lr
 ; CHECKV7M-NEXT:    pop {r7, pc}
 ;
 ; CHECKV7A-LABEL: icmp64_uge_m2:
 ; CHECKV7A:       @ %bb.0:
 ; CHECKV7A-NEXT:    .save {r4, lr}
 ; CHECKV7A-NEXT:    push {r4, lr}
 ; CHECKV7A-NEXT:    ldr r4, [sp, #16]
 ; CHECKV7A-NEXT:    mov r12, r0
 ; CHECKV7A-NEXT:    ldrd lr, r0, [sp, #8]
 ; CHECKV7A-NEXT:    lsls r4, r4, #31
 ; CHECKV7A-NEXT:    beq .LBB6_2
 ; CHECKV7A-NEXT:  @ %bb.1: @ %then
-; CHECKV7A-NEXT:    orrs r2, r3
-; CHECKV7A-NEXT:    lsrs r2, r2, #17
-; CHECKV7A-NEXT:    orr.w r2, r2, r3, lsl #15
-; CHECKV7A-NEXT:    orr.w r3, r12, r1
-; CHECKV7A-NEXT:    cmp r2, #0
+; CHECKV7A-NEXT:    subs.w r2, r2, #131072
+; CHECKV7A-NEXT:    sbcs r2, r3, #0
 ; CHECKV7A-NEXT:    mov r2, r0
-; CHECKV7A-NEXT:    lsr.w r3, r3, #17
-; CHECKV7A-NEXT:    orr.w r1, r3, r1, lsl #15
-; CHECKV7A-NEXT:    it ne
-; CHECKV7A-NEXT:    movne r2, lr
-; CHECKV7A-NEXT:    cmp r1, #0
-; CHECKV7A-NEXT:    it ne
-; CHECKV7A-NEXT:    movne r0, lr
+; CHECKV7A-NEXT:    it hs
+; CHECKV7A-NEXT:    movhs r2, lr
+; CHECKV7A-NEXT:    subs.w r3, r12, #131072
+; CHECKV7A-NEXT:    sbcs r1, r1, #0
+; CHECKV7A-NEXT:    it hs
+; CHECKV7A-NEXT:    movhs r0, lr
 ; CHECKV7A-NEXT:    add r0, r2
 ; CHECKV7A-NEXT:    pop {r4, pc}
 ; CHECKV7A-NEXT:  .LBB6_2: @ %else
-; CHECKV7A-NEXT:    orr.w r1, r2, r3
-; CHECKV7A-NEXT:    lsrs r1, r1, #17
-; CHECKV7A-NEXT:    orr.w r1, r1, r3, lsl #15
-; CHECKV7A-NEXT:    cmp r1, #0
-; CHECKV7A-NEXT:    it ne
-; CHECKV7A-NEXT:    movne r0, lr
+; CHECKV7A-NEXT:    subs.w r1, r2, #131072
+; CHECKV7A-NEXT:    sbcs r1, r3, #0
+; CHECKV7A-NEXT:    it hs
+; CHECKV7A-NEXT:    movhs r0, lr
 ; CHECKV7A-NEXT:    pop {r4, pc}
   br i1 %c, label %then, label %else
 then:
   %c1 = icmp uge i64 %x, 131072
   %c2 = icmp uge i64 %y, 131072
   %s1 = select i1 %c1, i32 %a, i32 %b
   %s2 = select i1 %c2, i32 %a, i32 %b
   %r = add i32 %s1, %s2

llvm/test/CodeGen/ARM/icmp-shift-opt.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=armv7 %s -o - | FileCheck %s
 
 ; Optimize expanded SRL/SHL used as an input of
 ; SETCC comparing it with zero by removing rotation.
 ;
 ; See https://bugs.llvm.org/show_bug.cgi?id=50197
 define i64 @opt_setcc_lt_power_of_2(i64 %a) nounwind {
 ; CHECK-LABEL: opt_setcc_lt_power_of_2:
 ; CHECK:       @ %bb.0:
 ; CHECK-NEXT:  .LBB0_1: @ %loop
 ; CHECK-NEXT:    @ =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    adds r0, r0, #1
 ; CHECK-NEXT:    adc r1, r1, #0
-; CHECK-NEXT:    orr r2, r0, r1
-; CHECK-NEXT:    uxth r3, r1
-; CHECK-NEXT:    orr r2, r3, r2, lsr #16
-; CHECK-NEXT:    cmp r2, #0
-; CHECK-NEXT:    bne .LBB0_1
+; CHECK-NEXT:    subs r2, r0, #65536
+; CHECK-NEXT:    sbcs r2, r1, #0
+; CHECK-NEXT:    bhs .LBB0_1
 ; CHECK-NEXT:  @ %bb.2: @ %exit
 ; CHECK-NEXT:    bx lr
   br label %loop
 
 loop:
   %phi.a = phi i64 [ %a, %0 ], [ %inc, %loop ]
   %inc = add i64 %phi.a, 1
   %cmp = icmp ult i64 %inc, 65536