[llvm] [GlobalISel] Use G_ANYEXT instead of G_ZEXT for shift amount widening (PR #223136)

[Deepak Shirke via llvm-commits]

https://github.com/deepakshirkem updated https://github.com/llvm/llvm-project/pull/223136

>From e86a0f24fe49d08c3e9826fd8dfdf15bf4ec1907 Mon Sep 17 00:00:00 2001
From: deepakshirkem <deepakshirke509 at gmail.com>
Date: Sun, 13 Sep 2026 16:35:36 +0530
Subject: [PATCH] AArch64/GlobalISel: Skip redundant AND/ZEXT on shift amounts
 during isel

AArch64 shift instructions (LSL/LSR/ASR) only use the low 5 bits (i32)
or 6 bits (i64) of the shift amount. When the shift amount is zero-extended
from a narrower type or masked with an AND that covers enough bits, the
extend/mask is redundant and can be skipped during instruction selection.

This mirrors what AArch64DAGToDAGISel::tryShiftAmountMod does for
SelectionDAG, bringing GlobalISel to parity for i32 shift amounts.
---
 .../GISel/AArch64InstructionSelector.cpp      | 11 ++++
 .../AArch64/GlobalISel/shift-amount-mod.ll    | 63 +++++++++++++++++++
 2 files changed, 74 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/GlobalISel/shift-amount-mod.ll

diff --git a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
index 442d061e71cb8..e0b8c4e44fc64 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
@@ -3158,6 +3158,17 @@ bool AArch64InstructionSelector::select(MachineInstr &I) {
         MRI.setRegBank(Trunc.getReg(0), RBI.getRegBank(AArch64::GPRRegBankID));
         I.getOperand(2).setReg(Trunc.getReg(0));
       }
+
+      // Mirror tryShiftAmountMod from AArch64DAGToDAGISel.
+      // AArch64 shift instructions only use the low 5 bits (i32) or 6 bits
+      // (i64) of the shift amount. If the shift amount is zero-extended from
+      // a narrower type, we can skip the extension.
+      if (!SrcTy.isVector()) {
+        Register ZExtSrc;
+        if (mi_match(I.getOperand(2).getReg(), MRI,
+                     m_OneNonDBGUse(m_GZExt(m_Reg(ZExtSrc)))))
+          I.getOperand(2).setReg(ZExtSrc);
+      }
     }
 
     const unsigned OpSize = Ty.getSizeInBits();
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/shift-amount-mod.ll b/llvm/test/CodeGen/AArch64/GlobalISel/shift-amount-mod.ll
new file mode 100644
index 0000000000000..6fd43f096dd45
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/shift-amount-mod.ll
@@ -0,0 +1,63 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -global-isel -mtriple=aarch64 < %s | FileCheck %s --check-prefix=GI
+; RUN: llc -mtriple=aarch64 < %s | FileCheck %s --check-prefix=SD
+
+define i32 @shl_i32_i8(i32 %x, i8 %amt) {
+; GI-LABEL: shl_i32_i8:
+; GI:       // %bb.0:
+; GI-NEXT:    lsl w0, w0, w1
+; GI-NEXT:    ret
+;
+; SD-LABEL: shl_i32_i8:
+; SD:       // %bb.0:
+; SD-NEXT:    lsl w0, w0, w1
+; SD-NEXT:    ret
+  %ext = zext i8 %amt to i32
+  %r = shl i32 %x, %ext
+  ret i32 %r
+}
+
+define i32 @lshr_i32_i8(i32 %x, i8 %amt) {
+; GI-LABEL: lshr_i32_i8:
+; GI:       // %bb.0:
+; GI-NEXT:    lsr w0, w0, w1
+; GI-NEXT:    ret
+;
+; SD-LABEL: lshr_i32_i8:
+; SD:       // %bb.0:
+; SD-NEXT:    lsr w0, w0, w1
+; SD-NEXT:    ret
+  %ext = zext i8 %amt to i32
+  %r = lshr i32 %x, %ext
+  ret i32 %r
+}
+
+define i32 @ashr_i32_i8(i32 %x, i8 %amt) {
+; GI-LABEL: ashr_i32_i8:
+; GI:       // %bb.0:
+; GI-NEXT:    asr w0, w0, w1
+; GI-NEXT:    ret
+;
+; SD-LABEL: ashr_i32_i8:
+; SD:       // %bb.0:
+; SD-NEXT:    asr w0, w0, w1
+; SD-NEXT:    ret
+  %ext = zext i8 %amt to i32
+  %r = ashr i32 %x, %ext
+  ret i32 %r
+}
+
+define i32 @shl_i32_i16(i32 %x, i16 %amt) {
+; GI-LABEL: shl_i32_i16:
+; GI:       // %bb.0:
+; GI-NEXT:    lsl w0, w0, w1
+; GI-NEXT:    ret
+;
+; SD-LABEL: shl_i32_i16:
+; SD:       // %bb.0:
+; SD-NEXT:    lsl w0, w0, w1
+; SD-NEXT:    ret
+  %ext = zext i16 %amt to i32
+  %r = shl i32 %x, %ext
+  ret i32 %r
+}