Skip to content

Commit b87d5de

Browse files
manuelcandalesmanuelcandales
committed
Update
[ghstack-poisoned]
1 parent f9c8989 commit b87d5de

File tree

1 file changed

+99
-30
lines changed

1 file changed

+99
-30
lines changed

‎backends/apple/metal/runtime/metal_backend.cpp‎

Lines changed: 99 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -58,33 +58,52 @@ using executorch::runtime::Result;
5858
using executorch::runtime::Span;
5959
using executorch::runtime::etensor::Tensor;
6060

61-
class ET_EXPERIMENTAL MetalBackend final : public ::executorch::runtime::BackendInterface {
61+
class ET_EXPERIMENTAL MetalBackend final
62+
: public ::executorch::runtime::BackendInterface {
6263
private:
6364
Error register_shared_library_functions(void* so_handle) const {
64-
ET_LOG(Info, "MetalBackend::register_shared_library_functions - Loading symbols");
65+
ET_LOG(
66+
Info,
67+
"MetalBackend::register_shared_library_functions - Loading symbols");
6568

6669
LOAD_SYMBOL(AOTInductorModelContainerCreateWithDevice, so_handle);
67-
ET_LOG(Info, "MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerCreateWithDevice");
70+
ET_LOG(
71+
Info,
72+
"MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerCreateWithDevice");
6873

6974
LOAD_SYMBOL(AOTInductorModelContainerDelete, so_handle);
70-
ET_LOG(Info, "MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerDelete");
75+
ET_LOG(
76+
Info,
77+
"MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerDelete");
7178

7279
LOAD_SYMBOL(AOTInductorModelContainerGetNumInputs, so_handle);
73-
ET_LOG(Info, "MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerGetNumInputs");
80+
ET_LOG(
81+
Info,
82+
"MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerGetNumInputs");
7483

7584
LOAD_SYMBOL(AOTInductorModelContainerGetNumConstants, so_handle);
76-
ET_LOG(Info, "MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerGetNumConstants");
85+
ET_LOG(
86+
Info,
87+
"MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerGetNumConstants");
7788

7889
LOAD_SYMBOL(AOTInductorModelContainerGetInputName, so_handle);
79-
ET_LOG(Info, "MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerGetInputName");
90+
ET_LOG(
91+
Info,
92+
"MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerGetInputName");
8093

8194
LOAD_SYMBOL(AOTInductorModelContainerGetNumOutputs, so_handle);
82-
ET_LOG(Info, "MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerGetNumOutputs");
95+
ET_LOG(
96+
Info,
97+
"MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerGetNumOutputs");
8398

8499
LOAD_SYMBOL(AOTInductorModelContainerRun, so_handle);
85-
ET_LOG(Info, "MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerRun");
100+
ET_LOG(
101+
Info,
102+
"MetalBackend::register_shared_library_functions - Loaded AOTInductorModelContainerRun");
86103

87-
ET_LOG(Info, "MetalBackend::register_shared_library_functions - All symbols loaded successfully");
104+
ET_LOG(
105+
Info,
106+
"MetalBackend::register_shared_library_functions - All symbols loaded successfully");
88107
return Error::Ok;
89108
}
90109

@@ -123,7 +142,10 @@ class ET_EXPERIMENTAL MetalBackend final : public ::executorch::runtime::Backend
123142
const NamedDataMap* named_data_map = context.get_named_data_map();
124143
ET_LOG(Info, "MetalBackend::init - Got named data map: %p", named_data_map);
125144

126-
ET_LOG(Info, "MetalBackend::init - Looking for blob key: %s", so_blob_key.c_str());
145+
ET_LOG(
146+
Info,
147+
"MetalBackend::init - Looking for blob key: %s",
148+
so_blob_key.c_str());
127149

128150
auto aoti_metal_buffer = named_data_map->get_data(so_blob_key.c_str());
129151
ET_CHECK_OR_RETURN_ERROR(
@@ -133,22 +155,29 @@ class ET_EXPERIMENTAL MetalBackend final : public ::executorch::runtime::Backend
133155
so_blob_key.c_str(),
134156
static_cast<uint32_t>(aoti_metal_buffer.error()));
135157

136-
ET_LOG(Info, "MetalBackend::init - Buffer is OK, size: %zu", aoti_metal_buffer->size());
158+
ET_LOG(
159+
Info,
160+
"MetalBackend::init - Buffer is OK, size: %zu",
161+
aoti_metal_buffer->size());
137162

138163
if (aoti_metal_buffer->data() == nullptr) {
139164
ET_LOG(Error, "MetalBackend::init - Buffer data is null");
140165
return Error::InvalidArgument;
141166
}
142167

143-
ET_LOG(Info, "MetalBackend::init - Buffer data pointer: %p", aoti_metal_buffer->data());
168+
ET_LOG(
169+
Info,
170+
"MetalBackend::init - Buffer data pointer: %p",
171+
aoti_metal_buffer->data());
144172

145173
// Generate dynamic temporary file path
146174
filesystem::path temp_dir = filesystem::temp_directory_path();
147175
filesystem::path so_path =
148176
temp_dir / (so_blob_key + to_string(getpid()) + ".so");
149177

150178
// Create a temporary file
151-
ET_LOG(Info, "MetalBackend::init - Creating temp file: %s", so_path.c_str());
179+
ET_LOG(
180+
Info, "MetalBackend::init - Creating temp file: %s", so_path.c_str());
152181
ofstream outfile(so_path.c_str(), ios::binary);
153182

154183
// Write the ELF buffer to the temporary file
@@ -183,7 +212,9 @@ class ET_EXPERIMENTAL MetalBackend final : public ::executorch::runtime::Backend
183212
ET_CHECK_OK_OR_RETURN_ERROR(register_shared_library_functions(so_handle));
184213

185214
AOTInductorModelContainerHandle container_handle = nullptr;
186-
ET_LOG(Info, "MetalBackend::init - About to create AOTI container with device='mps'");
215+
ET_LOG(
216+
Info,
217+
"MetalBackend::init - About to create AOTI container with device='mps'");
187218

188219
ET_CHECK_OK_OR_RETURN_ERROR(AOTInductorModelContainerCreateWithDevice(
189220
&container_handle, 1, "mps", nullptr));
@@ -262,7 +293,11 @@ class ET_EXPERIMENTAL MetalBackend final : public ::executorch::runtime::Backend
262293
auto cpu_tensor = &(args[i]->toTensor());
263294
auto sizes = cpu_tensor->sizes();
264295
auto scalar_type = cpu_tensor->scalar_type();
265-
ET_LOG(Debug, "MetalBackend input %d scalar_type=%d", i, static_cast<int32_t>(scalar_type));
296+
ET_LOG(
297+
Debug,
298+
"MetalBackend input %d scalar_type=%d",
299+
i,
300+
static_cast<int32_t>(scalar_type));
266301

267302
// Create GPU tensor with same shape
268303
std::vector<int64_t> sizes_vec(sizes.begin(), sizes.end());
@@ -283,17 +318,28 @@ class ET_EXPERIMENTAL MetalBackend final : public ::executorch::runtime::Backend
283318
}
284319

285320
// Log the created GPU tensor scalar type
286-
auto gpu_tensor = reinterpret_cast<executorch::runtime::etensor::Tensor*>(gpu_input_handle);
287-
ET_LOG(Debug, "MetalBackend created GPU tensor %d scalar_type=%d", i, static_cast<int32_t>(gpu_tensor->scalar_type()));
321+
auto gpu_tensor = reinterpret_cast<executorch::runtime::etensor::Tensor*>(
322+
gpu_input_handle);
323+
ET_LOG(
324+
Debug,
325+
"MetalBackend created GPU tensor %d scalar_type=%d",
326+
i,
327+
static_cast<int32_t>(gpu_tensor->scalar_type()));
288328

289329
gpu_inputs[i] = gpu_input_handle;
290330

291331
// Log the CPU tensor data before copying to GPU
292332
void* cpu_data = cpu_tensor->mutable_data_ptr();
293333
if (cpu_data && cpu_tensor->numel() > 0) {
294334
float* cpu_float_data = (float*)cpu_data;
295-
ET_LOG(Debug, "CPU input %d data before copy: [%.3f, %.3f, %.3f, ...] (numel=%zd)",
296-
i, cpu_float_data[0], cpu_float_data[1], cpu_float_data[2], cpu_tensor->numel());
335+
ET_LOG(
336+
Debug,
337+
"CPU input %d data before copy: [%.3f, %.3f, %.3f, ...] (numel=%zd)",
338+
i,
339+
cpu_float_data[0],
340+
cpu_float_data[1],
341+
cpu_float_data[2],
342+
cpu_tensor->numel());
297343
}
298344

299345
// Copy data from CPU to GPU
@@ -304,8 +350,14 @@ class ET_EXPERIMENTAL MetalBackend final : public ::executorch::runtime::Backend
304350
}
305351

306352
// Log the GPU tensor scalar type after copy
307-
auto gpu_tensor_after = reinterpret_cast<executorch::runtime::etensor::Tensor*>(gpu_inputs[i]);
308-
ET_LOG(Debug, "MetalBackend GPU tensor %d scalar_type after copy=%d", i, static_cast<int32_t>(gpu_tensor_after->scalar_type()));
353+
auto gpu_tensor_after =
354+
reinterpret_cast<executorch::runtime::etensor::Tensor*>(
355+
gpu_inputs[i]);
356+
ET_LOG(
357+
Debug,
358+
"MetalBackend GPU tensor %d scalar_type after copy=%d",
359+
i,
360+
static_cast<int32_t>(gpu_tensor_after->scalar_type()));
309361

310362
ET_LOG(Debug, "Successfully copied input %d from CPU to GPU", i);
311363
}
@@ -319,7 +371,11 @@ class ET_EXPERIMENTAL MetalBackend final : public ::executorch::runtime::Backend
319371
auto cpu_output_tensor = &(args[i + n_inputs]->toTensor());
320372
auto sizes = cpu_output_tensor->sizes();
321373
auto scalar_type = cpu_output_tensor->scalar_type();
322-
ET_LOG(Debug, "MetalBackend output %d scalar_type=%d", i, static_cast<int32_t>(scalar_type));
374+
ET_LOG(
375+
Debug,
376+
"MetalBackend output %d scalar_type=%d",
377+
i,
378+
static_cast<int32_t>(scalar_type));
323379

