Merge branch 'main' into toupstream/dim_pass

Author: SaoirseARM

backends/aoti/aoti_delegate_handle.h

@@ -31,6 +31,20 @@ using AOTInductorModelContainerHandle = AOTInductorModelContainerOpaque*;
 using AOTInductorStreamHandle = void*;
 using AOTIProxyExecutorHandle = void*;
 
+// Opaque types for AOTI constant management.
+// AtenTensorOpaque wraps at::Tensor* in the AOTI runtime — distinct from
+// AOTITensorHandle which wraps executorch::runtime::etensor::Tensor*.
+struct AtenTensorOpaque;
+using AtenTensorHandle = AtenTensorOpaque*;
+
+struct AOTInductorConstantMap;
+using AOTInductorConstantMapHandle = AOTInductorConstantMap*;
+
+struct AOTInductorConstantMapEntry {
+  const char* name;
+  AtenTensorHandle handle;
+};
+
 // Function pointer types for AOT Inductor model container operations
 using AOTInductorModelContainerCreateWithDeviceFunc = AOTIRuntimeError (*)(
     AOTInductorModelContainerHandle* container_handle,
@@ -77,6 +91,37 @@ using AOTInductorModelUpdateConstantsFromBlobFunc = AOTIRuntimeError (*)(
     AOTInductorModelContainerHandle container_handle,
     const uint8_t* weight_blob_ptr);
 
+// Retrieves a constant's AOTI internal name by index.
+using AOTInductorModelContainerGetConstantNameFunc = AOTIRuntimeError (*)(
+    AOTInductorModelContainerHandle container_handle,
+    size_t idx,
+    const char** name);
+
+// Retrieves a constant's original fully-qualified name by index.
+using AOTInductorModelContainerGetConstantOriginalFQNFunc =
+    AOTIRuntimeError (*)(
+        AOTInductorModelContainerHandle container_handle,
+        size_t idx,
+        const char** original_fqn);
+
+// Extracts the constants map from the container (active or inactive buffer).
+// constant_map_handle should point to a
+// std::unordered_map<std::string, AtenTensorHandle>.
+using AOTInductorModelContainerExtractConstantsMapFunc = AOTIRuntimeError (*)(
+    AOTInductorModelContainerHandle container_handle,
+    AOTInductorConstantMapHandle constant_map_handle,
+    bool use_inactive);
+
+// Updates the container's constants with user-managed tensor handles.
+// DLL-boundary safe — uses a flat C array instead of std::unordered_map.
+using AOTInductorModelContainerUpdateUserManagedConstantBufferPairsFunc =
+    AOTIRuntimeError (*)(
+        AOTInductorModelContainerHandle container_handle,
+        const AOTInductorConstantMapEntry* pairs,
+        size_t num_pairs,
+        bool use_inactive,
+        bool validate_full_update);
+
 } // extern "C"
 
 // AOTI Delegate Handle structure
@@ -93,6 +138,14 @@ struct AOTIDelegateHandle {
   AOTInductorModelContainerGetNumOutputsFunc get_num_outputs;
   AOTInductorModelContainerRunFunc run;
   AOTInductorModelUpdateConstantsFromBlobFunc update_constants_from_blob;
+
+  // Constant management function pointers (for cross-method buffer sharing)
+  AOTInductorModelContainerGetNumConstantsFunc get_num_constants;
+  AOTInductorModelContainerGetConstantNameFunc get_constant_name;
+  AOTInductorModelContainerGetConstantOriginalFQNFunc get_constant_original_fqn;
+  AOTInductorModelContainerExtractConstantsMapFunc extract_constants_map;
+  AOTInductorModelContainerUpdateUserManagedConstantBufferPairsFunc
+      update_user_managed_constant_buffer_pairs;
 };
 
 } // namespace aoti

backends/apple/coreml/runtime/delegate/backend_delegate.h

@@ -9,6 +9,9 @@
 
 #include "model_logging_options.h"
 
+#include <executorch/runtime/backend/options.h>
+#include <executorch/runtime/core/span.h>
+
 #include <system_error>
 #include <unordered_map>
 #include <vector>
