Skip to content

Commit

Permalink
[CoreML] Add support for running statefule model. (#5143)
Browse files Browse the repository at this point in the history
  • Loading branch information
cymbalrush authored Sep 7, 2024
1 parent 3268da2 commit 13da62b
Show file tree
Hide file tree
Showing 8 changed files with 96 additions and 52 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
26 changes: 20 additions & 6 deletions backends/apple/coreml/runtime/delegate/ETCoreMLModel.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;

Expand All @@ -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
57 changes: 54 additions & 3 deletions backends/apple/coreml/runtime/delegate/ETCoreMLModel.mm
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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
17 changes: 2 additions & 15 deletions backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
2 changes: 1 addition & 1 deletion backends/apple/coreml/runtime/delegate/MLModel_Prewarm.mm
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ - (BOOL)prewarmAndReturnError:(NSError * __autoreleasing *)error {
if (!inputs) {
return NO;
}

id<MLFeatureProvider> outputs = [self predictionFromFeatures:inputs error:error];
return outputs != nil;
}
Expand Down
7 changes: 3 additions & 4 deletions backends/apple/coreml/runtime/sdk/ETCoreMLModelAnalyzer.mm
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down
10 changes: 4 additions & 6 deletions backends/apple/coreml/runtime/sdk/ETCoreMLModelProfiler.h
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down
22 changes: 8 additions & 14 deletions backends/apple/coreml/runtime/sdk/ETCoreMLModelProfiler.mm
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#import "ETCoreMLModelProfiler.h"

#import "ETCoreMLAsset.h"
#import "ETCoreMLModel.h"
#import "ETCoreMLLogging.h"
#import "ETCoreMLModelStructurePath.h"
#import "ETCoreMLOperationProfilingInfo.h"
Expand Down Expand Up @@ -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
Expand All @@ -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];
Expand All @@ -280,7 +275,6 @@ - (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset *)compiledMod

self = [super init];
if (self) {
_outputNames = [outputNames copy];
_model = model;
_computePlan = computePlan;
_operationToPathMap = operationToPathMap;
Expand Down

0 comments on commit 13da62b

Please sign in to comment.