xref: /aosp_15_r20/external/executorch/exir/backend/test/demos/rpc/ExecutorBackend.cpp (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 #include <executorch/exir/backend/test/demos/rpc/ExecutorBackend.h>
10 
11 #include <cstdint>
12 #include <cstdio>
13 #include <cstdlib> /* strtol */
14 #include <memory>
15 
16 #include <executorch/extension/data_loader/buffer_data_loader.h>
17 #include <executorch/runtime/backend/interface.h>
18 #include <executorch/runtime/core/error.h>
19 #include <executorch/runtime/core/evalue.h>
20 #include <executorch/runtime/core/exec_aten/util/tensor_util.h>
21 #include <executorch/runtime/executor/method.h>
22 #include <executorch/runtime/executor/program.h>
23 
24 using ::executorch::aten::Tensor;
25 using ::executorch::extension::BufferDataLoader;
26 using ::executorch::runtime::ArrayRef;
27 using ::executorch::runtime::Backend;
28 using ::executorch::runtime::BackendExecutionContext;
29 using ::executorch::runtime::BackendInitContext;
30 using ::executorch::runtime::CompileSpec;
31 using ::executorch::runtime::DelegateHandle;
32 using ::executorch::runtime::Error;
33 using ::executorch::runtime::EValue;
34 using ::executorch::runtime::FreeableBuffer;
35 using ::executorch::runtime::HierarchicalAllocator;
36 using ::executorch::runtime::MemoryAllocator;
37 using ::executorch::runtime::MemoryManager;
38 using ::executorch::runtime::Method;
39 using ::executorch::runtime::MethodMeta;
40 using ::executorch::runtime::Program;
41 using ::executorch::runtime::Result;
42 using ::executorch::runtime::Span;
43 using ::executorch::runtime::Tag;
44 using ::executorch::runtime::internal::copy_tensor_data;
45 
46 namespace example {
47 
48 /**
49  * ExecutorBackend is a backend to execute an executorch program via delegate.
50  * In preprocess, the preprocesed bytes (delegate blob) is an executorch
51  * program. In ExecutorBackend, an executor backend is constructed in init and
52  * execute in execute. This backend can serve for 2 purposes
53  *
54  * 1. Serve as an RPC call to execute partial program on a different backend,
55  * for example, host executor in cpu and client executor in dsp.
56  * 2. Making incremental changes like experiment different different compiler
57  * front-end before having the actual backend ready.
58  */
59 
60 class ExecutorBackend final : public ::executorch::runtime::BackendInterface {
61  public:
62   ~ExecutorBackend() = default;
63 
is_available() const64   bool is_available() const override {
65     return true;
66   }
67 
init(BackendInitContext & context,FreeableBuffer * processed,ET_UNUSED ArrayRef<CompileSpec> compile_specs) const68   Result<DelegateHandle*> init(
69       BackendInitContext& context,
70       FreeableBuffer* processed,
71       ET_UNUSED ArrayRef<CompileSpec> compile_specs) const override {
72     // `processed` contains an executorch program. Wrap it in a DataLoader that
73     // will return the data directly without copying it.
74     MemoryAllocator* runtime_allocator = context.get_runtime_allocator();
75     auto loader = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
76         runtime_allocator, BufferDataLoader);
77     new (loader) BufferDataLoader(processed->data(), processed->size());
78     // Can't free `processed` because the program will point into that memory.
79 
80     // Try loading the program.
81     Result<Program> program_result = Program::load(loader);
82     if (!program_result.ok()) {
83       return program_result.error();
84     }
85 
86     // Move the Program off the stack.
87     auto client_program =
88         ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(runtime_allocator, Program);
89     new (client_program) Program(std::move(program_result.get()));
90 
91     Result<MethodMeta> method_meta = client_program->method_meta("forward");
92     if (!method_meta.ok()) {
93       ET_LOG(Error, "error constructing method meta");
94       return method_meta.error();
95     }
96 
97     // Building all different allocators for the client executor
98     auto num_memory_planned_buffers = method_meta->num_memory_planned_buffers();
99 
100     Span<uint8_t>* memory_planned_buffers = ET_ALLOCATE_LIST_OR_RETURN_ERROR(
101         runtime_allocator, Span<uint8_t>, num_memory_planned_buffers);
102 
103     for (size_t id = 0; id < num_memory_planned_buffers; ++id) {
104       size_t buffer_size = static_cast<size_t>(
105           method_meta->memory_planned_buffer_size(id).get());
106       uint8_t* buffer_i = ET_ALLOCATE_LIST_OR_RETURN_ERROR(
107           runtime_allocator, uint8_t, buffer_size);
108       memory_planned_buffers[id] = {buffer_i, buffer_size};
109     }
110 
111     auto client_planned_memory = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
112         runtime_allocator, HierarchicalAllocator);
113     new (client_planned_memory) HierarchicalAllocator(
114         {memory_planned_buffers, num_memory_planned_buffers});
115 
116     // Allocate some memory from runtime allocator for the client executor, in
117     // real case, like if it's an executor in dsp, it should allocate memory
118     // dedicated to this specific hardware
119     auto client_method_allocator = ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(
120         runtime_allocator, MemoryAllocator);
121     const size_t kClientRuntimeMemorySize = 4 * 1024U;
122     auto runtime_pool = ET_ALLOCATE_OR_RETURN_ERROR(
123         runtime_allocator, kClientRuntimeMemorySize);
124     new (client_method_allocator) MemoryAllocator(
125         kClientRuntimeMemorySize, static_cast<uint8_t*>(runtime_pool));
126 
127     auto client_memory_manager =
128         ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(runtime_allocator, MemoryManager);
129     new (client_memory_manager)
130         MemoryManager(client_method_allocator, client_planned_memory);
131 
132     // Construct the client Method
133     Result<Method> method_res =
134         client_program->load_method("forward", client_memory_manager);
135     if (!method_res.ok()) {
136       ET_LOG(
137           Error,
138           "Failed to load client method: 0x%x",
139           (unsigned int)method_res.error());
140       return method_res.error();
141     }
142 
143     auto client_method =
144         ET_ALLOCATE_INSTANCE_OR_RETURN_ERROR(runtime_allocator, Method);
145     new (client_method) Method(std::move(method_res.get()));
146 
147     // Return the client method so it will be passed to `execute()` as
148     // `handle`.
149     return client_method;
150   }
151 
execute(ET_UNUSED BackendExecutionContext & context,DelegateHandle * handle,EValue ** args) const152   Error execute(
153       ET_UNUSED BackendExecutionContext& context,
154       DelegateHandle* handle,
155       EValue** args) const override {
156     Method* client_method = static_cast<Method*>(handle);
157     auto num_inputs = client_method->inputs_size();
158     Error status = Error::Ok;
159 
160     // Receive client executor input
161     for (size_t input_idx = 0; input_idx < num_inputs; input_idx++) {
162       status = client_method->set_input(*args[input_idx], input_idx);
163     }
164     // Execute client executor
165     status = client_method->execute();
166 
167     auto output_sizes = client_method->outputs_size();
168     // Send the client executor output, we'd need to copy the data instead of
169     // assigning the Evalue pointer
170     for (int i = 0; i < output_sizes; i++) {
171       EValue output = client_method->get_output(i);
172       if (output.tag == Tag::Tensor) {
173         Tensor t_src = output.toTensor();
174         Tensor t_dst = args[num_inputs + i]->toTensor();
175         status = copy_tensor_data(t_dst, t_src);
176       }
177     }
178 
179     return status;
180   }
181 
destroy(DelegateHandle * handle) const182   void destroy(DelegateHandle* handle) const override {
183     if (handle != nullptr) {
184       Method* client_executor = static_cast<Method*>(handle);
185       client_executor->~Method();
186     }
187   }
188 };
189 
register_executor_backend()190 Error register_executor_backend() {
191   static auto cls = ExecutorBackend();
192   static Backend backend{"ExecutorBackend", &cls};
193   static auto success_with_compiler = register_backend(backend);
194   return success_with_compiler;
195 }
196 
197 } // namespace example
198