[InstCombine] C0 >>{ashr, exact} (X - C1) --> (C0 << C1) >>{ashr} X

[nico-abram]

Summary

As described in https://github.com/llvm/llvm-project/issues/55016
https://alive2.llvm.org/ce/z/pax7DF

Also includes C0 >> (X - C1) --> (C0 << C1) >> X for positive C0 , which was mentioned in the gh issue as the more common case
https://alive2.llvm.org/ce/z/sRq7w9

[spatel]

1. Please write a more general proof for these transforms. See examples in the comments here: https://github.com/llvm/llvm-project/issues/54890
2. Does the lshr transform require "exact" or not?
3. We need several tests - including "negative" tests that show that the transform does NOT occur when the pre-conditions are not met. You can use the existing tests for the 'shl' transform as a template (those are just below the tests added in this patch).
4. It would be better to create a small helper function rather than duplicate code for these transforms - can this live under InstCombinerImpl::commonShiftTransforms()? Would it include the existing shl transform too?

[spatel]

This is an alternate proof that I came up with for the previous patch:
https://alive2.llvm.org/ce/z/pnXQYR

Would it be easier to make a common shift transform if we implemented it using count leading/trailing zeros?

[nico-abram]

In D124369#3471465, @spatel wrote:

1. Please write a more general proof for these transforms. See examples in the comments here: https://github.com/llvm/llvm-project/issues/54890
2. Does the lshr transform require "exact" or not?

Is this enough proof for the ashr case? https://alive2.llvm.org/ce/z/iLuDyE
EDIT: This formulation is much simpler: https://alive2.llvm.org/ce/z/pax7DF

I don't think the lshr without exact is correct (It becomes more poisonous according to alive2), and it looks like we can only do it for positive numbers (which makes sense to me): https://alive2.llvm.org/ce/z/sRq7w9

[nico-abram]

Updating the diff with the condition change and more tests. Will see about avoiding the duplication and moving into commonShiftTransforms

[nico-abram]

I collapsed all three into InstCombinerImpl::commonShiftTransforms, right below the pre-existing pre-shift combine.

Had to add setHasNoSignedWrap(false); for the Shl case. Otherwise the @shl_nsw_add_negative test case was generating a shl nsw i32 1 instead of shl i32 1. Not entirely sure why, possibly because I'm genering the Op with BinaryOperator::Create(I.getOpcode()?

[spatel]

In D124369#3472995, @nico-abram wrote:

I collapsed all three into InstCombinerImpl::commonShiftTransforms, right below the pre-existing pre-shift combine.

Nice!

Had to add setHasNoSignedWrap(false); for the Shl case. Otherwise the @shl_nsw_add_negative test case was generating a shl nsw i32 1 instead of shl i32 1. Not entirely sure why, possibly because I'm genering the Op with BinaryOperator::Create(I.getOpcode()?

Hmm - I'm not sure what that is about. Let's pre-commit the tests with baseline CHECK lines to confirm that we are testing the expected patterns. First, we need to add at least 2 more tests though that do not have "exact" - we want to verify that this transform does not fire without "exact".
If you do not have commit access yet, let me know, and I'll push the tests for you.

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

204

If the constant is not negative, the ashr will get converted to lshr before we hit the new transform, so this test is not providing the coverage that we want. How about something like this:

%a = add i8 %x, -4
%r = ashr exact i8 -112, %a ; 0b1001_0000

[nico-abram]

Updated the diff with the 2 tests to ensure the opt is not applied on non-exact ashr/lshr ashr_add_negative_shift_no_exact and lshr_add_negative_shift_no_exact.

To be sure I understand correctly: You mean adding just the tests, and have the same kind of CHECK-NEXT FileCheck directives, but with the expected input IR i.e the same IR the tests contain, to verify the inputs are not being affected by unrelated optimizations? I'm attaching baseline.patch

baseline.patch11 KBDownload

with that changeset. If needed I can open a new revision for the baseline test cases, or just update the diff of this one. I ran the tests locally on the baseline and they passed.

I do not have commit access, so it would be great if you could push for me. Please use "Nicolas Abram Lujan [e-mail]" if a name/email is needed.

[spatel]

In D124369#3475837, @nico-abram wrote:

