[llvm] [AMDGPU] Accept sext addresses when folding image ops to a16 (PR #203189)

[Barbara Mitic via llvm-commits]

https://github.com/barbara-amd created https://github.com/llvm/llvm-project/pull/203189

canSafelyConvertTo16Bit() only accepts a zext when narrowing image address coordinates to 16 bits. Add an opt-in AllowI16SExt flag so a sext from i16 is accepted too, and enable it for sampler-less image instructions.
Coordinates of sampler-less loads/stores are unsigned, so sext and zext only disagree for a negative i16 (>= 0x8000), which is already out of bounds since the maximum image dimension is <= 0x8000. Accepting the sext therefore lets such coordinates fold to the a16 form, reducing VGPR pressure.
Together with the extractelement-of-widening-cast canonicalization in InstCombine (#201529), this also handles vector i16 coordinates that are extended and then split into per-dimension extracts.

>From 67d54c62b567092b1d21a081398ac0fb0b77afaf Mon Sep 17 00:00:00 2001
From: Barbara Mitic <Barbara.Mitic at amd.com>
Date: Wed, 3 Jun 2026 15:26:43 +0200
Subject: [PATCH] [AMDGPU] Accept sext addresses when folding image ops to a16

canSafelyConvertTo16Bit() only accepts a zext when narrowing image
address coordinates to 16 bits. Add an opt-in AllowI16SExt flag so a
sext from i16 is accepted too, and enable it for sampler-less image
instructions.
Coordinates of sampler-less loads/stores are unsigned, so sext and zext
only disagree for a negative i16 (>= 0x8000), which is already out of
bounds since the maximum image dimension is <= 0x8000. Accepting the
sext therefore lets such coordinates fold to the a16 form, reducing VGPR
pressure.
Together with the extractelement-of-widening-cast canonicalization in
InstCombine, this also handles vector i16 coordinates that are extended
and then split into per-dimension extracts.
---
 .../AMDGPU/AMDGPUInstCombineIntrinsic.cpp     | 19 ++++++++++++++-----
 .../InstCombine/AMDGPU/amdgcn-intrinsics.ll   |  7 +++----
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp b/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp
index b17ccef685b21..29bf18c4f11d0 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp
@@ -65,11 +65,13 @@ static APFloat fmed3AMDGCN(const APFloat &Src0, const APFloat &Src1,
   return maxnum(Src0, Src1);
 }
 
-// Check if a value can be converted to a 16-bit value without losing
-// precision.
+// Check if a value can be converted to a 16-bit value without losing precision.
 // The value is expected to be either a float (IsFloat = true) or an unsigned
-// integer (IsFloat = false).
-static bool canSafelyConvertTo16Bit(Value &V, bool IsFloat) {
+// integer (IsFloat = false). When AllowI16SExt is set, a sext from i16 is also
+// accepted: for unsigned addresses sext and zext only differ for a negative
+// i16, which is out of bounds anyway (see caller).
+static bool canSafelyConvertTo16Bit(Value &V, bool IsFloat,
+                                    bool AllowI16SExt = false) {
   Type *VTy = V.getType();
   if (VTy->isHalfTy() || VTy->isIntegerTy(16)) {
     // The value is already 16-bit, so we don't want to convert to 16-bit again!
@@ -97,6 +99,8 @@ static bool canSafelyConvertTo16Bit(Value &V, bool IsFloat) {
   Value *CastSrc;
   bool IsExt = IsFloat ? match(&V, m_FPExt(PatternMatch::m_Value(CastSrc)))
                        : match(&V, m_ZExt(PatternMatch::m_Value(CastSrc)));
+  if (!IsExt && !IsFloat && AllowI16SExt)
+    IsExt = match(&V, m_SExt(PatternMatch::m_Value(CastSrc)));
   if (IsExt) {
     Type *CastSrcTy = CastSrc->getType();
     if (CastSrcTy->isHalfTy() || CastSrcTy->isIntegerTy(16))
@@ -332,11 +336,16 @@ simplifyAMDGCNImageIntrinsic(const GCNSubtarget *ST,
   // true means derivatives can be converted to 16 bit, coordinates not
   bool OnlyDerivatives = false;
 
+  // Sampler-less addresses are unsigned, so a sext from i16 folds to a16 like a
+  // zext: they only disagree for a negative i16 (>= 0x8000), which is out of
+  // bounds while the max image dimension is <= 0x8000.
+  bool AllowI16SExt = !HasSampler;
+
   for (unsigned OperandIndex = ImageDimIntr->GradientStart;
        OperandIndex < ImageDimIntr->VAddrEnd; OperandIndex++) {
     Value *Coord = II.getOperand(OperandIndex);
     // If the values are not derived from 16-bit values, we cannot optimize.
-    if (!canSafelyConvertTo16Bit(*Coord, HasSampler)) {
+    if (!canSafelyConvertTo16Bit(*Coord, HasSampler, AllowI16SExt)) {
       if (OperandIndex < ImageDimIntr->CoordStart ||
           ImageDimIntr->GradientStart == ImageDimIntr->CoordStart) {
         return std::nullopt;
diff --git a/llvm/test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll b/llvm/test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll
index ed197ebde7dd8..0bedef0d93a59 100644
--- a/llvm/test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll
+++ b/llvm/test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll
@@ -3968,10 +3968,9 @@ define amdgpu_kernel void @image_load_a16_mip_1d(ptr addrspace(1) %out, <8 x i32
   ret void
 }
 
-define amdgpu_kernel void @image_load_a16_mip_1d_noopt(ptr addrspace(1) %out, <8 x i32> inreg %rsrc, i16 %s) {
-; CHECK-LABEL: @image_load_a16_mip_1d_noopt(
-; CHECK-NEXT:    [[S32:%.*]] = sext i16 [[S:%.*]] to i32
-; CHECK-NEXT:    [[RES:%.*]] = call <4 x float> @llvm.amdgcn.image.load.1d.v4f32.i32.v8i32(i32 15, i32 [[S32]], <8 x i32> [[RSRC:%.*]], i32 0, i32 0)
+define amdgpu_kernel void @image_load_a16_mip_1d_sext(ptr addrspace(1) %out, <8 x i32> inreg %rsrc, i16 %s) {
+; CHECK-LABEL: @image_load_a16_mip_1d_sext(
+; CHECK-NEXT:    [[RES:%.*]] = call <4 x float> @llvm.amdgcn.image.load.1d.v4f32.i16.v8i32(i32 15, i16 [[S:%.*]], <8 x i32> [[RSRC:%.*]], i32 0, i32 0)
 ; CHECK-NEXT:    store <4 x float> [[RES]], ptr addrspace(1) [[OUT:%.*]], align 16
 ; CHECK-NEXT:    ret void
 ;