[llvm] [X86] Widen zero-extended narrow lshr before isel (PR #207919)

[Jaeuk Lee via llvm-commits]

https://github.com/skku970412 updated https://github.com/llvm/llvm-project/pull/207919

>From 775cc41cf41dd3c5e649e4cd8f09fd2186d259ff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=EC=9D=B4=EC=9E=AC=EC=9A=B1?=
 <126692701+skku970412 at users.noreply.github.com>
Date: Tue, 7 Jul 2026 16:44:05 +0900
Subject: [PATCH 1/2] [X86] Widen zero-extended narrow lshr before isel

---
 llvm/lib/Target/X86/X86ISelLowering.cpp | 32 +++++++++++++++++++++++++
 llvm/test/CodeGen/X86/pr205898.ll       | 31 ++++++++++++++++++++++++
 2 files changed, 63 insertions(+)
 create mode 100644 llvm/test/CodeGen/X86/pr205898.ll

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 9209134a055c6..407f15688da30 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -57647,6 +57647,38 @@ static SDValue combineZext(SDNode *N, SelectionDAG &DAG,
                                                  DAG, DCI, Subtarget))
     return V;
 
+  // If an AssertZext-backed wide source is known to have zeros above the
+  // truncated type, prefer shifting the wide value directly. This avoids an
+  // 8/16-bit shift followed by a zero-extend when the result is used in a
+  // wider register, for example as an address index.
+  if (N->getOpcode() == ISD::ZERO_EXTEND &&
+      (VT == MVT::i32 || VT == MVT::i64) && N0.getOpcode() == ISD::SRL &&
+      N0.hasOneUse()) {
+    SDValue Trunc = N0.getOperand(0);
+    EVT NarrowVT = Trunc.getValueType();
+
+    if (Trunc.getOpcode() == ISD::TRUNCATE && Trunc.hasOneUse() &&
+        NarrowVT.isScalarInteger()) {
+      SDValue Wide = Trunc.getOperand(0);
+      EVT WideVT = Wide.getValueType();
+      unsigned NarrowBits = NarrowVT.getSizeInBits();
+      unsigned WideBits = WideVT.getSizeInBits();
+      std::optional<unsigned> ShiftAmt = DAG.getValidShiftAmount(N0);
+
+      if (WideVT == VT && Wide.getOpcode() == ISD::AssertZext && ShiftAmt &&
+          *ShiftAmt < NarrowBits) {
+        APInt HighBits =
+            APInt::getHighBitsSet(WideBits, WideBits - NarrowBits);
+        if (DAG.MaskedValueIsZero(Wide, HighBits)) {
+          SDValue WideShiftAmt =
+              DAG.getShiftAmountConstant(*ShiftAmt, VT, SDLoc(N0));
+          return DAG.getNode(ISD::SRL, SDLoc(N0), VT, Wide, WideShiftAmt,
+                             N0->getFlags());
+        }
+      }
+    }
+  }
+
   if (VT.isVector())
     if (SDValue R = PromoteMaskArithmetic(SDValue(N, 0), dl, DAG, Subtarget))
       return R;
diff --git a/llvm/test/CodeGen/X86/pr205898.ll b/llvm/test/CodeGen/X86/pr205898.ll
new file mode 100644
index 0000000000000..e51743cc30036
--- /dev/null
+++ b/llvm/test/CodeGen/X86/pr205898.ll
@@ -0,0 +1,31 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu | FileCheck %s
+
+define i64 @known_zero_trunc_lshr_index(i8 zeroext %b, ptr %p) {
+; CHECK-LABEL: known_zero_trunc_lshr_index:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    # kill: def $edi killed $edi def $rdi
+; CHECK-NEXT:    shrl $6, %edi
+; CHECK-NEXT:    movq (%rsi,%rdi,8), %rax
+; CHECK-NEXT:    retq
+  %shr = lshr i8 %b, 6
+  %idx = zext i8 %shr to i64
+  %gep = getelementptr i64, ptr %p, i64 %idx
+  %load = load i64, ptr %gep, align 8
+  ret i64 %load
+}
+
+define i64 @do_not_widen_trunc_lshr_without_known_high_zero(i32 %x, ptr %p) {
+; CHECK-LABEL: do_not_widen_trunc_lshr_without_known_high_zero:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    shrb $6, %dil
+; CHECK-NEXT:    movzbl %dil, %eax
+; CHECK-NEXT:    movq (%rsi,%rax,8), %rax
+; CHECK-NEXT:    retq
+  %tr = trunc i32 %x to i8
+  %shr = lshr i8 %tr, 6
+  %idx = zext i8 %shr to i64
+  %gep = getelementptr i64, ptr %p, i64 %idx
+  %load = load i64, ptr %gep, align 8
+  ret i64 %load
+}

>From 1a384b92a7487ebeae495dd721a14d960a2f2b53 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=EC=9D=B4=EC=9E=AC=EC=9A=B1?=
 <126692701+skku970412 at users.noreply.github.com>
Date: Tue, 7 Jul 2026 17:10:55 +0900
Subject: [PATCH 2/2] [X86] Format narrow shift combine

---
 llvm/lib/Target/X86/X86ISelLowering.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 407f15688da30..fe0d34b93e1f5 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -57667,8 +57667,7 @@ static SDValue combineZext(SDNode *N, SelectionDAG &DAG,
 
       if (WideVT == VT && Wide.getOpcode() == ISD::AssertZext && ShiftAmt &&
           *ShiftAmt < NarrowBits) {
-        APInt HighBits =
-            APInt::getHighBitsSet(WideBits, WideBits - NarrowBits);
+        APInt HighBits = APInt::getHighBitsSet(WideBits, WideBits - NarrowBits);
         if (DAG.MaskedValueIsZero(Wide, HighBits)) {
           SDValue WideShiftAmt =
               DAG.getShiftAmountConstant(*ShiftAmt, VT, SDLoc(N0));