Thread method-scoped kernel registry through Program and Method (pytorch#19561)

JacobSzwejbka · web-flow · commit a8cfe2b6ce23 · 2026-05-15T22:46:49.000Z
Differential Revision: D98080033 Pull Request resolved: pytorch#19561
diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp
@@ -802,8 +802,19 @@ Error Method::resolve_operator(
   }
 
   // Find a kernel with the matching name and tensor meta.
-  Result<OpFunction> op_function =
-      get_op_function_from_registry(operator_name, {meta, count});
+  // Try method-scoped registry first (if provided), then fall back to global.
+  auto resolve_op_function = [&]() -> Result<OpFunction> {
+    if (!kernel_registry_.empty()) {
+      Result<OpFunction> method_scoped_op_function =
+          get_op_function_from_registry(
+              operator_name, {meta, count}, kernel_registry_);
+      if (method_scoped_op_function.ok()) {
+        return method_scoped_op_function;
+      }
+    }
+    return get_op_function_from_registry(operator_name, {meta, count});
+  };
+  Result<OpFunction> op_function = resolve_op_function();
   if (!op_function.ok()) {
     ET_LOG(
         Error,
@@ -831,7 +842,8 @@ Result<Method> Method::load(
     MemoryManager* memory_manager,
     EventTracer* event_tracer,
     const NamedDataMap* external_data_map,
-    const LoadBackendOptionsMap* backend_options) {
+    const LoadBackendOptionsMap* backend_options,
+    Span<const Kernel> kernel_registry) {
   MemoryAllocator* temp_allocator = memory_manager->temp_allocator();
   if (temp_allocator == nullptr) {
     PlatformMemoryAllocator* platform_allocator =
@@ -844,7 +856,8 @@ Result<Method> Method::load(
     new (platform_allocator) PlatformMemoryAllocator();
     temp_allocator = platform_allocator;
   }
-  Method method(program, memory_manager, event_tracer, temp_allocator);
+  Method method(
+      program, memory_manager, event_tracer, temp_allocator, kernel_registry);
   ET_LOG(Debug, "Loading method: %s.", s_plan->name()->c_str());
   Error err = method.init(s_plan, external_data_map, backend_options);
   if (err != Error::Ok) {
diff --git a/runtime/executor/method.h b/runtime/executor/method.h
@@ -23,6 +23,7 @@
 #include <executorch/runtime/executor/memory_manager.h>
 #include <executorch/runtime/executor/merged_data_map.h>
 #include <executorch/runtime/executor/method_meta.h>
+#include <executorch/runtime/kernel/operator_registry.h>
 #include <executorch/runtime/platform/compiler.h>
 
 // Forward declare flatbuffer types. This is a public header and must not
@@ -82,6 +83,7 @@ class Method final {
         merged_data_map_(std::move(rhs.merged_data_map_)),
         external_constants_(rhs.external_constants_),
         n_external_constants_(rhs.n_external_constants_),
+        kernel_registry_(rhs.kernel_registry_),
         init_state_(rhs.init_state_) {
     // Required: clear out fields that the dtor looks at, so that we don't free
     // anything twice.
@@ -331,7 +333,8 @@ class Method final {
       const Program* program,
       MemoryManager* memory_manager,
       EventTracer* event_tracer,
-      MemoryAllocator* temp_allocator)
+      MemoryAllocator* temp_allocator,
+      Span<const Kernel> kernel_registry = {})
       : step_state_(),
         program_(program),
         memory_manager_(memory_manager),
@@ -348,6 +351,7 @@ class Method final {
         merged_data_map_(nullptr),
         external_constants_(nullptr),
         n_external_constants_(0),
+        kernel_registry_(kernel_registry),
         init_state_(InitializationState::Uninitialized) {}
 
   /// Static factory used by Program.
@@ -357,7 +361,8 @@ class Method final {
       MemoryManager* memory_manager,
       EventTracer* event_tracer,
       const NamedDataMap* named_data_map,
-      const LoadBackendOptionsMap* backend_options = nullptr);
+      const LoadBackendOptionsMap* backend_options = nullptr,
+      Span<const Kernel> kernel_registry = {});
 
   /**
    * Initialize the method from its serialized representation.
@@ -403,6 +408,8 @@ class Method final {
   NamedData* external_constants_;
   size_t n_external_constants_ = 0;
 
+  Span<const Kernel> kernel_registry_;
+
   InitializationState init_state_;
 
   /**
diff --git a/runtime/executor/program.cpp b/runtime/executor/program.cpp
@@ -371,7 +371,8 @@ Result<Method> Program::load_method(
     MemoryManager* memory_manager,
     EventTracer* event_tracer,
     const NamedDataMap* named_data_map,
-    const LoadBackendOptionsMap* backend_options) const {
+    const LoadBackendOptionsMap* backend_options,
+    Span<const Kernel> kernel_registry) const {
   EXECUTORCH_SCOPE_PROF("Program::load_method");
   internal::event_tracer_create_event_block(event_tracer, "Default");
   internal::EventTracerProfileMethodScope event_tracer_scope =
@@ -394,7 +395,8 @@ Result<Method> Program::load_method(
       memory_manager,
       event_tracer,
       named_data_map,
-      backend_options);
+      backend_options,
+      kernel_registry);
 }
 
 Result<MethodMeta> Program::method_meta(const char* method_name) const {
diff --git a/runtime/executor/program.h b/runtime/executor/program.h
@@ -21,6 +21,7 @@
 #include <executorch/runtime/executor/method.h>
 #include <executorch/runtime/executor/method_meta.h>
 #include <executorch/runtime/executor/pte_data_map.h>
+#include <executorch/runtime/kernel/operator_registry.h>
 #include <executorch/runtime/platform/compiler.h>
 
 // Forward declare flatbuffer types. This is a public header and must not
@@ -151,7 +152,8 @@ class Program final {
       MemoryManager* memory_manager,
       EventTracer* event_tracer = nullptr,
       const NamedDataMap* named_data_map = nullptr,
-      const LoadBackendOptionsMap* backend_options = nullptr) const;
+      const LoadBackendOptionsMap* backend_options = nullptr,
+      Span<const Kernel> kernel_registry = {}) const;
 
   /**
    * Gathers metadata for the named method.
diff --git a/runtime/kernel/test/operator_registry_test.cpp b/runtime/kernel/test/operator_registry_test.cpp
@@ -440,6 +440,69 @@ TEST_F(OperatorRegistryTest, GetOpFunctionUsesProvidedKernelList) {
   EXPECT_EQ(run_kernel(*fallback_func), 50);
 }
 
+TEST_F(OperatorRegistryTest, ProvidedKernelListMissCanFallBackToGlobal) {
+  std::array<char, kKernelKeyBufSize> buf{};
+  Error err = make_kernel_key(
+      {{ScalarType::Long, {0, 1, 2, 3}}}, buf.data(), buf.size());
+  ASSERT_EQ(err, Error::Ok);
+  KernelKey long_key = KernelKey(buf.data());
+
+  Kernel global_kernel = Kernel(
+      "test::provided_kernel_list_global_fallback",
+      KernelKey{},
+      [](KernelRuntimeContext& context, Span<EValue*> stack) {
+        (void)context;
+        *(stack[0]) = Scalar(50);
+      });
+  err = register_kernels({&global_kernel, 1});
+  ASSERT_EQ(err, Error::Ok);
+
+  Kernel scoped_kernel = Kernel(
+      "test::provided_kernel_list_global_fallback",
+      long_key,
+      [](KernelRuntimeContext& context, Span<EValue*> stack) {
+        (void)context;
+        *(stack[0]) = Scalar(100);
+      });
+  Span<const Kernel> scoped_registry(&scoped_kernel, 1);
+
+  std::array<Tensor::DimOrderType, 4> dims = {0, 1, 2, 3};
+  auto dim_order_type = Span<Tensor::DimOrderType>(dims.data(), dims.size());
+  std::array<TensorMeta, 1> long_meta = {
+      TensorMeta(ScalarType::Long, dim_order_type)};
+  Span<const TensorMeta> long_kernel_key(long_meta.data(), long_meta.size());
+
+  std::array<TensorMeta, 1> float_meta = {
+      TensorMeta(ScalarType::Float, dim_order_type)};
+  Span<const TensorMeta> float_kernel_key(float_meta.data(), float_meta.size());
+
+  auto run_kernel = [](OpFunction func) {
+    EValue value = Scalar(0);
+    std::array<EValue*, 1> stack = {&value};
+    KernelRuntimeContext context{};
+    func(context, Span<EValue*>(stack.data(), stack.size()));
+    return value.toScalar().to<int64_t>();
+  };
+
+  Result<OpFunction> scoped_func = get_op_function_from_registry(
+      "test::provided_kernel_list_global_fallback",
+      long_kernel_key,
+      scoped_registry);
+  ASSERT_EQ(scoped_func.error(), Error::Ok);
+  EXPECT_EQ(run_kernel(*scoped_func), 100);
+
+  Result<OpFunction> scoped_miss = get_op_function_from_registry(
+      "test::provided_kernel_list_global_fallback",
+      float_kernel_key,
+      scoped_registry);
+  ASSERT_EQ(scoped_miss.error(), Error::OperatorMissing);
+
+  Result<OpFunction> global_func = get_op_function_from_registry(
+      "test::provided_kernel_list_global_fallback", float_kernel_key);
+  ASSERT_EQ(global_func.error(), Error::Ok);
+  EXPECT_EQ(run_kernel(*global_func), 50);
+}
+
 TEST_F(OperatorRegistryTest, DoubleRegisterKernelsDies) {
   std::array<char, kKernelKeyBufSize> buf_long_contiguous;
   Error err = make_kernel_key(