324380
// Create GPU tensor with same shape for kernel output
325381
std::vector<int64_t> sizes_vec(sizes.begin(), sizes.end());
@@ -349,13 +405,21 @@ class ET_EXPERIMENTAL MetalBackend final : public ::executorch::runtime::Backend
349405
ET_LOG(Debug, "Passing to AOTInductorModelContainerRun:");
350406
for (int i = 0; i < n_inputs; i++) {
351407
void* gpu_input_data = gpu_inputs[i]->mutable_data_ptr();
352-
ET_LOG(Debug, " gpu_inputs[%d] = %p, data_ptr = %p",
353-
i, gpu_inputs[i], gpu_input_data);
408+
ET_LOG(
409+
Debug,
410+
" gpu_inputs[%d] = %p, data_ptr = %p",
411+
i,
412+
gpu_inputs[i],
413+
gpu_input_data);
354414
}
355415
for (int i = 0; i < n_outputs; i++) {
356416
void* gpu_output_data = gpu_outputs[i]->mutable_data_ptr();
357-
ET_LOG(Debug, " gpu_outputs[%d] = %p, data_ptr = %p",
358-
i, gpu_outputs[i], gpu_output_data);
417+
ET_LOG(
418+
Debug,
419+
" gpu_outputs[%d] = %p, data_ptr = %p",
420+
i,
421+
gpu_outputs[i],
422+
gpu_output_data);
359423
}
360424

361425
// Run AOTI container with GPU tensors
@@ -380,10 +444,15 @@ class ET_EXPERIMENTAL MetalBackend final : public ::executorch::runtime::Backend
380444
try {
381445
synchronize_metal_stream();
382446
} catch (const std::exception& e) {
383-
ET_LOG(Error, "Failed to synchronize Metal stream after kernel execution: %s", e.what());
447+
ET_LOG(
448+
Error,
449+
"Failed to synchronize Metal stream after kernel execution: %s",
450+
e.what());
384451
return Error::Internal;
385452
} catch (...) {
386-
ET_LOG(Error, "Failed to synchronize Metal stream after kernel execution: unknown exception");
453+
ET_LOG(
454+
Error,
455+
"Failed to synchronize Metal stream after kernel execution: unknown exception");
387456
return Error::Internal;
388457
}
389458

@@ -432,7 +501,7 @@ class ET_EXPERIMENTAL MetalBackend final : public ::executorch::runtime::Backend
432501
// which causes segmentation faults when attempting to delete other
433502
// containers. As a workaround, we skip explicit container deletion
434503
// and defer cleanup to the OS.
435-
// TODO(gasoonjia): Find a proper solution for safe container deletion.
504+
// TODO: Find a proper solution for safe container deletion.
436505
// AOTInductorModelContainerDelete(handle->container_handle);
437506

438507
// Now close the shared library

0 commit comments

Comments
 (0)