1 files changed, 26 insertions, 44 deletions

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 052300d6f72..54ca3721cf5 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -42572,16 +42572,17 @@ static SDValue combineGatherScatter(SDNode *N, SelectionDAG &DAG,
   SDValue Base = GorS->getBasePtr();
   SDValue Scale = GorS->getScale();
 
-  // Shrink constant indices if they are larger than 32-bits.
-  // Only do this before legalize types since v2i64 could become v2i32.
-  // FIXME: We could check that the type is legal if we're after legalize types,
-  // but then we would need to construct test cases where that happens.
-  // FIXME: We could support more than just constant vectors, but we need to
-  // careful with costing. A truncate that can be optimized out would be fine.
-  // Otherwise we might only want to create a truncate if it avoids a split.
   if (DCI.isBeforeLegalize()) {
+    unsigned IndexWidth = Index.getScalarValueSizeInBits();
+
+    // Shrink constant indices if they are larger than 32-bits.
+    // Only do this before legalize types since v2i64 could become v2i32.
+    // FIXME: We could check that the type is legal if we're after legalize
+    // types, but then we would need to construct test cases where that happens.
+    // FIXME: We could support more than just constant vectors, but we need to
+    // careful with costing. A truncate that can be optimized out would be fine.
+    // Otherwise we might only want to create a truncate if it avoids a split.
     if (auto *BV = dyn_cast<BuildVectorSDNode>(Index)) {
-      unsigned IndexWidth = Index.getScalarValueSizeInBits();
       if (BV->isConstant() && IndexWidth > 32 &&
           DAG.ComputeNumSignBits(Index) > (IndexWidth - 32)) {
         unsigned NumElts = Index.getValueType().getVectorNumElements();
@@ -42604,16 +42605,18 @@ static SDValue combineGatherScatter(SDNode *N, SelectionDAG &DAG,
                                     Scatter->getIndexType());
       }
     }
-  }
 
-  if (DCI.isBeforeLegalizeOps()) {
-    // Remove any sign extends from 32 or smaller to larger than 32.
-    // Only do this before LegalizeOps in case we need the sign extend for
-    // legalization.
-    if (Index.getOpcode() == ISD::SIGN_EXTEND &&
-        Index.getScalarValueSizeInBits() > 32 &&
-        Index.getOperand(0).getScalarValueSizeInBits() <= 32) {
-      Index = Index.getOperand(0);
+    // Shrink any sign/zero extends from 32 or smaller to larger than 32 if
+    // there are sufficient sign bits. Only do this before legalize types to
+    // avoid creating illegal types in truncate.
+    if ((Index.getOpcode() == ISD::SIGN_EXTEND ||
+         Index.getOpcode() == ISD::ZERO_EXTEND) &&
+        IndexWidth > 32 &&
+        Index.getOperand(0).getScalarValueSizeInBits() <= 32 &&
+        DAG.ComputeNumSignBits(Index) > (IndexWidth - 32)) {
+      unsigned NumElts = Index.getValueType().getVectorNumElements();
+      EVT NewVT = EVT::getVectorVT(*DAG.getContext(), MVT::i32, NumElts);
+      Index = DAG.getNode(ISD::TRUNCATE, DL, NewVT, Index);
       if (auto *Gather = dyn_cast<MaskedGatherSDNode>(GorS)) {
         SDValue Ops[] = { Chain, Gather->getPassThru(),
                           Mask, Base, Index, Scale } ;
@@ -42630,11 +42633,14 @@ static SDValue combineGatherScatter(SDNode *N, SelectionDAG &DAG,
                                   Ops, Scatter->getMemOperand(),
                                   Scatter->getIndexType());
     }
+  }
+
+  if (DCI.isBeforeLegalizeOps()) {
+    unsigned IndexWidth = Index.getScalarValueSizeInBits();
 
     // Make sure the index is either i32 or i64
-    unsigned ScalarSize = Index.getScalarValueSizeInBits();
-    if (ScalarSize != 32 && ScalarSize != 64) {
-      MVT EltVT = ScalarSize > 32 ? MVT::i64 : MVT::i32;
+    if (IndexWidth != 32 && IndexWidth != 64) {
+      MVT EltVT = IndexWidth > 32 ? MVT::i64 : MVT::i32;
       EVT IndexVT = EVT::getVectorVT(*DAG.getContext(), EltVT,
                                    Index.getValueType().getVectorNumElements());
       Index = DAG.getSExtOrTrunc(Index, DL, IndexVT);
@@ -42654,30 +42660,6 @@ static SDValue combineGatherScatter(SDNode *N, SelectionDAG &DAG,
                                   Ops, Scatter->getMemOperand(),
                                   Scatter->getIndexType());
     }
-
-    // Try to remove zero extends from 32->64 if we know the sign bit of
-    // the input is zero.
-    if (Index.getOpcode() == ISD::ZERO_EXTEND &&
-        Index.getScalarValueSizeInBits() == 64 &&
-        Index.getOperand(0).getScalarValueSizeInBits() == 32 &&
-        DAG.SignBitIsZero(Index.getOperand(0))) {
-      Index = Index.getOperand(0);
-      if (auto *Gather = dyn_cast<MaskedGatherSDNode>(GorS)) {
-        SDValue Ops[] = { Chain, Gather->getPassThru(),
-                          Mask, Base, Index, Scale } ;
-        return DAG.getMaskedGather(Gather->getVTList(),
-                                   Gather->getMemoryVT(), DL, Ops,
-                                   Gather->getMemOperand(),
-                                   Gather->getIndexType());
-      }
-      auto *Scatter = cast<MaskedScatterSDNode>(GorS);
-      SDValue Ops[] = { Chain, Scatter->getValue(),
-                        Mask, Base, Index, Scale };
-      return DAG.getMaskedScatter(Scatter->getVTList(),
-                                  Scatter->getMemoryVT(), DL,
-                                  Ops, Scatter->getMemOperand(),
-                                  Scatter->getIndexType());
-    }
   }
 
   // With vector masks we only demand the upper bit of the mask.