Manual kernel registration to include library names in API

[larryliu0820]

🚀 The feature, motivation and pitch

This request came from @shoumikhin.

For developers on Xcode, they wanted to be able to pick and use kernels_optimized, kernels_quantized and kernels_custom libraries in their application.

The caveat is that since they all use static initialization kernel registration mechanism, all of these libraries requires --force_load linker flag here:

https://github.com/pytorch/executorch/blob/main/examples/demo-apps/apple_ios/LLaMA/LLaMA/SupportingFiles/Release.xcconfig#L8-L14

This causes churns for app developers because they will have to add these linker flags manually and it's not an out of the box experience.

To fix this, we wanted to allow user to pass in a library name to codegen and generate a register_<lib_name>_kernels API.

Currently we trigger selective build using this API:

https://github.com/pytorch/executorch/blob/main/tools/cmake/Codegen.cmake#L63C10-L63C39

We want to add a new argument LIB_NAME and pass the value all the way down into the python script here:

https://github.com/pytorch/executorch/blob/main/tools/cmake/Codegen.cmake#L95

Then inside the python script, add an optional argument to take the lib name and optionally generate register_<lib_name>_kernels API.

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

cc @JacobSzwejbka @lucylq

[shoumikhin]

Thanks for addressing this!

Wdyt if we call the newly exposed symbol as register_kerlens_<name> to match the product naming?
Also, do we need a separate issue for explicit backend registration?

[Tanish2101]

@larryliu0820 can I work on this as it is tagged as good first issue. I am new to executorch this would be the first issue I would take up.

[Tanish2101]

After passing the library name from CMake to the Python codegen script

It will generate a named registration function, like for example: "void register_optimized_kernels()", instead of a generic: void register_kernels();

This allows users to explicitly call the registration function in their code avoiding the need for -force_load.

The generated void register_optimized_kernels() function should be called by the user or the framework otherwise the kernel code won’t be registered.

So my question is should we keep upto user to call the function?
Because if we never the call the function , then the kernels still won’t be registered at runtime.

Option 1: Let the app developers call it

Option 2: Call It Automatically in a Central Place

What should be done if someone can guide me?

@larryliu0820 @shoumikhin

[shoumikhin]

@larryliu0820 do you have any suggestions on how to implement it best?

@Tanish2101 the idea here is to have an option to expose some explicit symbol from the kernel libs, so that if built with that option, users could include some generated header with the symbol declaration and invoke it on-demand in their code to trigger kernels registration manually instead of relying on static initialization that may not be called at all for seemingly unused libs.

[Tanish2101]

@larryliu0820 do you have any suggestions on how to implement it best?

@Tanish2101 the idea here is to have an option to expose some explicit symbol from the kernel libs, so that if built with that option, users could include some generated header with the symbol declaration and invoke it on-demand in their code to trigger kernels registration manually instead of relying on static initialization that may not be called at all for seemingly unused libs.

I’m planning to add a --lib-name option to the codegen script, which will generate:

1. A C++ source file defining void register__kernels(), which calls internal registration functions like registerUnboxedKernels() and registerCustomOps()

2. A matching header file like register__kernels.h exposing the function for manual inclusion

This will allow developers to explicitly include and call register__kernels() in their apps, avoiding reliance on static initialization or -force_load.

Let me know if my approach is correct or if there is any different and better approach let me know.

@larryliu0820 @shoumikhin

[larryliu0820]

I’m planning to add a --lib-name option to the codegen script

Yep that's what it takes to do this. Basically this feature is an extension to the existing --manual-registration option: https://github.com/pytorch/executorch/blob/main/codegen/gen.py#L949-L955, so probably want to make sure that option is turned on.

With --manual-registration we will generate register_all_kernels() API in RegisterKernels.cpp generated source.

We want to change the template of RegisterKernels.cpp to be able to take the library name. Also do this for RegisterKernels.h as well.

Please let me know if you have further questions.