[WIP][ConstantRange] Add support for multiply and left shift with nowrap flags

[reames]

Summary

This patch adds support for multiply and left shift variants which use NoWrap flags to restrict the range returned. It directly follows the pattern already taken for add/sub, and extends it to mul/shl.

This is a WIP for two major reasons, both of which I'd appreciate some input on.

1. I don't understand the bits in add/sub about being required to return empty ranges on definite overflow. I can see that from a precision perspective, but since both the 2s complement and saturating variants are valid implementations of the no-wrap behavior, it seems like intersecting them should always be valid. Am I right here? Or is my code subtly wrong?
2. Suggestions on how to test this? I glanced at the ConstantRangeTest file, and am looking for specific guidance on which exhaustive test harness to use. This file has gotten quite complicated.

Once complete, I plan to use this in SCEV to help refine the ranges returned for multiplies.

[nikic]

I don't understand the bits in add/sub about being required to return empty ranges on definite overflow. I can see that from a precision perspective, but since both the 2s complement and saturating variants are valid implementations of the no-wrap behavior, it seems like intersecting them should always be valid. Am I right here? Or is my code subtly wrong?

There is no requirement to return an empty range. I believe the existing code is only pedantic about this so that tests can prove both correctness and optimality.

Suggestions on how to test this? I glanced at the ConstantRangeTest file, and am looking for specific guidance on which exhaustive test harness to use. This file has gotten quite complicated.

The generic testing framework is TestBinaryOpExhaustive. For the second callback, return an Optional<APInt> only if no overflow occurs. For the third callback, you probably want PreferSmallestUnsigned or PreferSmallestSigned depending on the PreferredRangeType. The fourth one depends on whether your implementation provides optimal results, at least for some subsets.

llvm/include/llvm/IR/ConstantRange.h

public:
   ConstantRange subWithNoWrap(const ConstantRange &Other, unsigned NoWrapKind,
                               PreferredRangeType RangeType = Smallest) const;
 
   /// Return a new range representing the possible values resulting
   /// from a multiplication of a value in this range and a value in \p Other,
   /// treating both this and \p Other as unsigned ranges.
   ConstantRange multiply(const ConstantRange &Other) const;
 
+  /// Return a new range representing the possible values resulting
+  /// from a multiplication with wrap type \p NoWrapKind of a value in this
+  /// range and a value in \p Other.
+  /// If the result range is disjoint, the preferred range is determined by the
+  /// \p PreferredRangeType.
+  ConstantRange mulWithNoWrap(const ConstantRange &Other, unsigned NoWrapKind,
+                              PreferredRangeType RangeType = Smallest) const;
+
   /// Return range of possible values for a signed multiplication of this and
   /// \p Other. However, if overflow is possible always return a full range
   /// rather than trying to determine a more precise result.
   ConstantRange smul_fast(const ConstantRange &Other) const;
 
   /// Return a new range representing the possible values resulting
   /// from a signed maximum of a value in this range and a value in \p Other.
   ConstantRange smax(const ConstantRange &Other) const;
public:
   /// from a binary-xor of a value in this range by a value in \p Other.
   ConstantRange binaryXor(const ConstantRange &Other) const;
 
   /// Return a new range representing the possible values resulting
   /// from a left shift of a value in this range by a value in \p Other.
   /// TODO: This isn't fully implemented yet.
   ConstantRange shl(const ConstantRange &Other) const;
 
+  /// Return a new range representing the possible values resulting
+  /// from a left shift with wrap type \p NoWrapKind of a value in this
+  /// range and a value in \p Other.
+  /// If the result range is disjoint, the preferred range is determined by the
+  /// \p PreferredRangeType.
+  ConstantRange shlWithNoWrap(const ConstantRange &Other, unsigned NoWrapKind,
+                              PreferredRangeType RangeType = Smallest) const;
+
+
   /// Return a new range representing the possible values resulting from a
   /// logical right shift of a value in this range and a value in \p Other.
   ConstantRange lshr(const ConstantRange &Other) const;
 
   /// Return a new range representing the possible values resulting from a
   /// arithmetic right shift of a value in this range and a value in \p Other.
   ConstantRange ashr(const ConstantRange &Other) const;
 

llvm/lib/IR/ConstantRange.cpp

ConstantRange ConstantRange::overflowingBinaryOp(Instruction::BinaryOps BinOp,
                                                  unsigned NoWrapKind) const {
   assert(Instruction::isBinaryOp(BinOp) && "Binary operators only!");
 
   switch (BinOp) {
   case Instruction::Add:
     return addWithNoWrap(Other, NoWrapKind);
   case Instruction::Sub:
     return subWithNoWrap(Other, NoWrapKind);
+  case Instruction::Mul:
+    return mulWithNoWrap(Other, NoWrapKind);
+  case Instruction::Shl:
+    return shlWithNoWrap(Other, NoWrapKind);
   default:
     // Don't know about this Overflowing Binary Operation.
     // Conservatively fallback to plain binop handling.
     return binaryOp(BinOp, Other);
   }
 }
 
 bool ConstantRange::isIntrinsicSupported(Intrinsic::ID IntrinsicID) {
auto L = {this_min * Other_min, this_min * Other_max,
             this_max * Other_min, this_max * Other_max};
   auto Compare = [](const APInt &A, const APInt &B) { return A.slt(B); };
   ConstantRange Result_sext(std::min(L, Compare), std::max(L, Compare) + 1);
   ConstantRange SR = Result_sext.truncate(getBitWidth());
 
   return UR.isSizeStrictlySmallerThan(SR) ? UR : SR;
 }
 
