xref: /aosp_15_r20/external/executorch/runtime/kernel/kernel_runtime_context.h (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1 /*
2  * Copyright (c) Meta Platforms, Inc. and affiliates.
3  * All rights reserved.
4  *
5  * This source code is licensed under the BSD-style license found in the
6  * LICENSE file in the root directory of this source tree.
7  */
8 
9 #pragma once
10 
11 #include <executorch/runtime/core/error.h>
12 #include <executorch/runtime/core/event_tracer_hooks.h>
13 #include <executorch/runtime/core/memory_allocator.h>
14 #include <executorch/runtime/core/result.h>
15 #include <executorch/runtime/platform/compiler.h>
16 
17 namespace executorch {
18 namespace runtime {
19 
20 /**
21  * Runtime state and functionality for kernel implementations.
22  *
23  * NOTE: Will not be passed to operators if running in ATen mode as those
24  * operators do not expect to receive a KernelRuntimeContext argument.
25  */
26 class KernelRuntimeContext {
27  public:
28   /**
29    * Construct a new kernel runtime context.
30    *
31    * KernelRuntimeContext does not take ownership
32    * of these pointers, so they must outlive the context instance.
33    *
34    * @param[in] event_tracer The optional EventTracer to use for
35    *     profiling/debugging
36    * @param[in] temp_allocator The optional MemoryAllocator used to allocate
37    *     temporary memory for the kernel. If not provided, an error will be
38    *     returned when calling allocate_temp.
39    */
40   KernelRuntimeContext(
41       EventTracer* event_tracer = nullptr,
42       MemoryAllocator* temp_allocator = nullptr)
event_tracer_(event_tracer)43       : event_tracer_(event_tracer), temp_allocator_(temp_allocator) {}
44   /**
45    * Tells the runtime that the kernel call has failed. Prefer this over
46    * ET_CHECK_*(), which fatally panics the process/system.
47    *
48    * If this is not called, the runtime will treat the kernel call as
49    * successful.
50    *
51    * This unusual error-propagation path is required because kernel signatures
52    * do not have a natural way to return errors directly. They are generally
53    * compatible with core PyTorch ATen kernel signatures, which use exceptions
54    * to report errors. But, ExecuTorch does not use exceptions.
55    */
fail(Error error)56   void fail(Error error) {
57     failure_state_ = error;
58   }
59 
60   /// Returns the current failure state.
failure_state()61   ET_NODISCARD Error failure_state() const {
62     return failure_state_;
63   }
64 
65   /**
66    * INTERNAL ONLY
67    *
68    * Returns a pointer to an instance of EventTracer to do profiling/debugging
69    * logging inside the codegen layer. This is only for internal usage inside
70    * the codegen layer and users should not be accessing this.
71    */
internal_event_tracer()72   EventTracer* internal_event_tracer() {
73     return event_tracer_;
74   }
75 
76   /**
77    * Allocates temporary memory that will be freed when the kernel returns. This
78    * returns a pointer to the allocated memory or an error if the allocation
79    * fails.
80    *
81    * @param[in] size Number of bytes to allocate.
82    * @param[in] alignment Minimum alignment for the returned pointer. Must be a
83    *     power of 2.
84    *
85    * @returns A result object containing either a pointer to the allocated
86    *     memory or an error to indicate failure
87    */
88   Result<void*> allocate_temp(
89       size_t size,
90       size_t alignment = MemoryAllocator::kDefaultAlignment) {
91     ET_CHECK_OR_RETURN_ERROR(
92         temp_allocator_ != nullptr, NotFound, "No temp allocator provided");
93     void* temp_memory = temp_allocator_->allocate(size, alignment);
94     ET_CHECK_OR_RETURN_ERROR(
95         temp_memory != nullptr,
96         MemoryAllocationFailed,
97         "Failed to allocate temp memory. Bytes requested: %zu",
98         size);
99     return temp_memory;
100   }
101 
102   // TODO(T147221312): Add a way to resize a tensor.
103 
104  private:
105   EventTracer* event_tracer_ = nullptr;
106   MemoryAllocator* temp_allocator_ = nullptr;
107   Error failure_state_ = Error::Ok;
108 };
109 
110 } // namespace runtime
111 } // namespace executorch
112 
113 // TODO(T197294990): Remove these deprecated aliases once all users have moved
114 // to the new `::executorch` namespaces.
115 namespace torch {
116 namespace executor {
117 /// DEPRECATED: Use ::executorch::runtime::KernelRuntimeContext instead.
118 using ::executorch::runtime::KernelRuntimeContext;
119 /// DEPRECATED: Use ::executorch::runtime::KernelRuntimeContext instead.
120 using RuntimeContext = ::executorch::runtime::KernelRuntimeContext;
121 } // namespace executor
122 } // namespace torch
123 namespace executorch {
124 namespace aten {
125 /// DEPRECATED: Use ::executorch::runtime::KernelRuntimeContext instead.
126 using RuntimeContext = ::executorch::runtime::KernelRuntimeContext;
127 } // namespace aten
128 } // namespace executorch
129 // DEPRECATED: The exec_aten:: namespace is deprecated. Use executorch::aten::
130 // instead.
131 namespace exec_aten = ::executorch::aten;
132