Fix MLIR conv-pointwise-layout fusion splitting#5064
Conversation
There was a problem hiding this comment.
Pull request overview
This PR prevents gpu::compile_ops from blindly reusing cached MLIR tuning solutions (perfConfig strings) that no longer apply to a fused MLIR submodule, by validating cached solutions against the current fused module before reuse. This fits into the GPU backend’s tuning/benchmarking pipeline and improves robustness when problem-cache keys collide across different MLIR module structures.
Changes:
- Add an applicability check for cached
gpu::mlir_opsolutions usingis_module_fusible()before reusing a cached perfConfig. - Fall back to configured tuning candidates (and benchmarking when enabled) when a cached solution is null or inapplicable.
- Add a verify regression that seeds
MIGRAPHX_PROBLEM_CACHEwith a known-bad MLIR perfConfig for a conv+pointwise fused path.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| test/verify/test_conv_add_tune.cpp | Adds a regression that seeds the problem cache with a bad MLIR perfConfig and exercises a fused conv/add MLIR path. |
| src/targets/gpu/compile_ops.cpp | Validates cached MLIR solutions against the fused submodule before reusing them; falls back to candidate selection/benchmarking otherwise. |
| CHANGELOG.md | Documents the fix for cached gpu::mlir_op perfConfig reuse causing compilation failures. |
umangyadav
left a comment
umangyadav
left a comment
There was a problem hiding this comment.
Changes look good to me but let's wait on someone from MIGraphX team to review
pfultz2
left a comment
pfultz2
left a comment
There was a problem hiding this comment.
I dont think this is a good approach. I think it would be better to just add is_module_fusible to the problem config.
I'm not sure that I follow. How would we add |
That shouldn't fail, we split the module when |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## develop #5064 +/- ##
========================================
Coverage 92.89% 92.89%
========================================
Files 603 603
Lines 32448 32448
========================================
Hits 30140 30140
Misses 2308 2308 🚀 New features to boost your workflow:
|
Regressions detected 🔴 |
|
justinrosner
changed the title
@pfultz2 I believe I tracked down the root cause to the selected MLIR perfConfig rejecting the full conv + pointwise + layout fused module, and then the existing fallback only knew how to split into MLIR kernel + pointwise kernel. That two-way fallback assumes the pointwise portion can write directly into the original final output buffer. Once there is a layout-only tail, that assumption breaks, because the pointwise output and final output layout/shape are separate stages. I am still worried about overall performance when trying to use the same perfConfigs for problemConfigs that index back to both fused and non-fused kernels. I know this is something that the rocMLIR team has talked about in the past before, so maybe we can try to come up with a better solution for this now. @dhernandez0 @umangyadav @pabloantoniom |
| migraphx::make_op("convolution", {{"padding", {1, 1, 1, 0}}, {"stride", {2, 2}}}), | ||
| x2, | ||
| w2); | ||
| auto add = mm->add_instruction(migraphx::make_op("add"), pooling, conv2); |
There was a problem hiding this comment.
this one is doing conv + add + reshape + transpose, right? I guess if you do (same conv) conv + relu it'd still be the same problem config but it will not work with split-k. Does the code take that into account as well? It'd be nice to have tests as well
There was a problem hiding this comment.
It should. Added a new test for this.
| return *it; | ||
| } | ||
|
|
||
| static optional<instruction_ref> find_layout_tail_split(instruction_ref pointwise_ins) |
There was a problem hiding this comment.
I dont think you need to write another function for this. reshape_lazy just needs to be added to the list in find_final_split.
There was a problem hiding this comment.
I saw your comment down here: #5064 (comment). Is the current function approach that I have the proper way to go about solving this? Or is this a different underlying root cause where there should never be a reshape_lazy in the first place?
| template struct test_conv_add_tune<migraphx::shape::fp8e4m3fn_type>; | ||
| template struct test_conv_add_tune<migraphx::shape::fp8e5m2_type>; | ||
|
|
||
| struct test_conv_add_layout_tune : verify_program<test_conv_add_layout_tune> |
There was a problem hiding this comment.
This should be named test_conv_add_reshape_lazy_transpose. It also be added to it own .cpp file.
| migraphx::make_op("convolution", {{"padding", {1, 1, 1, 0}}, {"stride", {2, 2}}}), | ||
| x1, | ||
| w1); | ||
| auto pooling = |
There was a problem hiding this comment.
I dont think this pooling op is needed for this test.
| auto w2 = mm->add_literal( | ||
| migraphx::generate_literal({migraphx::shape::half_type, {1, 256, 3, 2}}, 1)); | ||
|
|
||
| auto conv1 = mm->add_instruction( |
There was a problem hiding this comment.
I dont think the conv1 is needed for this test.
We already check for these "layout"-ops(which it more accurately called shape transforms as these ops do not change the memory layout) when we do the split, but we missed the check for
It needs to add |
|
Actually, we shouldn't be adding |
I think we should be careful with that because is_module_fusible depends on the selected perfConfig (the function’s solution param), not just the problem key/problem config. The same problem key can return either true or false depending on which perfConfig is being tested. For example, one config may fit the fused op into LDS while another does not. So adding that boolean to the problem key would encode a perfConfig-level property as if it were a problemConfig-level property (i.e., a property of the kernel itself). I do see the issue you’re trying to solve, though. Would it make more sense to add a problemConfig-level capability, like supportsSplitK, instead of using is_module_fusible directly? That check is independent of the chosen perfConfig, so it seems much safer to include in the problem key than the result of is_module_fusible. |
Motivation
Some
gpu::mlir_opconvolution + pointwise fusions can fail to compile when the fused module is followed by layout-only operations such asreshape_lazyortranspose. In these cases, MIGraphX was unable to use the selected perfConfig for the fused module, and was also failing to properly split the module into separate conv + pointwise kernels because of the layout tail.Technical Details
Update the GPU JIT MLIR compiler to detect layout-only tails after pointwise fusion and split the fused op into separate convolution MLIR, pointwise, and layout-copy kernels when needed. This keeps the MLIR portion compatible with the selected perf config while preserving the final output layout.
Changelog Category
Add a
CHANGELOG.mdentry for any option other thanNot ApplicableFollow the LLVM AI Tool Use Policy for contributions using AI.