Mangling automatic differentiation function names

Development Compiler
#1

Automatic differentiation creates new functions behind the scenes, and thus needs to come up with names for those functions. I'm using a hacky suffix right now to unique these functions, but I wonder what's the proper way to build this into the mangler.

Here's an example. When you differentiate a SIL function @foo, AD emits the following functions:

  • The primal function: @foo__primal_src_0_wrt_0_1, meaning "foo's primal function for computing partial derivatives from output 0 with respect to parameters 0 and 1.
  • The adjoint function: @foo__adjoint_src_0_wrt_0_1, meaning "foo's adjoint function for computing partial derivatives from output 0 with respect to parameters 0 and 1.

Above are mostly the source index and differentiation parameter indices. There are also things like gradient configurations

  • _s: seedable
  • _p: preserving result
  • _d: delayed

These hacky suffixes certainly aren't ready for upstreaming AD to the master branch in the future. In practice, they work for most cases, but sometimes interfere with passes like ObjectOutliner, which says

Can't demangle: $S4grad4testy10TensorFlow0C0VySfGAFF__grad_src_0_wrt_0_s_pTv_

Serialization experts, any thoughts on how I should proceed?

#2

cc @Erik_Eckstein , @Michael_Gottesman. The optimizer does something similar, so maybe there's a namespace you can add to? (And support in the tensorflow branch demangler.)

#3

There is a class hierarchy in SpecializationManagler.h that is used for these purposes. I imagine this would be similar to a function signature specialization in a certain sense. I would look at that code. That being said, I defer to @Erik_Eckstein he is the one who touched this code last.

#4

It doesn't need to follow the specialization mangling necessary, but it is generally good to come up with a unique mangling for these sorts of things, since that allows for the definition to be uniqued across translation units, gives you better symbolication in the debugger and other tools, etc. If you don't want to worry about mangling up front and just want to get your code off the ground, the verifier only checks that symbol names beginning with the $S prefix are demanglable. If you don't want the verifier to trip, generate symbol names without the $S prefix, or with something before it.

1 Like