To be sure I understand correctly: You mean adding just the tests, and have the same kind of CHECK-NEXT FileCheck directives, but with the expected input IR i.e the same IR the tests contain, to verify the inputs are not being affected by unrelated optimizations?

That is correct - we want to have all of the tests in place before this patch with the current output tested. That way if this patch has to be reverted we will still have test coverage in place for the transform. It also makes it easier to see exactly what is changing with this patch (and what is intentionally not changing).

I'm attaching baseline.patch

baseline.patch11 KBDownload

with that changeset. If needed I can open a new revision for the baseline test cases, or just update the diff of this one. I ran the tests locally on the baseline and they passed.

I do not have commit access, so it would be great if you could push for me. Please use "Nicolas Abram Lujan [e-mail]" if a name/email is needed.

I made several changes/additions to that set of tests and pushed that patch here:
fd9026131e6c

Please have a look, then update this patch by running update_test_checks.py on those tests and re-uploading.

We should not need to set NewShiftOp->setHasNoSignedWrap(false);. If you are seeing a test difference with that line in place, we need to investigate further.

[nico-abram]

Updated the patch with update_test_checks.py run with the changes applied on top of https://reviews.llvm.org/rGfd9026131e6c5de6f4e582d5c85c5b5588672dfb
Looks like what I'd expect. It's optimizing it through the opaque use call, but that doesn't look like a problem to me.

I also removed the setHasNoSignedWrap(false) and it did not change the test results or the update_test_checks changes. Not sure what I was seeing yesterday.

[spatel]

In D124369#3477338, @nico-abram wrote:

Updated the patch with update_test_checks.py run with the changes applied on top of https://reviews.llvm.org/rGfd9026131e6c5de6f4e582d5c85c5b5588672dfb
Looks like what I'd expect. It's optimizing it through the opaque use call, but that doesn't look like a problem to me.

I also removed the setHasNoSignedWrap(false) and it did not change the test results or the upodate_test_checks changes. Not sure what I was seeing yesterday.

Please rebase this patch on latest main - that way we'll just see the diffs in the tests.

[nico-abram]

Rebased

[spatel]

I think this is correct and sufficiently tested now. See inline comments for a couple of nits. Once those are addressed, I can push the patch on your behalf.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

426–427

This is more verbose than needed - if we're using &I, then we can just grab the opcode from it via I.getOpcode().
So we can either pass things in as params or use refs, instead of partly using both ways?

auto isSuitableForPreShift = [](Instruction &I, const APInt *AC, unsigned PosOffset) {
auto isSuitableForPreShift = [&I, AC, PosOffset]() {

458

Instead of a lambda + switch for setting the flags, I'd prefer the shorter and more direct:

if (I.getOpcode() == Instruction::Shl) {
   NewShiftOp->setHasNoUnsignedWrap(I.hasNoUnsignedWrap());
else
   NewShiftOp->setIsExact();

That might tilt the whole thing towards a lambda-less implementation, but you can decide which is easier to read.

[nico-abram]

Addressed comments, I kept the lambda.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

 }
 
 Instruction *InstCombinerImpl::commonShiftTransforms(BinaryOperator &I) {
   if (Instruction *Phi = foldBinopWithPhiOperands(I))
     return Phi;
 
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
   assert(Op0->getType() == Op1->getType());
+  Type *Ty = I.getType();
 
   // If the shift amount is a one-use `sext`, we can demote it to `zext`.
   Value *Y;
   if (match(Op1, m_OneUse(m_SExt(m_Value(Y))))) {
-    Value *NewExt = Builder.CreateZExt(Y, I.getType(), Op1->getName());
+    Value *NewExt = Builder.CreateZExt(Y, Ty, Op1->getName());
     return BinaryOperator::Create(I.getOpcode(), Op0, NewExt);
   }
 
   // See if we can fold away this shift.
   if (SimplifyDemandedInstructionBits(I))
     return &I;
 
   // Try to fold constant and into select arguments.
Instruction *InstCombinerImpl::commonShiftTransforms(BinaryOperator &I) {
   Value *A;
   Constant *C, *C1;
   if (match(Op0, m_Constant(C)) &&
       match(Op1, m_NUWAdd(m_Value(A), m_Constant(C1)))) {
     Constant *NewC = ConstantExpr::get(I.getOpcode(), C, C1);
     return BinaryOperator::Create(I.getOpcode(), NewC, A);
   }
 
+  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)) {
+    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));
+      }
+    };
+    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;
+    }
+  }
+
   // X shift (A srem C) -> X shift (A and (C - 1)) iff C is a power of 2.
   // Because shifts by negative values (which could occur if A were negative)
   // are undefined.
   if (Op1->hasOneUse() && match(Op1, m_SRem(m_Value(A), m_Constant(C))) &&
       match(C, m_Power2())) {
     // FIXME: Should this get moved into SimplifyDemandedBits by saying we don't
     // demand the sign bit (and many others) here??
-    Constant *Mask = ConstantExpr::getSub(C, ConstantInt::get(I.getType(), 1));
+    Constant *Mask = ConstantExpr::getSub(C, ConstantInt::get(Ty, 1));
     Value *Rem = Builder.CreateAnd(A, Mask, Op1->getName());
     return replaceOperand(I, 1, Rem);
   }
 
   if (Instruction *Logic = foldShiftOfShiftedLogic(I, Builder))
     return Logic;
 
   return nullptr;