@@ -75,12 +78,15 @@ class BackendDelegate {
     /// @param method_name The ExecuTorch method name for metadata lookup (optional, may be nullptr).
     /// @param function_name The CoreML function name to invoke (optional, may be nullptr).
     ///                      If nullptr, method_name is used as the function name.
+    /// @param runtime_specs Runtime options passed via BackendOptions (e.g., cache_dir).
     /// @retval An opaque handle to the initialized blob or `nullptr` if the
     /// initialization failed.
-    virtual Handle* init(Buffer processed,
-                         const std::unordered_map<std::string, Buffer>& specs,
-                         const char* method_name = nullptr,
-                         const char* function_name = nullptr) const noexcept = 0;
+    virtual Handle*
+    init(Buffer processed,
+         const std::unordered_map<std::string, Buffer>& specs,
+         const char* method_name = nullptr,
+         const char* function_name = nullptr,
+         executorch::runtime::Span<const executorch::runtime::BackendOption> runtime_specs = {}) const noexcept = 0;
 
     /// Must execute the CoreML model with the specified handle.
     ///

backends/apple/coreml/runtime/delegate/backend_delegate.mm

@@ -11,6 +11,7 @@
 #import "ETCoreMLAssetManager.h"
 #import "ETCoreMLLogging.h"
 #import "ETCoreMLModel.h"
+#import "ETCoreMLModelCache.h"
 #import "ETCoreMLModelManager.h"
 #import "ETCoreMLStrings.h"
 #import "model_event_logger.h"
@@ -100,7 +101,14 @@ - (ModelHandle*)loadModelFromAOTData:(NSData*)data
                        configuration:(MLModelConfiguration*)configuration
                           methodName:(nullable NSString*)methodName
                         functionName:(nullable NSString*)functionName
-                               error:(NSError* __autoreleasing*)error;
+                                error:(NSError* __autoreleasing*)error;
+
+- (ModelHandle*)loadModelFromAOTData:(NSData*)data
+                       configuration:(MLModelConfiguration*)configuration
+                          methodName:(nullable NSString*)methodName
+                        functionName:(nullable NSString*)functionName
+                           cachePath:(nullable NSString*)cachePath
+                                error:(NSError* __autoreleasing*)error;
 
 - (ModelHandle*)loadModelFromAOTData:(NSData*)data
                        configuration:(MLModelConfiguration*)configuration
@@ -119,6 +127,7 @@ - (BOOL)purgeModelsCacheAndReturnError:(NSError * _Nullable __autoreleasing *)er
 @property (assign, readonly, nonatomic) BackendDelegate::Config config;
 @property (strong, readonly, nonatomic) dispatch_queue_t syncQueue;
 @property (strong, nonatomic, nullable) ETCoreMLModelManager *impl;
+@property (strong, nonatomic, nullable) ETCoreMLModelCache *defaultCache;
 @property (assign, readonly, nonatomic) BOOL isAvailable;
 
 @end
@@ -157,6 +166,16 @@ - (BOOL)_loadAndReturnError:(NSError * _Nullable __autoreleasing *)error {
 
     self.impl = modelManager;
 
+    // Create default filesystem cache at the same location as assets
+    NSURL *defaultCacheURL = [NSURL fileURLWithPath:ETCoreMLStrings.assetsDirectoryPath isDirectory:YES];
+    ETCoreMLModelCache *defaultCache = [[ETCoreMLModelCache alloc] initWithCacheRootDirectory:defaultCacheURL];
+    if (defaultCache.isReady) {
+        self.defaultCache = defaultCache;
+    } else {
+        ETCoreMLLogError(defaultCache.initializationError,
+                         "Default cache initialization failed, will use asset manager as fallback");
+    }
+
     if (self.config.should_prewarm_asset) {
         [modelManager prewarmRecentlyUsedAssetsWithMaxCount:1];
     }
@@ -191,22 +210,86 @@ - (ModelHandle*)loadModelFromAOTData:(NSData*)data
                         configuration:configuration
                            methodName:nil
                          functionName:nil
+                            cachePath:nil
+                                error:error];
+}
+
+- (ModelHandle*)loadModelFromAOTData:(NSData*)data
+                       configuration:(MLModelConfiguration*)configuration
+                          methodName:(nullable NSString*)methodName
+                        functionName:(nullable NSString*)functionName
+                               error:(NSError* __autoreleasing*)error {
+    return [self loadModelFromAOTData:data
+                        configuration:configuration
+                           methodName:methodName
+                         functionName:functionName
+                            cachePath:nil
                                 error:error];
 }
 
 - (ModelHandle*)loadModelFromAOTData:(NSData*)data
                        configuration:(MLModelConfiguration*)configuration
                           methodName:(nullable NSString*)methodName
                         functionName:(nullable NSString*)functionName
