[InstCombine] Shift amount reassociation in shifty sign bit test (PR43595)

[lebedev.ri]

Summary

This problem consists of several parts:

• Basic sign bit extraction - trunc? (?shr %x, (bitwidth(x)-1)). This is trivial, and easy to do, we have a fold for it.
• Shift amount reassociation - if we have two identical shifts, and we can simplify-add their shift amounts together, then we likely can just perform them as a single shift. But this is finicky, has one-use restrictions, and shift opcodes must be identical.

But there is a super-pattern where both of these work together. to produce sign bit test from two shifts + comparison.
We do indeed already handle this in most cases. But since we get that fold transitively, it has one-use restrictions.
And what's worse, in this case the right-shifts aren't required to be identical, and we can't handle that transitively:

If the total shift amount is bitwidth-1, only a sign bit will remain in the output value.
But if we look at this from the perspective of two shifts, we can't fold - we can't possibly know what bit pattern
we'd produce via two shifts, it will be *some* kind of a mask produced from original sign bit, but we just can't
tell it's shape: https://rise4fun.com/Alive/cM0 https://rise4fun.com/Alive/9IN

But it will *only* contain sign bit and zeros. So from the perspective of sign bit test, we're good:
https://rise4fun.com/Alive/FRz https://rise4fun.com/Alive/qBU
Superb!

So the simplest solution is to extend reassociateShiftAmtsOfTwoSameDirectionShifts() to also have a
sudo-analysis mode that will ignore extra-uses, and will only check whether a) those are two right shifts
and b) they end up with bitwidth(x)-1 shift amount and return either the original value that we sign-checking,
or null.

This does not have any functionality change for the existing reassociateShiftAmtsOfTwoSameDirectionShifts().

https://bugs.llvm.org/show_bug.cgi?id=43595

[lebedev.ri]

Any thoughts on this?

[spatel]

In D68930#1713151, @lebedev.ri wrote:

Any thoughts on this?

It raises the usual questions: cost/complexity vs. benefit, and how much further do we want to take this pattern matching before moving it somewhere else (AggressiveInstCombine)?
I don't think we'll reach consensus on that in this patch, so I've just noted a small improvement.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

29–30

Code comment typo and difficult to parse. Maybe better:
AnalyzeForSignBitExtraction indicates that we will only analyze whether this pattern has any 2 right-shifts that sum to 1 less than original bit width.

[lebedev.ri]

In D68930#1714734, @spatel wrote:

In D68930#1713151, @lebedev.ri wrote:

Any thoughts on this?

It raises the usual questions: cost/complexity vs. benefit, and how much further do we want to take this pattern matching before moving it somewhere else (AggressiveInstCombine)?
I don't think we'll reach consensus on that in this patch

Indeed. I was hoping that it wouldn't come to this patch, i thought general folds
(reassociateShiftAmtsOfTwoSameDirectionShifts(), foldShiftIntoShiftInAnotherHandOfAndInICmp())
would handle it, but no, sadly i'm still seeing regressions (although smaller!) in target benchmark,
and analysis of IR points to *this* missing pattern. I can't yet tell what other folds are needed though.

so I've just noted a small improvement.

Thanks!

[lebedev.ri]

Adjust comment wording as proposed.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

29–30

Indeed, that is better; and fixes the param name.

[xbolva00]

sadly i'm still seeing regressions (although smaller!) in target benchmark, and analysis of IR points to *this* missing pattern

But does this or base pattern shows up somewhere else (test suite/clang/chromium) than in your benchmark?

I think this is very rare pattern and it is not very good that we are going to add another rare pattern to InstCombine which is a slow pass already. I still think AggresiveInstCombine is better solution - or improve your code of interest manually.

[lebedev.ri]

In D68930#1714856, @xbolva00 wrote:

sadly i'm still seeing regressions (although smaller!) in target benchmark, and analysis of IR points to *this* missing pattern

But does this or base pattern shows up somewhere else (test suite/clang/chromium) than in your benchmark?

It's not "'your benchmark'", that is a real code. I'm not pulling it out of a thin air.

I think this is very rare pattern and it is not very good that we are going to add another rare pattern to InstCombine which is a slow pass already. I still think AggresiveInstCombine is better solution - or improve your code of interest manually.

[spatel]

In D68930#1714878, @lebedev.ri wrote:

In D68930#1714856, @xbolva00 wrote:

sadly i'm still seeing regressions (although smaller!) in target benchmark, and analysis of IR points to *this* missing pattern

But does this or base pattern shows up somewhere else (test suite/clang/chromium) than in your benchmark?

It's not "'your benchmark'", that is a real code. I'm not pulling it out of a thin air.

I don't think the intent was to disrespect whatever code this pattern showed up in, but to ask if there are any stats for common clang benchmarks. That's the only way we can draw a line on these kinds of decisions. Every pattern is presumably showing up somewhere real for someone, so we can't distinguish where it belongs just based on existence.

My opinion (and others have disagreed with this) is that we always want the fold somewhere; it's just a question of where does it go to minimize cost/maximize benefit.

