[Unity][IR][UX] Privacy annotation in Relax #15140

slyubomirsky · 2023-06-21T22:17:37Z

This PR implements the proposal in #14899. Namely, the @R.function decorator now has an optional private attribute. If a function is marked as private, then it will not have a global symbol attached to it and thus will not be externally accessible. By default, functions are not private, so the parser does insert a global symbol for them.

Here is an example of how to use the privacy attribute:

@I.ir_module
class Module:
    # Will not have a global symbol! Often not very useful to make main private...
    @R.function(private=True)
    def main(...): ...

    # private is false by default. This function will have a global symbol ("other")
    @R.function
    def other(...): ...

Passes that introduce new functions should decide whether a global symbol should be included with those functions. Most of the time, this is not the case (e.g., the functions extracted by the lambda lifting pass probably should not be publicly accessible).

Note that there are alternative approaches we can consider. For example, we may have "private" as a field on a function node and fill in the global symbol attributes later. This may be important for some passes.

tvm-bot · 2023-06-21T22:17:40Z

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.

cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

tvm-bot · 2023-06-21T22:17:40Z

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.

cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

tvm-bot · 2023-06-21T22:17:41Z

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.

cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

tvm-bot · 2023-06-21T22:17:41Z

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.

cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

slyubomirsky · 2023-06-21T22:18:38Z

tests/python/relax/test_transform_merge_composite_functions.py

        return gv

-    @R.function
+    @R.function(private=True)


@masahi I requested your review because I am not sure why exactly I needed to insert the privacy annotations for these tests. Leaving it off prevents any of the functions from being merged, but I'm not sure why that would happen. Do you know where it might matter that the functions have a global_symbol attribute?

@yelite Masa said you might also be able to comment on this issue

With this diff, these tests will pass without making those function private

