[InstCombine] Fold umax/umin(nuw_shl(z, x), nuw_shl(z, y)) -> nuw_shl(z, umax/umin(x, y)) and umax/umin(nuw_shl(x, z), nuw_shl(y, z)) -> nuw_shl(umax/umin(x, y), z)

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

@@ -1554,7 +1554,11 @@ static bool leftDistributesOverRight(Instruction::BinaryOps LOp, bool HasNUW,
   switch (ROp) {
   case Intrinsic::umax:
   case Intrinsic::umin:
-    return HasNUW && LOp == Instruction::Add;
+    if (HasNUW && LOp == Instruction::Add)
+      return true;
+    if (HasNUW && LOp == Instruction::Shl)
+      return true;
+    return false;
   case Intrinsic::smax:
   case Intrinsic::smin:
     return HasNSW && LOp == Instruction::Add;
@@ -1592,29 +1596,34 @@ foldIntrinsicUsingDistributiveLaws(IntrinsicInst *II,
   if (!leftDistributesOverRight(InnerOpcode, HasNUW, HasNSW, TopLevelOpcode))
     return nullptr;
 
-  assert(II->isCommutative() && Op0->isCommutative() &&
-         "Only inner and outer commutative op codes are supported.");
-
   Value *A = Op0->getOperand(0);
   Value *B = Op0->getOperand(1);
   Value *C = Op1->getOperand(0);
   Value *D = Op1->getOperand(1);
 
-  // Attempts to swap variables such that A always equals C
-  if (A != C && A != D)
-    std::swap(A, B);
-  if (A == C || A == D) {
-    if (A != C)
-      std::swap(C, D);
+  // Attempts to swap variables such that A equals C or B equals D,
+  // if the inner operation is commutative.
+  if (Op0->isCommutative() && A != C && B != D && A == D)
+    std::swap(C, D);
+
+  if (A != C && B != D)
+    return nullptr;
+
+  BinaryOperator *NewBinop;
+  if (A == C) {
     Value *NewIntrinsic = Builder.CreateBinaryIntrinsic(TopLevelOpcode, B, D);
-    BinaryOperator *NewBinop =
-        cast<BinaryOperator>(Builder.CreateBinOp(InnerOpcode, NewIntrinsic, A));
-    NewBinop->setHasNoSignedWrap(HasNSW);
-    NewBinop->setHasNoUnsignedWrap(HasNUW);
-    return NewBinop;
+    NewBinop =
+        cast<BinaryOperator>(Builder.CreateBinOp(InnerOpcode, A, NewIntrinsic));
+  } else { // B == D
+    Value *NewIntrinsic = Builder.CreateBinaryIntrinsic(TopLevelOpcode, A, C);
+    NewBinop =
+        cast<BinaryOperator>(Builder.CreateBinOp(InnerOpcode, NewIntrinsic, B));
   }
 
-  return nullptr;
+  NewBinop->setHasNoUnsignedWrap(HasNUW);
+  NewBinop->setHasNoSignedWrap(HasNSW);
+
+  return NewBinop;
 }
 
 /// CallInst simplification. This mostly only handles folding of intrinsic
@@ -1887,6 +1896,7 @@ Instruction *InstCombinerImpl::visitCallInst(CallInst &CI) {
       if (Instruction *I = foldMaxMulShift(I1, I0))
         return I;
     }
+
     // If both operands of unsigned min/max are sign-extended, it is still ok
     // to narrow the operation.
     [[fallthrough]];

llvm/test/Transforms/InstCombine/div-shift.ll

@@ -144,11 +144,9 @@ define i8 @udiv_umin_(i8 %x, i8 %y, i8 %z) {
 ; Negative test, extra use
 define i8 @udiv_umin_extra_use(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @udiv_umin_extra_use(
-; CHECK-NEXT:    [[Y2:%.*]] = shl nuw i8 1, [[Y:%.*]]
-; CHECK-NEXT:    [[Z2:%.*]] = shl nuw i8 1, [[Z:%.*]]
-; CHECK-NEXT:    [[M:%.*]] = call i8 @llvm.umin.i8(i8 [[Y2]], i8 [[Z2]])
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.umin.i8(i8 [[Y:%.*]], i8 [[Z:%.*]])
+; CHECK-NEXT:    [[M:%.*]] = shl nuw i8 1, [[TMP1]]
 ; CHECK-NEXT:    call void @use(i8 [[M]])
-; CHECK-NEXT:    [[TMP1:%.*]] = call range(i8 0, 9) i8 @llvm.cttz.i8(i8 [[M]], i1 true)
 ; CHECK-NEXT:    [[D:%.*]] = lshr i8 [[X:%.*]], [[TMP1]]
 ; CHECK-NEXT:    ret i8 [[D]]
 ;