[InstCombine] Teach SimplifyDemandedBits that udiv doesn't demand low dividend bits that are zero in the divisor

[bkramer]

Summary

This is safe as long as the udiv is not exact. The pattern is not common in
C++ code, but comes up all the time in code generated by XLA's GPU backend.

[sanjoy]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

556 ↗

(On Diff #145954)

Can we also propagate the demanded bits of the division result (in some form) into DemandedMaskIn?

[bkramer]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

556 ↗

(On Diff #145954)

Maybe. This is hard for non-power of 2 divisors though and probably not worth the effort :|

[sanjoy]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

556 ↗

(On Diff #145954)

I think demanded-bits should compose which means you can treat a /u 12 as (a >> 2) /u 3. Since in udiv input bit at index i can only affect output bits at index i and below (i.e. less significant) we should be able to left-smear the demanded bits demanded bits for (a >> 2) /u 3 into the demanded bits for (a >> 2) and then apply the rule for >>.

For instance, if we demand bit only 5 for an i8 a /u 12 == (a >> 2) /u 3 then we can say:

• We only need bits 5, 6, 7 and from a >> 2
• Thus we only need bit 7 from a

I'd also be fine with a TODO here since it does seem complex.

[bkramer]

Add a FIXME for using the demanded mask of the udiv result.

[sanjoy]

lgtm

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

if (match(I->getOperand(1), m_APInt(SA))) {
         LShr->setIsExact(cast<BinaryOperator>(I)->isExact());
         return InsertNewInstWith(LShr, *I);
       } else if (Known.One[BitWidth-ShiftAmt-1]) { // New bits are known one.
         Known.One |= HighBits;
       }
     }
     break;
   }
+  case Instruction::UDiv: {
+    // UDiv doesn't demand low bits that are zero in the divisor.
+    const APInt *SA;
+    if (match(I->getOperand(1), m_APInt(SA))) {
+      // If the shift is exact, then it does demand the low bits.
+      if (cast<UDivOperator>(I)->isExact())
+        break;
+
+      // FIXME: Take the demanded mask of the result into account.
+      APInt DemandedMaskIn =
+          APInt::getHighBitsSet(BitWidth, BitWidth - SA->countTrailingZeros());
+      if (SimplifyDemandedBits(I, 0, DemandedMaskIn, Known, Depth + 1))
+        return I;
+    }
+    break;
+  }
   case Instruction::SRem:
     if (ConstantInt *Rem = dyn_cast<ConstantInt>(I->getOperand(1))) {
       // X % -1 demands all the bits because we don't want to introduce
       // INT_MIN % -1 (== undef) by accident.
       if (Rem->isMinusOne())
         break;
       APInt RA = Rem->getValue().abs();
       if (RA.isPowerOf2()) {

llvm/trunk/test/Transforms/InstCombine/udiv-simplify.ll

;
   %B9 = xor i177 %B4, %B3
   %B13 = ashr i177 %Y, %B2
   %B22 = add i177 %B9, %B13
   %B1 = udiv i177 %B5, %B6
   %C9 = icmp ult i177 %Y, %B22
   store i1 %C9, i1* undef
   ret i177 %B1
 }
+
+define i32 @udiv_demanded(i32 %a) {
+; CHECK-LABEL: @udiv_demanded(
+; CHECK-NEXT:    [[U:%.*]] = udiv i32 [[A:%.*]], 12
+; CHECK-NEXT:    ret i32 [[U]]
+;
+  %o = or i32 %a, 3
+  %u = udiv i32 %o, 12
+  ret i32 %u
+}
+
+define i32 @udiv_exact_demanded(i32 %a) {
+; CHECK-LABEL: @udiv_exact_demanded(
+; CHECK-NEXT:    [[O:%.*]] = and i32 [[A:%.*]], -3
+; CHECK-NEXT:    [[U:%.*]] = udiv exact i32 [[O]], 12
+; CHECK-NEXT:    ret i32 [[U]]
+;
+  %o = and i32 %a, -3
+  %u = udiv exact i32 %o, 12
+  ret i32 %u
+}