88#import < ETCoreMLModel.h>
99
1010#import < ETCoreMLAsset.h>
11+ #import < functional>
12+ #import < objc_array_util.h>
13+ #import < multiarray.h>
14+ #import < numeric>
15+
16+ #pragma mark - ETCoreMLMultiArrayDescriptor
17+ __attribute__ ((objc_subclassing_restricted))
18+ @interface ETCoreMLMultiArrayDescriptor: NSObject <NSCopying>
19+
20+ - (instancetype )init NS_UNAVAILABLE;
21+
22+ + (instancetype )new NS_UNAVAILABLE;
23+
24+ - (instancetype )initWithShape:(NSArray <NSNumber *> *)shape
25+ dataType:(MLMultiArrayDataType)dataType NS_DESIGNATED_INITIALIZER;
26+
27+ @property (copy, readonly, nonatomic) NSArray <NSNumber *> *shape;
28+
29+ @property (assign, readonly, nonatomic) MLMultiArrayDataType dataType;
30+
31+ @end
32+
33+ @implementation ETCoreMLMultiArrayDescriptor
34+
35+ - (instancetype )initWithShape : (NSArray <NSNumber *> *)shape
36+ dataType : (MLMultiArrayDataType)dataType {
37+ self = [super init ];
38+ if (self) {
39+ _shape = shape;
40+ _dataType = dataType;
41+ }
42+
43+ return self;
44+ }
45+
46+ - (BOOL )isEqual : (id )object {
47+ if (object == self) {
48+ return YES ;
49+ }
50+
51+ if (![object isKindOfClass: self .class ]) {
52+ return NO ;
53+ }
54+
55+ ETCoreMLMultiArrayDescriptor *other = (ETCoreMLMultiArrayDescriptor *)object;
56+ return [self .shape isEqualToArray: other.shape] && self.dataType == other.dataType ;
57+ }
58+
59+ - (NSUInteger )hash {
60+ return [self .shape hash ] ^ (NSUInteger )self.dataType ;
61+ }
62+
63+ - (instancetype )copyWithZone : (NSZone *)zone {
64+ return [[ETCoreMLMultiArrayDescriptor allocWithZone: zone] initWithShape: self .shape
65+ dataType: self .dataType];
66+ }
67+
68+ @end
69+
70+ namespace {
71+
72+ using namespace executorchcoreml ;
73+
74+ size_t get_number_of_bytes (MLMultiArrayDataType data_type) {
75+ switch (data_type) {
76+ case MLMultiArrayDataTypeFloat16: {
77+ return 2 ;
78+ }
79+ case MLMultiArrayDataTypeFloat32: {
80+ return 4 ;
81+ }
82+ case MLMultiArrayDataTypeInt32: {
83+ return 4 ;
84+ }
85+ case MLMultiArrayDataTypeFloat64: {
86+ return 8 ;
87+ }
88+ default : {
89+ return 0 ;
90+ }
91+ }
92+ }
93+
94+ std::vector<size_t > calculate_strides (const std::vector<size_t >& shape) {
95+ if (shape.size () == 0 ) {
96+ return {};
97+ }
98+
99+ if (shape.size () == 1 ) {
100+ return {1 };
101+ }
102+
103+ std::vector<size_t > strides (shape.size (), 1 );
104+ size_t product = 1 ;
105+ for (size_t i = shape.size (); i > 0 ; i--) {
106+ strides[i - 1 ] = product;
107+ product *= shape[i - 1 ];
108+ }
109+
110+ return strides;
111+ }
112+
113+ MLMultiArray * _Nullable make_ml_multi_array (const std::vector<size_t >& shape,
114+ MLMultiArrayDataType dataType,
115+ NSCache <ETCoreMLMultiArrayDescriptor *, NSMutableData *> *cache,
116+ NSError * __autoreleasing *error) {
117+ ETCoreMLMultiArrayDescriptor *descriptor = [[ETCoreMLMultiArrayDescriptor alloc ] initWithShape: to_array (shape)
118+ dataType: dataType];
119+ // Check the cache first otherwise allocate a new backing storage.
120+ NSMutableData *backing_storage = [cache objectForKey: descriptor];
121+ if (backing_storage) {
122+ [cache removeObjectForKey: descriptor];
123+ } else {
124+ size_t n = std::accumulate (shape.cbegin (), shape.cend (), 1 , std::multiplies<size_t >{});
125+ backing_storage = [[NSMutableData alloc ] initWithLength: n * get_number_of_bytes (dataType)];
126+ }
127+
128+ __weak NSCache <ETCoreMLMultiArrayDescriptor *, NSMutableData *> *weakCache = cache;
129+ // Add the storage back to the cache when it gets deallocated, the next prediction would use the same storage.
130+ MLMultiArray *result = [[MLMultiArray alloc ] initWithDataPointer: backing_storage.mutableBytes
131+ shape: descriptor.shape
132+ dataType: descriptor.dataType
133+ strides: to_array (calculate_strides (shape))
134+ deallocator: ^(void * _Nonnull bytes) {[weakCache setObject: backing_storage forKey: descriptor];}
135+ error: error];
136+
137+ return result;
138+ }
139+
140+ NSDictionary <NSString *, MLMultiArrayConstraint *> *
141+ get_multi_array_constraints_by_name (NSDictionary <NSString *, MLFeatureDescription *> *feature_descriptions) {
142+ NSMutableDictionary <NSString *, MLMultiArrayConstraint *> *result = [NSMutableDictionary dictionaryWithCapacity: feature_descriptions.count];
143+ [feature_descriptions enumerateKeysAndObjectsUsingBlock: ^(NSString *key, MLFeatureDescription *description, BOOL * _Nonnull stop) {
144+ result[key] = description.multiArrayConstraint ;
145+ }];
146+
147+ return result;
148+ }
149+
150+ NSDictionary <NSString *, MLMultiArrayConstraint *> *get_multi_array_input_constraints_by_name (MLModelDescription *description) {
151+ return get_multi_array_constraints_by_name (description.inputDescriptionsByName );
152+ }
153+
154+ NSDictionary <NSString *, MLMultiArrayConstraint *> *get_multi_array_output_constraints_by_name (MLModelDescription *description) {
155+ return get_multi_array_constraints_by_name (description.outputDescriptionsByName );
156+ }
157+
158+ }
159+
160+ #pragma mark - ETCoreMLModel
161+ @interface ETCoreMLModel ()
162+
163+ @property (strong , readonly , nonatomic ) NSCache <ETCoreMLMultiArrayDescriptor *, NSMutableData *> *cache;
164+ @property (copy , readonly , nonatomic ) NSDictionary <NSString *, MLMultiArrayConstraint *> *inputConstraintsByName;
165+ @property (copy , readonly , nonatomic ) NSDictionary <NSString *, MLMultiArrayConstraint *> *outputConstraintsByName;
166+
167+ @end
168+
11169
12170@implementation ETCoreMLModel
13171
@@ -33,13 +191,84 @@ - (nullable instancetype)initWithAsset:(ETCoreMLAsset *)asset
33191 _asset = asset;
34192 _orderedInputNames = [orderedInputNames copy ];
35193 _orderedOutputNames = [orderedOutputNames copy ];
194+ _cache = [[NSCache alloc ] init ];
195+ _inputConstraintsByName = get_multi_array_input_constraints_by_name (mlModel.modelDescription );
196+ _outputConstraintsByName = get_multi_array_output_constraints_by_name (mlModel.modelDescription );
36197 }
37-
198+
38199 return self;
39200}
40201
41202- (NSString *)identifier {
42203 return self.asset .identifier ;
43204}
44205
206+ - (nullable NSArray <MLMultiArray *> *)prepareArgs : (const std::vector<executorchcoreml::MultiArray>&)args
207+ argNames : (NSOrderedSet <NSString *> *)argNames
208+ argConstraintsByName : (NSDictionary <NSString *, MLMultiArrayConstraint *> *)argConstraintsByName
209+ copyData : (BOOL )copyData
210+ error : (NSError * __autoreleasing *)error {
211+ NSEnumerator *nameEnumerator = [argNames objectEnumerator ];
212+ NSMutableArray <MLMultiArray *> *result = [NSMutableArray arrayWithCapacity: args.size ()];
213+ for (const auto & arg : args) {
214+ NSString *argName = [nameEnumerator nextObject ];
215+ MLMultiArrayConstraint *constraint = argConstraintsByName[argName];
216+ const auto & layout = arg.layout ();
217+ auto dataType = to_ml_multiarray_data_type (layout.dataType ());
218+ MLMultiArray *multiArrayArg = nil ;
219+ if (dataType == constraint.dataType ) {
220+ // We can use the same data storage.
221+ multiArrayArg = [[MLMultiArray alloc ] initWithDataPointer: arg.data ()
222+ shape: to_array (layout.shape ())
223+ dataType: constraint.dataType
224+ strides: to_array (layout.strides ())
225+ deallocator: ^(void * _Nonnull bytes) {}
226+ error: error];
227+ copyData = NO ;
228+ } else {
229+ // We can' use the same data storage, data types are not the same.
230+ multiArrayArg = ::make_ml_multi_array (layout.shape (), constraint.dataType , self.cache , error);
231+ }
232+
233+ if (!multiArrayArg) {
234+ return nil ;
235+ }
236+
237+ if (multiArrayArg && copyData) {
238+ [multiArrayArg getMutableBytesWithHandler: ^(void *_Nonnull mutableBytes,
239+ NSInteger __unused size,
240+ NSArray <NSNumber *> *strides) {
241+ MultiArray buffer (mutableBytes, MultiArray::MemoryLayout (to_multiarray_data_type (constraint.dataType ).value (),
242+ layout.shape (),
243+ to_vector<ssize_t >(strides)));
244+ arg.copy (buffer);
245+ }];
246+ }
247+
248+ [result addObject: multiArrayArg];
249+ }
250+
251+ return result;
252+ }
253+
254+ - (nullable NSArray <MLMultiArray *> *)prepareInputs : (const std::vector<executorchcoreml::MultiArray>&)inputs
255+ error : (NSError * __autoreleasing *)error {
256+ return [self prepareArgs: inputs
257+ argNames: self .orderedInputNames
258+ argConstraintsByName: self .inputConstraintsByName
259+ copyData: YES
260+ error: error];
261+
262+ }
263+
264+ - (nullable NSArray <MLMultiArray *> *)prepareOutputBackings : (const std::vector<executorchcoreml::MultiArray>&)outputs
265+ error : (NSError * __autoreleasing *)error {
266+ return [self prepareArgs: outputs
267+ argNames: self .orderedOutputNames
268+ argConstraintsByName: self .outputConstraintsByName
269+ copyData: NO
270+ error: error];
271+
272+ }
273+
45274@end
0 commit comments