+ConstantRange ConstantRange::mulWithNoWrap(const ConstantRange &Other,
+                                           unsigned NoWrapKind,
+                                           PreferredRangeType RangeType) const {
+  // Calculate the range for "X * Y" which is guaranteed not to wrap(overflow).
+  // (X is from this, and Y is from Other)
+  if (isEmptySet() || Other.isEmptySet())
+    return getEmpty();
+  if (isFullSet() && Other.isFullSet())
+    return getFull();
+
+  using OBO = OverflowingBinaryOperator;
+  ConstantRange Result = multiply(Other);
+
+  if (NoWrapKind & OBO::NoSignedWrap)
+    Result = Result.intersectWith(smul_sat(Other), RangeType);
+
+  if (NoWrapKind & OBO::NoUnsignedWrap)
+    Result = Result.intersectWith(umul_sat(Other), RangeType);
+
+  return Result;
+}
+
 ConstantRange ConstantRange::smul_fast(const ConstantRange &Other) const {
   if (isEmptySet() || Other.isEmptySet())
     return getEmpty();
 
   APInt Min = getSignedMin();
   APInt Max = getSignedMax();
   APInt OtherMin = Other.getSignedMin();
   APInt OtherMax = Other.getSignedMax();
ConstantRange::shl(const ConstantRange &Other) const {
   // FIXME: implement the other tricky cases
 
   Min <<= Other.getUnsignedMin();
   Max <<= OtherMax;
 
   return ConstantRange::getNonEmpty(std::move(Min), std::move(Max) + 1);
 }
 
+ConstantRange ConstantRange::shlWithNoWrap(const ConstantRange &Other,
+                                           unsigned NoWrapKind,
+                                           PreferredRangeType RangeType) const {
+  // Calculate the range for "shl X Y" which is guaranteed not to wrap(overflow).
+  // (X is from this, and Y is from Other)
+  if (isEmptySet() || Other.isEmptySet())
+    return getEmpty();
+  if (isFullSet() && Other.isFullSet())
+    return getFull();
+
+  using OBO = OverflowingBinaryOperator;
+  ConstantRange Result = shl(Other);
+
+  if (NoWrapKind & OBO::NoSignedWrap)
+    Result = Result.intersectWith(sshl_sat(Other), RangeType);
+
+  if (NoWrapKind & OBO::NoUnsignedWrap)
+    Result = Result.intersectWith(ushl_sat(Other), RangeType);
+
+  return Result;
+}
+
 ConstantRange
 ConstantRange::lshr(const ConstantRange &Other) const {
   if (isEmptySet() || Other.isEmptySet())
     return getEmpty();
 
   APInt max = getUnsignedMax().lshr(Other.getUnsignedMin()) + 1;
   APInt min = getUnsignedMin().lshr(Other.getUnsignedMax());
   return getNonEmpty(std::move(min), std::move(max));

llvm/test/Transforms/CorrelatedValuePropagation/mul.ll

exit:
   ret i8 0
 }
 
 define i1 @nuw_range1(i8 %b) {
 ; CHECK-LABEL: @nuw_range1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[C:%.*]] = add nuw nsw i8 [[B:%.*]], 1
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i8 [[C]], 4
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[MUL]], 0
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   %c = add nuw nsw i8 %b, 1
   %mul = mul nuw i8 %c, 4
   %cmp = icmp eq i8 %mul, 0
   ret i1 %cmp
 }
 
 define i1 @nuw_range2(i8 %b) {
 ; CHECK-LABEL: @nuw_range2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[C:%.*]] = add nuw nsw i8 [[B:%.*]], 3
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i8 [[C]], 4
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[MUL]], 2
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   %c = add nuw nsw i8 %b, 3
   %mul = mul nuw i8 %c, 4
   %cmp = icmp ult i8 %mul, 2
   ret i1 %cmp
 }
 

llvm/test/Transforms/CorrelatedValuePropagation/shl.ll

 
 define i8 @test4(i8 %a, i8 %b) {
 ; CHECK-LABEL: @test4(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i8 [[B:%.*]], 7
 ; CHECK-NEXT:    br i1 [[CMP]], label [[BB:%.*]], label [[EXIT:%.*]]
 ; CHECK:       bb:
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw nsw i8 [[A:%.*]], [[B]]
-; CHECK-NEXT:    ret i8 [[SHL]]
+; CHECK-NEXT:    ret i8 -1
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i8 0
 ;
 entry:
   %cmp = icmp ugt i8 %b, 7
   br i1 %cmp, label %bb, label %exit
 
 bb:
exit:
   ret i8 0
 }
 
 define i1 @nuw_range1(i8 %b) {
 ; CHECK-LABEL: @nuw_range1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[C:%.*]] = add nuw nsw i8 [[B:%.*]], 1
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw i8 [[C]], 2
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i8 [[SHL]], 0
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   %c = add nuw nsw i8 %b, 1
   %shl = shl nuw i8 %c, 2
   %cmp = icmp eq i8 %shl, 0
   ret i1 %cmp
 }
 
 define i1 @nuw_range2(i8 %b) {
 ; CHECK-LABEL: @nuw_range2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[C:%.*]] = add nuw nsw i8 [[B:%.*]], 3
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw i8 [[C]], 2
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[SHL]], 2
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
 entry:
   %c = add nuw nsw i8 %b, 3
   %shl = shl nuw i8 %c, 2
   %cmp = icmp ult i8 %shl, 2
   ret i1 %cmp
 }