[AggressiveInstCombine] Add `udiv` and `urem` instrs to TruncInstCombine DAG

[anton-afanasyev]

Summary

Add udiv and urem instructions to the DAG post-dominated by trunc,
allowing TruncInstCombine to reduce bitwidth of expressions containing these
instructions. It is sufficient to require that all truncated bits of both
operands are zeros: https://alive2.llvm.org/ce/z/yiithn
(urem case is identical).

[lebedev.ri]

Seems reasonable.

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

152

sdiv, srem; see b289dc530632613edb3eb067895c1981cb77ccd0.

315–324

[anton-afanasyev]

Address comments

[anton-afanasyev]

Small fix

[lebedev.ri]

LG, thanks.

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

318–319

Hm, if you swap arguments around, does it lead to better clang-format?

[anton-afanasyev]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

152

Ok, added to todo.

315–324

Ok, done.

318–319

No(
It fixes even my "legal" trials like this one:

MinBitWidth = std::max(Known.getMaxValue().getActiveBits(),
                       MinBitWidth);

[spatel]

llvm/test/Transforms/AggressiveInstCombine/trunc_udivrem.ll

22

Typo: 'zext i16 %y to i32'

113

Typo: 'zext i16 %y to i32'

[anton-afanasyev]

llvm/test/Transforms/AggressiveInstCombine/trunc_udivrem.ll

22

Oops, thanks, fixed!

113

Oops, thanks, fixed!

[anton-afanasyev]

Fix test

[spatel]

LGTM

[anton-afanasyev]

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

152

Hi @lebedev.ri, @spatel , according to this patch rGb289dc530632, sdiv/srem (as well as udiv/urem) instrs are guarded with trunc/ext wherever it's possible. So it looks like support them here in AIC is redundant as we get always MinBitWidth == OrigBitWidth after correlated value propagation.
So I'm leaning towards reverting this patch. What do you think?

llvm/lib/Transforms/AggressiveInstCombine/TruncInstCombine.cpp

static void getRelevantOperands(Instruction *I, SmallVectorImpl<Value *> &Ops) {
   case Instruction::Sub:
   case Instruction::Mul:
   case Instruction::And:
   case Instruction::Or:
   case Instruction::Xor:
   case Instruction::Shl:
   case Instruction::LShr:
   case Instruction::AShr:
+  case Instruction::UDiv:
+  case Instruction::URem:
     Ops.push_back(I->getOperand(0));
     Ops.push_back(I->getOperand(1));
     break;
   case Instruction::Select:
     Ops.push_back(I->getOperand(1));
     Ops.push_back(I->getOperand(2));
     break;
   default:
while (!Worklist.empty()) {
     case Instruction::Sub:
     case Instruction::Mul:
     case Instruction::And:
     case Instruction::Or:
     case Instruction::Xor:
     case Instruction::Shl:
     case Instruction::LShr:
     case Instruction::AShr:
+    case Instruction::UDiv:
+    case Instruction::URem:
     case Instruction::Select: {
       SmallVector<Value *, 2> Operands;
       getRelevantOperands(I, Operands);
       append_range(Worklist, Operands);
       break;
     }
     default:
       // TODO: Can handle more cases here:
       // 1. shufflevector, extractelement, insertelement
-      // 2. udiv, urem
+      // 2. sdiv, srem
       // 3. phi node(and loop handling)
       // ...
       return false;
     }
   }
   return true;
 }
 
 unsigned TruncInstCombine::getMinBitWidth() {
   SmallVector<Value *, 8> Worklist;
if (I->isShift()) {
       if (I->getOpcode() == Instruction::AShr) {
         unsigned NumSignBits = ComputeNumSignBits(I->getOperand(0));
         MinBitWidth = std::max(MinBitWidth, OrigBitWidth - NumSignBits + 1);
       }
       if (MinBitWidth >= OrigBitWidth)
         return nullptr;
       Itr.second.MinBitWidth = MinBitWidth;
     }
+    if (I->getOpcode() == Instruction::UDiv ||
+        I->getOpcode() == Instruction::URem) {
+      unsigned MinBitWidth = 0;
+      for (const auto &Op : I->operands()) {
+        KnownBits Known = computeKnownBits(Op);
+        MinBitWidth =
+            std::max(Known.getMaxValue().getActiveBits(), MinBitWidth);
+        if (MinBitWidth >= OrigBitWidth)
+          return nullptr;
+      }
+      Itr.second.MinBitWidth = MinBitWidth;
+    }
   }
 
   // Calculate minimum allowed bit-width allowed for shrinking the currently
   // visited truncate's operand.
   unsigned MinBitWidth = getMinBitWidth();
 
   // Check that we can shrink to smaller bit-width than original one and that
   // it is similar to the DesiredBitWidth is such exists.
