[Encoding] Add encoding_dims flow/stream conversion #23322
Conversation
jtuyls
requested review from
IanWood1,
MaheshRavishankar,
benvanik and
hanhanW
as code owners
Signed-off-by: Jorn Tuyls <[email protected]>
jtuyls
force-pushed
the
users/jtuyls/encoding-dims-flow-stream-conversion
branch
from
aff97cf to
bcc050b
Compare
hanhanW
left a comment
hanhanW
left a comment
There was a problem hiding this comment.
TLDR: I think it will break encoding unification; I need to re-think about the design. See inline comments for a little more details.
| // Verify encoding_dims count matches what the encoding expects. | ||
| // Check both operand (for decode) and result (for encode) encodings. | ||
| for (RankedTensorType type : {operandType, resultType}) { | ||
| if (IREE::Encoding::SerializableAttr encodingAttr = | ||
| dyn_cast_or_null<IREE::Encoding::SerializableAttr>( | ||
| type.getEncoding())) { | ||
| std::optional<int64_t> expectedDims = | ||
| encodingAttr.getNumDynamicEncodingDims(); | ||
| if (expectedDims.has_value() && | ||
| static_cast<int64_t>(getEncodingDims().size()) != *expectedDims) { | ||
| return emitOpError() << "encoding expects " << *expectedDims | ||
| << " encoding dim(s), but " | ||
| << getEncodingDims().size() << " provided"; | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Maybe we can create a method in EncodingTypes.h and reuse it for Encoding/Flow/Stream ops.
| for (RankedTensorType type : {operandType, resultType}) { | ||
| if (IREE::Encoding::SerializableAttr encodingAttr = | ||
| dyn_cast_or_null<IREE::Encoding::SerializableAttr>( | ||
| type.getEncoding())) { | ||
| std::optional<int64_t> expectedDims = | ||
| encodingAttr.getNumDynamicEncodingDims(); |
There was a problem hiding this comment.
I've seen this pattern at least three times. Do we consider introducing a method in Encoding dialect? E.g., IREE::Encoding::getNumDynamicEncodingDims(Type).
| rewriter.replaceOpWithNewOp<IREE::Stream::AsyncDispatchOp>( | ||
| encodeOp, exportOp, | ||
| /*workload=*/dynamicDims, encodeOp.getResult().getType(), operands, | ||
| /*workload=*/workload, encodeOp.getResult().getType(), operands, |
There was a problem hiding this comment.
nit: remove the argument comment
| /*workload=*/workload, encodeOp.getResult().getType(), operands, | |
| workload, encodeOp.getResult().getType(), operands, |
There was a problem hiding this comment.
I think this is broken when an encoding op has encoding_dims while the other does not. (assuming the same encoding attributes are identical.)
The encodingSignature does not take encoding_dims into account.
| // encoding1 has dynamic M, encoding2 has static M=4. The GPU resolver's | ||
| // getUnifiedEncoding returns the first encoding, so both should be unified | ||
| // to #encoding1 (with iteration_sizes = [?, 4096, 4096]). |
There was a problem hiding this comment.
It looks broken to me. How does encoding dimensions get propagated to other functions (outside initializer)?
I feel that we'll need to create new globals and store these values; perform updates for dispatch ops. However, the executables are already outlined, which makes this pass broken when the number of encoding dims do not match.
I don't have a solution atm. We probably can't do encoding unification if the number of encoding_dims is not the same.
Then the next design question pops out in my mind is that we may have to carry all the values (including c128) in this case. I need to think more about it. Thanks for raising the PR and providing the examples.
There was a problem hiding this comment.
It looks broken to me. How does encoding dimensions get propagated to other functions (outside initializer)?
Do we really need that? This is before SpecializeEncodings, so we have unserialized encodings at this point and we don't update the internals of the executables right? After SpecializeEncodings, we have serialized encodings and shouldn't need the encoding dims anymore? Maybe I am missing something though and this may need a higher bw discussion?
There was a problem hiding this comment.
I had such question when I saw the result of MaterializeEncoding pass:
stream.async.dispatch on(#{{.+}}) @[[$EX]]::@[[$ENTRY]][%[[D0]], %[[D1]], %[[D0]], %[[D1]], %[[M]], %[[N]], %[[K]]]
I forgot that we'll need to pass in the encoding dims to dispatch op. If the unified encoding does not carry the same encoding dims, the dispatch op is broken, right?
In this example, the dispatch op that uses #encoding1 takes %M and the dispatch op that uses #encoding2 takes zero encoding dims. If the result of encoding unification is #encoding1, the latter dispatch op has to be updated with taking one more encoding dimension. Otherwise, you're not able to do your specialization later on.
I may be missing something. Your specialization (about layouts) is exercised inside the executable, no?
Implements phase 4 for adding specialization support on dynamic values for data-tiling: #22370. This carries the encoding dynamic values through Flow and Stream dialects.
Assisted-by: Claude