[llvm] [X86][SelectionDAG] Fix the Gather's base and index by modifying the Scale value (PR #134979)

[Rohit Aggarwal via llvm-commits]

================
@@ -56520,13 +56520,49 @@ static SDValue combineGatherScatter(SDNode *N, SelectionDAG &DAG,
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
 
   if (DCI.isBeforeLegalize()) {
+    // Attempt to move shifted index into the address scale, allows further
+    // index truncation below.
+    if (Index.getOpcode() == ISD::SHL && isa<ConstantSDNode>(Scale)) {
+      unsigned BitWidth = Index.getScalarValueSizeInBits();
+      unsigned MaskBits = BitWidth - Log2_32(Scale->getAsZExtVal());
+      APInt DemandedBits = APInt::getLowBitsSet(BitWidth, MaskBits);
+      if (TLI.SimplifyDemandedBits(Index, DemandedBits, DCI)) {
+        if (N->getOpcode() != ISD::DELETED_NODE)
+          DCI.AddToWorklist(N);
+        return SDValue(N, 0);
+      }
+      uint64_t ScaleAmt = Scale->getAsZExtVal();
+      if (auto MinShAmt = DAG.getValidMinimumShiftAmount(Index)) {
+        if (*MinShAmt >= 1 && (*MinShAmt + Log2_64(ScaleAmt)) < 4 &&
+            DAG.ComputeNumSignBits(Index.getOperand(0)) > 1) {
+          SDValue ShAmt = Index.getOperand(1);
+          SDValue NewShAmt =
+              DAG.getNode(ISD::SUB, DL, ShAmt.getValueType(), ShAmt,
+                          DAG.getConstant(1, DL, ShAmt.getValueType()));
+          SDValue NewIndex = DAG.getNode(ISD::SHL, DL, Index.getValueType(),
+                                         Index.getOperand(0), NewShAmt);
+          SDValue NewScale =
+              DAG.getConstant(ScaleAmt * 2, DL, Scale.getValueType());
+          return rebuildGatherScatter(GorS, NewIndex, Base, NewScale, DAG);
+        }
+      }
+    }
     unsigned IndexWidth = Index.getScalarValueSizeInBits();
 
     // Shrink 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.
-    if (IndexWidth > 32 && DAG.ComputeNumSignBits(Index) > (IndexWidth - 32)) {
+    unsigned ComputeNumSignBits = DAG.ComputeNumSignBits(Index);
+    if (Index.getOpcode() == ISD::SHL) {
+      if (auto MinShAmt = DAG.getValidMinimumShiftAmount(Index)) {
+        if (DAG.ComputeNumSignBits(Index.getOperand(0)) > 1) {
+          ComputeNumSignBits += *MinShAmt;
----------------
rohitaggarwal007 wrote:

As we are restricting shift for case where we can not completely remove the SHL as it has shift of 4 i.e.16 scale. Which is not supported scale value.
In our previous patch, we were allowing the folding for SHL and DAG.ComputeNumSignBits for Index was getting increasing from 31 to 32 after all the possible shift folding.
So taking a inspiration and understanding that number of shift contribute to calculate of the NumOfSignBits. So I recompute the value of ComputeNumSignBits using MinShAmt and try to enable the truncation folding.

I might not be fully correct  in my understanding.

https://github.com/llvm/llvm-project/pull/134979