[InstComb] Fold shr ([x]or x, y), x -> shr y, x

Proof: https://alive2.llvm.org/ce/z/yWCmMd

Co-authored-by: John Regehr <[email protected]>

Author: artagnon

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

@@ -512,19 +512,27 @@ Instruction *InstCombinerImpl::commonShiftTransforms(BinaryOperator &I) {
   if (match(Op1, m_Or(m_Value(), m_SpecificInt(BitWidth - 1))))
     return replaceOperand(I, 1, ConstantInt::get(Ty, BitWidth - 1));
 
-  Instruction *CmpIntr;
-  if ((I.getOpcode() == Instruction::LShr ||
-       I.getOpcode() == Instruction::AShr) &&
-      match(Op0, m_OneUse(m_Instruction(CmpIntr))) &&
-      isa<CmpIntrinsic>(CmpIntr) &&
-      match(Op1, m_SpecificInt(Ty->getScalarSizeInBits() - 1))) {
-    Value *Cmp =
-        Builder.CreateICmp(cast<CmpIntrinsic>(CmpIntr)->getLTPredicate(),
-                           CmpIntr->getOperand(0), CmpIntr->getOperand(1));
-    return CastInst::Create(I.getOpcode() == Instruction::LShr
-                                ? Instruction::ZExt
-                                : Instruction::SExt,
-                            Cmp, Ty);
+  if (I.getOpcode() == Instruction::LShr ||
+      I.getOpcode() == Instruction::AShr) {
+    // (l|a)shr ([x]or x, y), x --> (l|a)shr y, x
+    Value *X;
+    if (match(Op0, m_CombineOr(m_Xor(m_Value(X), m_Value(Y)),
+                               m_Or(m_Value(X), m_Value(Y)))) &&
+        Op1 == X)
+      return BinaryOperator::Create(I.getOpcode(), Y, X);
+
+    Instruction *CmpIntr;
+    if (match(Op0, m_OneUse(m_Instruction(CmpIntr))) &&
+        isa<CmpIntrinsic>(CmpIntr) &&
+        match(Op1, m_SpecificInt(Ty->getScalarSizeInBits() - 1))) {
+      Value *Cmp =
+          Builder.CreateICmp(cast<CmpIntrinsic>(CmpIntr)->getLTPredicate(),
+                             CmpIntr->getOperand(0), CmpIntr->getOperand(1));
+      return CastInst::Create(I.getOpcode() == Instruction::LShr
+                                  ? Instruction::ZExt
+                                  : Instruction::SExt,
+                              Cmp, Ty);
+    }
   }
 
   return nullptr;

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

@@ -550,8 +550,7 @@ define <2 x i64> @lshr_sub_poison(<2 x i64> %x, <2 x i64> %py) {
 
 define i32 @ashr_xor_operand_match(i32 %x, i32 %y) {
 ; CHECK-LABEL: @ashr_xor_operand_match(
-; CHECK-NEXT:    [[R:%.*]] = xor i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[RET:%.*]] = ashr i32 [[R]], [[X]]
+; CHECK-NEXT:    [[RET:%.*]] = ashr i32 [[R:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
   %r = xor i32 %x, %y
@@ -572,8 +571,7 @@ define i32 @ashr_xor_operand_mismtach(i32 %x, i32 %y) {
 
 define i32 @lshr_xor_operand_match(i32 %x, i32 %y) {
 ; CHECK-LABEL: @lshr_xor_operand_match(
-; CHECK-NEXT:    [[R:%.*]] = xor i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[RET:%.*]] = lshr i32 [[R]], [[X]]
+; CHECK-NEXT:    [[RET:%.*]] = lshr i32 [[R:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
   %r = xor i32 %x, %y
@@ -594,8 +592,7 @@ define i32 @lshr_xor_operand_mismtach(i32 %x, i32 %y) {
 
 define i32 @ashr_or_operand_match(i32 %x, i32 %y) {
 ; CHECK-LABEL: @ashr_or_operand_match(
-; CHECK-NEXT:    [[R:%.*]] = or i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[RET:%.*]] = ashr i32 [[R]], [[X]]
+; CHECK-NEXT:    [[RET:%.*]] = ashr i32 [[R:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
   %r = or i32 %x, %y
@@ -616,8 +613,7 @@ define i32 @ashr_or_operand_mismtach(i32 %x, i32 %y) {
 
 define i32 @lshr_or_operand_match(i32 %x, i32 %y) {
 ; CHECK-LABEL: @lshr_or_operand_match(
-; CHECK-NEXT:    [[R:%.*]] = or i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[RET:%.*]] = lshr i32 [[R]], [[X]]
+; CHECK-NEXT:    [[RET:%.*]] = lshr i32 [[R:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
   %r = or i32 %x, %y
@@ -638,10 +634,10 @@ define i32 @lshr_or_operand_mismtach(i32 %x, i32 %y) {
 
 define i32 @ashr_xor_operand_match_multiuse(i32 %x, i32 %y) {
 ; CHECK-LABEL: @ashr_xor_operand_match_multiuse(
-; CHECK-NEXT:    [[R:%.*]] = xor i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[Q:%.*]] = ashr i32 [[R]], [[X]]
+; CHECK-NEXT:    [[Q:%.*]] = ashr i32 [[R:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    [[RET:%.*]] = xor i32 [[R]], [[Q]]
-; CHECK-NEXT:    ret i32 [[RET]]
+; CHECK-NEXT:    [[RET1:%.*]] = xor i32 [[RET]], [[X]]
+; CHECK-NEXT:    ret i32 [[RET1]]
 ;
   %r = xor i32 %x, %y
   %q = ashr i32 %r, %x