[DAGCombiner] Bug 31275- Extract a shift from a constant mul or udiv if a rotate can be formed

[sameconrad]

Summary

Attempt to extract a shrl from a udiv or a shl from a mul if this allows a rotate to be formed. This targets cases where the input to a rotate pattern was a mul or udiv by a constant and InstCombine merged one of the shifts with the op.

Patch by: sameconrad (Sam Conrad)

[lebedev.ri]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5091–5102

Aren't we loosing by bailing out here? (And no tests show it?!)
Maybe you meant:

// If only one side has a shift, see if we can extract it
if (!(NeedRotLHS ^ NeedRotRHS)) {
  if (NeedRotRHS) {
    if (SDValue Shift = extractShiftFromMulOrUDiv(LHSShift, RHS, RHSMask, DL))
      RHSShift = Shift;
  }
  if (NeedRotLHS) {
    if (SDValue Shift = extractShiftFromMulOrUDiv(RHSShift, LHS, LHSMask, DL))
      LHSShift = Shift;
  }
}

Please don't blindly follow that snippet, i don't know what the other code does,
so your change may be correct as-is.

test/CodeGen/X86/rotate-extract.ll

2

You might want to use llvm/utils/update_llc_test_checks.py

[lebedev.ri]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5096–5100

Also, it is intentional that you extract from LHS and store it into a variable named RHS and vice-versa?
If yes, add a comment please.

[sameconrad]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5091–5102

I'll add some extra comments here and update in a bit, but I don't think we're losing out on anything.
The idea here is that if MatchRotateHalf matches one half of the OR as a shl/srl but not the other, we try to extract the needed shift from the unmatched side to complete the rotate. The case where neither side matched (ie NeedRotLHS && NeedRotRHS) is handled above, so at this point either no side is missing and neither of these if cases will fire, or only one side is missing, in which case we try to extract the shift from the missing side (using the shift that did match from the opposite side to compute the shift amount and opcode needed, hence we pass LHSShift when trying to extract from the RHS and vice versa). If extraction fails at that point then we definitely can't complete the rotate and bail.

For instance if NeedRotRHS==true and LHSShift is (srl v 10) and bitwidth(v)==32, then extractShiftFromMulOrDiv will compute that we need to extract (shl v 22) from RHS based on the shift amount and opcode in LHSShift.

[sameconrad]

Update tests with update_llc_test_checks
Improve comments in MatchRotate

[sameconrad]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5096–5100

Updated with some clarifying comments

[lebedev.ri]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5057

Is this always considered worth doing,
or the check whether this should be done is on the caller's side?

test/CodeGen/X86/rotate-extract.ll

3

You could also copy this into test/CodeGen/AArch64/rotate-extract.ll,
since that should also profit from the transform.

[sameconrad]

Added AArch64 test

[sameconrad]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5057

It looks like MatchRotate always gets called from DAGCombiner::visitOr, I don't think there's any sort of profitability check, so it seems to assume its always worthwhile (however I'm still pretty new to LLVM and don't know if there are other places this could be done). In any case, the new code would transform 4 instructions to 2 (2 mul/udiv ops, 1 shift, 1 or -> one mul/udiv op and 1 rotate).

[sameconrad]

Quick update to improve assert messages

[sameconrad]

Ping

[RKSimon]

The tests probably need committing with trunk's current codegen, so this patch then shows the diff.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4822

You can merge these if statements:

