Add backend runtime context to backend.execute (#211)

Summary:
Pull Request resolved: #211

We're adding backend runtime context which comes with mainly two extra functions
- Temp allocator (life span is per execute)
- Event tracer (profiling inside delegate)

Reviewed By: tarun292, dbort

Differential Revision: D48872100

fbshipit-source-id: 8e8019a4e1fa280ea54831d77a54e879c4f572d1

Author: cccclai

backends/qnnpack/QNNPackBackend.cpp

@@ -195,7 +195,10 @@ class QnnpackBackend final : public PyTorchBackendInterface {
     return executor;
   }
 
-  Error execute(DelegateHandle* handle, EValue** args) const override {
+  Error execute(
+      __ET_UNUSED BackendExecutionContext& context,
+      DelegateHandle* handle,
+      EValue** args) const override {
     static constexpr size_t kMaxDims = 16;
 
     QNNExecutor* etor = static_cast<QNNExecutor*>(handle);

backends/vulkan/VulkanBackend.cpp

@@ -270,7 +270,10 @@ class VulkanBackend final : public PyTorchBackendInterface {
     return compute_graph;
   }
 
-  Error execute(DelegateHandle* handle, EValue** args) const override {
+  Error execute(
+      __ET_UNUSED BackendExecutionContext& context,
+      DelegateHandle* handle,
+      EValue** args) const override {
     EXECUTORCH_SCOPE_PROF("VulkanBackend::execute");
 
     at::native::vulkan::ComputeGraph* compute_graph =

backends/xnnpack/runtime/XNNPACKBackend.cpp

@@ -49,7 +49,10 @@ class XnnpackBackend final : public PyTorchBackendInterface {
     return executor;
   }
 
-  Error execute(DelegateHandle* handle, EValue** args) const override {
+  Error execute(
+      __ET_UNUSED BackendExecutionContext& context,
+      DelegateHandle* handle,
+      EValue** args) const override {
     auto executor = static_cast<xnnpack::delegate::XNNExecutor*>(handle);
 
     std::vector<Tensor*> input_pointers;

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

@@ -140,7 +140,10 @@ class ExecutorBackend final : public PyTorchBackendInterface {
     return client_method;
   }
 
-  Error execute(DelegateHandle* handle, EValue** args) const override {
+  Error execute(
+      __ET_UNUSED BackendExecutionContext& context,
+      DelegateHandle* handle,
+      EValue** args) const override {
     Method* client_method = static_cast<Method*>(handle);
     auto num_inputs = client_method->inputs_size();
     Error status = Error::Ok;

runtime/backend/backend_registry.h

@@ -10,6 +10,7 @@
 
 #include <cstring>
 
+#include <executorch/runtime/backend/backend_execution_context.h>
 #include <executorch/runtime/core/array_ref.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/evalue.h>
@@ -88,8 +89,10 @@ class PyTorchBackendInterface {
    * @param[in] args The method’s inputs and outputs.
    * @retval Error::Ok if successful.
    */
-  __ET_NODISCARD virtual Error execute(DelegateHandle* handle, EValue** args)
-      const = 0;
+  __ET_NODISCARD virtual Error execute(
+      BackendExecutionContext& context,
+      DelegateHandle* handle,
+      EValue** args) const = 0;
 
   /**
    * Responsible for destroying a handle, if it's required for some backend.

runtime/executor/method.cpp

@@ -115,9 +115,11 @@ class BackendDelegate final {
     }
   }
 
-  Error Execute(EValue** args) const {
+  Error Execute(
+      BackendExecutionContext& backend_execution_context,
+      EValue** args) const {
     EXECUTORCH_SCOPE_PROF("delegate_execute");
-    return backend_->execute(handle_, args);
+    return backend_->execute(backend_execution_context, handle_, args);
   }
 
  private:
@@ -939,7 +941,9 @@ Error Method::execute_instruction() {
           delegate_idx,
           n_delegate_,
           step_state_.instr_idx);
+      BackendExecutionContext backend_execution_context;
       Error err = delegates_[delegate_idx].Execute(
+          backend_execution_context,
           chain.argument_lists_[step_state_.instr_idx].data());
       ET_CHECK_MSG(
           err == Error::Ok,

runtime/executor/test/backend_integration_test.cpp

@@ -29,6 +29,7 @@
 
 using namespace ::testing;
 using exec_aten::ArrayRef;
+using torch::executor::BackendExecutionContext;
 using torch::executor::CompileSpec;
 using torch::executor::DataLoader;
 using torch::executor::DelegateHandle;
@@ -91,7 +92,10 @@ class StubBackend final : public PyTorchBackendInterface {
     execute_fn_ = fn;
   }
 
-  Error execute(DelegateHandle* handle, EValue** args) const override {
+  Error execute(
+      __ET_UNUSED BackendExecutionContext& context,
+      DelegateHandle* handle,
+      EValue** args) const override {
     if (execute_fn_) {
       return execute_fn_.value()(handle, args);
     }

runtime/executor/test/test_backend_compiler_lib.cpp

@@ -136,7 +136,10 @@ class BackendWithCompiler final : public PyTorchBackendInterface {
   // execute and it only supports add, subtract, and constant. In a non toy
   // backend you can imagine how this function could be used to actually
   // dispatch the inputs to the relevant backend/device.
-  Error execute(DelegateHandle* handle, EValue** args) const override {
+  Error execute(
+      __ET_UNUSED BackendExecutionContext& context,
+      DelegateHandle* handle,
+      EValue** args) const override {
     EXECUTORCH_SCOPE_PROF("BackendWithCompiler::execute");
 
     // example: [('prim::Constant#1', 14), ('aten::add', 15)]

runtime/executor/test/test_backend_with_delegate_mapping.cpp

@@ -110,7 +110,10 @@ class BackendWithDelegateMapping final : public PyTorchBackendInterface {
 
   // This function doesn't actually execute the op but just prints out the op
   // name and the corresponding delegate debug identifier.
-  Error execute(DelegateHandle* handle, EValue** args) const override {
+  Error execute(
+      __ET_UNUSED BackendExecutionContext& context,
+      DelegateHandle* handle,
+      EValue** args) const override {
     (void)args;
     // example: [('prim::Constant#1', 14), ('aten::add', 15)]
     auto op_list = static_cast<const DemoOpList*>(handle);