for (auto &Itr : InstInfoMap) { // Forward
     case Instruction::Add:
     case Instruction::Sub:
     case Instruction::Mul:
     case Instruction::And:
     case Instruction::Or:
     case Instruction::Xor:
     case Instruction::Shl:
     case Instruction::LShr:
-    case Instruction::AShr: {
+    case Instruction::AShr:
+    case Instruction::UDiv:
+    case Instruction::URem: {
       Value *LHS = getReducedOperand(I->getOperand(0), SclTy);
       Value *RHS = getReducedOperand(I->getOperand(1), SclTy);
       Res = Builder.CreateBinOp((Instruction::BinaryOps)Opc, LHS, RHS);
       // Preserve `exact` flag since truncation doesn't change exactness
       if (auto *PEO = dyn_cast<PossiblyExactOperator>(I))
         if (auto *ResI = dyn_cast<Instruction>(Res))
           ResI->setIsExact(PEO->isExact());
       break;

llvm/test/Transforms/AggressiveInstCombine/trunc_udivrem.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -aggressive-instcombine -S | FileCheck %s
 
 define i16 @udiv_one_arg(i8 %x) {
 ; CHECK-LABEL: @udiv_one_arg(
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[X:%.*]] to i32
-; CHECK-NEXT:    [[DIV:%.*]] = udiv i32 [[ZEXT]], 42
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[DIV]] to i16
-; CHECK-NEXT:    ret i16 [[TRUNC]]
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[X:%.*]] to i16
+; CHECK-NEXT:    [[DIV:%.*]] = udiv i16 [[ZEXT]], 42
+; CHECK-NEXT:    ret i16 [[DIV]]
 ;
   %zext = zext i8 %x to i32
   %div = udiv i32 %zext, 42
   %trunc = trunc i32 %div to i16
   ret i16 %trunc
 }
 
 define i16 @udiv_two_args(i16 %x, i16 %y) {
 ; CHECK-LABEL: @udiv_two_args(
-; CHECK-NEXT:    [[ZEXTX:%.*]] = zext i16 [[X:%.*]] to i32
-; CHECK-NEXT:    [[ZEXTY:%.*]] = zext i16 [[X]] to i32
-; CHECK-NEXT:    [[I0:%.*]] = udiv i32 [[ZEXTX]], [[ZEXTY]]
-; CHECK-NEXT:    [[R:%.*]] = trunc i32 [[I0]] to i16
-; CHECK-NEXT:    ret i16 [[R]]
+; CHECK-NEXT:    [[I0:%.*]] = udiv i16 [[X:%.*]], [[X]]
+; CHECK-NEXT:    ret i16 [[I0]]
 ;
   %zextx = zext i16 %x to i32
   %zexty = zext i16 %x to i32
   %i0 = udiv i32 %zextx, %zexty
   %r = trunc i32 %i0 to i16
   ret i16 %r
 }
 
+; Negative test
 define i16 @udiv_big_const(i8 %x) {
 ; CHECK-LABEL: @udiv_big_const(
 ; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[X:%.*]] to i32
 ; CHECK-NEXT:    [[DIV:%.*]] = udiv i32 [[ZEXT]], 70000
 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[DIV]] to i16
 ; CHECK-NEXT:    ret i16 [[TRUNC]]
 ;
   %zext = zext i8 %x to i32
   %div = udiv i32 %zext, 70000
   %trunc = trunc i32 %div to i16
   ret i16 %trunc
 }
 
 define <2 x i16> @udiv_vector(<2 x i8> %x) {
 ; CHECK-LABEL: @udiv_vector(
-; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[S:%.*]] = udiv <2 x i32> [[Z]], <i32 4, i32 10>
-; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i32> [[S]] to <2 x i16>
-; CHECK-NEXT:    ret <2 x i16> [[T]]
+; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[S:%.*]] = udiv <2 x i16> [[Z]], <i16 4, i16 10>
+; CHECK-NEXT:    ret <2 x i16> [[S]]
 ;
   %z = zext <2 x i8> %x to <2 x i32>
   %s = udiv <2 x i32> %z, <i32 4, i32 10>
   %t = trunc <2 x i32> %s to <2 x i16>
   ret <2 x i16> %t
 }
 
