[pre-RFC] Name mangling in IRModules

Development pre-RFC
#1

Hi all,

In [mini-RFC] Name mangling in AOT, we discussed ways to accommodate the runtime::Module tree in an embedded environment, where we prefer to call functions directly in the generated C using their symbol name rather than looking them up at runtime via a string dictionary. Mangling was needed to namespace these functions.

I think initially we viewed mangling as necessary for the top-level function names; therefore, we modified TECompilerImpl::Lower to produce PrimFuncs that contained module names.

Since then, we’ve come to a number of challenges around mangled names:

  • tir.Constant may be declared as global variables by codegen; but then these names should really be mangled.
  • We introduced a new MetadataModule which defined a PrimFunc outside of TECompilerImpl::Lower which we use to support the PackedFunc-based AOTExecutor. Because this function was defined separately, its name wasn’t mangled. This broke the multiple-model-in-the-same program use case. We had to work around this by conditionally mangling the function name depending which runtime was in use.

Recall that mangling is needed in cases where two models are compiled into the same program–this turns out to surface itself mainly in the microTVM/C runtime use case, where we are not loading shared libraries. We were able to implement AOTExecutor in the C++ runtime because there, modules are compiled into .so and there is a mechanism to de-conflict symbol names which are duplicated across .so. In the C runtime we are not so lucky–all functions must be prefixed/mangled to avoid conflict.

I propose we push all mangling out of the central compiler and into the codegen layer. This means that codegen would need to rewrite references to unmangled names as necessary–e.g. to tir.Constant, to tir.Call, and e.g. when generating FuncRegistry. This is somewhat burdensome on the C and LLVM codegens, but there, mangling is truly a property of the runtime environment and not necessarily something that needs to live in the central compiler.

Would love to hear everyone’s thoughts!

cc @manupa-arm @Mousius @alanmacd @mbs-octoml @tqchen @junrushao @jroesch

Andrew

#2

This is somewhat related to linkage itself.

Specifically, depending on the progress of lowering, there can be a need to decide on name earlier if some of the external codegen involved(e.g. BYOC that pre-decides on “global_symbol” of the name, or user specifies that they would like to “lock” the name so they can be accessed from outside).

Aka we cannot force all he naming being deferred. This being said, we should have a clear critieria to tell that what names are “internal”, as a result can be changed, and what are “external”, as a result need to be fixed due to prior passes and agreements.

#3

One related thing was that we intended to consolidate all of the different functions into one IRModule (e.g. those passed to external codegen with those passed to ordinary TVM codegen). If we had done that, we could probably do the following:

  • If a Relay-to-Runtime hook is used, find a way to lock the name or provide a way for that codegen to mangle the function name internally.
  • If a TIR-to-Runtime hook is used, just mangle the names ourselves right before codegen.

We might still be able to get away with applying a mangle pass to a collection of IRModules if we did it at a single point.

#4

Agree that mangling should be deferred as long as possible. I think the keys is not to push all mangling to a certain stage, but have a clear spec about what can be mangled, and what cannot

#5

i agree with that direction, but I think to do that requires us to merge the IRModules so that mangling transforms are applied globally rather than to just a part of the program. do you agree?

#6

Other than the plumbing, is there an issue with threading a name supply so that globals have a unique and appropriately hinted name at birth? It’s not too hard to support name supply splitting such that names can be drawn from independent supplies without collision. It is also possible to refine the hints to, eg, include prefixes hinting at the backend or other context which can help keep the names a little self-describing.

(While doing this we could also ensure names begin in GlobalVar form instead of String form, since the latter causes multiple GlobalVars to be created with the same name hint but different identity, see our friend te_compiler.cc for hackery to work around that.)

I ended up needing a name supply in collage btw: https://github.com/mbs-octoml/mbs-tvm/blob/40b834911420c3c453e37cfe286b4c2c3e74b2ac/src/relay/collage/name_supply.h#L36

#7

Indeed the name_hint field in GlobalVar is supposed to work for that purpose(a hint rather than a fixed name). In TIR and likely broadly, we use global_symbol attr to indicate the externally enforced name(which can not change) and that name can be referenced externally and we cannot change it during compilation.

Followup on @areusch 's point merge the IRModules so that mangling transforms are applied globally

Agreed, alternatively, we still allow user to separate modules, where cross module interaction are done through a fixed global_symbol, which is less ideal but needed sometime. The encouraged path is to always hold as much as possible internally.

#8

One question in my mind is: there are some names which must be stable e.g. __main__, get_c_metadata. I presume there would be a way to tell the name supply about this?

How would we enforce this fixed global symbol? Also from name supply?

#9

