Add allocate_temp method to KernelRuntimeContext (#3209)

Summary:

This adds an `allocate_temp` method to KernelRuntimeContext, and passes the temporary memory allocator from `execute_instruction`. The method returns a result that errors if the temporary `MemoryAllocator` was not provided or the memory could not be allocated.

Differential Revision: D56421957

Author: David Lin

runtime/executor/method.cpp

@@ -1015,7 +1015,8 @@ Error Method::execute_instruction() {
           internal::EventTracerProfileScope(event_tracer_, "OPERATOR_CALL");
       // TODO(T147221312): Also expose the temp allocator and tensor resizer
       // via the context.
-      KernelRuntimeContext context(event_tracer_);
+      KernelRuntimeContext context(
+          event_tracer_, memory_manager_->temp_allocator());
       auto args = chain.argument_lists_[step_state_.instr_idx];
       chain.kernels_[step_state_.instr_idx](context, args.data());
       err = context.failure_state();

runtime/kernel/kernel_runtime_context.h

@@ -10,6 +10,8 @@
 
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/event_tracer_hooks.h>
+#include <executorch/runtime/core/memory_allocator.h>
+#include <executorch/runtime/core/result.h>
 #include <executorch/runtime/platform/compiler.h>
 
 namespace torch {
@@ -24,10 +26,13 @@ namespace executor {
 class KernelRuntimeContext {
  public:
   /**
-   * Construct a new kernel runtime context along with an optional event tracer.
+   * Construct a new kernel runtime context along with an optional event tracer
+   * and an optional memory allocator.
    */
-  KernelRuntimeContext(EventTracer* event_tracer = nullptr)
-      : event_tracer_(event_tracer) {}
+  KernelRuntimeContext(
+      EventTracer* event_tracer = nullptr,
+      MemoryAllocator* temp_allocator = nullptr)
+      : event_tracer_(event_tracer), temp_allocator_(temp_allocator) {}
   /**
    * Tells the runtime that the kernel call has failed. Prefer this over
    * ET_CHECK_*(), which fatally panics the process/system.
@@ -59,13 +64,29 @@ class KernelRuntimeContext {
   EventTracer* internal_event_tracer() {
     return event_tracer_;
   }
-
-  // TODO(T147221312): Add a way to allocate temporary memory.
+  /**
+   * Allocates temporary memory that will be freed when the kernel returns. This
+   * returns a pointer to the allocated memory, or an error if there is no
+   * MemoryAllocator provided or the allocation fails.
+   */
+  Result<void*> allocate_temp(
+      size_t size,
+      size_t alignment = MemoryAllocator::kDefaultAlignment) {
+    ET_CHECK_OR_RETURN_ERROR(
+        temp_allocator_ != nullptr, NotFound, "No temp allocator provided");
+    void* temp_memory = temp_allocator_->allocate(size, alignment);
+    ET_CHECK_OR_RETURN_ERROR(
+        temp_memory != nullptr,
+        MemoryAllocationFailed,
+        "Failed to allocate temp memory");
+    return temp_memory;
+  }
 
   // TODO(T147221312): Add a way to resize a tensor.
 
  private:
   EventTracer* event_tracer_ = nullptr;
+  MemoryAllocator* temp_allocator_ = nullptr;
   Error failure_state_ = Error::Ok;
 };
 

runtime/kernel/targets.bzl

@@ -55,6 +55,7 @@ def define_common_targets():
             exported_deps = [
                 "//executorch/runtime/core:core",
                 "//executorch/runtime/platform:platform",
+                "//executorch/runtime/core:memory_allocator",
                 "//executorch/runtime/core:event_tracer" + aten_suffix,
                 # TODO(T147221312): This will eventually depend on exec_aten
                 # once KernelRuntimeContext support tensor resizing, which is