[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 78406b7d0bbd40866d2dd9bcc43cf6c7fa68840d 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      | 20 ++++++
 .../AArch64/GlobalISel/shift-amount-mod.ll    | 63 +++++++++++++++++++
 2 files changed, 83 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..a9726e14f12cd 100644
--- a/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
@@ -3158,6 +3158,26 @@ 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, or masked with an AND that covers enough bits, we
+      // can remove the redundant operation.
+      if (!SrcTy.isVector()) {
+        int Bits = SrcTy.getSizeInBits() == 32 ? 5 : 6;
+        Register NewShiftReg = I.getOperand(2).getReg();
+
+        // Skip over G_ZEXT of the shift amount.
+        MachineInstr *ShiftAmtDef = MRI.getVRegDef(NewShiftReg);
+        if (ShiftAmtDef && ShiftAmtDef->getOpcode() == TargetOpcode::G_ZEXT &&
+            MRI.hasOneNonDBGUse(NewShiftReg)) {
+          NewShiftReg = ShiftAmtDef->getOperand(1).getReg();
+          ShiftAmtDef = MRI.getVRegDef(NewShiftReg);
+        }
+        if (NewShiftReg != I.getOperand(2).getReg())
+          I.getOperand(2).setReg(NewShiftReg);
+      }
     }
 
     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
+}