[llvm] [AMDGPU] Accept extractelement of a widening cast when folding image ops to a16 (PR #208207)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Barbara Mitic (barbara-amd)

<details>
<summary>Changes</summary>

canSafelyConvertTo16Bit() recognizes a scalar coordinate that is a direct sext/zext/fpext from a 16-bit value (sext gated on AllowI16SExt). Per-dimension coordinates can instead arrive as an extractelement of a widening vector cast (extractelement((s|z|fp)ext <N x i16/half> Vec), Idx). When the cast has more than one use, the extractelement(cast) -> cast(extractelement) canonicalization does not fire, so the cast is left in place and the coordinate is not recognized.

Strip a leading extractelement before the cast check so the same logic handles scalar and per-lane coordinates, and mirror this in convertTo16Bit() by re-extracting from the narrow vector, allowing the widening cast to be removed once it has no other uses.

---
Full diff: https://github.com/llvm/llvm-project/pull/208207.diff


2 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp (+22-4) 
- (modified) llvm/test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll (+49) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp b/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp
index f5d8e43dce9db..ee3fa7da86627 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp
@@ -96,13 +96,23 @@ static bool canSafelyConvertTo16Bit(Value &V, bool IsFloat,
     }
   }
 
+  // Coordinates may arrive as extractelement((s|z|fp)ext Vec), Idx. The
+  // widening cast has one use per lane, so it is never sunk into the extract;
+  // strip the extract here so the cast check below is common to scalar and
+  // vector coords.
+  Value *CastCandidate;
+  if (!match(&V, m_ExtractElt(PatternMatch::m_Value(CastCandidate),
+                              PatternMatch::m_Value())))
+    CastCandidate = &V;
+
   Value *CastSrc;
-  bool IsExt = IsFloat ? match(&V, m_FPExt(PatternMatch::m_Value(CastSrc)))
-                       : match(&V, m_ZExt(PatternMatch::m_Value(CastSrc)));
+  bool IsExt =
+      IsFloat ? match(CastCandidate, m_FPExt(PatternMatch::m_Value(CastSrc)))
+              : match(CastCandidate, m_ZExt(PatternMatch::m_Value(CastSrc)));
   if (!IsExt && !IsFloat && AllowI16SExt)
-    IsExt = match(&V, m_SExt(PatternMatch::m_Value(CastSrc)));
+    IsExt = match(CastCandidate, m_SExt(PatternMatch::m_Value(CastSrc)));
   if (IsExt) {
-    Type *CastSrcTy = CastSrc->getType();
+    Type *CastSrcTy = CastSrc->getType()->getScalarType();
     if (CastSrcTy->isHalfTy() || CastSrcTy->isIntegerTy(16))
       return true;
   }
