[Core ML] Implement intermediate tensor logging

Differential Revision: D60838372 Pull Request resolved: pytorch#4557
metascroy · Aug 15, 2024 · a9ed835 · a9ed835
1 parent 35a15a6
commit a9ed835
Show file tree

Hide file tree

Showing 33 changed files with 1,281 additions and 218 deletions.
diff --git a/backends/apple/coreml/CMakeLists.txt b/backends/apple/coreml/CMakeLists.txt
@@ -13,6 +13,11 @@ if(NOT EXECUTORCH_ROOT)
   set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..)
 endif()
 
+if(EXECUTORCH_BUILD_SDK)
+# protobuf requires frtti
+set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -frtti" )
+endif()
+
 option(COREML_BUILD_EXECUTOR_RUNNER "Build CoreML executor runner." OFF)
 
 # inmemoryfs sources
@@ -59,6 +64,7 @@ set(SDK_SOURCES
     runtime/sdk/ETCoreMLModelAnalyzer.mm
     runtime/sdk/ETCoreMLModelStructurePath.mm
     runtime/sdk/ETCoreMLOperationProfilingInfo.mm
+    runtime/sdk/ETCoreMLModelDebugInfo.mm
     runtime/sdk/ETCoreMLModelDebugger.mm
     runtime/sdk/ETCoreMLModelProfiler.mm
     runtime/sdk/ETCoreMLPair.mm
@@ -141,6 +147,8 @@ if(EXECUTORCH_BUILD_SDK)
   add_subdirectory(
     ${CMAKE_CURRENT_SOURCE_DIR}/third-party/coremltools/deps/protobuf/cmake
   )
+
+  target_link_options_shared_lib(libprotobuf-lite)
   target_link_libraries(coremldelegate PRIVATE libprotobuf-lite)
 endif()
 

diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm
@@ -5,36 +5,37 @@
 //
 // Please refer to the license found in the LICENSE file in the root directory of the source tree.
 
-#import <ETCoreMLAsset.h>
-#import <ETCoreMLAssetManager.h>
-#import <ETCoreMLDefaultModelExecutor.h>
-#import <ETCoreMLLogging.h>
-#import <ETCoreMLModel.h>
-#import <ETCoreMLModelCompiler.h>
-#import <ETCoreMLModelExecutor.h>
-#import <ETCoreMLModelLoader.h>
-#import <ETCoreMLModelManager.h>
-#import <ETCoreMLStrings.h>
-#import <MLModel_Prewarm.h>
-#import <MLMultiArray_Copy.h>
+#import "ETCoreMLAsset.h"
+#import "ETCoreMLAssetManager.h"
+#import "ETCoreMLDefaultModelExecutor.h"
+#import "ETCoreMLLogging.h"
+#import "ETCoreMLModel.h"
+#import "ETCoreMLModelCompiler.h"
+#import "ETCoreMLModelExecutor.h"
+#import "ETCoreMLModelLoader.h"
+#import "ETCoreMLModelManager.h"
+#import "ETCoreMLStrings.h"
+#import "MLModel_Prewarm.h"
+#import "MLMultiArray_Copy.h"
 #import <filesystem>
-#import <inmemory_filesystem_utils.hpp>
+#import "inmemory_filesystem_utils.hpp"
 #import <iostream>
 #import <memory>
-#import <model_metadata.h>
-#import <multiarray.h>
-#import <objc_array_util.h>
+#import "model_metadata.h"
+#import "multiarray.h"
+#import "objc_array_util.h"
 #import <optional>
 #import <os/lock.h>
-#import <serde_json.h>
+#import "serde_json.h"
 #import <string>
 #import <system_error>
 #import <vector>
 
 #if ET_EVENT_TRACER_ENABLED
-#import <ETCoreMLModelAnalyzer.h>
-#import <ETCoreMLModelStructurePath.h>
-#import <objc_safe_cast.h>
+#import "ETCoreMLModelAnalyzer.h"
+#import "ETCoreMLModelDebugInfo.h"
+#import "ETCoreMLModelStructurePath.h"
+#import "objc_safe_cast.h"
 #endif
 
 namespace {
@@ -317,31 +318,14 @@ void add_compute_unit(std::string& identifier, MLComputeUnits compute_units) {
     return [[ETCoreMLAsset alloc] initWithBackingAsset:std::move(backingAsset.value())];
 }
 
-NSDictionary<ETCoreMLModelStructurePath *, NSString *> * _Nullable get_operation_path_to_symbol_name_map(const inmemoryfs::InMemoryFileSystem *inMemoryFS,
-                                                                                                         NSError * __autoreleasing *error) {
+ETCoreMLModelDebugInfo * _Nullable get_model_debug_info(const inmemoryfs::InMemoryFileSystem *inMemoryFS,
+                                                        NSError * __autoreleasing *error) {
     NSData *file_data = get_file_data(inMemoryFS, ETCoreMLStrings.debugInfoFileRelativePath);
     if (!file_data) {
         return nil;
     }
-
-    id object = [NSJSONSerialization JSONObjectWithData:file_data options:(NSJSONReadingOptions)0 error:error];
-    if (!object) {
-        return nil;
-    }
-
-    NSDictionary<NSString *, id> *json_dict = SAFE_CAST(object, NSDictionary);
-    NSMutableDictionary<ETCoreMLModelStructurePath *, NSString *> *result = [NSMutableDictionary dictionaryWithCapacity:json_dict.count];
-    NSDictionary<NSString *, NSArray<id> *> *debug_symbol_to_operation_path_map = SAFE_CAST(json_dict[ETCoreMLStrings.debugSymbolToOperationPathKeyName], NSDictionary);
-    for (NSString *symbol_name in debug_symbol_to_operation_path_map) {
-        NSArray<NSDictionary<NSString *, id> *> *components = SAFE_CAST(debug_symbol_to_operation_path_map[symbol_name], NSArray);
-        if (components.count == 0) {
-            continue;
-        }
-        ETCoreMLModelStructurePath *path = [[ETCoreMLModelStructurePath alloc] initWithComponents:components];
-        result[path] = symbol_name;
-    }
-
-    return result;
+
+    return [ETCoreMLModelDebugInfo modelDebugInfoFromData:file_data error:error];
 }
 
 #endif
@@ -490,16 +474,16 @@ - (nullable NSURL *)compiledModelURLWithIdentifier:(NSString *)identifier
     }
 
     NSError *localError = nil;
-    NSDictionary<ETCoreMLModelStructurePath *, NSString *> *operation_path_to_symbol_name_map = get_operation_path_to_symbol_name_map(inMemoryFS,
-                                                                                                                                      &localError);
+    ETCoreMLModelDebugInfo *debug_info = get_model_debug_info(inMemoryFS, &localError);
     if (localError) {
         ETCoreMLLogError(localError, "Failed to parse debug info file");
     }
 
+
     return [[ETCoreMLModelAnalyzer alloc] initWithCompiledModelAsset:compiledModelAsset
                                                           modelAsset:modelAsset
+                                                      modelDebugInfo:debug_info
                                                             metadata:metadata
-                                       operationPathToDebugSymbolMap:operation_path_to_symbol_name_map
                                                        configuration:configuration
                                                         assetManager:self.assetManager
                                                                error:error];

diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLStrings.h b/backends/apple/coreml/runtime/delegate/ETCoreMLStrings.h
@@ -66,6 +66,8 @@ NS_ASSUME_NONNULL_BEGIN
 @property (class, copy, readonly, nonatomic, nullable) NSString* debugInfoFileRelativePath;
 /// The debug symbol to operation path key name.
 @property (class, copy, readonly, nonatomic, nullable) NSString* debugSymbolToOperationPathKeyName;
+/// The debug symbol to handles key name.
+@property (class, copy, readonly, nonatomic, nullable) NSString* debugSymbolToHandlesKeyName;
 
 @end
 

diff --git a/backends/apple/coreml/runtime/delegate/ETCoreMLStrings.mm b/backends/apple/coreml/runtime/delegate/ETCoreMLStrings.mm
@@ -95,6 +95,11 @@ + (NSString *)debugSymbolToOperationPathKeyName {
     return ETCoreMLDebugSymbolToOperationPathKeyName;
 }
 
++ (NSString *)debugSymbolToHandlesKeyName {
+    static NSString * const ETCoreMLDebugSymbolToHandlesKeyName = @"debugSymbolToHandles";
+    return ETCoreMLDebugSymbolToHandlesKeyName;
+}
+
 + (nullable NSString *)assetsDirectoryPath {
     static dispatch_once_t onceToken;
     static NSString *result = nil;

diff --git a/backends/apple/coreml/runtime/delegate/coreml_backend_delegate.mm b/backends/apple/coreml/runtime/delegate/coreml_backend_delegate.mm
@@ -124,7 +124,8 @@ ModelLoggingOptions get_logging_options(BackendExecutionContext& context) {
     auto event_tracer = context.event_tracer();
     if (event_tracer) {
         options.log_profiling_info = true;
-        options.log_intermediate_tensors = event_tracer->intermediate_outputs_logging_status();
+        auto debug_level = event_tracer->event_tracer_debug_level();
+        options.log_intermediate_tensors = (debug_level >= EventTracerDebugLogLevel::kIntermediateOutputs);
     }
 
     return options;

diff --git a/backends/apple/coreml/runtime/delegate/model_event_logger.h b/backends/apple/coreml/runtime/delegate/model_event_logger.h
@@ -34,8 +34,8 @@ class ModelEventLogger {
     ///
     /// @param op_path_to_value_map A dictionary with the operation path as the key and the operation's value as the
     /// value.
-    /// @param op_path_to_debug_symbol_name_map A dictionary with the operation path as the key and the symbol name as
-    /// the value. The symbol name is the delegate handle.
+    /// @param op_path_to_debug_symbol_name_map A dictionary with the operation path as the key and the debug symbol
+    /// name as the value.
     virtual void log_intermediate_tensors(
         NSDictionary<ETCoreMLModelStructurePath*, MLMultiArray*>* op_path_to_value_map,
         NSDictionary<ETCoreMLModelStructurePath*, NSString*>* op_path_to_debug_symbol_name_map) const noexcept = 0;

diff --git a/backends/apple/coreml/runtime/sdk/ETCoreMLModelAnalyzer.h b/backends/apple/coreml/runtime/sdk/ETCoreMLModelAnalyzer.h
@@ -7,14 +7,15 @@
 
 #import <CoreML/CoreML.h>
 
-#import <ETCoreMLModelExecutor.h>
+#import "ETCoreMLModelExecutor.h"
 
 namespace executorchcoreml {
 struct ModelMetadata;
 }
 
 @class ETCoreMLAsset;
 @class ETCoreMLAssetManager;
+@class ETCoreMLModelDebugInfo;
 @class ETCoreMLModelStructurePath;
 @protocol ETCoreMLModelEventLogger;
 
@@ -32,16 +33,15 @@ __attribute__((objc_subclassing_restricted))
 ///
 /// @param compiledModelAsset The compiled model asset (mlmodelc).
 /// @param modelAsset The model asset (mlpackage).
+/// @param modelDebugInfo The model debug info.
 /// @param metadata The model metadata.
-/// @param operationPathToDebugSymbolMap The operation path to debug symbol map.
 /// @param configuration The model configuration.
 /// @param assetManager The asset manager used to manage storage of compiled models.
 /// @param error   On failure, error is filled with the failure information.
 - (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset*)compiledModelAsset
                                          modelAsset:(nullable ETCoreMLAsset*)modelAsset
+                                     modelDebugInfo:(nullable ETCoreMLModelDebugInfo*)modelDebugInfo
                                            metadata:(const executorchcoreml::ModelMetadata&)metadata
-                      operationPathToDebugSymbolMap:
-                          (nullable NSDictionary<ETCoreMLModelStructurePath*, NSString*>*)operationPathToDebugSymbolMap
                                       configuration:(MLModelConfiguration*)configuration
                                        assetManager:(ETCoreMLAssetManager*)assetManager
                                               error:(NSError* __autoreleasing*)error NS_DESIGNATED_INITIALIZER;

diff --git a/backends/apple/coreml/runtime/sdk/ETCoreMLModelAnalyzer.mm b/backends/apple/coreml/runtime/sdk/ETCoreMLModelAnalyzer.mm
@@ -5,22 +5,24 @@
 //
 // Please refer to the license found in the LICENSE file in the root directory of the source tree.
 
-#import <ETCoreMLAsset.h>
-#import <ETCoreMLAssetManager.h>
-#import <ETCoreMLDefaultModelExecutor.h>
-#import <ETCoreMLLogging.h>
-#import <ETCoreMLModel.h>
-#import <ETCoreMLModelAnalyzer.h>
-#import <ETCoreMLModelLoader.h>
-#import <ETCoreMLModelStructurePath.h>
-#import <ETCoreMLModelDebugger.h>
-#import <ETCoreMLModelProfiler.h>
-#import <ETCoreMLStrings.h>
-#import <model_logging_options.h>
-#import <model_event_logger.h>
-#import <model_metadata.h>
-#import <model_package_info.h>
-#import <objc_safe_cast.h>
+#import "ETCoreMLModelAnalyzer.h"
+
+#import "ETCoreMLAsset.h"
+#import "ETCoreMLAssetManager.h"
+#import "ETCoreMLDefaultModelExecutor.h"
+#import "ETCoreMLLogging.h"
+#import "ETCoreMLModel.h"
+#import "ETCoreMLModelLoader.h"
+#import "ETCoreMLModelStructurePath.h"
+#import "ETCoreMLModelDebugInfo.h"
+#import "ETCoreMLModelDebugger.h"
+#import "ETCoreMLModelProfiler.h"
+#import "ETCoreMLStrings.h"
+#import "model_logging_options.h"
+#import "model_event_logger.h"
+#import "model_metadata.h"
+#import "model_package_info.h"
+#import "objc_safe_cast.h"
 
 namespace {
 using namespace executorchcoreml;
@@ -34,7 +36,7 @@ @interface ETCoreMLModelAnalyzer ()
 @property (strong, nonatomic, nullable) ETCoreMLModelProfiler *profiler;
 @property (strong, nonatomic, nullable) ETCoreMLModelDebugger *debugger;
 @property (strong, nonatomic, nullable) id<ETCoreMLModelExecutor> executor;
-@property (readonly, copy, nonatomic, nullable) NSDictionary<ETCoreMLModelStructurePath *, NSString *> *operationPathToDebugSymbolMap;
+@property (readonly, copy, nonatomic, nullable) ETCoreMLModelDebugInfo *modelDebugInfo;
 @property (readonly, strong, nonatomic) MLModelConfiguration *configuration;
 
 @end
@@ -43,8 +45,8 @@ @implementation ETCoreMLModelAnalyzer
 
 - (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset *)compiledModelAsset
                                          modelAsset:(nullable ETCoreMLAsset *)modelAsset
+                                     modelDebugInfo:(nullable ETCoreMLModelDebugInfo *)modelDebugInfo
                                            metadata:(const executorchcoreml::ModelMetadata&)metadata
-                      operationPathToDebugSymbolMap:(nullable NSDictionary<ETCoreMLModelStructurePath *, NSString *> *)operationPathToDebugSymbolMap
                                       configuration:(MLModelConfiguration *)configuration
                                        assetManager:(ETCoreMLAssetManager *)assetManager
                                               error:(NSError * __autoreleasing *)error {
@@ -72,9 +74,9 @@ - (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset *)compiledMod
     if (self) {
         _model = model;
         _modelAsset = modelAsset;
+        _modelDebugInfo = modelDebugInfo;
         _assetManager = assetManager;
         _configuration = configuration;
-        _operationPathToDebugSymbolMap = operationPathToDebugSymbolMap;
         _executor = [[ETCoreMLDefaultModelExecutor alloc] initWithModel:model];
     }
 
@@ -113,7 +115,7 @@ - (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset *)compiledMod
         return nil;
     }
 
-    eventLogger->log_profiling_infos(profilingInfos, self.operationPathToDebugSymbolMap);
+    eventLogger->log_profiling_infos(profilingInfos, self.modelDebugInfo.pathToDebugSymbolMap);
     return modelOutputs;
 }
 
@@ -131,6 +133,7 @@ - (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset *)compiledMod
 
     if (!self.debugger) {
         self.debugger = [[ETCoreMLModelDebugger alloc] initWithModelAsset:self.modelAsset
+                                                           modelDebugInfo:self.modelDebugInfo
                                                               outputNames:self.model.orderedOutputNames
                                                             configuration:self.configuration
                                                              assetManager:self.assetManager
@@ -143,6 +146,7 @@ - (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset *)compiledMod
 
     NSArray<MLMultiArray *> *modelOutputs = nil;
     NSArray<ETCoreMLModelStructurePath *> *operationPaths = self.debugger.operationPaths;
+    NSDictionary<ETCoreMLModelStructurePath *, NSString *> *operationPathToDebugSymbolMap = self.debugger.operationPathToDebugSymbolMap;
     NSInteger n = operationPaths.count/MAX_MODEL_OUTPUTS_COUNT + (operationPaths.count % MAX_MODEL_OUTPUTS_COUNT == 0 ? 0 : 1);
     for (NSInteger i = 0; i < n; i++) {
         @autoreleasepool {
@@ -157,7 +161,7 @@ - (nullable instancetype)initWithCompiledModelAsset:(ETCoreMLAsset *)compiledMod
             }
 
             if (outputs.count > 0) {
-                eventLogger->log_intermediate_tensors(outputs, self.operationPathToDebugSymbolMap);
+                eventLogger->log_intermediate_tensors(outputs, operationPathToDebugSymbolMap);
             }
         }
     }

diff --git a/backends/apple/coreml/runtime/sdk/ETCoreMLModelDebugInfo.h b/backends/apple/coreml/runtime/sdk/ETCoreMLModelDebugInfo.h
@@ -0,0 +1,47 @@
+//
+// ETCoreMLModelDebugInfo.h
+//
+// Copyright © 2024 Apple Inc. All rights reserved.
+//
+// Please refer to the license found in the LICENSE file in the root directory of the source tree.
+
+#import <Foundation/Foundation.h>
+
+@class ETCoreMLModelStructurePath;
+
+NS_ASSUME_NONNULL_BEGIN
+
+/// A class representing the profiling info of an operation.
+__attribute__((objc_subclassing_restricted))
+@interface ETCoreMLModelDebugInfo : NSObject
+
+- (instancetype)init NS_UNAVAILABLE;
+
++ (instancetype)new NS_UNAVAILABLE;
+
+
+/// Constructs an `ETCoreMLModelDebugInfo` instance.
+///
+/// @param pathToDebugSymbolMap Operation path to debug symbol map.
+/// @param pathToDebugHandlesMap Operation path to debug handles map.
+- (instancetype)initWithPathToDebugSymbolMap:(NSDictionary<ETCoreMLModelStructurePath*, NSString*>*)pathToDebugSymbolMap
+                       pathToDebugHandlesMap:
+                           (NSDictionary<ETCoreMLModelStructurePath*, NSArray<NSString*>*>*)pathToDebugHandlesMap
+    NS_DESIGNATED_INITIALIZER;
+
+/// Constructs an `ETCoreMLModelDebugInfo` instance.
+///
+/// @param data The json data.
+/// @param error   On failure, error is filled with the failure information.
++ (nullable instancetype)modelDebugInfoFromData:(NSData*)data error:(NSError* __autoreleasing*)error;
+
+/// Operation path to debug symbol map.
+@property (readonly, strong, nonatomic) NSDictionary<ETCoreMLModelStructurePath*, NSString*>* pathToDebugSymbolMap;
+
+/// Operation path to debug handles map.
+@property (readonly, strong, nonatomic)
+    NSDictionary<ETCoreMLModelStructurePath*, NSArray<NSString*>*>* pathToDebugHandlesMap;
+
+@end
+
+NS_ASSUME_NONNULL_END