1 files changed, 22 insertions, 1 deletions
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 5c2e13134df..4ce1987e857 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -35836,7 +35836,7 @@ static SDValue combineGatherScatter(SDNode *N, SelectionDAG &DAG,
   // Pre-shrink oversized index elements to avoid triggering scalarization.
   if (DCI.isBeforeLegalize()) {
     SDValue Index = N->getOperand(4);
-    if (Index.getValueType().getScalarSizeInBits() > 64) {
+    if (Index.getScalarValueSizeInBits() > 64) {
       EVT IndexVT = EVT::getVectorVT(*DAG.getContext(), MVT::i64,
                                    Index.getValueType().getVectorNumElements());
       SDValue Trunc = DAG.getNode(ISD::TRUNCATE, DL, IndexVT, Index);
@@ -35848,6 +35848,27 @@ static SDValue combineGatherScatter(SDNode *N, SelectionDAG &DAG,
     }
   }
 
+  // Try to remove sign extends from i32 to i64 on the index.
+  // Only do this before legalize in case we are relying on it for
+  // legalization.
+  // TODO: We should maybe remove any sign extend once we learn how to sign
+  // extend narrow index during lowering.
+  if (DCI.isBeforeLegalizeOps()) {
+    SDValue Index = N->getOperand(4);
+    if (Index.getScalarValueSizeInBits() == 64 &&
+        Index.getOpcode() == ISD::SIGN_EXTEND &&
+        Index.getOperand(0).getScalarValueSizeInBits() == 32) {
+      SmallVector<SDValue, 5> NewOps(N->op_begin(), N->op_end());
+      NewOps[4] = Index.getOperand(0);
+      DAG.UpdateNodeOperands(N, NewOps);
+      // The original sign extend has less users, add back to worklist in case
+      // it needs to be removed.
+      DCI.AddToWorklist(Index.getNode());
+      DCI.AddToWorklist(N);
+      return SDValue(N, 0);
+    }
+  }
+
   // Gather and Scatter instructions use k-registers for masks. The type of
   // the masks is v*i1. So the mask will be truncated anyway.
   // The SIGN_EXTEND_INREG my be dropped.