Thread kernel_registry through Module::load_method (#19641)

Summary:
Pull Request resolved: #19641

D98080033 added a method-scoped kernel registry to Program::load_method and Method, allowing callers to override specific kernels for a single method without affecting the global registry. However, the Module facade class did not expose this parameter, forcing consumers to bypass Module and manage memory manually.

This adds an optional `Span<const Kernel> kernel_registry` parameter (defaulting to empty) to Module::load_method and Module::load_forward, and forwards it to Program::load_method. Existing callers are completely unaffected — the default empty span causes the runtime to fall back to the global kernel registry, exactly as before.

Differential Revision: D104433196

Author: wliuyx

extension/module/module.cpp

@@ -20,6 +20,7 @@ namespace extension {
 namespace ET_MODULE_NAMESPACE {
 
 using ET_MERGED_DATA_MAP_NAMESPACE::MergedDataMap;
+using ET_RUNTIME_NAMESPACE::Kernel;
 using ET_RUNTIME_NAMESPACE::MethodMeta;
 using ET_RUNTIME_NAMESPACE::Program;
 
@@ -365,7 +366,8 @@ runtime::Error Module::load_method(
     const std::string& method_name,
     runtime::HierarchicalAllocator* planned_memory,
     torch::executor::EventTracer* event_tracer,
-    const LoadBackendOptionsMap* backend_options) {
+    const LoadBackendOptionsMap* backend_options,
+    std::vector<Kernel> kernel_registry) {
   if (!is_method_loaded(method_name)) {
     ET_CHECK_OK_OR_RETURN_ERROR(load());
 
@@ -402,12 +404,16 @@ runtime::Error Module::load_method(
 
     method_holder.memory_manager = std::make_unique<runtime::MemoryManager>(
         memory_allocator_.get(), planned_memory, temp_allocator_.get());
+    method_holder.kernel_registry = std::move(kernel_registry);
     auto res_method = program_->load_method(
         method_name.c_str(),
         method_holder.memory_manager.get(),
         event_tracer ? event_tracer : this->event_tracer(),
         merged_data_map_.get(),
-        effective_backend_options);
+        effective_backend_options,
+        runtime::Span<const Kernel>(
+            method_holder.kernel_registry.data(),
+            method_holder.kernel_registry.size()));
     if (!res_method.ok()) {
       return res_method.error();
     }

extension/module/module.h

@@ -25,6 +25,7 @@
 namespace executorch {
 namespace extension {
 
+using ET_RUNTIME_NAMESPACE::Kernel;
 using ET_RUNTIME_NAMESPACE::Method;
 using ET_RUNTIME_NAMESPACE::MethodMeta;
 using ET_RUNTIME_NAMESPACE::NamedDataMap;
@@ -255,7 +256,8 @@ class Module {
       const std::string& method_name,
       runtime::HierarchicalAllocator* planned_memory = nullptr,
       torch::executor::EventTracer* event_tracer = nullptr,
-      const LoadBackendOptionsMap* backend_options = nullptr);
+      const LoadBackendOptionsMap* backend_options = nullptr,
+      std::vector<Kernel> kernel_registry = {});
 
   ET_DEPRECATED ET_NODISCARD runtime::Error inline load_method(
       const std::string& method_name,
@@ -303,9 +305,11 @@ class Module {
   ET_NODISCARD inline runtime::Error load_forward(
       runtime::HierarchicalAllocator* planned_memory = nullptr,
       torch::executor::EventTracer* event_tracer = nullptr,
-      const LoadBackendOptionsMap* backend_options = nullptr) {
+      const LoadBackendOptionsMap* backend_options = nullptr,
+      std::vector<Kernel> kernel_registry = {}) {
     return load_method(
-        "forward", planned_memory, event_tracer, backend_options);
+        "forward", planned_memory, event_tracer, backend_options,
+        std::move(kernel_registry));
   }
 
   ET_DEPRECATED ET_NODISCARD inline runtime::Error load_forward(
@@ -698,6 +702,7 @@ class Module {
     std::unique_ptr<PlannedMemory> planned_memory;
     std::unique_ptr<runtime::MemoryManager> memory_manager;
     std::unique_ptr<Method> method;
+    std::vector<Kernel> kernel_registry;
   };
 
   std::string file_path_;