1 files changed, 29 insertions, 9 deletions

diff --git a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
index c98cea8a74b..f6c82cb2b35 100644
--- a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
@@ -3976,22 +3976,36 @@ void X86DAGToDAGISel::Select(SDNode *Node) {
 
     // For operations of the form (x << C1) op C2, check if we can use a smaller
     // encoding for C2 by transforming it into (x op (C2>>C1)) << C1.
-    SDValue N0 = Node->getOperand(0);
+    SDValue Shift = Node->getOperand(0);
     SDValue N1 = Node->getOperand(1);
 
-    if (N0->getOpcode() != ISD::SHL || !N0->hasOneUse())
+    ConstantSDNode *Cst = dyn_cast<ConstantSDNode>(N1);
+    if (!Cst)
+      break;
+
+    // If we have an any_extend feeding the AND, look through it to see if there
+    // is a shift behind it. But only if the AND doesn't use the extended bits.
+    // FIXME: Generalize this to other ANY_EXTEND than i32 to i64?
+    int64_t Val = Cst->getSExtValue();
+    bool FoundAnyExtend = false;
+    if (Shift.getOpcode() == ISD::ANY_EXTEND && Shift.hasOneUse() &&
+        Shift.getOperand(0).getSimpleValueType() == MVT::i32 &&
+        isUInt<32>(Val)) {
+      FoundAnyExtend = true;
+      Shift = Shift.getOperand(0);
+    }
+
+    if (Shift.getOpcode() != ISD::SHL || !Shift.hasOneUse())
       break;
 
     // i8 is unshrinkable, i16 should be promoted to i32.
     if (NVT != MVT::i32 && NVT != MVT::i64)
       break;
 
-    ConstantSDNode *Cst = dyn_cast<ConstantSDNode>(N1);
-    ConstantSDNode *ShlCst = dyn_cast<ConstantSDNode>(N0->getOperand(1));
-    if (!Cst || !ShlCst)
+    ConstantSDNode *ShlCst = dyn_cast<ConstantSDNode>(Shift.getOperand(1));
+    if (!ShlCst)
       break;
 
-    int64_t Val = Cst->getSExtValue();
     uint64_t ShAmt = ShlCst->getZExtValue();
 
     // Make sure that we don't change the operation by removing bits.
@@ -4028,13 +4042,19 @@ void X86DAGToDAGISel::Select(SDNode *Node) {
 
     int64_t ShiftedVal;
     if (CanShrinkImmediate(ShiftedVal)) {
+      SDValue X = Shift.getOperand(0);
+      if (FoundAnyExtend) {
+        SDValue NewX = CurDAG->getNode(ISD::ANY_EXTEND, dl, NVT, X);
+        insertDAGNode(*CurDAG, SDValue(Node, 0), NewX);
+        X = NewX;
+      }
+
       SDValue NewCst = CurDAG->getConstant(ShiftedVal, dl, NVT);
       insertDAGNode(*CurDAG, SDValue(Node, 0), NewCst);
-      SDValue NewBinOp = CurDAG->getNode(Opcode, dl, NVT, N0->getOperand(0),
-                                         NewCst);
+      SDValue NewBinOp = CurDAG->getNode(Opcode, dl, NVT, X, NewCst);
       insertDAGNode(*CurDAG, SDValue(Node, 0), NewBinOp);
       SDValue NewSHL = CurDAG->getNode(ISD::SHL, dl, NVT, NewBinOp,
-                                       N0->getOperand(1));
+                                       Shift.getOperand(1));
       ReplaceNode(Node, NewSHL.getNode());
       SelectCode(NewSHL.getNode());
       return;