Expose method name as part of backend init context (#6622)

Summary:


Provide the method name to backend so they can load the corresponding method name accordingly.

The most immediate need is that the qnn context binary can include two methods, one for prefill and one for decode. Since we don't allow backend access multi methods at the moment, we do it in a hacky way via following

## AOT:
```
class LLama_transformer():
      
    def prefill()


    def decode()

```
Then we will have two custom ops from two to_backends ops, and both will have two context binary
```
QAT (prefill) -> to_backend(...) => prefill.qcir flatbuffers
QAT (decode) -> to_backend(...) => decode.qcir flatbuffers

=> 

graph prefill(
    
    custom_op_prefill() -> context_binary (two graphs)

)

graph decode()

    custom_op_decode() -> context_binary (two graphs)
)
```
Since two context binary from these two customs ops will be exactly the same and they can be deduplicate during emit via these two lines https://github.com/pytorch/executorch/blob/d4a9ca01eb5bb786ecbfbcd8302253eb7797e8bb/exir/emit/_emitter.py#L136 and here https://github.com/pytorch/executorch/blob/d4a9ca01eb5bb786ecbfbcd8302253eb7797e8bb/exir/emit/_emitter.py#L1065-L1066
```
.pte instrucions
[
"prefill" [instructions: call_delegate(prefill_input)]

"decode": [instructions: call_delegate(decode_input)]

"delegate_payload:: Dict[bytes, index])

]
```
## Runtime
After we expose the method name via this change, the backend can access the method name, and load the same method as the top level method 
```
Result<DelegateHandle*> QNNBackend::init(
      BackendInitContext& context,
      FreeableBuffer* processed,
      ArrayRef<CompileSpec> compile_specs) {
     const char* method_name = context.get_method_name() // for example, "prefill"
     handle = qnn_backend.load(method_name)
     return handle
}
```

This is to unblock sharing weight between prefill and decode for using htp backend.

Reviewed By: kimishpatel, iseeyuan

Differential Revision: D65386597

Author: cccclai

runtime/backend/backend_execution_context.h

@@ -21,8 +21,9 @@ class BackendExecutionContext final {
  public:
   BackendExecutionContext(
       EventTracer* event_tracer = nullptr,
-      MemoryAllocator* temp_allocator = nullptr)
-      : event_tracer_(event_tracer), temp_allocator_(temp_allocator) {}
+      MemoryAllocator* temp_allocator = nullptr,
+      const char* method_name = nullptr)
+      : event_tracer_(event_tracer), temp_allocator_(temp_allocator), method_name_(method_name) {}
 
   /**
    * Returns a pointer to an instance of EventTracer to do profiling/debugging
@@ -52,9 +53,17 @@ class BackendExecutionContext final {
     return temp_allocator_;
   }
 
+  /** 
+   * Get the loaded method name from ExecuTorch runtime.
+   */
+  const char* get_method_name() {
+    return method_name_;
+  }
+
  private:
   EventTracer* event_tracer_ = nullptr;
   MemoryAllocator* temp_allocator_ = nullptr;
+  const char* method_name_ = nullptr;
 };
 
 } // namespace runtime

runtime/backend/backend_init_context.h

@@ -18,8 +18,8 @@ namespace runtime {
  */
 class BackendInitContext final {
  public:
-  explicit BackendInitContext(MemoryAllocator* runtime_allocator)
-      : runtime_allocator_(runtime_allocator) {}
+  explicit BackendInitContext(MemoryAllocator* runtime_allocator, const char* method_name = nullptr)
+      : runtime_allocator_(runtime_allocator), method_name_(method_name) {}
 
   /** Get the runtime allocator passed from Method. It's the same runtime
    * executor used by the standard executor runtime and the life span is the
@@ -28,9 +28,21 @@ class BackendInitContext final {
   MemoryAllocator* get_runtime_allocator() {
     return runtime_allocator_;
   }
+  
+  /** Get the loaded method name from ExecuTorch runtime. Usually it's "forward",
+   * however, if there are multiple methods in the .pte file, it can be different. 
+   * One example is that we may have prefill and decode methods in the same .pte file.
+   * In this case, when client loads "prefill" method, the `get_method_name` function will 
+   * return "prefill", when client loads "decode" method, the `get_method_name` function will 
+   * return "decode".
+   */
+  const char* get_method_name() {
+    return method_name_;
+  }
 
  private:
   MemoryAllocator* runtime_allocator_ = nullptr;
+  const char* method_name_ = nullptr;
 };
 
 } // namespace runtime

runtime/executor/method.cpp

