Fix test failures (pytorch#3876)

Summary:
Fixes two issues

-  Add model is now `x + y` earlier it used to be
```
int z = x + y;
z = z + x;
z = z + x;
z = z + z;
```

-  CoreML delegate was loading model manager at static initialization and there is a chance that executorch runtime might not have been initialized. `ETCoreMLModelManagerDelegate` loads it lazily, so that the runtime is guaranteed to be initialized.

Pull Request resolved: pytorch#3876

Reviewed By: cccclai

Differential Revision: D58257664

Pulled By: mergennachin

fbshipit-source-id: bd8197d552572134b79ef777da85799e07b257dd

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

@@ -104,7 +104,7 @@ NS_ASSUME_NONNULL_BEGIN
 ///
 /// @param error   On failure, error is filled with the failure information.
 /// @retval `YES` is the assets are purged otherwise `NO`.
-- (BOOL)purge:(NSError* __autoreleasing*)error;
+- (BOOL)purgeAndReturnError:(NSError* __autoreleasing*)error;
 
 /// The estimated size of the assets store. The returned value might not correct if the asset
 /// directory is tampered externally.

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

@@ -740,7 +740,7 @@ - (BOOL)_purge:(NSError * __autoreleasing *)error {
     return static_cast<BOOL>(status);
 }
 
-- (BOOL)purge:(NSError * __autoreleasing *)error {
+- (BOOL)purgeAndReturnError:(NSError * __autoreleasing *)error {
     __block BOOL result = 0;
     dispatch_sync(self.syncQueue, ^{
         result = [self _purge:error];

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

@@ -103,8 +103,11 @@ __attribute__((objc_subclassing_restricted))
 /// @retval `YES` if the model was pre-warmed otherwise `NO`.
 - (BOOL)prewarmModelWithHandle:(ModelHandle*)handle error:(NSError* __autoreleasing*)error;
 
-/// The `ETCoreMLAssetManager` instance used to manage models cache.
-@property (strong, readonly, nonatomic) ETCoreMLAssetManager* assetManager;
+/// Purges model cache.
+///
+/// @param error   On failure, error is filled with the failure information.
+/// @retval `YES` if the cache is purged otherwise `NO`.
+- (BOOL)purgeModelsCacheAndReturnError:(NSError* __autoreleasing*)error;
 
 @end
 

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

@@ -352,6 +352,7 @@ @interface ETCoreMLModelManager () {
 }
 
 @property (nonatomic, readonly, strong) NSFileManager *fileManager;
+@property (strong, readonly, nonatomic) ETCoreMLAssetManager* assetManager;
 @property (nonatomic, readonly, strong) NSMutableDictionary<NSValue *, id<ETCoreMLModelExecutor>> *handleToExecutorMap;
 @property (nonatomic, readonly, strong) NSMapTable<NSString *, dispatch_queue_t> *modelIdentifierToLoadingQueueMap;
 @property (nonatomic, readonly, strong) NSMutableDictionary<NSString *, ETCoreMLAsset *> *modelIdentifierToPrewarmedAssetMap;
@@ -832,4 +833,8 @@ - (BOOL)unloadModelWithHandle:(ModelHandle *)handle {
     return result;
 }
 
+- (BOOL)purgeModelsCacheAndReturnError:(NSError *__autoreleasing *)error {
+    return [self.assetManager purgeAndReturnError:error];
+}
+
 @end

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

@@ -10,7 +10,6 @@
 #import <ETCoreMLModel.h>
 #import <ETCoreMLModelManager.h>
 #import <ETCoreMLStrings.h>
-#import <atomic>
 #import <backend_delegate.h>
 #import <model_event_logger.h>
 #import <multiarray.h>
@@ -88,6 +87,153 @@ MLComputeUnits get_compute_units(const Buffer& buffer) {
 }
 } //namespace
 
+@interface ETCoreMLModelManagerDelegate : NSObject
+
+- (instancetype)init NS_UNAVAILABLE;
+
++ (instancetype)new NS_UNAVAILABLE;
+
+- (instancetype)initWithConfig:(BackendDelegate::Config)config NS_DESIGNATED_INITIALIZER;
+
+- (BOOL)loadAndReturnError:(NSError * _Nullable __autoreleasing *)error;
+
+- (void)loadAsynchronously;
+
+- (ModelHandle*)loadModelFromAOTData:(NSData*)data
+                       configuration:(MLModelConfiguration*)configuration
+                               error:(NSError* __autoreleasing*)error;
+
+- (BOOL)executeModelWithHandle:(ModelHandle*)handle
+                       argsVec:(const std::vector<executorchcoreml::MultiArray>&)argsVec
+                loggingOptions:(const executorchcoreml::ModelLoggingOptions&)loggingOptions
+                   eventLogger:(const executorchcoreml::ModelEventLogger* _Nullable)eventLogger
+                         error:(NSError* __autoreleasing*)error;
+
+- (BOOL)unloadModelWithHandle:(ModelHandle*)handle;
+
+- (BOOL)purgeModelsCacheAndReturnError:(NSError * _Nullable __autoreleasing *)error;
+
+@property (assign, readonly, nonatomic) BackendDelegate::Config config;
+@property (strong, readonly, nonatomic) dispatch_queue_t syncQueue;
+@property (strong, nonatomic, nullable) ETCoreMLModelManager *impl;
+@property (assign, readonly, nonatomic) BOOL isAvailable;
+
+@end
+
+@implementation ETCoreMLModelManagerDelegate
+
+- (instancetype)initWithConfig:(BackendDelegate::Config)config {
+    self = [super init];
+    if (self) {
+        _config = std::move(config);
+        _syncQueue = dispatch_queue_create("com.executorchcoreml.modelmanagerdelegate.sync", DISPATCH_QUEUE_SERIAL_WITH_AUTORELEASE_POOL);
+    }
+
+    return self;
+}
+
+- (BOOL)_loadAndReturnError:(NSError * _Nullable __autoreleasing *)error {
+    if (self.impl != nil) {
+        return YES;
+    }
+
+    ETCoreMLAssetManager *assetManager = create_asset_manager(ETCoreMLStrings.assetsDirectoryPath,
+                                                              ETCoreMLStrings.trashDirectoryPath,
+                                                              ETCoreMLStrings.databaseDirectoryPath,
+                                                              ETCoreMLStrings.databaseName,
+                                                              self.config.max_models_cache_size,
+                                                              error);
+    if (!assetManager) {
+        return NO;
+    }
+
+    ETCoreMLModelManager *modelManager = [[ETCoreMLModelManager alloc] initWithAssetManager:assetManager];
+    if (!modelManager) {
+        return NO;
+    }
+
+    self.impl = modelManager;
+
+    if (self.config.should_prewarm_asset) {
+        [modelManager prewarmRecentlyUsedAssetsWithMaxCount:1];
+    }
+
+    return YES;
+}
+
+- (BOOL)loadAndReturnError:(NSError * _Nullable __autoreleasing *)error {
+    __block NSError *localError = nil;
+    __block BOOL result = NO;
+    dispatch_sync(self.syncQueue, ^{
+        result = [self _loadAndReturnError:&localError];
+    });
+
+    if (error) {
+        *error = localError;
+    }
+
+    return result;
+}
+
+- (void)loadAsynchronously {
+    dispatch_async(self.syncQueue, ^{
+        (void)[self _loadAndReturnError:nil];
+    });
+}
+
+- (ModelHandle*)loadModelFromAOTData:(NSData*)data
+                       configuration:(MLModelConfiguration*)configuration
+                               error:(NSError* __autoreleasing*)error {
+    if (![self loadAndReturnError:error]) {
+        return nil;
+    }
+
+    return [self.impl loadModelFromAOTData:data
+                             configuration:configuration
+                                     error:error];
+}
+
+- (BOOL)executeModelWithHandle:(ModelHandle*)handle
+                       argsVec:(const std::vector<executorchcoreml::MultiArray>&)argsVec
+                loggingOptions:(const executorchcoreml::ModelLoggingOptions&)loggingOptions
+                   eventLogger:(const executorchcoreml::ModelEventLogger* _Nullable)eventLogger
+                         error:(NSError* __autoreleasing*)error {
+    assert(self.impl != nil && "Impl must not be nil");
+    return [self.impl executeModelWithHandle:handle
+                                     argsVec:argsVec
+                              loggingOptions:loggingOptions
+                                 eventLogger:eventLogger
+                                       error:error];
+}
+
+- (nullable ETCoreMLModel*)modelWithHandle:(ModelHandle*)handle {
+    assert(self.impl != nil && "Impl must not be nil");
+    return [self.impl modelWithHandle:handle];
+}
+
+- (BOOL)unloadModelWithHandle:(ModelHandle*)handle {
+    assert(self.impl != nil && "Impl must not be nil");
+    return [self.impl unloadModelWithHandle:handle];
+}
+
+- (BOOL)purgeModelsCacheAndReturnError:(NSError * _Nullable __autoreleasing *)error {
+    if (![self loadAndReturnError:error]) {
+        return NO;
+    }
+
+    return [self.impl purgeModelsCacheAndReturnError:error];;
+}
+
+- (BOOL)isAvailable {
+    if (![self loadAndReturnError:nil]) {
+        return NO;
+    }
+
+    return YES;
+}
+
+@end
+
 namespace executorchcoreml {
 
 std::string BackendDelegate::ErrorCategory::message(int code) const {
@@ -114,20 +260,9 @@ MLComputeUnits get_compute_units(const Buffer& buffer) {
 class BackendDelegateImpl: public BackendDelegate {
 public:
     explicit BackendDelegateImpl(const Config& config) noexcept
-    :BackendDelegate(), config_(config) {
-        NSError *localError = nil;
-        ETCoreMLAssetManager *asset_manager = create_asset_manager(ETCoreMLStrings.assetsDirectoryPath,
-                                                                   ETCoreMLStrings.trashDirectoryPath,
-                                                                   ETCoreMLStrings.databaseDirectoryPath,
-                                                                   ETCoreMLStrings.databaseName,
-                                                                   config.max_models_cache_size,
-                                                                   &localError);
-
-        model_manager_ = (asset_manager != nil) ? [[ETCoreMLModelManager alloc] initWithAssetManager:asset_manager] : nil;
-        if (model_manager_ != nil && config_.should_prewarm_asset) {
-            [model_manager_ prewarmRecentlyUsedAssetsWithMaxCount:1];
-        }
-        available_.store(model_manager_ != nil, std::memory_order_seq_cst);
+    :BackendDelegate(), model_manager_([[ETCoreMLModelManagerDelegate alloc] initWithConfig:config])
+    {
+        [model_manager_ loadAsynchronously];
     }
 
     BackendDelegateImpl(BackendDelegateImpl const&) = delete;
@@ -142,11 +277,6 @@ explicit BackendDelegateImpl(const Config& config) noexcept
         ModelHandle *modelHandle = [model_manager_ loadModelFromAOTData:data
                                                           configuration:configuration
                                                                   error:&localError];
-        if (modelHandle && config_.should_prewarm_model) {
-            NSError *localError = nil;
-            [model_manager_ prewarmModelWithHandle:modelHandle error:&localError];
-        }
-
         return modelHandle;
     }
 
@@ -158,7 +288,7 @@ bool execute(Handle* handle,
         NSError *error = nil;
         if (![model_manager_ executeModelWithHandle:handle
                                             argsVec:args
-                                    loggingOptions:logging_options
+                                     loggingOptions:logging_options
                                         eventLogger:event_logger
                                               error:&error]) {
             ec = static_cast<ErrorCode>(error.code);
@@ -173,7 +303,7 @@ bool is_valid_handle(Handle* handle) const noexcept override {
     }
 
     bool is_available() const noexcept override {
-        return available_.load(std::memory_order_acquire);
+        return static_cast<bool>(model_manager_.isAvailable);
     }
 
     std::pair<size_t, size_t> get_num_arguments(Handle* handle) const noexcept override {
@@ -187,12 +317,11 @@ void destroy(Handle* handle) const noexcept override {
 
     bool purge_models_cache() const noexcept override {
         NSError *localError = nil;
-        bool result = static_cast<bool>([model_manager_.assetManager purge:&localError]);
+        bool result = static_cast<bool>([model_manager_ purgeModelsCacheAndReturnError:&localError]);
         return result;
     }
 
-    ETCoreMLModelManager *model_manager_;
-    std::atomic<bool> available_;
+    ETCoreMLModelManagerDelegate *model_manager_;
     Config config_;
 };
 

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

@@ -11,6 +11,7 @@
 #import <backend_delegate.h>
 #import <coreml_backend/delegate.h>
 #import <executorch/runtime/core/evalue.h>
+#import <executorch/runtime/platform/log.h>
 #import <memory>
 #import <model_event_logger.h>
 #import <model_logging_options.h>

backends/apple/coreml/runtime/test/BackendDelegateTests.mm

@@ -15,6 +15,7 @@
 #import <model_logging_options.h>
 #import <multiarray.h>
 #import <objc_array_util.h>
+#import <executorch/runtime/platform/runtime.h>
 
 using namespace executorchcoreml;
 
@@ -58,6 +59,10 @@ + (nullable NSURL *)bundledResourceWithName:(NSString *)name extension:(NSString
     return [bundle URLForResource:name withExtension:extension];
 }
 
++ (void)setUp {
+    torch::executor::runtime_init();
+}
+
 - (void)setUp {
     @autoreleasepool {
         _delegate = BackendDelegate::make(BackendDelegate::Config());
@@ -151,9 +156,6 @@ - (void)testAddModelExecution {
     int y = 50;
     // add_coreml_all does the following operations.
     int z = x + y;
-    z = z + x;
-    z = z + x;
-    z = z + z;
 
     NSArray<MLMultiArray *> *inputs = [ETCoreMLTestUtils inputsForModel:model repeatedValues:@[@(x), @(y)] error:&localError];
     XCTAssertNotNil(inputs);

backends/apple/coreml/runtime/test/CoreMLBackendDelegateTests.mm

@@ -6,14 +6,12 @@
 // Please refer to the license found in the LICENSE file in the root directory of the source tree.
 
 #import <XCTest/XCTest.h>
-
-#import <string>
-
 #import <executorch/runtime/core/data_loader.h>
 #import <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
 #import <executorch/runtime/executor/method.h>
 #import <executorch/runtime/executor/program.h>
 #import <executorch/runtime/platform/runtime.h>
+#import <string>
 
 static constexpr size_t kRuntimeMemorySize = 50 * 1024U * 1024U; // 50 MB
 
@@ -128,7 +126,7 @@ @interface CoreMLBackendDelegateTests : XCTestCase
 @implementation CoreMLBackendDelegateTests
 
 + (void)setUp {
-    runtime_init();
+    torch::executor::runtime_init();
 }
 
 + (nullable NSURL *)bundledResourceWithName:(NSString *)name extension:(NSString *)extension {

backends/apple/coreml/runtime/test/ETCoreMLAssetManagerTests.mm

@@ -6,12 +6,11 @@
 //
 // Please refer to the license found in the LICENSE file in the root directory of the source tree.
 
-#import <XCTest/XCTest.h>
-
+#import "ETCoreMLTestUtils.h"
 #import <ETCoreMLAsset.h>
 #import <ETCoreMLAssetManager.h>
-
-#import "ETCoreMLTestUtils.h"
+#import <XCTest/XCTest.h>
+#import <executorch/runtime/platform/runtime.h>
 
 @interface ETCoreMLAssetManagerTests : XCTestCase
 
@@ -23,6 +22,10 @@ @interface ETCoreMLAssetManagerTests : XCTestCase
 
 @implementation ETCoreMLAssetManagerTests
 
++ (void)setUp {
+    torch::executor::runtime_init();
+}
+
 - (void)setUp {
     @autoreleasepool {
         NSError *localError = nil;
@@ -145,7 +148,7 @@ - (void)testPurge {
         // Close the asset so that it could be deleted.
         [asset close];
     }
-    XCTAssertTrue([self.assetManager purge:&localError]);
+    XCTAssertTrue([self.assetManager purgeAndReturnError:&localError]);
 }
 
 @end