1 // 2 // ETCoreMLModelManager.h 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 <CoreML/CoreML.h> 9 10 #import <vector> 11 12 NS_ASSUME_NONNULL_BEGIN 13 14 namespace executorchcoreml { 15 struct ModelLoggingOptions; 16 class ModelEventLogger; 17 class MultiArray; 18 }; 19 20 @class ETCoreMLModel; 21 @class ETCoreMLAssetManager; 22 23 typedef void ModelHandle; 24 25 /// A class responsible for managing the models loaded by the delegate. 26 __attribute__((objc_subclassing_restricted)) 27 @interface ETCoreMLModelManager : NSObject 28 29 + (instancetype)new NS_UNAVAILABLE; 30 31 - (instancetype)init NS_UNAVAILABLE; 32 33 /// Constructs an `ETCoreMLModelManager` instance. 34 /// 35 /// @param assetManager The asset manager used to manage storage of compiled models. 36 - (instancetype)initWithAssetManager:(ETCoreMLAssetManager*)assetManager NS_DESIGNATED_INITIALIZER; 37 38 /// Loads the model from the AOT data. 39 /// 40 /// The data is the AOT blob stored in the executorch Program. The method first parses the model 41 /// metadata stored in the blob and extracts the identifier. If the asset store contains an asset 42 /// with the same identifier then the model is directly loaded from the asset otherwise the model is 43 /// saved to the filesystem, compiled, moved to the assets store, and then loaded from there. 44 /// 45 /// @param data The AOT blob data. 46 /// @param configuration The model configuration that will be used to load the model. 47 /// @param error On failure, error is filled with the failure information. 48 /// @retval An opaque handle that points to the loaded model. 49 - (ModelHandle*)loadModelFromAOTData:(NSData*)data 50 configuration:(MLModelConfiguration*)configuration 51 error:(NSError* __autoreleasing*)error; 52 53 /// Executes the loaded model. 54 /// 55 /// @param handle The handle to the loaded model. 56 /// @param args The arguments (inputs and outputs) of the model. 57 /// @param loggingOptions The model logging options. 58 /// @param error On failure, error is filled with the failure information. 59 /// @retval `YES` if the execution succeeded otherwise `NO`. 60 - (BOOL)executeModelWithHandle:(ModelHandle*)handle 61 args:(NSArray<MLMultiArray*>*)args 62 loggingOptions:(const executorchcoreml::ModelLoggingOptions&)loggingOptions 63 eventLogger:(const executorchcoreml::ModelEventLogger* _Nullable)eventLogger 64 error:(NSError* __autoreleasing*)error; 65 66 /// Executes the loaded model. 67 /// 68 /// @param handle The handle to the loaded model. 69 /// @param argsVec The arguments (inputs and outputs) of the model. 70 /// @param loggingOptions The model logging options. 71 /// @param error On failure, error is filled with the failure information. 72 /// @retval `YES` if the execution succeeded otherwise `NO`. 73 - (BOOL)executeModelWithHandle:(ModelHandle*)handle 74 argsVec:(const std::vector<executorchcoreml::MultiArray>&)argsVec 75 loggingOptions:(const executorchcoreml::ModelLoggingOptions&)loggingOptions 76 eventLogger:(const executorchcoreml::ModelEventLogger* _Nullable)eventLogger 77 error:(NSError* __autoreleasing*)error; 78 79 /// Unloads the loaded model. 80 /// 81 /// @param handle The handle to the loaded model. 82 /// @retval `YES` if the model is unloaded otherwise `NO`. 83 - (BOOL)unloadModelWithHandle:(ModelHandle*)handle; 84 85 /// Returns the loaded model associated with the handle. 86 /// 87 /// @param handle The handle to the loaded model. 88 /// @retval The loaded model. 89 - (nullable ETCoreMLModel*)modelWithHandle:(ModelHandle*)handle; 90 91 /// Pre-warms most recently used assets. This does an advisory read ahead of the asset (compiled 92 /// model) files and could potentially improve the read time if the data is already available in the 93 /// system cache. 94 /// 95 /// @param maxCount The maximum count of assets to be pre-warmed. 96 - (void)prewarmRecentlyUsedAssetsWithMaxCount:(NSUInteger)maxCount; 97 98 /// Pre-warms the model associated with the handle. This could potentially improve the model 99 /// execution time. 100 /// 101 /// @param handle The handle to the loaded model. 102 /// @param error On failure, error is filled with the failure information. 103 /// @retval `YES` if the model was pre-warmed otherwise `NO`. 104 - (BOOL)prewarmModelWithHandle:(ModelHandle*)handle error:(NSError* __autoreleasing*)error; 105 106 /// Purges model cache. 107 /// 108 /// @param error On failure, error is filled with the failure information. 109 /// @retval `YES` if the cache is purged otherwise `NO`. 110 - (BOOL)purgeModelsCacheAndReturnError:(NSError* __autoreleasing*)error; 111 112 @end 113 114 NS_ASSUME_NONNULL_END 115