diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLDefaultModelExecutor.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLDefaultModelExecutor.mm index 57316e2801..226307f3c8 100644 --- 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 outputs = [self.model.mlModel predictionFromFeatures:inputs - options:predictionOptions - error:error]; + + id 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 index 9bf3183e65..0bbd1132e9 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLModel.h +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLModel.h @@ -8,6 +8,10 @@ #import #import +#if !defined(MODEL_STATE_IS_SUPPORTED) && __has_include() +#define MODEL_STATE_IS_SUPPORTED 1 +#endif + NS_ASSUME_NONNULL_BEGIN @class ETCoreMLAsset; @@ -37,15 +41,12 @@ __attribute__((objc_subclassing_restricted)) orderedOutputNames:(NSOrderedSet*)orderedOutputNames error:(NSError* __autoreleasing*)error NS_DESIGNATED_INITIALIZER; -- (nullable NSArray*)prepareInputs:(const std::vector&)inputs - error:(NSError* __autoreleasing*)error; - -- (nullable NSArray*)prepareOutputBackings:(const std::vector&)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* orderedOutputNames; + +- (nullable id)predictionFromFeatures:(id)input + options:(MLPredictionOptions*)options + error:(NSError* __autoreleasing*)error; + +- (nullable NSArray*)prepareInputs:(const std::vector&)inputs + error:(NSError* __autoreleasing*)error; + +- (nullable NSArray*)prepareOutputBackings:(const std::vector&)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 index ee7218bd27..250d5cd951 100644 --- a/backends/apple/coreml/runtime/delegate/ETCoreMLModel.mm +++ b/backends/apple/coreml/runtime/delegate/ETCoreMLModel.mm @@ -7,10 +7,12 @@ #import -#import +#import "ETCoreMLAsset.h" +#import "ETCoreMLLogging.h" +#import "multiarray.h" +#import "objc_array_util.h" +#import "MLModel_Prewarm.h" #import -#import -#import #import #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)predictionFromFeatures:(id)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 index 8d6d537385..2a5c3ed696 100644 --- 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 index 71ce967ac3..97d0400796 100644 --- 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 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 index 1740faf00e..988b5d808a 100644 --- 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 index 07a384a516..7a43a30d75 100644 --- 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*)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 index c9ad324a6c..5998701eb0 100644 --- 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 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 *outputNames; #if MODEL_PROFILING_IS_AVAILABLE @@ -240,25 +241,19 @@ @interface ETCoreMLModelProfiler () @implementation ETCoreMLModelProfiler -- (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset *)compiledModelAsset - outputNames:(NSOrderedSet *)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 *operationPaths = [NSMutableArray array]; __block NSMutableDictionary *operationToPathMap = [NSMutableDictionary dictionary]; __block NSMutableArray *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;