+                           cachePath:(nullable NSString*)cachePath
+                               error:(NSError* __autoreleasing*)error {
+    // Default to using the old cache (useNewCache = NO)
+    return [self loadModelFromAOTData:data
+                        configuration:configuration
+                           methodName:methodName
+                         functionName:functionName
+                            cachePath:cachePath
+                          useNewCache:NO
+                                error:error];
+}
+
+- (ModelHandle*)loadModelFromAOTData:(NSData*)data
+                       configuration:(MLModelConfiguration*)configuration
+                          methodName:(nullable NSString*)methodName
+                        functionName:(nullable NSString*)functionName
+                           cachePath:(nullable NSString*)cachePath
+                         useNewCache:(BOOL)useNewCache
                                error:(NSError* __autoreleasing*)error {
     if (![self loadAndReturnError:error]) {
         return nil;
     }
 
+    id<ETCoreMLCache> cache = nil;
+    if (cachePath != nil) {
+        // Use NEW filesystem cache at specified path
+        NSURL *cacheURL = [NSURL fileURLWithPath:cachePath isDirectory:YES];
+        ETCoreMLModelCache *modelCache = [[ETCoreMLModelCache alloc] initWithCacheRootDirectory:cacheURL];
+        if (!modelCache.isReady) {
+            // Fallback error if initializationError is unexpectedly nil
+            NSError *cacheError = modelCache.initializationError
+                ?: [NSError errorWithDomain:ETCoreMLModelCacheErrorDomain
+                                       code:ETCoreMLModelCacheErrorCodeInitializationFailed
+                                   userInfo:@{NSLocalizedDescriptionKey: @"Cache initialization failed"}];
+            if (error) *error = cacheError;
+            return nil;
+        }
+        cache = modelCache;
+    } else if (useNewCache) {
+        if (self.defaultCache != nil) {
+            // Use default filesystem cache
+            cache = self.defaultCache;
+} else {
+            // Fallback: useNewCache requested but default cache unavailable
+            NSError *fallbackError = [NSError errorWithDomain:ETCoreMLErrorDomain
+                                                         code:ETCoreMLErrorInternalError
+                                                     userInfo:@{NSLocalizedDescriptionKey: @"Default cache unavailable"}];
+            ETCoreMLLogError(fallbackError,
+                             "useNewCache=YES but default cache is unavailable, falling back to asset manager");
+        }
+    }
+    // If useNewCache is false or defaultCache is nil, cache remains nil
+    // and loadModelFromAOTData will use the asset manager path
+
     auto handle = [self.impl loadModelFromAOTData:data
                                     configuration:configuration
                                        methodName:methodName
                                      functionName:functionName
+                                            cache:cache
                                             error:error];
     if ((handle != NULL) && self.config.should_prewarm_model) {
         [self.impl prewarmModelWithHandle:handle error:nil];
@@ -291,9 +374,10 @@ explicit BackendDelegateImpl(const Config& config) noexcept
     BackendDelegateImpl& operator=(BackendDelegateImpl const&) = delete;
 
 Handle *init(Buffer processed,
-                     const std::unordered_map<std::string, Buffer>& specs,
-                     const char* method_name = nullptr,
-                     const char* function_name = nullptr) const noexcept override {
+             const std::unordered_map<std::string, Buffer>& specs,
+             const char* method_name = nullptr,
+             const char* function_name = nullptr,
+             executorch::runtime::Span<const executorch::runtime::BackendOption> runtime_specs = {}) const noexcept override {
         NSError *localError = nil;
         MLModelConfiguration *configuration = get_model_configuration(specs, &localError);
         if (configuration == nil) {
@@ -304,13 +388,31 @@ explicit BackendDelegateImpl(const Config& config) noexcept
         NSString *methodNameStr = method_name ? @(method_name) : nil;
         NSString *functionNameStr = function_name ? @(function_name) : nil;
 
+        // Parse cache_dir and _use_new_cache from runtime_specs
+        NSString *cachePath = nil;
+        BOOL useNewCache = NO; // Default to using the old cache (asset manager)
+        for (size_t i = 0; i < runtime_specs.size(); ++i) {
+            const auto& opt = runtime_specs[i];
+            if (std::strcmp(opt.key, "cache_dir") == 0) {
+                if (auto* arr = std::get_if<std::array<char, executorch::runtime::kMaxOptionValueLength>>(&opt.value)) {
+                    cachePath = @(arr->data());
+                }
+            } else if (std::strcmp(opt.key, "_use_new_cache") == 0) {
+                if (auto* val = std::get_if<bool>(&opt.value)) {
+                    useNewCache = *val ? YES : NO;
+                }
+            }
+        }
+
         NSData *data = [NSData dataWithBytesNoCopy:const_cast<void *>(processed.data())
                                             length:processed.size()
                                       freeWhenDone:NO];
         ModelHandle *modelHandle = [model_manager_ loadModelFromAOTData:data
                                                           configuration:configuration
                                                              methodName:methodNameStr
                                                            functionName:functionNameStr
+                                                              cachePath:cachePath
+                                                            useNewCache:useNewCache
                                                                   error:&localError];
         if (localError != nil) {
             ETCoreMLLogError(localError, "Model init failed");

backends/apple/coreml/runtime/delegate/coreml_backend_delegate.mm

@@ -336,7 +336,7 @@ ModelLoggingOptions get_logging_options(BackendExecutionContext& context) {
 
     std::error_code error;
     const char* function_name_cstr = functionName.empty() ? nullptr : functionName.c_str();
-    auto handle = impl_->init(std::move(buffer), specs_map, method_name, function_name_cstr);
+    auto handle = impl_->init(std::move(buffer), specs_map, method_name, function_name_cstr, runtime_specs);
     ET_CHECK_OR_RETURN_ERROR(handle != nullptr,
                              InvalidProgram,
                              "%s: Failed to init the model.", ETCoreMLStrings.delegateIdentifier.UTF8String);

backends/apple/coreml/runtime/include/coreml_backend/coreml_backend_options.h

@@ -75,6 +75,41 @@ class LoadOptionsBuilder {
         return *this;
     }
 
+    /**
+     * Sets a custom cache directory for compiled model storage.
+     *
+     * When specified, the backend uses the new filesystem-based cache
+     * (ETCoreMLModelCache) at the given path instead of the default
+     * SQLite-based cache (ETCoreMLAssetManager).
+     *
+     * This enables per-model cache selection for experimentation:
+     * - Specify cache_dir to opt-in to the new filesystem cache
+     * - Omit cache_dir to use the default (legacy) cache
+     *
+     * @param path The directory path for the cache. Must be a valid
+     *             filesystem path with write permissions.
+     * @return Reference to this builder for chaining.
+     */
+    LoadOptionsBuilder& setCacheDirectory(const char* path) {
+        options_.set_option("cache_dir", path);
+        return *this;
+    }
+
+    /**
+     * Controls whether to use the new filesystem cache (ETCoreMLModelCache).
+     *
+     * This is a temporary runtime option for A/B testing the new cache
+     * implementation. It will be removed once the new cache is fully rolled out.
+     *
+     * @param enabled If true, uses the new filesystem cache.
+     *                If false (default), uses the legacy asset manager.
+     * @return Reference to this builder for chaining.
+     */
+    LoadOptionsBuilder& setUseNewCache(bool enabled) {
+        options_.set_option("_use_new_cache", enabled);
+        return *this;
+    }
+
     /**
      * Returns the backend identifier for this options builder.
      */

backends/apple/coreml/runtime/sdk/ETCoreMLModelProfiler.mm

@@ -276,6 +276,7 @@ - (nullable instancetype)initWithModel:(ETCoreMLModel *)model
         self = [super init];
         if (self) {
             _model = model;
+            _outputNames = model.orderedOutputNames;
             _computePlan = computePlan;
             _operationToPathMap = operationToPathMap;
             _topologicallySortedOperations = topologicallySortedOperations;