Add unique method ID (#12692)

Summary:

Add an auto-incremented identifier field to the Method class and pass into BackendInitContext. This allows backends to disambiguated between methods with the same name (two models with a "forward" method, for example).

Differential Revision: D78697798

Author: GregoryComer

runtime/backend/backend_init_context.h

@@ -21,6 +21,7 @@ class BackendInitContext final {
  public:
   explicit BackendInitContext(
       MemoryAllocator* runtime_allocator,
+      size_t method_id,
       EventTracer* event_tracer = nullptr,
       const char* method_name = nullptr,
       const NamedDataMap* named_data_map = nullptr)
@@ -31,6 +32,7 @@ class BackendInitContext final {
         event_tracer_(nullptr),
 #endif
         method_name_(method_name),
+        method_id_(method_id),
         named_data_map_(named_data_map) {
   }
 
@@ -61,6 +63,14 @@ class BackendInitContext final {
   const char* get_method_name() const {
     return method_name_;
   }
+  
+  /** Get the method identifier of the loaded method. This corresponds to the
+   * id() method on the Method object, and can be used to disambiguate between
+   * different methods with the same name.
+   */
+  size_t method_id() const {
+    return method_id_;
+  }
 
   /** Get the named data map from ExecuTorch runtime.
    * This provides a way for backends to retrieve data blobs by key.
@@ -73,6 +83,7 @@ class BackendInitContext final {
   MemoryAllocator* runtime_allocator_ = nullptr;
   EventTracer* event_tracer_ = nullptr;
   const char* method_name_ = nullptr;
+  size_t method_id_ = 0;
   const NamedDataMap* named_data_map_ = nullptr;
 };
 

runtime/executor/method.cpp

@@ -293,6 +293,9 @@ Result<bool> parse_cond_value(const EValue& cond_value) {
 
 } // namespace
 
+// Initialize static method ID counter.
+size_t Method::next_id_ = 0;
+
 Result<size_t> Method::get_num_external_constants() {
   auto flatbuffer_values = serialization_plan_->values();
   size_t n_value = flatbuffer_values->size();
@@ -842,6 +845,7 @@ Error Method::init(
       const auto& delegate = *delegates->Get(i);
       BackendInitContext backend_init_context(
           method_allocator,
+          /*method_id=*/id(),
           /*event_tracer=*/event_tracer_,
           /*method_name=*/serialization_plan_->name()->c_str(),
           /*named_data_map=*/named_data_map);

runtime/executor/method.h

@@ -65,7 +65,8 @@ class Method final {
    * and leaves `rhs` in an uninitialized state.
    */
   Method(Method&& rhs) noexcept
-      : step_state_(rhs.step_state_),
+      : id_(rhs.id_),
+        step_state_(rhs.step_state_),
         program_(rhs.program_),
         memory_manager_(rhs.memory_manager_),
         temp_allocator_(rhs.temp_allocator_),
@@ -103,6 +104,18 @@ class Method final {
     rhs.n_chains_ = 0;
     rhs.chains_ = nullptr;
   }
+  
+  /**
+   * Retrieve the unique identifier for this method. This is guaranteed to be
+   * unique among all methods loaded by the instance of the runtime. Note that
+   * this value is only meaningful in the context of a specific instance of the
+   * executorch runtime and will vary run to run.
+   * 
+   * @returns The unique identifier for this method instance.
+   */
+  size_t id() const {
+    return id_;
+  }
 
   /**
    * Sets the internal input value to be equivalent to the to the provided
@@ -306,7 +319,8 @@ class Method final {
       MemoryManager* memory_manager,
       EventTracer* event_tracer,
       MemoryAllocator* temp_allocator)
-      : step_state_(),
+      : id_(next_id_++),
+        step_state_(),
         program_(program),
         memory_manager_(memory_manager),
         temp_allocator_(temp_allocator),
@@ -352,7 +366,10 @@ class Method final {
 
   // Executes a single instruction using the state in step_state_
   ET_NODISCARD Error execute_instruction();
+  
+  static size_t next_id_;
 
+  size_t id_;
   StepState step_state_;
   const Program* program_;
   MemoryManager* memory_manager_;

runtime/executor/test/method_test.cpp

@@ -365,6 +365,27 @@ TEST_F(MethodTest, MethodGetAttributeTest) {
   EXPECT_EQ(res->const_data_ptr<int32_t>()[0], 1);
 }
 
+TEST_F(MethodTest, MethodIdTest) {
+  // Verify that methods are assigned unique ids. Load three methods and verify
+  // that they have different ids.
+  ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
+
+  Result<Method> add_method = programs_["add"]->load_method("forward", &mmm.get());
+  ASSERT_EQ(add_method.error(), Error::Ok);
+
+  Result<Method> add_mul_method =
+      programs_["add_mul"]->load_method("forward", &mmm.get());
+  ASSERT_EQ(add_mul_method.error(), Error::Ok);
+
+  Result<Method> stateful_method =
+      programs_["stateful"]->load_method("forward", &mmm.get());
+  ASSERT_EQ(stateful_method.error(), Error::Ok);
+  
+  EXPECT_NE(add_method->id(), add_mul_method->id());
+  EXPECT_NE(add_method->id(), stateful_method->id());
+  EXPECT_NE(add_mul_method->id(), stateful_method->id());
+}
+
 /*
  * TODO(T161163608): Test is disabled due to a resize bug in tensor_index_out of
  * the portable op lib