From a631ed77c012dbf0a86f3147e6fbe1c236f08c02 Mon Sep 17 00:00:00 2001 From: Yulong Wang <7679871+fs-eire@users.noreply.github.com> Date: Fri, 24 Feb 2023 15:50:15 -0800 Subject: [PATCH] [js/web] support flag 'optimizedModelFilePath' in session options (#14355) ### Description * Support flag 'optimizedModelFilePath' in session options. In Node.js, the model will be saved into filesystem just like its behaviour on native platforms. In browser, the new model is not saved to filesystem. the file path is ignored. Instead, a new pop-up window will be launched in browser and user can 'save' the file as onnx model. * Add corresponding commandline args for the following session option flags: - optimizedModelFilePath - graphOptimizationLevel --- cmake/CMakeLists.txt | 5 ++++ cmake/onnxruntime_webassembly.cmake | 4 +-- js/common/lib/inference-session.ts | 8 ++++++ js/web/lib/wasm/binding/ort-wasm.d.ts | 2 +- js/web/lib/wasm/session-options.ts | 7 +++++- js/web/script/test-runner-cli-args.ts | 29 +++++++++++++++++++++ js/web/script/test-runner-cli.ts | 2 ++ js/web/test/test-main.ts | 3 ++- js/web/test/test-runner.ts | 12 ++++++--- js/web/test/test-types.ts | 1 + onnxruntime/core/graph/model.cc | 36 +++++++++++++++++++++++++++ onnxruntime/wasm/api.cc | 7 +++++- onnxruntime/wasm/api.h | 4 ++- 13 files changed, 109 insertions(+), 11 deletions(-) diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 7a475be35e..20ea1d6095 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -160,6 +160,7 @@ option(onnxruntime_ENABLE_WEBASSEMBLY_API_EXCEPTION_CATCHING "Enable this option option(onnxruntime_ENABLE_WEBASSEMBLY_EXCEPTION_THROWING "Enable this option to turn on exception throwing even if the build disabled exceptions support" OFF) option(onnxruntime_ENABLE_WEBASSEMBLY_DEBUG_INFO "Enable this option to turn on DWARF format debug info" OFF) option(onnxruntime_ENABLE_WEBASSEMBLY_PROFILING "Enable this option to turn on WebAssembly profiling and preserve function names" OFF) +option(onnxruntime_ENABLE_WEBASSEMBLY_OUTPUT_OPTIMIZED_MODEL "Enable this option to allow WebAssembly to output optimized model" OFF) # Enable bitcode for iOS option(onnxruntime_ENABLE_BITCODE "Enable bitcode for iOS only" OFF) @@ -1546,6 +1547,10 @@ endif() if (onnxruntime_BUILD_WEBASSEMBLY) message(STATUS "WebAssembly Build is enabled") list(APPEND ONNXRUNTIME_CMAKE_FILES onnxruntime_webassembly) + + if (onnxruntime_ENABLE_WEBASSEMBLY_OUTPUT_OPTIMIZED_MODEL) + add_compile_definitions(ORT_ENABLE_WEBASSEMBLY_OUTPUT_OPTIMIZED_MODEL) + endif() endif() if(onnxruntime_BUILD_KERNEL_EXPLORER) diff --git a/cmake/onnxruntime_webassembly.cmake b/cmake/onnxruntime_webassembly.cmake index 25761aa841..82154efb45 100644 --- a/cmake/onnxruntime_webassembly.cmake +++ b/cmake/onnxruntime_webassembly.cmake @@ -203,13 +203,13 @@ else() -s \"EXPORTED_FUNCTIONS=_malloc,_free\" \ -s MAXIMUM_MEMORY=4294967296 \ -s WASM=1 \ - -s NO_EXIT_RUNTIME=0 \ + -s EXIT_RUNTIME=0 \ -s ALLOW_MEMORY_GROWTH=1 \ -s MODULARIZE=1 \ -s EXPORT_ALL=0 \ -s LLD_REPORT_UNDEFINED \ -s VERBOSE=0 \ - -s NO_FILESYSTEM=1 \ + -s FILESYSTEM=0 \ ${WASM_API_EXCEPTION_CATCHING} \ --no-entry") diff --git a/js/common/lib/inference-session.ts b/js/common/lib/inference-session.ts index 1f2f855a3e..638cb90f36 100644 --- a/js/common/lib/inference-session.ts +++ b/js/common/lib/inference-session.ts @@ -94,6 +94,14 @@ export declare namespace InferenceSession { */ executionMode?: 'sequential'|'parallel'; + /** + * Optimized model file path. + * + * If this setting is specified, the optimized model will be dumped. In browser, a blob will be created + * with a pop-up window. + */ + optimizedModelFilePath?: string; + /** * Wether enable profiling. * diff --git a/js/web/lib/wasm/binding/ort-wasm.d.ts b/js/web/lib/wasm/binding/ort-wasm.d.ts index fd82a83bd7..e7bf279f0e 100644 --- a/js/web/lib/wasm/binding/ort-wasm.d.ts +++ b/js/web/lib/wasm/binding/ort-wasm.d.ts @@ -36,7 +36,7 @@ export interface OrtWasmModule extends EmscriptenModule { _OrtCreateSessionOptions( graphOptimizationLevel: number, enableCpuMemArena: boolean, enableMemPattern: boolean, executionMode: number, enableProfiling: boolean, profileFilePrefix: number, logId: number, logSeverityLevel: number, - logVerbosityLevel: number): number; + logVerbosityLevel: number, optimizedModelFilePath: number): number; _OrtAppendExecutionProvider(sessionOptionsHandle: number, name: number): number; _OrtAddSessionConfigEntry(sessionOptionsHandle: number, configKey: number, configValue: number): number; _OrtReleaseSessionOptions(sessionOptionsHandle: number): void; diff --git a/js/web/lib/wasm/session-options.ts b/js/web/lib/wasm/session-options.ts index 6d4d8eeb34..b63ec78a85 100644 --- a/js/web/lib/wasm/session-options.ts +++ b/js/web/lib/wasm/session-options.ts @@ -122,10 +122,15 @@ export const setSessionOptions = (options?: InferenceSession.SessionOptions): [n sessionOptions.enableProfiling = false; } + let optimizedModelFilePathOffset = 0; + if (typeof options?.optimizedModelFilePath === 'string') { + optimizedModelFilePathOffset = allocWasmString(options.optimizedModelFilePath, allocs); + } + sessionOptionsHandle = wasm._OrtCreateSessionOptions( graphOptimizationLevel, !!sessionOptions.enableCpuMemArena!, !!sessionOptions.enableMemPattern!, executionMode, !!sessionOptions.enableProfiling!, 0, logIdDataOffset, sessionOptions.logSeverityLevel!, - sessionOptions.logVerbosityLevel!); + sessionOptions.logVerbosityLevel!, optimizedModelFilePathOffset); if (sessionOptionsHandle === 0) { throw new Error('Can\'t create session options'); } diff --git a/js/web/script/test-runner-cli-args.ts b/js/web/script/test-runner-cli-args.ts index c008933bc4..8d02f81e11 100644 --- a/js/web/script/test-runner-cli-args.ts +++ b/js/web/script/test-runner-cli-args.ts @@ -50,6 +50,10 @@ Options: This flag can be used with a number as value, specifying the total count of test cases to run. The test cases may be used multiple times. Default value is 10. -c, --file-cache Enable file cache. +*** Session Options *** + -u=<...>, --optimized-model-file-path=<...> Specify whether to dump the optimized model. + -o=<...>, --graph-optimization-level=<...> Specify graph optimization level. + Default is 'all'. Valid values are 'disabled', 'basic', 'extended', 'all'. *** Logging Options *** --log-verbose=<...> Set log level to verbose @@ -149,6 +153,16 @@ export interface TestRunnerCliArgs { */ times?: number; + /** + * whether to dump the optimized model + */ + optimizedModelFilePath?: string; + + /** + * Specify graph optimization level + */ + graphOptimizationLevel: 'disabled'|'basic'|'extended'|'all'; + cpuOptions?: InferenceSession.CpuExecutionProviderOption; cudaOptions?: InferenceSession.CudaExecutionProviderOption; cudaFlags?: Record; @@ -378,6 +392,19 @@ export function parseTestRunnerCliArgs(cmdlineArgs: string[]): TestRunnerCliArgs logConfig.push({category: 'TestRunner.Perf', config: {minimalSeverity: 'verbose'}}); } + // Option: -u, --optimized-model-file-path + const optimizedModelFilePath = args['optimized-model-file-path'] || args.u || undefined; + if (typeof optimizedModelFilePath !== 'undefined' && typeof optimizedModelFilePath !== 'string') { + throw new Error('Flag "optimized-model-file-path" need to be either empty or a valid file path.'); + } + + // Option: -o, --graph-optimization-level + const graphOptimizationLevel = args['graph-optimization-level'] || args.o || 'all'; + if (typeof graphOptimizationLevel !== 'string' || + ['disabled', 'basic', 'extended', 'all'].indexOf(graphOptimizationLevel) === -1) { + throw new Error(`graph optimization level is invalid: ${graphOptimizationLevel}`); + } + // Option: -c, --file-cache const fileCache = parseBooleanArg(args['file-cache'] || args.c, false); @@ -405,6 +432,8 @@ export function parseTestRunnerCliArgs(cmdlineArgs: string[]): TestRunnerCliArgs logConfig, profile, times: perf ? times : undefined, + optimizedModelFilePath, + graphOptimizationLevel: graphOptimizationLevel as TestRunnerCliArgs['graphOptimizationLevel'], fileCache, cpuOptions, webglOptions, diff --git a/js/web/script/test-runner-cli.ts b/js/web/script/test-runner-cli.ts index f4d2d0a90b..6b3f980384 100644 --- a/js/web/script/test-runner-cli.ts +++ b/js/web/script/test-runner-cli.ts @@ -146,6 +146,8 @@ run({ log: args.logConfig, profile: args.profile, options: { + sessionOptions: + {graphOptimizationLevel: args.graphOptimizationLevel, optimizedModelFilePath: args.optimizedModelFilePath}, debug: args.debug, cpuOptions: args.cpuOptions, webglOptions: args.webglOptions, diff --git a/js/web/test/test-main.ts b/js/web/test/test-main.ts index 64cef55bd1..2610cbe1d8 100644 --- a/js/web/test/test-main.ts +++ b/js/web/test/test-main.ts @@ -103,7 +103,8 @@ for (const group of ORT_WEB_TEST_CONFIG.model) { let context: ModelTestContext; before('prepare session', async () => { - context = await ModelTestContext.create(test, ORT_WEB_TEST_CONFIG.profile); + context = await ModelTestContext.create( + test, ORT_WEB_TEST_CONFIG.profile, ORT_WEB_TEST_CONFIG.options.sessionOptions); }); after('release session', () => { diff --git a/js/web/test/test-runner.ts b/js/web/test/test-runner.ts index 4bdea197bb..3f1fa0e0f8 100644 --- a/js/web/test/test-runner.ts +++ b/js/web/test/test-runner.ts @@ -114,7 +114,7 @@ async function loadTensors( } async function initializeSession( - modelFilePath: string, backendHint: string, profile: boolean, + modelFilePath: string, backendHint: string, profile: boolean, sessionOptions: ort.InferenceSession.SessionOptions, fileCache?: FileCacheBuffer): Promise { const preloadModelData: Uint8Array|undefined = fileCache && fileCache[modelFilePath] ? fileCache[modelFilePath] : undefined; @@ -124,7 +124,8 @@ async function initializeSession( preloadModelData ? ` [preloaded(${preloadModelData.byteLength})]` : ''}`); const profilerConfig = profile ? {maxNumberEvents: 65536} : undefined; - const sessionConfig = {executionProviders: [backendHint], profiler: profilerConfig, enableProfiling: profile}; + const sessionConfig = + {...sessionOptions, executionProviders: [backendHint], profiler: profilerConfig, enableProfiling: profile}; let session: ort.InferenceSession; try { @@ -197,7 +198,9 @@ export class ModelTestContext { /** * create a ModelTestContext object that used in every test cases in the given ModelTest. */ - static async create(modelTest: Test.ModelTest, profile: boolean): Promise { + static async create( + modelTest: Test.ModelTest, profile: boolean, + sessionOptions?: ort.InferenceSession.SessionOptions): Promise { if (this.initializing) { throw new Error('cannot create a ModelTestContext object when the previous creation is not done'); } @@ -206,7 +209,8 @@ export class ModelTestContext { this.initializing = true; const initStart = now(); - const session = await initializeSession(modelTest.modelUrl, modelTest.backend!, profile, this.cache); + const session = + await initializeSession(modelTest.modelUrl, modelTest.backend!, profile, sessionOptions || {}, this.cache); const initEnd = now(); for (const testCase of modelTest.cases) { diff --git a/js/web/test/test-types.ts b/js/web/test/test-types.ts index a7ab9d7025..966b1e704a 100644 --- a/js/web/test/test-types.ts +++ b/js/web/test/test-types.ts @@ -104,6 +104,7 @@ export declare namespace Test { */ export interface Options { debug?: boolean; + sessionOptions?: InferenceSession.SessionOptions; cpuOptions?: InferenceSession.CpuExecutionProviderOption; cpuFlags?: Record; cudaOptions?: InferenceSession.CudaExecutionProviderOption; diff --git a/onnxruntime/core/graph/model.cc b/onnxruntime/core/graph/model.cc index 8af9f99ed1..feec7d5407 100644 --- a/onnxruntime/core/graph/model.cc +++ b/onnxruntime/core/graph/model.cc @@ -29,6 +29,10 @@ #include "core/graph/function_utils.h" #endif +#if defined(__wasm__) +#include +#endif + using namespace ONNX_NAMESPACE; using namespace onnxruntime; using namespace onnxruntime::common; @@ -500,6 +504,37 @@ static Status LoadModel(const T& file_path, std::shared_ptr& p_model, template static Status SaveModel(Model& model, const T& file_path) { +#if defined(__wasm__) && defined(ORT_ENABLE_WEBASSEMBLY_OUTPUT_OPTIMIZED_MODEL) + ORT_RETURN_IF_ERROR(model.MainGraph().Resolve()); + auto model_proto = model.ToProto(); + auto buffer_size = model_proto.ByteSizeLong(); + void* buffer = malloc(buffer_size); + model_proto.SerializeToArray(buffer, buffer_size); + + EM_ASM(({ + const buffer = $0; + const buffer_size = $1; + const file_path = UTF8ToString($2); + const bytes = new Uint8Array(buffer_size); + bytes.set(HEAPU8.subarray(buffer, buffer + buffer_size)); + if (typeof process == 'object' && typeof process.versions == 'object' && typeof process.versions.node == 'string') { + // Node.js + require('fs').writeFileSync(file_path, bytes); + } else { + // Browser + const file = new File([bytes], file_path, {type: "application/octet-stream" }); + const url = URL.createObjectURL(file); + window.open(url, '_blank'); + } + }), + reinterpret_cast(buffer), + static_cast(buffer_size), + reinterpret_cast(file_path.c_str())); + + free(buffer); + return Status::OK(); + +#else int fd; Status status = Env::Default().FileOpenWr(file_path, fd); ORT_RETURN_IF_ERROR(status); @@ -518,6 +553,7 @@ static Status SaveModel(Model& model, const T& file_path) { return status; } return Env::Default().FileClose(fd); +#endif } #ifdef _WIN32 diff --git a/onnxruntime/wasm/api.cc b/onnxruntime/wasm/api.cc index fb73826800..62856c1bdc 100644 --- a/onnxruntime/wasm/api.cc +++ b/onnxruntime/wasm/api.cc @@ -68,10 +68,15 @@ OrtSessionOptions* OrtCreateSessionOptions(size_t graph_optimization_level, const char* /*profile_file_prefix*/, const char* log_id, size_t log_severity_level, - size_t log_verbosity_level) { + size_t log_verbosity_level, + const char* optimized_model_filepath) { OrtSessionOptions* session_options = nullptr; RETURN_NULLPTR_IF_ERROR(CreateSessionOptions, &session_options); + if (optimized_model_filepath) { + RETURN_NULLPTR_IF_ERROR(SetOptimizedModelFilePath, session_options, optimized_model_filepath); + } + // assume that a graph optimization level is checked and properly set at JavaScript RETURN_NULLPTR_IF_ERROR(SetSessionGraphOptimizationLevel, session_options, diff --git a/onnxruntime/wasm/api.h b/onnxruntime/wasm/api.h index d3435f2958..80466ecd87 100644 --- a/onnxruntime/wasm/api.h +++ b/onnxruntime/wasm/api.h @@ -46,6 +46,7 @@ int EMSCRIPTEN_KEEPALIVE OrtInit(int num_threads, int logging_level); * @param log_id logger id for session output * @param log_severity_level verbose, info, warning, error or fatal * @param log_verbosity_level vlog level + * @param optimized_model_filepath filepath of the optimized model to dump. * @returns a pointer to a session option handle and must be freed by calling OrtReleaseSessionOptions(). */ ort_session_options_handle_t EMSCRIPTEN_KEEPALIVE OrtCreateSessionOptions(size_t graph_optimization_level, @@ -56,7 +57,8 @@ ort_session_options_handle_t EMSCRIPTEN_KEEPALIVE OrtCreateSessionOptions(size_t const char* profile_file_prefix, const char* log_id, size_t log_severity_level, - size_t log_verbosity_level); + size_t log_verbosity_level, + const char* optimized_model_filepath); /** * append an execution provider for a session.