xref: /aosp_15_r20/external/executorch/backends/apple/coreml/runtime/sdk/ETCoreMLModelProfiler.mm (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1//
2// ETCoreMLModelProfiler.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 "ETCoreMLModelProfiler.h"
9
10#import "ETCoreMLAsset.h"
11#import "ETCoreMLModel.h"
12#import "ETCoreMLLogging.h"
13#import "ETCoreMLModelStructurePath.h"
14#import "ETCoreMLOperationProfilingInfo.h"
15#import "ETCoreMLPair.h"
16#import "ETCoreMLStrings.h"
17#import <mach/mach_time.h>
18#import <math.h>
19#import "program_path.h"
20
21namespace  {
22using namespace executorchcoreml::modelstructure;
23
24#if MODEL_PROFILING_IS_AVAILABLE
25
26API_AVAILABLE(macos(14.4), ios(17.4), tvos(17.4), watchos(10.4))
27MLComputePlan *_Nullable get_compute_plan_of_model_at_url(NSURL *model_url,
28                                                          MLModelConfiguration *configuration,
29                                                          NSError* __autoreleasing *error) {
30    __block NSError *local_error = nil;
31    __block MLComputePlan *result = nil;
32    dispatch_semaphore_t sema = dispatch_semaphore_create(0);
33    [MLComputePlan loadContentsOfURL:model_url configuration:configuration completionHandler:^(MLComputePlan * _Nullable compute_plan,
34                                                                                               NSError * _Nullable compute_plan_error) {
35        result = compute_plan;
36        local_error = compute_plan_error;
37        dispatch_semaphore_signal(sema);
38    }];
39
40    long status = dispatch_semaphore_wait(sema, dispatch_time(DISPATCH_TIME_NOW, (int64_t)(5 * 60 * NSEC_PER_SEC)));
41    if (status != 0) {
42        ETCoreMLLogUnderlyingErrorAndSetNSError(error,
43                                                ETCoreMLErrorCompilationFailed,
44                                                local_error,
45                                                "%@: Failed to get compute plan of model with name=%@.",
46                                                NSStringFromClass(ETCoreMLModelProfiler.class),
47                                                model_url.lastPathComponent);
48        return nil;
49    }
50
51    return result;
52}
53
54API_AVAILABLE(macos(14.4), ios(17.4), tvos(17.4), watchos(10.4))
55void visit_program_operation(MLModelStructureProgramBlock *block,
56                             const Path& block_path,
57                             BOOL (^handler)(MLModelStructureProgramOperation *operation, ETCoreMLModelStructurePath *path)) {
58    for (MLModelStructureProgramOperation *operation in block.operations) {
59        Path operation_path = block_path;
60        operation_path.append_component(Path::Program::Operation(operation.outputs.firstObject.name.UTF8String));
61        if (!handler(operation, [[ETCoreMLModelStructurePath alloc] initWithUnderlyingValue:operation_path])) {
62            return;
63        }
64
65        for (size_t i = 0; i < operation.blocks.count; ++i) {
66            Path nested_block_path = operation_path;
67            nested_block_path.append_component(Path::Program::Block(i));
68            visit_program_operation(operation.blocks[i], nested_block_path,handler);
69        }
70    }
71}
72
73API_AVAILABLE(macos(14.4), ios(17.4), tvos(17.4), watchos(10.4))
74void visit_program_operation(MLModelStructure *modelStructure, BOOL (^handler)(MLModelStructureProgramOperation *operation, ETCoreMLModelStructurePath *path)) {
75    using namespace executorchcoreml::modelstructure;
76    [modelStructure.program.functions enumerateKeysAndObjectsUsingBlock:^(NSString *function_name,
77                                                                          MLModelStructureProgramFunction *function,
78                                                                          BOOL * _Nonnull __unused stop) {
79        Path path;
80        path.append_component(Path::Program());
81        path.append_component(Path::Program::Function(function_name.UTF8String));
82        path.append_component(Path::Program::Block(-1));
83        visit_program_operation(function.block, path, handler);
84    }];
85}
86
87API_AVAILABLE(macos(14.4), ios(17.4), tvos(17.4), watchos(10.4))
88ETCoreMLComputeUnits to_compute_unit(id<MLComputeDeviceProtocol> compute_device) {
89    if ([compute_device isKindOfClass:MLCPUComputeDevice.class]) {
90        return ETCoreMLComputeUnitCPU;
91    } else if ([compute_device isKindOfClass:MLGPUComputeDevice.class]) {
92        return ETCoreMLComputeUnitGPU;
93    } else if ([compute_device isKindOfClass:MLNeuralEngineComputeDevice.class]) {
94        return ETCoreMLComputeUnitNeuralEngine;
95    } else {
96        return ETCoreMLComputeUnitUnknown;
97    }
98}
99
100API_AVAILABLE(macos(14.4), ios(17.4), tvos(17.4), watchos(10.4))
101ETCoreMLComputeUnits to_compute_units(NSArray<id<MLComputeDeviceProtocol>> *compute_devices) {
102    ETCoreMLComputeUnits units = ETCoreMLComputeUnitUnknown;
103    for (id<MLComputeDeviceProtocol> compute_device in compute_devices) {
104        units |= to_compute_unit(compute_device);
105    }
106
107    return units;
108}
109
110API_AVAILABLE(macos(14.4), ios(17.4), tvos(17.4), watchos(10.4))
111ETCoreMLOperationProfilingInfo *get_profiling_info(MLComputePlanDeviceUsage *device_usage,
112                                                   MLModelStructureProgramOperation *operation,
113                                                   uint64_t op_execution_start_time,
114                                                   uint64_t op_execution_end_time,
115                                                   double estimatedCost) {
116    NSMutableArray<NSString *> *outputNames = [[NSMutableArray alloc] initWithCapacity:operation.outputs.count];
117    for (MLModelStructureProgramNamedValueType *output in operation.outputs) {
118        [outputNames addObject:output.name];
119    }
120
121    ETCoreMLComputeUnits preferred_compute_unit = to_compute_unit(device_usage.preferredComputeDevice);
122    ETCoreMLComputeUnits supported_compute_units = to_compute_units(device_usage.supportedComputeDevices);
123    ETCoreMLOperationProfilingInfo *info = [[ETCoreMLOperationProfilingInfo alloc] initWithPreferredComputeUnit:preferred_compute_unit
124                                                                                          supportedComputeUnits:supported_compute_units
125                                                                                    estimatedExecutionStartTime:op_execution_start_time
126                                                                                      estimatedExecutionEndTime:op_execution_end_time
127                                                                                                  estimatedCost:estimatedCost
128                                                                                                    outputNames:outputNames
129                                                                                                   operatorName:operation.operatorName];
130    return info;
131}
132
133API_AVAILABLE(macos(14.4), ios(17.4), tvos(17.4), watchos(10.4))
134NSDictionary<ETCoreMLModelStructurePath *, ETCoreMLOperationProfilingInfo *> *
135prepare_profiling_infos(NSArray<MLModelStructureProgramOperation *> *operations,
136                        NSDictionary<NSValue *, ETCoreMLModelStructurePath *> *operation_to_path_map,
137                        MLComputePlan *compute_plan) {
138    NSMutableDictionary<ETCoreMLModelStructurePath *, ETCoreMLOperationProfilingInfo *> *result = [NSMutableDictionary dictionaryWithCapacity:operation_to_path_map.count];
139    for (MLModelStructureProgramOperation *operation in operations) {
140        MLComputePlanCost *estimated_cost = [compute_plan estimatedCostOfMLProgramOperation:operation];
141        if (!estimated_cost || std::isnan(estimated_cost.weight)) {
142            continue;
143        }
144
145        NSValue *key = [NSValue valueWithPointer:(const void*)operation];
146        ETCoreMLModelStructurePath *path = operation_to_path_map[key];
147        MLComputePlanDeviceUsage *device_usage = [compute_plan computeDeviceUsageForMLProgramOperation:operation];
148        if (path && device_usage) {
149            ETCoreMLOperationProfilingInfo *profiling_info = get_profiling_info(device_usage,
150                                                                                operation,
151                                                                                0,
152                                                                                0,
153                                                                                estimated_cost.weight);
154            result[path] = profiling_info;
155        }
156    }
157
158    return result;
159}
160
161API_AVAILABLE(macos(14.4), ios(17.4), tvos(17.4), watchos(10.4))
162NSDictionary<ETCoreMLModelStructurePath *, ETCoreMLOperationProfilingInfo *> *
163get_profiling_infos_for_paths(NSSet<ETCoreMLModelStructurePath *> *paths,
164                              NSArray<MLModelStructureProgramOperation *> *topologically_sorted_operations,
165                              NSDictionary<NSValue *, ETCoreMLModelStructurePath *> *operation_to_path_map,
166                              NSDictionary<ETCoreMLModelStructurePath *, ETCoreMLOperationProfilingInfo *> *profiling_infos,
167                              uint64_t model_execution_start_time,
168                              uint64_t model_execution_end_time) {
169    NSMutableDictionary<ETCoreMLModelStructurePath *, ETCoreMLOperationProfilingInfo *> *result = [NSMutableDictionary dictionaryWithCapacity:paths.count];
170    uint64_t op_execution_start_time = model_execution_start_time;
171    uint64_t op_execution_end_time = model_execution_start_time;
172    // `model_execution_end_time` >= `model_execution_start_time`.
173    uint64_t model_execution_duration = model_execution_end_time - model_execution_start_time;
174    for (MLModelStructureProgramOperation *operation in topologically_sorted_operations) {
175        NSValue *key = [NSValue valueWithPointer:(const void*)operation];
176        ETCoreMLModelStructurePath *path = operation_to_path_map[key];
177        ETCoreMLOperationProfilingInfo *profiling_info = profiling_infos[path];
178        if (!profiling_info) {
179            continue;
180        }
181
182        op_execution_end_time = op_execution_start_time + static_cast<uint64_t>(static_cast<double>(model_execution_duration) * profiling_info.estimatedCost);
183        if (path && [paths containsObject:path]) {
184            ETCoreMLOperationProfilingInfo *profiling_info_new = [[ETCoreMLOperationProfilingInfo alloc] initWithPreferredComputeUnit:profiling_info.preferredComputeUnit
185                                                                                                                supportedComputeUnits:profiling_info.supportedComputeUnits
186                                                                                                          estimatedExecutionStartTime:op_execution_start_time
187                                                                                                            estimatedExecutionEndTime:op_execution_end_time
188                                                                                                                        estimatedCost:profiling_info.estimatedCost
189                                                                                                                          outputNames:profiling_info.outputNames
190                                                                                                                         operatorName:profiling_info.operatorName
191                                                                                                                             metadata:profiling_info.metadata];
192            result[path] = profiling_info_new;
193        }
194        op_execution_start_time = op_execution_end_time;
195    }
196
197    return result;
198}
199
200API_AVAILABLE(macos(14.4), ios(17.4), tvos(17.4), watchos(10.4))
201BOOL is_const_operation(MLModelStructureProgramOperation *operation) {
202    return [operation.operatorName isEqualToString:@"const"];
203}
204
205void set_model_outputs(id<MLFeatureProvider> output_features,
206                       NSOrderedSet<NSString *> *output_names,
207                       NSArray<MLMultiArray *> *_Nullable __autoreleasing *_Nonnull model_outputs) {
208    NSMutableArray<MLMultiArray *> *values = [NSMutableArray arrayWithCapacity:output_names.count];
209    for (NSString *output_name in output_names) {
210        MLFeatureValue *feature_value = [output_features featureValueForName:output_name];
211        NSCAssert(feature_value.multiArrayValue != nil, @"%@: Expected a multiarray value for output name=%@.",
212                  NSStringFromClass(ETCoreMLModelProfiler.class),
213                  output_name);
214        [values addObject:feature_value.multiArrayValue];
215    }
216
217    *model_outputs = values;
218}
219
220#endif
221
222}
223
224@interface ETCoreMLModelProfiler ()
225/// The model.
226@property (readonly, strong, nonatomic) ETCoreMLModel *model;
227/// The model output names.
228@property (readonly, copy, nonatomic) NSOrderedSet<NSString *> *outputNames;
229#if MODEL_PROFILING_IS_AVAILABLE
230/// The compute plan.
231@property (readonly, strong, nonatomic) MLComputePlan *computePlan API_AVAILABLE(macos(14.4), ios(17.4), tvos(17.4), watchos(10.4));
232/// The topologically sorted operations.
233@property (readonly, copy, nonatomic) NSArray<MLModelStructureProgramOperation *> *topologicallySortedOperations API_AVAILABLE(macos(14.4), ios(17.4), tvos(17.4), watchos(10.4));
234#endif
235/// The mapping from operation to it's path in the model structure.
236@property (readonly, strong, nonatomic) NSDictionary<NSValue *, ETCoreMLModelStructurePath *> *operationToPathMap;
237/// The profiling infos for all the operations.
238@property (readonly, copy, nonatomic) NSDictionary<ETCoreMLModelStructurePath *, ETCoreMLOperationProfilingInfo *> *profilingInfos;
239
240@end
241
242@implementation ETCoreMLModelProfiler
243
244- (nullable instancetype)initWithModel:(ETCoreMLModel *)model
245                         configuration:(MLModelConfiguration *)configuration
246                                 error:(NSError * __autoreleasing *)error  {
247#if MODEL_PROFILING_IS_AVAILABLE
248    if (@available(macOS 14.4, iOS 17.4, tvOS 17.4, watchOS 10.4, *)) {
249        NSURL *compiledModelURL = model.asset.contentURL;
250        MLComputePlan *computePlan = get_compute_plan_of_model_at_url(compiledModelURL,
251                                                                      configuration,
252                                                                      error);
253        if (!computePlan) {
254            return nil;
255        }
256
257        __block NSMutableArray<ETCoreMLModelStructurePath *> *operationPaths = [NSMutableArray array];
258        __block NSMutableDictionary<NSValue *, ETCoreMLModelStructurePath *> *operationToPathMap = [NSMutableDictionary dictionary];
259        __block NSMutableArray<MLModelStructureProgramOperation *> *topologicallySortedOperations = [NSMutableArray new];
260        visit_program_operation(computePlan.modelStructure, ^BOOL(MLModelStructureProgramOperation *operation, ETCoreMLModelStructurePath *operationPath) {
261            if (is_const_operation(operation)) {
262                return YES;
263            }
264
265            [topologicallySortedOperations addObject:operation];
266            NSValue *key = [NSValue valueWithPointer:(const void*)operation];
267            operationToPathMap[key] = operationPath;
268            [operationPaths addObject:operationPath];
269            return YES;
270        });
271
272        NSDictionary<ETCoreMLModelStructurePath *, ETCoreMLOperationProfilingInfo *> *profilingInfos = prepare_profiling_infos(topologicallySortedOperations,
273                                                                                                                               operationToPathMap,
274                                                                                                                               computePlan);
275
276        self = [super init];
277        if (self) {
278            _model = model;
279            _computePlan = computePlan;
280            _operationToPathMap = operationToPathMap;
281            _topologicallySortedOperations = topologicallySortedOperations;
282            _operationPaths = operationPaths;
283            _profilingInfos = profilingInfos;
284        }
285
286        return self;
287    }
288#endif
289    ETCoreMLLogErrorAndSetNSError(error,
290                                  ETCoreMLErrorModelProfilingNotSupported,
291                                  "%@: Model profiling is only available for macOS >= 14.4, iOS >= 17.4, tvOS >= 17.4 and watchOS >= 10.4.",
292                                  NSStringFromClass(self.class));
293    return nil;
294}
295
296- (nullable ETCoreMLModelProfilingResult *)profilingInfoForOperationsAtPaths:(NSArray<ETCoreMLModelStructurePath *> *)paths
297                                                                     options:(MLPredictionOptions *)options
298                                                                      inputs:(id<MLFeatureProvider>)inputs
299                                                                modelOutputs:(NSArray<MLMultiArray *> *_Nullable __autoreleasing *_Nonnull)modelOutputs
300                                                                       error:(NSError* __autoreleasing *)error {
301#if MODEL_PROFILING_IS_AVAILABLE
302    uint64_t modelExecutionStartTime = mach_absolute_time();
303    id<MLFeatureProvider> outputFeatures = [self.model predictionFromFeatures:inputs options:options error:error];
304    uint64_t modelExecutionEndTime = mach_absolute_time();
305    if (!modelOutputs) {
306        return nil;
307    }
308
309    if (@available(macOS 14.4, iOS 17.4, tvOS 17.4, watchOS 10.4, *)) {
310        ETCoreMLModelProfilingResult *profilingInfos = get_profiling_infos_for_paths([NSSet setWithArray:paths],
311                                                                                     self.topologicallySortedOperations,
312                                                                                     self.operationToPathMap,
313                                                                                     self.profilingInfos,
314                                                                                     modelExecutionStartTime,
315                                                                                     modelExecutionEndTime);
316
317
318
319        if (outputFeatures) {
320            set_model_outputs(outputFeatures, self.outputNames, modelOutputs);
321        }
322
323        return profilingInfos;
324    }
325#endif
326    return nil;
327}
328
329- (nullable ETCoreMLModelProfilingResult *)profilingInfoForOperationsAtPaths:(MLPredictionOptions *)options
330                                                                      inputs:(id<MLFeatureProvider>)inputs
331                                                                modelOutputs:(NSArray<MLMultiArray *> *_Nullable __autoreleasing *_Nonnull)modelOutputs
332                                                                       error:(NSError* __autoreleasing *)error {
333#if MODEL_PROFILING_IS_AVAILABLE
334    if (@available(macOS 14.4, iOS 17.4, tvOS 17.4, watchOS 10.4, *)) {
335        __block NSMutableArray<ETCoreMLModelStructurePath *> *paths = [NSMutableArray array];
336        visit_program_operation(self.computePlan.modelStructure, ^BOOL(MLModelStructureProgramOperation *operation, ETCoreMLModelStructurePath *path) {
337            if (!is_const_operation(operation)) {
338                [paths addObject:path];
339            }
340            return YES;
341        });
342
343        return [self profilingInfoForOperationsAtPaths:paths
344                                               options:options
345                                                inputs:inputs
346                                          modelOutputs:modelOutputs
347                                                 error:error];
348    }
349#endif
350    return nil;
351}
352
353@end
354