xref: /aosp_15_r20/external/executorch/backends/apple/coreml/runtime/sdk/ETCoreMLModelDebugger.mm (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1//
2// ETCoreMLModelDebugger.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#import "ETCoreMLModelDebugger.h"
9
10#import <CoreML/CoreML.h>
11#import "ETCoreMLAsset.h"
12#import "ETCoreMLAssetManager.h"
13#import "ETCoreMLLogging.h"
14#import "ETCoreMLModelCompiler.h"
15#import "ETCoreMLModelDebugInfo.h"
16#import "ETCoreMLModelStructurePath.h"
17#import "ETCoreMLPair.h"
18#import "ETCoreMLStrings.h"
19#import <format/MIL.pb.h>
20#import <format/Model.pb.h>
21#import <fstream>
22#import <iostream>
23#import "model_package_info.h"
24#import "objc_json_serde.h"
25#import <string>
26#import <unordered_map>
27
28typedef ETCoreMLPair<MLModel *, NSArray<ETCoreMLModelStructurePath *> *> DebuggableModel;
29
30namespace {
31using namespace executorchcoreml;
32using namespace executorchcoreml::modelstructure;
33using namespace CoreML::Specification;
34
35NSURL * _Nullable get_model_spec_url(NSURL *model_package_url,
36                                     NSFileManager *file_manager,
37                                     NSError* __autoreleasing *error) {
38    auto info = ModelPackageInfo::make(model_package_url, file_manager, error);
39    if (!info) {
40        return nil;
41    }
42
43    const auto& info_value = info.value();
44    auto it = info_value.items.find(info_value.root_model_identifier);
45    if (it == info_value.items.end()) {
46        ETCoreMLLogErrorAndSetNSError(error, 0, "%@ is broken, root model info doesn't exist.", model_package_url.lastPathComponent);
47        return nil;
48    }
49
50    auto path = it->second.path;
51    if (path.empty()) {
52        ETCoreMLLogErrorAndSetNSError(error, 0, "%@ is broken, root model path doesn't exist.", model_package_url.lastPathComponent);
53        return nil;
54    }
55
56    return [model_package_url URLByAppendingPathComponent:[NSString stringWithFormat:@"Data/%s", path.c_str()]];
57}
58
59std::optional<int> index_of_output(const MILSpec::Operation& operation, const std::string& output_name) {
60    for (int i = 0; i < operation.outputs_size(); i++) {
61        if (operation.outputs(i).name() == output_name) {
62            return i;
63        }
64    }
65
66    return std::nullopt;
67}
68
69BOOL is_const_operation(const MILSpec::Operation& operation) {
70    return operation.type() == "const";
71}
72
73BOOL is_datatype_supported_as_model_output(MILSpec::DataType datatype) {
74    switch (datatype) {
75        case MILSpec::DataType::INT32:
76            return true;
77        case MILSpec::DataType::FLOAT16:
78            return true;
79        case MILSpec::DataType::FLOAT32:
80            return true;
81        case MILSpec::DataType::FLOAT64:
82            return true;
83        default:
84            return false;
85    }
86}
87
88BOOL is_output_type_supported_as_model_output(const MILSpec::ValueType& type) {
89    return type.has_tensortype() && is_datatype_supported_as_model_output(type.tensortype().datatype());
90}
91
92BOOL is_operation_output_supported_as_model_output(const MILSpec::Operation& operation) {
93    if (is_const_operation(operation)) {
94        return NO;
95    }
96
97    return YES;
98}
99
100const MILSpec::NamedValueType *add_output(MILSpec::Block& block, const Path& path, size_t block_component_index) {
101    const auto& components = path.components();
102    auto block_component = std::get_if<Path::Program::Block>(&components[block_component_index]);
103    NSCAssert(block_component != nullptr, @"%@: Invalid path, component doesn't refer to a block.", NSStringFromClass(ETCoreMLModelDebugger.class));
104    // Next component must be an operation.
105    size_t operation_component_index = block_component_index + 1;
106    auto operation_component = std::get_if<Path::Program::Operation>(&components[operation_component_index]);
107    const auto& output_name = operation_component->output_name;
108
109    for (int i = 0; i < block.operations_size(); i++) {
110        auto& operation = *(block.mutable_operations(i));
111        auto output_index = index_of_output(operation, output_name);
112        if (!output_index) {
113            continue;
114        }
115
116        if (components.size() == operation_component_index + 1) {
117            const auto& output = operation.outputs(output_index.value());
118            if (!is_output_type_supported_as_model_output(output.type())) {
119                return nullptr;
120            }
121
122            block.add_outputs(output.name());
123            return &output;
124        }
125
126        // Handle nested block.
127        size_t nested_block_index = operation_component_index + 1;
128        auto nested_block_component = std::get_if<Path::Program::Block>(&components[nested_block_index]);
129        NSCAssert(nested_block_component != nullptr, @"%@: Invalid path, component doesn't refer to a nested block.", NSStringFromClass(ETCoreMLModelDebugger.class));
130        auto& nested_block = *(operation.mutable_blocks(static_cast<int>(nested_block_component->index)));
131        return add_output(nested_block, path, nested_block_index);
132    }
133
134    return nullptr;
135}
136
137const MILSpec::NamedValueType *add_output(MILSpec::Function& function, const Path& path, size_t function_component_index) {
138    size_t block_component_index = function_component_index + 1;
139    const auto& block_name = function.opset();
140    auto& block = (*function.mutable_block_specializations())[block_name];
141
142    return add_output(block, path, block_component_index);
143}
144
145const MILSpec::NamedValueType *add_output(MILSpec::Program& program, const Path& path) {
146    size_t function_component_index = 1;
147    const auto& components = path.components();
148    auto function_component = std::get_if<Path::Program::Function>(&components[function_component_index]);
149    NSCAssert(function_component != nullptr, @"%@: Invalid path, path doesn't refer to a function.", NSStringFromClass(ETCoreMLModelDebugger.class));
150    auto& functions = *(program.mutable_functions());
151    auto& function = functions[function_component->name];
152
153    return add_output(function, path, function_component_index);
154}
155
156std::optional<ArrayFeatureType_ArrayDataType> to_array_datatype(MILSpec::DataType datatype,
157                                                                int model_specification_version) {
158    switch (datatype) {
159        case MILSpec::DataType::INT32:
160            return ArrayFeatureType_ArrayDataType::ArrayFeatureType_ArrayDataType_INT32;
161        case MILSpec::DataType::FLOAT16:
162            return ArrayFeatureType_ArrayDataType::ArrayFeatureType_ArrayDataType_FLOAT16;
163        case MILSpec::DataType::FLOAT32:
164            return ArrayFeatureType_ArrayDataType::ArrayFeatureType_ArrayDataType_FLOAT32;
165        case MILSpec::DataType::FLOAT64:
166            return ArrayFeatureType_ArrayDataType::ArrayFeatureType_ArrayDataType_DOUBLE;
167        default:
168            return std::nullopt;
169    }
170}
171
172const MILSpec::NamedValueType *add_output(Model& model, const Path& path) {
173    NSCAssert(model.has_mlprogram(), @"%@: Model is not a ML Program.", NSStringFromClass(ETCoreMLModelDebugger.class));
174    auto& program = *(model.mutable_mlprogram());
175    auto output = add_output(program, path);
176    if (!output) {
177        return nullptr;
178    }
179
180    auto& description = *(model.mutable_description());
181    auto& output_feature = *(description.add_output());
182    output_feature.mutable_name()->assign(output->name());
183    auto& multi_array_type = *(output_feature.mutable_type()->mutable_multiarraytype());
184    NSCAssert(output->type().has_tensortype(), @"%@: Only a tensor type can be model output.", NSStringFromClass(ETCoreMLModelDebugger.class));
185    auto tensor_type = output->type().tensortype();
186    auto feature_data_type = to_array_datatype(tensor_type.datatype(), model.specificationversion());
187    NSCAssert(feature_data_type.has_value(), @"%@: Unsupported datatype.", NSStringFromClass(ETCoreMLModelDebugger.class));
188    multi_array_type.set_datatype(feature_data_type.value());
189
190    return output;
191}
192
193void visit_program_operation(const MILSpec::Block& block,
194                             const Path& block_path,
195                             BOOL (^handler)(const MILSpec::Operation& operation, ETCoreMLModelStructurePath *path)) {
196    for (int i = 0; i < block.operations_size(); ++i) {
197        const MILSpec::Operation& operation = block.operations(i);
198        Path operation_path = block_path;
199        if (operation.outputs_size() == 0) {
200            continue;
201        }
202        operation_path.append_component(Path::Program::Operation(operation.outputs(0).name()));
203        if (!handler(operation, [[ETCoreMLModelStructurePath alloc] initWithUnderlyingValue:operation_path])) {
204            return;
205        }
206
207        for (int j = 0; j < operation.blocks_size(); ++j) {
208            Path nested_block_path = operation_path;
209            nested_block_path.append_component(Path::Program::Block(j));
210            visit_program_operation(operation.blocks(j), nested_block_path, handler);
211        }
212    }
213}
214
215void visit_program_operation(Model& model, BOOL (^handler)(const MILSpec::Operation& operation, ETCoreMLModelStructurePath *path)) {
216    const auto& functions = model.mlprogram().functions();
217    for (const auto& [function_name, function] : functions) {
218        Path path;
219        path.append_component(Path::Program());
220        path.append_component(Path::Program::Function(function_name));
221        path.append_component(Path::Program::Block(-1));
222        const auto& blocks = function.block_specializations();
223        const auto& specialization = blocks.at(function.opset());
224        visit_program_operation(specialization, path, handler);
225    }
226}
227
228NSString *to_string(MLComputeUnits compute_units) {
229    switch (compute_units) {
230        case MLComputeUnitsAll: {
231            return ETCoreMLStrings.allComputeUnitsName;
232        }
233        case MLComputeUnitsCPUOnly: {
234            return ETCoreMLStrings.cpuComputeUnitName;
235        }
236        case MLComputeUnitsCPUAndGPU: {
237            return ETCoreMLStrings.cpuAndGpuComputeUnitsName;
238        }
239        case MLComputeUnitsCPUAndNeuralEngine: {
240            return ETCoreMLStrings.cpuAndNeuralEngineComputeUnitsName;
241        }
242        default: {
243            return ETCoreMLStrings.allComputeUnitsName;
244        }
245    }
246}
247
248NSString *get_asset_identifier(NSString *asset_identifier,
249                               MLComputeUnits compute_units,
250                               NSArray<ETCoreMLModelStructurePath *> *paths) {
251    size_t paths_hash = 0;
252    for (ETCoreMLModelStructurePath *path in paths) {
253        executorchcoreml::hash_combine(paths_hash, path.underlyingValue);
254    }
255
256    return [NSString stringWithFormat:@"%@_%zu_%@", asset_identifier, paths_hash, to_string(compute_units)];
257}
258
259std::unique_ptr<Model> parse_model_spec(NSURL *model_spec_url,
260                                        NSError * __autoreleasing *error) {
261    NSData *data = [NSData dataWithContentsOfURL:model_spec_url options:NSDataReadingMappedIfSafe error:error];
262    if (!data) {
263        return nullptr;
264    }
265
266    auto model = std::make_unique<Model>();
267    if (!model->ParseFromArray(data.bytes, (int)data.length)) {
268        return nullptr;
269    }
270
271    return model;
272}
273
274std::unique_ptr<Model> copy_model_spec(const Model& model_spec) {
275    auto model_spec_copy = std::make_unique<Model>();
276    model_spec_copy->CopyFrom(model_spec);
277
278    return model_spec_copy;
279}
280
281void update_model_spec_version_to_include_fp16_output(Model& model_spec) {
282    constexpr int minimum_spec_version_with_fp16_support = 7;
283    int spec_version = MAX(model_spec.specificationversion(), minimum_spec_version_with_fp16_support);
284    model_spec.set_specificationversion(spec_version);
285}
286
287NSURL * _Nullable get_compiled_model_url_with_intermediate_outputs(NSURL *model_url,
288                                                                   NSURL *model_spec_url,
289                                                                   const Model& model_spec,
290                                                                   NSOrderedSet<NSString *> *outputNames,
291                                                                   NSArray<ETCoreMLModelStructurePath *> *paths,
292                                                                   NSError * __autoreleasing *error) {
293    // Update model asset spec.
294    auto model_spec_copy = copy_model_spec(model_spec);
295    for (ETCoreMLModelStructurePath *path in paths) {
296        if ([outputNames containsObject:path.operationOutputName]) {
297            continue;
298        }
299        add_output(*model_spec_copy.get(), path.underlyingValue);
300    }
301
302    update_model_spec_version_to_include_fp16_output(*model_spec_copy);
303    int size = model_spec_copy->ByteSize();
304    NSMutableData *data = [NSMutableData dataWithLength:size];
305    if (!model_spec_copy->SerializeToArray(data.mutableBytes, size)) {
306        return nil;
307    }
308
309    if (![data writeToURL:model_spec_url options:NSDataWritingAtomic error:error]) {
310        return nil;
311    }
312
313    return [ETCoreMLModelCompiler compileModelAtURL:model_url
314                               maxWaitTimeInSeconds:(5 * 60)
315                                              error:error];
316}
317
318ETCoreMLAsset * _Nullable make_asset(NSURL *asset_url,
319                                     NSString *identifier,
320                                     NSFileManager *fm,
321                                     NSError * __autoreleasing *error) {
322    auto underlying_asset = Asset::make(asset_url, identifier, fm, error);
323    if (!underlying_asset) {
324        return nil;
325    }
326
327    ETCoreMLAsset *asset = [[ETCoreMLAsset alloc] initWithBackingAsset:std::move(underlying_asset.value())];
328    if (![asset keepAliveAndReturnError:error]) {
329        return nil;
330    }
331
332    return asset;
333}
334
335NSArray<NSString *> *get_output_names(NSArray<ETCoreMLModelStructurePath *> *paths) {
336    NSMutableArray<NSString *> *result = [NSMutableArray arrayWithCapacity:paths.count];
337    for (ETCoreMLModelStructurePath *path in paths) {
338        NSString *output_name = path.operationOutputName;
339        if (output_name) {
340            [result addObject:output_name];
341        }
342    }
343
344    return result;
345}
346
347void set_model_outputs(id<MLFeatureProvider> output_features,
348                       NSOrderedSet<NSString *> *output_names,
349                       NSArray<MLMultiArray *> *_Nullable __autoreleasing *_Nonnull model_outputs) {
350    NSMutableArray<MLMultiArray *> *values = [NSMutableArray arrayWithCapacity:output_names.count];
351    for (NSString *output_name in output_names) {
352        MLFeatureValue *feature_value = [output_features featureValueForName:output_name];
353        NSCAssert(feature_value.multiArrayValue != nil, @"%@: Expected a multiarray value for output name=%@.",
354                  NSStringFromClass(ETCoreMLModelDebugger.class),
355                  output_name);
356        [values addObject:feature_value.multiArrayValue];
357    }
358
359    *model_outputs = values;
360}
361
362void set_intermediate_outputs(id<MLFeatureProvider> output_features,
363                              NSArray<ETCoreMLModelStructurePath *> *paths,
364                              NSMutableDictionary<ETCoreMLModelStructurePath *, MLMultiArray *> *result) {
365    for (ETCoreMLModelStructurePath *path in paths) {
366        NSString *output_name = path.operationOutputName;
367        if (!output_name) {
368            continue;
369        }
370
371        MLFeatureValue *feature_value = [output_features featureValueForName:output_name];
372        if (!feature_value) {
373            continue;
374        }
375        MLMultiArray *multi_array = feature_value.multiArrayValue;
376        result[path] = multi_array;
377    }
378}
379
380NSArray<ETCoreMLModelStructurePath *> *get_operation_dependencies(const MILSpec::Operation &operation,
381                                                                  ETCoreMLModelStructurePath *path,
382                                                                  NSSet<ETCoreMLModelStructurePath *> *paths) {
383    const auto& inputs = operation.inputs();
384    const auto cppPath = path.underlyingValue;
385    NSMutableArray<ETCoreMLModelStructurePath *> *deps = [NSMutableArray arrayWithCapacity:inputs.size()];
386    for (const auto& [_, arg] : inputs) {
387        const auto& bindings = arg.arguments();
388        for (const auto& binding : bindings) {
389            if (binding.has_value()) {
390                continue;
391            }
392
393            const auto& name = binding.name();
394            auto dep = cppPath;
395            dep.remove_last_component();
396            dep.append_component(Path::Program::Operation(name));
397            ETCoreMLModelStructurePath *path = [[ETCoreMLModelStructurePath alloc] initWithUnderlyingValue:dep];
398            if ([paths containsObject:path]) {
399                [deps addObject:path];
400            }
401        }
402    }
403
404    return deps;
405}
406
407NSDictionary<NSString *, NSArray<ETCoreMLModelStructurePath *> *> *get_debug_handle_to_operation_paths_map(ETCoreMLModelDebugInfo *debug_info) {
408    NSUInteger capacity = debug_info.pathToDebugHandlesMap.count;
409    NSMutableDictionary<NSString *, NSMutableArray<ETCoreMLModelStructurePath *> *> *result = [NSMutableDictionary dictionaryWithCapacity:capacity];
410    [debug_info.pathToDebugHandlesMap enumerateKeysAndObjectsUsingBlock:^(ETCoreMLModelStructurePath *path,
411                                                                          NSArray<NSString *> *debug_handles,
412                                                                          BOOL * _Nonnull __unused stop) {
413        for (NSString *debug_handle in debug_handles) {
414            NSMutableArray<ETCoreMLModelStructurePath *> *paths = result[debug_handle];
415            if (!paths) {
416                paths = [NSMutableArray array];
417                result[debug_handle] = paths;
418            }
419
420            [paths addObject:path];
421        }
422
423    }];
424
425    return result;
426}
427
428BOOL is_node_terminal_node(ETCoreMLModelStructurePath *node,
429                           NSArray<ETCoreMLModelStructurePath *> *nodes,
430                           NSDictionary<ETCoreMLModelStructurePath *, NSArray<ETCoreMLModelStructurePath *> *> *dependencies) {
431    NSMutableSet<ETCoreMLModelStructurePath *> *nodes_dependencies = [NSMutableSet set];
432    for (ETCoreMLModelStructurePath *current_node in nodes) {
433        if ([current_node isEqual:node]) {
434            continue;
435        }
436        NSArray<ETCoreMLModelStructurePath *> *node_dependencies = dependencies[current_node];
437        if (node_dependencies.count > 0) {
438            [nodes_dependencies addObjectsFromArray:node_dependencies];
439        }
440    }
441
442    return ![nodes_dependencies containsObject:node];
443}
444
445ETCoreMLModelStructurePath *_Nullable find_terminal_node_from_nodes(NSArray<ETCoreMLModelStructurePath *> *nodes,
446                                                                    NSDictionary<ETCoreMLModelStructurePath *, NSArray<ETCoreMLModelStructurePath *> *> *dependencies) {
447    if (nodes.count < 2) {
448        return nodes.firstObject;
449    }
450
451    for (ETCoreMLModelStructurePath *node in nodes) {
452        if (is_node_terminal_node(node, nodes, dependencies)) {
453            return node;
454        }
455    }
456
457    return nil;
458}
459
460NSDictionary<ETCoreMLModelStructurePath *, NSString *> *get_operation_path_to_debug_symbol_map(ETCoreMLModelDebugInfo *model_debug_info,
461                                                                                               NSDictionary<NSString *, NSArray<ETCoreMLModelStructurePath *> *> *debug_handle_to_operation_paths_map,
462                                                                                               NSDictionary<ETCoreMLModelStructurePath *, NSArray<ETCoreMLModelStructurePath *> *> *dependencies) {
463    // When decomposing an EXIR operation into a MIL graph, it is essential to report the output of the terminal node of the MIL graph.
464    // This output corresponds directly to the output of the original EXIR operation.
465    NSUInteger capacity = debug_handle_to_operation_paths_map.count;
466    NSMutableDictionary<ETCoreMLModelStructurePath *, NSString *> *operation_path_to_debug_symbol_map = [NSMutableDictionary dictionaryWithCapacity:capacity];
467    [debug_handle_to_operation_paths_map enumerateKeysAndObjectsUsingBlock:^(NSString *debug_handle,
468                                                                             NSArray<ETCoreMLModelStructurePath *> *operation_paths,
469                                                                             BOOL * __unused stop) {
470        ETCoreMLModelStructurePath *operation_path = find_terminal_node_from_nodes(operation_paths, dependencies);
471        NSString *debug_symbol = (operation_path != nil) ? model_debug_info.pathToDebugSymbolMap[operation_path] : nil;
472        if (debug_symbol) {
473            operation_path_to_debug_symbol_map[operation_path] = debug_symbol;
474        }
475
476    }];
477
478    return operation_path_to_debug_symbol_map;
479}
480
481}
482
483@interface ETCoreMLModelDebugger ()
484/// The model output names.
485@property (readonly, copy, nonatomic) NSOrderedSet<NSString *> *outputNames;
486/// The model asset.
487@property (readonly, copy, nonatomic) ETCoreMLAsset *modelAsset;
488/// The model debug info.
489@property (readonly, copy, nonatomic, nullable) ETCoreMLModelDebugInfo *modelDebugInfo;
490/// The asset manager.
491@property (readonly, copy, nonatomic) ETCoreMLAssetManager *assetManager;
492/// The model configuration.
493@property (readonly, strong, nonatomic) MLModelConfiguration *configuration;
494/// The url to the model specification.
495@property (readonly, copy, nonatomic) NSURL *modelSpecURL;
496
497@end
498
499@implementation ETCoreMLModelDebugger {
500    std::unique_ptr<Model> _modelSpec;
501}
502
503- (nullable instancetype)initWithModelAsset:(ETCoreMLAsset *)modelAsset
504                             modelDebugInfo:(nullable ETCoreMLModelDebugInfo *)modelDebugInfo
505                                outputNames:(NSOrderedSet<NSString *> *)outputNames
506                              configuration:(MLModelConfiguration *)configuration
507                               assetManager:(ETCoreMLAssetManager *)assetManager
508                                      error:(NSError * __autoreleasing *)error {
509    if (![modelAsset keepAliveAndReturnError:error]) {
510        return nil;
511    }
512
513    NSFileManager *fileManager = [[NSFileManager alloc] init];
514    NSURL *modelSpecURL = get_model_spec_url(modelAsset.contentURL, fileManager, error);
515    if (!modelSpecURL) {
516        return nil;
517    }
518
519    auto modelSpec = parse_model_spec(modelSpecURL, error);
520    if (!modelSpec) {
521        return nil;
522    }
523
524    __block NSMutableDictionary<ETCoreMLModelStructurePath *, NSArray<ETCoreMLModelStructurePath *> *> *dependencies = [NSMutableDictionary dictionary];
525    __block NSMutableArray<ETCoreMLModelStructurePath *> *operationPaths = [NSMutableArray array];
526    __block NSMutableSet<ETCoreMLModelStructurePath *> *allOperationPaths = [NSMutableSet set];
527    visit_program_operation(*modelSpec, ^BOOL(const MILSpec::Operation &operation, ETCoreMLModelStructurePath *operationPath) {
528        dependencies[operationPath] = get_operation_dependencies(operation, operationPath, allOperationPaths);
529        [allOperationPaths addObject:operationPath];
530        if (is_operation_output_supported_as_model_output(operation)) {
531            [operationPaths addObject:operationPath];
532        }
533
534        return YES;
535    });
536
537
538    NSDictionary<NSString *, NSArray<ETCoreMLModelStructurePath *> *> *debugHandleToOperationPathsMap = get_debug_handle_to_operation_paths_map(modelDebugInfo);
539
540    NSDictionary<ETCoreMLModelStructurePath *, NSString *> *operationPathToDebugSymbolMap = get_operation_path_to_debug_symbol_map(modelDebugInfo,
541                                                                                                                                   debugHandleToOperationPathsMap,
542                                                                                                                                   dependencies);
543
544    self = [super init];
545    if (self) {
546        _modelAsset = modelAsset;
547        _configuration = configuration;
548        _outputNames = [outputNames copy];
549        _assetManager = assetManager;
550        _modelSpec = std::move(modelSpec);
551        _modelSpecURL = modelSpecURL;
552        _operationPaths = operationPaths;
553        _operationPathToDebugSymbolMap = operationPathToDebugSymbolMap;
554        _modelDebugInfo = modelDebugInfo;
555    }
556
557    return self;
558}
559
560- (nullable ETCoreMLAsset *)compiledModelAssetWithOutputsAtPaths:(NSArray<ETCoreMLModelStructurePath *> *)paths
561                                                           error:(NSError* __autoreleasing *)error {
562    NSString *identifier = get_asset_identifier(self.modelAsset.identifier,
563                                                self.configuration.computeUnits,
564                                                paths);
565    NSError *localError = nil;
566    ETCoreMLAsset *compiledAsset = [self.assetManager assetWithIdentifier:identifier error:&localError];
567    if (compiledAsset) {
568        return compiledAsset;
569    }
570
571    if (localError) {
572        ETCoreMLLogError(localError,
573                         "%@: Failed to retrieve asset with identifier=%@",
574                         NSStringFromClass(ETCoreMLModelDebugger.class),
575                         identifier);
576    }
577
578    NSURL *compiledModelURL = get_compiled_model_url_with_intermediate_outputs(self.modelAsset.contentURL,
579                                                                               self.modelSpecURL,
580                                                                               *(_modelSpec.get()),
581                                                                               self.outputNames,
582                                                                               paths,
583                                                                               error);
584    if (!compiledModelURL) {
585        return nil;
586    }
587
588    compiledAsset = [self.assetManager storeAssetAtURL:compiledModelURL
589                                        withIdentifier:identifier
590                                                 error:&localError];
591
592    if (compiledAsset) {
593        return compiledAsset;
594    }
595
596    if (localError) {
597        ETCoreMLLogError(localError,
598                         "%@: Failed to store asset with identifier=%@",
599                         NSStringFromClass(ETCoreMLModelDebugger.class),
600                         identifier);
601    }
602
603    return make_asset(compiledModelURL, identifier, self.assetManager.fileManager, error);
604}
605
606- (nullable NSArray<DebuggableModel *> *)_modelsWithOutputsOfOperationsAtPath:(NSArray<ETCoreMLModelStructurePath *> *)paths
607                                                                        error:(NSError* __autoreleasing *)error {
608    if (paths.count == 0) {
609        return @[];
610    }
611
612    ETCoreMLAsset *compiledAsset = [self compiledModelAssetWithOutputsAtPaths:paths error:error];
613    if (!compiledAsset) {
614        return nil;
615    }
616
617    NSError *localError = nil;
618    MLModel *model = [MLModel modelWithContentsOfURL:compiledAsset.contentURL
619                                       configuration:self.configuration
620                                               error:&localError];
621    if (model) {
622        DebuggableModel *pair = [[ETCoreMLPair alloc] initWithFirst:model second:paths];
623        return @[pair];
624    }
625
626    if (localError) {
627        ETCoreMLLogError(localError, "%@: Failed to load model with outputs=%@",
628                         NSStringFromClass(ETCoreMLModelDebugger.class),
629                         get_output_names(paths));
630    }
631
632    if ([self.assetManager removeAssetWithIdentifier:compiledAsset.identifier error:&localError]) {
633        ETCoreMLLogError(localError, "%@: Failed to remove compiled asset with identifier=%@",
634                         NSStringFromClass(ETCoreMLModelDebugger.class),
635                         compiledAsset.identifier);
636    }
637
638    if (paths.count == 1) {
639        *error = localError;
640        return nil;
641    }
642
643    // There is a chance that the model compilation fails because of the number of outputs. In this case, we divide the paths into two and try again.
644    NSArray<ETCoreMLModelStructurePath *> *leftPaths = [paths subarrayWithRange:NSMakeRange(0, paths.count/2)];
645    NSArray<ETCoreMLModelStructurePath *> *rightPaths = [paths subarrayWithRange:NSMakeRange(paths.count/2, paths.count - paths.count/2)];
646    NSArray<DebuggableModel *> *leftModels = [self modelsWithOutputsOfOperationsAtPath:leftPaths error:&localError];
647    NSArray<DebuggableModel *> *rightModels = [self modelsWithOutputsOfOperationsAtPath:rightPaths error:&localError];
648    if (leftModels.count == 0 && rightModels.count == 0) {
649        *error = localError;
650        return nil;
651    }
652
653    NSArray<DebuggableModel *> *models = [(leftModels == nil ? @[] : leftModels) arrayByAddingObjectsFromArray:(rightModels == nil ? @[] : rightModels)];
654    return models;
655}
656
657- (nullable NSArray<DebuggableModel *> *)modelsWithOutputsOfOperationsAtPath:(NSArray<ETCoreMLModelStructurePath *> *)paths
658                                                                       error:(NSError* __autoreleasing *)error {
659    @autoreleasepool {
660        return [self _modelsWithOutputsOfOperationsAtPath:paths error:error];
661    }
662}
663
664- (nullable ETCoreMLModelOutputs *)outputsOfOperationsAtPaths:(NSArray<ETCoreMLModelStructurePath *> *)paths
665                                                      options:(MLPredictionOptions *)options
666                                                       inputs:(id<MLFeatureProvider>)inputs
667                                                 modelOutputs:(NSArray<MLMultiArray *> *_Nullable __autoreleasing *_Nonnull)modelOutputs
668                                                        error:(NSError* __autoreleasing *)error {
669    NSArray<MLMultiArray *> *lModelOutputs = nil;
670    NSMutableDictionary<ETCoreMLModelStructurePath *, MLMultiArray *> *result = [NSMutableDictionary dictionaryWithCapacity:paths.count];
671    @autoreleasepool {
672        NSArray<DebuggableModel *> *models = [self modelsWithOutputsOfOperationsAtPath:paths error:error];
673        if (!models) {
674            return nil;
675        }
676
677        for (DebuggableModel *pair in models) {
678            id<MLFeatureProvider> outputFeatures = [pair.first predictionFromFeatures:inputs options:options error:error];
679            set_intermediate_outputs(outputFeatures, paths, result);
680            if (modelOutputs) {
681                set_model_outputs(outputFeatures, self.outputNames, &lModelOutputs);
682            }
683        }
684    }
685
686    if (modelOutputs) {
687        *modelOutputs = lModelOutputs;
688    }
689
690    return result;
691}
692
693@end
694