diff --git a/src/relax/transform/dead_code_elimination.cc b/src/relax/transform/dead_code_elimination.cc index fe36eb28ef..596c68ef3e 100644 --- a/src/relax/transform/dead_code_elimination.cc +++ b/src/relax/transform/dead_code_elimination.cc @@ -96,7 +96,7 @@ IRModule RemoveUnusedFunctions(IRModule mod_, Array<runtime::String> entry_funcs for (auto f : existing_functions) { // If a function has an external linkage type, we do not remove it. // Otherwise, we check the function and remove it if it is not used anywhere. - if (f.second->GetLinkageType() == LinkageType::kInternal && !tracer.check_if_called(f.first)) { + if (!tracer.check_if_called(f.first)) { mod_->Remove(f.first); } } diff --git a/src/relax/transform/merge_composite_functions.cc b/src/relax/transform/merge_composite_functions.cc index 0bc92ba923..bdae719e31 100644 --- a/src/relax/transform/merge_composite_functions.cc +++ b/src/relax/transform/merge_composite_functions.cc @@ -60,6 +60,8 @@ #include <tvm/tir/function.h> #include "../../support/arena.h" +#include "tvm/ir/attrs.h" +#include "tvm/ir/function.h" #include "utils.h" namespace tvm { @@ -311,7 +313,8 @@ class CompositeInliner : public ExprMutator { if (func->GetAttr<String>(attr::kComposite)) { if (!inlined_functions_.count(func)) { - inlined_functions_.Set(func, CopyWithNewVars(func)); + inlined_functions_.Set(func, + WithoutAttr(CopyWithNewVars(func), tvm::attr::kGlobalSymbol)); } return Call(inlined_functions_[func], call->args); }

At the end of MergeCompositeFunction, a DCE pass will be called to removed the original copies of the functions that are merged. However, DCE will not remove public functions, i.e. those with a global symbol. Making those functions public prevents them from being cleaned up by DCE.

I think the logic here is correct. Those fused_xxx functions should be private as they are not meant to be called externally. So we probably don't want to change anything here.

Thank you very much for the explanation, @yelite. So, to confirm, you don't think there needs to be a change to the merge composite pass?

Yes, the merge composite pass shouldn't be changed. We might need to update the fusion pass if it produces fused functions with global symbol.

I think by default fusion pass should not produce functions with global symbol. Ideally we should prevent fusing functions that with global symbol --- because user expects to access the original function

slyubomirsky · 2023-06-21T22:20:00Z

src/relax/transform/lift_transform_params.cc


    // Step 2: Add the lifted function to the module
-    builder_->AddFunction(lift_plan_.f_transform_params, new_func_name);
+    // (The lifted function should be public so we add a global symbol to it)


I am assuming we want them to be public.

tests/python/relax/test_tvmscript_printer_relax.py

tqchen · 2023-06-21T22:30:46Z

cc @yongwww

slyubomirsky · 2023-06-22T02:55:43Z

src/script/printer/relax/function.cc

-                            ->Call({d->AsDoc<ExprDoc>(n->attrs, n_p->Attr("attrs"))})));
+        // if the function is a global function and has a global symbol,
+        // then don't print the global symbol (it will be implicit from not being private)
+        if (d->frames.size() == 3 && n->attrs->dict.count("global_symbol")) {


Not fond of having this magic number, but it was the only way I could figure out to check that the function is at the top level of a module. Possibly a hack--is there a better way I could check for that?

@yongwww might there be a more robust way to check if we're inside a module?

tqchen · 2023-06-22T18:30:35Z

tests/python/relax/test_tvmscript_parser.py

    bb = relax.BlockBuilder()
    x = relax.Var("x", relax.TensorStructInfo([10, 20], "float32"))
-    with bb.function("main", [x]):
+    with bb.function("main", [x], {"global_symbol": "main"}):


how about we also have private=False flag here(default to False)

private=False in the parser (which is the default) means that there's a global symbol in the attributes.

Do you mean the block builder should also have a private flag when building a function?

Right, this way it is closer to what function provide

I've implemented that, good idea 👍

tqchen · 2023-06-22T18:33:19Z

Thanks @slyubomirsky on the default behavior, would be great if we default private=False if not specified, and ask for explicit private=True.

It will involve more handling in printing(like not printing global symbol, and when global symbol do not exist, print as private=True), and modifying the test cases that comes with global_symbol (by removing them). But it would be much cleaner from user pov

slyubomirsky · 2023-06-22T18:39:31Z

Do you mean that private should be another attribute? Or that it should be a field in the function AST node? RIght now, privacy is handled entirely in the parser (if private is False, which is the default, the parser adds a global symbol).

slyubomirsky · 2023-06-22T18:41:35Z

The tough part is dealing with nested functions. There is no way to have a public nested function, but the BlockBuilder/printer/etc would have to know whether we're inside another function or not. (Though it would not be hard to modify the BlockBuilder to track that if that's what we want.)

tqchen · 2023-06-22T20:36:30Z

sorry i am not talking about changing the AST. I only mean the interface

@R.function
def public_func():
    pass
    
    
@R.function(private=True):
def private_func():
    pass

And make it happen also for tir

slyubomirsky · 2023-06-22T20:56:45Z

Then, yes, that is already how I had implemented it. False is the default value for private.

I will investigate doing the same for TIR (I've never touched the TVMScript for TIR).

tqchen · 2023-06-23T14:06:42Z

We can get relax tests in first. Can we detect that printer prints out things correctly (aka skips the global symbol and instead use private annotation)?

slyubomirsky · 2023-06-23T16:55:00Z

Yeah, there are printer tests included. I am making some further tweaks because the implementation right now is rather hacky.

slyubomirsky · 2023-06-23T18:03:01Z

tests/python/relax/test_tvmscript_printer_relax.py

    )


+def test_nested_function():


This actually uncovered a seemingly unrelated bug in the printer, where the nested function would be printed with a suffix after its name (see the changes in src/script/printer/relax/function.cc) because the LHS variable in the VarBinding was treated as already having been encountered. Also fixed.

src/script/printer/utils.h

slyubomirsky · 2023-06-23T20:36:27Z

I will implement the same changes for TIR in a separate PR

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.

Use privacy annotation to decide whether to include global_symbol

1600441

slyubomirsky added the branch: unity label Jun 21, 2023

