[LVI] Handle icmp of ashr.

[aemerson]

Summary

This handles the case where this combine:

icmp sgt (ashr X, ShAmtC), C --> icmp sgt X, ((C + 1) << ShAmtC) - 1

wasn't performed by instcombine.

This is my first patch to LVI and I'm not really sure this is right.

Proof of the original combine: https://alive2.llvm.org/ce/z/SfpsvX

[aemerson]

ping

[nikic]

Can you please add an alive2 proof to the patch description?

llvm/lib/Analysis/LazyValueInfo.cpp

1154

For symmetry reasons, this needs to handle all signed predicates (sgt, sge, slt, sle). The non-strict predicates are non-canonical by themselves, but will appear when the inverse predicate is taken.

[aemerson]

The original proof from 7ac2db6a489f3f8e588589a69164882df7973d34 isn't alive (ha) anymore, so I re-created it: https://alive2.llvm.org/ce/z/SfpsvX

[nikic]

In D151911#4416351, @aemerson wrote:

The original proof from 7ac2db6a489f3f8e588589a69164882df7973d34 isn't alive (ha) anymore, so I re-created it: https://alive2.llvm.org/ce/z/SfpsvX

Thanks! However, that proof has a bunch of preconditions -- don't we need to check them here as well?

llvm/lib/Analysis/LazyValueInfo.cpp

1155

Don't we have to adjust the range based on the predicate? E.g. the result for sgt and sle would be inverts of each other, while this returns the same range for both.

[aemerson]

Fix the ranges for the other predicates.

Also add precondition checks. However I wasn't able to construct a test that actually covers the checks because those all get optimized away before the checks are hit.

[aemerson]

ping

[aemerson]

ping @nikic does this look right?

[aemerson]

ping

[nikic]

llvm/lib/Analysis/LazyValueInfo.cpp

1152

