[ConstraintElimination] Fix undefined behaviour in shl decomposition

[zjaffal]

Summary

Add checks to prevent decomposing constants bigger than 64.
relates to https://github.com/llvm/llvm-project/issues/61127

[fhahn]

LGTM. thanks! But please include the full test (from D145676) in the commit with the fix, as the tests will otherwise fail on the UBSan bots.

[zjaffal]

Add the tests from from D145676 herew

[NathanielMcVicar]

This change appears to have broken the Windows MLIR bot (I suspect others Windows builders).
https://lab.llvm.org/buildbot/#/builders/13/builds/32862/steps/6/logs/stdio

C:\buildbot\mlir-x64-windows-ninja\llvm-project\llvm\lib\Transforms\Scalar\ConstraintElimination.cpp(382): error C2220: the following warning is treated as an error
C:\buildbot\mlir-x64-windows-ninja\llvm-project\llvm\lib\Transforms\Scalar\ConstraintElimination.cpp(382): warning C4334: '<<': result of 32-bit shift implicitly converted to 64 bits (was 64-bit shift intended?)

[vitalybuka]

In D145677#4182592, @NathanielMcVicar wrote:

This change appears to have broken the Windows MLIR bot (I suspect others Windows builders).
https://lab.llvm.org/buildbot/#/builders/13/builds/32862/steps/6/logs/stdio

C:\buildbot\mlir-x64-windows-ninja\llvm-project\llvm\lib\Transforms\Scalar\ConstraintElimination.cpp(382): error C2220: the following warning is treated as an error
C:\buildbot\mlir-x64-windows-ninja\llvm-project\llvm\lib\Transforms\Scalar\ConstraintElimination.cpp(382): warning C4334: '<<': result of 32-bit shift implicitly converted to 64 bits (was 64-bit shift intended?)

UBSAN as well https://lab.llvm.org/buildbot/#/builders/238/builds/1810

[fhahn]

Thanks for the fix @vitalybuka ! I am surprised this wasn't covered by the added tests. Any idea why that is?

[vitalybuka]

In D145677#4184405, @fhahn wrote:

Thanks for the fix @vitalybuka ! I am surprised this wasn't covered by the added tests. Any idea why that is?

I guess you need shift by [32; 63] to hit this case

[fhahn]

In D145677#4185360, @vitalybuka wrote:

In D145677#4184405, @fhahn wrote:

Thanks for the fix @vitalybuka ! I am surprised this wasn't covered by the added tests. Any idea why that is?

I guess you need shift by [32; 63] to hit this case

Right, @zjaffal could you add a test case which shifts by an offset between [32; 63] ? Ideally one where ConstraintElimination can perform a simplification.

llvm/lib/Transforms/Scalar/ConstraintElimination.cpp

if (match(V, m_Add(m_Value(Op0), m_ConstantInt(CI))) && CI->isNegative() &&
       canUseSExt(CI)) {
     Preconditions.emplace_back(
         CmpInst::ICMP_UGE, Op0,
         ConstantInt::get(Op0->getType(), CI->getSExtValue() * -1));
     return MergeResults(Op0, CI, true);
   }
 
   if (match(V, m_NUWShl(m_Value(Op1), m_ConstantInt(CI))) && canUseSExt(CI)) {
-    int64_t Mult = int64_t(std::pow(int64_t(2), CI->getSExtValue()));
+    if (CI->getSExtValue() < 0 || CI->getSExtValue() >= 64)
+      return {V, IsKnownNonNegative};
+    int64_t Mult = 1 << CI->getSExtValue();
     auto Result = decompose(Op1, Preconditions, IsSigned, DL);
     Result.mul(Mult);
     return Result;
   }
 
   if (match(V, m_NUWMul(m_Value(Op1), m_ConstantInt(CI))) && canUseSExt(CI) &&
       (!CI->isNegative())) {
     auto Result = decompose(Op1, Preconditions, IsSigned, DL);

llvm/test/Transforms/ConstraintElimination/shl.ll

 define i1 @shl_overflow(i64 %start) {
 ; CHECK-LABEL: @shl_overflow(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[PRE_COND:%.*]] = icmp ugt i64 [[START:%.*]], 0
 ; CHECK-NEXT:    br i1 [[PRE_COND]], label [[MAIN:%.*]], label [[EXIT:%.*]]
 ; CHECK:       main:
 ; CHECK-NEXT:    [[TMP0:%.*]] = shl nuw nsw i64 [[START]], -1
 ; CHECK-NEXT:    [[TMP1:%.*]] = icmp uge i64 [[TMP0]], [[START]]
-; CHECK-NEXT:    ret i1 false
+; CHECK-NEXT:    ret i1 [[TMP1]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i1 false
 ;
 entry:
   %pre.cond = icmp ugt i64 %start, 0
   br i1 %pre.cond, label %main, label %exit
 
 main:
   %0 = shl nuw nsw i64 %start, -1
   %1 = icmp uge i64 %0, %start
   ret i1 %1
 
 exit:
   ret i1 0
 }
 
 
 define i1 @shl_overflow_2() {
 ; CHECK-LABEL: @shl_overflow_2(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SHL_UB:%.*]] = shl nuw nsw i256 0, 64
 ; CHECK-NEXT:    [[SHL_CMP:%.*]] = icmp uge i256 [[SHL_UB]], 0
-; CHECK-NEXT:    ret i1 true
+; CHECK-NEXT:    ret i1 [[SHL_CMP]]
 ;
 entry:
   %shl.ub = shl nuw nsw i256 0, 64
   %shl.cmp = icmp uge i256 %shl.ub, 0
   ret i1 %shl.cmp
 }
 
 define i1 @shl_overflow_3() {
 ; CHECK-LABEL: @shl_overflow_3(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[SHL_UB:%.*]] = shl nuw nsw i256 0, 65
 ; CHECK-NEXT:    [[SHL_CMP:%.*]] = icmp uge i256 [[SHL_UB]], 0
-; CHECK-NEXT:    ret i1 true
+; CHECK-NEXT:    ret i1 [[SHL_CMP]]
 ;
 entry:
   %shl.ub = shl nuw nsw i256 0, 65
   %shl.cmp = icmp uge i256 %shl.ub, 0
   ret i1 %shl.cmp
 }