Add backend init context to backend.init (#276)

Summary:
Pull Request resolved: #276

Create `BackendInitContext` to wrap the runtime allocator. We may inject `EventTracer` to `BackendInitContext` to profile init time later

Reviewed By: dbort

Differential Revision: D48872105

fbshipit-source-id: 5022cec0b3dbf63562a6397d71417474eae10a37

Author: cccclai

backends/qnnpack/QNNPackBackend.cpp

@@ -112,9 +112,10 @@ class QnnpackBackend final : public PyTorchBackendInterface {
   }
 
   Result<DelegateHandle*> init(
+      BackendInitContext& context,
       FreeableBuffer* processed,
-      ArrayRef<CompileSpec> compile_specs,
-      MemoryAllocator* runtime_allocator) const override {
+      ArrayRef<CompileSpec> compile_specs) const override {
+    MemoryAllocator* runtime_allocator = context.get_runtime_allocator();
     auto dynamic_linear = fb_qnnpack::GetQNNDynamicLinear(processed->data());
     auto bias = dynamic_linear->bias();
 

backends/vulkan/VulkanBackend.cpp

@@ -245,9 +245,9 @@ class VulkanBackend final : public PyTorchBackendInterface {
   }
 
   Result<DelegateHandle*> init(
+      BackendInitContext& context,
       FreeableBuffer* processed,
-      ArrayRef<CompileSpec>,
-      MemoryAllocator* runtime_allocator) const override {
+      ArrayRef<CompileSpec>) const override {
     ET_CHECK_OR_RETURN_ERROR(
         at::vulkan::delegate::VkGraphBufferHasIdentifier(processed->data()),
         DelegateInvalidCompatibility,
@@ -257,7 +257,7 @@ class VulkanBackend final : public PyTorchBackendInterface {
 
     at::native::vulkan::ComputeGraph* compute_graph =
         ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
-            runtime_allocator, at::native::vulkan::ComputeGraph);
+            context.get_runtime_allocator(), at::native::vulkan::ComputeGraph);
 
     new (compute_graph) at::native::vulkan::ComputeGraph(generate_config());
 

backends/xnnpack/runtime/XNNPACKBackend.cpp

@@ -25,11 +25,11 @@ class XnnpackBackend final : public PyTorchBackendInterface {
   }
 
   Result<DelegateHandle*> init(
+      BackendInitContext& context,
       FreeableBuffer* processed,
-      ArrayRef<CompileSpec> compile_specs,
-      MemoryAllocator* runtime_allocator) const override {
+      ArrayRef<CompileSpec> compile_specs) const override {
     auto executor = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
-        runtime_allocator, xnnpack::delegate::XNNExecutor);
+        context.get_runtime_allocator(), xnnpack::delegate::XNNExecutor);
 
     // Executor has been allocated but not constructed, ensure that runtime_ is
     // nullptr by constructing it in place here. NOTE: Since we use placement
@@ -38,7 +38,10 @@ class XnnpackBackend final : public PyTorchBackendInterface {
     new (executor) xnnpack::delegate::XNNExecutor;
 
     Error err = xnnpack::delegate::XNNCompiler::compileModel(
-        processed->data(), processed->size(), executor, runtime_allocator);
+        processed->data(),
+        processed->size(),
+        executor,
+        context.get_runtime_allocator());
     if (err != Error::Ok) {
       ET_LOG(Error, "XNNCompiler::compleModel failed: 0x%x", (unsigned int)err);
     }

exir/backend/test/demos/rpc/ExecutorBackend.cpp

@@ -44,11 +44,12 @@ class ExecutorBackend final : public PyTorchBackendInterface {
   }
 
   Result<DelegateHandle*> init(
+      BackendInitContext& context,
       FreeableBuffer* processed,
-      __ET_UNUSED ArrayRef<CompileSpec> compile_specs,
-      MemoryAllocator* runtime_allocator) const override {
+      __ET_UNUSED ArrayRef<CompileSpec> compile_specs) const override {
     // `processed` contains an executorch program. Wrap it in a DataLoader that
     // will return the data directly without copying it.
+    MemoryAllocator* runtime_allocator = context.get_runtime_allocator();
     auto loader = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
         runtime_allocator, util::BufferDataLoader);
     new (loader) util::BufferDataLoader(processed->data(), processed->size());

runtime/backend/backend_registry.h

@@ -11,6 +11,7 @@
 #include <cstring>
 
 #include <executorch/runtime/backend/backend_execution_context.h>
+#include <executorch/runtime/backend/backend_init_context.h>
 #include <executorch/runtime/core/array_ref.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/evalue.h>
@@ -75,9 +76,9 @@ class PyTorchBackendInterface {
    * @returns On error, a value other than Error:Ok.
    */
   __ET_NODISCARD virtual Result<DelegateHandle*> init(
+      BackendInitContext& context,
       FreeableBuffer* processed,
-      ArrayRef<CompileSpec> compile_specs,
-      MemoryAllocator* memory_allocator) const = 0;
+      ArrayRef<CompileSpec> compile_specs) const = 0;
 
   /**
    * Responsible for executing the given method’s handle, as it was produced

runtime/executor/method.cpp

@@ -52,7 +52,7 @@ class BackendDelegate final {
   static Error Init(
       const executorch_flatbuffer::BackendDelegate& delegate,
       const Program* program,
-      MemoryAllocator* runtime_allocator,
+      BackendInitContext& backend_init_context,
       BackendDelegate* out) {
     // Look up the backend.
     const char* backend_id = delegate.id()->c_str();
@@ -78,7 +78,7 @@ class BackendDelegate final {
     // Parse compilation specs from program
     CompileSpec* compile_specs;
     Error err = PopulateCompileSpecs(
-        delegate.compile_specs(), runtime_allocator, &compile_specs);
+        delegate.compile_specs(), backend_init_context, &compile_specs);
     if (err != Error::Ok) {
       ET_LOG(Error, "Failed to get compile specs for backend %s", backend_id);
       return err;
@@ -93,9 +93,9 @@ class BackendDelegate final {
 
     // Initialize the delegate.
     Result<DelegateHandle*> handle = backend->init(
+        backend_init_context,
         &out->segment_,
-        ArrayRef<CompileSpec>(compile_specs, num_compile_specs),
-        runtime_allocator);
+        ArrayRef<CompileSpec>(compile_specs, num_compile_specs));
     if (!handle.ok()) {
       ET_LOG(
           Error,
@@ -135,12 +135,14 @@ class BackendDelegate final {
   static Error PopulateCompileSpecs(
       const flatbuffers::Vector<flatbuffers::Offset<
           executorch_flatbuffer::CompileSpec>>* compile_specs_in_program,
-      torch::executor::MemoryAllocator* runtime_allocator,
+      BackendInitContext& backend_init_context,
       CompileSpec** out_spec) {
     auto number_of_compile_specs = compile_specs_in_program->size();
 
     CompileSpec* compile_specs_list = ET_ALLOCATE_LIST_OR_RETURN_ERROR(
-        runtime_allocator, CompileSpec, number_of_compile_specs);
+        backend_init_context.get_runtime_allocator(),
+        CompileSpec,
+        number_of_compile_specs);
 
     // Initialize the spec list for each method spec
     for (size_t j = 0; j < number_of_compile_specs; j++) {
@@ -537,14 +539,15 @@ 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(runtime_allocator);
       Error err = BackendDelegate::Init(
-          delegate, program_, runtime_allocator, &delegates_[i]);
+          delegate, program_, backend_init_context, &delegates_[i]);
       if (err != Error::Ok) {
         return err;
       }
       // ~Method() will try to clean up n_delegate_ entries in the delegates_
-      // array. Only increment this once we know the entry is valid, so that we
-      // don't try to clean up an uninitialized entry.
+      // array. Only increment this once we know the entry is valid, so that
+      // we don't try to clean up an uninitialized entry.
       n_delegate_ = i + 1;
     }
   }
@@ -1035,8 +1038,8 @@ Error Method::execute() {
       NotSupported,
       "Cannot execute until method has been initialized.");
 
-  // Chains are executed sequentially today, but future async designs may branch
-  // and run many in parallel or out of order.
+  // Chains are executed sequentially today, but future async designs may
+  // branch and run many in parallel or out of order.
   for (step_state_.chain_idx = 0; step_state_.chain_idx < n_chains_;
        ++step_state_.chain_idx) {
     Chain& chain = chains_[step_state_.chain_idx];

runtime/executor/test/backend_integration_test.cpp

@@ -30,6 +30,7 @@
 using namespace ::testing;
 using exec_aten::ArrayRef;
 using torch::executor::BackendExecutionContext;
+using torch::executor::BackendInitContext;
 using torch::executor::CompileSpec;
 using torch::executor::DataLoader;
 using torch::executor::DelegateHandle;
@@ -78,11 +79,12 @@ class StubBackend final : public PyTorchBackendInterface {
   }
 
   Result<DelegateHandle*> init(
+      BackendInitContext& context,
       FreeableBuffer* processed,
-      ArrayRef<CompileSpec> compile_specs,
-      MemoryAllocator* runtime_allocator) const override {
+      ArrayRef<CompileSpec> compile_specs) const override {
     if (init_fn_) {
-      return init_fn_.value()(processed, compile_specs, runtime_allocator);
+      return init_fn_.value()(
+          processed, compile_specs, context.get_runtime_allocator());
     }
     // Return a benign value otherwise.
     return nullptr;

runtime/executor/test/test_backend_compiler_lib.cpp

@@ -94,9 +94,10 @@ class BackendWithCompiler final : public PyTorchBackendInterface {
   }
 
   Result<DelegateHandle*> init(
+      BackendInitContext& context,
       FreeableBuffer* processed,
-      ArrayRef<CompileSpec> compile_specs,
-      MemoryAllocator* runtime_allocator) const override {
+      ArrayRef<CompileSpec> compile_specs) const override {
+    MemoryAllocator* runtime_allocator = context.get_runtime_allocator();
     int shape = *(int*)(compile_specs.at(0).value.buffer);
     ET_CHECK_OR_RETURN_ERROR(
         shape <= max_shape,

runtime/executor/test/test_backend_with_delegate_mapping.cpp

@@ -77,9 +77,10 @@ class BackendWithDelegateMapping final : public PyTorchBackendInterface {
   }
 
   Result<DelegateHandle*> init(
+      BackendInitContext& context,
       FreeableBuffer* processed,
-      ArrayRef<CompileSpec> compile_specs,
-      MemoryAllocator* runtime_allocator) const override {
+      ArrayRef<CompileSpec> compile_specs) const override {
+    MemoryAllocator* runtime_allocator = context.get_runtime_allocator();
     (void)compile_specs;
     const char* kSignLiteral = "#";
     // The first number is the number of total instruction