Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Stream] Implement SpecializeEncodings pass (1/n) #19502

Merged
merged 6 commits into from
Jan 9, 2025

Conversation

hanhanW
Copy link
Contributor

@hanhanW hanhanW commented Dec 17, 2024

There are three major changes in the revision:

  • Introduce AffinityAnalysisDialectInterface Stream dialect interface. It is used to fetch attributes that are defined by other dialects. In the revision, HAL implements the dialect interface, and it can return whatever attribute attached in HAL::ExecutableTarget attributes. The main idea of the dialect interface is that Stream does not need to depend on HAL to get the layout information.
  • Add cloneWithLayouts method to the EncodingAttr. It is used in the encoding specialization pass where it can resolve the layout requirements and add it to the layouts field. The other optional parameters are dropped because the layout is already resolved. It can be a new Encoding dialect attribute because it is just describing the layout. The stream tensor ops do not need to know the op_type, element_types and operand_index parameters. It only needs the layout information, and the attribute should implement the interface method.
  • Partially implement the SpecializeEncodings pass. The responsibility of the pass is large, so I decide to implement it incrementally. This revision only implements the mechanism of updating stream tensor ops' encoding, and only stream.tensor.sizeof op is supported. The rest of the support for other stream tensor op can be added later on. The executable duplication and the update of dispatch ops will be implemented in subsequent PRs.

@hanhanW hanhanW requested review from lialan and Max191 and removed request for bjacob December 17, 2024 14:46
using AffinityAnalysisDialectInterface::AffinityAnalysisDialectInterface;
IREE::Stream::LayoutAttrSolverFn
makeLayoutAttrSolver(ModuleOp moduleOp) const {
return [=](IREE::Stream::AffinityAttr aff, Operation *op,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

always use full names for things wherever possible - it's useful hints to readers as to what the type is and how it is treated - shortening to arbitrary forms makes it impossible to find/replace in codebases and more difficult for non-native speakers (here and anywhere else you've shortened things).

Suggested change
return [=](IREE::Stream::AffinityAttr aff, Operation *op,
return [=](IREE::Stream::AffinityAttr affinityAttr, Operation *op,

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, done.

compiler/src/iree/compiler/Dialect/HAL/IR/HALDialect.cpp Outdated Show resolved Hide resolved
compiler/src/iree/compiler/Dialect/HAL/IR/HALDialect.cpp Outdated Show resolved Hide resolved
// TODO(hanchung): We should use the default device in this case. However,
// it is not guaranteed that default device attribute will always be set in
// the IR. (Is the statement correct?)
auto affAttr = affinityOp.getAffinityAttr();
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

as mentioned, affinityAttr (here and elsewhere)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done. @benvanik does my comment above make sense? If so, I'm going to remove the (Is the statement correct?).

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

At this point in the IR I don't believe there is a concept of default device, so the whole TODO is not required - in a program with multiple functions there may be many routes to a function through the call graph and we can't locally make decisions about the whole program like that.

@hanhanW hanhanW requested a review from lialan January 6, 2025 07:39
@hanhanW
Copy link
Contributor Author

hanhanW commented Jan 6, 2025

@benvanik can you take a second look at this?

@hanhanW hanhanW force-pushed the specialize-encodings-1-n branch from a76fb9a to 86dd9b0 Compare January 7, 2025 05:08
Signed-off-by: hanhanW <[email protected]>
@hanhanW hanhanW enabled auto-merge (squash) January 9, 2025 03:19
@hanhanW hanhanW merged commit 02d145e into iree-org:main Jan 9, 2025
32 checks passed
@hanhanW hanhanW deleted the specialize-encodings-1-n branch January 9, 2025 03:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants