[InstCombine] Fold zext(X) + C2 pred C -> X + C3 pred C4

[dtcxzyw]

Motivating case from https://github.com/torvalds/linux/blob/9852d85ec9d492ebef56dc5f229416c925758edc/drivers/gpu/drm/drm_edid.c#L5238-L5240:

define i1 @src(i8 noundef %v13) {
entry:
  %conv1 = zext i8 %v13 to i32
  %add = add nsw i32 %conv1, -4
  %cmp = icmp ult i32 %add, 3
  %cmp4 = icmp slt i8 %v13, 4
  %cond = select i1 %cmp4, i1 true, i1 %cmp
  ret i1 %cond
}

define i1 @tgt(i8 noundef %v13) {
entry:
  %cmp4 = icmp slt i8 %v13, 7
  ret i1 %cmp4
}

[goldsteinn]

Think the 'iff ...' is pretty misleading given that you also need 'C2' and 'C' to fit in the new bitwidth.

[dtcxzyw]

We need a new ConstantRange API to convert ranges for zext(V) to ranges for V.
For example, we can convert zext(i8 X to i32) - 255 u< -4 to X + 1 u< -4. But current implementation cannot achieve this.

[goldsteinn]

Is this intentionally still in draft form? Or are you ready for review?

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

Motivating case from https://github.com/torvalds/linux/blob/9852d85ec9d492ebef56dc5f229416c925758edc/drivers/gpu/drm/drm_edid.c#L5238-L5240:

define i1 @<!-- -->src(i8 noundef %v13) {
entry:
  %conv1 = zext i8 %v13 to i32
  %add = add nsw i32 %conv1, -4
  %cmp = icmp ult i32 %add, 3
  %cmp4 = icmp slt i8 %v13, 4
  %cond = select i1 %cmp4, i1 true, i1 %cmp
  ret i1 %cond
}

define i1 @<!-- -->tgt(i8 noundef %v13) {
entry:
  %cmp4 = icmp slt i8 %v13, 7
  ret i1 %cmp4
}

---

Full diff: https://github.com/llvm/llvm-project/pull/110511.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp (+25)
• (modified) llvm/test/Transforms/InstCombine/icmp-add.ll (+91)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 6c3fc987d9add2..0fbf446480d0dc 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -3168,6 +3168,31 @@ Instruction *InstCombinerImpl::foldICmpAddConstant(ICmpInst &Cmp,
                         Builder.CreateAdd(X, ConstantInt::get(Ty, *C2 - C - 1)),
                         ConstantInt::get(Ty, ~C));
 
+  // zext(V) + C2 pred C -> V + C3 pred' C4
+  Value *V;
+  if (match(X, m_ZExt(m_Value(V)))) {
+    Type *NewCmpTy = V->getType();
+    unsigned CmpBW = Ty->getScalarSizeInBits();
+    unsigned NewCmpBW = NewCmpTy->getScalarSizeInBits();
+    if (shouldChangeType(Ty, NewCmpTy)) {
+      if (auto ZExtCR = CR.exactIntersectWith(ConstantRange(
+              APInt::getZero(CmpBW), APInt::getOneBitSet(CmpBW, NewCmpBW)))) {
+        ConstantRange SrcCR = ZExtCR->truncate(NewCmpBW);
+        CmpInst::Predicate EquivPred;
+        APInt EquivInt;
+        APInt EquivOffset;
+
+        SrcCR.getEquivalentICmp(EquivPred, EquivInt, EquivOffset);
+        return new ICmpInst(
+            EquivPred,
+            EquivOffset.isZero()
+                ? V
+                : Builder.CreateAdd(V, ConstantInt::get(NewCmpTy, EquivOffset)),
+            ConstantInt::get(NewCmpTy, EquivInt));
+      }
+    }
+  }
+
   return nullptr;
 }
 
diff --git a/llvm/test/Transforms/InstCombine/icmp-add.ll b/llvm/test/Transforms/InstCombine/icmp-add.ll
index 0c141d4b8e73aa..2239e48468ee04 100644
--- a/llvm/test/Transforms/InstCombine/icmp-add.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-add.ll
@@ -3183,3 +3183,94 @@ define i1 @icmp_of_ucmp_plus_const_with_const(i32 %x, i32 %y) {
   %cmp2 = icmp ult i8 %add, 2
   ret i1 %cmp2
 }
+
+define i1 @zext_range_check_ult(i8 %x) {
+; CHECK-LABEL: @zext_range_check_ult(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i8 [[X:%.*]], -4
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[TMP0]], 3
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+entry:
+  %conv = zext i8 %x to i32
+  %add = add i32 %conv, -4
+  %cmp = icmp ult i32 %add, 3
+  ret i1 %cmp
+}
+
+; TODO: should be canonicalized to (x - 4) u> 2
+define i1 @zext_range_check_ugt(i8 %x) {
+; CHECK-LABEL: @zext_range_check_ugt(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[X:%.*]] to i32
+; CHECK-NEXT:    [[TMP0:%.*]] = add nsw i32 [[CONV]], -7
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[TMP0]], -3
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+entry:
+  %conv = zext i8 %x to i32
+  %add = add i32 %conv, -4
+  %cmp = icmp ugt i32 %add, 2
+  ret i1 %cmp
+}
+
+; TODO: should be canonicalized to (x - 4) u> 2
+define i1 @zext_range_check_ult_alter(i8 %x) {
+; CHECK-LABEL: @zext_range_check_ult_alter(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[X:%.*]] to i32
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[CONV]], -7
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[ADD]], -3
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+entry:
+  %conv = zext i8 %x to i32
+  %add = add i32 %conv, -7
+  %cmp = icmp ult i32 %add, -3
+  ret i1 %cmp
+}
+
+define i1 @zext_range_check_mergable(i8 %x) {
+; CHECK-LABEL: @zext_range_check_mergable(
+; CHECK-NEXT:    [[COND:%.*]] = icmp slt i8 [[X:%.*]], 7
+; CHECK-NEXT:    ret i1 [[COND]]
+;
+  %conv = zext i8 %x to i32
+  %add = add nsw i32 %conv, -4
+  %cmp1 = icmp ult i32 %add, 3
+  %cmp2 = icmp slt i8 %x, 4
+  %cond = select i1 %cmp2, i1 true, i1 %cmp1
+  ret i1 %cond
+}
+
+; Negative tests
+
+define i1 @sext_range_check_ult(i8 %x) {
+; CHECK-LABEL: @sext_range_check_ult(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[X:%.*]] to i32
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[CONV]], -4
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[ADD]], 3
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+entry:
+  %conv = sext i8 %x to i32
+  %add = add i32 %conv, -4
+  %cmp = icmp ult i32 %add, 3
+  ret i1 %cmp
+}
+
+define i1 @zext_range_check_ult_illegal_type(i7 %x) {
+; CHECK-LABEL: @zext_range_check_ult_illegal_type(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CONV:%.*]] = zext i7 [[X:%.*]] to i32
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[CONV]], -4
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[ADD]], 3
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+entry:
+  %conv = zext i7 %x to i32
+  %add = add i32 %conv, -4
+  %cmp = icmp ult i32 %add, 3
+  ret i1 %cmp
+}

[dtcxzyw]

Is this intentionally still in draft form? Or are you ready for review?

It is ready for review now. I believe I cannot generalize it further without new ConstantRange API.