[llvm] [X86][SelectionDAG] Handle the case for gather where index is SHL (PR #139703)

[Rohit Aggarwal via llvm-commits]

https://github.com/rohitaggarwal007 created https://github.com/llvm/llvm-project/pull/139703

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.

>From 2b79bbe52acef964b57a47c8aa1a52df0431874b Mon Sep 17 00:00:00 2001
From: Rohit Aggarwal <Rohit.Aggarwal at amd.com>
Date: Tue, 13 May 2025 16:04:10 +0530
Subject: [PATCH] Handle the case for gather where index is SHL

---
 llvm/lib/Target/X86/X86ISelLowering.cpp | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

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