[X86][SelectionDAG] Handle the case for gather where index is SHL

[rohitaggarwal007]

Fix the Gather's Index for SHL Opcode in which shift amount is 4 or greater.

It is in the continuity of #137813

@RKSimon Please review this.

[llvmbot]

@llvm/pr-subscribers-backend-x86

Author: Rohit Aggarwal (rohitaggarwal007)

Changes

Fix the Gather's Index for SHL Opcode in which shift amount is 4 or greater.

It is in the continuity of #137813

@RKSimon Please review this.

---

Full diff: https://github.com/llvm/llvm-project/pull/139703.diff

1 Files Affected:

• (modified) llvm/lib/Target/X86/X86ISelLowering.cpp (+13-2)

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index ac4fb157a6026..4bb23ced2bc42 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -56710,12 +56710,23 @@ static SDValue combineGatherScatter(SDNode *N, SelectionDAG &DAG,
 
   if (DCI.isBeforeLegalize()) {
     unsigned IndexWidth = Index.getScalarValueSizeInBits();
-
+    // If the index is a left shift, \ComputeNumSignBits we are recomputing
+    // the number of sign bits from the shifted value. We are trying to enable
+    // the optimization in which we can shrink indices if they are larger than
+    // 32-bits. Using the existing fold techniques implemented below.
+    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;
+        }
+      }
+    }
     // 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)) {
+    if (IndexWidth > 32 && ComputeNumSignBits > (IndexWidth - 32)) {
       EVT NewVT = IndexVT.changeVectorElementType(MVT::i32);
 
       // FIXME: We could support more than just constant fold, but we need to

[RKSimon]

@rohitaggarwal007 please can you rebase against trunk?

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[rohitaggarwal007]

@rohitaggarwal007 please can you rebase against trunk?

Done

[RKSimon]

This still seems wrong - afaict you are doing this:

define i64 @src(i32 noundef %x) {
#0:
  %and = and i32 noundef %x, 536870911
  %zext = zext i32 %and to i64
  %hi = shl i64 %zext, 4
  ret i64 %hi
}
=>
define i64 @tgt(i32 noundef %x) {
#0:
  %shl = shl i32 noundef %x, 4
  %ext = sext i32 %shl to i64
  ret i64 %ext
}
Transformation doesn't verify!

[rohitaggarwal007]

So the problem is that we can overflow the value on left shift, right?