I think the better way to explain this is icmp slt (ashr X, ShAmtC), C --> icmp slt X, C << ShAmtC (https://alive2.llvm.org/ce/z/ww5u4i). In this case the only precondition is ((C << ShAmtC) >> ShAmtC) == C. Same for sge.

All the +1/-1 get introduced by using sgt/sle because we're effectively trying to reduce them to that base case. This also introduces the extra preconditions to avoid overflow on +/-1.

With this in mind, I think we can introduce an abstraction like this to transparently handle all the signed predicates while only specifying the fold for slt (without knowledge of anything about ashr in particular):

static std::optional<ConstantRange> getRangeViaSLT(
    CmpInst::Predicate Pred, APInt RHS,
    function_ref<std::optional<ConstantRange>(const APInt &)> Fn) {
  bool Invert = false;
  if (Pred == ICmpInst::SGT || Pred == ICmpInst::SGE) {
    Pred = ICmpInst::getInversePredicate(Pred);
    Invert = true;
  }
  if (Pred == ICmpInst::SLE) {
    Pred = ICmpInst::SLT;
    if (RHS.isSignedMax())
      return std::nullopt; // Could also return full/empty here, if we wanted.
    ++RHS;
  }
  assert(Pred == ICmpInst::SLT && "Must be signed predicate");
  if (auto CR = Fn(RHS))
    return Invert ? CR->inverse() : CR;
  return std::nullopt;
}

And then do something like this:

const APInt *ShAmtC;
if (CmpInst::isSigned(EdgePred) &&
    match(LHS, m_AShr(m_Specific(Val), m_APInt(ShAmtC))) &&
    match(RHS, m_APInt(C))) {
  return getRangeViaSLT(EdgePred, *C, [&](const APInt &RHS) {
    APInt New = RHS << *ShAmtC;
    if ((New.ashr(*ShAmtC)) != New)
      return std::nullopt;
    return ConstantRange::getNonEmpty(APInt:.getSignedMinValue(BW), New + 1);
  });
}

Code is untested, but I think the general idea should work.

llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll

1257

Can we have a test where the RHS is not zero? We should also have one where the shift round trip precondition does not hold.

[aemerson]

llvm/lib/Analysis/LazyValueInfo.cpp

1152

I've been looking at this and despite my efforts, I can't get this approach to generate the same constant ranges as my patch. Here's your implementation with some fixes and modifications to make it work for sgt but I can't see how you can just use one implementation that covers the permutations needed.

static std::optional<ConstantRange>
getRangeViaSLT(CmpInst::Predicate Pred, APInt RHS,
               function_ref<std::optional<ConstantRange>(const APInt &)> Fn) {
  bool Invert = false;
  if (Pred == ICmpInst::ICMP_SGT || Pred == ICmpInst::ICMP_SGE) {
    Pred = ICmpInst::getInversePredicate(Pred);
    Invert = true;
  }
  if (Pred == ICmpInst::ICMP_SLE) {
    Pred = ICmpInst::ICMP_SLT;
    if (RHS.isMaxSignedValue())
      return std::nullopt;
    ++RHS;
  }
  assert(Pred == ICmpInst::ICMP_SLT && "Must be signed predicate");
  if (auto CR = Fn(RHS)) {
    dbgs() << "Init CR: " << *CR << "\n";
    if (Invert) {
      CR = ConstantRange(CR->getUpper(), ~CR->getLower());
    }
    return CR;
  }
  return std::nullopt;
}

// Recognize:
// icmp sgt (ashr X, ShAmtC), C --> icmp sgt X, ((C + 1) << ShAmtC) - 1
// and friends.
// Preconditions: (C != SIGNED_MAX) &&
//                ((C+1) << ShAmtC != SIGNED_MIN) &&
//                (((C+1) << ShAmtC) >> ShAmtC) == (C+1)
const APInt *ShAmtC;
if (CmpInst::isSigned(EdgePred) &&
    match(LHS, m_AShr(m_Specific(Val), m_APInt(ShAmtC))) &&
    match(RHS, m_APInt(C))) {
  auto CR = getRangeViaSLT(
      EdgePred, *C, [&](const APInt &RHS) -> std::optional<ConstantRange> {
        APInt New = RHS << *ShAmtC;
        if ((New.ashr(*ShAmtC)) != RHS)
          return std::nullopt;
        return ConstantRange::getNonEmpty(
            APInt::getSignedMinValue(C->getBitWidth()), New/* + 1*/);
      });
  if (CR) {
    dbgs() << ICI->getFunction()->getName() << ": " << *ICI << "\n";
    dbgs() << "CR is " << *CR << "\n";
    return ValueLatticeElement::getRange(*CR);
  }
}

I may have just misunderstood completely though. But from what I can tell, with the changes needed to make this correct, you'd just end up with code that's harder to follow and just as verbose as before?

[aemerson]

@nikic do you have any more thoughts?

[aemerson]

ping

[aemerson]

ping

[aemerson]

Committed in 9189822f2c2eca55f6e7adcd26e3f15d65797b76

llvm/lib/Analysis/LazyValueInfo.cpp

 #include "llvm/Analysis/ValueLattice.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/AssemblyAnnotationWriter.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
+#include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/Debug.h"
if (match(LHS, m_CombineOr(m_URem(m_Specific(Val), m_Value()),
       match(RHS, m_APInt(C))) {
     // Use the icmp region so we don't have to deal with different predicates.
     ConstantRange CR = ConstantRange::makeExactICmpRegion(EdgePred, *C);
     if (!CR.isEmptySet())
       return ValueLatticeElement::getRange(ConstantRange::getNonEmpty(
           CR.getUnsignedMin().zext(BitWidth), APInt(BitWidth, 0)));
   }
 
+  // Recognize:
+  // icmp sgt (ashr X, ShAmtC), C --> icmp sgt X, ((C + 1) << ShAmtC) - 1
+  // and friends.
+  // Preconditions: (C != SIGNED_MAX) &&
+  //                ((C+1) << ShAmtC != SIGNED_MIN) &&
+  //                (((C+1) << ShAmtC) >> ShAmtC) == (C+1)
+  const APInt *ShAmtC;
+  if (CmpInst::isSigned(EdgePred) &&
+      match(LHS, m_AShr(m_Specific(Val), m_APInt(ShAmtC))) &&
+      match(RHS, m_APInt(C))) {
+    APInt New = ((*C + 1) << *ShAmtC) - 1;
+    APInt MaxSigned = APInt::getSignedMaxValue(New.getBitWidth());
+    APInt MinSigned = APInt::getSignedMinValue(New.getBitWidth());
+    auto CheckPreConds = [&]() {
+      if (*C == MaxSigned)
+        return false;
+      APInt Shifted = (*C + 1) << *ShAmtC;
+      if (Shifted == MinSigned)
+        return false;
+      if ((Shifted.ashr(*ShAmtC)) != (*C + 1))
+        return false;
+      return true;
+    };
+    if (!CheckPreConds())
+      return ValueLatticeElement::getOverdefined();
+    APInt Lower, Upper;
+    switch (EdgePred) {
+    default:
+      llvm_unreachable("Unknown signed predicate!");
+    case ICmpInst::ICMP_SGT:
+      Lower = New + 1;
+      Upper = MaxSigned;
+      break;
+    case ICmpInst::ICMP_SLE:
+      Lower = MinSigned;
+      Upper = New + 1;
+      break;
+    case ICmpInst::ICMP_SGE:
+      Lower = New;
+      Upper = MaxSigned;
+      break;
+    case ICmpInst::ICMP_SLT:
+      Lower = MinSigned;
+      Upper = New;
+      break;
+    }
+
+    return ValueLatticeElement::getRange(
+        ConstantRange::getNonEmpty(Lower, Upper));
+  }
+
   return ValueLatticeElement::getOverdefined();
 }
 
 // Handle conditions of the form
 // extractvalue(op.with.overflow(%x, C), 1).
 static ValueLatticeElement getValueFromOverflowCondition(
     Value *Val, WithOverflowInst *WO, bool IsTrueDest) {
   // TODO: This only works with a constant RHS for now. We could also compute

llvm/test/Transforms/CorrelatedValuePropagation/icmp.ll

 ; CHECK-NEXT:    ret <2 x i1> [[CMP1]]
 ;
   %a2 = call <2 x i8> @llvm.umin.v2i8(<2 x i8> %a, <2 x i8> <i8 10, i8 10>)
   %b2 = call <2 x i8> @llvm.umax.v2i8(<2 x i8> %b, <2 x i8> <i8 11, i8 11>)
   %cmp1 = icmp ult <2 x i8> %a2, %b2
   ret <2 x i1> %cmp1
 }
 
+define void @ashr_sgt(i8 %x) {
+; CHECK-LABEL: @ashr_sgt(
+; CHECK-NEXT:    [[S:%.*]] = ashr i8 [[X:%.*]], 2
+; CHECK-NEXT:    [[C:%.*]] = icmp sgt i8 [[S]], 0
+; CHECK-NEXT:    br i1 [[C]], label [[IF:%.*]], label [[ELSE:%.*]]
+; CHECK:       if:
+; CHECK-NEXT:    call void @check1(i1 true)
+; CHECK-NEXT:    [[C3:%.*]] = icmp ugt i8 [[X]], 4
+; CHECK-NEXT:    call void @check1(i1 [[C3]])
+; CHECK-NEXT:    ret void
+; CHECK:       else:
+; CHECK-NEXT:    ret void
+;
+  %s = ashr i8 %x, 2
+  %c = icmp sgt i8 %s, 0
+  br i1 %c, label %if, label %else
+if:
+  %c2 = icmp sgt i8 %x, 3
+  call void @check1(i1 %c2)
+  %c3 = icmp sgt i8 %x, 4
+  call void @check1(i1 %c3)
+  ret void
+else:
+  ret void
+}
+
+define void @ashr_sge(i8 %x) {
+; CHECK-LABEL: @ashr_sge(
+; CHECK-NEXT:    [[S:%.*]] = ashr i8 [[X:%.*]], 2
+; CHECK-NEXT:    [[C:%.*]] = icmp sge i8 [[S]], 0
+; CHECK-NEXT:    br i1 [[C]], label [[IF:%.*]], label [[ELSE:%.*]]
+; CHECK:       if:
+; CHECK-NEXT:    call void @check1(i1 true)
+; CHECK-NEXT:    [[C3:%.*]] = icmp uge i8 [[X]], 4
+; CHECK-NEXT:    call void @check1(i1 [[C3]])
+; CHECK-NEXT:    ret void
+; CHECK:       else:
+; CHECK-NEXT:    ret void
+;
+  %s = ashr i8 %x, 2
+  %c = icmp sge i8 %s, 0
+  br i1 %c, label %if, label %else
+if:
+  %c2 = icmp sge i8 %x, 3
+  call void @check1(i1 %c2)
+  %c3 = icmp sge i8 %x, 4
+  call void @check1(i1 %c3)
+  ret void
+else:
+  ret void
+}
+
+define void @ashr_slt(i8 %x) {
+; CHECK-LABEL: @ashr_slt(
+; CHECK-NEXT:    [[S:%.*]] = ashr i8 [[X:%.*]], 2
+; CHECK-NEXT:    [[C:%.*]] = icmp slt i8 [[S]], 0
+; CHECK-NEXT:    br i1 [[C]], label [[IF:%.*]], label [[ELSE:%.*]]
+; CHECK:       if:
+; CHECK-NEXT:    call void @check1(i1 true)
+; CHECK-NEXT:    [[C3:%.*]] = icmp slt i8 [[X]], 2
+; CHECK-NEXT:    call void @check1(i1 [[C3]])
+; CHECK-NEXT:    ret void
+; CHECK:       else:
+; CHECK-NEXT:    ret void
+;
+  %s = ashr i8 %x, 2
+  %c = icmp slt i8 %s, 0
+  br i1 %c, label %if, label %else
+if:
+  %c2 = icmp slt i8 %x, 3
+  call void @check1(i1 %c2)
+  %c3 = icmp slt i8 %x, 2
+  call void @check1(i1 %c3)
+  ret void
+else:
+  ret void
+}
+
+define void @ashr_sle(i8 %x) {
+; CHECK-LABEL: @ashr_sle(
+; CHECK-NEXT:    [[S:%.*]] = ashr i8 [[X:%.*]], 2
+; CHECK-NEXT:    [[C:%.*]] = icmp sle i8 [[S]], 0
+; CHECK-NEXT:    br i1 [[C]], label [[IF:%.*]], label [[ELSE:%.*]]
+; CHECK:       if:
+; CHECK-NEXT:    call void @check1(i1 true)
+; CHECK-NEXT:    [[C3:%.*]] = icmp sle i8 [[X]], 2
+; CHECK-NEXT:    call void @check1(i1 [[C3]])
+; CHECK-NEXT:    ret void
+; CHECK:       else:
+; CHECK-NEXT:    ret void
+;
+  %s = ashr i8 %x, 2
+  %c = icmp sle i8 %s, 0
+  br i1 %c, label %if, label %else
+if:
+  %c2 = icmp sle i8 %x, 3
+  call void @check1(i1 %c2)
+  %c3 = icmp sle i8 %x, 2
+  call void @check1(i1 %c3)
+  ret void
+else:
+  ret void
+}
+
 declare i8 @llvm.umin.i8(i8, i8)
 declare i8 @llvm.umax.i8(i8, i8)
 declare <2 x i8> @llvm.umin.v2i8(<2 x i8>, <2 x i8>)
 declare <2 x i8> @llvm.umax.v2i8(<2 x i8>, <2 x i8>)
 
 attributes #4 = { noreturn }