integrate pytorch#6622, will drop it once merged

Author: haowhsu-quic

backends/qualcomm/runtime/QnnExecuTorchBackend.cpp

@@ -71,6 +71,7 @@ Result<DelegateHandle*> QnnExecuTorchBackend::init(
           "Fail to allocate tensor");
     }
   }
+  method_name_ = context.get_method_name();
   return qnn_manager;
 }
 
@@ -79,13 +80,10 @@ Error QnnExecuTorchBackend::execute(
     DelegateHandle* handle,
     EValue** args) const {
   QnnManager* qnn_manager = static_cast<QnnManager*>(handle);
-  // change this after BackendExecutionContext change lands
-  auto graph_name = qnn_manager->GetGraphNames()[0];
-
   std::vector<std::shared_ptr<TensorWrapper>> input_tensors =
-      qnn_manager->GetGraphInputs(graph_name);
+      qnn_manager->GetGraphInputs(method_name_);
   std::vector<std::shared_ptr<TensorWrapper>> output_tensors =
-      qnn_manager->GetGraphOutputs(graph_name);
+      qnn_manager->GetGraphOutputs(method_name_);
   std::vector<Qnn_Tensor_t> input_tensor_structs;
   std::vector<Qnn_Tensor_t> output_tensor_structs;
 
@@ -118,14 +116,14 @@ Error QnnExecuTorchBackend::execute(
 
   ET_CHECK_OR_RETURN_ERROR(
       qnn_manager->Execute(
-          graph_name,
+          method_name_,
           input_tensor_structs,
           output_tensor_structs,
           context.event_tracer()) == Error::Ok,
       Internal,
       "Fail to execute graph");
   ET_CHECK_OR_RETURN_ERROR(
-      qnn_manager->ProfileExecuteData(graph_name, context.event_tracer()) ==
+      qnn_manager->ProfileExecuteData(method_name_, context.event_tracer()) ==
           Error::Ok,
       Internal,
       "Fail to profile graph");

backends/qualcomm/runtime/QnnExecuTorchBackend.h

@@ -34,6 +34,9 @@ class QnnExecuTorchBackend final
   void destroy(executorch::runtime::DelegateHandle* handle) const override;
 
   bool is_available() const override;
+
+ private:
+  mutable std::string method_name_;
 };
 
 } // namespace qnn

runtime/backend/backend_init_context.h

@@ -18,8 +18,10 @@ namespace runtime {
  */
 class BackendInitContext final {
  public:
-  explicit BackendInitContext(MemoryAllocator* runtime_allocator)
-      : runtime_allocator_(runtime_allocator) {}
+  explicit BackendInitContext(
+      MemoryAllocator* runtime_allocator,
+      const char* method_name)
+      : 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
@@ -29,8 +31,13 @@ class BackendInitContext final {
     return runtime_allocator_;
   }
 
+  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

@@ -55,7 +55,7 @@ class StubBackend final : public BackendInterface {
   using InitFn = std::function<Result<DelegateHandle*>(
       FreeableBuffer*,
       ArrayRef<CompileSpec>,
-      MemoryAllocator*)>;
+      BackendInitContext)>;
   using ExecuteFn = std::function<Error(DelegateHandle*, EValue**)>;
   using DestroyFn = std::function<void(DelegateHandle*)>;
 
@@ -83,8 +83,7 @@ class StubBackend final : public BackendInterface {
       FreeableBuffer* processed,
       ArrayRef<CompileSpec> compile_specs) const override {
     if (init_fn_) {
-      return init_fn_.value()(
-          processed, compile_specs, context.get_runtime_allocator());
+      return init_fn_.value()(processed, compile_specs, context);
     }
     // Return a benign value otherwise.
     return nullptr;
@@ -351,7 +350,7 @@ TEST_P(BackendIntegrationTest, FreeingProcessedBufferSucceeds) {
   StubBackend::singleton().install_init(
       [&](FreeableBuffer* processed,
           ET_UNUSED ArrayRef<CompileSpec> compile_specs,
-          ET_UNUSED MemoryAllocator* runtime_allocator)
+          ET_UNUSED BackendInitContext backend_init_context)
           -> Result<DelegateHandle*> {
         init_called = true;
         processed_data = processed->data();
@@ -395,7 +394,7 @@ TEST_P(BackendIntegrationTest, EndToEndTestWithProcessedAsHandle) {
   StubBackend::singleton().install_init(
       [&](FreeableBuffer* processed,
           ET_UNUSED ArrayRef<CompileSpec> compile_specs,
-          ET_UNUSED MemoryAllocator* runtime_allocator)
+          ET_UNUSED BackendInitContext backend_init_context)
           -> Result<DelegateHandle*> {
         init_processed = processed;
         return processed;
@@ -492,7 +491,7 @@ TEST_P(BackendIntegrationTest, SegmentInfoIsPassedIntoDataLoader) {
   StubBackend::singleton().install_init(
       [&](FreeableBuffer* processed,
           ET_UNUSED ArrayRef<CompileSpec> compile_specs,
-          ET_UNUSED MemoryAllocator* runtime_allocator)
+          ET_UNUSED BackendInitContext backend_init_context)
           -> Result<DelegateHandle*> {
         processed_data = processed->data();
         processed->Free();
@@ -528,6 +527,24 @@ 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()
@@ -606,7 +623,7 @@ TEST_P(DelegateDataAlignmentTest, ExpectedDataAlignment) {
   StubBackend::singleton().install_init(
       [&](FreeableBuffer* processed,
           ET_UNUSED ArrayRef<CompileSpec> compile_specs,
-          ET_UNUSED MemoryAllocator* runtime_allocator)
+          ET_UNUSED BackendInitContext backend_init_context)
           -> Result<DelegateHandle*> {
         processed_data = processed->data();
         return nullptr;