[InstCombine] shrink type of sdiv if dividend is sexted and constant divisor is small enough (PR28153)

[spatel]

Summary

This is intended to fix PR28153:
https://llvm.org/bugs/show_bug.cgi?id=28153

[majnemer]

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1156 ↗

(On Diff #62013)

Say X is i16 -32768 and C is i32 65535, the result of the sdiv will be -32768/65535 which is zero.
However, trunc of C to i16 is 65535. the result of X sdiv C will be -32768/-1 which is UB because it overflows.

[spatel]

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1156 ↗

(On Diff #62013)

The code comment may be misleading. Your example won't fold because getMinSignedBits() for 65535 is 17. It may be clearer when you see the limits in the test cases. Any ideas about how to make that code comment more specific? Should I put your example case in there?

[spatel]

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1156 ↗

(On Diff #62013)

Sorry if this comment wasn't clear as well. :)
I checked in a couple of negative test cases ahead of this patch; they're the last ones in test/Transforms/InstCombine/div.ll.

[majnemer]

LGTM

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

1156 ↗

(On Diff #62013)

I missed the MinSignedBits part of your patch, seems fine to me. I think it'd be good to include by testcase.

llvm/trunk/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

if (match(Op1, m_APInt(Op1C))) {
     if (Op1C->isAllOnesValue())
       return BinaryOperator::CreateNeg(Op0);
 
     // sdiv exact X, C  -->  ashr exact X, log2(C)
     if (I.isExact() && Op1C->isNonNegative() && Op1C->isPowerOf2()) {
       Value *ShAmt = ConstantInt::get(Op1->getType(), Op1C->exactLogBase2());
       return BinaryOperator::CreateExactAShr(Op0, ShAmt, I.getName());
     }
+
+    // If the dividend is sign-extended and the constant divisor is small enough
+    // to fit in the source type, shrink the division to the narrower type:
+    // (sext X) sdiv C --> sext (X sdiv C)
+    Value *Op0Src;
+    if (match(Op0, m_OneUse(m_SExt(m_Value(Op0Src)))) &&
+        Op0Src->getType()->getScalarSizeInBits() >= Op1C->getMinSignedBits()) {
+
+      // In the general case, we need to make sure that the dividend is not the
+      // minimum signed value because dividing that by -1 is UB. But here, we
+      // know that the -1 divisor case is already handled above.
+
+      Constant *NarrowDivisor =
+          ConstantExpr::getTrunc(cast<Constant>(Op1), Op0Src->getType());
+      Value *NarrowOp = Builder->CreateSDiv(Op0Src, NarrowDivisor);
+      return new SExtInst(NarrowOp, Op0->getType());
+    }
   }
 
   if (Constant *RHS = dyn_cast<Constant>(Op1)) {
     // X/INT_MIN -> X == INT_MIN
     if (RHS->isMinSignedValue())
       return new ZExtInst(Builder->CreateICmpEQ(Op0, Op1), I.getType());
 
     // -X/C  -->  X/-C  provided the negation doesn't overflow.

llvm/trunk/test/Transforms/InstCombine/div.ll

lor.rhs:                                          ; preds = %entry
   br label %lor.end
 
 lor.end:                                          ; preds = %lor.rhs, %entry
   %t.0 = phi i32 [ %0, %entry ], [ %mul, %lor.rhs ]
   %div = sdiv i32 %t.0, 2
   ret i32 %div
 }
 
-; FIXME:
 ; We can perform the division in the smaller type.
 
 define i32 @shrink(i8 %x) {
 ; CHECK-LABEL: @shrink(
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 %x to i32
-; CHECK-NEXT:    [[DIV:%.*]] = sdiv i32 [[CONV]], 127
+; CHECK-NEXT:    [[TMP1:%.*]] = sdiv i8 %x, 127
+; CHECK-NEXT:    [[DIV:%.*]] = sext i8 [[TMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[DIV]]
 ;
   %conv = sext i8 %x to i32
   %div = sdiv i32 %conv, 127
   ret i32 %div
 }
 
 ; Division in the smaller type can lead to more optimizations.
 
 define i32 @zap(i8 %x) {
 ; CHECK-LABEL: @zap(
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 %x to i32
-; CHECK-NEXT:    [[DIV:%.*]] = sdiv i32 [[CONV]], -128
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i8 %x, -128
+; CHECK-NEXT:    [[DIV:%.*]] = zext i1 [[TMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[DIV]]
 ;
   %conv = sext i8 %x to i32
   %div = sdiv i32 %conv, -128
   ret i32 %div
 }
 
 ; Splat constant divisors should get the same folds.
 
 define <3 x i32> @shrink_vec(<3 x i8> %x) {
 ; CHECK-LABEL: @shrink_vec(
-; CHECK-NEXT:    [[CONV:%.*]] = sext <3 x i8> %x to <3 x i32>
-; CHECK-NEXT:    [[DIV:%.*]] = sdiv <3 x i32> [[CONV]], <i32 127, i32 127, i32 127>
+; CHECK-NEXT:    [[TMP1:%.*]] = sdiv <3 x i8> %x, <i8 127, i8 127, i8 127>
+; CHECK-NEXT:    [[DIV:%.*]] = sext <3 x i8> [[TMP1]] to <3 x i32>
 ; CHECK-NEXT:    ret <3 x i32> [[DIV]]
 ;
   %conv = sext <3 x i8> %x to <3 x i32>
   %div = sdiv <3 x i32> %conv, <i32 127, i32 127, i32 127>
   ret <3 x i32> %div
 }
 
 define <2 x i32> @zap_vec(<2 x i8> %x) {
 ; CHECK-LABEL: @zap_vec(
-; CHECK-NEXT:    [[CONV:%.*]] = sext <2 x i8> %x to <2 x i32>
-; CHECK-NEXT:    [[DIV:%.*]] = sdiv <2 x i32> [[CONV]], <i32 -128, i32 -128>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq <2 x i8> %x, <i8 -128, i8 -128>
+; CHECK-NEXT:    [[DIV:%.*]] = zext <2 x i1> [[TMP1]] to <2 x i32>
 ; CHECK-NEXT:    ret <2 x i32> [[DIV]]
 ;
   %conv = sext <2 x i8> %x to <2 x i32>
   %div = sdiv <2 x i32> %conv, <i32 -128, i32 -128>
   ret <2 x i32> %div
 }
 
 ; But we can't do this if the signed constant won't fit in the original type.

 ; CHECK-NEXT:    [[DIV:%.*]] = sdiv i32 [[CONV]], -129
 ; CHECK-NEXT:    ret i32 [[DIV]]
 ;
   %conv = sext i8 %x to i32
   %div = sdiv i32 %conv, -129
   ret i32 %div
 }
 
+; 17 bits are needed to represent 65535 as a signed value, so this shouldn't fold.
+
+define i32 @shrink_no3(i16 %x) {
+; CHECK-LABEL: @shrink_no3(
+; CHECK-NEXT:    [[CONV:%.*]] = sext i16 %x to i32
+; CHECK-NEXT:    [[DIV:%.*]] = sdiv i32 [[CONV]], 65535
+; CHECK-NEXT:    ret i32 [[DIV]]
+;
+  %conv = sext i16 %x to i32
+  %div = sdiv i32 %conv, 65535
+  ret i32 %div
+}
+