[InstCombine] try harder to narrow bitwise logic with cast operands

This works with any logic + extend:
https://alive2.llvm.org/ce/z/vzsqQD

The motivating case is from issue #56294, but that's still not optimal
(it should simplify completely).

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

@@ -1430,12 +1430,34 @@ Instruction *InstCombinerImpl::foldCastedBitwiseLogic(BinaryOperator &I) {
   if (!Cast1)
     return nullptr;
 
-  // Both operands of the logic operation are casts. The casts must be of the
+  // Both operands of the logic operation are casts. The casts must be the
-  // same type for reduction.
+  // same kind for reduction.
-  auto CastOpcode = Cast0->getOpcode();
+  Instruction::CastOps CastOpcode = Cast0->getOpcode();
-  if (CastOpcode != Cast1->getOpcode() || SrcTy != Cast1->getSrcTy())
+  if (CastOpcode != Cast1->getOpcode())
     return nullptr;
 
+  // If the source types do not match, but the casts are matching extends, we
+  // can still narrow the logic op.
+  if (SrcTy != Cast1->getSrcTy()) {
+    Value *X, *Y;
+    if (match(Cast0, m_OneUse(m_ZExtOrSExt(m_Value(X)))) &&
+        match(Cast1, m_OneUse(m_ZExtOrSExt(m_Value(Y))))) {
+      // Cast the narrower source to the wider source type.
+      unsigned XNumBits = X->getType()->getScalarSizeInBits();
+      unsigned YNumBits = Y->getType()->getScalarSizeInBits();
+      if (XNumBits < YNumBits)
+        X = Builder.CreateCast(CastOpcode, X, Y->getType());
+      else
+        Y = Builder.CreateCast(CastOpcode, Y, X->getType());
+      // Do the logic op in the intermediate width, then widen more.
+      Value *NarrowLogic = Builder.CreateBinOp(LogicOpc, X, Y);
+      return CastInst::Create(CastOpcode, NarrowLogic, DestTy);
+    }
+
+    // Give up for other cast opcodes.
+    return nullptr;
+  }
+
   Value *Cast0Src = Cast0->getOperand(0);
   Value *Cast1Src = Cast1->getOperand(0);
 

llvm/test/Transforms/InstCombine/and-xor-or.ll

@@ -4209,9 +4209,9 @@ define i32 @trunc_trunc_xor_uses(i65 %x, i65 %y) {
 define i16 @and_zext_zext(i8 %x, i4 %y) {
 ; CHECK-LABEL: define {{[^@]+}}@and_zext_zext
 ; CHECK-SAME: (i8 [[X:%.*]], i4 [[Y:%.*]]) {
-; CHECK-NEXT:    [[ZX:%.*]] = zext i8 [[X]] to i16
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i4 [[Y]] to i8
-; CHECK-NEXT:    [[ZY:%.*]] = zext i4 [[Y]] to i16
+; CHECK-NEXT:    [[TMP2:%.*]] = and i8 [[TMP1]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = and i16 [[ZX]], [[ZY]]
+; CHECK-NEXT:    [[R:%.*]] = zext i8 [[TMP2]] to i16
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %zx = zext i8 %x to i16
@@ -4223,9 +4223,9 @@ define i16 @and_zext_zext(i8 %x, i4 %y) {
 define i16 @or_zext_zext(i8 %x, i4 %y) {
 ; CHECK-LABEL: define {{[^@]+}}@or_zext_zext
 ; CHECK-SAME: (i8 [[X:%.*]], i4 [[Y:%.*]]) {
-; CHECK-NEXT:    [[ZX:%.*]] = zext i8 [[X]] to i16
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i4 [[Y]] to i8
-; CHECK-NEXT:    [[ZY:%.*]] = zext i4 [[Y]] to i16
+; CHECK-NEXT:    [[TMP2:%.*]] = or i8 [[TMP1]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = or i16 [[ZY]], [[ZX]]
+; CHECK-NEXT:    [[R:%.*]] = zext i8 [[TMP2]] to i16
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %zx = zext i8 %x to i16
@@ -4237,9 +4237,9 @@ define i16 @or_zext_zext(i8 %x, i4 %y) {
 define <2 x i16> @xor_zext_zext(<2 x i8> %x, <2 x i4> %y) {
 ; CHECK-LABEL: define {{[^@]+}}@xor_zext_zext
 ; CHECK-SAME: (<2 x i8> [[X:%.*]], <2 x i4> [[Y:%.*]]) {
-; CHECK-NEXT:    [[ZX:%.*]] = zext <2 x i8> [[X]] to <2 x i16>
+; CHECK-NEXT:    [[TMP1:%.*]] = zext <2 x i4> [[Y]] to <2 x i8>
-; CHECK-NEXT:    [[ZY:%.*]] = zext <2 x i4> [[Y]] to <2 x i16>
+; CHECK-NEXT:    [[TMP2:%.*]] = xor <2 x i8> [[TMP1]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = xor <2 x i16> [[ZX]], [[ZY]]
+; CHECK-NEXT:    [[R:%.*]] = zext <2 x i8> [[TMP2]] to <2 x i16>
 ; CHECK-NEXT:    ret <2 x i16> [[R]]
 ;
   %zx = zext <2 x i8> %x to <2 x i16>
@@ -4251,9 +4251,9 @@ define <2 x i16> @xor_zext_zext(<2 x i8> %x, <2 x i4> %y) {
 define i16 @and_sext_sext(i8 %x, i4 %y) {
 ; CHECK-LABEL: define {{[^@]+}}@and_sext_sext
 ; CHECK-SAME: (i8 [[X:%.*]], i4 [[Y:%.*]]) {
-; CHECK-NEXT:    [[SX:%.*]] = sext i8 [[X]] to i16
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i4 [[Y]] to i8
-; CHECK-NEXT:    [[SY:%.*]] = sext i4 [[Y]] to i16
+; CHECK-NEXT:    [[TMP2:%.*]] = and i8 [[TMP1]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = and i16 [[SY]], [[SX]]
+; CHECK-NEXT:    [[R:%.*]] = sext i8 [[TMP2]] to i16
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %sx = sext i8 %x to i16
@@ -4265,9 +4265,9 @@ define i16 @and_sext_sext(i8 %x, i4 %y) {
 define i16 @or_sext_sext(i8 %x, i4 %y) {
 ; CHECK-LABEL: define {{[^@]+}}@or_sext_sext
 ; CHECK-SAME: (i8 [[X:%.*]], i4 [[Y:%.*]]) {
-; CHECK-NEXT:    [[SX:%.*]] = sext i8 [[X]] to i16
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i4 [[Y]] to i8
-; CHECK-NEXT:    [[SY:%.*]] = sext i4 [[Y]] to i16
+; CHECK-NEXT:    [[TMP2:%.*]] = or i8 [[TMP1]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = or i16 [[SX]], [[SY]]
+; CHECK-NEXT:    [[R:%.*]] = sext i8 [[TMP2]] to i16
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %sx = sext i8 %x to i16
@@ -4279,9 +4279,9 @@ define i16 @or_sext_sext(i8 %x, i4 %y) {
 define i16 @xor_sext_sext(i8 %x, i4 %y) {
 ; CHECK-LABEL: define {{[^@]+}}@xor_sext_sext
 ; CHECK-SAME: (i8 [[X:%.*]], i4 [[Y:%.*]]) {
-; CHECK-NEXT:    [[SX:%.*]] = sext i8 [[X]] to i16
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i4 [[Y]] to i8
-; CHECK-NEXT:    [[SY:%.*]] = sext i4 [[Y]] to i16
+; CHECK-NEXT:    [[TMP2:%.*]] = xor i8 [[TMP1]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = xor i16 [[SX]], [[SY]]
+; CHECK-NEXT:    [[R:%.*]] = sext i8 [[TMP2]] to i16
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %sx = sext i8 %x to i16
@@ -4290,6 +4290,8 @@ define i16 @xor_sext_sext(i8 %x, i4 %y) {
   ret i16 %r
 }
 
+; negative test - mismatched casts
+
 define i16 @and_zext_sext(i8 %x, i4 %y) {
 ; CHECK-LABEL: define {{[^@]+}}@and_zext_sext
 ; CHECK-SAME: (i8 [[X:%.*]], i4 [[Y:%.*]]) {
@@ -4304,6 +4306,8 @@ define i16 @and_zext_sext(i8 %x, i4 %y) {
   ret i16 %r
 }
 
+; negative test - don't create an extra instruction
+
 define i32 @and_zext_zext_use1(i8 %x, i4 %y) {
 ; CHECK-LABEL: define {{[^@]+}}@and_zext_zext_use1
 ; CHECK-SAME: (i8 [[X:%.*]], i4 [[Y:%.*]]) {
@@ -4320,6 +4324,8 @@ define i32 @and_zext_zext_use1(i8 %x, i4 %y) {
   ret i32 %r
 }
 
+; negative test - don't create an extra instruction
+
 define i32 @or_sext_sext_use1(i8 %x, i4 %y) {
 ; CHECK-LABEL: define {{[^@]+}}@or_sext_sext_use1
 ; CHECK-SAME: (i8 [[X:%.*]], i4 [[Y:%.*]]) {
@@ -4335,3 +4341,21 @@ define i32 @or_sext_sext_use1(i8 %x, i4 %y) {
   %r = or i32 %sx, %sy
   ret i32 %r
 }
+
+define i1 @PR56294(i8 %x) {
+; CHECK-LABEL: define {{[^@]+}}@PR56294
+; CHECK-SAME: (i8 [[X:%.*]]) {
+; CHECK-NEXT:    [[T2:%.*]] = icmp eq i8 [[X]], 2
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i1 [[T2]] to i8
+; CHECK-NEXT:    [[TMP2:%.*]] = and i8 [[TMP1]], [[X]]
+; CHECK-NEXT:    [[T7:%.*]] = icmp ne i8 [[TMP2]], 0
+; CHECK-NEXT:    ret i1 [[T7]]
+;
+  %t2 = icmp eq i8 %x, 2
+  %t3 = and i8 %x, 1
+  %t4 = zext i1 %t2 to i32
+  %t5 = zext i8 %t3 to i32
+  %t6 = and i32 %t4, %t5
+  %t7 = icmp ne i32 %t6, 0
+  ret i1 %t7
+}