tvm icon indicating copy to clipboard operation
tvm copied to clipboard

Unify name mangling in TVM

Open gigiblender opened this issue 2 years ago • 7 comments

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.

gigiblender avatar Jul 12 '22 13:07 gigiblender

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.

tqchen avatar Jul 12 '22 13:07 tqchen

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

gigiblender avatar Jul 12 '22 16:07 gigiblender

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.

tqchen avatar Jul 12 '22 16:07 tqchen

Built docs for commit 96a7135b3c5d6c9260c303059fc6a71cd728be30 can be found here.

github-actions[bot] avatar Jul 19 '22 12:07 github-actions[bot]

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

areusch avatar Jul 26 '22 01:07 areusch

@gigiblender is this one ready for another look?

areusch avatar Aug 02 '22 18:08 areusch

@gigiblender is this one ready for another look?

Yes. All comments addressed so far.

gigiblender avatar Aug 02 '22 18:08 gigiblender