[llvm] [InstCombine] Fold shift of a constant into a reverse shift (PR #192982)

[Piotr Fusik via llvm-commits]

================
@@ -460,42 +460,69 @@ Instruction *InstCombinerImpl::commonShiftTransforms(BinaryOperator &I) {
 
   unsigned BitWidth = Ty->getScalarSizeInBits();
 
-  const APInt *AC, *AddC;
-  // Try to pre-shift a constant shifted by a variable amount added with a
-  // negative number:
-  // C << (X - AddC) --> (C >> AddC) << X
-  // and
-  // C >> (X - AddC) --> (C << AddC) >> X
-  if (match(Op0, m_APInt(AC)) && match(Op1, m_Add(m_Value(A), m_APInt(AddC))) &&
-      AddC->isNegative() && (-*AddC).ult(BitWidth)) {
+  const APInt *AC;
+  if (match(Op0, m_APInt(AC))) {
     assert(!AC->isZero() && "Expected simplify of shifted zero");
-    unsigned PosOffset = (-*AddC).getZExtValue();
 
-    auto isSuitableForPreShift = [PosOffset, &I, AC]() {
-      switch (I.getOpcode()) {
-      default:
-        return false;
-      case Instruction::Shl:
-        return (I.hasNoSignedWrap() || I.hasNoUnsignedWrap()) &&
-               AC->eq(AC->lshr(PosOffset).shl(PosOffset));
-      case Instruction::LShr:
-        return I.isExact() && AC->eq(AC->shl(PosOffset).lshr(PosOffset));
-      case Instruction::AShr:
-        return I.isExact() && AC->eq(AC->shl(PosOffset).ashr(PosOffset));
+    // Try to pre-shift a constant shifted by a variable amount added with a
+    // negative number:
+    // C << (X - AddC) --> (C >> AddC) << X
+    // and
+    // C >> (X - AddC) --> (C << AddC) >> X
+    const APInt *AddC;
+    if (match(Op1, m_Add(m_Value(A), m_APInt(AddC))) && AddC->isNegative() &&
+        (-*AddC).ult(BitWidth)) {
+      unsigned PosOffset = (-*AddC).getZExtValue();
+
+      auto isSuitableForPreShift = [PosOffset, &I, AC]() {
+        switch (I.getOpcode()) {
+        default:
+          return false;
+        case Instruction::Shl:
+          return (I.hasNoSignedWrap() || I.hasNoUnsignedWrap()) &&
+                 AC->eq(AC->lshr(PosOffset).shl(PosOffset));
+        case Instruction::LShr:
+          return I.isExact() && AC->eq(AC->shl(PosOffset).lshr(PosOffset));
+        case Instruction::AShr:
+          return I.isExact() && AC->eq(AC->shl(PosOffset).ashr(PosOffset));
+        }
+      };
+      if (isSuitableForPreShift()) {
+        Constant *NewC = ConstantInt::get(Ty, I.getOpcode() == Instruction::Shl
+                                                  ? AC->lshr(PosOffset)
+                                                  : AC->shl(PosOffset));
+        BinaryOperator *NewShiftOp =
+            BinaryOperator::Create(I.getOpcode(), NewC, A);
+        if (I.getOpcode() == Instruction::Shl) {
+          NewShiftOp->setHasNoUnsignedWrap(I.hasNoUnsignedWrap());
+        } else {
+          NewShiftOp->setIsExact();
+        }
+        return NewShiftOp;
       }
-    };
-    if (isSuitableForPreShift()) {
-      Constant *NewC = ConstantInt::get(Ty, I.getOpcode() == Instruction::Shl
-                                                ? AC->lshr(PosOffset)
-                                                : AC->shl(PosOffset));
-      BinaryOperator *NewShiftOp =
-          BinaryOperator::Create(I.getOpcode(), NewC, A);
+    }
+
+    // C1 << (C2 - X) -> (C1 << C2) >> X
+    // C1 >> (C2 - X) -> (C1 >> C2) << X
+    // X must be >=0 (checked by NUWSub).
+    // Also match (X ^ C2) if equivalent to (C2 - X).
+    uint64_t C2;
+    Value *X;
+    if (match(Op1, m_NUWSub(m_ConstantInt(C2), m_Value(X))) ||
+        (match(Op1, m_Xor(m_Value(X), m_ConstantInt(C2))) &&
+         (C2 | computeKnownBits(X, &I).Zero).isAllOnes())) {
       if (I.getOpcode() == Instruction::Shl) {
-        NewShiftOp->setHasNoUnsignedWrap(I.hasNoUnsignedWrap());
-      } else {
-        NewShiftOp->setIsExact();
+        if (AC->countl_zero() >= C2)
+          return BinaryOperator::CreateExactLShr(
+              ConstantInt::get(Ty, AC->shl(C2)), X);
+        if (AC->countl_one() > C2)
+          return BinaryOperator::CreateExactAShr(
+              ConstantInt::get(Ty, AC->shl(C2)), X);
+      } else if (AC->countr_zero() >= C2) {
+        APInt C =
+            I.getOpcode() == Instruction::LShr ? AC->lshr(C2) : AC->ashr(C2);
+        return BinaryOperator::CreateShl(ConstantInt::get(Ty, C), X);
----------------
pfusik wrote:

Indeed. If C1>=0, we can set both nuw and nsw.

https://github.com/llvm/llvm-project/pull/192982