8000 [Codegen][Tuner] merge the default td specs by bangtianliu · Pull Request #20127 · iree-org/iree · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

[Codegen][Tuner] merge the default td specs #20127

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.

Sign up for GitHub

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

Jump to bottom
Merged
merged 13 commits into from
Mar 21, 2025
Merged

[Codegen][Tuner] merge the default td specs #20127

merged 13 commits into from
Mar 21, 2025

Conversation

bangtianliu
Copy link
Contributor
@bangtianliu bangtianliu commented Feb 28, 2025 •
edited
Loading

This PR adds the implementation of merging all the default td specs (if all inner modules are) in the LinkTuningSpecsPass. It mainly merges multiple transform.foreach_match operations found across different inner modules into a single consolidated operation inside a newly created top-level __kernel_config NamedSequenceOp.

Issue: nod-ai/shark-ai#810

@bangtianliu bangtianliu requested a review from hanhanW as a code owner February 28, 2025 01:01
@bangtianliu bangtianliu marked this pull request as draft February 28, 2025 01:01
@bangtianliu bangtianliu requested review from kuhar and Max191 February 28, 2025 01:01
@bangtianliu bangtianliu force-pushed the merge_default_specs branch 4 times, most recently from 5e55989 to 30e7c43 Compare February 28, 2025 17:40
@bangtianliu bangtianliu marked this pull request as ready for review February 28, 2025 18:09
@bangtianliu bangtianliu force-pushed the merge_default_specs branch from eeaefb0 to 668fad2 Compare March 3, 2025 00:27
kuhar
kuhar requested changes Mar 3, 2025
View reviewed changes
Copy link
Member
@kuhar kuhar left a comment

Choose a reason for hiding this comment

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

To thoroughly test this code, we can start by assuming that tuning specs with default entypoints should always link, and then emitting warning whenever we notice they can't be linked. This is something we can test with verify-diagnostics.

The way we usually do this in compilers is that we try to have two phases:

  1. Analysis that determines transformation legality
  2. Transformation that cannot bail out

@bangtianliu bangtianliu force-pushed the merge_default_specs branch from 665cf37 to 0004cc1 Compare March 4, 2025 21:14
@bangtianliu bangtianliu requested a review from kuhar March 4, 2025 21:21
Max191
Max191 reviewed Mar 5, 2025
View reviewed changes
Copy link
Contributor
@Max191 Max191 left a comment

Choose a reason for hiding this comment

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

Leaving initial comments. Coming back to look at the big emitLinkedDefaultTuningSpec later this morning.

Max191
Max191 requested changes Mar 5, 2025
View reviewed changes
Copy link
Contributor
@Max191 Max191 left a comment

Choose a reason for hiding this comment

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

High level question: Do we need to be checking for the default tuning spec attr before merging? This seems useful for non-default specs too, and it is already matching for an expected format of the nested modules.

@bangtianliu
Copy link
Contributor Author

High level question: Do we need to be checking for the default tuning spec attr before merging? This seems useful for non-default specs too, and it is already matching for an expected format of the nested modules.

Yes, I think it could be another todo item if we could use the attribute as a hint to simplify some implementations. CC @kuhar.

@kuhar
Copy link
Member
kuhar commented Mar 5, 2025

@Max191 note that this checks for attributes that say that the tuning spec has a default entry point, not that it's a default spec in the compiler. A default spec is a spec with a default entrypoint that we decide to ship with the compiler.

So to directly answer your question: yes, I think it makes sense to restrict this to a subset of transform scripts and also switch the tuner to use that format.

@Max191
Copy link
Contributor
Max191 commented Mar 5, 2025

@Max191 note that this checks for attributes that say that the tuning spec has a default entry point, not that it's a default spec in the compiler. A default spec is a spec with a default entrypoint that we decide to ship with the compiler.

So to directly answer your question: yes, I think it makes sense to restrict this to a subset of transform scripts and also switch the tuner to use that format.

I see, thanks for the clarification!

@kuhar kuhar requested a review from andfau-amd March 5, 2025 16:08
@Max191
Copy link
Contributor
Max191 commented Mar 5, 2025

Actually, followup question: Do we need the extra attribute on the module? We could just check that there is a named_sequence with iree_codegen.tuning_spec_entrypoint inside the module, right?

Seems like it's just an extra hint for the linking pass that isn't really necessary to me. Am I missing something?

@kuhar
Copy link
Member
kuhar commented Mar 5, 2025

We use it to trigger the attribute verifier over the whole module

@kuhar
Copy link
Member
kuhar commented Mar 5, 2025 •
edited
Loading

We could just check that there is a named_sequence with iree_codegen.tuning_spec_entrypoint inside the module, right?

There's more than that: we also verify the entry point function signature (name and types) and expect a single foreach_match op in the body. And that there's exactly one entry point like that.

@Max191
Copy link
Contributor
Max191 commented Mar 5, 2025

We use it to trigger the attribute verifier over the whole module

I see, this is what I was missing, thanks!

There's more than that: we also verify the entry point function signature (name and types) and expect a single foreach_match op in the body. And that there's exactly one entry point like that.

