[InstCombine] Treat lshr nneg as ashr in getBinOpsForFactorization

[dtcxzyw]

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

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

Fixes #70582.

---

Full diff: https://github.com/llvm/llvm-project/pull/75521.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+10-3)
• (modified) llvm/test/Transforms/InstCombine/xor.ll (+45)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 6ac1fdb9252bff..a7ddadc25de43c 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/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.
diff --git a/llvm/test/Transforms/InstCombine/xor.ll b/llvm/test/Transforms/InstCombine/xor.ll
index 7c61401fca07c7..388755b832096c 100644
--- a/llvm/test/Transforms/InstCombine/xor.ll
+++ b/llvm/test/Transforms/InstCombine/xor.ll
@@ -1395,3 +1395,48 @@ 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_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
+}

[nikic]

LGTM, but please add a negative test with ashr + lshr with negative LHS.