[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

202

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

425–426

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]() {

457

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

if (Op1->hasOneUse() && match(Op1, m_SRem(m_Value(A), m_Constant(C))) &&
     // 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));
     Value *Rem = Builder.CreateAnd(A, Mask, Op1->getName());
     return replaceOperand(I, 1, Rem);
   }
 
   if (Instruction *Logic = foldShiftOfShiftedLogic(I, Builder))
     return Logic;
 
   return nullptr;
 }
 
 /// Return true if we can simplify two logical (either left or right) shifts
 /// that have constant shift amounts: OuterShift (InnerShift X, C1), C2.
 static bool canEvaluateShiftedShift(unsigned OuterShAmt, bool IsOuterShl,
                                     Instruction *InnerShift,
                                     InstCombinerImpl &IC, Instruction *CxtI) {
static bool canEvaluateShiftedShift(unsigned OuterShAmt, bool IsOuterShl,
   // Equal shift amounts in opposite directions become bitwise 'and':
   // lshr (shl X, C), C --> and X, C'
   // shl (lshr X, C), C --> and X, C'
   if (*InnerShiftConst == OuterShAmt)
     return true;
 
   // If the 2nd shift is bigger than the 1st, we can fold:
   // lshr (shl X, C1), C2 -->  and (shl X, C1 - C2), C3
   // shl (lshr X, C1), C2 --> and (lshr X, C1 - C2), C3
   // but it isn't profitable unless we know the and'd out bits are already zero.
   // Also, check that the inner shift is valid (less than the type width) or
   // we'll crash trying to produce the bit mask for the 'and'.
   unsigned TypeWidth = InnerShift->getType()->getScalarSizeInBits();
   if (InnerShiftConst->ugt(OuterShAmt) && InnerShiftConst->ult(TypeWidth)) {
     unsigned InnerShAmt = InnerShiftConst->getZExtValue();
     unsigned MaskShift =
         IsInnerShl ? TypeWidth - InnerShAmt : InnerShAmt - OuterShAmt;
Instruction *InstCombinerImpl::visitLShr(BinaryOperator &I) {
   // Transform  (x << y) >> y  to  x & (-1 >> y)
   Value *X;
   if (match(Op0, m_OneUse(m_Shl(m_Value(X), m_Specific(Op1))))) {
     Constant *AllOnes = ConstantInt::getAllOnesValue(Ty);
     Value *Mask = Builder.CreateLShr(AllOnes, Op1);
     return BinaryOperator::CreateAnd(Mask, X);
   }
 
+  // Try to pre-shift a constant shifted by a variable amount:
+  // C >> (X - AddC) --> (C << AddC) >> X
+  // This requires a negative offset constant.
+  const APInt *AddC;
+  unsigned BitWidth = Ty->getScalarSizeInBits();
+  if (I.isExact() && 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->shl(PosOffset).lshr(PosOffset))) {
+      Constant *NewC = ConstantInt::get(Ty, C->shl(PosOffset));
+      BinaryOperator *NewLShr = BinaryOperator::CreateLShr(NewC, X);
+      NewLShr->setIsExact();
+      return NewLShr;
+    }
+  }
+
   return nullptr;
 }
 
 Instruction *
 InstCombinerImpl::foldVariableSignZeroExtensionOfVariableHighBitExtract(
     BinaryOperator &OldAShr) {
   assert(OldAShr.getOpcode() == Instruction::AShr &&
          "Must be called with arithmetic right-shift instruction only.");
Mask = Constant::mergeUndefsWith(
         cast<Constant>(cast<Instruction>(Op0)->getOperand(1)));
     X = Builder.CreateAnd(X, Mask);
     return BinaryOperator::CreateNeg(X);
   }
 
   if (Instruction *R = foldVariableSignZeroExtensionOfVariableHighBitExtract(I))
     return R;
 
+  // Try to pre-shift a constant shifted by a variable amount:
+  // C >> (X - AddC) --> (C << AddC) >> X
+  // This requires the exact flag and a negative offset constant.
+  const APInt *AddC;
+  const APInt *C;
+  if (I.isExact() && 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->shl(PosOffset).ashr(PosOffset))) {
+      Constant *NewC = ConstantInt::get(Ty, C->shl(PosOffset));
+      BinaryOperator *NewAShr = BinaryOperator::CreateAShr(NewC, X);
+      NewAShr->setIsExact();
+      return NewAShr;
+    }
+  }
+
   // See if we can turn a signed shr into an unsigned shr.
   if (MaskedValueIsZero(Op0, APInt::getSignMask(BitWidth), 0, &I))
     return BinaryOperator::CreateLShr(Op0, Op1);
 
   // ashr (xor %x, -1), %y  -->  xor (ashr %x, %y), -1
   if (match(Op0, m_OneUse(m_Not(m_Value(X))))) {
     // Note that we must drop 'exact'-ness of the shift!
     // Note that we can't keep undef's in -1 vector constant!
     auto *NewAShr = Builder.CreateAShr(X, Op1, Op0->getName() + ".not");
     return BinaryOperator::CreateNot(NewAShr);
   }
 
   return nullptr;
 }

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

 ; CHECK-NEXT:    ret i32 -1
 ;
   %a = add nuw i32 %x, 5
   store i32 %a, i32* %p
   %r = ashr i32 -6, %a
   ret i32 %r
 }
 
+; Negative offset precondition check for lshr_add_negative_shift_positive
+define i32 @lshr_add_positive_shift_positive(i32 %x) {
+; CHECK-LABEL: @lshr_add_positive_shift_positive(
+; 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
+}
+
+; Positive constant precondition check for lshr_add_negative_shift_positive
+define i32 @lshr_add_negative_shift_negative(i32 %x) {
+; CHECK-LABEL: @lshr_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
+}
+
+define i32 @lshr_add_negative_shift_positive(i32 %x) {
+; CHECK-LABEL: @lshr_add_negative_shift_positive(
+; 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
+}
+
+; Precondition check for ashr_exact_add_negative_shift_positive
+define i32 @ashr_exact_add_negative_shift_no_trailing_zeros(i32 %x) {
+; CHECK-LABEL: @ashr_exact_add_negative_shift_no_trailing_zeros(
+; CHECK-NEXT:    [[A:%.*]] = add i32 [[X:%.*]], -1
+; CHECK-NEXT:    [[R:%.*]] = lshr i32 1073741824, [[A:%.*]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %a = add i32 %x, -1
+  %r = ashr exact i32 1073741824, %a
+  ret i32 %r
+}
+
+define i32 @ashr_exact_add_negative_shift_negative(i32 %x) {
+; CHECK-LABEL: @ashr_exact_add_negative_shift_negative(
+; 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 i32 @ashr_exact_add_negative_shift_positive(i32 %x) {
+; CHECK-LABEL: @ashr_exact_add_negative_shift_positive(
+; CHECK-NEXT:    [[R:%.*]] = lshr i32 4, [[X:%.*]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %a = add i32 %x, -1
+  %r = ashr exact i32 2, %a
+  ret i32 %r
+}
+
 ; negative test - must have 'nuw'
 
 define i32 @shl_add_nsw(i32 %x) {
 ; CHECK-LABEL: @shl_add_nsw(
 ; CHECK-NEXT:    [[A:%.*]] = add nsw i32 [[X:%.*]], 5
 ; CHECK-NEXT:    [[R:%.*]] = shl i32 6, [[A]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;