[CoreML] Add support for running statefule model. (#5143)

cymbalrush · web-flow · commit 13da62b00ac5 · 2024-09-07T12:53:51.000-07:00
diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLDefaultModelExecutor.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLDefaultModelExecutor.mm
@@ -29,9 +29,10 @@ - (instancetype)initWithModel:(ETCoreMLModel *)model {
     if (self.ignoreOutputBackings) {
         predictionOptions.outputBackings = @{};
     }
-    id<MLFeatureProvider> outputs = [self.model.mlModel predictionFromFeatures:inputs
-                                                                       options:predictionOptions
-                                                                         error:error];
+
+    id<MLFeatureProvider> outputs = [self.model predictionFromFeatures:inputs
+                                                               options:predictionOptions
+                                                                 error:error];
     if (!outputs) {
         return nil;
     }
diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLModel.h b/backends/apple/coreml/runtime/delegate/ETCoreMLModel.h
@@ -8,6 +8,10 @@
 #import <CoreML/CoreML.h>
 #import <vector>
 
+#if !defined(MODEL_STATE_IS_SUPPORTED) && __has_include(<CoreML/MLModel+MLState.h>)
+#define MODEL_STATE_IS_SUPPORTED 1
+#endif
+
 NS_ASSUME_NONNULL_BEGIN
 
 @class ETCoreMLAsset;
@@ -37,15 +41,12 @@ __attribute__((objc_subclassing_restricted))
                     orderedOutputNames:(NSOrderedSet<NSString*>*)orderedOutputNames
                                  error:(NSError* __autoreleasing*)error NS_DESIGNATED_INITIALIZER;
 
-- (nullable NSArray<MLMultiArray*>*)prepareInputs:(const std::vector<executorchcoreml::MultiArray>&)inputs
-                                            error:(NSError* __autoreleasing*)error;
-
-- (nullable NSArray<MLMultiArray*>*)prepareOutputBackings:(const std::vector<executorchcoreml::MultiArray>&)outputs
-                                                    error:(NSError* __autoreleasing*)error;
-
 /// The underlying MLModel.
 @property (strong, readonly, nonatomic) MLModel* mlModel;
 
+/// The model state.
+@property (strong, readonly, nonatomic) id state API_AVAILABLE(macos(15.0), ios(18.0), tvos(18.0), watchos(11.0));
+
 /// The asset from which the model is loaded.
 @property (strong, readonly, nonatomic) ETCoreMLAsset* asset;
 
@@ -58,6 +59,19 @@ __attribute__((objc_subclassing_restricted))
 /// The ordered output names of the model.
 @property (copy, readonly, nonatomic) NSOrderedSet<NSString*>* orderedOutputNames;
 
+
+- (nullable id<MLFeatureProvider>)predictionFromFeatures:(id<MLFeatureProvider>)input
+                                                 options:(MLPredictionOptions*)options
+                                                   error:(NSError* __autoreleasing*)error;
+
+- (nullable NSArray<MLMultiArray*>*)prepareInputs:(const std::vector<executorchcoreml::MultiArray>&)inputs
+                                            error:(NSError* __autoreleasing*)error;
+
+- (nullable NSArray<MLMultiArray*>*)prepareOutputBackings:(const std::vector<executorchcoreml::MultiArray>&)outputs
+                                                    error:(NSError* __autoreleasing*)error;
+
+- (BOOL)prewarmAndReturnError:(NSError* __autoreleasing*)error;
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLModel.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLModel.mm
@@ -7,10 +7,12 @@
 
 #import <ETCoreMLModel.h>
 
-#import <ETCoreMLAsset.h>
+#import "ETCoreMLAsset.h"
+#import "ETCoreMLLogging.h"
+#import "multiarray.h"
+#import "objc_array_util.h"
+#import "MLModel_Prewarm.h"
 #import <functional>
-#import <objc_array_util.h>
-#import <multiarray.h>
 #import <numeric>
 
 #pragma mark - ETCoreMLMultiArrayDescriptor
@@ -194,6 +196,11 @@ - (nullable instancetype)initWithAsset:(ETCoreMLAsset *)asset
         _cache = [[NSCache alloc] init];
         _inputConstraintsByName = get_multi_array_input_constraints_by_name(mlModel.modelDescription);
         _outputConstraintsByName = get_multi_array_output_constraints_by_name(mlModel.modelDescription);
+#if MODEL_STATE_IS_SUPPORTED
+        if (@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, *)) {
+            _state = mlModel.modelDescription.stateDescriptionsByName.count > 0 ? [_mlModel newState] : nil;
+        }
+#endif
     }
     
     return self;