@@ -598,6 +598,7 @@ Error Method::init(executorch_flatbuffer::ExecutionPlan* s_plan) {
   init_state_ =
       InitializationState::InitializationFailed; // Until proven otherwise
   serialization_plan_ = s_plan;
+  method_name_ = s_plan->name()->str().c_str();
   auto method_allocator = memory_manager_->method_allocator();
 
   {
@@ -626,7 +627,7 @@ Error Method::init(executorch_flatbuffer::ExecutionPlan* s_plan) {
 
     for (size_t i = 0; i < n_delegate; ++i) {
       const auto& delegate = *delegates->Get(i);
-      BackendInitContext backend_init_context(method_allocator);
+      BackendInitContext backend_init_context(method_allocator, method_name_);
       Error err = BackendDelegate::Init(
           delegate, program_, backend_init_context, &delegates_[i]);
       if (err != Error::Ok) {
@@ -1098,7 +1099,8 @@ Error Method::execute_instruction() {
           step_state_.instr_idx);
       BackendExecutionContext backend_execution_context(
           /*event_tracer*/ event_tracer_,
-          /*temp_allocator*/ temp_allocator_);
+          /*temp_allocator*/ temp_allocator_,
+          /*method_name_*/ method_name_);
       err = delegates_[delegate_idx].Execute(
           backend_execution_context,
           chain.argument_lists_[step_state_.instr_idx].data());

runtime/executor/method.h

@@ -328,6 +328,7 @@ class Method final {
 
   size_t n_chains_;
   Chain* chains_;
+  const char* method_name_;
 
   InitializationState init_state_;
 

runtime/executor/test/backend_integration_test.cpp

@@ -56,7 +56,7 @@ class StubBackend final : public BackendInterface {
       FreeableBuffer*,
       ArrayRef<CompileSpec>,
       BackendInitContext&)>;
-  using ExecuteFn = std::function<Error(DelegateHandle*, EValue**)>;
+  using ExecuteFn = std::function<Error(DelegateHandle*, EValue**, BackendExecutionContext&)>;
   using DestroyFn = std::function<void(DelegateHandle*)>;
 
   // Default name that this backend is registered as.
@@ -98,7 +98,7 @@ class StubBackend final : public BackendInterface {
       DelegateHandle* handle,
       EValue** args) const override {
     if (execute_fn_) {
-      return execute_fn_.value()(handle, args);
+      return execute_fn_.value()(handle, args, context);
     }
     // Return a benign value otherwise.
     return Error::Ok;
@@ -404,7 +404,7 @@ TEST_P(BackendIntegrationTest, EndToEndTestWithProcessedAsHandle) {
   // FreeableBuffer.
   DelegateHandle* execute_handle = nullptr;
   StubBackend::singleton().install_execute(
-      [&](DelegateHandle* handle, ET_UNUSED EValue** args) -> Error {
+      [&](DelegateHandle* handle, ET_UNUSED EValue** args, ET_UNUSED BackendExecutionContext& context) -> Error {
         execute_handle = handle;
         auto* processed = reinterpret_cast<FreeableBuffer*>(handle);
 
@@ -527,6 +527,40 @@ TEST_P(BackendIntegrationTest, SegmentInfoIsPassedIntoDataLoader) {
   EXPECT_EQ(backend_load_was_called, using_segments());
 }
 
+TEST_P(BackendIntegrationTest, GetMethodNameDuringInitSuccess) {
+  Result<FileDataLoader> loader = FileDataLoader::from(program_path());
+  ASSERT_EQ(loader.error(), Error::Ok);
+  const void* processed_data = nullptr;
+  StubBackend::singleton().install_init(
+      [&](FreeableBuffer* processed,
+          ET_UNUSED ArrayRef<CompileSpec> compile_specs,
+          ET_UNUSED BackendInitContext backend_init_context)
+          -> Result<DelegateHandle*> {
+        auto method_name = backend_init_context.get_method_name();
+        EXPECT_EQ(method_name, "forward");
+        processed_data = processed->data();
+        return nullptr;
+      });
+  Result<Program> program = Program::load(&loader.get());
+  ASSERT_EQ(program.error(), Error::Ok);
+
+}
+
+TEST_P(BackendIntegrationTest, GetMethodNameDuringExecuteSuccess) {
+  Result<FileDataLoader> loader = FileDataLoader::from(program_path());
+  ASSERT_EQ(loader.error(), Error::Ok);
+  DelegateHandle* execute_handle = nullptr;
+  StubBackend::singleton().install_execute(
+      [&](DelegateHandle* handle, ET_UNUSED EValue** args, ET_UNUSED BackendExecutionContext& backend_execution_context)-> Error {
+        auto method_name = backend_execution_context.get_method_name();
+        EXPECT_EQ(method_name, "forward");
+        return Error::Ok;
+      });
+  Result<Program> program = Program::load(&loader.get());
+  ASSERT_EQ(program.error(), Error::Ok);
+
+}
+
 // TODO: Add more tests for the runtime-to-backend interface. E.g.:
 // - Errors during init() or execute() result in runtime init/execution failures
 // - Correct values are passed to init()/execute()