Sorry, this is what I had meant (i.e., run the verification of the full module), but I see now that we have this attribute plumbed through in verifyOperationAttribute. It all makes sense to me now!

@bangtianliu
Copy link
Contributor Author

I've just submitted a commit addressing all of Max's comments. I'll follow up with another PR to add verification for iree_codegen.tuning_spec_entrypoint, ensuring that a single foreach_match op is present in the body. Once that verification PR is merged, this PR can be further simplified.

@bangtianliu bangtianliu force-pushed the merge_default_specs branch from ab6039e to f97aeda Compare March 6, 2025 21:07
@bangtianliu
Copy link
Contributor Author

I will further simplify this PR after #20173 is landed.

@kuhar
Copy link
Member
kuhar commented Mar 6, 2025

I will further simplify this PR after #20173 is landed.

You can rebase this PR on top of 20173 before it lands. This will make code review easier.

@bangtianliu bangtianliu force-pushed the merge_default_specs branch from f97aeda to a25cafd Compare March 6, 2025 21:56
@bangtianliu bangtianliu requested a review from kuhar March 18, 2025 22:38
@bangtianliu bangtianliu force-pushed the merge_default_specs branch from e9351ef to 8c50308 Compare March 18, 2025 23:33
kuhar
kuhar reviewed Mar 19, 2025
View reviewed changes
compiler/src/iree/compiler/Codegen/Common/LinkTuningSpecsPass.cpp Outdated
llvm::DenseMap<NamedSequenceOp, ForeachMatchOp>
&namedSequenceToForeachMatch) {
StringRef specName = op.getSymName();
std::string newSpecName = getUniqueSpecName(specName, specNameCounts);
Copy link
Member

Choose a reason for hiding this comment

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

I'd like to avoid setting the suffix in one place just to forget it and try to recover this information in getBaseName and update again in getUniqueSpecName. This is bad design.

I'd rather do one of:

  1. only set it under a flag in emitLinkedTuningSpec and decide it in the code that does merging if it has different naming requirements
  2. add a different name uniqing mode to emitLinkedTuningSpec to do the right thing for merging
  3. update all names as a pre-processing steps that's common across both code paths (linking and merging)

@kuhar
Copy link
Member
kuhar commented Mar 19, 2025

Overall looks good, the only remaining issue I see is the way we handle names.

@bangtianliu
[Codegen][Tuner] merge the default td specs
376f97b 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu
plumb through MaterializeTuningSpecsPass
b738c6e 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu
solve the name confliction
21e272d 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu
address reviewer comments
91c7a04 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu
address reviewer comments: further solve name confliction, format cod…
121b3d0 
…e and add test

Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu
rebase and simplfiy the code
b988c3b 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu
simplify the code based on improved verifier
51d1c3b 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu
format the code
a098efb 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu
collect foreachMatchOps at step-2
76d9e10 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu
add helper functions and rename the map
2277ab1 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu
add a new method to address name conflicts for merging td specs
4aa2fff 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu bangtianliu force-pushed the merge_default_specs branch from 8c50308 to 4aa2fff Compare March 19, 2025 21:26
@bangtianliu bangtianliu requested a review from kuhar March 19, 2025 21:26
kuhar
kuhar requested changes Mar 20, 2025
View reviewed changes
@bangtianliu
reorder named squence ops in test
bad539b 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu bangtianliu requested a review from kuhar March 20, 2025 19:50
Max191
Max191 approved these changes Mar 20, 2025
View reviewed changes
Copy link
Contributor
@Max191 Max191 left a comment

Choose a reason for hiding this comment

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

Overall LGTM, please address the comments before landing, though.

kuhar
kuhar reviewed Mar 20, 2025
View reviewed changes
@Max191
Copy link
Contributor
Max191 commented Mar 20, 2025

(Also wait for Jakub's review before landing)
58E8

@bangtianliu
address reviewer comments: format the code
bb276c2 
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
@bangtianliu bangtianliu force-pushed the merge_default_specs branch from c3c599c to bb276c2 Compare March 20, 2025 21:22
@bangtianliu bangtianliu requested a review from kuhar March 20, 2025 21:29
kuhar
kuhar approved these changes Mar 21, 2025
View reviewed changes
Copy link
Member
@kuhar kuhar left a comment

Choose a reason for hiding this comment

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

LGTM

@kuhar
Copy link
Member
kuhar commented Mar 21, 2025 •
edited
Loading

Please update the PR description before merging, this PR does not fix the issue:

fixes nod-ai/shark-ai#810

You can replace this with Issue: <url> to link to it but not close after merging

@kuhar kuhar merged commit 8e4ce9c into iree-org:main Mar 21, 2025
212 of 228 checks passed
@ScottTodd ScottTodd mentioned this pull request May 5, 2025
Closed
6 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@kuhar kuhar kuhar approved these changes

@Max191 Max191 Max191 approved these changes

@hanhanW hanhanW Awaiting requested review from hanhanW hanhanW is a code owner

@MaheshRavishankar MaheshRavishankar Awaiting requested review from MaheshRavishankar

@andfau-amd andfau-amd Awaiting requested review from andfau-amd

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@bangtianliu @kuhar @Max191 @andfau-amd
0