[InstCombine] Treat lshr nneg as ashr in getBinOpsForFactorization (#75521)

This patch reinterprets `lshr nneg C, X` as `ashr nneg C, X` to allow
more factorization opportunities.
Fixes #70582.

Author: dtcxzyw

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -610,7 +610,7 @@ static Value *getIdentityValue(Instruction::BinaryOps Opcode, Value *V) {
 /// allow more factorization opportunities.
 static Instruction::BinaryOps
 getBinOpsForFactorization(Instruction::BinaryOps TopOpcode, BinaryOperator *Op,
-                          Value *&LHS, Value *&RHS) {
+                          Value *&LHS, Value *&RHS, BinaryOperator *OtherOp) {
   assert(Op && "Expected a binary operator");
   LHS = Op->getOperand(0);
   RHS = Op->getOperand(1);
@@ -623,6 +623,13 @@ getBinOpsForFactorization(Instruction::BinaryOps TopOpcode, BinaryOperator *Op,
     }
     // TODO: We can add other conversions e.g. shr => div etc.
   }
+  if (Instruction::isBitwiseLogicOp(TopOpcode)) {
+    if (OtherOp && OtherOp->getOpcode() == Instruction::AShr &&
+        match(Op, m_LShr(m_NonNegative(), m_Value()))) {
+      // lshr nneg C, X --> ashr nneg C, X
+      return Instruction::AShr;
+    }
+  }
   return Op->getOpcode();
 }
 
@@ -963,9 +970,9 @@ Value *InstCombinerImpl::tryFactorizationFolds(BinaryOperator &I) {
   Instruction::BinaryOps LHSOpcode, RHSOpcode;
 
   if (Op0)
-    LHSOpcode = getBinOpsForFactorization(TopLevelOpcode, Op0, A, B);
+    LHSOpcode = getBinOpsForFactorization(TopLevelOpcode, Op0, A, B, Op1);
   if (Op1)
-    RHSOpcode = getBinOpsForFactorization(TopLevelOpcode, Op1, C, D);
+    RHSOpcode = getBinOpsForFactorization(TopLevelOpcode, Op1, C, D, Op0);
 
   // The instruction has the form "(A op' B) op (C op' D)".  Try to factorize
   // a common term.

llvm/test/Transforms/InstCombine/xor.ll

@@ -1395,3 +1395,61 @@ define i32 @ctlz_pow2_wrong_const(i32 %x) {
   %r = xor i32 %z, 30
   ret i32 %r
 }
+
+; Tests from PR70582
+define i32 @tryFactorization_xor_ashr_lshr(i32 %a) {
+; CHECK-LABEL: @tryFactorization_xor_ashr_lshr(
+; CHECK-NEXT:    [[XOR:%.*]] = ashr i32 -8, [[A:%.*]]
+; CHECK-NEXT:    ret i32 [[XOR]]
+;
+  %not = ashr i32 -3, %a
+  %shr1 = lshr i32 5, %a
+  %xor = xor i32 %not, %shr1
+  ret i32 %xor
+}
+
+define i32 @tryFactorization_xor_lshr_ashr(i32 %a) {
+; CHECK-LABEL: @tryFactorization_xor_lshr_ashr(
+; CHECK-NEXT:    [[XOR:%.*]] = ashr i32 -8, [[A:%.*]]
+; CHECK-NEXT:    ret i32 [[XOR]]
+;
+  %not = ashr i32 -3, %a
+  %shr1 = lshr i32 5, %a
+  %xor = xor i32 %shr1, %not
+  ret i32 %xor
+}
+
+define i32 @tryFactorization_xor_ashr_lshr_negative_lhs(i32 %a) {
+; CHECK-LABEL: @tryFactorization_xor_ashr_lshr_negative_lhs(
+; CHECK-NEXT:    [[NOT:%.*]] = ashr i32 -3, [[A:%.*]]
+; CHECK-NEXT:    [[SHR1:%.*]] = lshr i32 -5, [[A]]
+; CHECK-NEXT:    [[XOR:%.*]] = xor i32 [[NOT]], [[SHR1]]
+; CHECK-NEXT:    ret i32 [[XOR]]
+;
+  %not = ashr i32 -3, %a
+  %shr1 = lshr i32 -5, %a
+  %xor = xor i32 %not, %shr1
+  ret i32 %xor
+}
+
+define i32 @tryFactorization_xor_lshr_lshr(i32 %a) {
+; CHECK-LABEL: @tryFactorization_xor_lshr_lshr(
+; CHECK-NEXT:    [[XOR:%.*]] = lshr i32 -8, [[A:%.*]]
+; CHECK-NEXT:    ret i32 [[XOR]]
+;
+  %not = lshr i32 -3, %a
+  %shr1 = lshr i32 5, %a
+  %xor = xor i32 %not, %shr1
+  ret i32 %xor
+}
+
+define i32 @tryFactorization_xor_ashr_ashr(i32 %a) {
+; CHECK-LABEL: @tryFactorization_xor_ashr_ashr(
+; CHECK-NEXT:    [[XOR:%.*]] = lshr i32 6, [[A:%.*]]
+; CHECK-NEXT:    ret i32 [[XOR]]
+;
+  %not = ashr i32 -3, %a
+  %shr1 = ashr i32 -5, %a
+  %xor = xor i32 %not, %shr1
+  ret i32 %xor
+}