Update base for Update on "[executorch] Propagate device metadata from partitioner result onto TensorSpecs"

Add end-to-end device type annotation support from export to runtime. Currently we only support one device per graph

The overall pipeline is:
a. Partitioner use `compile_spec` to determine which device the partitoned blob is runing on
b. after lowered partitioned graph to backend, the new-introed propagate_device_pass will annotate the input and output tensors of delegate blob as target device.

Differential Revision: [D95842511](https://our.internmc.facebook.com/intern/diff/D95842511/)

[ghstack-poisoned]

Author: Gasoonjia

.ci/scripts/build-qnn-sdk.sh

@@ -40,6 +40,8 @@ set_up_aot() {
       -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \
       -DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=ON \
       -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
+      -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
+      -DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON \
       -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \
       -DPYTHON_EXECUTABLE=python3
   cmake --build $PWD --target "PyQnnManagerAdaptor" -j$(nproc)

.github/workflows/android-release-artifacts.yml

@@ -53,6 +53,9 @@ jobs:
         run: |
           VERSION="${{ inputs.version }}"
           FLAVOR="${{ inputs.flavor }}"
+          if [ -z "$FLAVOR" ]; then
+            FLAVOR="xnnpack"
+          fi
           if [ -z "$VERSION" ]; then
             echo "No version name specified. Will create a snapshot AAR"
             echo "should-skip=false" >> $GITHUB_OUTPUT
@@ -115,9 +118,10 @@ jobs:
         fi
 
         FLAVOR="${{ inputs.flavor }}"
-        if [ ! -z "$FLAVOR" ]; then
-          GRADLE_ARGS+=" -Dflavor=${FLAVOR}"
+        if [ -z "$FLAVOR" ]; then
+          FLAVOR="xnnpack"
         fi
+        GRADLE_ARGS+=" -Dflavor=${FLAVOR}"
 
         if [[ "$FLAVOR" == "vulkan" || "$FLAVOR" == "all" || -z "$FLAVOR" ]]; then
           curl -O https://sdk.lunarg.com/sdk/download/1.4.321.1/linux/vulkansdk-linux-x86_64-1.4.321.1.tar.xz

.gitignore

@@ -18,6 +18,7 @@ cmake-out*
 cmake-out-android/
 build-android/
 build-x86/
+build-hexagon/
 dist/
 arm-scratch/
 executorch.egg-info

CMakeLists.txt

@@ -50,6 +50,17 @@ project(executorch)
 
 set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR})
 
+# Hexagon toolchain with release build complains about code in third party
+# libraries.
+if("${CMAKE_SYSTEM_PROCESSOR}" MATCHES "Hexagon" AND "${CMAKE_BUILD_TYPE}"
+                                                     STREQUAL "Release"
+)
+  add_compile_options(
+    -Wno-error=format -Wno-error=implicit-int-conversion
+    -Wno-error=unused-variable -Wno-error=unused-function
+  )
+endif()
+
 # --- ExecuTorch Version ---
 # Parse version from version.txt (single source of truth)
 file(READ "${EXECUTORCH_ROOT}/version.txt" ET_VERSION_STRING)

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.
      */