[xbolva00]

I don't think the intent was to disrespect whatever code this pattern showed up in, but to ask if there are any stats for common clang benchmarks. That's the only way we can draw a line on these kinds of decisions. Every pattern is presumably showing up somewhere real for someone, so we can't distinguish where it belongs just based on existence.

Strong +1.

[lebedev.ri]

In D68930#1715144, @spatel wrote:

In D68930#1714878, @lebedev.ri wrote:

In D68930#1714856, @xbolva00 wrote:

sadly i'm still seeing regressions (although smaller!) in target benchmark, and analysis of IR points to *this* missing pattern

But does this or base pattern shows up somewhere else (test suite/clang/chromium) than in your benchmark?

It's not "'your benchmark'", that is a real code. I'm not pulling it out of a thin air.

I don't think the intent was to disrespect whatever code this pattern showed up in,

Sorry, yes, i think i agree.

but to ask if there are any stats for common clang benchmarks. That's the only way we can draw a line on these kinds of decisions. Every pattern is presumably showing up somewhere real for someone, so we can't distinguish where it belongs just based on existence.

This doesn't show up in test-suite with no externals, though we have already previously
established that test-suite plain has poor coverage for the code patterns in question..
To be noted, this 'benchmark' is proposed for addition as per rL345166, so it isn't entirely outlandish to consider it.

To be noted, said library, to my knowledge, is only one of two in it's class (camera raw image decoding;
counting only open source libs), with other one using this one, so it's moderately widely used in-the-wild.
Not sure if that counts...

My opinion (and others have disagreed with this) is that we always want the fold somewhere;
it's just a question of where does it go to minimize cost/maximize benefit.

Ignoring the -O3 issue, i'm not really confident AggressiceInstCombine is a great fit for these problems.
The thing is, these folds are not one-off "folds - yay, no fold - shrug".
There isn't much interest in this fold happening in itself.
The really important part is that it will pave the road for other folds to happen.
In particular, this fold will allow e.g. rL373964 to happen.
The situation is likely similar with all other "aggressive" folds i'we added.

[spatel]

In D68930#1715209, @lebedev.ri wrote:

Ignoring the -O3 issue, i'm not really confident AggressiceInstCombine is a great fit for these problems.
The thing is, these folds are not one-off "folds - yay, no fold - shrug".
There isn't much interest in this fold happening in itself.
The really important part is that it will pave the road for other folds to happen.
In particular, this fold will allow e.g. rL373964 to happen.
The situation is likely similar with all other "aggressive" folds i'we added.

Yes, I understand this patch alone doesn't look like a big change. (We'd move the whole reassociateShiftAmtsOfTwoSameDirectionShifts() if we decide to move anything?)

But it should be ok. We can assume that regular and aggressive will work together, and we can make sure that doesn't break by adding a test within PhaseOrdering. Obviously, it's more work to do it that way, so that's probably why we haven't done it often.

An alternative could be to guard this code (and possibly many more) with:

if (ExpensiveCombines)

?

[lebedev.ri]

In D68930#1715645, @spatel wrote:

In D68930#1715209, @lebedev.ri wrote:

Ignoring the -O3 issue, i'm not really confident AggressiceInstCombine is a great fit for these problems.
The thing is, these folds are not one-off "folds - yay, no fold - shrug".
There isn't much interest in this fold happening in itself.
The really important part is that it will pave the road for other folds to happen.
In particular, this fold will allow e.g. rL373964 to happen.
The situation is likely similar with all other "aggressive" folds i'we added.

Yes, I understand this patch alone doesn't look like a big change.

To recap, the main problem is not that they use InstSimplify,
but that SimplifyAssociativeBinOp() uses recursive reasoning, correct?

In this case before we ever call SimplifyAddInst() we need to match 3 instructions - != 0 + 2 right shifts.
reassociateShiftAmtsOfTwoSameDirectionShifts() for example only requires just 2 instructions,
foldShiftIntoShiftInAnotherHandOfAndInICmp() - 4 instructions,
dropRedundantMaskingOfLeftShiftInput() - 3/4 instructions.

