From 8b8264828cbc4de633dfce90a7185c52a83687dc Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper@intel.com>
Date: Sat, 20 Apr 2019 04:38:49 +0000
Subject: [PATCH] [X86] Turn (and (anyextend (shl X, C1), C2)) into (shl (and
 (anyextend X), (C1 >> C2), C2) if the AND could match a movzx.

There's one slight regression in here because we don't check that the immediate
already allowed movzx before the shift. I'll fix that next.

llvm-svn: 358804
---
 llvm/lib/Target/X86/X86ISelDAGToDAG.cpp | 38 +++++++++++++++++++++++++--------
 llvm/test/CodeGen/X86/narrow-shl-cst.ll | 15 ++++++-------
 2 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
index c98cea8..f6c82cb 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;
diff --git a/llvm/test/CodeGen/X86/narrow-shl-cst.ll b/llvm/test/CodeGen/X86/narrow-shl-cst.ll
index 86fc7a1..1d28cb9 100644
--- a/llvm/test/CodeGen/X86/narrow-shl-cst.ll
+++ b/llvm/test/CodeGen/X86/narrow-shl-cst.ll
@@ -152,8 +152,9 @@ define i32 @test12(i32 %x, i32* %y) nounwind {
 define i64 @test13(i64 %x, i64* %y) nounwind {
 ; CHECK-LABEL: test13:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    addl %edi, %edi
-; CHECK-NEXT:    movzbl %dil, %eax
+; CHECK-NEXT:    movq %rdi, %rax
+; CHECK-NEXT:    andl $127, %eax
+; CHECK-NEXT:    addq %rax, %rax
 ; CHECK-NEXT:    movq %rax, (%rsi)
 ; CHECK-NEXT:    retq
   %and = shl i64 %x, 1
@@ -212,9 +213,8 @@ define i32 @test17(i32 %x) nounwind {
 define i64 @test18(i64 %x) nounwind {
 ; CHECK-LABEL: test18:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movq %rdi, %rax
-; CHECK-NEXT:    shll $10, %eax
-; CHECK-NEXT:    andl $261120, %eax # imm = 0x3FC00
+; CHECK-NEXT:    movzbl %dil, %eax
+; CHECK-NEXT:    shlq $10, %rax
 ; CHECK-NEXT:    retq
   %and = shl i64 %x, 10
   %shl = and i64 %and, 261120
@@ -235,9 +235,8 @@ define i32 @test19(i32 %x) nounwind {
 define i64 @test20(i64 %x) nounwind {
 ; CHECK-LABEL: test20:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movq %rdi, %rax
-; CHECK-NEXT:    shll $10, %eax
-; CHECK-NEXT:    andl $67107840, %eax # imm = 0x3FFFC00
+; CHECK-NEXT:    movzwl %di, %eax
+; CHECK-NEXT:    shlq $10, %rax
 ; CHECK-NEXT:    retq
   %and = shl i64 %x, 10
   %shl = and i64 %and, 67107840
-- 
2.7.4