[AMDGPU] Accept extractelement of a widening cast when folding image ops to a16#208207
[AMDGPU] Accept extractelement of a widening cast when folding image ops to a16#208207barbara-amd wants to merge 1 commit into
Conversation
…ops to a16 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.
|
@llvm/pr-subscribers-backend-amdgpu Author: Barbara Mitic (barbara-amd) ChangescanSafelyConvertTo16Bit() 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:
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
; --------------------------------------------------------------------
|
|
@llvm/pr-subscribers-llvm-transforms Author: Barbara Mitic (barbara-amd) ChangescanSafelyConvertTo16Bit() 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:
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
; --------------------------------------------------------------------
|
|
@Flakebi, could you review this since we already discussed the issue, or add more reviewers? |
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.