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

[Relax][Transform] Handle identical PrimFunc with distinct VDevice #16959

Merged
merged 2 commits into from May 13, 2024

Conversation

Lunderberg
Copy link
Contributor

Prior to this commit, if an IRModule contained two expressions, where the types of the arguments differed only by the VDevice, these would be legalized to produce a single PrimFunc. This PrimFunc would have the a tvm::attr::kTarget annotation specific to one of those expressions, and would be incorrect for use in the other location.

This commit updates the LegalizeOps transform to handle this case, producing multiple TIR PrimFuncs if required by the VDevice annotations.

Prior to this commit, if an `IRModule` contained two expressions,
where the types of the arguments differed only by the `VDevice`, these
would be legalized to produce a single PrimFunc.  This PrimFunc would
have the a `tvm::attr::kTarget` annotation specific to one of those
expressions, and would be incorrect for use in the other location.

This commit updates the `LegalizeOps` transform to handle this case,
producing multiple TIR PrimFuncs if required by the `VDevice`
annotations.
@Lunderberg Lunderberg force-pushed the relax_legalize_primfuncs_with_vdevice branch from 2d7e065 to 2bc7f47 Compare April 30, 2024 19:19
@Lunderberg
Copy link
Contributor Author

CI is stalled on tests/python/relax/test_transform_few_shot_tuning.py::test_funcs[True-Softmax], with no updates shown in the last 2 days. This test was not modified by this PR, and passes when run locally, so it looks like an issue with the CI environment.

Restarting CI to see if it resolves the issue.

@Lunderberg
Copy link
Contributor Author

@tvm-bot rerun

@masahi masahi merged commit c2d14ae into apache:main May 13, 2024
19 checks passed
@Lunderberg Lunderberg deleted the relax_legalize_primfuncs_with_vdevice branch May 14, 2024 14:38
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.

None yet

2 participants