[BypassSlowDivision] Improve our handling of divisions by constants

[sanjoy]

Summary

Don't bail out on constant divisors for divisions that can be narrowed without
introducing control flow . This gives us a 32 bit multiply instead of an
emulated 64 bit multiply in the generated PTX assembly.

[jlebar]

lib/Transforms/Utils/BypassSlowDivision.cpp

359

The divisor also stays a constant in the other two cases, so is this really the thing that makes us want to do this transformation but not the other ones when the divisor is a constant?

[sanjoy]

lib/Transforms/Utils/BypassSlowDivision.cpp

359

The way I worked this out in my head is that, for the optimization to be worth it, the perf improvement from doing a shorter op must be more than the perf regression due to control flow. When the op is divide the assumption is that this tradeoff is worth it, but when the op is multiply (division by constant) this tradeoff is not worth it.

However, if we can narrow the op (divide or multiply) without any control flow then there is no tradeoff -- we should always narrow the op.

[sanjoy]

• update comments

lib/Transforms/Utils/BypassSlowDivision.cpp

 }
 
 /// Substitutes the div/rem instruction with code that checks the value of the
 /// operands and uses a shorter-faster div/rem instruction when possible.
 Optional<QuotRemPair> FastDivInsertionTask::insertFastDivAndRem() {
   Value *Dividend = SlowDivOrRem->getOperand(0);
   Value *Divisor = SlowDivOrRem->getOperand(1);
 
-  if (isa<ConstantInt>(Divisor)) {
-    // Keep division by a constant for DAGCombiner.
-    return None;
-  }
-
   VisitedSetTy SetL;
   ValueRange DividendRange = getValueRange(Dividend, SetL);
   if (DividendRange == VALRNG_LIKELY_LONG)
     return None;
 
   VisitedSetTy SetR;
   ValueRange DivisorRange = getValueRange(Divisor, SetR);
   if (DivisorRange == VALRNG_LIKELY_LONG)
     return None;
 
   bool DividendShort = (DividendRange == VALRNG_KNOWN_SHORT);
   bool DivisorShort = (DivisorRange == VALRNG_KNOWN_SHORT);
 
   if (DividendShort && DivisorShort) {
     // If both operands are known to be short then just replace the long
-    // division with a short one in-place.
+    // division with a short one in-place.  If the divisor is a constant, then
+    // it stays a constant so we don't have to worry about breaking DAGCombiner
+    // optimizations on divisions by constants.
 
     IRBuilder<> Builder(SlowDivOrRem);
     Value *TruncDividend = Builder.CreateTrunc(Dividend, BypassType);
     Value *TruncDivisor = Builder.CreateTrunc(Divisor, BypassType);
     Value *TruncDiv = Builder.CreateUDiv(TruncDividend, TruncDivisor);
     Value *TruncRem = Builder.CreateURem(TruncDividend, TruncDivisor);
     Value *ExtDiv = Builder.CreateZExt(TruncDiv, getSlowType());
     Value *ExtRem = Builder.CreateZExt(TruncRem, getSlowType());
     return QuotRemPair(ExtDiv, ExtRem);
-  } else if (DividendShort && !isSignedOp()) {
+  }
+
+  if (isa<ConstantInt>(Divisor)) {
+    // Keep division by a constant for DAGCombiner.
+    return None;
+  }
+
+  if (DividendShort && !isSignedOp()) {
     // If the division is unsigned and Dividend is known to be short, then
     // either
     // 1) Divisor is less or equal to Dividend, and the result can be computed
     //    with a short division.
     // 2) Divisor is greater than Dividend. In this case, no division is needed
     //    at all: The quotient is 0 and the remainder is equal to Dividend.
     //
     // So instead of checking at runtime whether Divisor fits into BypassType,

test/Transforms/CodeGenPrepare/NVPTX/bypass-slow-div.ll

define void @rem_only(i64 %a, i64 %b, i64* %retptr) {
   ; CHECK: urem i32
   ; CHECK-NOT: div
   ; CHECK: rem i64
   ; CHECK-NOT: div
   %d = srem i64 %a, %b
   store i64 %d, i64* %retptr
   ret void
 }
+
+; CHECK-LABEL: @udiv_by_constant(
+define i64 @udiv_by_constant(i32 %a) {
+; CHECK-NEXT:    [[A_ZEXT:%.*]] = zext i32 [[A:%.*]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i64 [[A_ZEXT]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = udiv i32 [[TMP1]], 50
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    ret i64 [[TMP3]]
+
+  %a.zext = zext i32 %a to i64
+  %wide.div = udiv i64 %a.zext, 50
+  ret i64 %wide.div
+}
+
+; CHECK-LABEL: @urem_by_constant(
+define i64 @urem_by_constant(i32 %a) {
+; CHECK-NEXT:    [[A_ZEXT:%.*]] = zext i32 [[A:%.*]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i64 [[A_ZEXT]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = urem i32 [[TMP1]], 50
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i32 [[TMP2]] to i64
+; CHECK-NEXT:    ret i64 [[TMP3]]
+
+  %a.zext = zext i32 %a to i64
+  %wide.div = urem i64 %a.zext, 50
+  ret i64 %wide.div
+}
+
+; Negative test: instead of emitting a runtime check on %a, we prefer to let the
+; DAGCombiner transform this division by constant into a multiplication (with a
+; "magic constant").
+;
+; CHECK-LABEL: @udiv_by_constant_negative_0(
+define i64 @udiv_by_constant_negative_0(i64 %a) {
+; CHECK-NEXT:    [[WIDE_DIV:%.*]] = udiv i64 [[A:%.*]], 50
+; CHECK-NEXT:    ret i64 [[WIDE_DIV]]
+
+  %wide.div = udiv i64 %a, 50
+  ret i64 %wide.div
+}
+
+; Negative test: while we know the dividend is short, the divisor isn't.  This
+; test is here for completeness, but instcombine will optimize this to return 0.
+;
+; CHECK-LABEL: @udiv_by_constant_negative_1(
+define i64 @udiv_by_constant_negative_1(i32 %a) {
+; CHECK-NEXT:    [[A_ZEXT:%.*]] = zext i32 [[A:%.*]] to i64
+; CHECK-NEXT:    [[WIDE_DIV:%.*]] = udiv i64 [[A_ZEXT]], 8589934592
+; CHECK-NEXT:    ret i64 [[WIDE_DIV]]
+
+  %a.zext = zext i32 %a to i64
+  %wide.div = udiv i64 %a.zext, 8589934592 ;; == 1 << 33
+  ret i64 %wide.div
+}
+
+; URem version of udiv_by_constant_negative_0
+;
+; CHECK-LABEL: @urem_by_constant_negative_0(
+define i64 @urem_by_constant_negative_0(i64 %a) {
+; CHECK-NEXT:    [[WIDE_DIV:%.*]] = urem i64 [[A:%.*]], 50
+; CHECK-NEXT:    ret i64 [[WIDE_DIV]]
+
+  %wide.div = urem i64 %a, 50
+  ret i64 %wide.div
+}
+
+; URem version of udiv_by_constant_negative_1
+;
+; CHECK-LABEL: @urem_by_constant_negative_1(
+define i64 @urem_by_constant_negative_1(i32 %a) {
+; CHECK-NEXT:    [[A_ZEXT:%.*]] = zext i32 [[A:%.*]] to i64
+; CHECK-NEXT:    [[WIDE_DIV:%.*]] = urem i64 [[A_ZEXT]], 8589934592
+; CHECK-NEXT:    ret i64 [[WIDE_DIV]]
+
+  %a.zext = zext i32 %a to i64
+  %wide.div = urem i64 %a.zext, 8589934592 ;; == 1 << 33
+  ret i64 %wide.div
+}