[InstCombine] Don't crash on out of bounds shifts

[igor-laevsky]

Summary

This is branched from the D40390.

We can't crash instcombine when we encounter shift by unreasonable amount. Even though it's an undefined behaviour, compiler crash is incorrect.

[spatel]

We should stub this out sooner and in one place. I think we can extend what we did in D38637:

Index: lib/Analysis/ValueTracking.cpp
===================================================================
--- lib/Analysis/ValueTracking.cpp	(revision 319435)
+++ lib/Analysis/ValueTracking.cpp	(working copy)
@@ -800,8 +800,14 @@
   unsigned BitWidth = Known.getBitWidth();
 
   if (auto *SA = dyn_cast<ConstantInt>(I->getOperand(1))) {
+    if (SA->getZExtValue() >= BitWidth) {
+      // This shift produces poison because the shift amount is too big. We can
+      // return anything we want. Choose 0 for the best folding opportunity.
+      Known.setAllZero();
+      return;
+    }
+
     unsigned ShiftAmt = SA->getLimitedValue(BitWidth-1);
-
     computeKnownBits(I->getOperand(0), Known, Depth + 1, Q);
     Known.Zero = KZF(Known.Zero, ShiftAmt);
     Known.One  = KOF(Known.One, ShiftAmt);

[igor-laevsky]

Update test case

[igor-laevsky]

I see your point. Unfortunately I can't find any good place to common this out. Problem is that code is structured as a consecutive matching of unrelated patterns which may sometimes contain invalid shift operations. But I don't know the code base enough so maybe I'm missing something. Any advice?

[craig.topper]

Have you checked for compile time impact?

[craig.topper]

Oops wrong review. Ignore that comment

[spatel]

In D40649#942228, @igor-laevsky wrote:

I see your point. Unfortunately I can't find any good place to common this out. Problem is that code is structured as a consecutive matching of unrelated patterns which may sometimes contain invalid shift operations. But I don't know the code base enough so maybe I'm missing something. Any advice?

Not sure if you saw my earlier patch draft. Here's a hopefully better version - does this solve the crashing you're seeing?

Index: lib/Analysis/ValueTracking.cpp
===================================================================
--- lib/Analysis/ValueTracking.cpp	(revision 319542)
+++ lib/Analysis/ValueTracking.cpp	(working copy)
@@ -800,7 +800,16 @@
   unsigned BitWidth = Known.getBitWidth();
 
   if (auto *SA = dyn_cast<ConstantInt>(I->getOperand(1))) {
-    unsigned ShiftAmt = SA->getLimitedValue(BitWidth-1);
+    // We can't use getZExtValue() here because the shift amount could be bigger
+    // than a 64-bit, but anything over 32-bit (the knownbits width is
+    // 32-bit) would result in poison, so it's fine to saturate to a 64-bit.
+    uint64_t ShiftAmt = SA->getLimitedValue();
+    if (ShiftAmt >= BitWidth) {
+      // This shift produces poison because the shift amount is too big. We can
+      // return anything we want. Choose 0 for the best folding opportunity.
+      Known.setAllZero();
+      return;
+    }
 
     computeKnownBits(I->getOperand(0), Known, Depth + 1, Q);
     Known.Zero = KZF(Known.Zero, ShiftAmt);

[igor-laevsky]

No, unfortunately your patch doesn't fix the problem. I believe those are two different paths in the code.

[spatel]

In D40649#942349, @igor-laevsky wrote:

No, unfortunately your patch doesn't fix the problem. I believe those are two different paths in the code.

Ah - I see what happened. I was only trying my reduced case. In the example here, we're getting into computeKnownBitsFromAssume() from the 'and'.

But I think my draft wasn't a waste. Try this test too:

define i128 @test1(i128 %a) {

%and1 = and i128 %a, 3
%B = lshr i128 %and1, -1
%cmp = icmp eq i128 %B, 1
tail call void @llvm.assume(i1 %cmp)
ret i128 %and1

}

[igor-laevsky]

Thanks! Please see updated change with this issue fixed.

[spatel]

LGTM.

llvm/trunk/lib/Analysis/ValueTracking.cpp

if (Depth == MaxDepth)
       continue;
 
     Value *A, *B;
     auto m_V = m_CombineOr(m_Specific(V),
                            m_CombineOr(m_PtrToInt(m_Specific(V)),
                            m_BitCast(m_Specific(V))));
 
     CmpInst::Predicate Pred;
-    ConstantInt *C;
+    uint64_t C;
     // assume(v = a)
     if (match(Arg, m_c_ICmp(Pred, m_V, m_Value(A))) &&
         Pred == ICmpInst::ICMP_EQ && isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
       KnownBits RHSKnown(BitWidth);
       computeKnownBits(A, RHSKnown, Depth+1, Query(Q, I));
       Known.Zero |= RHSKnown.Zero;
       Known.One  |= RHSKnown.One;
     // assume(v & b = a)
// assume(~(v ^ b) = a)
       Known.Zero |= RHSKnown.One  & BKnown.Zero;
       Known.One  |= RHSKnown.Zero & BKnown.Zero;
       Known.Zero |= RHSKnown.Zero & BKnown.One;
       Known.One  |= RHSKnown.One  & BKnown.One;
     // assume(v << c = a)
     } else if (match(Arg, m_c_ICmp(Pred, m_Shl(m_V, m_ConstantInt(C)),
                                    m_Value(A))) &&
                Pred == ICmpInst::ICMP_EQ &&
-               isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
+               isValidAssumeForContext(I, Q.CxtI, Q.DT) &&
+               C < BitWidth) {
       KnownBits RHSKnown(BitWidth);
       computeKnownBits(A, RHSKnown, Depth+1, Query(Q, I));
       // For those bits in RHS that are known, we can propagate them to known
       // bits in V shifted to the right by C.
-      RHSKnown.Zero.lshrInPlace(C->getZExtValue());
+      RHSKnown.Zero.lshrInPlace(C);
       Known.Zero |= RHSKnown.Zero;
-      RHSKnown.One.lshrInPlace(C->getZExtValue());
+      RHSKnown.One.lshrInPlace(C);
       Known.One  |= RHSKnown.One;
     // assume(~(v << c) = a)
     } else if (match(Arg, m_c_ICmp(Pred, m_Not(m_Shl(m_V, m_ConstantInt(C))),
                                    m_Value(A))) &&
                Pred == ICmpInst::ICMP_EQ &&
-               isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
+               isValidAssumeForContext(I, Q.CxtI, Q.DT) &&
+               C < BitWidth) {
       KnownBits RHSKnown(BitWidth);
       computeKnownBits(A, RHSKnown, Depth+1, Query(Q, I));
       // For those bits in RHS that are known, we can propagate them inverted
       // to known bits in V shifted to the right by C.
-      RHSKnown.One.lshrInPlace(C->getZExtValue());
+      RHSKnown.One.lshrInPlace(C);
       Known.Zero |= RHSKnown.One;
-      RHSKnown.Zero.lshrInPlace(C->getZExtValue());
+      RHSKnown.Zero.lshrInPlace(C);
       Known.One  |= RHSKnown.Zero;
     // assume(v >> c = a)
     } else if (match(Arg,
                      m_c_ICmp(Pred, m_Shr(m_V, m_ConstantInt(C)),
                               m_Value(A))) &&
                Pred == ICmpInst::ICMP_EQ &&
-               isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
+               isValidAssumeForContext(I, Q.CxtI, Q.DT) &&
+               C < BitWidth) {
       KnownBits RHSKnown(BitWidth);
       computeKnownBits(A, RHSKnown, Depth+1, Query(Q, I));
       // For those bits in RHS that are known, we can propagate them to known
       // bits in V shifted to the right by C.
-      Known.Zero |= RHSKnown.Zero << C->getZExtValue();
-      Known.One  |= RHSKnown.One  << C->getZExtValue();
+      Known.Zero |= RHSKnown.Zero << C;
+      Known.One  |= RHSKnown.One  << C;
     // assume(~(v >> c) = a)
     } else if (match(Arg, m_c_ICmp(Pred, m_Not(m_Shr(m_V, m_ConstantInt(C))),
                                    m_Value(A))) &&
                Pred == ICmpInst::ICMP_EQ &&
-               isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
+               isValidAssumeForContext(I, Q.CxtI, Q.DT) &&
+               C < BitWidth) {
       KnownBits RHSKnown(BitWidth);
       computeKnownBits(A, RHSKnown, Depth+1, Query(Q, I));
       // For those bits in RHS that are known, we can propagate them inverted
       // to known bits in V shifted to the right by C.
-      Known.Zero |= RHSKnown.One  << C->getZExtValue();
-      Known.One  |= RHSKnown.Zero << C->getZExtValue();
+      Known.Zero |= RHSKnown.One  << C;
+      Known.One  |= RHSKnown.Zero << C;
     // assume(v >=_s c) where c is non-negative
     } else if (match(Arg, m_ICmp(Pred, m_V, m_Value(A))) &&
                Pred == ICmpInst::ICMP_SGE &&
                isValidAssumeForContext(I, Q.CxtI, Q.DT)) {
       KnownBits RHSKnown(BitWidth);
       computeKnownBits(A, RHSKnown, Depth+1, Query(Q, I));
 
       if (RHSKnown.isNonNegative()) {

llvm/trunk/test/Transforms/InstCombine/out-of-bounds-indexes.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instcombine -S | FileCheck %s
+; Check that we don't crash on unreasonable constant indexes
+
+define i32 @test_out_of_bounds(i32 %a, i1 %x, i1 %y) {
+; CHECK-LABEL: @test_out_of_bounds(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AND1:%.*]] = and i32 [[A:%.*]], 3
+; CHECK-NEXT:    tail call void @llvm.assume(i1 false)
+; CHECK-NEXT:    ret i32 [[AND1]]
+;
+entry:
+  %and1 = and i32 %a, 3
+  %B = lshr i32 %and1, -2147483648
+  %cmp = icmp eq i32 %B, 1
+  tail call void @llvm.assume(i1 %cmp)
+  ret i32 %and1
+}
+
+define i128 @test_non64bit(i128 %a) {
+; CHECK-LABEL: @test_non64bit(
+; CHECK-NEXT:    [[AND1:%.*]] = and i128 [[A:%.*]], 3
+; CHECK-NEXT:    tail call void @llvm.assume(i1 false)
+; CHECK-NEXT:    ret i128 [[AND1]]
+;
+  %and1 = and i128 %a, 3
+  %B = lshr i128 %and1, -1
+  %cmp = icmp eq i128 %B, 1
+  tail call void @llvm.assume(i1 %cmp)
+  ret i128 %and1
+}
+
+declare void @llvm.assume(i1)