Instruction *InstCombinerImpl::visitShl(BinaryOperator &I) {
   }
 
   // (1 << (C - x)) -> ((1 << C) >> x) if C is bitwidth - 1
   if (match(Op0, m_One()) &&
       match(Op1, m_Sub(m_SpecificInt(BitWidth - 1), m_Value(X))))
     return BinaryOperator::CreateLShr(
         ConstantInt::get(Ty, APInt::getSignMask(BitWidth)), X);
 
-  // Try to pre-shift a constant shifted by a variable amount:
-  // C << (X + AddC) --> (C >> -AddC) << X
-  // This requires a no-wrap flag and negative offset constant.
-  const APInt *AddC;
-  if ((I.hasNoSignedWrap() || I.hasNoUnsignedWrap()) &&
-      match(Op0, m_APInt(C)) && match(Op1, m_Add(m_Value(X), m_APInt(AddC))) &&
-      AddC->isNegative() && (-*AddC).ult(BitWidth)) {
-    assert(!C->isZero() && "Expected simplify of shifted zero");
-    unsigned PosOffset = (-*AddC).getZExtValue();
-    if (C->eq(C->lshr(PosOffset).shl(PosOffset))) {
-      Constant *NewC = ConstantInt::get(Ty, C->lshr(PosOffset));
-      Instruction *NewShl = BinaryOperator::CreateShl(NewC, X);
-      NewShl->setHasNoUnsignedWrap(I.hasNoUnsignedWrap());
-      return NewShl;
-    }
-  }
-
   return nullptr;
 }
 
 Instruction *InstCombinerImpl::visitLShr(BinaryOperator &I) {
   if (Value *V = SimplifyLShrInst(I.getOperand(0), I.getOperand(1), I.isExact(),
                                   SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 

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

;
   ret i32 %r
 }
 
 ; leading zeros for shifted constant precondition check for lshr_exact_add_negative_shift_positive
 
 define i32 @lshr_exact_add_negative_shift_negative(i32 %x) {
 ; CHECK-LABEL: @lshr_exact_add_negative_shift_negative(
 ; CHECK-NEXT:    [[A:%.*]] = add i32 [[X:%.*]], -1
 ; CHECK-NEXT:    [[R:%.*]] = lshr exact i32 -2, [[A]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %a = add i32 %x, -1
   %r = lshr exact i32 -2, %a
   ret i32 %r
 }
 
 ; exact precondition check for lshr_exact_add_negative_shift_positive
 
 define i32 @lshr_add_negative_shift_no_exact(i32 %x) {
 ; CHECK-LABEL: @lshr_add_negative_shift_no_exact(
 ; CHECK-NEXT:    [[A:%.*]] = add i32 [[X:%.*]], -1
 ; CHECK-NEXT:    [[R:%.*]] = lshr i32 2, [[A]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %a = add i32 %x, -1
   %r = lshr i32 2, %a
   ret i32 %r
 }
 
 define i32 @lshr_exact_add_negative_shift_positive(i32 %x) {
 ; CHECK-LABEL: @lshr_exact_add_negative_shift_positive(
-; CHECK-NEXT:    [[A:%.*]] = add i32 [[X:%.*]], -1
-; CHECK-NEXT:    [[R:%.*]] = lshr exact i32 2, [[A]]
+; CHECK-NEXT:    [[R:%.*]] = lshr exact i32 4, [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %a = add i32 %x, -1
   %r = lshr exact i32 2, %a
   ret i32 %r
 }
 
 define i8 @lshr_exact_add_negative_shift_positive_extra_use(i8 %x) {
 ; CHECK-LABEL: @lshr_exact_add_negative_shift_positive_extra_use(
 ; CHECK-NEXT:    [[A:%.*]] = add i8 [[X:%.*]], -1
 ; CHECK-NEXT:    call void @use(i8 [[A]])
-; CHECK-NEXT:    [[R:%.*]] = lshr exact i8 64, [[A]]
+; CHECK-NEXT:    [[R:%.*]] = lshr exact i8 -128, [[X]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %a = add i8 %x, -1
   call void @use(i8 %a)
   %r = lshr exact i8 64, %a
   ret i8 %r
 }
 
 define <2 x i9> @lshr_exact_add_negative_shift_positive_vec(<2 x i9> %x) {
 ; CHECK-LABEL: @lshr_exact_add_negative_shift_positive_vec(
-; CHECK-NEXT:    [[A:%.*]] = add <2 x i9> [[X:%.*]], <i9 -7, i9 -7>
-; CHECK-NEXT:    [[R:%.*]] = lshr exact <2 x i9> <i9 2, i9 2>, [[A]]
+; CHECK-NEXT:    [[R:%.*]] = lshr exact <2 x i9> <i9 -256, i9 -256>, [[X:%.*]]
 ; CHECK-NEXT:    ret <2 x i9> [[R]]
 ;
   %a = add <2 x i9> %x, <i9 -7, i9 -7>
   %r = lshr exact <2 x i9> <i9 2, i9 2>, %a
   ret <2 x i9> %r
 }
 
 ; not enough leading zeros in shift constant

 ;
   %a = add i32 %x, -1
   %r = ashr i32 -2, %a
   ret i32 %r
 }
 
 define i32 @ashr_exact_add_negative_shift_negative(i32 %x) {
 ; CHECK-LABEL: @ashr_exact_add_negative_shift_negative(
-; CHECK-NEXT:    [[A:%.*]] = add i32 [[X:%.*]], -1
-; CHECK-NEXT:    [[R:%.*]] = ashr exact i32 -2, [[A]]
+; CHECK-NEXT:    [[R:%.*]] = ashr exact i32 -4, [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %a = add i32 %x, -1
   %r = ashr exact i32 -2, %a
   ret i32 %r
 }
 
 define i8 @ashr_exact_add_negative_shift_negative_extra_use(i8 %x) {
 ; CHECK-LABEL: @ashr_exact_add_negative_shift_negative_extra_use(
 ; CHECK-NEXT:    [[A:%.*]] = add i8 [[X:%.*]], -2
 ; CHECK-NEXT:    call void @use(i8 [[A]])
-; CHECK-NEXT:    [[R:%.*]] = ashr exact i8 -32, [[A]]
+; CHECK-NEXT:    [[R:%.*]] = ashr exact i8 -128, [[X]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %a = add i8 %x, -2
   call void @use(i8 %a)
   %r = ashr exact i8 -32, %a
   ret i8 %r
 }
 
 define <2 x i7> @ashr_exact_add_negative_shift_negative_vec(<2 x i7> %x) {
 ; CHECK-LABEL: @ashr_exact_add_negative_shift_negative_vec(
-; CHECK-NEXT:    [[A:%.*]] = add <2 x i7> [[X:%.*]], <i7 -5, i7 -5>
-; CHECK-NEXT:    [[R:%.*]] = ashr exact <2 x i7> <i7 -2, i7 -2>, [[A]]
+; CHECK-NEXT:    [[R:%.*]] = ashr exact <2 x i7> <i7 -64, i7 -64>, [[X:%.*]]
 ; CHECK-NEXT:    ret <2 x i7> [[R]]
 ;
   %a = add <2 x i7> %x, <i7 -5, i7 -5>
   %r = ashr exact <2 x i7> <i7 -2, i7 -2>, %a
   ret <2 x i7> %r
 }
 
 ; not enough leading ones in shift constant