are all the changes here part of this PR? (or is there any chance you can split the eager changes into their own PR for review?)

sorry, i have been exporting these PRs from a stack in phabricator, and I didn't realize that the way github would render them is unusable. If you click on just the top commit shown here, it will show you the small set of changes that actually goes with this PR.

"external backends" -> "everything that isn't LTC"? (Technically I think this applies both our eager backends in-tree like CPU/CUDA, and also any other backends out-of-tree like Intel XPU)

yes, it was meant as in external LTC backends such as one developed out of tree. I should probably update it to say more clearly that only the in-tree lazy torchscript backend gets special behavior, and others don't

I'm a little confused by this split between DispatchKeyNativeFunctions.h and RegisterDispatchKey.cpp.

Don't we always end up #includeing DispatchKeyNativeFunctions.h in this file, so the Register${BackendName}${DispatchKey}Modules function will already be declared? (and the macro above will already be defined).

Maybe in practice this might be true. But its not explicitly clear just from the template definition that this is the case. Which is why I added the #ifndef clause. If it happens that there is a case (either now or in the future) where the DispatchKeyNativeFunctions.h is not included, it would avoid a compile error here and revert back to the default behaviour.

I think the "delay registration using a lambda" feels like a niche enough use case to me that we should try to implement it with less clutter to this file. Instead of using the macros, how do you feel about unconditionally generating the Register{...}Modules lambda? And still using ${eager_registration} to decide in the codegen whether or not we call it directly in the file or not (to run the registrations at load time).

The initTSBackend() function wouldn't need to check a macro - it would just always call the lambda to run the registrations. You might run into problems if we got the codegen wrong and accidentally registered the TS backend on startup, but you could also just add a unit test for this (assert in one of the C++ tests that there aren't any kernels registered to the Lazy key on startup).

I think that might not be a bad idea. But perhaps better left to a follow PR seeing as this is just a port of a previously merged PR?

Perhaps what could be done is if the Register{...}Modules function is called at c++ static initialization, we could also reset the value to nullptr so it can't be called again. Then we can just run the function only if it hasn't been called before. This would eliminate the need of the EAGER_REGISTRATION macro. Thoughts?

yep, that sounds reasonable.

I kind of don't want to clutter up this file with the extra macros, mostly because this is a pretty core template used in a bunch of different contexts (used for most of our eager mode dispatcher registrations too) and so anyone looking at this file for non-LTC purposes would now have to reason about the macros.

@wconstab would you rather make changes to this PR directly, or address feedback later?

I think we can make the changes here. I already landed this as-is on the staging branch, but we are almost at the point of abandoning the staging branch so I think we can start to do things differently here instead of striving to stay in sync there.

It also depends how we make the changes; if you have some clear ideas of what i should change I can just make the changes. If you actually want to try to modify it yourself, I could either land it first and you change it, or maybe you could commandeer this diff.

I think what @antoniojkim put sounds reasonable:

(1) kill the macro and unconditionally generate the Register${BackendName}${DispatchKey}Modules();

(2) In the codegen, if the eager_registration argument is true, then emit the code to run the function immediately (which this PR already does), and then also set it to re-assign it to a lambda that no-ops (that way you won't accidentally call it and register things multiple times - although I think the risk of this being a problem in practice is pretty small, since nobody should ever call this function except for the torchscript one. So if this is hard to do then I think it's ok if you don't add it).

technically this should always be set to 1 for TS, right? Maybe we should just assert it.

Yea, I think that's true.

This might just be from a different PR, but I think we technically don't need BackendName here? (the function name will be unique without it). Plumbing a new BackendName argument through the codegen for the name seems like overkill to me, but maybe there are other places in the codegen where BackendName actually is required.

@bdhirsh This is something I added to differentiate between backends. There is no guarantee as far as I am aware as to the order in which the symbols are statically initialized. This means that if we had just a RegisterLazyModules function being defined by both the TS backend and a custom backend, then there is no way of knowing who assigned the value of that function pointer last. To avoid ambiguity, I plumbed through the backend name to make things clear and unambiguous

yeah good point, I can't think of an easy way to get around this. (cc @ezyang if you happen to have a different opinion. Context = this PR updates the codegen'd dispatcher registrations to be able to be delayed, and registered through a lambda instead of at load time).

@wconstab I think there's a problem here. According to the function definition of gen_dispatchkey_nativefunc_headers, the positional argument after autograd_dispatch_key is eager_registration, but instead you passed in backend_name. Since you also set the value of eager_registration using a keyword, the following error occurs:

TypeError: gen_dispatchkey_nativefunc_headers() got multiple values for argument 'eager_registration'

This is pretty invasive, I'm not sure I like it.

nit: you can probably just use one argument here instead of two in the template file (dispatch_registrations), since right now one of these two arguments is always an empty string.

but I need to put the content in 2 different places, one place is inside anonymous namespace and the other is inside at:: namespace. Maybe I'm missing something?

ah yep, i forgot about that. makes sense!

minor nit: most of our codegen uses f-strings now instead of CodeTemplate, e.g:

deferred_dispatch_registrations = f"""\
TORCH_API void Register{backend_name}{dispatch_key}NativeFunctions() {{
...

is it bad to use CodeTemplate though? I started off using f-strings, but since I am substituting a multi-line string (with indentation) CodeTemplate actually provides value here.

I think f-strings are pretty much superior to CodeTemplate (you can do multi-line f-strings just fine, like in the codegen here). It's not too big of a deal though so i don't wanna get hung up on it.

It looks like ASAN isn't happy (the ASAN test is complaining about initialization-order-fiasco: https://github.com/pytorch/pytorch/runs/5746632034?check_suite_focus=true)

I think the ordering problem might be between this static bool (which eventually calls RegisterTorchScriptLazyNativeFunctions), and the static library object that that function uses in the codegen: static auto m = MAKE_TORCH_LIBRARY_IMPL(aten, $dispatch_key);.

You might need to either access m through a function, or maybe not make this bool static?