Skip to content

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
@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.

2 participants