Unify name mangling in TVM #12066

gigiblender · 2022-07-12T13:31:50Z

This PR is an implementation for this RFC discussion.
Tracking issue: #12181
Forum discussion: https://discuss.tvm.apache.org/t/pre-rfc-name-mangling-in-irmodules/12944/7

It implements a NameSupply and GlobalVarSupply to be used in TVM for obtaining unique names / global vars. It also replaces the previous usages of a name_map and the GetUniqueName method with a NameSupply, or GlobalVarSupply depending on the case.

tqchen

Given that we can recreate GlobalVarSupply from the current IRModule, and we don't need global var renaming for most of the passes(that only do local edits of functions).

It might make sense to create helper method to generate GlobalVarSupply rather than maintaining as a member of IRModule when needed, to keep the main IRModule data structure simple.

Doing so would also simplify the serialization part of IRModule, right now the unordered map data structure is not serialized, as a result save_json/load_json of IRModule will no longer work if we make it as a member.

tqchen · 2022-07-12T13:40:33Z

include/tvm/ir/module.h

@@ -64,6 +65,8 @@ class IRModuleNode : public Object {
  /* \brief Additional attributes storing meta-data about the module. */
  DictAttrs attrs;

+  GlobalVarSupply global_var_supply;


Given that we can recreate GlobalVarSupply from the current IRModule, and we don't need global var renaming for most of the passes(that only do local edits of functions).

It might make sense to create helper method to generate GlobalVarSupply rather than maintaining as a member of IRModule when needed, to keep the main IRModule data structure simple.

Doing so would also simplify the serialization part of IRModule, right now the unordered map data structure is not serialized, as a result save_json/load_json of IRModule will no longer work if we make it as a member.

Given that we can recreate GlobalVarSupply from the current IRModule

I'm not sure this is so trivial. A couple of cases to think through:

right now we still lower per-backend separately, so the NameSupply needs to track names across a collection of IRModules in the compiler

Suppose we needed to reserve an extern name not mentioned in the IR anywhere. We would surely need to serialize something to track that.

The first case here seems to be the one most at tension with the idea of making NameSupply idempotent. I agree that NameSupply isn't necessarily attached to any given IRModule now, so having a way to access it should mainly be a convenience method and not participate in IRModule serialization at least right now. I'm not so sure it means that it's valid to serialize an IRModule and destroy a TVM instance, then load it back in another TVM instance and expect the NameSupply to be the same.

I agree @areusch , both of your pts seems to point that we should bring NameSupply outside IRModule, and have additional functionalities to re-create from multple (collection) of IRModules

tqchen · 2022-07-12T13:54:58Z

Alternatively, we can make name supply an idempotent object, which means it is not part of the serialization and will be nullptr initially. Calling an API IRModule->GetOrCreateNameSupply() will recreate it and store it in the internal private field if it starts as nullptr, and reuse the created object.

Even in this alternative case, we need to work carefully with cases where we copy on write the IRModule. In nature NameSupply is a mutable object while IRModule is moving towards more immutable style. e.g. if we copy an IRModule and mutate it once, we need to understand the API impact on the original IRModule. Allowing recreation of NameSupply from fresh might be the best way to avoid such problems.

So possible options:

F0: IRModule->CreateNameSupply() which recreates NameSupply when needed, always fresh, not all passes need it and minimize IRModule data structure for serialization and other purpoes
F1: IRModule->GetOrCreateNameSupply() initialize NameSupply in first use, get around serialization problem. Need to be careful when we Copy and branch out IRModules (where the new IRModule needs to effectively start from a nullptr again and GetOrCreateNameSupply recreates it lazily).

Example to demonstrate the branching case

def transform_example(moda):
     modb = tvm.transform.SomePass1()(moda)
     modc = tvm.transform.SomePass2()(moda)

In the above example, we can either choose to persist NameSupply during transformation, as a result modb, modc share the same NameSupply, causing bugs since they really are two different pathways. Or we duplicate/recreate the NameSupply in each pass(since we do not know whether or not the NameSupply will be used in another mod), which leads to overhead.

This is of course follows the design rationale of keeping things immutable, in this world, likely likely F0 and F1 will have similar effects and recreation when needed(in few times that it is necessary) may not be a bad idea.

gigiblender · 2022-07-12T16:13:46Z

This is of course follows the design rationale of keeping things immutable, in this world, likely likely F0 and F1 will have similar effects and recreation when needed(in few times that it is necessary) may not be a bad idea.

Thank you @tqchen for the review. I also prefer F0, where we always create for immediate use a fresh NameSupply or GlobalVarSupply from the IRModule.

One possible issue I can see with both F0 and F1 is that usually the module name is used as a prefix to the GlobalVars that belong to the same module. Currently, the module name is not a member of IRModule and is only sometimes passed along (for example in te_compiler.cc LowerTE).
One solution I was thinking about is to add the module name as an attribute to the IRModule. If no such attribute is present, then we can fallback to a default (as is the case currently in te_compiler.cc when a module name is not provided by the user).

tqchen · 2022-07-12T16:29:29Z

Thanks @gigiblender Adding module name to IRModule as an optional attr sounds great. BTW, we still need to distinguish the case where the user expect "the name". This is different from cases when a pass wants to allocate a new name that will be used later, but there is no expectation of external reference(as a result the name can be flexible). This happens in cases where user expects an external function name that they can call during runtime.

In which case NameSupply can not do anything but to check duplication.

github-actions · 2022-07-19T12:31:35Z

Built docs for commit 96a7135 can be found here.

include/tvm/ir/global_var_supply.h

include/tvm/ir/name_supply.h

areusch

thanks @gigiblender, left a round of comments

include/tvm/ir/module.h

include/tvm/ir/name_supply.h

src/target/source/codegen_source_base.cc

src/ir/name_supply.cc

src/relay/backend/te_compiler.cc

tests/cpp/name_supply_test.cc

src/ir/name_supply.cc

areusch · 2022-07-26T01:43:25Z

cc @kparzysz-quic @Mousius @manupa-arm @junrushao1994 @Hzfengsy @tqchen @comaniac please review

kparzysz-quic

I reviewed half of it so far.

Make sure that your algorithm can handle this (pseudo code):

get_new_name("name2");   // creates "name2"
get_new_name("name");    // creates "name"
get_new_name("name");    // creates "name1"
get_new_name("name");    // does not create "name2" again!

src/ir/global_var_supply.cc

src/ir/module.cc

src/ir/name_supply.cc

kparzysz-quic · 2022-07-27T16:34:16Z

src/ir/name_supply.cc

+  for (size_t i = 0; i < prefix.size(); ++i) {
+    if (prefix[i] == '.') prefix[i] = '_';
+  }


Why are you doing this replacement?

I think the dot operator is used to select a member of a class/struct and therefore should be mangled in all names.

src/ir/name_supply.cc

src/relay/backend/graph_executor_codegen.cc

src/relay/backend/te_compiler.cc

areusch · 2022-08-02T18:04:15Z

@gigiblender is this one ready for another look?

gigiblender · 2022-08-02T18:09:58Z

@gigiblender is this one ready for another look?

Yes. All comments addressed so far.

… an IRModule.

* [TVM] Update Submodule * [Compatible] Fix apache/tvm#12066 * [Compatible] Fix apache/tvm#12382 Co-authored-by: SubmoduleUpdaterBot <submodule-updater-bot@users.noreply.github.com> Co-authored-by: Cody Yu <comaniac0422@gmail.com>

* Add NameSupply and GlobalVarSupply * Build GlobalVarSupply from IRModules instead of having it attached to an IRModule. * Pass GlobalVarSupply when lowering shape funcs * Partially replace instantiations of GlobalVar with GlobalVarSupply * Construct GlobalVarSupply from IRModule * Add tests for supply * Add documentation for NameSupply and GlobalVarSupply Co-authored-by: Florin-Gabriel Blanaru <fgb@system76-pc.localdomain>

gigiblender mentioned this pull request Jul 12, 2022

[RFC] Name mangling in IRModules apache/tvm-rfcs#84

Merged

tqchen reviewed Jul 12, 2022

View reviewed changes

gigiblender force-pushed the name-mangling branch from a2fc6db to fb27659 Compare July 12, 2022 16:08

gigiblender force-pushed the name-mangling branch from fb27659 to a7a9278 Compare July 12, 2022 16:17

gigiblender force-pushed the name-mangling branch 3 times, most recently from 4ccc004 to 96a7135 Compare July 19, 2022 09:54

gigiblender force-pushed the name-mangling branch 2 times, most recently from 84ae2ef to 428a598 Compare July 20, 2022 17:19

gigiblender marked this pull request as ready for review July 21, 2022 00:47

gigiblender force-pushed the name-mangling branch from c3182af to 6962c2f Compare July 21, 2022 14:43

kparzysz-quic reviewed Jul 24, 2022

View reviewed changes

include/tvm/ir/global_var_supply.h Show resolved Hide resolved

include/tvm/ir/name_supply.h Show resolved Hide resolved

gigiblender mentioned this pull request Jul 25, 2022

[RFC][Tracking Issue] Name mangling in IRModules #12181

Closed

areusch reviewed Jul 26, 2022

View reviewed changes

gigiblender force-pushed the name-mangling branch 2 times, most recently from e1177da to 30c27af Compare July 27, 2022 15:46

kparzysz-quic reviewed Jul 27, 2022

View reviewed changes

src/relay/backend/graph_executor_codegen.cc Show resolved Hide resolved

src/relay/backend/te_compiler.cc Outdated Show resolved Hide resolved

src/relay/backend/te_compiler.cc Outdated Show resolved Hide resolved

gigiblender force-pushed the name-mangling branch 2 times, most recently from f4ee874 to 25ca250 Compare July 28, 2022 20:34

gigiblender and others added 3 commits August 9, 2022 07:10

Add NameSupply and GlobalVarSupply

a752322

Build GlobalVarSupply from IRModules instead of having it attached to…

11418f2

… an IRModule.

Pass GlobalVarSupply when lowering shape funcs

3ad2d08

Florin-Gabriel Blanaru added 4 commits August 9, 2022 07:10

Partially replace instantiations of GlobalVar with GlobalVarSupply

8593707

Construct GlobalVarSupply from IRModule

b606261

Add tests for supply

806d415

Add documentation for NameSupply and GlobalVarSupply

5df20bb

gigiblender force-pushed the name-mangling branch from 25ca250 to 5df20bb Compare August 9, 2022 14:11

kparzysz-quic approved these changes Aug 11, 2022

View reviewed changes

areusch approved these changes Aug 11, 2022

View reviewed changes

areusch merged commit ecfd969 into apache:main Aug 11, 2022

comaniac added a commit to awslabs/raf that referenced this pull request Aug 15, 2022

[Compatible] Fix apache/tvm#12066

d79d0d9

AndrewZhaoLuo mentioned this pull request Oct 4, 2022

TVM v0.10.0.rc0 Release Candidate Notes #12979

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Unify name mangling in TVM #12066

Unify name mangling in TVM #12066

gigiblender commented Jul 12, 2022 •

edited

Loading

tqchen left a comment

tqchen Jul 12, 2022

areusch Jul 12, 2022

tqchen Jul 12, 2022

tqchen commented Jul 12, 2022 •

edited

Loading

gigiblender commented Jul 12, 2022

tqchen commented Jul 12, 2022 •

edited

Loading

github-actions bot commented Jul 19, 2022

areusch left a comment

areusch commented Jul 26, 2022

kparzysz-quic left a comment

kparzysz-quic Jul 27, 2022

gigiblender Jul 28, 2022

areusch commented Aug 2, 2022

gigiblender commented Aug 2, 2022

Unify name mangling in TVM #12066

Unify name mangling in TVM #12066

Conversation

gigiblender commented Jul 12, 2022 • edited Loading

tqchen left a comment

Choose a reason for hiding this comment

tqchen Jul 12, 2022

Choose a reason for hiding this comment

areusch Jul 12, 2022

Choose a reason for hiding this comment

tqchen Jul 12, 2022

Choose a reason for hiding this comment

tqchen commented Jul 12, 2022 • edited Loading

gigiblender commented Jul 12, 2022

tqchen commented Jul 12, 2022 • edited Loading

github-actions bot commented Jul 19, 2022

areusch left a comment

Choose a reason for hiding this comment

areusch commented Jul 26, 2022

kparzysz-quic left a comment

Choose a reason for hiding this comment

kparzysz-quic Jul 27, 2022

Choose a reason for hiding this comment

gigiblender Jul 28, 2022

Choose a reason for hiding this comment

areusch commented Aug 2, 2022

gigiblender commented Aug 2, 2022

gigiblender commented Jul 12, 2022 •

edited

Loading

tqchen commented Jul 12, 2022 •

edited

Loading

tqchen commented Jul 12, 2022 •

edited

Loading