Layout conversion bypass for blocked to dotOperand #4538
Conversation
| int kDim = dotOperandLayout.getOpIdx() == 0 ? rank - 1 : rank - 2; | ||
| int nonKDim = dotOperandLayout.getOpIdx() == 0 ? rank - 2 : rank - 1; | ||
| auto ctaLayout = blockedLayout.getCTALayout(); | ||
|
|
There was a problem hiding this comment.
One issue we have in the codebase is lots of mysterious layout/indexing--it's not easy for others reading the code to pick up the intent. The following might not be that tricky; but still can we add a comment to explain what the following checks are doing in a high level?
There was a problem hiding this comment.
Thanks for adding the comments! but the wording is quite confusing to me right now. What about something like
The following logic checks that a source blocked layout B matches a destination dot operand layout with blocked layout parent P. It's considered match if 1) each thread holds a whole copy of all elements along the K dimension for B, and 2) distribution along all other non-K dimensions match between S and B. This is to guarantee that each thread have all the data needed for reduction without exchange with other threads. (And/or whatever other reasons why we want this kind of match.)
alefimov-amd
force-pushed
the
fma_lds_bypass
branch
from
de080f3
to
44c50c5
Compare
This PR extends shared memory bypass for blocked->dot operand conversions and adds bypass check in DecomposeUnsupportedConversions and ReduceDataDuplication.
alefimov-amd
force-pushed
the
fma_lds_bypass
branch
from
44c50c5
to
5e50eb0
Compare
| // i.e. tensor<64x32xf16, #dot_op<{opIdx=0, parent=#blocked}>> will have sizePerThread = [<depends on #blocked>, 32] | ||
| // and tensor<64x32xf16, #dot_op<{opIdx=1, parent=#blocked}>> will have sizePerThread = [64, <depends on #blocked>] | ||
| // | ||
| // For example tensor<64x32xf16, #dot_op<{opIdx=0, parent=#blocked<{sizePerThread = [2, 8], threadsPerWarp = [32, 1]}>>> |
There was a problem hiding this comment.
This is going from dot operand to blocked layout? Isn't it the reverse of what we are doing in this function? I'm also not sure the distribution is correct? Isn't this contradicting to the check at L571?
There was a problem hiding this comment.
Yes, this is misleading, let me change this comment.
I mean that these dot and blocked layouts are equal? I should not use "converted" here
| int kDim = dotOperandLayout.getOpIdx() == 0 ? rank - 1 : rank - 2; | ||
| int nonKDim = dotOperandLayout.getOpIdx() == 0 ? rank - 2 : rank - 1; | ||
| auto ctaLayout = blockedLayout.getCTALayout(); | ||
|
|
There was a problem hiding this comment.
Thanks for adding the comments! but the wording is quite confusing to me right now. What about something like
The following logic checks that a source blocked layout B matches a destination dot operand layout with blocked layout parent P. It's considered match if 1) each thread holds a whole copy of all elements along the K dimension for B, and 2) distribution along all other non-K dimensions match between S and B. This is to guarantee that each thread have all the data needed for reduction without exchange with other threads. (And/or whatever other reasons why we want this kind of match.)
This PR extends shared memory bypass for blocked->dot operand conversions and adds bypass check in DecomposeUnsupportedConversions and ReduceDataDuplication.
This PR is a part of PR series. Final goal is to improve efficiency of small dot operations and bypass as much shared memory accesses as possible.
Rough list of PRs: