[AIRADSW-732] Fix reshape layout propagation - #5096
Conversation
ivarusic-amd
added
bugfix
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
There was a problem hiding this comment.
This prevents us from propagating the shape because another reshape_lazy downstream is not contiguous.
A much simpler approach is to not lower reshape -> reshape_lazy until after eliminate_contiguous. And remove the propagate_reshape_layout pass and replace it with a lower_reshape pass.
| // Singleton dimensions can make a stride order ambiguous. In that case, | ||
| // find_permutation may produce a packed shape different from relayout that cannot | ||
| // alias the reshape output. Keep the standardizing contiguous instead. | ||
| auto reshaped = reshape_dims(layout_shape, rl->get_shape().sym_dims(), {.lazy = true}); |
There was a problem hiding this comment.
Why are you applying the reshape to the output?
|
|
||
| auto layout_op = make_op("layout", {{"permutation", find_permutation(*relayout)}}); | ||
| auto layout_shape = layout_op.compute_shape({s}); | ||
| // Singleton dimensions can make a stride order ambiguous. In that case, |
There was a problem hiding this comment.
What is a singleton dimension?
There was a problem hiding this comment.
I think she means one dim here?
| // find_permutation may produce a packed shape different from relayout that cannot | ||
| // alias the reshape output. Keep the standardizing contiguous instead. | ||
| auto reshaped = reshape_dims(layout_shape, rl->get_shape().sym_dims(), {.lazy = true}); | ||
| if(not reshaped or not can_propagate_shape(m, rl, *reshaped)) |
There was a problem hiding this comment.
If it throws an error downstream, then this wont propagate the shape, which it should be.
There was a problem hiding this comment.
Yes, the current fix was intentionally conservative: if propagating the new shape would invalidate a downstream operation, it retains the original contiguous path instead of failing compilation. This was preventing the regression; but i see the problem with it now
Check flagged results 🔆 * No develop baseline was found for this PR's branch point; compared against the latest available develop run instead. |
|
TedThemistokleous
left a comment
TedThemistokleous
left a comment
There was a problem hiding this comment.
Few comments. Feel free to reach out
| shapes.emplace(&ins, | ||
| ins.get_operator().compute_shape(input_shapes, ins.module_inputs())); | ||
| } | ||
| catch(const exception&) |
There was a problem hiding this comment.
So catching any exception here? We should look for a specific exception instead of doing a catch all
There was a problem hiding this comment.
There is no specific exception type for the different compute_shape failures; that's the reason why. Agree
| std::back_inserter(input_shapes), | ||
| [&](instruction_ref input) { | ||
| auto iter = shapes.find(&*input); | ||
| return iter == shapes.end() ? input->get_shape() : iter->second; |
There was a problem hiding this comment.
Can we rewrite this to remove the find? We're doing a search and if that fails you're still searching O(n) here. The de reference and reference looks really odd here too.
There was a problem hiding this comment.
Yes, i though average would be O(1).. I will change it
Hi @pfultz2 , I tried your approach locally by leaving reshape generic through eliminate_contiguous, then lowering it in a new lower_reshape pass. |
3 participants
Motivation
Prevent gpu::propagate_reshape_layout from producing layouts that cause downstream reshape_lazy operations to fail with an unpacked-axis error.
Technical Details
A potential follow-up is to traverse only dependent instructions to reduce compile-time overhead.This adds traversal and topological-ordering complexity, may not improve dense graphs.
Changelog Category
Add a
CHANGELOG.mdentry for any option other thanNot ApplicableFollow the LLVM AI Tool Use Policy for contributions using AI.