+; Negative test: can only fold to <2 x i16>, requiring new vector type
 define <2 x i8> @udiv_vector_need_new_vector_type(<2 x i8> %x) {
 ; CHECK-LABEL: @udiv_vector_need_new_vector_type(
 ; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
 ; CHECK-NEXT:    [[S:%.*]] = udiv <2 x i32> [[Z]], <i32 4, i32 500>
 ; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i32> [[S]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[T]]
 ;
   %z = zext <2 x i8> %x to <2 x i32>
   %s = udiv <2 x i32> %z, <i32 4, i32 500>
   %t = trunc <2 x i32> %s to <2 x i8>
   ret <2 x i8> %t
 }
 
+; Negative test
 define <2 x i16> @udiv_vector_big_const(<2 x i8> %x) {
 ; CHECK-LABEL: @udiv_vector_big_const(
 ; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
 ; CHECK-NEXT:    [[S:%.*]] = udiv <2 x i32> [[Z]], <i32 16, i32 70000>
 ; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i32> [[S]] to <2 x i16>
 ; CHECK-NEXT:    ret <2 x i16> [[T]]
 ;
   %z = zext <2 x i8> %x to <2 x i32>
   %s = udiv <2 x i32> %z, <i32 16, i32 70000>
   %t = trunc <2 x i32> %s to <2 x i16>
   ret <2 x i16> %t
 }
 
 define i16 @udiv_exact(i16 %x, i16 %y) {
 ; CHECK-LABEL: @udiv_exact(
-; CHECK-NEXT:    [[ZEXTX:%.*]] = zext i16 [[X:%.*]] to i32
-; CHECK-NEXT:    [[ZEXTY:%.*]] = zext i16 [[X]] to i32
-; CHECK-NEXT:    [[I0:%.*]] = udiv exact i32 [[ZEXTX]], [[ZEXTY]]
-; CHECK-NEXT:    [[R:%.*]] = trunc i32 [[I0]] to i16
-; CHECK-NEXT:    ret i16 [[R]]
+; CHECK-NEXT:    [[I0:%.*]] = udiv exact i16 [[X:%.*]], [[X]]
+; CHECK-NEXT:    ret i16 [[I0]]
 ;
   %zextx = zext i16 %x to i32
   %zexty = zext i16 %x to i32
   %i0 = udiv exact i32 %zextx, %zexty
   %r = trunc i32 %i0 to i16
   ret i16 %r
 }
 
 
 define i16 @urem_one_arg(i8 %x) {
 ; CHECK-LABEL: @urem_one_arg(
-; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[X:%.*]] to i32
-; CHECK-NEXT:    [[DIV:%.*]] = urem i32 [[ZEXT]], 42
-; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[DIV]] to i16
-; CHECK-NEXT:    ret i16 [[TRUNC]]
+; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[X:%.*]] to i16
+; CHECK-NEXT:    [[DIV:%.*]] = urem i16 [[ZEXT]], 42
+; CHECK-NEXT:    ret i16 [[DIV]]
 ;
   %zext = zext i8 %x to i32
   %div = urem i32 %zext, 42
   %trunc = trunc i32 %div to i16
   ret i16 %trunc
 }
 
 define i16 @urem_two_args(i16 %x, i16 %y) {
 ; CHECK-LABEL: @urem_two_args(
-; CHECK-NEXT:    [[ZEXTX:%.*]] = zext i16 [[X:%.*]] to i32
-; CHECK-NEXT:    [[ZEXTY:%.*]] = zext i16 [[X]] to i32
-; CHECK-NEXT:    [[I0:%.*]] = urem i32 [[ZEXTX]], [[ZEXTY]]
-; CHECK-NEXT:    [[R:%.*]] = trunc i32 [[I0]] to i16
-; CHECK-NEXT:    ret i16 [[R]]
+; CHECK-NEXT:    [[I0:%.*]] = urem i16 [[X:%.*]], [[X]]
+; CHECK-NEXT:    ret i16 [[I0]]
 ;
   %zextx = zext i16 %x to i32
   %zexty = zext i16 %x to i32
   %i0 = urem i32 %zextx, %zexty
   %r = trunc i32 %i0 to i16
   ret i16 %r
 }
 
+; Negative test
 define i16 @urem_big_const(i8 %x) {
 ; CHECK-LABEL: @urem_big_const(
 ; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[X:%.*]] to i32
 ; CHECK-NEXT:    [[DIV:%.*]] = urem i32 [[ZEXT]], 70000
 ; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i32 [[DIV]] to i16
 ; CHECK-NEXT:    ret i16 [[TRUNC]]
 ;
   %zext = zext i8 %x to i32
   %div = urem i32 %zext, 70000
   %trunc = trunc i32 %div to i16
   ret i16 %trunc
 }
 
 define <2 x i16> @urem_vector(<2 x i8> %x) {
 ; CHECK-LABEL: @urem_vector(
-; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[S:%.*]] = urem <2 x i32> [[Z]], <i32 4, i32 10>
-; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i32> [[S]] to <2 x i16>
-; CHECK-NEXT:    ret <2 x i16> [[T]]
+; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i16>
+; CHECK-NEXT:    [[S:%.*]] = urem <2 x i16> [[Z]], <i16 4, i16 10>
+; CHECK-NEXT:    ret <2 x i16> [[S]]
 ;
   %z = zext <2 x i8> %x to <2 x i32>
   %s = urem <2 x i32> %z, <i32 4, i32 10>
   %t = trunc <2 x i32> %s to <2 x i16>
   ret <2 x i16> %t
 }
 
+; Negative test: can only fold to <2 x i16>, requiring new vector type
 define <2 x i8> @urem_vector_need_new_vector_type(<2 x i8> %x) {
 ; CHECK-LABEL: @urem_vector_need_new_vector_type(
 ; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
 ; CHECK-NEXT:    [[S:%.*]] = urem <2 x i32> [[Z]], <i32 500, i32 10>
 ; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i32> [[S]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[T]]
 ;
   %z = zext <2 x i8> %x to <2 x i32>
   %s = urem <2 x i32> %z, <i32 500, i32 10>
   %t = trunc <2 x i32> %s to <2 x i8>
   ret <2 x i8> %t
 }
 
+; Negative test
 define <2 x i16> @urem_vector_big_const(<2 x i8> %x) {
 ; CHECK-LABEL: @urem_vector_big_const(
 ; CHECK-NEXT:    [[Z:%.*]] = zext <2 x i8> [[X:%.*]] to <2 x i32>
 ; CHECK-NEXT:    [[S:%.*]] = urem <2 x i32> [[Z]], <i32 16, i32 70000>
 ; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i32> [[S]] to <2 x i16>
 ; CHECK-NEXT:    ret <2 x i16> [[T]]
 ;
   %z = zext <2 x i8> %x to <2 x i32>
   %s = urem <2 x i32> %z, <i32 16, i32 70000>
   %t = trunc <2 x i32> %s to <2 x i16>
   ret <2 x i16> %t
 }