xref: /aosp_15_r20/external/executorch/backends/apple/coreml/runtime/delegate/backend_delegate.mm (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1//
2// backend_delegate.mm
3//
4// Copyright © 2024 Apple Inc. All rights reserved.
5//
6// Please refer to the license found in the LICENSE file in the root directory of the source tree.
7
8
9#import <ETCoreMLAssetManager.h>
10#import <ETCoreMLModel.h>
11#import <ETCoreMLModelManager.h>
12#import <ETCoreMLStrings.h>
13#import <backend_delegate.h>
14#import <model_event_logger.h>
15#import <multiarray.h>
16
17namespace  {
18using namespace executorchcoreml;
19
20MLComputeUnits get_compute_units(const Buffer& buffer) {
21    std::string value(reinterpret_cast<const char *>(buffer.data()), buffer.size());
22    if (value == std::string(ETCoreMLStrings.cpuComputeUnitName.UTF8String)) {
23        return MLComputeUnitsCPUOnly;
24    } else if (value == std::string(ETCoreMLStrings.cpuAndGpuComputeUnitsName.UTF8String)) {
25        return MLComputeUnitsCPUAndGPU;
26    } else if (value == std::string(ETCoreMLStrings.cpuAndNeuralEngineComputeUnitsName.UTF8String)) {
27        return MLComputeUnitsCPUAndNeuralEngine;
28    } else {
29        return MLComputeUnitsAll;
30    }
31}
32
33MLModelConfiguration *get_model_configuration(const std::unordered_map<std::string, Buffer>& specs) {
34    std::string key_name(ETCoreMLStrings.computeUnitsKeyName.UTF8String);
35    MLModelConfiguration *configuration = [[MLModelConfiguration alloc] init];
36    for (const auto& [key, buffer] : specs) {
37        if (key == key_name) {
38            configuration.computeUnits = get_compute_units(buffer);
39            break;
40        }
41    }
42
43    return configuration;
44}
45
46NSURL * _Nullable create_directory_if_needed(NSURL *url,
47                                             NSFileManager *fileManager,
48                                             NSError * __autoreleasing *error) {
49    if (![fileManager fileExistsAtPath:url.path] &&
50        ![fileManager createDirectoryAtURL:url withIntermediateDirectories:YES attributes:@{} error:error]) {
51        return nil;
52    }
53
54    return url;
55}
56
57ETCoreMLAssetManager * _Nullable create_asset_manager(NSString *assets_directory_path,
58                                                      NSString *trash_directory_path,
59                                                      NSString *database_directory_path,
60                                                      NSString *database_name,
61                                                      NSInteger max_assets_size_in_bytes,
62                                                      NSError * __autoreleasing *error) {
63    NSFileManager *fm  = [[NSFileManager alloc] init];
64
65    NSURL *assets_directory_url = [NSURL fileURLWithPath:assets_directory_path];
66    if (!create_directory_if_needed(assets_directory_url, fm, error)) {
67        return nil;
68    }
69
70    NSURL *trash_directory_url = [NSURL fileURLWithPath:trash_directory_path];
71    if (!create_directory_if_needed(trash_directory_url, fm, error)) {
72        return nil;
73    }
74
75    NSURL *database_directory_url = [NSURL fileURLWithPath:database_directory_path];
76    if (!create_directory_if_needed(database_directory_url, fm, error)) {
77        return nil;
78    }
79
80    NSURL *database_url = [database_directory_url URLByAppendingPathComponent:database_name];
81    ETCoreMLAssetManager *manager = [[ETCoreMLAssetManager alloc] initWithDatabaseURL:database_url
82                                                                   assetsDirectoryURL:assets_directory_url
83                                                                    trashDirectoryURL:trash_directory_url
84                                                                 maxAssetsSizeInBytes:max_assets_size_in_bytes
85                                                                                error:error];
86    return manager;
87}
88} //namespace
89
90@interface ETCoreMLModelManagerDelegate : NSObject
91
92- (instancetype)init NS_UNAVAILABLE;
93
94+ (instancetype)new NS_UNAVAILABLE;
95
96- (instancetype)initWithConfig:(BackendDelegate::Config)config NS_DESIGNATED_INITIALIZER;
97
98- (BOOL)loadAndReturnError:(NSError * _Nullable __autoreleasing *)error;
99
100- (void)loadAsynchronously;
101
102- (ModelHandle*)loadModelFromAOTData:(NSData*)data
103                       configuration:(MLModelConfiguration*)configuration
104                               error:(NSError* __autoreleasing*)error;
105
106- (BOOL)executeModelWithHandle:(ModelHandle*)handle
107                       argsVec:(const std::vector<executorchcoreml::MultiArray>&)argsVec
108                loggingOptions:(const executorchcoreml::ModelLoggingOptions&)loggingOptions
109                   eventLogger:(const executorchcoreml::ModelEventLogger* _Nullable)eventLogger
110                         error:(NSError* __autoreleasing*)error;
111
112- (BOOL)unloadModelWithHandle:(ModelHandle*)handle;
113
114- (BOOL)purgeModelsCacheAndReturnError:(NSError * _Nullable __autoreleasing *)error;
115
116@property (assign, readonly, nonatomic) BackendDelegate::Config config;
117@property (strong, readonly, nonatomic) dispatch_queue_t syncQueue;
118@property (strong, nonatomic, nullable) ETCoreMLModelManager *impl;
119@property (assign, readonly, nonatomic) BOOL isAvailable;
120
121@end
122
123@implementation ETCoreMLModelManagerDelegate
124
125- (instancetype)initWithConfig:(BackendDelegate::Config)config {
126    self = [super init];
127    if (self) {
128        _config = std::move(config);
129        _syncQueue = dispatch_queue_create("com.executorchcoreml.modelmanagerdelegate.sync", DISPATCH_QUEUE_SERIAL_WITH_AUTORELEASE_POOL);
130    }
131
132    return self;
133}
134
135- (BOOL)_loadAndReturnError:(NSError * _Nullable __autoreleasing *)error {
136    if (self.impl != nil) {
137        return YES;
138    }
139
140    ETCoreMLAssetManager *assetManager = create_asset_manager(ETCoreMLStrings.assetsDirectoryPath,
141                                                              ETCoreMLStrings.trashDirectoryPath,
142                                                              ETCoreMLStrings.databaseDirectoryPath,
143                                                              ETCoreMLStrings.databaseName,
144                                                              self.config.max_models_cache_size,
145                                                              error);
146    if (!assetManager) {
147        return NO;
148    }
149
150    ETCoreMLModelManager *modelManager = [[ETCoreMLModelManager alloc] initWithAssetManager:assetManager];
151    if (!modelManager) {
152        return NO;
153    }
154
155    self.impl = modelManager;
156
157    if (self.config.should_prewarm_asset) {
158        [modelManager prewarmRecentlyUsedAssetsWithMaxCount:1];
159    }
160
161    return YES;
162}
163
164- (BOOL)loadAndReturnError:(NSError * _Nullable __autoreleasing *)error {
165    __block NSError *localError = nil;
166    __block BOOL result = NO;
167    dispatch_sync(self.syncQueue, ^{
168        result = [self _loadAndReturnError:&localError];
169    });
170
171    if (error) {
172        *error = localError;
173    }
174
175    return result;
176}
177
178- (void)loadAsynchronously {
179    dispatch_async(self.syncQueue, ^{
180        (void)[self _loadAndReturnError:nil];
181    });
182}
183
184- (ModelHandle*)loadModelFromAOTData:(NSData*)data
185                       configuration:(MLModelConfiguration*)configuration
186                               error:(NSError* __autoreleasing*)error {
187    if (![self loadAndReturnError:error]) {
188        return nil;
189    }
190
191    auto handle = [self.impl loadModelFromAOTData:data
192                                    configuration:configuration
193                                            error:error];
194    if ((handle != NULL) && self.config.should_prewarm_model) {
195        [self.impl prewarmModelWithHandle:handle error:nil];
196    }
197
198    return handle;
199}
200
201- (BOOL)executeModelWithHandle:(ModelHandle*)handle
202                       argsVec:(const std::vector<executorchcoreml::MultiArray>&)argsVec
203                loggingOptions:(const executorchcoreml::ModelLoggingOptions&)loggingOptions
204                   eventLogger:(const executorchcoreml::ModelEventLogger* _Nullable)eventLogger
205                         error:(NSError* __autoreleasing*)error {
206    assert(self.impl != nil && "Impl must not be nil");
207    return [self.impl executeModelWithHandle:handle
208                                     argsVec:argsVec
209                              loggingOptions:loggingOptions
210                                 eventLogger:eventLogger
211                                       error:error];
212}
213
214- (nullable ETCoreMLModel*)modelWithHandle:(ModelHandle*)handle {
215    assert(self.impl != nil && "Impl must not be nil");
216    return [self.impl modelWithHandle:handle];
217}
218
219- (BOOL)unloadModelWithHandle:(ModelHandle*)handle {
220    assert(self.impl != nil && "Impl must not be nil");
221    return [self.impl unloadModelWithHandle:handle];
222}
223
224- (BOOL)purgeModelsCacheAndReturnError:(NSError * _Nullable __autoreleasing *)error {
225    if (![self loadAndReturnError:error]) {
226        return NO;
227    }
228
229    return [self.impl purgeModelsCacheAndReturnError:error];;
230}
231
232- (BOOL)isAvailable {
233    if (![self loadAndReturnError:nil]) {
234        return NO;
235    }
236
237    return YES;
238}
239
240@end
241
242namespace executorchcoreml {
243
244std::string BackendDelegate::ErrorCategory::message(int code) const {
245    switch (static_cast<ErrorCode>(code)) {
246        case ErrorCode::CorruptedData:
247            return "AOT blob can't be parsed";
248        case ErrorCode::CorruptedMetadata:
249            return "AOT blob has incorrect or missing metadata.";
250        case ErrorCode::CorruptedModel:
251            return "AOT blob has incorrect or missing CoreML model.";
252        case ErrorCode::BrokenModel:
253            return "CoreML model doesn't match the input and output specifications.";
254        case ErrorCode::CompilationFailed:
255            return "Failed to compile CoreML model.";
256        case ErrorCode::ModelSaveFailed:
257            return "Failed to write CoreML model to disk.";
258        case ErrorCode::ModelCacheCreationFailed:
259            return "Failed to create model cache.";
260        default:
261            return "Unexpected error.";
262    }
263}
264
265class BackendDelegateImpl: public BackendDelegate {
266public:
267    explicit BackendDelegateImpl(const Config& config) noexcept
268    :BackendDelegate(), model_manager_([[ETCoreMLModelManagerDelegate alloc] initWithConfig:config])
269    {
270        [model_manager_ loadAsynchronously];
271    }
272
273    BackendDelegateImpl(BackendDelegateImpl const&) = delete;
274    BackendDelegateImpl& operator=(BackendDelegateImpl const&) = delete;
275
276    Handle *init(Buffer processed,const std::unordered_map<std::string, Buffer>& specs) const noexcept override {
277        NSError *localError = nil;
278        MLModelConfiguration *configuration = get_model_configuration(specs);
279        NSData *data = [NSData dataWithBytesNoCopy:const_cast<void *>(processed.data())
280                                            length:processed.size()
281                                      freeWhenDone:NO];
282        ModelHandle *modelHandle = [model_manager_ loadModelFromAOTData:data
283                                                          configuration:configuration
284                                                                  error:&localError];
285        return modelHandle;
286    }
287
288    bool execute(Handle* handle,
289                 const std::vector<MultiArray>& args,
290                 const ModelLoggingOptions& logging_options,
291                 ModelEventLogger *event_logger,
292                 std::error_code& ec) const noexcept override {
293        NSError *error = nil;
294        if (![model_manager_ executeModelWithHandle:handle
295                                            argsVec:args
296                                     loggingOptions:logging_options
297                                        eventLogger:event_logger
298                                              error:&error]) {
299            ec = static_cast<ErrorCode>(error.code);
300            return false;
301        }
302
303        return true;
304    }
305
306    bool is_valid_handle(Handle* handle) const noexcept override {
307        return [model_manager_ modelWithHandle:handle] != nil;
308    }
309
310    bool is_available() const noexcept override {
311        return static_cast<bool>(model_manager_.isAvailable);
312    }
313
314    std::pair<size_t, size_t> get_num_arguments(Handle* handle) const noexcept override {
315        ETCoreMLModel *model = [model_manager_ modelWithHandle:handle];
316        return {model.orderedInputNames.count, model.orderedOutputNames.count};
317    }
318
319    void destroy(Handle* handle) const noexcept override {
320        [model_manager_ unloadModelWithHandle:handle];
321    }
322
323    bool purge_models_cache() const noexcept override {
324        NSError *localError = nil;
325        bool result = static_cast<bool>([model_manager_ purgeModelsCacheAndReturnError:&localError]);
326        return result;
327    }
328
329    ETCoreMLModelManagerDelegate *model_manager_;
330    Config config_;
331};
332
333std::shared_ptr<BackendDelegate> BackendDelegate::make(const Config& config) {
334    return std::make_shared<BackendDelegateImpl>(config);
335}
336} //namespace executorchcoreml
337