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

[TIR][Driver] Use BindTarget to specify target for FP8 legalization #16767

Merged
merged 3 commits into from
Mar 25, 2024

Conversation

slyubomirsky
Copy link
Contributor

This PR cleans up the interface for the FP8 legalization passes by having them use the target attributes that are added by the BindTarget pass. This not only removes the need to pass the target as a parameter to those passes but also, in principle, allows for PrimFuncs with different targets to coexist and be handled correctly by the passes. Thanks to @Lunderberg for giving the suggestion.

Copy link
Contributor

@Lunderberg Lunderberg left a comment

Choose a reason for hiding this comment

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

Looks good to me, and thank you for the follow-up!

@@ -569,15 +569,14 @@ transform::Sequential MixedModulePassManager(IRModule mixed_mod, Target target)

Array<Pass> mixed_pass_list;

mixed_pass_list.push_back(tir::transform::FP8ComputeLegalize(target));
mixed_pass_list.push_back(tir::transform::BindTarget(target));
Copy link
Contributor

Choose a reason for hiding this comment

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

Can we add a comment here specifying that the BindTarget should occur first in this sequence, so that later passes can rely on the target attribute being present? (It looks like both VerifyVTCMLimit and LowerVtcmAlloc were implemented more recently than BindTarget, and probably should have been placed after BindTarget at that point.)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Yeah will do

Copy link
Contributor

Choose a reason for hiding this comment

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

but also, in principle, allows for PrimFuncs with different targets to coexist and be handled correctly by the passes.

Here BindTarget is applying the target to all functions in the mixed IRModule, no?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Yeah that's what it does. In principle, we could set the attributes differently.

Copy link
Contributor

Choose a reason for hiding this comment

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

Slight nitpick: All functions that do not already have the target attribute. A function could already have the target attribute, such as defining a module that contains kernels for multiple different targets.

@Lunderberg
Copy link
Contributor

@csullivan @JosephTheOctonaut FYI on this change, as it relates to your work on FP8. I don't expect it to impact your work, but wanting to make sure you're aware of it.

@yongwww yongwww merged commit cb08f0d into apache:main Mar 25, 2024
22 checks passed
thaisacs pushed a commit to thaisacs/tvm that referenced this pull request Apr 3, 2024
…apache#16767)

* Do not pass target explicitly to FP8 legalization, use BindTarget instead

* Lint: Remove unused import

* Add comment on pass ordering
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.

4 participants