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 93.09% 93.09%
========================================
Files 623 623
Lines 33073 33073
========================================
Hits 30789 30789
Misses 2284 2284 🚀 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 |
| 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?
There was a problem hiding this comment.
Ok I looked into more. The problem happens with reshape as well as reshape_lazy. I dont think we need to replace the find_final_split function, we can just add reshape_lazy to the list. It doesnt completely fix it though.
The problem is that we cant compile the pointwise module with reshapes in them, so we need to split the tail module and compile the pointwise and then insert the last half.
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. |
|
There are now models and usecases where this PR would resolve failures. @pfultz2 please work with @justinrosner to resolve |
7df8254 to
d118e6f
Compare
It looks like they are being inserted by |
| })) | ||
| return nullopt; | ||
| return *std::prev(it); | ||
| } |
There was a problem hiding this comment.
This doesnt split after an add or mul like find_final_split does which prevents fusing these operators with the gemm.
Why arent you using find_final_split and just add reshape_lazy to the list?
There was a problem hiding this comment.
Updated to use find_final_split again
| std::array<module_with_inputs, 2> mods = {std::move(mod_splits3[0]), | ||
| std::move(mod_splits3[1])}; | ||
| return insert(cops, mods, ins, split_ins); | ||
| } |
There was a problem hiding this comment.
Whats the reason for changing all this code? What is this trying to fix?
There was a problem hiding this comment.
This is modified slightly in the most recent version, but my intention with this was to handle the case where after splitting off the trailing shape transform, the pointwise kernel can no longer write directly to the final output allocation (because of the transpose/reshapes)
| { | ||
| auto tail_path = get_output_path(tail_split.value()); | ||
| if(std::find(tail_path.begin(), tail_path.end(), split_ins) != tail_path.end()) | ||
| split_ins = tail_split.value(); |
There was a problem hiding this comment.
The tail_split should not replace the split_ins. You need both split instructions. The split_ins is the barrier between the mlir and pointwise kernel and the tail_split is the barrier between the pointwise and "reshapes".
| if(std::find(tail_path.begin(), tail_path.end(), split_ins) != tail_path.end()) | ||
| split_ins = tail_split.value(); | ||
| } | ||
| std::array<module_with_inputs, 2> mod_splits = smod->split(input_args, {split_ins}); |
There was a problem hiding this comment.
This should be std::array<module_with_inputs, 3> mod_splits = smod->split(input_args, {split_ins}, tail_split) since we are now splitting into 3 modules.
| })) | ||
| return nullopt; | ||
| return *std::prev(it); | ||
| } |
There was a problem hiding this comment.
This should traverse up from the returns to skip over the layouts. This also probably needs to return a vector instead of a single instruction because of multiple returns.
| auto copy_input_shape = copy_op.expected_inputs.front(); | ||
| auto copy_input = m.insert_instruction( | ||
| ins, | ||
| migraphx::make_op("as_shape", {{"shape", to_value(copy_input_shape)}}), |
There was a problem hiding this comment.
Inserting as_shape is not correct.
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.