[AMDGPU] Optimize image sample followed by llvm.amdgcn.cvt.pkrtz into d16 variant #145203
+170
−0
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-llvm-transforms @llvm/pr-subscribers-backend-amdgpu Author: Harrison Hao (harrisonGPU) ChangesThis patch folds the pattern where the result of an Folding is only applied if:
Example matched pattern: %sample = call float @<!-- -->llvm.amdgcn.image.sample
%pack = call <2 x half> @<!-- -->llvm.amdgcn.cvt.pkrtz(float %sample, float %any)
%h0 = extractelement <2 x half> %pack, i64 0This gets folded into: %sample = call half @<!-- -->llvm.amdgcn.image.sampleFull diff: https://github.com/llvm/llvm-project/pull/145203.diff 2 Files Affected:
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp b/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp
index 5477c5eae9392..8a944c58a0ea7 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp
@@ -247,6 +247,42 @@ simplifyAMDGCNImageIntrinsic(const GCNSubtarget *ST,
ArgTys[0] = User->getType();
});
}
+
+ // Fold image.sample + cvt.pkrtz -> extractelement idx0 into a single
+ // d16 image sample.
+ // Pattern to match:
+ // %sample = call float @llvm.amdgcn.image.sample...
+ // %pack = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float %sample,
+ // float %any)
+ // %low = extractelement <2 x half> %pack, i64 0
+ // Replacement:
+ // call half @llvm.amdgcn.image.sample
+ //
+ // Folding criteria:
+ // 1. The only user of the image.sample intrinsic is amdgcn.cvt.pkrtz.
+ // 2. That cvt.pkrtz call has exactly one use.
+ // 3. Its sole user is an extractelement instruction with index zero.
+ // Otherwise, folding is not performed, because D16 sampling only
+ // guarantees that the element at index 0 is defined; index 1 is
+ // undefined and using it will result in poison.
+ if (auto *CvtPkrtzCall = dyn_cast<CallInst>(User)) {
+ if (CvtPkrtzCall->getIntrinsicID() == Intrinsic::amdgcn_cvt_pkrtz &&
+ CvtPkrtzCall->hasOneUse()) {
+ // Unique use must be extractelement idx == 0
+ if (auto *Ext =
+ dyn_cast<ExtractElementInst>(*CvtPkrtzCall->user_begin())) {
+ if (isa<ConstantInt>(Ext->getIndexOperand()) &&
+ cast<ConstantInt>(Ext->getIndexOperand())->isZero()) {
+
+ return modifyIntrinsicCall(
+ II, *CvtPkrtzCall, ImageDimIntr->Intr, IC,
+ [&](auto &Args, auto &ArgTys) {
+ ArgTys[0] = CvtPkrtzCall->getType();
+ });
+ }
+ }
+ }
+ }
}
// Only perform D16 folding if every user of the image sample is
diff --git a/llvm/test/Transforms/InstCombine/AMDGPU/image-d16.ll b/llvm/test/Transforms/InstCombine/AMDGPU/image-d16.ll
index ee5ccf5af987d..f4f74a84bcb8b 100644
--- a/llvm/test/Transforms/InstCombine/AMDGPU/image-d16.ll
+++ b/llvm/test/Transforms/InstCombine/AMDGPU/image-d16.ll
@@ -239,6 +239,140 @@ main_body:
ret bfloat %res
}
+define amdgpu_ps float @image_sample_2d_single_pkrtz_to_d16(<8 x i32> %surf_desc, <4 x i32> %samp, float %u, float %v) {
+; GFX7-LABEL: @image_sample_2d_single_pkrtz_to_d16(
+; GFX7-NEXT: main_body:
+; GFX7-NEXT: [[SAMPLE:%.*]] = call float @llvm.amdgcn.image.sample.lz.2d.f32.f32.v8i32.v4i32(i32 1, float [[U:%.*]], float [[V:%.*]], <8 x i32> [[SURF_DESC:%.*]], <4 x i32> [[SAMP:%.*]], i1 false, i32 0, i32 0)
+; GFX7-NEXT: [[PACK:%.*]] = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float [[SAMPLE]], float 0.000000e+00)
+; GFX7-NEXT: [[H0:%.*]] = extractelement <2 x half> [[PACK]], i64 0
+; GFX7-NEXT: [[MUL:%.*]] = fmul reassoc arcp contract afn half [[H0]], [[H0]]
+; GFX7-NEXT: [[DIV:%.*]] = fdiv reassoc arcp contract afn half [[MUL]], [[H0]]
+; GFX7-NEXT: [[ADD:%.*]] = fadd reassoc arcp contract afn half [[DIV]], [[H0]]
+; GFX7-NEXT: [[RES:%.*]] = fpext half [[ADD]] to float
+; GFX7-NEXT: ret float [[RES]]
+;
+; GFX81PLUS-LABEL: @image_sample_2d_single_pkrtz_to_d16(
+; GFX81PLUS-NEXT: main_body:
+; GFX81PLUS-NEXT: [[SAMPLE:%.*]] = call half @llvm.amdgcn.image.sample.lz.2d.f16.f32.v8i32.v4i32(i32 1, float [[U:%.*]], float [[V:%.*]], <8 x i32> [[SURF_DESC:%.*]], <4 x i32> [[SAMP:%.*]], i1 false, i32 0, i32 0)
+; GFX81PLUS-NEXT: [[MUL:%.*]] = fmul reassoc arcp contract afn half [[SAMPLE]], [[SAMPLE]]
+; GFX81PLUS-NEXT: [[DIV:%.*]] = fdiv reassoc arcp contract afn half [[MUL]], [[SAMPLE]]
+; GFX81PLUS-NEXT: [[ADD:%.*]] = fadd reassoc arcp contract afn half [[DIV]], [[SAMPLE]]
+; GFX81PLUS-NEXT: [[RES:%.*]] = fpext half [[ADD]] to float
+; GFX81PLUS-NEXT: ret float [[RES]]
+;
+main_body:
+ %sample = call float @llvm.amdgcn.image.sample.lz.2d.f32.f32.v8i32.v4i32(i32 1, float %u, float %v, <8 x i32> %surf_desc, <4 x i32> %samp, i1 false, i32 0, i32 0)
+ %pack = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float %sample, float 0.000000e+00)
+ %h0 = extractelement <2 x half> %pack, i64 0
+ %mul = fmul reassoc arcp contract afn half %h0, %h0
+ %div = fdiv reassoc arcp contract afn half %mul, %h0
+ %add = fadd reassoc arcp contract afn half %div, %h0
+ %res = fpext half %add to float
+ ret float %res
+}
+
+define amdgpu_ps float @image_sample_2d_pkrtz_variable_no_d16(<8 x i32> %surf_desc, <4 x i32> %samp, float %u, float %v) {
+; GFX7-LABEL: @image_sample_2d_pkrtz_variable_no_d16(
+; GFX7-NEXT: main_body:
+; GFX7-NEXT: [[SAMPLE:%.*]] = call float @llvm.amdgcn.image.sample.lz.2d.f32.f32.v8i32.v4i32(i32 1, float [[U:%.*]], float [[V:%.*]], <8 x i32> [[SURF_DESC:%.*]], <4 x i32> [[SAMP:%.*]], i1 false, i32 0, i32 0)
+; GFX7-NEXT: [[PACK:%.*]] = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float [[SAMPLE]], float [[V]])
+; GFX7-NEXT: [[H0:%.*]] = extractelement <2 x half> [[PACK]], i64 0
+; GFX7-NEXT: [[H1:%.*]] = extractelement <2 x half> [[PACK]], i64 1
+; GFX7-NEXT: [[MUL:%.*]] = fmul half [[H0]], [[H1]]
+; GFX7-NEXT: [[ADD:%.*]] = fadd half [[MUL]], [[H0]]
+; GFX7-NEXT: [[RES:%.*]] = fpext half [[ADD]] to float
+; GFX7-NEXT: ret float [[RES]]
+;
+; GFX81PLUS-LABEL: @image_sample_2d_pkrtz_variable_no_d16(
+; GFX81PLUS-NEXT: main_body:
+; GFX81PLUS-NEXT: [[SAMPLE:%.*]] = call float @llvm.amdgcn.image.sample.lz.2d.f32.f32.v8i32.v4i32(i32 1, float [[U:%.*]], float [[V:%.*]], <8 x i32> [[SURF_DESC:%.*]], <4 x i32> [[SAMP:%.*]], i1 false, i32 0, i32 0)
+; GFX81PLUS-NEXT: [[PACK:%.*]] = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float [[SAMPLE]], float [[V]])
+; GFX81PLUS-NEXT: [[H0:%.*]] = extractelement <2 x half> [[PACK]], i64 0
+; GFX81PLUS-NEXT: [[H1:%.*]] = extractelement <2 x half> [[PACK]], i64 1
+; GFX81PLUS-NEXT: [[MUL:%.*]] = fmul half [[H0]], [[H1]]
+; GFX81PLUS-NEXT: [[ADD:%.*]] = fadd half [[MUL]], [[H0]]
+; GFX81PLUS-NEXT: [[RES:%.*]] = fpext half [[ADD]] to float
+; GFX81PLUS-NEXT: ret float [[RES]]
+;
+main_body:
+ %sample = call float @llvm.amdgcn.image.sample.lz.2d.f32.f32.v8i32.v4i32(i32 1, float %u, float %v, <8 x i32> %surf_desc, <4 x i32> %samp, i1 false, i32 0, i32 0)
+ %pack = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float %sample, float %v)
+ %h0 = extractelement <2 x half> %pack, i64 0
+ %h1 = extractelement <2 x half> %pack, i64 1
+ %mul = fmul half %h0, %h1
+ %add = fadd half %mul, %h0
+ %res = fpext half %add to float
+ ret float %res
+}
+
+define amdgpu_ps float @image_sample_2d_pkrtz_constant_no_fold(<8 x i32> %surf_desc, <4 x i32> %samp, float %u, float %v) {
+; GFX7-LABEL: @image_sample_2d_pkrtz_constant_no_fold(
+; GFX7-NEXT: main_body:
+; GFX7-NEXT: [[SAMPLE:%.*]] = call float @llvm.amdgcn.image.sample.lz.2d.f32.f32.v8i32.v4i32(i32 1, float [[U:%.*]], float [[V:%.*]], <8 x i32> [[SURF_DESC:%.*]], <4 x i32> [[SAMP:%.*]], i1 false, i32 0, i32 0)
+; GFX7-NEXT: [[PACK:%.*]] = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float [[SAMPLE]], float 0.000000e+00)
+; GFX7-NEXT: [[H0:%.*]] = extractelement <2 x half> [[PACK]], i64 0
+; GFX7-NEXT: [[H1:%.*]] = extractelement <2 x half> [[PACK]], i64 1
+; GFX7-NEXT: [[MUL:%.*]] = fmul half [[H0]], [[H1]]
+; GFX7-NEXT: [[ADD:%.*]] = fadd half [[MUL]], [[H0]]
+; GFX7-NEXT: [[RES:%.*]] = fpext half [[ADD]] to float
+; GFX7-NEXT: ret float [[RES]]
+;
+; GFX81PLUS-LABEL: @image_sample_2d_pkrtz_constant_no_fold(
+; GFX81PLUS-NEXT: main_body:
+; GFX81PLUS-NEXT: [[SAMPLE:%.*]] = call float @llvm.amdgcn.image.sample.lz.2d.f32.f32.v8i32.v4i32(i32 1, float [[U:%.*]], float [[V:%.*]], <8 x i32> [[SURF_DESC:%.*]], <4 x i32> [[SAMP:%.*]], i1 false, i32 0, i32 0)
+; GFX81PLUS-NEXT: [[PACK:%.*]] = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float [[SAMPLE]], float 0.000000e+00)
+; GFX81PLUS-NEXT: [[H0:%.*]] = extractelement <2 x half> [[PACK]], i64 0
+; GFX81PLUS-NEXT: [[H1:%.*]] = extractelement <2 x half> [[PACK]], i64 1
+; GFX81PLUS-NEXT: [[MUL:%.*]] = fmul half [[H0]], [[H1]]
+; GFX81PLUS-NEXT: [[ADD:%.*]] = fadd half [[MUL]], [[H0]]
+; GFX81PLUS-NEXT: [[RES:%.*]] = fpext half [[ADD]] to float
+; GFX81PLUS-NEXT: ret float [[RES]]
+;
+main_body:
+ %sample = call float @llvm.amdgcn.image.sample.lz.2d.f32.f32.v8i32.v4i32(i32 1, float %u, float %v, <8 x i32> %surf_desc, <4 x i32> %samp, i1 false, i32 0, i32 0)
+ %pack = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float %sample, float 0.000000e+00)
+ %h0 = extractelement <2 x half> %pack, i64 0
+ %h1 = extractelement <2 x half> %pack, i64 1
+ %mul = fmul half %h0, %h1
+ %add = fadd half %mul, %h0
+ %res = fpext half %add to float
+ ret float %res
+}
+
+define amdgpu_ps float @image_sample_2d_single_pkrtz_high_no_d16(<8 x i32> %surf_desc, <4 x i32> %samp, float %u, float %v) {
+; GFX7-LABEL: @image_sample_2d_single_pkrtz_high_no_d16(
+; GFX7-NEXT: main_body:
+; GFX7-NEXT: [[SAMPLE:%.*]] = call float @llvm.amdgcn.image.sample.lz.2d.f32.f32.v8i32.v4i32(i32 1, float [[U:%.*]], float [[V:%.*]], <8 x i32> [[SURF_DESC:%.*]], <4 x i32> [[SAMP:%.*]], i1 false, i32 0, i32 0)
+; GFX7-NEXT: [[PACK:%.*]] = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float 0.000000e+00, float [[SAMPLE]])
+; GFX7-NEXT: [[H0:%.*]] = extractelement <2 x half> [[PACK]], i64 1
+; GFX7-NEXT: [[MUL:%.*]] = fmul reassoc arcp contract afn half [[H0]], [[H0]]
+; GFX7-NEXT: [[DIV:%.*]] = fdiv reassoc arcp contract afn half [[MUL]], [[H0]]
+; GFX7-NEXT: [[ADD:%.*]] = fadd reassoc arcp contract afn half [[DIV]], [[H0]]
+; GFX7-NEXT: [[RES:%.*]] = fpext half [[ADD]] to float
+; GFX7-NEXT: ret float [[RES]]
+;
+; GFX81PLUS-LABEL: @image_sample_2d_single_pkrtz_high_no_d16(
+; GFX81PLUS-NEXT: main_body:
+; GFX81PLUS-NEXT: [[SAMPLE:%.*]] = call float @llvm.amdgcn.image.sample.lz.2d.f32.f32.v8i32.v4i32(i32 1, float [[U:%.*]], float [[V:%.*]], <8 x i32> [[SURF_DESC:%.*]], <4 x i32> [[SAMP:%.*]], i1 false, i32 0, i32 0)
+; GFX81PLUS-NEXT: [[PACK:%.*]] = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float 0.000000e+00, float [[SAMPLE]])
+; GFX81PLUS-NEXT: [[H0:%.*]] = extractelement <2 x half> [[PACK]], i64 1
+; GFX81PLUS-NEXT: [[MUL:%.*]] = fmul reassoc arcp contract afn half [[H0]], [[H0]]
+; GFX81PLUS-NEXT: [[DIV:%.*]] = fdiv reassoc arcp contract afn half [[MUL]], [[H0]]
+; GFX81PLUS-NEXT: [[ADD:%.*]] = fadd reassoc arcp contract afn half [[DIV]], [[H0]]
+; GFX81PLUS-NEXT: [[RES:%.*]] = fpext half [[ADD]] to float
+; GFX81PLUS-NEXT: ret float [[RES]]
+;
+main_body:
+ %sample = call float @llvm.amdgcn.image.sample.lz.2d.f32.f32.v8i32.v4i32(i32 1, float %u, float %v, <8 x i32> %surf_desc, <4 x i32> %samp, i1 false, i32 0, i32 0)
+ %pack = call <2 x half> @llvm.amdgcn.cvt.pkrtz(float 0.000000e+00, float %sample)
+ %h0 = extractelement <2 x half> %pack, i64 1
+ %mul = fmul reassoc arcp contract afn half %h0, %h0
+ %div = fdiv reassoc arcp contract afn half %mul, %h0
+ %add = fadd reassoc arcp contract afn half %div, %h0
+ %res = fpext half %add to float
+ ret float %res
+}
+
define amdgpu_ps half @image_gather4_2d_v4f32(<8 x i32> inreg %rsrc, <4 x i32> inreg %samp, half %s, half %t) {
; GFX7-LABEL: @image_gather4_2d_v4f32(
; GFX7-NEXT: main_body:
|
DadSchoorse
reviewed
Jun 22, 2025
Comment on lines
+251
to
+252
| // Fold image.sample + cvt.pkrtz -> extractelement idx0 into a single | ||
| // d16 image sample. |
There was a problem hiding this comment.
We had a similar optimization in mesa and it broke tests because D16 does not round the same as a 32bit load followed by v_cvt_pk_f16_f32 for fixed point formats. D16 directly rounds towards nearest even in fp16, while the unoptimized pattern will first round towards nearest even in fp32, and then rounds towards zero in fp16.
There was a problem hiding this comment.
Thanks! Could you share some test cases for me to verify this issue? So far, I haven’t encountered the problem you mentioned. Since image sampling with D16 only uses the lower 16 bits, and I ensure that the first argument to cvt.pkrtz is the result of the image sample, I believe this folding is valid. I’ve tested a few cases and haven’t observed any issues so far, but I’ll continue testing more to be safe.
Additionally, we've used this optimization in other compilers (not LLVM) for years without encountering any problems.
There was a problem hiding this comment.
The mesa change was in https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/28730, but sadly the CI pipeline data was already deleted, so finding the exact test case that failed back then isn't trivial. But it resulted in disabling D16 optimizations if a fp16 rounding mode is required.
There was a problem hiding this comment.
Thanks, I will test it again. :-)
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This patch folds the pattern where the result of an
image sampleis passed tollvm.amdgcn.cvt.pkrtz, and the packed result is accessed withextractelement <2 x half>, index 0. In this case, we can emit the d16 variantof the image intrinsic directly, improving code clarity, avoiding unnecessary
v_cvt_pk_rtz_f16_f32_e64instructions, and reducing VGPR usage.Folding is only applied if:
image sampleresult has exactly one use: a call toamdgcn.cvt.pkrtzpkrtzhas one use: an extractelement withindex 0Example matched pattern:
This gets folded into: