[Unity][UX][TIR] Implement privacy annotation for the @prim_func decorator #15171
Conversation
|
Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.
Generated by tvm-bot |
1 similar comment
|
Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.
Generated by tvm-bot |
python/tvm/tir/schedule/testing.py
Outdated
| from tvm.tir.schedule import Schedule, Trace | ||
|
|
||
|
|
||
| def assert_structural_equal_gs(f1: PrimFunc, f2: PrimFunc, *args: Any, **kwargs: Any) -> None: |
There was a problem hiding this comment.
I don't like the name but I didn't want it to be too verbose because this has to be used in a lot of places. Also, is there a better file where this should reside? It's used in tests outside of those related to scheduling.
|
|
||
| def _check(original, transformed): | ||
| mod = tvm.IRModule.from_expr(original) | ||
| mod = tvm.IRModule.from_expr(original.with_attr("global_symbol", "main")) |
There was a problem hiding this comment.
Should from_expr include a global symbol by default? The only reason I hesitated on implementing that is because Relay and Relax also use from_expr and I didn't want to mess with that.
slyubomirsky
force-pushed
the
privacy-annotation-tir
branch
from
e528cb1 to
dcbaff6
Compare
| @I.ir_module | ||
| class Expected: | ||
| @T.prim_func | ||
| @T.prim_func(private=True) |
There was a problem hiding this comment.
Throughout these Relax tests, I assumed that most PrimFuncs we will add through fusion, etc, are likely not intended to be public. Please let me know if this assumption is mistaken.
| Note that the below keys are reserved: | ||
| * "primfunc_name_hint" is reserved for passing a name hint | ||
| to the PrimFunc that gets generated. | ||
| * "primfunc_public" is reserved for indicating whether the |
There was a problem hiding this comment.
For emitting PrimFuncs from Relax, I assume private is the default we want. Happy to do it the other way if that's desirable.
|
edit: One suggestion was that it had to do with the edit 2: That was right! Thank you, @csullivan |
… the BlockBuilder
slyubomirsky
force-pushed
the
privacy-annotation-tir
branch
from
78dae74 to
5693c75
Compare
|
Thank you @slyubomirsky , do you mind sending this PR to main instead? we can then merge it back to unity |
|
Okay. I'll have to get rid of the Relax test changes to have it in main. |
|
Reposted for mainline: #15214. There seem to be some different tests in mainline, so I'll probably have to chase after some new failures. |
Implements proposal #14899 for TIR. Adapted from Unity PR #15171, by the same method as PR #15140. Namely, a TIR PrimFunc can be specified to be private (without a global symbol attribute) in TVMScript in the `prim_func` decorator. By default, `PrimFunc`s are not private, so they will have a `global_symbol` attribute that is mapped to their name. Example usage: ```python # not private: its global symbol will be "func" @T.prim_func def func(...): ... # no global symbol included @T.prim_func(private=True) def func(...): ... ``` This did require changing very, very many tests, unfortunately.
3 participants
Implements proposal #14899 for TIR, by the same method as PR #15140. Namely, a TIR PrimFunc can be specified to be private (without a global symbol attribute) in TVMScript in the
prim_funcdecorator. By default,PrimFuncs are not private, so they will have aglobal_symbolattribute that is mapped to their name.Example usage:
This did require changing very, very many tests, unfortunately.