[X86] Fold some (truncate (srl (add X, C1), C2)) patterns to (add (truncate (srl X, C2), C1'))

C1' will be smaller than C1 so we are able to avoid generating code with MOVABS and large constants in certain cases.

Author: joaotgouveia

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -48472,6 +48472,64 @@ static SDValue combineSetCCMOVMSK(SDValue EFLAGS, X86::CondCode &CC,
   return SDValue();
 }
 
+// Attempt to fold some (truncate (srl (add X, C1), C2)) patterns to
+// (add (truncate (srl X, C2), C1')). C1' will be smaller than C1 so we are able
+// to avoid generating code with MOVABS and large constants in certain cases.
+static SDValue combineSetCCTruncAdd(SDValue EFLAGS, X86::CondCode &CC,
+                                    SelectionDAG &DAG) {
+  if (!(CC == X86::COND_E || CC == X86::COND_NE || CC == X86::COND_AE ||
+        CC == X86::COND_B))
+    return SDValue();
+
+  EVT VT = EFLAGS.getValueType();
+  if (EFLAGS.getOpcode() == X86ISD::SUB && VT == MVT::i32) {
+    SDValue CmpLHS = EFLAGS.getOperand(0);
+    auto *CmpConstant = dyn_cast<ConstantSDNode>(EFLAGS.getOperand(1));
+
+    if (CmpLHS.getOpcode() != ISD::TRUNCATE || !CmpConstant)
+      return SDValue();
+
+    SDValue Srl = CmpLHS.getOperand(0);
+    EVT SrlVT = Srl.getValueType();
+    if (Srl.getOpcode() != ISD::SRL || SrlVT != MVT::i64)
+      return SDValue();
+
+    SDValue Add = Srl.getOperand(0);
+    // Avoid changing the ADD if it is used elsewhere.
+    if (Add.getOpcode() != ISD::ADD || !Add.hasOneUse())
+      return SDValue();
+
+    auto *AddConstant = dyn_cast<ConstantSDNode>(Add.getOperand(1));
+    auto *SrlConstant = dyn_cast<ConstantSDNode>(Srl.getOperand(1));
+    if (!AddConstant || !SrlConstant)
+      return SDValue();
+
+    APInt AddConstVal = AddConstant->getAPIntValue();
+    APInt SrlConstVal = SrlConstant->getAPIntValue();
+    if (!SrlConstVal.ugt(VT.getSizeInBits()))
+      return SDValue();
+
+    APInt CmpConstVal = CmpConstant->getAPIntValue();
+    APInt ShiftedAddConst = AddConstVal.lshr(SrlConstVal);
+    if (!CmpConstVal.ult(ShiftedAddConst.trunc(VT.getSizeInBits())) ||
+        (ShiftedAddConst.shl(SrlConstVal)) != AddConstVal)
+      return SDValue();
+
+    SDLoc DL(EFLAGS);
+    SDValue AddLHSSrl =
+        DAG.getNode(ISD::SRL, DL, SrlVT, Add.getOperand(0), Srl.getOperand(1));
+    SDValue Trunc = DAG.getNode(ISD::TRUNCATE, DL, VT, AddLHSSrl);
+
+    APInt NewAddConstVal =
+        (~((~AddConstVal).lshr(SrlConstVal))).trunc(VT.getSizeInBits());
+    SDValue NewAddConst = DAG.getConstant(NewAddConstVal, DL, VT);
+    SDValue NewAddNode = DAG.getNode(ISD::ADD, DL, VT, Trunc, NewAddConst);
+    return DAG.getNode(X86ISD::CMP, DL, VT, NewAddNode, EFLAGS.getOperand(1));
+  }
+
+  return SDValue();
+}
+
 /// Optimize an EFLAGS definition used according to the condition code \p CC
 /// into a simpler EFLAGS value, potentially returning a new \p CC and replacing
 /// uses of chain values.
@@ -48494,6 +48552,9 @@ static SDValue combineSetCCEFLAGS(SDValue EFLAGS, X86::CondCode &CC,
   if (SDValue R = combineSetCCMOVMSK(EFLAGS, CC, DAG, Subtarget))
     return R;
 
+  if (SDValue R = combineSetCCTruncAdd(EFLAGS, CC, DAG))
+    return R;
+
   return combineSetCCAtomicArith(EFLAGS, CC, DAG, Subtarget);
 }
 

llvm/test/CodeGen/X86/combine-setcc-trunc-add.ll

@@ -6,10 +6,9 @@
 define i1 @test_ult_trunc_add(i64 %x) {
 ; X64-LABEL: test_ult_trunc_add:
 ; X64:       # %bb.0: # %entry
-; X64-NEXT:    movabsq $3940649673949184, %rax # imm = 0xE000000000000
-; X64-NEXT:    addq %rdi, %rax
-; X64-NEXT:    shrq $48, %rax
-; X64-NEXT:    cmpl $3, %eax
+; X64-NEXT:    shrq $48, %rdi
+; X64-NEXT:    addl $-65522, %edi # imm = 0xFFFF000E
+; X64-NEXT:    cmpl $3, %edi
 ; X64-NEXT:    setb %al
 ; X64-NEXT:    retq
 entry:
@@ -23,10 +22,9 @@ entry:
 define i1 @test_ult_add(i64 %x) {
 ; X64-LABEL: test_ult_add:
 ; X64:       # %bb.0: # %entry
-; X64-NEXT:    movabsq $3940649673949184, %rax # imm = 0xE000000000000
-; X64-NEXT:    addq %rdi, %rax
-; X64-NEXT:    shrq $48, %rax
-; X64-NEXT:    cmpl $3, %eax
+; X64-NEXT:    shrq $48, %rdi
+; X64-NEXT:    addl $-65522, %edi # imm = 0xFFFF000E
+; X64-NEXT:    cmpl $3, %edi
 ; X64-NEXT:    setb %al
 ; X64-NEXT:    retq
 entry:
@@ -38,10 +36,9 @@ entry:
 define i1 @test_ugt_trunc_add(i64 %x) {
 ; X64-LABEL: test_ugt_trunc_add:
 ; X64:       # %bb.0: # %entry
-; X64-NEXT:    movabsq $3940649673949184, %rax # imm = 0xE000000000000
-; X64-NEXT:    addq %rdi, %rax
-; X64-NEXT:    shrq $48, %rax
-; X64-NEXT:    cmpl $4, %eax
+; X64-NEXT:    shrq $48, %rdi
+; X64-NEXT:    addl $-65522, %edi # imm = 0xFFFF000E
+; X64-NEXT:    cmpl $4, %edi
 ; X64-NEXT:    setae %al
 ; X64-NEXT:    retq
 entry:
@@ -70,10 +67,9 @@ entry:
 define i1 @test_eq_trunc_add(i64 %x) {
 ; X64-LABEL: test_eq_trunc_add:
 ; X64:       # %bb.0: # %entry
-; X64-NEXT:    movabsq $3940649673949184, %rax # imm = 0xE000000000000
-; X64-NEXT:    addq %rdi, %rax
-; X64-NEXT:    shrq $48, %rax
-; X64-NEXT:    cmpl $3, %eax
+; X64-NEXT:    shrq $48, %rdi
+; X64-NEXT:    addl $-65522, %edi # imm = 0xFFFF000E
+; X64-NEXT:    cmpl $3, %edi
 ; X64-NEXT:    sete %al
 ; X64-NEXT:    retq
 entry:
@@ -100,10 +96,9 @@ entry:
 define i1 @test_ne_trunc_add(i64 %x) {
 ; X64-LABEL: test_ne_trunc_add:
 ; X64:       # %bb.0: # %entry
-; X64-NEXT:    movabsq $3940649673949184, %rax # imm = 0xE000000000000
-; X64-NEXT:    addq %rdi, %rax
-; X64-NEXT:    shrq $48, %rax
-; X64-NEXT:    cmpl $3, %eax
+; X64-NEXT:    shrq $48, %rdi
+; X64-NEXT:    addl $-65522, %edi # imm = 0xFFFF000E
+; X64-NEXT:    cmpl $3, %edi
 ; X64-NEXT:    setne %al
 ; X64-NEXT:    retq
 entry: