xref: /aosp_15_r20/external/executorch/backends/apple/coreml/runtime/sdk/ETCoreMLOperationProfilingInfo.mm (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1//
2// ETCoreMLOperationProfilingInfo.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 "ETCoreMLOperationProfilingInfo.h"
9
10#import "hash_util.h"
11#import "model_event_logger_impl.h"
12
13namespace  {
14NSString *const kPreferredComputeUnitKey = @"preferredComputeUnit";
15NSString *const kSupportedComputeUnitsKey = @"supportedComputeUnits";
16NSString *const kOutputNamesKey = @"outputNames";
17NSString *const kOperatorNameKey = @"operatorName";
18NSString *const kEstimatedCostKey = @"estimatedCost";
19
20NSArray<NSString *> *compute_unit_names(ETCoreMLComputeUnits compute_units) {
21    NSMutableArray<NSString *> *result = [NSMutableArray array];
22    if ((compute_units & ETCoreMLComputeUnitNeuralEngine) > 0) {
23        [result addObject:@"NE"];
24    }
25    if ((compute_units & ETCoreMLComputeUnitCPU) > 0) {
26        [result addObject:@"CPU"];
27    }
28    if ((compute_units & ETCoreMLComputeUnitGPU) > 0) {
29        [result addObject:@"GPU"];
30    }
31
32    return result;
33}
34
35NSData *get_metadata(ETCoreMLComputeUnits preferredComputeUnit,
36                     ETCoreMLComputeUnits supportedComputeUnits,
37                     NSArray<NSString *> *outputNames,
38                     NSString *operatorName,
39                     double estimatedCost) {
40    NSMutableDictionary<NSString *, id> *result = [NSMutableDictionary new];
41    result[kPreferredComputeUnitKey] = compute_unit_names(preferredComputeUnit).firstObject;
42    result[kSupportedComputeUnitsKey] = compute_unit_names(supportedComputeUnits);
43    result[kOutputNamesKey] = outputNames;
44    result[kOperatorNameKey] = operatorName;
45    result[kEstimatedCostKey] = @(estimatedCost);
46    NSError *local_error = nil;
47    NSData *data = [NSJSONSerialization dataWithJSONObject:result options:NSJSONWritingOptions(0) error:&local_error];
48    NSCAssert(data != nil, @"%@: Failed to serialize metadata.", NSStringFromClass(ETCoreMLOperationProfilingInfo.class));
49
50    return data;
51};
52}
53
54@implementation ETCoreMLOperationProfilingInfo
55
56- (instancetype)initWithPreferredComputeUnit:(ETCoreMLComputeUnits)preferredComputeUnit
57                       supportedComputeUnits:(ETCoreMLComputeUnits)supportedComputeUnits
58                 estimatedExecutionStartTime:(uint64_t)estimatedExecutionStartTime
59                   estimatedExecutionEndTime:(uint64_t)estimatedExecutionEndTime
60                               estimatedCost:(double)estimatedCost
61                                 outputNames:(NSArray<NSString *> *)outputNames
62                                operatorName:(NSString *)operatorName
63                                    metadata:(NSData *)metadata {
64    self = [super init];
65    if (self) {
66        _preferredComputeUnit = preferredComputeUnit;
67        _supportedComputeUnits = supportedComputeUnits;
68        _estimatedExecutionStartTime = estimatedExecutionStartTime;
69        _estimatedExecutionEndTime = estimatedExecutionEndTime;
70        _estimatedCost = estimatedCost;
71        _outputNames = [outputNames copy];
72        _operatorName = [operatorName copy];
73        _metadata = get_metadata(preferredComputeUnit, supportedComputeUnits, outputNames, operatorName, estimatedCost);
74    }
75
76    return self;
77}
78
79- (instancetype)initWithPreferredComputeUnit:(ETCoreMLComputeUnits)preferredComputeUnit
80                       supportedComputeUnits:(ETCoreMLComputeUnits)supportedComputeUnits
81                 estimatedExecutionStartTime:(uint64_t)estimatedExecutionStartTime
82                   estimatedExecutionEndTime:(uint64_t)estimatedExecutionEndTime
83                               estimatedCost:(double)estimatedCost
84                                 outputNames:(NSArray<NSString *> *)outputNames
85                                operatorName:(NSString *)operatorName {
86    NSData *metadata = get_metadata(preferredComputeUnit, supportedComputeUnits, outputNames, operatorName, estimatedCost);
87    return [self initWithPreferredComputeUnit:preferredComputeUnit
88                        supportedComputeUnits:supportedComputeUnits
89                  estimatedExecutionStartTime:estimatedExecutionStartTime
90                    estimatedExecutionEndTime:estimatedExecutionEndTime
91                                estimatedCost:estimatedCost
92                                  outputNames:outputNames
93                                 operatorName:operatorName
94                                     metadata:metadata];
95}
96
97- (BOOL)isEqual:(id)object {
98    if (object == self) {
99        return YES;
100    }
101
102    if (![object isKindOfClass:self.class]) {
103        return NO;
104    }
105
106    ETCoreMLOperationProfilingInfo *other = (ETCoreMLOperationProfilingInfo *)object;
107
108    return self.preferredComputeUnit == other.preferredComputeUnit &&
109    self.supportedComputeUnits == other.supportedComputeUnits &&
110    self.estimatedExecutionStartTime == other.estimatedExecutionStartTime &&
111    self.estimatedExecutionEndTime == other.estimatedExecutionEndTime &&
112    [self.outputNames isEqualToArray:other.outputNames] &&
113    [self.operatorName isEqualToString:other.operatorName] &&
114    [self.metadata isEqualToData:other.metadata];
115}
116
117- (NSUInteger)hash {
118    size_t seed = 0;
119    executorchcoreml::hash_combine(seed, self.preferredComputeUnit);
120    executorchcoreml::hash_combine(seed, self.supportedComputeUnits);
121    executorchcoreml::hash_combine(seed, self.estimatedExecutionStartTime);
122    executorchcoreml::hash_combine(seed, self.estimatedExecutionEndTime);
123    executorchcoreml::hash_combine(seed, self.outputNames.hash);
124    executorchcoreml::hash_combine(seed, self.operatorName.hash);
125    executorchcoreml::hash_combine(seed, self.metadata.hash);
126
127    return seed;
128}
129
130- (instancetype)copyWithZone:(NSZone *)zone {
131    return [[ETCoreMLOperationProfilingInfo allocWithZone:zone] initWithPreferredComputeUnit:self.preferredComputeUnit
132                                                                       supportedComputeUnits:self.supportedComputeUnits
133                                                                 estimatedExecutionStartTime:self.estimatedExecutionStartTime
134                                                                   estimatedExecutionEndTime:self.estimatedExecutionEndTime
135                                                                               estimatedCost:self.estimatedCost
136                                                                                 outputNames:self.outputNames
137                                                                                operatorName:self.operatorName
138                                                                                    metadata:self.metadata];
139}
140
141@end
142