How would we enforce this fixed global symbol? Also from name supply?

When a function need to fix a global name, it will attach an attr “global_symbol” with the name. That attr can be used say by BYOC to pick a name, attach the global symbol, so followup passes respect that constraint.

We can have a stable set of reserved names.

#10

Name supplies usually have both a cache lookup:

name_supply.UniqueGlobalFor("__main__")

and a hinted fresh name generator:

name_supply.FreshGlobalWithPrefix("my_module", "my_var", "any_other_prefix")
#11

Thanks for raising this @areusch!

Just to clarify the two approaches in my mind here.

RemoveNameCollisions Pass

A Pass which looks for variables with matching name_hints within a module and updates it to remove collisions with any other names in the module. With the additional rule that a global_symbol which strictly enforce the name, un-resolvable collisions would be an error.

NameSupply Generator

This is similar to our current GetUniqueName function:

std::string GetUniqueName(std::string name, std::unordered_map<std::string, int>* name_map_);

Except we use it to maintain not only indexes of unique names, but also hold some context of the module and sub-structures. In the case of CMSIS-NN, this means all variables could likely be scoped neatly:

name_supply.FreshGlobalWithPrefix(“ultimate_cat_spotter”, “filter”, “cmsisnn”)

Thoughts

The nice thing about the Pass approach is that it is essentially magic, a developer doesn’t have to sprinkle GetUniqueName everywhere as just defining a Var does that for them. The downside is that the meaning of the names and the rules to generate them are essentially down to whomever defines the Var so you will likely end up with more nonsensical names? (“constant_0_1_2”, as we keep generating “constant”s for things).

@mbs-octoml, with the NameSupply, we could create a sub-NameSupply correct? Something like:

std::string my_module_name = “ultimate_cat_spotter”;
NameSupply module_name_supply(my_module_name);
module_name_supply.FreshGlobal(“constant”); // “tvmgen_ultimate_cat_spotter_constant_0”

NameSupply sub_module_name_supply = module_name_supply.WithPrefix(“sub_module”);
// Send that to some sub-process to generate things
sub_module_name_supply.FreshGlobal(“constant”); // “tvmgen_ultimate_cat_spotter_sub_module_constant_0”

That’d mean for BYOC, we pass it a module and a name supply to use, and we can hierarchically define a prefix? Potentially UniqueGlobal would instead revert to the central prefix for shared variables.

As far as I can tell, the main difference here is that the NameSupply approach would give us nicer outputs, with a bit more developer overhead and the Pass would conversely simplify the developer experience at the expense of the eventual output?

#12

we could create a sub-NameSupply correct?

Yes, that’s the way it’s usually done, both to ‘refine’ the supply of names (eg for a particular external codegen), and to support nested scopes (ideally Vars would have a unique name at birth so that we can match them up when debugging).

Pass would conversely simplify the developer

I think there’s two things to check: i) please take a look at the hackery in te_compiler which tries to resestablish unique GlobalVars despite names being generated via strings. Any change we make in this space should repair that mess. ii) We have both strings and GlobalVar forms of names kicking about in attributes as well in Exprs. So any mangling would need to account for those.

Thanks, -m

#13

Is the idea here to use pointer equality to determine when two Var with the same name reference the same logical thing? not sure how we differentiate between uses across the IRModule otherwise.

#14

I think that’s the only way this could work unless I’m missing something? Even within an IRModule you’d need to do pointer comparison so you can have Var('constant') and Var('constant') resolve to constant_0 and constant_1.

One of the goals should be to strip te_compiler of all of this logic in favour of something smarter dealing with it, I think there’s more than just te_compiler that tries to figure out GlobalVar from String and pointer value at the same time and it really shouldn’t exist in any of these places imho? Does that agree with your view @mbs-octoml ?

Are these at all consistent in the attrs? It’d be good to see examples but I get the impression it may just be a free-for-all? Unsure quite how to deal with it if people hide all the values in different attrs for the standalone Pass, unless we can have a more stable representation in the IR?

#15

something smarter dealing with it,

Amen to that, in this case the smarter thing IMHO is to push the NameSupply down so that all global definitions are assigned a unique GlobalVar object at birth.

It’d be good to see examples but I get the impression it may just be a free-for-all

Our friend te_compiler is again an example of this ugliness since it stores GlobalVars recording the connections between Relay functions, their lowered PrimFunc and any supporting PrimFuncs, and ditto but for the dynamic shape PrimFunc. Better would be to put those in an official structure, but nevertheless it is an example of global vars hiding outside of the usual definitional and referential places. Simply ensuring GlobalVars are in their final form at birth would mean this is all just fine.