[llvm] e283ef7 - [X86] matchAddressRecursively - add foldMaskAndShiftToScale handling to ZERO_EXTEND nodes.

[Simon Pilgrim via llvm-commits]

Author: Simon Pilgrim
Date: 2023-08-24T11:47:07+01:00
New Revision: e283ef7e936b8fca9e7216ecc61003e7f9790923

URL: https://github.com/llvm/llvm-project/commit/e283ef7e936b8fca9e7216ecc61003e7f9790923
DIFF: https://github.com/llvm/llvm-project/commit/e283ef7e936b8fca9e7216ecc61003e7f9790923.diff

LOG: [X86] matchAddressRecursively - add foldMaskAndShiftToScale handling to ZERO_EXTEND nodes.

Added: 
    

Modified: 
    llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
    llvm/test/CodeGen/X86/fold-and-shift-x86_64.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
index f1eb1f4e85fb5b..600e7c85a811e9 100644
--- a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
@@ -2121,11 +2121,14 @@ static bool foldMaskAndShiftToScale(SelectionDAG &DAG, SDValue N,
     insertDAGNode(DAG, N, NewX);
     X = NewX;
   }
+
+  MVT XVT = X.getSimpleValueType();
   SDLoc DL(N);
   SDValue NewSRLAmt = DAG.getConstant(ShiftAmt + AMShiftAmt, DL, MVT::i8);
-  SDValue NewSRL = DAG.getNode(ISD::SRL, DL, VT, X, NewSRLAmt);
+  SDValue NewSRL = DAG.getNode(ISD::SRL, DL, XVT, X, NewSRLAmt);
+  SDValue NewExt = DAG.getZExtOrTrunc(NewSRL, DL, VT);
   SDValue NewSHLAmt = DAG.getConstant(AMShiftAmt, DL, MVT::i8);
-  SDValue NewSHL = DAG.getNode(ISD::SHL, DL, VT, NewSRL, NewSHLAmt);
+  SDValue NewSHL = DAG.getNode(ISD::SHL, DL, VT, NewExt, NewSHLAmt);
 
   // Insert the new nodes into the topological ordering. We must do this in
   // a valid topological ordering as nothing is going to go back and re-sort
@@ -2134,13 +2137,14 @@ static bool foldMaskAndShiftToScale(SelectionDAG &DAG, SDValue N,
   // hierarchy left to express.
   insertDAGNode(DAG, N, NewSRLAmt);
   insertDAGNode(DAG, N, NewSRL);
+  insertDAGNode(DAG, N, NewExt);
   insertDAGNode(DAG, N, NewSHLAmt);
   insertDAGNode(DAG, N, NewSHL);
   DAG.ReplaceAllUsesWith(N, NewSHL);
   DAG.RemoveDeadNode(N.getNode());
 
   AM.Scale = 1 << AMShiftAmt;
-  AM.IndexReg = NewSRL;
+  AM.IndexReg = NewExt;
   return false;
 }
 
@@ -2619,11 +2623,17 @@ bool X86DAGToDAGISel::matchAddressRecursively(SDValue N, X86ISelAddressMode &AM,
       return false;
     }
 
-    // Try to fold the mask and shift into an extract and scale.
-    if (Src.getOpcode() == ISD::SRL && !Mask.isAllOnes() &&
-        !foldMaskAndShiftToExtract(*CurDAG, N, Mask.getZExtValue(), Src,
+    if (Src.getOpcode() == ISD::SRL && !Mask.isAllOnes()) {
+      // Try to fold the mask and shift into an extract and scale.
+      if (!foldMaskAndShiftToExtract(*CurDAG, N, Mask.getZExtValue(), Src,
+                                     Src.getOperand(0), AM))
+        return false;
+
+      // Try to fold the mask and shift directly into the scale.
+      if (!foldMaskAndShiftToScale(*CurDAG, N, Mask.getZExtValue(), Src,
                                    Src.getOperand(0), AM))
-      return false;
+        return false;
+    }
 
     break;
   }

diff  --git a/llvm/test/CodeGen/X86/fold-and-shift-x86_64.ll b/llvm/test/CodeGen/X86/fold-and-shift-x86_64.ll
index 964221d2434969..fabf7581e9b3cf 100644
--- a/llvm/test/CodeGen/X86/fold-and-shift-x86_64.ll
+++ b/llvm/test/CodeGen/X86/fold-and-shift-x86_64.ll
@@ -96,9 +96,8 @@ define i32 @t7(<16 x i8> %a0, ptr %p0) {
 ; CHECK-LABEL: t7:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    pmovmskb %xmm0, %eax
-; CHECK-NEXT:    shrl %eax
-; CHECK-NEXT:    andl $-4, %eax
-; CHECK-NEXT:    movzbl (%rdi,%rax), %eax
+; CHECK-NEXT:    shrl $3, %eax
+; CHECK-NEXT:    movzbl (%rdi,%rax,4), %eax
 ; CHECK-NEXT:    retq
   %i = call i32 @llvm.x86.sse2.pmovmskb.128(<16 x i8> %a0)
   %index = lshr i32 %i, 1