So *this* case is certainly not worst,
i`m honestly not sure what's re-triggered this disscussion yet again...

(We'd move the whole reassociateShiftAmtsOfTwoSameDirectionShifts() if we decide to move anything?)
But it should be ok. We can assume that regular and aggressive will work together,
and we can make sure that doesn't break by adding a test within PhaseOrdering.
Obviously, it's more work to do it that way, so that's probably why we haven't done it often.

To be perfectly honest i don't really see/don't look forward that happening,
because i don't really believe just a single pass of AIC followed by IC will be enough to still handle everything.
i suspect it would need to be IC-AIC-IC-AIC-IC, and right now AIC only runs once.

An alternative could be to guard this code (and possibly many more) with:

if (ExpensiveCombines)

?

Hm, opt -instcombine does result in ExpensiveCombines=1, opt/clang -O0/-O1/-O2 doesn't, opt/clang -O3 does.
So InstCombine's ExpensiveCombines is basically in-pass AIC (i.e. they run at same optlevels),
just without all the issues that will arise from not running combines till completion.
I suppose this could work, but i *really* don't like -O3 limitation (should be -O2 in these cases,
but not for the sole existing transform guarded by ExpensiveCombines).

All that being said, if it's really this much of a concern to warrant immediate action,
can't we simply expose "Depth" in the llvm::SimplifyAddInst()
and call it with such a depth that either disallows recursive reasoning,
or severely limits it? Some fastpaths will likely needed to be added then most likely.
That may be the best outcome considering the options..

[spatel]

In D68930#1715680, @lebedev.ri wrote:

In D68930#1715645, @spatel wrote:

In D68930#1715209, @lebedev.ri wrote:

Ignoring the -O3 issue, i'm not really confident AggressiceInstCombine is a great fit for these problems.
The thing is, these folds are not one-off "folds - yay, no fold - shrug".
There isn't much interest in this fold happening in itself.
The really important part is that it will pave the road for other folds to happen.
In particular, this fold will allow e.g. rL373964 to happen.
The situation is likely similar with all other "aggressive" folds i'we added.

Yes, I understand this patch alone doesn't look like a big change.

To recap, the main problem is not that they use InstSimplify,
but that SimplifyAssociativeBinOp() uses recursive reasoning, correct?

Any time we call into InstSimplify, we may be incurring expensive analysis either by recursion directly or via ValueTracking, so that's part of it. The other part is that we add a few lines in every patch, and then end up with 100+ lines of code to deal with some rare patterns.

In this case before we ever call SimplifyAddInst() we need to match 3 instructions - != 0 + 2 right shifts.
reassociateShiftAmtsOfTwoSameDirectionShifts() for example only requires just 2 instructions,
foldShiftIntoShiftInAnotherHandOfAndInICmp() - 4 instructions,
dropRedundantMaskingOfLeftShiftInput() - 3/4 instructions.

So *this* case is certainly not worst,
i`m honestly not sure what's re-triggered this disscussion yet again...

It's the accumulating effect, and I agree it's out-of-scope to beat up on this patch alone, so I won't hold it up. But we should take a step back, do some measurements, and decide how things can be split up.

(We'd move the whole reassociateShiftAmtsOfTwoSameDirectionShifts() if we decide to move anything?)
But it should be ok. We can assume that regular and aggressive will work together,
and we can make sure that doesn't break by adding a test within PhaseOrdering.
Obviously, it's more work to do it that way, so that's probably why we haven't done it often.

To be perfectly honest i don't really see/don't look forward that happening,
because i don't really believe just a single pass of AIC followed by IC will be enough to still handle everything.
i suspect it would need to be IC-AIC-IC-AIC-IC, and right now AIC only runs once.

Hmm...I don't remember it being that way, but you're right. I thought it was running interleaved with instcombine at some point.

After all that, LGTM. :)

[lebedev.ri]

In D68930#1715721, @spatel wrote:

In D68930#1715680, @lebedev.ri wrote:

In D68930#1715645, @spatel wrote:

In D68930#1715209, @lebedev.ri wrote:

Ignoring the -O3 issue, i'm not really confident AggressiceInstCombine is a great fit for these problems.
The thing is, these folds are not one-off "folds - yay, no fold - shrug".
There isn't much interest in this fold happening in itself.
The really important part is that it will pave the road for other folds to happen.
In particular, this fold will allow e.g. rL373964 to happen.
The situation is likely similar with all other "aggressive" folds i'we added.

Yes, I understand this patch alone doesn't look like a big change.

To recap, the main problem is not that they use InstSimplify,
but that SimplifyAssociativeBinOp() uses recursive reasoning, correct?

Any time we call into InstSimplify, we may be incurring expensive analysis either by recursion directly or via ValueTracking, so that's part of it. The other part is that we add a few lines in every patch, and then end up with 100+ lines of code to deal with some rare patterns.

Right.

In this case before we ever call SimplifyAddInst() we need to match 3 instructions - != 0 + 2 right shifts.
reassociateShiftAmtsOfTwoSameDirectionShifts() for example only requires just 2 instructions,
foldShiftIntoShiftInAnotherHandOfAndInICmp() - 4 instructions,
dropRedundantMaskingOfLeftShiftInput() - 3/4 instructions.

