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.

Differential Revision: D65386597

Author: cccclai

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) {

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

@@ -528,6 +528,25 @@ TEST_P(BackendIntegrationTest, SegmentInfoIsPassedIntoDataLoader) {
   EXPECT_EQ(backend_load_was_called, using_segments());
 }
 
+TEST_P(BackendIntegrationTest, GetMethodNameSuccess) {
+  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);
+
+}
+
 // 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()