@@ -272,4 +279,48 @@ MultiArray buffer(mutableBytes, MultiArray::MemoryLayout(to_multiarray_data_type
     
 }
 
+- (nullable id<MLFeatureProvider>)predictionFromFeatures:(id<MLFeatureProvider>)input
+                                                 options:(MLPredictionOptions *)options
+                                                   error:(NSError **)error {
+
+#if MODEL_STATE_IS_SUPPORTED
+    if (@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, *)) {
+        if (self.state != nil) {
+            return [self.mlModel predictionFromFeatures:input
+                                             usingState:(MLState *)self.state
+                                                options:options
+                                                  error:error];
+        }
+    }
+#endif
+
+    return [self.mlModel predictionFromFeatures:input
+                                        options:options
+                                          error:error];
+}
+
+- (BOOL)prewarmAndReturnError:(NSError* __autoreleasing*)error {
+    BOOL prewarm = YES;
+#if MODEL_STATE_IS_SUPPORTED
+    if (@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, *)) {
+        prewarm = (self.mlModel.modelDescription.stateDescriptionsByName.count == 0);
+    }
+#endif
+
+    NSError *localError = nil;
+    BOOL result = prewarm ? [self.mlModel prewarmAndReturnError:&localError] : NO;
+    if (!result) {
+        ETCoreMLLogError(localError,
+                         "%@: Failed to prewarm model with identifier = %@",
+                         NSStringFromClass(self.class),
+                         self.identifier);
+    }
+
+    if (error) {
+        *error = localError;
+    }
+
+    return result;
+}
+
 @end
diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm
@@ -598,21 +598,8 @@ - (BOOL)prewarmModelWithHandle:(ModelHandle *)handle
     if (!model) {
         return NO;
     }