if (Op.getOpcode() == ISD::AND &&
    DAG.isConstantIntBuildVectorOrConstantInt(Op.getOperand(1))) {

4831

Does this need to be a DAGCombiner method any more? Just reduce it to a static.

4842

Again, what do you gain from this being a DAGCombiner method?

4850

Do we have any tests that use the Mask?

4878

This is more convoluted than it needs to be - the GetConstant isn't really helping and comparing against APInt() is a bad idea

Would it be better to just do (not sure about this):

// Attempt to find non-zero constants for opposing shift and the mul/div +shift.
ConstantSDNode *OppShiftCst = isConstOrConstSplat(OpposingShift.getOperand(1));
ConstantSDNode *ShiftCst = isConstOrConstSplat(ShiftLHS.getOperand(1));
ConstantSDNode *MulOrDivCst = isConstOrConstSplat(MulOrDiv.getOperand(1));
if (!OppShiftCst || !ShiftCst !! MulOrDivCst ||
    !OppShiftCst->getAPIntValue() || !ShiftCst->getAPIntValue() || !MulOrDivCst->getAPIntValue())
  return SDValue();

Also, we should bail if the OppShiftCst is out of range.

4885

Is this better?

APInt ShiftExtractDiv = APInt::getOneBitSet(VTWidth, NewShiftAmt;.getZextValue());

4901

Worth doing this earlier?

5115

This could be reduced to:

if (NeedRotRHS)
  RHSShift = extractShiftFromMulOrUDiv(LHSShift, RHS, RHSMask, DL);
if (NeedRotLHS)
  LHSShift = extractShiftFromMulOrUDiv(RHSShift, LHS, LHSMask, DL);
if (!RHSShift || !LHSShift)
  return nullptr;

[lebedev.ri]

BTW, that reminds me.
@sameconrad while you are looking at rotates in backend, can you please also check (at least add tests,
if there aren't any yet) that backend properly handles the cases where mul/div were transformed into shifts?

[lebedev.ri]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

490

In other words, the cases where c0, c1, c2 were power-of-two, and thus instcombine turned them into shifts.
So additional [straight-forward] cases are:

///   (or (shrl (shl v c0') c1) (shl v c2')) ->
///   (or (shrl (shl v c0') c1) (shl (shl v c0') c3))
/// and
///   (or (shrl v c0') (shl (shrl v c1') c2)) ->
///   (or (shrl (shrl v c1') c3) (shl (shrl v c1') c2))

There is also an additional problem when we are/will convert shrl+shl / shl+rhrl into masking...

[sameconrad]

Clean up code quality
Add tests with bitmask case

[sameconrad]

Small bit of code cleanup that was missed

[sameconrad]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4850

Added a mask test

4878

Fixed up most of this but wasn't sure what you meant by out of range, do you mean if the existing shift amount is greater than the bitwidth?

[sameconrad]

In D47681#1131359, @lebedev.ri wrote:

BTW, that reminds me.
@sameconrad while you are looking at rotates in backend, can you please also check (at least add tests,
if there aren't any yet) that backend properly handles the cases where mul/div were transformed into shifts?

I can take a look at that, will do that in a separate patch.

Actually nvm, I see what you're saying. I'll try to update this patch in a bit for that case.

[lebedev.ri]

In D47681#1131982, @sameconrad wrote:

In D47681#1131359, @lebedev.ri wrote:

BTW, that reminds me.
@sameconrad while you are looking at rotates in backend, can you please also check (at least add tests,
if there aren't any yet) that backend properly handles the cases where mul/div were transformed into shifts?

I can take a look at that, will do that in a separate patch.

Actually nvm, I see what you're saying. I'll try to update this patch in a bit for that case.

You can refer to the tests in D48229 for the examples of patterns that are related.

[sameconrad]

Update to also handle cases where we need to extract a smaller shift from a shl/srl.

[sameconrad]

In D47681#1134801, @lebedev.ri wrote:

In D47681#1131982, @sameconrad wrote:

In D47681#1131359, @lebedev.ri wrote:

BTW, that reminds me.
@sameconrad while you are looking at rotates in backend, can you please also check (at least add tests,
if there aren't any yet) that backend properly handles the cases where mul/div were transformed into shifts?

I can take a look at that, will do that in a separate patch.

Actually nvm, I see what you're saying. I'll try to update this patch in a bit for that case.

You can refer to the tests in D48229 for the examples of patterns that are related.

Updated the patch to handle extracting from a shl/srl.
For handling srl+shl->bitmask cases I'd prefer to do a follow-up patch since I need a bit more time to look at it.

[lebedev.ri]

In D47681#1134871, @sameconrad wrote:

In D47681#1134801, @lebedev.ri wrote:

In D47681#1131982, @sameconrad wrote:

In D47681#1131359, @lebedev.ri wrote:

BTW, that reminds me.
@sameconrad while you are looking at rotates in backend, can you please also check (at least add tests,
if there aren't any yet) that backend properly handles the cases where mul/div were transformed into shifts?

I can take a look at that, will do that in a separate patch.

Actually nvm, I see what you're saying. I'll try to update this patch in a bit for that case.

You can refer to the tests in D48229 for the examples of patterns that are related.

Updated the patch to handle extracting from a shl/srl.

For handling srl+shl->bitmask cases I'd prefer to do a follow-up patch since I need a bit more time to look at it.

Yes, absolutely. Currently we don't produce that pattern yet in all the cases, so it is of less importance!

[lebedev.ri]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4847–4857

This should be:

//   (or (shrl (mul v c0) c1)      (mul v c2)) ->
//   (or (shrl (mul v c0) c1) (shl (mul v c3) c4))
//
//   (or       (udiv v c0)     (shl (udiv v c1) c2)) ->
//   (or (shrl (udiv v c3) c4) (shl (udiv v c1) c2))
//
//   (or (shrl (shl v c0) c1)      (shl v c2)) ->
//   (or (shrl (shl v c0) c1) (shl (shl v c3) c4))
//
//   (or       (shrl v c0)     (shl (shrl v c1) c2)) ->
//   (or (shrl (shrl v c1) c3) (shl (shrl v c1) c2))
//
// NEW:
//   (or (shrl (shl v c0) c1)      (mul v c2)) ->
//   (or (shrl (shl v c0) c1) (shl (mul v c3) c4))
//
//   (or       (shrl v c0)     (shl (div v c1) c2)) ->
//   (or (shrl (shrl v c1) c3) (shl (div v c1) c2))

We may not always succeed in converting mul/div into a shift,

4858–4859

You want to add a comment that it will be called for both the LHS and RHS,
so it does not have to worry about looking at the other side here.

[sameconrad]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4847–4857

Could you clarify a bit? The four cases shown in the original code comment were meant to be the valid expansions the function may produce, whereas the added cases don't seem to be valid expansions that could be used to form rotates. For example in your comment the first case was changed to show the LHS of the original shift as (mul v c0) but the LHS of the new shift as (mul v c3) which seems to imply c0 != c3. If the function encountered that it would return an empty SDValue because a rotate can't be formed if the LHS of the existing shift and the leftover LHS of the extracted shift don't match.

(or (shrl (shl v c0) c1) (mul v c2)) ->
// (or (shrl (shl v c0) c1) (shl (mul v c3) c4))

Would InstCombine would generate a case like this that would still provide a valid expansion after extracting the shift (ie where the resulting (mul v c3) is equivalent to (shl v c0))?
In its current state this function should see that the existing shift's LHS (shl v c0) has an opcode that doesn't match the extract op (mul v c3) and just return the SDValue(), because it expects that if InstCombine had merged a shl with another shl the result should have been a greater shl rather than a mul (and thus just assumes that after extracting we couldn't have a leftover op that matches (shl v c0)). I had assumed that if InstCombine merges a mul by a power of 2 with a shift it would produce a greater shift rather than a mul by a greater power of 2 but I would have to test that, if it doesn't and produces a mul then I could probably add some logic to check if the leftover mul can be converted to a shift after extraction to match the existing shift's LHS.

[lebedev.ri]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4847–4857

If the function encountered that it would return an empty SDValue because a rotate can't be formed if the LHS of the existing shift and the leftover LHS of the extracted shift don't match.

So that is the requirement on the end result? That explains it. Then i was wrong.
Can you please state that in the comment, somehow, in addition to the Attempts to expand:.

Would InstCombine would generate a case like this that would still provide a valid expansion after extracting the shift (ie where the resulting (mul v c3) is equivalent to (shl v c0))?

Now that i actually think about, no, i think it's not possible.

I had assumed that if InstCombine merges a mul by a power of 2 with a shift it would produce a greater shift

Yes, pretty sure that is what will happen.

[sameconrad]

Try to fix up clarity of comments

[sameconrad]

Friendly ping. Any more issues with this?

[lebedev.ri]

I can add some nits, but i'm not confident reviewing this as a whole..

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4879–4892

This does basically the same (but opposite) thing twice.
I would personally try to do something like

auto parse = [&IsMulOrDiv, &Opcode, OppShift, ExtractFrom](unsigned Op0, unsigned Op1, unsigned Op2) -> bool {
  if (OppShift.getOpcode() != Op0)
    return false;
  IsMulOrDiv = ExtractFrom.getOpcode() == Op1;
  if (IsMulOrDiv || ExtractFrom.getOpcode() == Op2)
    Opcode = Op2;
  return true;
};
if(!(parse(ISD::SRL, ISD::MUL,  ISD::SHL) ||
     parse(ISD::SHL, ISD::UDIV, ISD::SRL)) ||
   Opcode == ISD::DELETED_NODE)
  return SDValue();

Even if it's the same LOC, it significantly reduced duplication,
allowing to focus on the actual details..

4901

+1 this should be done right ExtractFrom is modified.

4911–4913

What does getAPIntValue() check? We don't want shifts/mul/div by 0?

[RKSimon]

test/CodeGen/X86/rotate-extract.ll

11

Please can you add tests for a mixture of i8, i16, i32 and i64 integers - no need for duplication of all these tests just have decent range.

Ideally we'd include some vector tests as well, possibly in a separate test file.

[sameconrad]

Address code style nits
Update scalar tests to use a variety of types/bitwidths
Add X86 vector tests

[sameconrad]

Forgot to add vector test file

[lebedev.ri]

Some more nits.
Overall, this seems ok, but i don't feel confident reviewing this as a whole :/

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4807

Nit: unneeded {}

4879–4881

Can't you check OppShiftLHS.getOperand(0) != ExtractFrom.getOperand(0) here?

4901

Not done.

4911–4913

I didn't receive a reply here

4922

bitwidth

5091–5093

Do you need those extra variables to to track this?
Maybe you can s/NeedRotLHS/!LHSShift/, and drop the variables?
Otherwise i fear later on they may start meaning different things.

test/CodeGen/X86/rotate-extract.ll

84

These are all rol's, and in aarch64 they are all ror's, even though the tests appear identical, interesting.
Any way to produce both ror and rol instructions for both arches?

[efriedma]

test/CodeGen/X86/rotate-extract.ll

84

AArch64 doesn't have a rol instruction.

[sameconrad]

Address comments

[sameconrad]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4858–4859

This is mentioned in the comments I added in MatchRotate

4879–4881

I moved the opcode check to right after stripConstantMask like you mentioned in your other comment, so this operand(0) check comes right after that now. I think we should verify the ExtractFrom opcode is a shift/mul/udiv before checking operands, otherwise we don't know what the operand count of the instruction is.

4901

I moved the opcode check to the beginning of the function, I believe that was the original request here? The review comments get misaligned after updating the code

4911–4913

Yes, just checks that is not 0. A rotate by zero would just be a noop, so it didn't seem to make sense to try and extract anything in that case.

5091–5093

Thanks for pointing that out. Those bools were completely unneeded.

test/CodeGen/X86/rotate-extract.ll

84

The ISD selection in MatchRotate just looks like HasROTL ? ISD::ROTL : ISD::ROTR, so it looks will just always choose a rotl if available and fall back to a rotr like it does on AArch64. Since X86 has rotl it will always select that, I don't see any way to vary it within the same ISA.

[lebedev.ri]

It seems everyone else is busy with other differentials, and don't have time to spare to review this :/

Have you run this on a llvm test suite?

I don't have further comments, and i could accept this, if testsuite passes,
but it would be more comfortable if someone else could review this, too.

[sameconrad]

In D47681#1162716, @lebedev.ri wrote:

It seems everyone else is busy with other differentials, and don't have time to spare to review this :/

Have you run this on a llvm test suite?

I don't have further comments, and i could accept this, if testsuite passes,
but it would be more comfortable if someone else could review this, too.

Sorry about the late response, missed this comment before. I just ran the test-suite and it looks like all 2555 passed without issue.

[lebedev.ri]

In D47681#1167516, @sameconrad wrote:

In D47681#1162716, @lebedev.ri wrote:

It seems everyone else is busy with other differentials, and don't have time to spare to review this :/

Have you run this on a llvm test suite?

I don't have further comments, and i could accept this, if testsuite passes,
but it would be more comfortable if someone else could review this, too.

Sorry about the late response, missed this comment before. I just ran the test-suite and it looks like all 2555 passed without issue.

Ok then, LG.
Please wait a bit (until monday?) just in case others may want to comment, and then feel free to commit.

[RKSimon]

I've added your tests to trunk with current codegen so you need to rebase, which will show the effects of your patch.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

4894

(style) full stops at the end of comments.

4908

It'd be great if this could support non-uniform constant vectors but a TODO comment is fine for now.

[sameconrad]

Rebase after tests were committed
Add full stops to comments
Add todo for handling non-splat vectors

[sameconrad]

Updated after rebasing
If this looks good someone else will have to commit since I don't have commit access
Thanks!

[lebedev.ri]

Last-minute thought - should this be limited to some specific value types?
On X86, the RCL/RCR/ROL/ROR only exist for 8/16/32/64 bit widths.

[RKSimon]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

5065

@lebedev.ri This should ensure legality of the ROT

[lebedev.ri]

In D47681#1169179, @sameconrad wrote:

Updated after rebasing
If this looks good someone else will have to commit since I don't have commit access
Thanks!

Two other landed differentials over the last year, might as well ask for the commit rights and land it yourself :)

[RKSimon]

@sameconrad Are you getting commit access (and in time for the Aug 1 branch)?

[sameconrad]

In D47681#1179976, @RKSimon wrote:

@sameconrad Are you getting commit access (and in time for the Aug 1 branch)?

I requested it but haven't heard anything back yet, so I'm not sure if I'll be able to commit by August 1st.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

private:
     ///
     /// Notice that other chains may still be replaced even if the function
     /// returns false.
     bool findBetterNeighborChains(StoreSDNode *St);
 
     /// Match "(X shl/srl V1) & V2" where V2 may not be present.
     bool MatchRotateHalf(SDValue Op, SDValue &Shift, SDValue &Mask);
 
+    /// Helper function for visitOR to extract one side of a rotate idiom
+    /// from a mul or div operation.  This is meant to handle cases where
+    /// InstCombine merged some outside mul or div with one of the shifts from
+    /// the rotate pattern. 
+    /// Attempts to fold: 
+    ///   (or (shrl (mul v c0) c1) (mul v c2)) ->
+    ///   (or (shrl (mul v c0) c1) (shl (mul v c0) c3))
+    /// and
+    ///   (or (udiv v c0) (shl (udiv v c1) c2)) ->
+    ///   (or (shrl (udiv v c1) c3) (shl (udiv v c1) c2))
+    SDValue extractShiftFromMulOrUDiv(SDValue OpposingShift, SDValue MulOrDiv,
+                                      SDValue &Mask, const SDLoc &DL);
+
     /// Holds a pointer to an LSBaseSDNode as well as information on where it
     /// is located in a sequence of memory operations connected by a chain.
     struct MemOpLink {
       // Ptr to the mem node.
       LSBaseSDNode *MemNode;
 
       // Offset from the base ptr.
       int64_t OffsetFromBase;
SDValue DAGCombiner::visitOR(SDNode *N) {
   }
 
   // Simplify: (or (op x...), (op y...))  -> (op (or x, y))
   if (N0.getOpcode() == N1.getOpcode())
     if (SDValue Tmp = SimplifyBinOpWithSameOpcodeHands(N))
       return Tmp;
 
   // See if this is some rotate idiom.
   if (SDNode *Rot = MatchRotate(N0, N1, SDLoc(N)))
     return SDValue(Rot, 0);
 
   if (SDValue Load = MatchLoadCombine(N))
     return Load;
 
   // Simplify the operands using demanded-bits information.
   if (SimplifyDemandedBits(SDValue(N, 0)))
     return SDValue(N, 0);
 
   return SDValue();
 }
 
-/// Match "(X shl/srl V1) & V2" where V2 may not be present.
-bool DAGCombiner::MatchRotateHalf(SDValue Op, SDValue &Shift, SDValue &Mask) {
+static SDValue stripConstantMask(SelectionDAG &DAG, SDValue Op, SDValue &Mask) {
   if (Op.getOpcode() == ISD::AND) {
     if (DAG.isConstantIntBuildVectorOrConstantInt(Op.getOperand(1))) {
       Mask = Op.getOperand(1);
-      Op = Op.getOperand(0);
-    } else {
-      return false;
+      return Op.getOperand(0);
     }
   }
+  return Op;
+}
+
+/// Match "(X shl/srl V1) & V2" where V2 may not be present.
+bool DAGCombiner::MatchRotateHalf(SDValue Op, SDValue &Shift, SDValue &Mask) {
+  Op = stripConstantMask(DAG, Op, Mask);
 
   if (Op.getOpcode() == ISD::SRL || Op.getOpcode() == ISD::SHL) {
     Shift = Op;
     return true;
   }
 
   return false;
 }
 
+SDValue DAGCombiner::extractShiftFromMulOrUDiv(SDValue OpposingShift,
+                                               SDValue MulOrDiv, SDValue &Mask,
+                                               const SDLoc &DL) {
+  assert(OpposingShift && MulOrDiv &&
+         (OpposingShift.getOpcode() == ISD::SHL ||
+          OpposingShift.getOpcode() == ISD::SRL));
+
+  MulOrDiv = stripConstantMask(DAG, MulOrDiv, Mask);
+  SDValue ShiftLHS = OpposingShift.getOperand(0);
+  EVT ShiftedVT = ShiftLHS.getValueType();
+
+  // Preconditions:
+  // 
+  //    or( (shiftl/r (op0 v c0) c1) (op1 v c2) )
+  //
+  // op1 is a mul or udiv and is the same as op0
+  // op1 and op0 have v as the LHS input and produce the same value type
+  if ((ShiftLHS.getOpcode() != ISD::MUL && ShiftLHS.getOpcode() != ISD::UDIV) ||
+      ShiftLHS.getOpcode() != MulOrDiv.getOpcode() ||
+      ShiftLHS.getOperand(0) != MulOrDiv.getOperand(0) ||
+      ShiftedVT != MulOrDiv.getValueType())
+    return SDValue();
+
+  auto GetConstant = [](SDValue Operand) {
+    if (auto *Node = isConstOrConstSplat(Operand))
+      return Node->getAPIntValue();
+    return APInt();
+  };
+  // Amount of the existing shift
+  APInt OppShiftAmt = GetConstant(OpposingShift.getOperand(1));
+  // Constant mul/div amount from the RHS of the shift's LHS op
+  APInt ShiftLHSMulOrDivAmount = GetConstant(ShiftLHS.getOperand(1));
+  // Constant mul/div amount from the RHS of the MulOrDiv op
+  APInt MulOrDivAmount = GetConstant(MulOrDiv.getOperand(1));
+  if (OppShiftAmt == 0 | ShiftLHSMulOrDivAmount == 0 || MulOrDivAmount == 0)
+    return SDValue();
+
+  // Check:
+  //     c2 / (1 << (bitwidth(op0 v c0) - c1)) == c0
+  //     c2 % (1 << (bitwidth(op0 v c0) - c1)) == 0
+  const unsigned VTWidth = ShiftedVT.getScalarSizeInBits(); 
+  const APInt NewShiftAmt = VTWidth - OppShiftAmt;
+  const APInt ShiftExtractDiv = APInt(VTWidth, 1) << NewShiftAmt;
+  APInt NewMulOrDivAmount;
+  APInt Rem;
+  APInt::udivrem(MulOrDivAmount, ShiftExtractDiv, NewMulOrDivAmount, Rem);
+  if (Rem != 0 || NewMulOrDivAmount != ShiftLHSMulOrDivAmount)
+    return SDValue();
+
+  unsigned Opcode = 0;
+  // If both are muls and the opposing is an shrl, we can extract a shl
+  if (MulOrDiv.getOpcode() == ISD::MUL && OpposingShift.getOpcode() == ISD::SRL)
+    Opcode = ISD::SHL;
+  // If both are udivs and the opposing is an shl, we can extract a shrl
+  else if (MulOrDiv.getOpcode() == ISD::UDIV && OpposingShift.getOpcode() == ISD::SHL)
+    Opcode = ISD::SRL;
+  // Can't extract
+  else
+    return SDValue();
+
+  // Fold:
+  //
+  //     or( (shiftl/r (op0 v c0) c1) (op0 v c2) ) ->
+  //     or( (shiftl/r (op0 v c0) c1) (shiftr/l (op0 v c0) c3) )
+  //
+  // such that c1 + c3 == bitwidth(op0 v c0) and a rotate can be formed
+  EVT ShiftVT = OpposingShift.getOperand(1).getValueType();
+  EVT ResVT = MulOrDiv.getValueType();
+  SDValue NewShiftNode = DAG.getConstant(NewShiftAmt, DL, ShiftVT);
+  return DAG.getNode(Opcode, DL, ResVT, ShiftLHS, NewShiftNode);
+}
+
 // Return true if we can prove that, whenever Neg and Pos are both in the
 // range [0, EltSize), Neg == (Pos == 0 ? 0 : EltSize - Pos).  This means that
 // for two opposing shifts shift1 and shift2 and a value X with OpBits bits:
 //
 //     (or (shift1 X, Neg), (shift2 X, Pos))
 //
 // reduces to a rotate in direction shift2 by Pos or (equivalently) a rotate
 // in direction shift1 by Neg.  The range [0, EltSize) means that we only need
 // to consider shift amounts with defined behavior.
 static bool matchRotateSub(SDValue Pos, SDValue Neg, unsigned EltSize,
                            SelectionDAG &DAG) {
   // If EltSize is a power of 2 then:
   //
   //  (a) (Pos == 0 ? 0 : EltSize - Pos) == (EltSize - Pos) & (EltSize - 1)
   //  (b) Neg == Neg & (EltSize - 1) whenever Neg is in [0, EltSize).
   //
   // So if EltSize is a power of 2 and Neg is (and Neg', EltSize-1), we check
SDNode *DAGCombiner::MatchRotatePosNeg(SDValue Shifted, SDValue Pos,
   }
 
   return nullptr;
 }
 
 // MatchRotate - Handle an 'or' of two operands.  If this is one of the many
 // idioms for rotate, and if the target supports rotation instructions, generate
 // a rot[lr].
 SDNode *DAGCombiner::MatchRotate(SDValue LHS, SDValue RHS, const SDLoc &DL) {
   // Must be a legal type.  Expanded 'n promoted things won't work with rotates.
   EVT VT = LHS.getValueType();
   if (!TLI.isTypeLegal(VT)) return nullptr;
 
   // The target must have at least one rotate flavor.
   bool HasROTL = TLI.isOperationLegalOrCustom(ISD::ROTL, VT);
   bool HasROTR = TLI.isOperationLegalOrCustom(ISD::ROTR, VT);
   if (!HasROTL && !HasROTR) return nullptr;
 
   // Check for truncated rotate.
   if (LHS.getOpcode() == ISD::TRUNCATE && RHS.getOpcode() == ISD::TRUNCATE &&
       LHS.getOperand(0).getValueType() == RHS.getOperand(0).getValueType()) {
     assert(LHS.getValueType() == RHS.getValueType());
     if (SDNode *Rot = MatchRotate(LHS.getOperand(0), RHS.getOperand(0), DL)) {
       return DAG.getNode(ISD::TRUNCATE, SDLoc(LHS), LHS.getValueType(),
                          SDValue(Rot, 0)).getNode();
     }
   }
 
+  bool NeedRotLHS = true;
+  bool NeedRotRHS = true;
+
   // Match "(X shl/srl V1) & V2" where V2 may not be present.
   SDValue LHSShift;   // The shift.
   SDValue LHSMask;    // AND value if any.
-  if (!MatchRotateHalf(LHS, LHSShift, LHSMask))
-    return nullptr; // Not part of a rotate.
+  if (MatchRotateHalf(LHS, LHSShift, LHSMask))
+    NeedRotLHS = false;
 
   SDValue RHSShift;   // The shift.
   SDValue RHSMask;    // AND value if any.
-  if (!MatchRotateHalf(RHS, RHSShift, RHSMask))
-    return nullptr; // Not part of a rotate.
+  if (MatchRotateHalf(RHS, RHSShift, RHSMask))
+    NeedRotRHS = false;
+
+  // If neither side matched a rotate half, bail
+  if (NeedRotLHS && NeedRotRHS)
+    return nullptr;
+
+  // Iff one side matched a rotate half, we may be able to extract the
+  // needed shift from the opposite side if it is a constant mul or udiv.
+  // We also need to pass the matched shift to compute the needed shift amount
+  // to extract (example: if the side with the shift is (shl v c0), and v has a
+  // bitwidth of 32, then the needed shift to extract from the opposite side is
+  // computed to be (srl v (32 - c0)))
+
+  // If we only got the LHS side of the rotate, try to extract the needed
+  // shift from the RHS
+  if (NeedRotRHS) {
+    if (!(RHSShift = extractShiftFromMulOrUDiv(
+              /*matched shift*/ LHSShift, /*extract from*/ RHS, RHSMask, DL)))
+      return nullptr;
+  }
+  // If we only got the RHS side of the rotate, try to extract the needed
+  // shift from the LHS
+  if (NeedRotLHS) {
+    if (!(LHSShift = extractShiftFromMulOrUDiv(
+              /*matched shift*/ RHSShift, /*extract from*/ LHS, LHSMask, DL)))
+      return nullptr;
+  }
+
+  // At this point we've matched a shift op on each side
 
   if (LHSShift.getOperand(0) != RHSShift.getOperand(0))
     return nullptr;   // Not shifting the same value.
 
   if (LHSShift.getOpcode() == RHSShift.getOpcode())
     return nullptr;   // Shifts must disagree.
 
   // Canonicalize shl to left side in a shl/srl pair.

test/CodeGen/AArch64/rotate-extract.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=aarch64-unknown-unknown | FileCheck %s
+
+; Check that under certain conditions we can factor out a rotate
+; from the following idioms:
+;   (a*c0) >> s1 | (a*c1)
+;   (a/c0) << s1 | (a/c1)
+; This targets cases where instcombine has folded a mul or
+; udiv with one of the shifts of the normal rotate idiom
+
+define i64 @roll_extract_mul(i64 %i) nounwind {
+; CHECK-LABEL: roll_extract_mul:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x0, x0, lsl #3
+; CHECK-NEXT:    ror x0, x8, #57
+; CHECK-NEXT:    ret
+  %lhs_mul = mul i64 %i, 9
+  %rhs_mul = mul i64 %i, 1152
+  %lhs_shift = lshr i64 %lhs_mul, 57
+  %out = or i64 %lhs_shift, %rhs_mul
+  ret i64 %out
+}
+
+define i64 @roll_extract_udiv(i64 %i) nounwind {
+; CHECK-LABEL: roll_extract_udiv:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov x8, #-6148914691236517206
+; CHECK-NEXT:    movk x8, #43691
+; CHECK-NEXT:    umulh x8, x0, x8
+; CHECK-NEXT:    lsr x8, x8, #1
+; CHECK-NEXT:    ror x0, x8, #4
+; CHECK-NEXT:    ret
+  %lhs_div = udiv i64 %i, 3
+  %rhs_div = udiv i64 %i, 48
+  %lhs_shift = shl i64 %lhs_div, 60
+  %out = or i64 %lhs_shift, %rhs_div
+  ret i64 %out
+}
+
+; Can factor 128 from 2304, but result is 18 instead of 9
+define i64 @no_extract_mul(i64 %i) nounwind {
+; CHECK-LABEL: no_extract_mul:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    add x8, x0, x0, lsl #3
+; CHECK-NEXT:    lsr x0, x8, #57
+; CHECK-NEXT:    bfi x0, x8, #8, #56
+; CHECK-NEXT:    ret
+  %lhs_mul = mul i64 %i, 9
+  %rhs_mul = mul i64 %i, 2304
+  %lhs_shift = lshr i64 %lhs_mul, 57
+  %out = or i64 %lhs_shift, %rhs_mul
+  ret i64 %out
+}
+
+; Can't evenly factor 16 from 49
+define i64 @no_extract_udiv(i64 %i) nounwind {
+; CHECK-LABEL: no_extract_udiv:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov x9, #38787
+; CHECK-NEXT:    movk x9, #61523, lsl #16
+; CHECK-NEXT:    movk x9, #2674, lsl #32
+; CHECK-NEXT:    movk x9, #20062, lsl #48
+; CHECK-NEXT:    mov x8, #-6148914691236517206
+; CHECK-NEXT:    umulh x9, x0, x9
+; CHECK-NEXT:    movk x8, #43691
+; CHECK-NEXT:    sub x10, x0, x9
+; CHECK-NEXT:    umulh x8, x0, x8
+; CHECK-NEXT:    add x9, x9, x10, lsr #1
+; CHECK-NEXT:    lsr x8, x8, #1
+; CHECK-NEXT:    lsr x0, x9, #5
+; CHECK-NEXT:    bfi x0, x8, #60, #4
+; CHECK-NEXT:    ret
+  %lhs_div = udiv i64 %i, 3
+  %rhs_div = udiv i64 %i, 49
+  %lhs_shift = shl i64 %lhs_div, 60
+  %out = or i64 %lhs_shift, %rhs_div
+  ret i64 %out
+}

test/CodeGen/X86/rotate-extract.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s -mtriple=x86_64-unknown-unknown | FileCheck %s
+
+; Check that under certain conditions we can factor out a rotate
+; from the following idioms:
+;   (a*c0) >> s1 | (a*c1)
+;   (a/c0) << s1 | (a/c1)
+; This targets cases where instcombine has folded a mul or
+; udiv with one of the shifts of the normal rotate idiom
+
+define i64 @roll_extract_mul(i64 %i) nounwind {
+; CHECK-LABEL: roll_extract_mul:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    leaq (%rdi,%rdi,8), %rax
+; CHECK-NEXT:    rolq $7, %rax
+; CHECK-NEXT:    retq
+  %lhs_mul = mul i64 %i, 9
+  %rhs_mul = mul i64 %i, 1152
+  %lhs_shift = lshr i64 %lhs_mul, 57
+  %out = or i64 %lhs_shift, %rhs_mul
+  ret i64 %out
+}
+
+define i64 @roll_extract_udiv(i64 %i) nounwind {
+; CHECK-LABEL: roll_extract_udiv:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    movabsq $-6148914691236517205, %rcx # imm = 0xAAAAAAAAAAAAAAAB
+; CHECK-NEXT:    movq %rdi, %rax
+; CHECK-NEXT:    mulq %rcx
+; CHECK-NEXT:    shrq %rdx
+; CHECK-NEXT:    rolq $60, %rdx
+; CHECK-NEXT:    movq %rdx, %rax
+; CHECK-NEXT:    retq
+  %lhs_div = udiv i64 %i, 3
+  %rhs_div = udiv i64 %i, 48
+  %lhs_shift = shl i64 %lhs_div, 60
+  %out = or i64 %lhs_shift, %rhs_div
+  ret i64 %out
+}
+
+; Can factor 128 from 2304, but result is 18 instead of 9
+define i64 @no_extract_mul(i64 %i) nounwind {
+; CHECK-LABEL: no_extract_mul:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    leaq (%rdi,%rdi,8), %rax
+; CHECK-NEXT:    shlq $8, %rdi
+; CHECK-NEXT:    leaq (%rdi,%rdi,8), %rcx
+; CHECK-NEXT:    shrq $57, %rax
+; CHECK-NEXT:    orq %rcx, %rax
+; CHECK-NEXT:    retq
+  %lhs_mul = mul i64 %i, 9
+  %rhs_mul = mul i64 %i, 2304
+  %lhs_shift = lshr i64 %lhs_mul, 57
+  %out = or i64 %lhs_shift, %rhs_mul
+  ret i64 %out
+}
+
+; Can't evenly factor 16 from 49
+define i64 @no_extract_udiv(i64 %i) nounwind {
+; CHECK-LABEL: no_extract_udiv:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    movq %rdi, %rcx
+; CHECK-NEXT:    movabsq $-6148914691236517205, %rdx # imm = 0xAAAAAAAAAAAAAAAB
+; CHECK-NEXT:    movq %rdi, %rax
+; CHECK-NEXT:    mulq %rdx
+; CHECK-NEXT:    movq %rdx, %rsi
+; CHECK-NEXT:    andq $-2, %rsi
+; CHECK-NEXT:    shlq $59, %rsi
+; CHECK-NEXT:    movabsq $5646962471543740291, %rdx # imm = 0x4E5E0A72F0539783
+; CHECK-NEXT:    movq %rdi, %rax
+; CHECK-NEXT:    mulq %rdx
+; CHECK-NEXT:    subq %rdx, %rcx
+; CHECK-NEXT:    shrq %rcx
+; CHECK-NEXT:    addq %rdx, %rcx
+; CHECK-NEXT:    shrq $5, %rcx
+; CHECK-NEXT:    leaq (%rcx,%rsi), %rax
+; CHECK-NEXT:    retq
+  %lhs_div = udiv i64 %i, 3
+  %rhs_div = udiv i64 %i, 49
+  %lhs_shift = shl i64 %lhs_div, 60
+  %out = or i64 %lhs_shift, %rhs_div
+  ret i64 %out
+}