slyubomirsky requested review from masahi and tqchen June 21, 2023 22:17

slyubomirsky commented Jun 21, 2023

View reviewed changes

Remove unused import

4cd11a2

slyubomirsky commented Jun 21, 2023

View reviewed changes

slyubomirsky added 2 commits June 21, 2023 18:22

Simplify function frame construction

dabcb1e

Remove debug prints

ebd1b80

slyubomirsky commented Jun 21, 2023

View reviewed changes

tests/python/relax/test_tvmscript_printer_relax.py Outdated Show resolved Hide resolved

Whitespace fix

e2ecadc

slyubomirsky commented Jun 22, 2023

View reviewed changes

Include global_symbol in importers

f2c999d

tqchen reviewed Jun 22, 2023

View reviewed changes

Formatting fix

e012683

Add private field to function builder in BlockBuilder

24769c0

Whitespace fix

aaee0cd

Set the global symbol for functions outside a module too

801000b

slyubomirsky added 2 commits June 23, 2023 13:45

Formatting fix

d1f7c01

Print privacy attribute for functions not in an IRModule

5b6c814

slyubomirsky commented Jun 23, 2023

View reviewed changes

src/script/printer/utils.h Outdated Show resolved Hide resolved

slyubomirsky added 3 commits June 23, 2023 16:11

Fix placement of pylint overrides in ir.py

954dfd3

Fix line length

249ec0e

Check for nested function case in function.cc instead

886a63e

slyubomirsky changed the title ~~[Unity][IR][UX] Privacy annotation~~ [Unity][IR][UX] Privacy annotation in Relax Jun 23, 2023

slyubomirsky added 11 commits June 23, 2023 16:45

Print the global symbol if it doesn't match the name for some reason

4b4b52f

Correctly coerce the attribute type

8064cfe

Tweak pylint override again

25e4712

Add pylint override for trailing whitespace in printer tests

8c7bab1

Fix whitespace

b815449

Remove trailing whitespace altogether instead of trying to override it

68142e1

Fix test_utils

4e272ae

Fix test_tvmscript_parser_op_datatype

5314469

Fix global symbols in torch dynamo importer

a0df076

Fix test_dataflow_pattern

d2a4902

Use tvm::attr::kGlobalSymbol instead of string literal

5ac02e8

tqchen approved these changes Jun 27, 2023

View reviewed changes

tqchen merged commit daf9c20 into apache:unity Jun 27, 2023

This was referenced Jun 27, 2023

[Unity][UX][Tweak] Make it an error to mark a function private and specify a global symbol #15170

Merged

[Unity][UX][TIR] Implement privacy annotation for the @prim_func decorator #15171

Closed

[TIR][UX] Implement privacy annotations in TIR #15214

Merged

[Unity][IR][UX] Privacy annotation in Relax #15140

[Unity][IR][UX] Privacy annotation in Relax #15140

Uh oh!

Conversation

slyubomirsky commented Jun 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tvm-bot commented Jun 21, 2023

Uh oh!

tvm-bot commented Jun 21, 2023

Uh oh!

tvm-bot commented Jun 21, 2023

Uh oh!

tvm-bot commented Jun 21, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yelite Jun 27, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tqchen commented Jun 21, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

slyubomirsky Jun 22, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tqchen commented Jun 22, 2023

Uh oh!

slyubomirsky commented Jun 22, 2023

Uh oh!

slyubomirsky commented Jun 22, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tqchen commented Jun 22, 2023

Uh oh!

slyubomirsky commented Jun 22, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tqchen commented Jun 23, 2023

Uh oh!

slyubomirsky commented Jun 23, 2023

Uh oh!

slyubomirsky Jun 23, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

slyubomirsky commented Jun 23, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

slyubomirsky commented Jun 21, 2023 •

edited

Loading

yelite Jun 27, 2023 •

edited

Loading

slyubomirsky Jun 22, 2023 •

edited

Loading

slyubomirsky commented Jun 22, 2023 •

edited

Loading

slyubomirsky commented Jun 22, 2023 •

edited

Loading

slyubomirsky Jun 23, 2023 •

edited

Loading