-    
-    NSError *localError = nil;
-    BOOL result = [model.mlModel prewarmAndReturnError:&localError];
-    if (!result) {
-        ETCoreMLLogError(localError,
-                         "%@: Failed to prewarm model with identifier = %@",
-                         NSStringFromClass(self.assetManager.class),
-                         model.identifier);
-    }
-    
-    if (error) {
-        *error = localError;
-    }
-    
-    return result;
+
+    return [model prewarmAndReturnError:error];
 }
 
 - (void)prewarmRecentlyUsedAssetsWithMaxCount:(NSUInteger)maxCount {
diff --git a/backends/apple/coreml/runtime/delegate/MLModel_Prewarm.mm b/backends/apple/coreml/runtime/delegate/MLModel_Prewarm.mm
@@ -77,7 +77,7 @@ - (BOOL)prewarmAndReturnError:(NSError * __autoreleasing *)error {
         if (!inputs) {
             return NO;
         }
-        
+
         id<MLFeatureProvider> outputs = [self predictionFromFeatures:inputs error:error];
         return outputs != nil;
     }
diff --git a/backends/apple/coreml/runtime/sdk/ETCoreMLModelAnalyzer.mm b/backends/apple/coreml/runtime/sdk/ETCoreMLModelAnalyzer.mm
@@ -88,10 +88,9 @@ - (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset *)compiledMod
                                                  eventLogger:(const executorchcoreml::ModelEventLogger *)eventLogger
                                                        error:(NSError * __autoreleasing *)error {
     if (self.profiler == nil) {
-        ETCoreMLModelProfiler *profiler = [[ETCoreMLModelProfiler alloc] initWithCompiledModelAsset:self.model.asset
-                                                                                        outputNames:self.model.orderedOutputNames
-                                                                                      configuration:self.configuration
-                                                                                              error:error];
+        ETCoreMLModelProfiler *profiler = [[ETCoreMLModelProfiler alloc] initWithModel:self.model
+                                                                         configuration:self.configuration
+                                                                                 error:error];
         self.profiler = profiler;
     }
        
diff --git a/backends/apple/coreml/runtime/sdk/ETCoreMLModelProfiler.h b/backends/apple/coreml/runtime/sdk/ETCoreMLModelProfiler.h
@@ -31,14 +31,12 @@ __attribute__((objc_subclassing_restricted))
 
 /// Constructs an `ETCoreMLModelProfiler` instance.
 ///
-/// @param compiledModelAsset The compiled model asset (mlmodelc).
-/// @param outputNames The model output names.
+/// @param model The model.
 /// @param configuration The model configuration.
 /// @param error   On failure, error is filled with the failure information.
-- (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset*)compiledModelAsset
-                                        outputNames:(NSOrderedSet<NSString*>*)outputNames
-                                      configuration:(MLModelConfiguration*)configuration
-                                              error:(NSError* __autoreleasing*)error NS_DESIGNATED_INITIALIZER;
+- (nullable instancetype)initWithModel:(ETCoreMLModel*)model
+                         configuration:(MLModelConfiguration*)configuration
+                                 error:(NSError* __autoreleasing*)error NS_DESIGNATED_INITIALIZER;
 
 /// Returns profiling info of operations at the specified paths.
 ///
diff --git a/backends/apple/coreml/runtime/sdk/ETCoreMLModelProfiler.mm b/backends/apple/coreml/runtime/sdk/ETCoreMLModelProfiler.mm
@@ -8,6 +8,7 @@
 #import "ETCoreMLModelProfiler.h"
 
 #import "ETCoreMLAsset.h"
+#import "ETCoreMLModel.h"
 #import "ETCoreMLLogging.h"
 #import "ETCoreMLModelStructurePath.h"
 #import "ETCoreMLOperationProfilingInfo.h"
@@ -221,8 +222,8 @@ void set_model_outputs(id<MLFeatureProvider> output_features,
 }
 
 @interface ETCoreMLModelProfiler ()
-/// The CoreML model.
-@property (readonly, strong, nonatomic) MLModel *model;
+/// The model.
+@property (readonly, strong, nonatomic) ETCoreMLModel *model;
 /// The model output names.
 @property (readonly, copy, nonatomic) NSOrderedSet<NSString *> *outputNames;
 #if MODEL_PROFILING_IS_AVAILABLE
@@ -240,25 +241,19 @@ @interface ETCoreMLModelProfiler ()
 
 @implementation ETCoreMLModelProfiler
 
-- (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset *)compiledModelAsset
-                                        outputNames:(NSOrderedSet<NSString *> *)outputNames
-                                      configuration:(MLModelConfiguration *)configuration
-                                              error:(NSError * __autoreleasing *)error  {
+- (nullable instancetype)initWithModel:(ETCoreMLModel *)model
+                         configuration:(MLModelConfiguration *)configuration
+                                 error:(NSError * __autoreleasing *)error  {
 #if MODEL_PROFILING_IS_AVAILABLE
     if (@available(macOS 14.4, iOS 17.4, tvOS 17.4, watchOS 10.4, *)) {
-        NSURL *compiledModelURL = compiledModelAsset.contentURL;
+        NSURL *compiledModelURL = model.asset.contentURL;
         MLComputePlan *computePlan = get_compute_plan_of_model_at_url(compiledModelURL,
                                                                       configuration,
                                                                       error);
         if (!computePlan) {
             return nil;
         }
-        
-        MLModel *model = [MLModel modelWithContentsOfURL:compiledModelURL error:error];
-        if (!model) {
-            return nil;
-        }
-        
+
         __block NSMutableArray<ETCoreMLModelStructurePath *> *operationPaths = [NSMutableArray array];
         __block NSMutableDictionary<NSValue *, ETCoreMLModelStructurePath *> *operationToPathMap = [NSMutableDictionary dictionary];
         __block NSMutableArray<MLModelStructureProgramOperation *> *topologicallySortedOperations = [NSMutableArray new];
@@ -280,7 +275,6 @@ - (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset *)compiledMod
         
         self = [super init];
         if (self) {
-            _outputNames = [outputNames copy];
             _model = model;
             _computePlan = computePlan;
             _operationToPathMap = operationToPathMap;

Original file line number	Diff line number	Diff line change
`@@ -29,9 +29,10 @@ - (instancetype)initWithModel:(ETCoreMLModel *)model {`
`29`	`29`	`if (self.ignoreOutputBackings) {`
`30`	`30`	`predictionOptions.outputBackings = @{};`
`31`	`31`	`}`
`32`		`- id<MLFeatureProvider> outputs = [self.model.mlModel predictionFromFeatures:inputs`
`33`		`- options:predictionOptions`
`34`		`- error:error];`
	`32`	`+`
	`33`	`+ id<MLFeatureProvider> outputs = [self.model predictionFromFeatures:inputs`
	`34`	`+ options:predictionOptions`
	`35`	`+ error:error];`
`35`	`36`	`if (!outputs) {`
`36`	`37`	`return nil;`
`37`	`38`	`}`
Original file line number	Diff line number	Diff line change
`@@ -77,7 +77,7 @@ - (BOOL)prewarmAndReturnError:(NSError * __autoreleasing *)error {`
`77`	`77`	`if (!inputs) {`
`78`	`78`	`return NO;`
`79`	`79`	`}`
`80`		`-`
	`80`	`+`
`81`	`81`	`id<MLFeatureProvider> outputs = [self predictionFromFeatures:inputs error:error];`
`82`	`82`	`return outputs != nil;`
`83`	`83`	`}`