So *this* case is certainly not worst,
i`m honestly not sure what's re-triggered this disscussion yet again...

It's the accumulating effect

Certainly true, okay.

, and I agree it's out-of-scope to beat up on this patch alone, so I won't hold it up. But we should take a step back, do some measurements, and decide how things can be split up.

Any thoughts on the approach i suggested:

All that being said, if it's really this much of a concern to warrant immediate action,
can't we simply expose "Depth" in the llvm::SimplifyAddInst()
and call it with such a depth that either disallows recursive reasoning,
or severely limits it? Some fastpaths will likely needed to be added then most likely.
That may be the best outcome considering the options..

?

(We'd move the whole reassociateShiftAmtsOfTwoSameDirectionShifts() if we decide to move anything?)
But it should be ok. We can assume that regular and aggressive will work together,
and we can make sure that doesn't break by adding a test within PhaseOrdering.
Obviously, it's more work to do it that way, so that's probably why we haven't done it often.

To be perfectly honest i don't really see/don't look forward that happening,
because i don't really believe just a single pass of AIC followed by IC will be enough to still handle everything.
i suspect it would need to be IC-AIC-IC-AIC-IC, and right now AIC only runs once.

Hmm...I don't remember it being that way, but you're right. I thought it was running interleaved with instcombine at some point.

Ah, i see why it was suggested then.
Good that i believe we reached understanding on why such move might be troubling..

After all that, LGTM. :)

[lebedev.ri]

@spatel

In D68930#1715726, @lebedev.ri wrote:

In D68930#1715721, @spatel wrote:

In D68930#1715680, @lebedev.ri wrote:

, and I agree it's out-of-scope to beat up on this patch alone, so I won't hold it up. But we should take a step back, do some measurements, and decide how things can be split up.

Any thoughts on the approach i suggested:

All that being said, if it's really this much of a concern to warrant immediate action,
can't we simply expose "Depth" in the llvm::SimplifyAddInst()
and call it with such a depth that either disallows recursive reasoning,
or severely limits it? Some fastpaths will likely needed to be added then most likely.
That may be the best outcome considering the options..

?

Up

[spatel]

In D68930#1720610, @lebedev.ri wrote:

@spatel

In D68930#1715726, @lebedev.ri wrote:

In D68930#1715721, @spatel wrote:

In D68930#1715680, @lebedev.ri wrote:

, and I agree it's out-of-scope to beat up on this patch alone, so I won't hold it up. But we should take a step back, do some measurements, and decide how things can be split up.

Any thoughts on the approach i suggested:

All that being said, if it's really this much of a concern to warrant immediate action,
can't we simply expose "Depth" in the llvm::SimplifyAddInst()
and call it with such a depth that either disallows recursive reasoning,
or severely limits it? Some fastpaths will likely needed to be added then most likely.
That may be the best outcome considering the options..

?

Up

Still catching up on mail backlog...
Exposing/varying the depth param doesn't seem like the right way to approach this to me. My fear is that we'd end up with semi-random uses of caller depths that wouldn't age well. So we might handle all of the cases in *this* patch with depth=1, but then someone notices their code needs depth=2 to get the optimization and changes it. The depth param is just a hack based on a magic number, so it feels wrong to expand on that hack. I don't think we can distinguish compile-time cost at that level anyway (but it would be great to have data). Apart from all of that, messing with recursion depth doesn't address the underlying problem that the code is just too big/complicated.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

Instruction *InstCombiner::foldIRemByPowerOfTwoToBitTest(ICmpInst &I) {
   // This may increase instruction count, we don't enforce that Y is a constant.
   Value *Mask = Builder.CreateAdd(Y, Constant::getAllOnesValue(Y->getType()));
   Value *Masked = Builder.CreateAnd(X, Mask);
   return ICmpInst::Create(Instruction::ICmp, Pred, Masked, Zero);
 }
 
 /// Fold equality-comparison between zero and any (maybe truncated) right-shift
 /// by one-less-than-bitwidth into a sign test on the original value.
-Instruction *foldSignBitTest(ICmpInst &I) {
+Instruction *InstCombiner::foldSignBitTest(ICmpInst &I) {
+  Instruction *Val;
   ICmpInst::Predicate Pred;
-  Value *X;
-  Constant *C;
-  if (!I.isEquality() ||
-      !match(&I, m_ICmp(Pred, m_TruncOrSelf(m_Shr(m_Value(X), m_Constant(C))),
-                        m_Zero())))
+  if (!I.isEquality() || !match(&I, m_ICmp(Pred, m_Instruction(Val), m_Zero())))
     return nullptr;
 
-  Type *XTy = X->getType();
+  Value *X;
+  Type *XTy;
+
+  Constant *C;
+  if (match(Val, m_TruncOrSelf(m_Shr(m_Value(X), m_Constant(C))))) {
+    XTy = X->getType();
     unsigned XBitWidth = XTy->getScalarSizeInBits();
     if (!match(C, m_SpecificInt_ICMP(ICmpInst::Predicate::ICMP_EQ,
                                      APInt(XBitWidth, XBitWidth - 1))))
       return nullptr;
+  } else if (isa<BinaryOperator>(Val) &&
+             (X = reassociateShiftAmtsOfTwoSameDirectionShifts(
+                  cast<BinaryOperator>(Val), SQ.getWithInstruction(Val),
+                  /*AnalyzeForSignBitExtraction=*/true))) {
+    XTy = X->getType();
+  } else
+    return nullptr;
 
   return ICmpInst::Create(Instruction::ICmp,
                           Pred == ICmpInst::ICMP_EQ ? ICmpInst::ICMP_SGE
                                                     : ICmpInst::ICMP_SLT,
                           X, ConstantInt::getNullValue(XTy));
 }
 
 // Handle  icmp pred X, 0

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   Instruction *visitUDiv(BinaryOperator &I);
   Instruction *visitSDiv(BinaryOperator &I);
   Instruction *visitFDiv(BinaryOperator &I);
   Value *simplifyRangeCheck(ICmpInst *Cmp0, ICmpInst *Cmp1, bool Inverted);
   Instruction *visitAnd(BinaryOperator &I);
   Instruction *visitOr(BinaryOperator &I);
   Instruction *visitXor(BinaryOperator &I);
   Instruction *visitShl(BinaryOperator &I);
+  Value *reassociateShiftAmtsOfTwoSameDirectionShifts(
+      BinaryOperator *Sh0, const SimplifyQuery &SQ,
+      bool AnalyzeForSignBitExtraction = false);
   Instruction *foldVariableSignZeroExtensionOfVariableHighBitExtract(
       BinaryOperator &OldAShr);
   Instruction *visitAShr(BinaryOperator &I);
   Instruction *visitLShr(BinaryOperator &I);
   Instruction *commonShiftTransforms(BinaryOperator &I);
   Instruction *visitFCmpInst(FCmpInst &I);
   Instruction *visitICmpInst(ICmpInst &I);
   Instruction *FoldShiftByConstant(Value *Op0, Constant *Op1,
private:
   Instruction *foldICmpUsingKnownBits(ICmpInst &Cmp);
   Instruction *foldICmpWithDominatingICmp(ICmpInst &Cmp);
   Instruction *foldICmpWithConstant(ICmpInst &Cmp);
   Instruction *foldICmpInstWithConstant(ICmpInst &Cmp);
   Instruction *foldICmpInstWithConstantNotInt(ICmpInst &Cmp);
   Instruction *foldICmpBinOp(ICmpInst &Cmp, const SimplifyQuery &SQ);
   Instruction *foldICmpEquality(ICmpInst &Cmp);
   Instruction *foldIRemByPowerOfTwoToBitTest(ICmpInst &I);
+  Instruction *foldSignBitTest(ICmpInst &I);
   Instruction *foldICmpWithZero(ICmpInst &Cmp);
 
   Value *foldUnsignedMultiplicationOverflowCheck(ICmpInst &Cmp);
 
   Instruction *foldICmpSelectConstant(ICmpInst &Cmp, SelectInst *Select,
                                       ConstantInt *C);
   Instruction *foldICmpTruncConstant(ICmpInst &Cmp, TruncInst *Trunc,
                                      const APInt &C);

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

 
 #define DEBUG_TYPE "instcombine"
 
 // Given pattern:
 //   (x shiftopcode Q) shiftopcode K
 // we should rewrite it as
 //   x shiftopcode (Q+K)  iff (Q+K) u< bitwidth(x)
 // This is valid for any shift, but they must be identical.
-static Instruction *
-reassociateShiftAmtsOfTwoSameDirectionShifts(BinaryOperator *Sh0,
-                                             const SimplifyQuery &SQ,
-                                             InstCombiner::BuilderTy &Builder) {
+//
+// AnalysisIsSignbitExtract is special - we will only analyze whether this
+// pattern has two right-shifts that shift my one less than original bit width,
+// thus extracting original sign bit, and return the base value X if so.
+Value *InstCombiner::reassociateShiftAmtsOfTwoSameDirectionShifts(
+    BinaryOperator *Sh0, const SimplifyQuery &SQ,
+    bool AnalyzeForSignBitExtraction) {
   // Look for a shift of some instruction, ignore zext of shift amount if any.
   Instruction *Sh0Op0;
   Value *ShAmt0;
   if (!match(Sh0,
              m_Shift(m_Instruction(Sh0Op0), m_ZExtOrSelf(m_Value(ShAmt0)))))
     return nullptr;
 
   // If there is a truncation between the two shifts, we must make note of it
Value *InstCombiner::reassociateShiftAmtsOfTwoSameDirectionShifts(
   if (!match(Sh1, m_Shift(m_Value(X), m_ZExtOrSelf(m_Value(ShAmt1)))))
     return nullptr;
 
   // We have two shift amounts from two different shifts. The types of those
   // shift amounts may not match. If that's the case let's bailout now..
   if (ShAmt0->getType() != ShAmt1->getType())
     return nullptr;
 
-  // The shift opcodes must be identical.
+  // We are only looking for signbit extraction if we have two right shifts.
+  bool HadTwoRightShifts = match(Sh0, m_Shr(m_Value(), m_Value())) &&
+                           match(Sh1, m_Shr(m_Value(), m_Value()));
+  // ... and if it's not two right-shifts, we know the answer already.
+  if (AnalyzeForSignBitExtraction && !HadTwoRightShifts)
+    return nullptr;
+
+  // The shift opcodes must be identical, unless we are just checking whether
+  // this pattern can be interpreted as a sign-bit-extraction.
   Instruction::BinaryOps ShiftOpcode = Sh0->getOpcode();
-  if (ShiftOpcode != Sh1->getOpcode())
+  bool IdenticalShOpcodes = Sh0->getOpcode() == Sh1->getOpcode();
+  if (!IdenticalShOpcodes && !AnalyzeForSignBitExtraction)
     return nullptr;
 
   // If we saw truncation, we'll need to produce extra instruction,
-  // and for that one of the operands of the shift must be one-use.
-  if (Trunc && !match(Sh0, m_c_BinOp(m_OneUse(m_Value()), m_Value())))
+  // and for that one of the operands of the shift must be one-use,
+  // unless of course we don't actually plan to produce any instructions here.
+  if (Trunc && !AnalyzeForSignBitExtraction &&
+      !match(Sh0, m_c_BinOp(m_OneUse(m_Value()), m_Value())))
     return nullptr;
 
   // Can we fold (ShAmt0+ShAmt1) ?
   auto *NewShAmt = dyn_cast_or_null<Constant>(
       SimplifyAddInst(ShAmt0, ShAmt1, /*isNSW=*/false, /*isNUW=*/false,
                       SQ.getWithInstruction(Sh0)));
   if (!NewShAmt)
     return nullptr; // Did not simplify.
   unsigned NewShAmtBitWidth = NewShAmt->getType()->getScalarSizeInBits();
   unsigned XBitWidth = X->getType()->getScalarSizeInBits();
   // Is the new shift amount smaller than the bit width of inner/new shift?
   if (!match(NewShAmt, m_SpecificInt_ICMP(ICmpInst::Predicate::ICMP_ULT,
                                           APInt(NewShAmtBitWidth, XBitWidth))))
     return nullptr; // FIXME: could perform constant-folding.
 
   // If there was a truncation, and we have a right-shift, we can only fold if
-  // we are left with the original sign bit.
+  // we are left with the original sign bit. Likewise, if we were just checking
+  // that this is a sighbit extraction, this is the place to check it.
   // FIXME: zero shift amount is also legal here, but we can't *easily* check
   // more than one predicate so it's not really worth it.
-  if (Trunc && ShiftOpcode != Instruction::BinaryOps::Shl &&
-      !match(NewShAmt,
+  if (HadTwoRightShifts && (Trunc || AnalyzeForSignBitExtraction)) {
+    // If it's not a sign bit extraction, then we're done.
+    if (!match(NewShAmt,
                m_SpecificInt_ICMP(ICmpInst::Predicate::ICMP_EQ,
                                   APInt(NewShAmtBitWidth, XBitWidth - 1))))
       return nullptr;
+    // If it is, and that was the question, return the base value.
+    if (AnalyzeForSignBitExtraction)
+      return X;
+  }
+
+  assert(IdenticalShOpcodes && "Should not get here with different shifts.");
 
   // All good, we can do this fold.
   NewShAmt = ConstantExpr::getZExtOrBitCast(NewShAmt, X->getType());
 
   BinaryOperator *NewShift = BinaryOperator::Create(ShiftOpcode, X, NewShAmt);
 
   // The flags can only be propagated if there wasn't a trunc.
   if (!Trunc) {
if (isa<Constant>(Op0))
     if (SelectInst *SI = dyn_cast<SelectInst>(Op1))
       if (Instruction *R = FoldOpIntoSelect(I, SI))
         return R;
 
   if (Constant *CUI = dyn_cast<Constant>(Op1))
     if (Instruction *Res = FoldShiftByConstant(Op0, CUI, I))
       return Res;
 
-  if (Instruction *NewShift =
-          reassociateShiftAmtsOfTwoSameDirectionShifts(&I, SQ, Builder))
+  if (auto *NewShift = cast_or_null<Instruction>(
+          reassociateShiftAmtsOfTwoSameDirectionShifts(&I, SQ)))
     return NewShift;
 
   // (C1 shift (A add C2)) -> (C1 shift C2) shift A)
   // iff A and C2 are both positive.
   Value *A;
   Constant *C;
   if (match(Op0, m_Constant()) && match(Op1, m_Add(m_Value(A), m_Constant(C))))
     if (isKnownNonNegative(A, DL, 0, &AC, &I, &DT) &&

llvm/test/Transforms/InstCombine/sign-bit-test-via-right-shifting-all-other-bits.ll

 ; CHECK-NEXT:    [[SKIP_ALL_BITS_TILL_SIGNBIT:%.*]] = add i32 [[NBITS]], -1
 ; CHECK-NEXT:    [[SIGNBIT:%.*]] = lshr i32 [[HIGH_BITS_EXTRACTED_NARROW]], [[SKIP_ALL_BITS_TILL_SIGNBIT]]
 ; CHECK-NEXT:    call void @use32(i32 [[NUM_LOW_BITS_TO_SKIP]])
 ; CHECK-NEXT:    call void @use64(i64 [[NUM_LOW_BITS_TO_SKIP_WIDE]])
 ; CHECK-NEXT:    call void @use64(i64 [[HIGH_BITS_EXTRACTED]])
 ; CHECK-NEXT:    call void @use32(i32 [[HIGH_BITS_EXTRACTED_NARROW]])
 ; CHECK-NEXT:    call void @use32(i32 [[SKIP_ALL_BITS_TILL_SIGNBIT]])
 ; CHECK-NEXT:    call void @use32(i32 [[SIGNBIT]])
-; CHECK-NEXT:    [[ISNEG:%.*]] = icmp ne i32 [[SIGNBIT]], 0
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i64 [[DATA]], 0
 ; CHECK-NEXT:    ret i1 [[ISNEG]]
 ;
   %num_low_bits_to_skip = sub i32 64, %nbits
   %num_low_bits_to_skip_wide = zext i32 %num_low_bits_to_skip to i64
   %high_bits_extracted = lshr i64 %data, %num_low_bits_to_skip_wide
   %high_bits_extracted_narrow = trunc i64 %high_bits_extracted to i32
   %skip_all_bits_till_signbit = sub i32 %nbits, 1
   %signbit = lshr i32 %high_bits_extracted_narrow, %skip_all_bits_till_signbit

 ; CHECK-NEXT:    [[SKIP_ALL_BITS_TILL_SIGNBIT:%.*]] = add i32 [[NBITS]], -1
 ; CHECK-NEXT:    [[SIGNBIT:%.*]] = ashr i32 [[HIGH_BITS_EXTRACTED_NARROW]], [[SKIP_ALL_BITS_TILL_SIGNBIT]]
 ; CHECK-NEXT:    call void @use32(i32 [[NUM_LOW_BITS_TO_SKIP]])
 ; CHECK-NEXT:    call void @use64(i64 [[NUM_LOW_BITS_TO_SKIP_WIDE]])
 ; CHECK-NEXT:    call void @use64(i64 [[HIGH_BITS_EXTRACTED]])
 ; CHECK-NEXT:    call void @use32(i32 [[HIGH_BITS_EXTRACTED_NARROW]])
 ; CHECK-NEXT:    call void @use32(i32 [[SKIP_ALL_BITS_TILL_SIGNBIT]])
 ; CHECK-NEXT:    call void @use32(i32 [[SIGNBIT]])
-; CHECK-NEXT:    [[ISNEG:%.*]] = icmp ne i32 [[SIGNBIT]], 0
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i64 [[DATA]], 0
 ; CHECK-NEXT:    ret i1 [[ISNEG]]
 ;
   %num_low_bits_to_skip = sub i32 64, %nbits
   %num_low_bits_to_skip_wide = zext i32 %num_low_bits_to_skip to i64
   %high_bits_extracted = ashr i64 %data, %num_low_bits_to_skip_wide
   %high_bits_extracted_narrow = trunc i64 %high_bits_extracted to i32
   %skip_all_bits_till_signbit = sub i32 %nbits, 1
   %signbit = ashr i32 %high_bits_extracted_narrow, %skip_all_bits_till_signbit

 ; CHECK-NEXT:    [[NUM_LOW_BITS_TO_SKIP:%.*]] = sub i32 32, [[NBITS:%.*]]
 ; CHECK-NEXT:    [[HIGH_BITS_EXTRACTED:%.*]] = lshr i32 [[DATA:%.*]], [[NUM_LOW_BITS_TO_SKIP]]
 ; CHECK-NEXT:    [[SKIP_ALL_BITS_TILL_SIGNBIT:%.*]] = add i32 [[NBITS]], -1
 ; CHECK-NEXT:    [[SIGNBIT:%.*]] = ashr i32 [[HIGH_BITS_EXTRACTED]], [[SKIP_ALL_BITS_TILL_SIGNBIT]]
 ; CHECK-NEXT:    call void @use32(i32 [[NUM_LOW_BITS_TO_SKIP]])
 ; CHECK-NEXT:    call void @use32(i32 [[HIGH_BITS_EXTRACTED]])
 ; CHECK-NEXT:    call void @use32(i32 [[SKIP_ALL_BITS_TILL_SIGNBIT]])
 ; CHECK-NEXT:    call void @use32(i32 [[SIGNBIT]])
-; CHECK-NEXT:    [[ISNEG:%.*]] = icmp ne i32 [[SIGNBIT]], 0
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[DATA]], 0
 ; CHECK-NEXT:    ret i1 [[ISNEG]]
 ;
   %num_low_bits_to_skip = sub i32 32, %nbits
   %high_bits_extracted = lshr i32 %data, %num_low_bits_to_skip
   %skip_all_bits_till_signbit = sub i32 %nbits, 1
   %signbit = ashr i32 %high_bits_extracted, %skip_all_bits_till_signbit
 
   call void @use32(i32 %num_low_bits_to_skip)

 ; CHECK-NEXT:    [[SKIP_ALL_BITS_TILL_SIGNBIT:%.*]] = add i32 [[NBITS]], -1
 ; CHECK-NEXT:    [[SIGNBIT:%.*]] = ashr i32 [[HIGH_BITS_EXTRACTED_NARROW]], [[SKIP_ALL_BITS_TILL_SIGNBIT]]
 ; CHECK-NEXT:    call void @use32(i32 [[NUM_LOW_BITS_TO_SKIP]])
 ; CHECK-NEXT:    call void @use64(i64 [[NUM_LOW_BITS_TO_SKIP_WIDE]])
 ; CHECK-NEXT:    call void @use64(i64 [[HIGH_BITS_EXTRACTED]])
 ; CHECK-NEXT:    call void @use32(i32 [[HIGH_BITS_EXTRACTED_NARROW]])
 ; CHECK-NEXT:    call void @use32(i32 [[SKIP_ALL_BITS_TILL_SIGNBIT]])
 ; CHECK-NEXT:    call void @use32(i32 [[SIGNBIT]])
-; CHECK-NEXT:    [[ISNEG:%.*]] = icmp ne i32 [[SIGNBIT]], 0
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i64 [[DATA]], 0
 ; CHECK-NEXT:    ret i1 [[ISNEG]]
 ;
   %num_low_bits_to_skip = sub i32 64, %nbits
   %num_low_bits_to_skip_wide = zext i32 %num_low_bits_to_skip to i64
   %high_bits_extracted = lshr i64 %data, %num_low_bits_to_skip_wide
   %high_bits_extracted_narrow = trunc i64 %high_bits_extracted to i32
   %skip_all_bits_till_signbit = sub i32 %nbits, 1
   %signbit = ashr i32 %high_bits_extracted_narrow, %skip_all_bits_till_signbit

 ; CHECK-NEXT:    [[NUM_LOW_BITS_TO_SKIP:%.*]] = sub i32 32, [[NBITS:%.*]]
 ; CHECK-NEXT:    [[HIGH_BITS_EXTRACTED:%.*]] = ashr i32 [[DATA:%.*]], [[NUM_LOW_BITS_TO_SKIP]]
 ; CHECK-NEXT:    [[SKIP_ALL_BITS_TILL_SIGNBIT:%.*]] = add i32 [[NBITS]], -1
 ; CHECK-NEXT:    [[SIGNBIT:%.*]] = lshr i32 [[HIGH_BITS_EXTRACTED]], [[SKIP_ALL_BITS_TILL_SIGNBIT]]
 ; CHECK-NEXT:    call void @use32(i32 [[NUM_LOW_BITS_TO_SKIP]])
 ; CHECK-NEXT:    call void @use32(i32 [[HIGH_BITS_EXTRACTED]])
 ; CHECK-NEXT:    call void @use32(i32 [[SKIP_ALL_BITS_TILL_SIGNBIT]])
 ; CHECK-NEXT:    call void @use32(i32 [[SIGNBIT]])
-; CHECK-NEXT:    [[ISNEG:%.*]] = icmp ne i32 [[SIGNBIT]], 0
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i32 [[DATA]], 0
 ; CHECK-NEXT:    ret i1 [[ISNEG]]
 ;
   %num_low_bits_to_skip = sub i32 32, %nbits
   %high_bits_extracted = ashr i32 %data, %num_low_bits_to_skip
   %skip_all_bits_till_signbit = sub i32 %nbits, 1
   %signbit = lshr i32 %high_bits_extracted, %skip_all_bits_till_signbit
 
   call void @use32(i32 %num_low_bits_to_skip)

 ; CHECK-NEXT:    [[SKIP_ALL_BITS_TILL_SIGNBIT:%.*]] = add i32 [[NBITS]], -1
 ; CHECK-NEXT:    [[SIGNBIT:%.*]] = lshr i32 [[HIGH_BITS_EXTRACTED_NARROW]], [[SKIP_ALL_BITS_TILL_SIGNBIT]]
 ; CHECK-NEXT:    call void @use32(i32 [[NUM_LOW_BITS_TO_SKIP]])
 ; CHECK-NEXT:    call void @use64(i64 [[NUM_LOW_BITS_TO_SKIP_WIDE]])
 ; CHECK-NEXT:    call void @use64(i64 [[HIGH_BITS_EXTRACTED]])
 ; CHECK-NEXT:    call void @use32(i32 [[HIGH_BITS_EXTRACTED_NARROW]])
 ; CHECK-NEXT:    call void @use32(i32 [[SKIP_ALL_BITS_TILL_SIGNBIT]])
 ; CHECK-NEXT:    call void @use32(i32 [[SIGNBIT]])
-; CHECK-NEXT:    [[ISNEG:%.*]] = icmp ne i32 [[SIGNBIT]], 0
+; CHECK-NEXT:    [[ISNEG:%.*]] = icmp slt i64 [[DATA]], 0
 ; CHECK-NEXT:    ret i1 [[ISNEG]]
 ;
   %num_low_bits_to_skip = sub i32 64, %nbits
   %num_low_bits_to_skip_wide = zext i32 %num_low_bits_to_skip to i64
   %high_bits_extracted = ashr i64 %data, %num_low_bits_to_skip_wide
   %high_bits_extracted_narrow = trunc i64 %high_bits_extracted to i32
   %skip_all_bits_till_signbit = sub i32 %nbits, 1
   %signbit = lshr i32 %high_bits_extracted_narrow, %skip_all_bits_till_signbit