[ConstantRange] Handle `Intrinsic::ctlz`

[antoniofrighetto]

Summary

Introduce support for ctlz intrinsic, including exhaustive testing. Among other things, LVI may now be able to propagate information about constant ranges lattice values. This change makes sure to provide a precise result where possible, or a more conservative one (e.g., should the first argument value be zero and is_zero_poison true), but always correct.

[arsenm]

llvm/test/Analysis/LazyValueAnalysis/lvi-for-ctlz.ll

32 ↗

(On Diff #490893)

Test looks pretty small for all the different paths in the patch. You don't test the i1 false case for example

[antoniofrighetto]

llvm/test/Analysis/LazyValueAnalysis/lvi-for-ctlz.ll

32 ↗

(On Diff #490893)

For the second argument i1 false the result should always be more precise than the current one, but I'll make sure to test it soon as well.

[nikic]

This needs to be implemented inside ConstantRange and tested exhaustively inside ConstantRangeTest. See the handling of Intrinsic::abs as a possible sample to follow (it also has a semantics-controlling argument).

[antoniofrighetto]

Moved the implementation in ConstantRange.cpp, and tested exhaustively. For the sake of completeness and rigor, I left the test in LVI, further improved: it shouldn’t harm, and there are similar unit tests in ConstantRangeTest.cpp which are also in test/Analysis.

[nikic]

The unit test currently fails both for correctness and optimality checks:

/var/lib/buildkite-agent/builds/llvm-project/llvm/unittests/IR/ConstantRangeTest.cpp:137
Value of: rangeContains(CR, APInt(BitWidth, Elem), Inputs)
  Actual: false (empty-set does not contain 0 for inputs: full-set, )
Expected: true
/var/lib/buildkite-agent/builds/llvm-project/llvm/unittests/IR/ConstantRangeTest.cpp:137
Value of: rangeContains(CR, APInt(BitWidth, Elem), Inputs)
  Actual: false (empty-set does not contain -1 for inputs: full-set, )
Expected: true
/var/lib/buildkite-agent/builds/llvm-project/llvm/unittests/IR/ConstantRangeTest.cpp:145
Value of: CR.isEmptySet()
  Actual: false
Expected: true
/var/lib/buildkite-agent/builds/llvm-project/llvm/unittests/IR/ConstantRangeTest.cpp:177
Value of: NotPreferred(PossibleCR)
  Actual: false (Inputs = full-set, CR = full-set, BetterCR = [0,-1))
Expected: true
/var/lib/buildkite-agent/builds/llvm-project/llvm/unittests/IR/ConstantRangeTest.cpp:145
Value of: CR.isEmptySet()
  Actual: false
Expected: true
/var/lib/buildkite-agent/builds/llvm-project/llvm/unittests/IR/ConstantRangeTest.cpp:177
Value of: NotPreferred(PossibleCR)
  Actual: false (Inputs = [0,2), CR = full-set, BetterCR = [3,4))
Expected: true
/var/lib/buildkite-agent/builds/llvm-project/llvm/unittests/IR/ConstantRangeTest.cpp:177
Value of: NotPreferred(PossibleCR)
  Actual: false (Inputs = [0,3), CR = full-set, BetterCR = [2,4))
Expected: true
/var/lib/buildkite-agent/builds/llvm-project/llvm/unittests/IR/ConstantRangeTest.cpp:177
Value of: NotPreferred(PossibleCR)
  Actual: false (Inputs = [0,4), CR = full-set, BetterCR = [2,4))
Expected: true
/var/lib/buildkite-agent/builds/llvm-project/llvm/unittests/IR/ConstantRangeTest.cpp:177
Value of: NotPreferred(PossibleCR)
  Actual: false (Inputs = [0,5), CR = full-set, BetterCR = [1,4))
Expected: true
[...]

You can run the unit tests using ninja -Cbuild check-llvm-unit. Or ninja -Cbuild IRTests && build/unittests/IR/IRTests to run just this test suite.

llvm/include/llvm/ADT/APInt.h

1741 ↗

(On Diff #492170)

I'd rather not have this API on APInt. Getting leading zeros as an APInt is pretty unusual.

llvm/lib/IR/ConstantRange.cpp

1684

Constant ranges are understood modulo poison. The value is in the range or it is poison. As such, if ZeroIsPoison is set, we should return a better range that doesn't have to account for zero input, not a worse range.

1700

Can Lower be larger than Upper for a non-wrapped set?

llvm/test/Analysis/LazyValueAnalysis/lvi-for-ctlz.ll

1 ↗

(On Diff #492170)

Having an IR test is fine, but please do not test LVI debug output. Just check the resulting IR change using update_test_checks.py.

[antoniofrighetto]

@nikic, thank you, I'm having a look at these now. Just a few considerations.

1. I'm not entirely sure about the unit test failures: should the returned range simply not include zero, if the incoming interval may contain zero and ZeroIsPoison is true, as you noted? I reckoned that being more conservative (unknown/overdefined) could be better here. Same applies for isEmptySet, we may want to return overdefined, I guess (I'm looking at Value of: CR.isEmptySet(), Actual: false, Expected: true)?.
2. If we don't want to have ctlz API on APInt, could it be fine to overload TestUnaryOpExhaustive so that it accepts a callable reference to std::optional<unsigned>(const APInt &) too?
3. range.ll now fails as the metadata's range is no longer being used now. I'm gradually realizing that this needs to be handled properly in CVP as well. Could it be fine to revert the ret i1 true to this change, considering that the range is now [0,16] (verification here) for this time, and postpone the handling in CVP for a separate dedicated patch?

[nikic]

In D142234#4082670, @antoniofrighetto wrote:

@nikic, thank you, I'm having a look at these now. Just a few considerations.

1. I'm not entirely sure about the unit test failures: should the returned range simply not include zero, if the incoming interval may contain zero and ZeroIsPoison is true, as you noted? I reckoned that being more conservative (unknown/overdefined) could be better here. Same applies for isEmptySet, we may want to return overdefined, I guess (I'm looking at Value of: CR.isEmptySet(), Actual: false, Expected: true)?.

It's not necessary to make any conservative choices here. If ZeroIsPoison is set, you can assume that the input range does not contain zero, e.g. by intersecting it out.

2. If we don't want to have ctlz API on APInt, could it be fine to overload TestUnaryOpExhaustive so that it accepts a callable reference to std::optional<unsigned>(const APInt &) too?

I think we should just create the APInt in the two places where it's needed.

3. range.ll now fails as the metadata's range is no longer being used now. I'm gradually realizing that this needs to be handled properly in CVP as well. Could it be fine to revert the ret i1 true to this change, considering that the range is now [0,16] (verification here) for this time, and postpone the handling in CVP for a separate dedicated patch?

This issue should already be fixed by https://github.com/llvm/llvm-project/commit/2e9bc1b8614c9422573cf2f4728525787b0cb0cb.

[antoniofrighetto]

llvm/lib/IR/ConstantRange.cpp

1700

It seems it is possible with negative numbers as well ([0,-7)), non-wrapped set.

[nikic]

This looks on the right track now.

llvm/lib/IR/ConstantRange.cpp

1737

This looks basically fine, but a bit overly complicated. This should be sufficient:

ConstantRange ConstantRange::ctlz(bool ZeroIsPoison) const {
  if (isEmptySet())
    return getEmpty();

  APInt Zero = APInt::getZero(getBitWidth());
  if (ZeroIsPoison && contains(Zero)) {
    // ZeroIsPoison is set, and zero is contained. We discern three cases, in
    // which a zero can appear:
    // 1) Lower is zero, handling cases of kind [0, 1), [0, 2), etc.
    // 2) Upper is zero, wrapped set, handling cases of kind [3, 0], etc.
    // 3) Zero contained in a wrapped set, e.g., [3, 2), [3, 1), etc.
    
    if (getLower().isZero()) {
      if ((getUpper() - 1).isZero()) {
        // We have in input interval of kind [0, 1). In this case we cannot
        // really help but return empty-set.
        return getEmpty();
      }

      // Compute the resulting range by excluding zero from Lower.
      return ConstantRange(
          APInt(getBitWidth(), (getUpper() - 1).countLeadingZeros()), 
          APInt(getBitWidth(), (getLower() + 1).countLeadingZeros() + 1));
    } else if ((getUpper() - 1).isZero()) {
      // Compute the resulting range by excluding zero from Upper.
      return ConstantRange(
          Zero, APInt(getBitWidth(), getLower().countLeadingZeros() + 1));
    } else { 
      return ConstantRange(Zero, APInt(getBitWidth(), getBitWidth()));
    }   
  } 
  
  // Zero is either safe or not in the range. The output range is composed by
  // the result of countLeadingZero of the two extremes.
  return getNonEmpty(
      APInt(getBitWidth(), getUnsignedMax().countLeadingZeros()), 
      APInt(getBitWidth(), getUnsignedMin().countLeadingZeros() + 1));
}

llvm/test/Analysis/LazyValueAnalysis/lvi-for-ctlz.ll

1 ↗

(On Diff #492170)

I don't think that as written, these tests really test anything. There needs to be a comparison involving the ctlz that can be folded away, or similar.

[antoniofrighetto]

llvm/test/Analysis/LazyValueAnalysis/lvi-for-ctlz.ll

1 ↗

(On Diff #492170)

I feel like a suitable test could be the following one:

int lol(int b) {
    if (b < 65536) {
        int n = __builtin_clz(b);
        if (n < 8)
    	    return 0;
        else
            return 1;
   }
   return 2;
}

which now could be simplified in a return (b < 65536) ? 1 : 2;. However, as far as I understand this, CVP needs to be extended as well in order to be able to obtain the above constant folding. Is that correct?

[StephenFan]

llvm/lib/IR/ConstantRange.cpp

1686

IIUC, is [3, 0] the same as the [3, 1) in case 3?

llvm/test/Analysis/LazyValueAnalysis/lvi-for-ctlz.ll

1 ↗

(On Diff #492170)

Is the simplification to return (b < 65536) ? 1:2; correct? Since if b is negative, clz returns 0.

[antoniofrighetto]

llvm/lib/IR/ConstantRange.cpp

1686

Correct.

llvm/test/Analysis/LazyValueAnalysis/lvi-for-ctlz.ll

1 ↗

(On Diff #492170)

True, sorry, I originally intended the first argument to be unsigned, whose optimization seems to occur successfully in this case.

[antoniofrighetto]

llvm/test/Analysis/LazyValueAnalysis/lvi-for-ctlz.ll

1 ↗

(On Diff #492170)

@nikic, perhaps a better example could be the following one? That would still require a change in CVP if I'm not mistaken though, correct?

int test_ctlz(int b) {
    if (b < 65536) {
        int n = __builtin_clz(b);
        if (n < 8 && n > 2)
    	    return 0;
        else
            return 1;
   }
   return 2;
}

[nikic]

llvm/test/Analysis/LazyValueAnalysis/lvi-for-ctlz.ll

1 ↗

(On Diff #492170)

Wouldn't something like this work as a test?

int test_ctlz(unsigned b) {
    if (b < 65536) {
        unsigned n = __builtin_clz(b);
        return n >= 16; // Fold to 1
   }
   return 2;
}

[nikic]

Can you please rebase over https://github.com/llvm/llvm-project/commit/14fcdd7f9d7b3973661efc5a426da18e077155bf? I think that should work as a CVP test for this functionality, if I got the numbers right.

[antoniofrighetto]

@nikic, rebased and removed the old test. Might be wrong, but it seems like O2 is already handling the test case you suggested; whereas the latter one I previously suggested (the one with if (n < 8 && n > 2) return 0;, which can never happen) does not seem to be reduced to a icmp, select, ret as of now.

[nikic]

Can you please rerun update_test_checks.py on the CorrelatedPropagation/range.ll test? There should be some diffs there.

In D142234#4128565, @antoniofrighetto wrote:

@nikic, rebased and removed the old test. Might be wrong, but it seems like O2 is already handling the test case you suggested; whereas the latter one I previously suggested (the one with if (n < 8 && n > 2) return 0;, which can never happen) does not seem to be reduced to a icmp, select, ret as of now.

Right. It's not really important that individual pass tests aren't folded by -O2. One construct more complex cases that LVI can handle that InstCombine can't. It's just more straightforward to test the basic case.

[nikic]

llvm/test/Transforms/CorrelatedValuePropagation/range.ll

954

Did you use the right build to produce this? This is not the expected diff, and it still fails in pre-merge checks.

[nikic]

LGTM

[antoniofrighetto]

llvm/test/Transforms/CorrelatedValuePropagation/range.ll

954

Sorry, my bad, confused builds and must have accidentally passed the wrong --opt-bin path. Checked twice, we should be good now.

llvm/include/llvm/IR/ConstantRange.h

public:
   /// Return a new range that is the logical not of the current set.
   ConstantRange inverse() const;
 
   /// Calculate absolute value range. If the original range contains signed
   /// min, then the resulting range will contain signed min if and only if
   /// \p IntMinIsPoison is false.
   ConstantRange abs(bool IntMinIsPoison = false) const;
 
+  /// Calculate ctlz range. If \p ZeroIsPoison is set, the range is computed
+  /// ignoring a possible zero value contained in the input range.
+  ConstantRange ctlz(bool ZeroIsPoison = false) const;
+
   /// Represents whether an operation on the given constant range is known to
   /// always or never overflow.
   enum class OverflowResult {
     /// Always overflows in the direction of signed/unsigned min value.
     AlwaysOverflowsLow,
     /// Always overflows in the direction of signed/unsigned max value.
     AlwaysOverflowsHigh,
     /// May or may not overflow.

llvm/lib/IR/ConstantRange.cpp

bool ConstantRange::isIntrinsicSupported(Intrinsic::ID IntrinsicID) {
   case Intrinsic::usub_sat:
   case Intrinsic::sadd_sat:
   case Intrinsic::ssub_sat:
   case Intrinsic::umin:
   case Intrinsic::umax:
   case Intrinsic::smin:
   case Intrinsic::smax:
   case Intrinsic::abs:
+  case Intrinsic::ctlz:
     return true;
   default:
     return false;
   }
 }
 
 ConstantRange ConstantRange::intrinsic(Intrinsic::ID IntrinsicID,
                                        ArrayRef<ConstantRange> Ops) {
ConstantRange ConstantRange::intrinsic(Intrinsic::ID IntrinsicID,
   case Intrinsic::smax:
     return Ops[0].smax(Ops[1]);
   case Intrinsic::abs: {
     const APInt *IntMinIsPoison = Ops[1].getSingleElement();
     assert(IntMinIsPoison && "Must be known (immarg)");
     assert(IntMinIsPoison->getBitWidth() == 1 && "Must be boolean");
     return Ops[0].abs(IntMinIsPoison->getBoolValue());
   }
+  case Intrinsic::ctlz: {
+    const APInt *ZeroIsPoison = Ops[1].getSingleElement();
+    assert(ZeroIsPoison && "Must be known (immarg)");
+    assert(ZeroIsPoison->getBitWidth() == 1 && "Must be boolean");
+    return Ops[0].ctlz(ZeroIsPoison->getBoolValue());
+  }
   default:
     assert(!isIntrinsicSupported(IntrinsicID) && "Shouldn't be supported");
     llvm_unreachable("Unsupported intrinsic");
   }
 }
 
 ConstantRange
 ConstantRange::add(const ConstantRange &Other) const {
ConstantRange ConstantRange::abs(bool IntMinIsPoison) const {
   if (SMax.isNegative())
     return ConstantRange(-SMax, -SMin + 1);
 
   // Range crosses zero.
   return ConstantRange::getNonEmpty(APInt::getZero(getBitWidth()),
                                     APIntOps::umax(-SMin, SMax) + 1);
 }
 
+ConstantRange ConstantRange::ctlz(bool ZeroIsPoison) const {
+  if (isEmptySet())
+    return getEmpty();
+
+  APInt Zero = APInt::getZero(getBitWidth());
+  if (ZeroIsPoison && contains(Zero)) {
+    // ZeroIsPoison is set, and zero is contained. We discern three cases, in
+    // which a zero can appear:
+    // 1) Lower is zero, handling cases of kind [0, 1), [0, 2), etc.
+    // 2) Upper is zero, wrapped set, handling cases of kind [3, 0], etc.
+    // 3) Zero contained in a wrapped set, e.g., [3, 2), [3, 1), etc.
+
+    if (getLower().isZero()) {
+      if ((getUpper() - 1).isZero()) {
+        // We have in input interval of kind [0, 1). In this case we cannot
+        // really help but return empty-set.
+        return getEmpty();
+      }
+
+      // Compute the resulting range by excluding zero from Lower.
+      return ConstantRange(
+          APInt(getBitWidth(), (getUpper() - 1).countLeadingZeros()),
+          APInt(getBitWidth(), (getLower() + 1).countLeadingZeros() + 1));
+    } else if ((getUpper() - 1).isZero()) {
+      // Compute the resulting range by excluding zero from Upper.
+      return ConstantRange(
+          Zero, APInt(getBitWidth(), getLower().countLeadingZeros() + 1));
+    } else {
+      return ConstantRange(Zero, APInt(getBitWidth(), getBitWidth()));
+    }
+  }
+
+  // Zero is either safe or not in the range. The output range is composed by
+  // the result of countLeadingZero of the two extremes.
+  return getNonEmpty(
+      APInt(getBitWidth(), getUnsignedMax().countLeadingZeros()),
+      APInt(getBitWidth(), getUnsignedMin().countLeadingZeros() + 1));
+}
+
 ConstantRange::OverflowResult ConstantRange::unsignedAddMayOverflow(
     const ConstantRange &Other) const {
   if (isEmptySet() || Other.isEmptySet())
     return OverflowResult::MayOverflow;
 
   APInt Min = getUnsignedMin(), Max = getUnsignedMax();
   APInt OtherMin = Other.getUnsignedMin(), OtherMax = Other.getUnsignedMax();
 
   // a u+ b overflows high iff a u> ~b.
   if (Min.ugt(~OtherMin))
     return OverflowResult::AlwaysOverflowsHigh;
   if (Max.ugt(~OtherMax))
     return OverflowResult::MayOverflow;
   return OverflowResult::NeverOverflows;
 }
 
 ConstantRange::OverflowResult ConstantRange::signedAddMayOverflow(
     const ConstantRange &Other) const {
   if (isEmptySet() || Other.isEmptySet())
     return OverflowResult::MayOverflow;
 
   APInt Min = getSignedMin(), Max = getSignedMax();
   APInt OtherMin = Other.getSignedMin(), OtherMax = Other.getSignedMax();
 
   APInt SignedMin = APInt::getSignedMinValue(getBitWidth());
   APInt SignedMax = APInt::getSignedMaxValue(getBitWidth());
 
   // a s+ b overflows high iff a s>=0 && b s>= 0 && a s> smax - b.
   // a s+ b overflows low iff a s< 0 && b s< 0 && a s< smin - b.
   if (Min.isNonNegative() && OtherMin.isNonNegative() &&

llvm/test/Transforms/CorrelatedValuePropagation/range.ll

;
   ret i1 %res
 }
 
 define i1 @abs_with_range_metadata(i16 %x) {
 ; CHECK-LABEL: @abs_with_range_metadata(
 ; CHECK-NEXT:    [[ABS:%.*]] = call i16 @llvm.abs.i16(i16 [[X:%.*]], i1 false), !range [[RNG5]]
 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i16 [[ABS]] to i8
 ; CHECK-NEXT:    ret i1 true
 ;
   %abs = call i16 @llvm.abs.i16(i16 %x, i1 false), !range !{i16 0, i16 8}
   %trunc = trunc i16 %abs to i8
   %res = icmp ult i8 %trunc, 8
   ret i1 %res
 }
 
 define i1 @ctlz_fold(i16 %x) {
 ; CHECK-LABEL: @ctlz_fold(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i16 [[X:%.*]], 256
 ; CHECK-NEXT:    br i1 [[CMP]], label [[IF:%.*]], label [[ELSE:%.*]]
 ; CHECK:       if:
 ; CHECK-NEXT:    [[CTLZ:%.*]] = call i16 @llvm.ctlz.i16(i16 [[X]], i1 false)
-; CHECK-NEXT:    [[RES:%.*]] = icmp uge i16 [[CTLZ]], 8
-; CHECK-NEXT:    ret i1 [[RES]]
+; CHECK-NEXT:    ret i1 true
 ; CHECK:       else:
 ; CHECK-NEXT:    [[CTLZ2:%.*]] = call i16 @llvm.ctlz.i16(i16 [[X]], i1 false)
-; CHECK-NEXT:    [[RES2:%.*]] = icmp ult i16 [[CTLZ2]], 8
-; CHECK-NEXT:    ret i1 [[RES2]]
+; CHECK-NEXT:    ret i1 true
 ;
   %cmp = icmp ult i16 %x, 256
   br i1 %cmp, label %if, label %else
 
 if:
   %ctlz = call i16 @llvm.ctlz.i16(i16 %x, i1 false)
   %res = icmp uge i16 %ctlz, 8
   ret i1 %res

llvm/unittests/IR/ConstantRangeTest.cpp

TestUnaryOpExhaustive(
       [](const ConstantRange &CR) { return CR.abs(/*IntMinIsPoison=*/true); },
       [](const APInt &N) -> std::optional<APInt> {
         if (N.isMinSignedValue())
           return std::nullopt;
         return N.abs();
       });
 }
 
+TEST_F(ConstantRangeTest, Ctlz) {
+  TestUnaryOpExhaustive([](const ConstantRange &CR) { return CR.ctlz(); },
+                        [](const APInt &N) {
+                          return APInt(N.getBitWidth(), N.countLeadingZeros());
+                        });
+
+  TestUnaryOpExhaustive(
+      [](const ConstantRange &CR) { return CR.ctlz(/*ZeroIsPoison=*/true); },
+      [](const APInt &N) -> std::optional<APInt> {
+        if (N.isZero())
+          return std::nullopt;
+        return APInt(N.getBitWidth(), N.countLeadingZeros());
+      });
+}
+
 TEST_F(ConstantRangeTest, castOps) {
   ConstantRange A(APInt(16, 66), APInt(16, 128));
   ConstantRange FpToI8 = A.castOp(Instruction::FPToSI, 8);
   EXPECT_EQ(8u, FpToI8.getBitWidth());
   EXPECT_TRUE(FpToI8.isFullSet());
 
   ConstantRange FpToI16 = A.castOp(Instruction::FPToSI, 16);
   EXPECT_EQ(16u, FpToI16.getBitWidth());