@@ -115,6 +125,14 @@ static Value *convertTo16Bit(Value &V, InstCombiner::BuilderTy &Builder) {
   Type *VTy = V.getType();
   if (isa<FPExtInst, SExtInst, ZExtInst>(&V))
     return cast<Instruction>(&V)->getOperand(0);
+  // Vector form: extractelement((s|z|fp)ext Vec), Idx -> extractelement(Vec,
+  // Idx), taking the narrow lane directly so the widening cast can be removed.
+  Value *VecCast, *Idx;
+  if (match(&V, m_ExtractElt(PatternMatch::m_Value(VecCast),
+                             PatternMatch::m_Value(Idx))) &&
+      isa<FPExtInst, SExtInst, ZExtInst>(VecCast))
+    return Builder.CreateExtractElement(
+        cast<Instruction>(VecCast)->getOperand(0), Idx);
   if (VTy->isIntegerTy())
     return Builder.CreateIntCast(&V, Type::getInt16Ty(V.getContext()), false);
   if (VTy->isFloatingPointTy())
diff --git a/llvm/test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll b/llvm/test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll
index 0bedef0d93a59..849866b4a117d 100644
--- a/llvm/test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll
+++ b/llvm/test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll
@@ -4018,6 +4018,55 @@ define amdgpu_kernel void @image_load_a16_mip_2d_const_noopt(ptr addrspace(1) %o
   ret void
 }
 
+define amdgpu_kernel void @image_load_a16_mip_2d_sext_vec(ptr addrspace(1) %out, <8 x i32> inreg %rsrc, <2 x i16> %coord) {
+; CHECK-LABEL: @image_load_a16_mip_2d_sext_vec(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x i16> [[COORD:%.*]], i64 0
+; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <2 x i16> [[COORD]], i64 1
+; CHECK-NEXT:    [[RES:%.*]] = call <4 x float> @llvm.amdgcn.image.load.2d.v4f32.i16.v8i32(i32 15, i16 [[TMP1]], i16 [[TMP2]], <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
+;
+  %coord32 = sext <2 x i16> %coord to <2 x i32>
+  %s32 = extractelement <2 x i32> %coord32, i64 0
+  %t32 = extractelement <2 x i32> %coord32, i64 1
+  %res = call <4 x float> @llvm.amdgcn.image.load.mip.2d.v4f32.i32.v8i32(i32 15, i32 %s32, i32 %t32, i32 0, <8 x i32> %rsrc, i32 0, i32 0)
+  store <4 x float> %res, ptr addrspace(1) %out
+  ret void
+}
+
+define amdgpu_kernel void @image_load_a16_mip_2d_zext_vec(ptr addrspace(1) %out, <8 x i32> inreg %rsrc, <2 x i16> %coord) {
+; CHECK-LABEL: @image_load_a16_mip_2d_zext_vec(
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x i16> [[COORD:%.*]], i64 0
+; CHECK-NEXT:    [[TMP2:%.*]] = extractelement <2 x i16> [[COORD]], i64 1
+; CHECK-NEXT:    [[RES:%.*]] = call <4 x float> @llvm.amdgcn.image.load.2d.v4f32.i16.v8i32(i32 15, i16 [[TMP1]], i16 [[TMP2]], <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
+;
+  %coord32 = zext <2 x i16> %coord to <2 x i32>
+  %s32 = extractelement <2 x i32> %coord32, i64 0
+  %t32 = extractelement <2 x i32> %coord32, i64 1
+  %res = call <4 x float> @llvm.amdgcn.image.load.mip.2d.v4f32.i32.v8i32(i32 15, i32 %s32, i32 %t32, i32 0, <8 x i32> %rsrc, i32 0, i32 0)
+  store <4 x float> %res, ptr addrspace(1) %out
+  ret void
+}
+
+define amdgpu_kernel void @image_load_a16_mip_2d_zext_vec_i8_noopt(ptr addrspace(1) %out, <8 x i32> inreg %rsrc, <2 x i8> %coord) {
+; CHECK-LABEL: @image_load_a16_mip_2d_zext_vec_i8_noopt(
+; CHECK-NEXT:    [[COORD32:%.*]] = zext <2 x i8> [[COORD:%.*]] to <2 x i32>
+; CHECK-NEXT:    [[S32:%.*]] = extractelement <2 x i32> [[COORD32]], i64 0
+; CHECK-NEXT:    [[T32:%.*]] = extractelement <2 x i32> [[COORD32]], i64 1
+; CHECK-NEXT:    [[RES:%.*]] = call <4 x float> @llvm.amdgcn.image.load.2d.v4f32.i32.v8i32(i32 15, i32 [[S32]], i32 [[T32]], <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
+;
+  %coord32 = zext <2 x i8> %coord to <2 x i32>
+  %s32 = extractelement <2 x i32> %coord32, i64 0
+  %t32 = extractelement <2 x i32> %coord32, i64 1
+  %res = call <4 x float> @llvm.amdgcn.image.load.mip.2d.v4f32.i32.v8i32(i32 15, i32 %s32, i32 %t32, i32 0, <8 x i32> %rsrc, i32 0, i32 0)
+  store <4 x float> %res, ptr addrspace(1) %out
+  ret void
+}
+
 ; --------------------------------------------------------------------
 ; llvm.amdgcn.image.sample g16
 ; --------------------------------------------------------------------

``````````

</details>


https://github.com/llvm/llvm-project/pull/208207