[InstCombine] simplify average of lsb (#95684)

close: #94737
alive2: https://alive2.llvm.org/ce/z/WF_7mX

In this patch, we combine `(X + Y) / 2` into `(X & Y)` only when both X
and Y are less than or equal to 1.

Author: c8ef

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

@@ -1284,6 +1284,12 @@ Instruction *InstCombinerImpl::visitLShr(BinaryOperator &I) {
     return NewSub;
   }
 
+  // Fold (X + Y) / 2 --> (X & Y) iff (X u<= 1) && (Y u<= 1)
+  if (match(Op0, m_Add(m_Value(X), m_Value(Y))) && match(Op1, m_One()) &&
+      computeKnownBits(X, /*Depth=*/0, &I).countMaxActiveBits() <= 1 &&
+      computeKnownBits(Y, /*Depth=*/0, &I).countMaxActiveBits() <= 1)
+    return BinaryOperator::CreateAnd(X, Y);
+
   // (sub nuw X, (Y << nuw Z)) >>u exact Z --> (X >>u exact Z) sub nuw Y
   if (I.isExact() &&
       match(Op0, m_OneUse(m_NUWSub(m_Value(X),

llvm/test/Transforms/InstCombine/avg-lsb.ll

@@ -0,0 +1,46 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i8 @avg_lsb(i8 %a, i8 %b) {
+; CHECK-LABEL: define i8 @avg_lsb(
+; CHECK-SAME: i8 [[A:%.*]], i8 [[B:%.*]]) {
+; CHECK-NEXT:    [[REM:%.*]] = and i8 [[A]], 1
+; CHECK-NEXT:    [[DIV2:%.*]] = and i8 [[REM]], [[B]]
+; CHECK-NEXT:    ret i8 [[DIV2]]
+;
+  %rem = and i8 %a, 1
+  %rem1 = and i8 %b, 1
+  %add = add nuw nsw i8 %rem1, %rem
+  %div2 = lshr i8 %add, 1
+  ret i8 %div2
+}
+
+define i8 @avg_lsb_mismatch(i8 %a, i8 %b) {
+; CHECK-LABEL: define i8 @avg_lsb_mismatch(
+; CHECK-SAME: i8 [[A:%.*]], i8 [[B:%.*]]) {
+; CHECK-NEXT:    [[REM:%.*]] = and i8 [[A]], 1
+; CHECK-NEXT:    [[REM1:%.*]] = and i8 [[B]], 3
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i8 [[REM1]], [[REM]]
+; CHECK-NEXT:    [[DIV2:%.*]] = lshr i8 [[ADD]], 1
+; CHECK-NEXT:    ret i8 [[DIV2]]
+;
+  %rem = and i8 %a, 1
+  %rem1 = and i8 %b, 3
+  %add = add nuw nsw i8 %rem1, %rem
+  %div2 = lshr i8 %add, 1
+  ret i8 %div2
+}
+
+define <2 x i8> @avg_lsb_vector(<2 x i8> %a, <2 x i8> %b) {
+; CHECK-LABEL: define <2 x i8> @avg_lsb_vector(
+; CHECK-SAME: <2 x i8> [[A:%.*]], <2 x i8> [[B:%.*]]) {
+; CHECK-NEXT:    [[REM:%.*]] = and <2 x i8> [[A]], <i8 1, i8 1>
+; CHECK-NEXT:    [[DIV2:%.*]] = and <2 x i8> [[REM]], [[B]]
+; CHECK-NEXT:    ret <2 x i8> [[DIV2]]
+;
+  %rem = and <2 x i8> %a, <i8 1, i8 1>
+  %rem1 = and <2 x i8> %b, <i8 1, i8 1>
+  %add = add nuw nsw <2 x i8> %rem1, %rem
+  %div2 = lshr <2 x i8> %add, <i8 1, i8 1>
+  ret <2 x i8> %div2
+}

llvm/test/Transforms/InstCombine/lshr.ll

@@ -1390,8 +1390,8 @@ define i8 @not_signbit_trunc(i16 %x) {
 
 define i2 @bool_add_lshr(i1 %a, i1 %b) {
 ; CHECK-LABEL: @bool_add_lshr(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i1 [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[LSHR:%.*]] = zext i1 [[TMP1]] to i2
+; CHECK-NEXT:    [[LSHR1:%.*]] = and i1 [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[LSHR:%.*]] = zext i1 [[LSHR1]] to i2
 ; CHECK-NEXT:    ret i2 [[LSHR]]
 ;
   %zext.a = zext i1 %a to i2
@@ -1436,8 +1436,8 @@ define i2 @bool_add_ashr(i1 %a, i1 %b) {
 
 define <2 x i8> @bool_add_lshr_vec(<2 x i1> %a, <2 x i1> %b) {
 ; CHECK-LABEL: @bool_add_lshr_vec(
-; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i1> [[A:%.*]], [[B:%.*]]
-; CHECK-NEXT:    [[LSHR:%.*]] = zext <2 x i1> [[TMP1]] to <2 x i8>
+; CHECK-NEXT:    [[LSHR1:%.*]] = and <2 x i1> [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[LSHR:%.*]] = zext <2 x i1> [[LSHR1]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[LSHR]]
 ;
   %zext.a = zext <2 x i1> %a to <2 x i8>
@@ -1453,8 +1453,7 @@ define i32 @bool_add_lshr_uses(i1 %a, i1 %b) {
 ; CHECK-NEXT:    call void @use(i32 [[ZEXT_A]])
 ; CHECK-NEXT:    [[ZEXT_B:%.*]] = zext i1 [[B:%.*]] to i32
 ; CHECK-NEXT:    call void @use(i32 [[ZEXT_B]])
-; CHECK-NEXT:    [[TMP1:%.*]] = and i1 [[A]], [[B]]
-; CHECK-NEXT:    [[LSHR:%.*]] = zext i1 [[TMP1]] to i32
+; CHECK-NEXT:    [[LSHR:%.*]] = and i32 [[ZEXT_A]], [[ZEXT_B]]
 ; CHECK-NEXT:    ret i32 [[LSHR]]
 ;
   %zext.a = zext i1 %a to i32
@@ -1473,8 +1472,7 @@ define i32 @bool_add_lshr_uses2(i1 %a, i1 %b) {
 ; CHECK-NEXT:    call void @use(i32 [[ZEXT_B]])
 ; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i32 [[ZEXT_A]], [[ZEXT_B]]
 ; CHECK-NEXT:    call void @use(i32 [[ADD]])
-; CHECK-NEXT:    [[TMP1:%.*]] = and i1 [[A]], [[B]]
-; CHECK-NEXT:    [[LSHR:%.*]] = zext i1 [[TMP1]] to i32
+; CHECK-NEXT:    [[LSHR:%.*]] = and i32 [[ZEXT_A]], [[ZEXT_B]]
 ; CHECK-NEXT:    ret i32 [[LSHR]]
 ;
   %zext.a = zext i1 %a to i32
@@ -1496,7 +1494,7 @@ define i32 @bool_add_lshr_uses3(i1 %a, i1 %b) {
 ; CHECK-NEXT:    call void @use(i32 [[ZEXT_B]])
 ; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i32 [[ZEXT_A]], [[ZEXT_B]]
 ; CHECK-NEXT:    call void @use(i32 [[ADD]])
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[ADD]], 1
+; CHECK-NEXT:    [[LSHR:%.*]] = and i32 [[ZEXT_A]], [[ZEXT_B]]
 ; CHECK-NEXT:    ret i32 [[LSHR]]
 ;
   %zext.a = zext i1 %a to i32