[LAYOUTS] Generalise HoistLayoutConversion to work with arbitrary layouts and chains of ops #5673
+425
−407
Conversation
This file contains 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
lezcano
force-pushed
the
remat_dot
branch
2 times, most recently
from
1eed10c to
a77a54a
Compare
lezcano
changed the title
lezcano
commented
Jan 27, 2025
Comment on lines
+3278
to
+3187
| #blocked = #ttg.blocked<{sizePerThread = [1, 2], threadsPerWarp = [4, 8], warpsPerCTA = [4, 1], order = [1, 0]}> | ||
| #mma = #ttg.nvidia_mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [1, 4]}> | ||
| module attributes {"ttg.num-warps" = 4 : i32, "ttg.target" = "cuda:80"} { | ||
| tt.func @dot_op_hoisted_to_load_with_unsupported_op_and_initializer_above_slice( |
There was a problem hiding this comment.
This is all codemovement + adding this test that was proposed but not merged in #5349 (comment)
as we now hoist everything as expected
lezcano
commented
Jan 27, 2025
ThomasRaoux
approved these changes
Jan 28, 2025
Mogball
reviewed
Jan 28, 2025
| // This could be generalised if necessary | ||
| if (!loadOp) { | ||
| auto op = v.getDefiningOp(); | ||
| if (isa<arith::ConstantOp>(op) || noDataMovement(op)) { |
There was a problem hiding this comment.
Should ConstantOp just be put inside noDataMovement?
There was a problem hiding this comment.
Probably yeah, but leaving it as-is for now because it currently works and we are probably going to refactor this pass in the near future so whatever.
We now support all layouts as LL, and reductions support any layout as input. As such, at least in theory, we should be able to propagate layouts freely, even DotOperands, similar to what we do with other layouts. This PR is a bit tentative. Let's see if anything interesting breaks
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.
We generalise
HoistLayoutConversionto lift a givenconvert_layout dot_operandabove any chain of operations that do not require data movement. We
could totally generalise this in the future to lift it over other ops. We do
this as a first step to keep the code somewhat similar to the previous
one.
Regarding the previous limitations of
canHoistDotOpEncV2I did a bit of archeology:UIToFpOp, this is now supported after [BACKEND] Get rid of unpack/pack I32 #5044convert_layoutrework.We also add proper support for
isPureforelementwise_inline_asmopsOn the location of the code, we just leave it in
RemoveLayoutConversion.cpptotake advantage of the rather generic implementation of
rewriteSlice. We could totallymove this pass outside of
remove-layout-conversion, as it's probably enough to runit once. This code will go through further changes in the near future, so we'll assess this
then.