Skip to content

Add prefill/decode multifunction support in ET #16552

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
metascroy wants to merge 1 commit into from
Closed

Add prefill/decode multifunction support in ET #16552

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
493 changes: 430 additions & 63 deletions backends/apple/coreml/compiler/coreml_preprocess.py
View file
Open in desktop

Large diffs are not rendered by default.

24 changes: 24 additions & 0 deletions backends/apple/coreml/runtime/delegate/ETCoreMLModel.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,12 +200,27 @@ - (nullable instancetype)initWithAsset:(ETCoreMLAsset *)asset
return nil;
}

if (orderedInputNames == nil) {
ETCoreMLLogErrorAndSetNSError(error,
ETCoreMLErrorCorruptedModel,
"orderedInputNames must not be nil");
return nil;
}

if (orderedOutputNames == nil) {
ETCoreMLLogErrorAndSetNSError(error,
ETCoreMLErrorCorruptedModel,
"orderedOutputNames must not be nil");
return nil;
}

self = [super init];
if (self) {
_mlModel = mlModel;
_asset = asset;
_orderedInputNames = [orderedInputNames copy];
_orderedOutputNames = [orderedOutputNames copy];

_cache = [[NSCache alloc] init];
_inputConstraintsByName = get_multi_array_input_constraints_by_name(mlModel.modelDescription);
_outputConstraintsByName = get_multi_array_output_constraints_by_name(mlModel.modelDescription);
Expand Down Expand Up @@ -234,6 +249,15 @@ - (NSString *)identifier {
BOOL lCopyData = copyData;
NSString *argName = [nameEnumerator nextObject];
MLMultiArrayConstraint *constraint = argConstraintsByName[argName];

if (constraint == nil) {
ETCoreMLLogErrorAndSetNSError(error,
ETCoreMLErrorCorruptedModel,
"No constraint found for arg '%@'. Model may have mismatched input/output names.",
argName);
return nil;
}

const auto& layout = arg.layout();
auto dataType = to_ml_multiarray_data_type(layout.dataType());
MLMultiArray *multiArrayArg = nil;
Expand Down
32 changes: 17 additions & 15 deletions backends/apple/coreml/runtime/delegate/ETCoreMLModelLoader.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,20 +26,6 @@

return result;
}

ETCoreMLModel * _Nullable get_model_from_asset(ETCoreMLAsset *asset,
MLModelConfiguration *configuration,
const executorchcoreml::ModelMetadata& metadata,
NSError * __autoreleasing *error) {
NSOrderedSet<NSString *> *orderedInputNames = ::get_ordered_set(metadata.input_names);
NSOrderedSet<NSString *> *orderedOutputNames = ::get_ordered_set(metadata.output_names);
ETCoreMLModel *model = [[ETCoreMLModel alloc] initWithAsset:asset
configuration:configuration
orderedInputNames:orderedInputNames
orderedOutputNames:orderedOutputNames
error:error];
return model;
}
} // namespace

@implementation ETCoreMLModelLoader
Expand All @@ -48,14 +34,30 @@ + (nullable ETCoreMLModel *)loadModelWithCompiledAsset:(ETCoreMLAsset *)compiled
configuration:(MLModelConfiguration *)configuration
metadata:(const executorchcoreml::ModelMetadata&)metadata
error:(NSError * __autoreleasing *)error {
if (compiledAsset == nil) {
return nil;
}

// Use the metadata's ordered input/output names.
// For multifunction models, the caller should load the per-method metadata
// which contains the correct input/output names for that method.
NSOrderedSet<NSString *> *orderedInputNames = ::get_ordered_set(metadata.input_names);
NSOrderedSet<NSString *> *orderedOutputNames = ::get_ordered_set(metadata.output_names);

NSError *localError = nil;
ETCoreMLModel *model = (compiledAsset != nil) ? get_model_from_asset(compiledAsset, configuration, metadata, &localError) : nil;
ETCoreMLModel *model = [[ETCoreMLModel alloc] initWithAsset:compiledAsset
configuration:configuration
orderedInputNames:orderedInputNames
orderedOutputNames:orderedOutputNames
error:&localError];
if (model) {
return model;
}

if (error) {
*error = localError;
}

return nil;
}

Expand Down
21 changes: 21 additions & 0 deletions backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,27 @@ __attribute__((objc_subclassing_restricted))
configuration:(MLModelConfiguration*)configuration
error:(NSError* __autoreleasing*)error;

/// Loads the model from the AOT data with an optional method name for multifunction model support.
///
/// The data is the AOT blob stored in the executorch Program. The method first parses the model
/// metadata stored in the blob and extracts the identifier. For multifunction CoreML models,
/// the methodName is used to populate the correct input/output names from the method's metadata,
/// while functionName specifies the CoreML function to invoke (via MLModelConfiguration.functionName).
/// The method/function names are NOT included in the cache key since the multifunction
/// model contains all functions and should only be compiled once.
///
/// @param data The AOT blob data.
/// @param configuration The model configuration that will be used to load the model.
/// @param methodName Optional method name (e.g., "forward", "prefill") for metadata lookup.
/// @param functionName Optional CoreML function name to invoke. If nil, methodName is used.
/// @param error On failure, error is filled with the failure information.
/// @retval An opaque handle that points to the loaded model.
- (ModelHandle*)loadModelFromAOTData:(NSData*)data
configuration:(MLModelConfiguration*)configuration
methodName:(nullable NSString*)methodName
functionName:(nullable NSString*)functionName
error:(NSError* __autoreleasing*)error;

/// Executes the loaded model.
///
/// @param handle The handle to the loaded model.
Expand Down
Loading
Loading