diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index b81df151562dc..a841b17a2a571 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -180,7 +180,6 @@ option(onnxruntime_ENABLE_TRAINING_OPS "Include training operators but no traini option(onnxruntime_ENABLE_TRAINING_E2E_TESTS "Enable training end-to-end tests." OFF) option(onnxruntime_ENABLE_CPU_FP16_OPS "Build with advanced instruction sets" ON) option(onnxruntime_USE_NCCL "Build with NCCL support" OFF) -option(onnxruntime_USE_MPI "Build with MPI support" OFF) # WebAssembly options option(onnxruntime_BUILD_WEBASSEMBLY_STATIC_LIB "Enable this option to create WebAssembly static library" OFF) @@ -1267,6 +1266,7 @@ function(onnxruntime_set_compile_flags target_name) # Unsupported by Clang 18 yet. list(REMOVE_ITEM ORT_HIP_WARNING_FLAGS -Wno-dangling-reference) + list(REMOVE_ITEM ORT_HIP_WARNING_FLAGS -Wno-interference-size) # float16.h:90:12: error: ‘tmp’ is used uninitialized list(APPEND ORT_HIP_WARNING_FLAGS -Wno-uninitialized) list(APPEND ORT_HIP_WARNING_FLAGS -Wno-deprecated-copy) @@ -1650,29 +1650,6 @@ if (onnxruntime_ENABLE_DLPACK) endif() if (UNIX OR onnxruntime_USE_NCCL) - # MPI is INDEPENDENT of NCCL for now. You can build NCLL without MPI and launch multi-GPU with your own launcher. - if (onnxruntime_USE_MPI) - if (EXISTS "${onnxruntime_MPI_HOME}") - set(MPI_HOME "${onnxruntime_MPI_HOME}") - elseif (EXISTS "/bert_ort/openmpi") - set(MPI_HOME "/bert_ort/openmpi") - endif() - find_package(MPI) - - if (MPI_CXX_FOUND) - message( STATUS "MPI Version: ${MPI_CXX_VERSION}") - message( STATUS "MPI (include: ${MPI_CXX_INCLUDE_DIRS}, library: ${MPI_CXX_LIBRARIES})" ) - mark_as_advanced(MPI_CXX_INCLUDE_DIRS MPI_CXX_LIBRARIES) - list(APPEND onnxruntime_EXTERNAL_LIBRARIES ${MPI_CXX_LIBRARIES} ${MPI_CXX_LINK_FLAGS}) - else () - message( - FATAL_ERROR - "MPI is not found. Please define onnxruntime_MPI_HOME to specify the path of MPI. Otherwise, NCCL will be disabled." - "or you can remove --use_mpi from build args to disable MPI." - ) - endif() - endif() - # Find NCCL if (onnxruntime_USE_NCCL) if (onnxruntime_USE_CUDA) @@ -1742,13 +1719,9 @@ if (UNIX OR onnxruntime_USE_NCCL) endif() else() set(onnxruntime_USE_NCCL OFF) - set(onnxruntime_USE_MPI OFF) - message( WARNING "MPI and NCCL are disabled because build is on Windows or USE_NCCL is set to OFF." ) + message( WARNING "NCCL is disabled because build is on Windows or USE_NCCL is set to OFF." ) endif() -if (onnxruntime_USE_MPI) - add_definitions(-DUSE_MPI=1) -endif() # Default version parts for Microsoft.AI.MachineLearning.dll, onnxruntime.dll, onnxruntime_providers_openvino.dll and onnxruntime_providers_shared.dll in non-ADO pipeline local builds set(VERSION_MAJOR_PART 0 CACHE STRING "First part of numeric file/product version.") diff --git a/cmake/external/onnxruntime_external_deps.cmake b/cmake/external/onnxruntime_external_deps.cmake index 552b8b5329fb0..0e0083af9ed54 100644 --- a/cmake/external/onnxruntime_external_deps.cmake +++ b/cmake/external/onnxruntime_external_deps.cmake @@ -719,6 +719,11 @@ if (onnxruntime_USE_WEBGPU) # - (private) Remove hard-coded CMAKE_OSX_DEPLOYMENT_TARGET in Dawn's CMake files # https://github.com/microsoft/onnxruntime/pull/23729 # + # - (private) Reduce unsafe buffer usage warning in aligned_storage.h + # https://github.com/microsoft/onnxruntime/pull/24308 + # The patch disables the UNSAFE_BUFFER_USAGE warning around the AlignedStorage struct in aligned_storage.h. This is done + # by using TINT_BEGIN_DISABLE_WARNING and TINT_END_DISABLE_WARNING macros, which helps in warnings related to unsafe buffer usage + # usage when compiling the code, making the build process cleaner and faster. # PATCH_COMMAND ${Patch_EXECUTABLE} --binary --ignore-whitespace -p1 < ${PROJECT_SOURCE_DIR}/patches/dawn/dawn.patch EXCLUDE_FROM_ALL diff --git a/cmake/onnxruntime_framework.cmake b/cmake/onnxruntime_framework.cmake index 9c9a25f8ee77e..d248f3652e064 100644 --- a/cmake/onnxruntime_framework.cmake +++ b/cmake/onnxruntime_framework.cmake @@ -77,9 +77,6 @@ if (onnxruntime_ENABLE_TRAINING_OPS) onnxruntime_add_include_to_target(onnxruntime_framework Python::Module dlpack::dlpack) endif() endif() -if (onnxruntime_USE_MPI) - target_include_directories(onnxruntime_framework PUBLIC ${MPI_CXX_INCLUDE_DIRS}) -endif() if (onnxruntime_ENABLE_ATEN) onnxruntime_add_include_to_target(onnxruntime_framework dlpack::dlpack) diff --git a/cmake/onnxruntime_providers_cpu.cmake b/cmake/onnxruntime_providers_cpu.cmake index 3a5131947b917..b9c810e1250a0 100644 --- a/cmake/onnxruntime_providers_cpu.cmake +++ b/cmake/onnxruntime_providers_cpu.cmake @@ -206,7 +206,7 @@ if (onnxruntime_ENABLE_TRAINING) onnxruntime_add_include_to_target(onnxruntime_providers Python::Module) endif() - if (onnxruntime_USE_NCCL OR onnxruntime_USE_MPI) + if (onnxruntime_USE_NCCL) target_include_directories(onnxruntime_providers PUBLIC ${MPI_CXX_INCLUDE_DIRS}) endif() endif() diff --git a/cmake/onnxruntime_providers_cuda.cmake b/cmake/onnxruntime_providers_cuda.cmake index 795f019f6d70b..b688e61f53915 100644 --- a/cmake/onnxruntime_providers_cuda.cmake +++ b/cmake/onnxruntime_providers_cuda.cmake @@ -252,10 +252,6 @@ target_include_directories(${target} PRIVATE ${ORTTRAINING_ROOT} ${MPI_CXX_INCLUDE_DIRS}) endif() - if(onnxruntime_USE_MPI) - target_link_libraries(${target} PRIVATE ${MPI_LIBRARIES} ${MPI_CXX_LINK_FLAGS}) - endif() - if (onnxruntime_USE_NCCL) target_include_directories(${target} PRIVATE ${NCCL_INCLUDE_DIRS}) target_link_libraries(${target} PRIVATE ${NCCL_LIBRARIES}) diff --git a/cmake/onnxruntime_providers_rocm.cmake b/cmake/onnxruntime_providers_rocm.cmake index 0ee5a6e32a434..108b8b46deb27 100644 --- a/cmake/onnxruntime_providers_rocm.cmake +++ b/cmake/onnxruntime_providers_rocm.cmake @@ -158,9 +158,6 @@ if (onnxruntime_ENABLE_TRAINING) target_include_directories(onnxruntime_providers_rocm PRIVATE ${ORTTRAINING_ROOT} ${CMAKE_CURRENT_BINARY_DIR}/amdgpu/orttraining ${MPI_CXX_INCLUDE_DIRS}) - if(onnxruntime_USE_MPI) - target_link_libraries(onnxruntime_providers_rocm PRIVATE ${MPI_LIBRARIES} ${MPI_CXX_LINK_FLAGS}) - endif() # RCCL is enabled by default for ROCM builds #if (onnxruntime_USE_NCCL) diff --git a/cmake/patches/dawn/dawn.patch b/cmake/patches/dawn/dawn.patch index 1ce9c75f750f1..1fe66d2cf917d 100644 --- a/cmake/patches/dawn/dawn.patch +++ b/cmake/patches/dawn/dawn.patch @@ -36,3 +36,24 @@ index 5bfac41dcc..71a153daaa 100644 void WGPUBufferImpl::Destroy() { emwgpuBufferDestroy(this); AbortPendingMap("Buffer was destroyed before mapping was resolved."); +diff --git a/src/tint/utils/memory/aligned_storage.h b/src/tint/utils/memory/aligned_storage.h +index c532c4fc38..19c950af4c 100644 +--- a/src/tint/utils/memory/aligned_storage.h ++++ b/src/tint/utils/memory/aligned_storage.h +@@ -31,6 +31,9 @@ + #include + + #include "src/tint/utils/memory/bitcast.h" ++#include "src/tint/utils/macros/compiler.h" ++ ++TINT_BEGIN_DISABLE_WARNING(UNSAFE_BUFFER_USAGE); + + namespace tint { + +@@ -50,4 +53,6 @@ struct alignas(alignof(T)) AlignedStorage { + + } // namespace tint + ++TINT_END_DISABLE_WARNING(UNSAFE_BUFFER_USAGE); ++ + #endif // SRC_TINT_UTILS_MEMORY_ALIGNED_STORAGE_H_ diff --git a/java/build.gradle b/java/build.gradle index 845121dd17a48..8452daab72872 100644 --- a/java/build.gradle +++ b/java/build.gradle @@ -211,6 +211,7 @@ test { 'USE_TENSORRT', 'USE_QNN', 'USE_XNNPACK', + 'USE_WEBGPU', ]) testLogging { events "passed", "skipped", "failed" diff --git a/java/src/main/java/ai/onnxruntime/OrtSession.java b/java/src/main/java/ai/onnxruntime/OrtSession.java index bd988e2bb7468..f15ad938463a7 100644 --- a/java/src/main/java/ai/onnxruntime/OrtSession.java +++ b/java/src/main/java/ai/onnxruntime/OrtSession.java @@ -1339,6 +1339,18 @@ public void addCoreML(Map providerOptions) throws OrtException { addExecutionProvider(CoreMLProviderName, providerOptions); } + /** + * Adds WebGPU as an execution backend. + * + * @param providerOptions Configuration options for the WebGPU backend. Refer to the WebGPU + * execution provider's documentation. + * @throws OrtException If there was an error in native code. + */ + public void addWebGPU(Map providerOptions) throws OrtException { + String webGpuProviderName = "WebGPU"; + addExecutionProvider(webGpuProviderName, providerOptions); + } + private native void setExecutionMode(long apiHandle, long nativeHandle, int mode) throws OrtException; diff --git a/java/src/test/android/app/src/androidTest/java/ai/onnxruntime/example/javavalidator/SimpleTest.kt b/java/src/test/android/app/src/androidTest/java/ai/onnxruntime/example/javavalidator/SimpleTest.kt index 0af8b6f1a8ab1..6e8957d05c49b 100644 --- a/java/src/test/android/app/src/androidTest/java/ai/onnxruntime/example/javavalidator/SimpleTest.kt +++ b/java/src/test/android/app/src/androidTest/java/ai/onnxruntime/example/javavalidator/SimpleTest.kt @@ -52,6 +52,11 @@ class SimpleTest { runSigmoidModelTestImpl(1, OrtProvider.QNN) } + @Test + fun runSigmoidModelTestWEBGPU() { + runSigmoidModelTestImpl(1, OrtProvider.WEBGPU) + } + @Throws(IOException::class) private fun readModel(fileName: String): ByteArray { return InstrumentationRegistry.getInstrumentation().context.assets.open(fileName) @@ -90,6 +95,16 @@ class SimpleTest { } } + OrtProvider.WEBGPU -> { + if (OrtEnvironment.getAvailableProviders().contains(OrtProvider.WEBGPU)) { + val providerOptions = Collections.emptyMap() + opts.addWebGPU(providerOptions) + } else { + Log.println(Log.INFO, TAG, "NO WEBGPU EP available, skip the test") + return + } + } + OrtProvider.CPU -> { // No additional configuration is needed for CPU } diff --git a/java/src/test/java/ai/onnxruntime/InferenceTest.java b/java/src/test/java/ai/onnxruntime/InferenceTest.java index 83cdc32dbaeb5..41c4f5e02bd69 100644 --- a/java/src/test/java/ai/onnxruntime/InferenceTest.java +++ b/java/src/test/java/ai/onnxruntime/InferenceTest.java @@ -749,6 +749,12 @@ public void testQNN() throws OrtException { runProvider(OrtProvider.QNN); } + @Test + @EnabledIfSystemProperty(named = "USE_WEBGPU", matches = "1") + public void testWEBGPU() throws OrtException { + runProvider(OrtProvider.WEBGPU); + } + private void runProvider(OrtProvider provider) throws OrtException { EnumSet providers = OrtEnvironment.getAvailableProviders(); assertTrue(providers.size() > 1); diff --git a/js/common/lib/env.ts b/js/common/lib/env.ts index d6d9f7fa48790..112c8a1f78851 100644 --- a/js/common/lib/env.ts +++ b/js/common/lib/env.ts @@ -41,16 +41,22 @@ export declare namespace Env { numThreads?: number; /** - * set or get a boolean value indicating whether to enable SIMD. If set to false, SIMD will be forcely disabled. + * set a value indicating whether to enable SIMD. + * + * ONNX Runtime will perform feature detection based on the value of this property. Specifically, when the value is + * set to: + * - `undefined`, `true` or `"fixed"`: will check availability of Fixed-width SIMD. + * - `"relaxed"`: will check availability of Relaxed SIMD. + * - `false`: will not perform SIMD feature checking. + * + * Setting this property does not make ONNX Runtime to switch to the corresponding runtime automatically. User need + * to set `wasmPaths` or `wasmBinary` property to load the corresponding runtime. * * This setting is available only when WebAssembly SIMD feature is available in current context. * * @defaultValue `true` - * - * @deprecated This property is deprecated. Since SIMD is supported by all major JavaScript engines, non-SIMD - * build is no longer provided. This property will be removed in future release. */ - simd?: boolean; + simd?: boolean | 'fixed' | 'relaxed'; /** * set or get a boolean value indicating whether to enable trace. diff --git a/js/web/docs/webnn-operators.md b/js/web/docs/webnn-operators.md index a6a2ecdf6f467..33906e291d423 100644 --- a/js/web/docs/webnn-operators.md +++ b/js/web/docs/webnn-operators.md @@ -49,6 +49,7 @@ platforms. Check the [WebNN status](https://webmachinelearning.github.io/webnn-s | GlobalLpPool| ai.onnx(7+) | l2Pool2d | Only supports 4-D input, 'p' value is 2 | | Greater | ai.onnx(7-8, 9-12, 13+) | greater | | | GreaterOrEqual | ai.onnx(12-15, 16+) | greaterOrEqual | | +| GroupQueryAttention | com.microsoft(1+) | add, cast, concat, constant, cumulativeSum, div, expand, lesser, matmul, reshape, scatterND, softmax, transpose, where | Only supports input total_sequence_length is constant and past_sequence_length of past kv equals to present_sequence_length of present kv. Does not support cos_cache and sin_cache inputs | | GRU | ai.onnx(7-13, 14-21, 22+) | gru | Only supports 'layout' == 0. 'clip' is not supported. The activation functions in 'activations' must be one of 'Relu', 'Tanh', 'Sigmoid'. Forward and backward activations must be the same if bidirectional. 'sequence_lens' if present should be constant with values equal to the first dimension length of input 'X' | | HardSigmoid | ai.onnx(7+) | hardSigmoid | | | HardSwish | ai.onnx(14+) | hardSwish | | @@ -63,6 +64,7 @@ platforms. Check the [WebNN status](https://webmachinelearning.github.io/webnn-s | LRN | ai.onnx(7-12, 13+) | pad, averagePool2d, transpose, add, mul, pow, div | | | LSTM | ai.onnx(7-13, 14-21, 22+) | lstm | Only supports 'layout' == 0, 'input_forget' == 0. 'clip' is not supported. The activation functions in 'activations' must be one of 'Relu', 'Tanh', 'Sigmoid'. Forward and backward activations must be the same if bidirectional. 'sequence_lens' if present should be constant with values equal to the first dimension length of input 'X' | | MatMul | ai.onnx(7-8, 9-12, 13+) | matmul | | +| MatMulNBits | com.microsoft(1+) | add, dequantizeLinear, matmul, reshape, transpose | Inputs 'B' and 'zero_points' (if present) should be constants, input 'g_idx' is not supported, only bits=4 is supported | | Max | ai.onnx(7, 8-11, 12, 13+) | max | | | MaxPool | ai.onnx(7, 8-9, 10, 11, 12+) | maxPool2d | Only supports 4-D input, 2-D 'kernel_shape', 'storage_order' != 1, one output | | Min | ai.onnx(7, 8-11, 12, 13+) | min | | diff --git a/js/web/lib/backend-wasm.ts b/js/web/lib/backend-wasm.ts index 5e5f4804aed92..fe8bfaa188ea9 100644 --- a/js/web/lib/backend-wasm.ts +++ b/js/web/lib/backend-wasm.ts @@ -17,12 +17,13 @@ export const initializeFlags = (): void => { env.wasm.initTimeout = 0; } - if (env.wasm.simd === false) { + const simd = env.wasm.simd; + if (typeof simd !== 'boolean' && simd !== undefined && simd !== 'fixed' && simd !== 'relaxed') { // eslint-disable-next-line no-console console.warn( - 'Deprecated property "env.wasm.simd" is set to false. ' + - 'non-SIMD build is no longer provided, and this setting will be ignored.', + `Property "env.wasm.simd" is set to unknown value "${simd}". Reset it to \`false\` and ignore SIMD feature checking.`, ); + env.wasm.simd = false; } if (typeof env.wasm.proxy !== 'boolean') { diff --git a/js/web/lib/wasm/wasm-factory.ts b/js/web/lib/wasm/wasm-factory.ts index 0f49d25040409..23eb2f0978feb 100644 --- a/js/web/lib/wasm/wasm-factory.ts +++ b/js/web/lib/wasm/wasm-factory.ts @@ -64,6 +64,34 @@ const isSimdSupported = (): boolean => { } }; +const isRelaxedSimdSupported = (): boolean => { + try { + // Test for WebAssembly Relaxed SIMD capability (for both browsers and Node.js) + // This typed array is a WebAssembly program containing Relaxed SIMD instructions. + + // The binary data is generated from the following code by wat2wasm: + // (module + // (func (result v128) + // i32.const 1 + // i8x16.splat + // i32.const 2 + // i8x16.splat + // i32.const 3 + // i8x16.splat + // i32x4.relaxed_dot_i8x16_i7x16_add_s + // ) + // ) + return WebAssembly.validate( + new Uint8Array([ + 0, 97, 115, 109, 1, 0, 0, 0, 1, 5, 1, 96, 0, 1, 123, 3, 2, 1, 0, 10, 19, 1, 17, 0, 65, 1, 253, 15, 65, 2, 253, + 15, 65, 3, 253, 15, 253, 147, 2, 11, + ]), + ); + } catch (e) { + return false; + } +}; + export const initializeWebAssembly = async (flags: Env.WebAssemblyFlags): Promise => { if (initialized) { return Promise.resolve(); @@ -82,7 +110,14 @@ export const initializeWebAssembly = async (flags: Env.WebAssemblyFlags): Promis let numThreads = flags.numThreads!; // ensure SIMD is supported - if (!isSimdSupported()) { + if (flags.simd === false) { + // skip SIMD feature checking as it is disabled explicitly by user + } else if (flags.simd === 'relaxed') { + // check if relaxed SIMD is supported + if (!isRelaxedSimdSupported()) { + throw new Error('Relaxed WebAssembly SIMD is not supported in the current environment.'); + } + } else if (!isSimdSupported()) { throw new Error('WebAssembly SIMD is not supported in the current environment.'); } diff --git a/js/web/package.json b/js/web/package.json index af02453b0870a..0af2513632472 100644 --- a/js/web/package.json +++ b/js/web/package.json @@ -68,6 +68,7 @@ "browser": "dist/ort.min.js", "exports": { ".": { + "types": "./types.d.ts", "node": { "import": "./dist/ort.node.min.mjs", "require": "./dist/ort.node.min.js" @@ -79,6 +80,7 @@ "require": "./dist/ort.min.js" }, "./all": { + "types": "./types.d.ts", "import": { "onnxruntime-web-use-extern-wasm": "./dist/ort.all.min.mjs", "default": "./dist/ort.all.bundle.min.mjs" @@ -86,6 +88,7 @@ "require": "./dist/ort.all.min.js" }, "./wasm": { + "types": "./types.d.ts", "import": { "onnxruntime-web-use-extern-wasm": "./dist/ort.wasm.min.mjs", "default": "./dist/ort.wasm.bundle.min.mjs" @@ -93,10 +96,12 @@ "require": "./dist/ort.wasm.min.js" }, "./webgl": { + "types": "./types.d.ts", "import": "./dist/ort.webgl.min.mjs", "require": "./dist/ort.webgl.min.js" }, "./webgpu": { + "types": "./types.d.ts", "import": { "onnxruntime-web-use-extern-wasm": "./dist/ort.webgpu.min.mjs", "default": "./dist/ort.webgpu.bundle.min.mjs" diff --git a/js/web/test/e2e/run.js b/js/web/test/e2e/run.js index d508625a04a5c..da02c26b87a4e 100644 --- a/js/web/test/e2e/run.js +++ b/js/web/test/e2e/run.js @@ -83,6 +83,12 @@ async function main() { await testExportsMain(PRESERVE, PACKAGES_TO_INSTALL); } + // perform type (typescript related) testing + { + const testTypeMain = require(path.join(TEST_E2E_RUN_FOLDER, './type/main')); + await testTypeMain(PRESERVE, PACKAGES_TO_INSTALL); + } + // prepare .wasm files for path override testing prepareWasmPathOverrideFiles(); diff --git a/js/web/test/e2e/type/main.js b/js/web/test/e2e/type/main.js new file mode 100644 index 0000000000000..697c6f5bc5d98 --- /dev/null +++ b/js/web/test/e2e/type/main.js @@ -0,0 +1,23 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +'use strict'; + +const path = require('path'); +const { installOrtPackages, runShellCmd } = require('./utils'); + +/** + * Entry point for type tests. + * + * @param {string[]} packagesToInstall + */ +module.exports = async function main(PRESERVE, PACKAGES_TO_INSTALL) { + console.log('Running type tests...'); + + // testcases/module-resolution + { + await installOrtPackages('module-resolution', PRESERVE, PACKAGES_TO_INSTALL); + + await runShellCmd('npx tsc', { wd: path.join(__dirname, 'testcases', 'module-resolution') }); + } +}; diff --git a/js/web/test/e2e/type/testcases/module-resolution.md b/js/web/test/e2e/type/testcases/module-resolution.md new file mode 100644 index 0000000000000..525122dfaf571 --- /dev/null +++ b/js/web/test/e2e/type/testcases/module-resolution.md @@ -0,0 +1,11 @@ +## module-resolution + +### Summary + +This is a very simple TypeScript application that uses the `onnxruntime-web` package. + +This test case is created to make sure future changes to the `onnxruntime-web` package do not break the module resolution of TypeScript. + +For details, see: + +- https://github.com/microsoft/onnxruntime/issues/16242 diff --git a/js/web/test/e2e/type/testcases/module-resolution/index.ts b/js/web/test/e2e/type/testcases/module-resolution/index.ts new file mode 100644 index 0000000000000..68a88cc2b419b --- /dev/null +++ b/js/web/test/e2e/type/testcases/module-resolution/index.ts @@ -0,0 +1,8 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import * as ort from 'onnxruntime-web'; + +export const print = () => { + console.log('onnxruntime-web version:', ort.env.versions); +}; diff --git a/js/web/test/e2e/type/testcases/module-resolution/package-lock.json b/js/web/test/e2e/type/testcases/module-resolution/package-lock.json new file mode 100644 index 0000000000000..990388706349a --- /dev/null +++ b/js/web/test/e2e/type/testcases/module-resolution/package-lock.json @@ -0,0 +1,29 @@ +{ + "name": "ort-type-test", + "version": "0.0.0", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "name": "ort-type-test", + "version": "0.0.0", + "devDependencies": { + "typescript": "^5.8.3" + } + }, + "node_modules/typescript": { + "version": "5.8.3", + "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.8.3.tgz", + "integrity": "sha512-p1diW6TqL9L07nNxvRMM7hMMw4c5XOo/1ibL4aAIGmSAt9slTE1Xgw5KWuof2uTOvCg9BY7ZRi+GaF+7sfgPeQ==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "tsc": "bin/tsc", + "tsserver": "bin/tsserver" + }, + "engines": { + "node": ">=14.17" + } + } + } +} diff --git a/js/web/test/e2e/type/testcases/module-resolution/package.json b/js/web/test/e2e/type/testcases/module-resolution/package.json new file mode 100644 index 0000000000000..e2a490f425d6c --- /dev/null +++ b/js/web/test/e2e/type/testcases/module-resolution/package.json @@ -0,0 +1,9 @@ +{ + "name": "ort-type-test-module-resolution", + "private": true, + "version": "0.0.0", + "type": "module", + "devDependencies": { + "typescript": "^5.8.3" + } +} diff --git a/js/web/test/e2e/type/testcases/module-resolution/tsconfig.json b/js/web/test/e2e/type/testcases/module-resolution/tsconfig.json new file mode 100644 index 0000000000000..89afe32431c2f --- /dev/null +++ b/js/web/test/e2e/type/testcases/module-resolution/tsconfig.json @@ -0,0 +1,7 @@ +{ + "compilerOptions": { + "module": "ES2020", + "moduleResolution": "bundler", + "strict": true + } +} diff --git a/js/web/test/e2e/type/utils.js b/js/web/test/e2e/type/utils.js new file mode 100644 index 0000000000000..9ce2240f6a7e7 --- /dev/null +++ b/js/web/test/e2e/type/utils.js @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +'use strict'; + +const path = require('path'); +const { spawn } = require('child_process'); +const treeKill = require('tree-kill'); + +async function installOrtPackages(testCaseName, PRESERVE, PACKAGES_TO_INSTALL) { + if (!PRESERVE) { + const wd = path.join(__dirname, 'testcases', testCaseName); + if (PACKAGES_TO_INSTALL.length === 0) { + await runShellCmd('npm ci', { wd }); + } else { + await runShellCmd(`npm install ${PACKAGES_TO_INSTALL.map((i) => `"${i}"`).join(' ')}`, { wd }); + } + } +} + +async function runShellCmd(cmd, { wd = __dirname, event = null, ready = null, ignoreExitCode = false }) { + console.log('==============================================================='); + console.log(' Running command in shell:'); + console.log(' > ' + cmd); + console.log('==============================================================='); + + return new Promise((resolve, reject) => { + const childProcess = spawn(cmd, { shell: true, stdio: ['ignore', 'pipe', 'inherit'], cwd: wd }); + childProcess.on('close', function (code, signal) { + if (code === 0 || ignoreExitCode) { + resolve(); + } else { + reject(`Process exits with code ${code}`); + } + }); + childProcess.stdout.on('data', (data) => { + process.stdout.write(data); + + if (ready && event && data.toString().includes(ready)) { + event.emit('serverReady'); + } + }); + if (event) { + event.on('kill', () => { + childProcess.stdout.destroy(); + treeKill(childProcess.pid); + console.log('killing process...'); + }); + } + }); +} + +module.exports = { + runShellCmd, + installOrtPackages, +}; diff --git a/js/web/tsconfig.json b/js/web/tsconfig.json index 80d0cd0642b80..8eec6a723101f 100644 --- a/js/web/tsconfig.json +++ b/js/web/tsconfig.json @@ -6,5 +6,5 @@ "typeRoots": ["./node_modules/@webgpu/types", "./node_modules/@types", "../node_modules/@types"] }, "include": ["lib", "test"], - "exclude": ["lib/wasm/proxy-worker", "test/ort.test.js", "test/ort.test.min.js"] + "exclude": ["lib/wasm/proxy-worker", "test/ort.test.js", "test/ort.test.min.js", "test/e2e"] } diff --git a/onnxruntime/contrib_ops/webgpu/bert/bias_add.cc b/onnxruntime/contrib_ops/webgpu/bert/bias_add.cc index e822f8764b63f..0f2e6d725007f 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/bias_add.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/bias_add.cc @@ -21,16 +21,16 @@ ONNX_OPERATOR_KERNEL_EX( BiasAdd); Status BiasAddProgram::GenerateShaderCode(ShaderHelper& shader) const { - const ShaderVariableHelper& input = shader.AddInput("input"); - const ShaderVariableHelper& bias = shader.AddInput("bias"); - const ShaderVariableHelper& residual = shader.AddInput("residual"); - const ShaderVariableHelper& output = shader.AddOutput("output"); + const ShaderVariableHelper& input = shader.AddInput("input", ShaderUsage::UseUniform); + const ShaderVariableHelper& bias = shader.AddInput("bias", ShaderUsage::UseUniform); + const ShaderVariableHelper& residual = shader.AddInput("residual", ShaderUsage::UseUniform); + const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform); shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size") - << "let value = " << input.GetByOffset("global_idx") + << " let value = " << input.GetByOffset("global_idx") << " + " << bias.GetByOffset("global_idx % uniforms.channels") << " + " << residual.GetByOffset("global_idx") << ";\n" - << output.SetByOffset("global_idx", "value"); + << " " + output.SetByOffset("global_idx", "value"); return Status::OK(); } @@ -47,23 +47,26 @@ Status BiasAdd::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) co } int64_t channels = input_shape[2]; - int64_t components = GetMaxComponents(channels); - channels /= components; - TensorShape bias_shape = bias->Shape(); if (bias_shape.NumDimensions() != 1 || bias_shape[0] != channels) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "BiasAdd bias should have 1 dimension with size equal to the number of channels."); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "BiasAdd bias should have 1 dimension with size equal to the number of channels."); } + int components = GetMaxComponents(channels); + channels /= components; + auto* output = context.Output(0, input_shape); int64_t output_size = output->Shape().Size() / components; BiasAddProgram program{}; - program.AddInputs({{input}, {bias}, {residual}}) - .AddOutput({output}) + program + .AddInputs({{input, ProgramTensorMetadataDependency::None, components}, + {bias, ProgramTensorMetadataDependency::None, components}, + {residual, ProgramTensorMetadataDependency::None, components}}) + .AddOutput({output, ProgramTensorMetadataDependency::None, components}) .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) - .AddUniformVariables({{static_cast(output_size)}, - {static_cast(channels)}}); + .AddUniformVariables({{static_cast(output_size)}, {static_cast(channels)}}); return context.RunProgram(program); } diff --git a/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc b/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc index 99cd643423400..4656fc486fccf 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/bias_split_gelu.cc @@ -22,20 +22,21 @@ ONNX_OPERATOR_KERNEL_EX( BiasSplitGelu); Status BiasSplitGeluProgram::GenerateShaderCode(ShaderHelper& shader) const { - const ShaderVariableHelper& input = shader.AddInput("input"); - const ShaderVariableHelper& bias = shader.AddInput("bias"); - const ShaderVariableHelper& output = shader.AddOutput("output"); + const ShaderVariableHelper& x = + shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); + const ShaderVariableHelper& bias = shader.AddInput("bias", ShaderUsage::UseUniform); + const ShaderVariableHelper& output = shader.AddOutput("output", ShaderUsage::UseUniform); shader.AdditionalImplementation() << ErfImpl; shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size") << "const M_SQRT2: f32 = sqrt(2.0);\n" - << "const halfChannels = uniforms.channels / 2u;\n" + << "let halfChannels = uniforms.channels / 2u;\n" << "let biasIdx = global_idx % halfChannels;\n" << "let batchIndex = global_idx / halfChannels;\n" << "let inputOffset = biasIdx + batchIndex * halfChannels * 2;\n" - << "let valueLeft = " << input.GetByOffset("inputOffset") << " + " << bias.GetByOffset("biasIdx") << ";\n" - << "let valueRight = " << input.GetByOffset("inputOffset + halfChannels") << " + " << bias.GetByOffset("biasIdx + halfChannels") << ";\n" + << "let valueLeft = " << x.GetByOffset("inputOffset") << " + " << bias.GetByOffset("biasIdx") << ";\n" + << "let valueRight = " << x.GetByOffset("inputOffset + halfChannels") << " + " << bias.GetByOffset("biasIdx + halfChannels") << ";\n" << "let geluRight = valueRight * 0.5 * (erf_v(valueRight / M_SQRT2) + 1);\n" << output.SetByOffset("global_idx", "valueLeft * geluRight"); @@ -53,25 +54,27 @@ Status BiasSplitGelu::ComputeInternal(onnxruntime::webgpu::ComputeContext& conte } int64_t channels = input_shape[2]; - int64_t components = GetMaxComponents(channels); - channels /= components; input_shape[2] = channels / 2; // for output shape calculation (N,S,D) -> (N,S,D/2) TensorShape bias_shape = bias->Shape(); if (bias_shape.NumDimensions() != 1 || bias_shape[0] != channels) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "BiasSplitGelu bias should have 1 dimension with size equal to the number of channels."); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "BiasSplitGelu bias should have 1 dimension with size equal to the number of channels."); } + int components = GetMaxComponents(channels); + channels /= components; + auto* output = context.Output(0, input_shape); int64_t output_size = output->Shape().Size() / components; BiasSplitGeluProgram program{}; - program.AddInputs({{input, ProgramTensorMetadataDependency::TypeAndRank}, - {bias}}) - .AddOutput({output}) + program + .AddInputs({{input, ProgramTensorMetadataDependency::None, components}, + {bias, ProgramTensorMetadataDependency::None, components}}) + .AddOutput({output, ProgramTensorMetadataDependency::None, components}) .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) - .AddUniformVariables({{static_cast(output_size)}, - {static_cast(channels)}}); + .AddUniformVariables({{static_cast(output_size)}, {static_cast(channels)}}); return context.RunProgram(program); } diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 26b08159a3b84..bb5de40eb27c5 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -28,29 +28,28 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { // Attention bias is in BN(total_sequence_length) const auto& key = shader.AddInput("key", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias | ShaderUsage::UseIndicesTypeAlias); shader.AddInput("value", ShaderUsage::UseUniform); - const auto& present_key = shader.AddOutput("present_key", ShaderUsage::UseUniform); + const auto& present_key = shader.AddOutput("present_key", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias); const auto& present_value = shader.AddOutput("present_value", ShaderUsage::UseUniform); - const auto& valid_present_shape = shader.AddIndices("valid_present_shape"); + const auto& copy_kv_shape = shader.AddIndices("copy_kv_shape"); - shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.valid_present_size") - << " let output_indices = " << valid_present_shape.OffsetToIndices("global_idx") << ";\n" + shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.copy_size") + << " let output_indices = " << copy_kv_shape.OffsetToIndices("global_idx") << ";\n" << " let head_size_id = output_indices[3];\n" " let sequence_id = output_indices[2];\n" " let num_head_id = output_indices[1];\n" - " let batch = output_indices[0];\n" - << " let present_offset = " << (past_present_share_buffer_ ? present_key.IndicesToOffset("output_indices") : "global_idx") << ";\n"; + " let batch = output_indices[0];\n"; if (has_past_) { shader.MainFunctionBody() << "let past_sequence_length = uniforms.past_sequence_length;\n"; if (past_present_share_buffer_) { - shader.MainFunctionBody() << "if (sequence_id >= past_sequence_length) {\n" - << " let offset = " << key.IndicesToOffset(kv_BNSH_ ? "key_indices_t(batch, num_head_id, sequence_id - past_sequence_length, head_size_id)" : "key_indices_t(batch, sequence_id - past_sequence_length, num_head_id, head_size_id)") << ";\n" + shader.MainFunctionBody() << " let present_offset = " << present_key.IndicesToOffset("present_key_indices_t(batch, num_head_id, past_sequence_length + sequence_id, head_size_id)") << ";\n" + << " let offset = " << key.IndicesToOffset(kv_BNSH_ ? "key_indices_t(batch, num_head_id, sequence_id, head_size_id)" : "key_indices_t(batch, sequence_id, num_head_id, head_size_id)") << ";\n" << " " << present_key.SetByOffset("present_offset", "key[offset]") << ";\n" - << " " << present_value.SetByOffset("present_offset", "value[offset]") << ";\n" - << "}"; + << " " << present_value.SetByOffset("present_offset", "value[offset]") << ";\n"; } else { const auto& past_key = shader.AddInput("past_key", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias | ShaderUsage::UseIndicesTypeAlias); shader.AddInput("past_value", ShaderUsage::UseUniform); - shader.MainFunctionBody() << "if (sequence_id < past_sequence_length) {\n" + shader.MainFunctionBody() << "let present_offset = global_idx;" + << "if (sequence_id < past_sequence_length) {\n" << " let pastOffset = " << past_key.IndicesToOffset("past_key_indices_t(batch, num_head_id, sequence_id, head_size_id)") << ";\n" << " " << present_key.SetByOffset("present_offset", "past_key[pastOffset]") << ";\n" << " " << present_value.SetByOffset("present_offset", "past_value[pastOffset]") << ";\n" @@ -61,7 +60,8 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { << "}"; } } else { - shader.MainFunctionBody() << "let offset = " << key.IndicesToOffset(kv_BNSH_ ? "key_indices_t(batch, num_head_id, sequence_id, head_size_id)" : "key_indices_t(batch, sequence_id, num_head_id, head_size_id)") << ";\n" + shader.MainFunctionBody() << " let present_offset = " << (past_present_share_buffer_ ? present_key.IndicesToOffset("output_indices") : "global_idx") << ";\n" + << "let offset = " << key.IndicesToOffset(kv_BNSH_ ? "key_indices_t(batch, num_head_id, sequence_id, head_size_id)" : "key_indices_t(batch, sequence_id, num_head_id, head_size_id)") << ";\n" << present_key.SetByOffset("present_offset", "key[offset]") << ";\n" << present_value.SetByOffset("present_offset", "value[offset]") << ";\n"; } @@ -75,18 +75,21 @@ Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAtt // This makes it so that FlashAttention only needs to look at present key and value, and saves // number of input buffers in the shader, which we run out of (<=8) without this optimization. const int components = parameters.head_size_ % 4 == 0 ? 4 : (parameters.head_size_ % 2 == 0 ? 2 : 1); + bool has_past = (parameters.total_sequence_length_ - parameters.kv_sequence_length_) > 0; // parameters.total_sequence_length_ is past_sequence_length + kv_sequence_length. // parameters.kv_num_heads_ may be smaller than parameters.num_heads_ when parameters.is_gqa_ is true. int num_heads = parameters.is_gqa_ ? parameters.kv_num_heads_ : parameters.num_heads_; - TensorShape valid_present_shape{parameters.batch_size_, num_heads, parameters.total_sequence_length_, parameters.head_size_ / components}; - int64_t valid_kv_size = valid_present_shape.Size(); - bool has_past = (parameters.past_sequence_length_ > 0); + // Only copy the new kv data for static kv cache + int copy_sequence_length = has_past && parameters.past_present_share_buffer_ ? parameters.kv_sequence_length_ : parameters.total_sequence_length_; + TensorShape copy_kv_shape{parameters.batch_size_, num_heads, copy_sequence_length, parameters.head_size_ / components}; + int64_t copy_size = copy_kv_shape.Size(); CopyKVCacheProgram program{"CopyKVCache", has_past, parameters.qkv_format_ == Q_K_V_BSNH_BNSH_BNSH, parameters.past_present_share_buffer_}; if (parameters.qkv_format_ == Q_K_V_BSNH_BNSH_BNSH) { program.AddInputs({{K, ProgramTensorMetadataDependency::TypeAndRank, components}, {V, ProgramTensorMetadataDependency::TypeAndRank, components}}); } else { ORT_ENFORCE(parameters.qkv_format_ == Q_K_V_BSNH, "qkv format ", parameters.qkv_format_, " is not supported yet in CopyKVCache."); + // Reshape (batch_size, kv_sequence_length, kv_hidden_size) to (batch_size, kv_sequence_length, num_head, head_size) TensorShape reshaped_KV_shape{parameters.batch_size_, parameters.kv_sequence_length_, num_heads, parameters.head_size_ / components}; program.AddInputs({{K, ProgramTensorMetadataDependency::TypeAndRank, reshaped_KV_shape, components}, {V, ProgramTensorMetadataDependency::TypeAndRank, reshaped_KV_shape, components}}); @@ -97,11 +100,11 @@ Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAtt } program.AddOutputs({{present_key, ProgramTensorMetadataDependency::Rank, components}, {present_value, ProgramTensorMetadataDependency::Rank, components}}) - .AddIndices(std::move(valid_present_shape)); - program.SetDispatchGroupSize(onnxruntime::narrow((valid_kv_size + 63) / 64)) + .AddIndices(std::move(copy_kv_shape)); + program.SetDispatchGroupSize(static_cast((copy_size + 63) / 64)) .SetWorkgroupSize(64) .CacheHint(has_past, parameters.qkv_format_, parameters.past_present_share_buffer_) - .AddUniformVariables({{static_cast(valid_kv_size)}, + .AddUniformVariables({{static_cast(copy_size)}, // Note that when parameters.past_present_share_buffer_ is true, parameters.past_sequence_length_ will become to // max_sequence_length. To get a valid past_sequence_length, we use total_sequence_length - kv_sequence_length. {static_cast(parameters.total_sequence_length_ - parameters.kv_sequence_length_)}}); @@ -425,49 +428,397 @@ Status FlashAttentionProgram::GenerateShaderCode(ShaderHelper& shader) const { return Status::OK(); } -Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, const Tensor* attention_bias, - Tensor* output, const Tensor* past_key, Tensor* present_key, const Tensor* past_value, Tensor* present_value, - const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context) { - ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value)); +Status FlashAttentionDecodeQKTProgram::GenerateShaderCode(ShaderHelper& shader) const { + shader.AddInput("q", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); + shader.AddInput("present_key", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); + if (has_attention_bias_) { + shader.AddInput("attention_bias", ShaderUsage::UseUniform); + } + shader.AddOutput("output", ShaderUsage::UseUniform); + shader.AddOutput("metadata", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); + // Note that this shader adopts similar algorithm with dp4a generation shader. + // + // This algorithm works to compute dot product of keys with queries parallelly, by processing on the k (head_size) dimension at each step amongst tile_size_k_vec threads, + // and utilizing the remaining threads in the workgroup to process additional rows of |present_key| in parallel (such that the values in shared memory (tile_q) for |q| can be reused). + // For each load of q, the tile_size_k_vec threads also reload |present_key| tile_size/sub_tile_count times to compute partial dot products of other |present_key| rows + // in order to complete all tile_size |present_key| rows in this workgroup and also reusing the loaded in register values of |q|. + constexpr int tile_size_k_vec = 8; - const uint32_t tile_size = 64; - bool has_attention_bias = attention_bias != nullptr; - FlashAttentionProgram program{"FlashAttention", has_attention_bias, parameters.head_size_, parameters.num_heads_}; - program.AddInputs({{Q, ProgramTensorMetadataDependency::TypeAndRank, 4}, - {present_key, ProgramTensorMetadataDependency::TypeAndRank, 4}, - {present_value, ProgramTensorMetadataDependency::TypeAndRank, 4}}); - if (has_attention_bias) { - program.AddInputs({{attention_bias, ProgramTensorMetadataDependency::TypeAndRank}}); + // 1. Each workgroup processes one row of |q| and tile_size rows of |present_key| + // + // 2. Computation Process: + // - Reads [tile_size][tile_size_k_vec] block of |present_key| data at a time + // - Each thread within workgroup computes dot products of 4 A*B elements since each k represents 4 elements of |present_key| + // - Stores intermediate results in shared memory (inner_qk_values) + // - Iterates through columns (head_size_vec) accumulating results in inner_qk_values + // - Performs final reduction sum in inner_qk_values for output + shader.AdditionalImplementation() << "const tile_size = " << tile_size_ << "u;\n" + << "const tile_size_k_vec = " << tile_size_k_vec << "u;\n" + << "const sub_tile_count = " << WorkgroupSizeX() / tile_size_k_vec << "u;\n"; + shader.AdditionalImplementation() << R"ADDNL_FN( +var tile_q: array; +var inner_qk_values: array, tile_size>; +var tile_qk: array; +)ADDNL_FN"; + + if (has_attention_bias_) { + shader.AdditionalImplementation() << R"HELPER_FN( + fn loadAttentionBias(idx: u32) -> q_element_t + { + return attention_bias[idx]; + } + )HELPER_FN"; + } else { + shader.AdditionalImplementation() << R"HELPER_FN( + fn loadAttentionBias(idx: u32) -> q_element_t + { + return q_element_t(0); + } + )HELPER_FN"; } - program.AddOutputs({{output, ProgramTensorMetadataDependency::TypeAndRank, 4}}); + + shader.MainFunctionBody() << R"MAIN_FN( + let local_row = u32(local_idx / tile_size_k_vec); + let local_col = local_idx % tile_size_k_vec; + let total_seq_offset = (workgroup_idx % uniforms.num_total_seq_length_tile) * tile_size; + let head_idx = u32(workgroup_idx / uniforms.num_total_seq_length_tile); + let q_offset = head_idx * uniforms.head_size_vec; + var total_sequence_length = uniforms.total_sequence_length; + let present_offset = u32(head_idx / uniforms.n_reps) * uniforms.present_sequence_length * uniforms.head_size_vec; + for (var k: u32 = 0u; k < uniforms.head_size_vec; k += tile_size_k_vec) { + if (local_idx < tile_size_k_vec && k + local_idx < uniforms.head_size_vec) { + tile_q[local_idx] = q[q_offset + k + local_idx]; + } + workgroupBarrier(); + let q_data = tile_q[local_col]; + if (k + local_col < uniforms.head_size_vec) { + for (var row_offset = 0u; row_offset < tile_size; row_offset += sub_tile_count) { + if (total_seq_offset + row_offset + local_row < total_sequence_length) { + inner_qk_values[row_offset + local_row][local_col] += dot(present_key[present_offset + (total_seq_offset + row_offset + local_row) * uniforms.head_size_vec + k + local_col], q_data); + } + } + } + workgroupBarrier(); + } + + if (local_idx < tile_size && total_seq_offset + local_idx < total_sequence_length) { + var sum = q_element_t(0); + for (var i = 0u; i < tile_size_k_vec; i++) { + sum += inner_qk_values[local_idx][i]; + } + + let output_idx = head_idx * total_sequence_length + total_seq_offset + local_idx; + sum = sum * q_element_t(uniforms.alpha) + loadAttentionBias(output_idx); + tile_qk[local_idx] = sum; + output[output_idx] = sum; + } + workgroupBarrier(); + + if (local_idx == 0u) { + // Calculate the max and sum in current split. + var l_max = f32(-3.402823e+38f); + var l_sum = f32(0); + for (var i = 0u; i < tile_size && (total_seq_offset + i) < total_sequence_length; i++) { + l_max = max(l_max, f32(tile_qk[i])); + } + for (var i = 0u; i < tile_size && (total_seq_offset + i) < total_sequence_length; i++) { + l_sum += exp(f32(tile_qk[i]) - l_max); + } + let meta_offset = head_idx * uniforms.num_total_seq_length_tile + workgroup_idx % uniforms.num_total_seq_length_tile; + metadata[meta_offset] = metadata_value_t(metadata_element_t(l_max), metadata_element_t(l_sum)); + } +)MAIN_FN"; + + return Status::OK(); +} + +Status ComputeFlashAttentionDecodeQKT(onnxruntime::webgpu::ComputeContext& context, const Tensor* Q, + const Tensor* attention_bias, Tensor* output, Tensor* present_key, Tensor* metadata, + const WebgpuAttentionParameters& parameters, uint32_t num_total_seq_length_tile, uint32_t tile_size) { const float alpha = parameters.scale_ == 0.0f ? 1.f / sqrt(static_cast(parameters.head_size_)) : parameters.scale_; - std::string cache_hint = std::to_string(has_attention_bias) + - std::to_string(parameters.head_size_) + - std::to_string(parameters.num_heads_); - const uint32_t num_seq_tile = (parameters.sequence_length_ + tile_size - 1) / tile_size; - program.SetDispatchGroupSize(parameters.num_heads_ * num_seq_tile) - .SetWorkgroupSize(tile_size) - .CacheHint(cache_hint) - .AddUniformVariables({{static_cast(parameters.sequence_length_)}, + + const bool has_attention_bias = attention_bias != nullptr; + const int components = 4; + + FlashAttentionDecodeQKTProgram program{"FlashAttentionDecodeQKT", has_attention_bias, tile_size}; + program.AddInputs({{Q, ProgramTensorMetadataDependency::TypeAndRank, components}, + {present_key, ProgramTensorMetadataDependency::TypeAndRank, components}}); + if (has_attention_bias) { + program.AddInput({attention_bias, ProgramTensorMetadataDependency::TypeAndRank}); + } + program.AddOutputs({{output, ProgramTensorMetadataDependency::Rank}, + {metadata, ProgramTensorMetadataDependency::Rank, 2}}); + + const uint32_t vectorized_head_size = parameters.head_size_ / components; + program.SetDispatchGroupSize(parameters.num_heads_ * num_total_seq_length_tile) + .SetWorkgroupSize(64) + .CacheHint(tile_size, has_attention_bias) + .AddUniformVariables({{static_cast(vectorized_head_size)}, {static_cast(parameters.total_sequence_length_)}, - {static_cast(parameters.past_present_share_buffer_ ? parameters.past_sequence_length_ : parameters.total_sequence_length_)}, - {static_cast(parameters.total_sequence_length_ - parameters.kv_sequence_length_)}, - {static_cast(parameters.is_gqa_ ? 1 : 0)}, + {static_cast(alpha)}, + // present_sequence_length is used to index into the KV cache, for static kv cache it is the max sequence length. + {static_cast(parameters.is_gqa_ ? parameters.seqlen_present_kv_cache_ : parameters.total_sequence_length_)}, {static_cast(parameters.n_reps)}, - {alpha}, - {num_seq_tile}}); + {num_total_seq_length_tile}}); return context.RunProgram(program); } +Status FlashAttentionDecodeSplitVxProgram::GenerateShaderCode(ShaderHelper& shader) const { + shader.AddInput("metadata", ShaderUsage::UseUniform); + shader.AddInput("qk", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); + shader.AddInput("present_value", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); + shader.AddOutput("out_split_vx", ShaderUsage::UseUniform); + + // Note that this shader adopts similar algorithm with dp4a generation shader. + // + // This algorithm works to compute dot product of v with qk parallelly, by processing on the head_size dimension at each step amongst tile_size_k_vec threads, + // and utilizing the remaining threads in the workgroup to process additional rows of |present_value| in parallel (such that the values in shared memory (tile_qk) for |qk| can be reused). + // The tile_size_k_vec threads also reload |present_value| tile_size/sub_tile_count times to compute partial dot products of other |present_value| rows + // in order to complete all tile_size |present_value| rows in this workgroup and also reusing the values in tile_qk. + // + // The difference with FlashAttentionDecodeQKTProgram is that the dot products go through the rows (total_sequence_length) of |present_value| instead of columns (head_size_vec). + // And each workgroup only calculate current tile_size's dot products instead of iterating the whole row |total_sequence_length|. + // That's why this shader is a split shader. The final reduce will be done in FlashAttentionDecodeReduceProgram. + constexpr int tile_size_k_vec = 8; + + shader.AdditionalImplementation() << "const head_size_vec = " << head_size_vec_ << "u;\n" + << "const tile_size = " << tile_size_ << "u;\n" + << "const tile_size_k_vec = " << tile_size_k_vec << "u;\n" + << "const sub_tile_count = " << WorkgroupSizeX() / tile_size_k_vec << "u;\n"; + shader.AdditionalImplementation() << R"HELPER_FN( +var tile_qk: array; +var tile_output: array; +var qkv_values: array, sub_tile_count>; + + )HELPER_FN"; + + // TODO: Ideally, there should only be two shaders FlashAttentionDecodeSplitVx and FlashAttentionDecodeVxReduce, which can also reduce the intermediate memory. + // The FlashAttentionDecodeQKT can be merged into split shader and do the final softmax adjustment in the reduce shader. However, some issues are met that when + // the total sequence length exceeds some value, the result will become garbage. Since it can't be resolved in a short time, leave it as TODO to fix it in future. + shader.MainFunctionBody() << R"MAIN_FN( + let local_row = u32(local_idx / tile_size_k_vec); + let local_col = local_idx % tile_size_k_vec; + let total_seq_offset = (workgroup_idx % uniforms.num_total_seq_length_tile) * tile_size; + let head_idx = u32(workgroup_idx / uniforms.num_total_seq_length_tile); + var total_sequence_length = uniforms.total_sequence_length; + let present_offset = u32(head_idx / uniforms.n_reps) * uniforms.head_size_vec * uniforms.present_sequence_length; + + // Calculate the global max and sum in qk. + var g_max = f32(-3.402823e+38f); + for (var i = 0u; i < uniforms.num_total_seq_length_tile; i++) + { + let meta_offset = head_idx * uniforms.num_total_seq_length_tile + i; + g_max = max(g_max, f32(metadata[meta_offset].x)); + } + var g_sum = f32(0); + for (var i = 0u; i < uniforms.num_total_seq_length_tile; i++) + { + let meta_offset = head_idx * uniforms.num_total_seq_length_tile + i; + let m_value = metadata[meta_offset]; + g_sum += exp(f32(m_value.x) - g_max) * f32(m_value.y); + } + + if (total_seq_offset + local_idx < total_sequence_length) { + tile_qk[local_idx] = qk_value_t(exp(f32(qk[head_idx * total_sequence_length + total_seq_offset + local_idx]) - g_max) / g_sum); + } + + for (var k: u32 = 0u; k < uniforms.head_size_vec; k += tile_size_k_vec) { + var value = present_value_value_t(0); + qkv_values[local_row][local_col] = present_value_value_t(0); + workgroupBarrier(); + + if (k + local_col < uniforms.head_size_vec) { + for (var row_offset = 0u; row_offset < tile_size; row_offset += sub_tile_count) { + if (total_seq_offset + row_offset + local_row < total_sequence_length) { + value += present_value[present_offset + (total_seq_offset + row_offset + local_row) * uniforms.head_size_vec + k + local_col] * tile_qk[row_offset + local_row]; + } + } + } + + qkv_values[local_row][local_col] = value; + workgroupBarrier(); + + if (local_idx < tile_size_k_vec) { + for (var i = 0u; i < sub_tile_count; i++) { + tile_output[k + local_idx] += qkv_values[i][local_idx]; + } + } + workgroupBarrier(); + } + + for (var i = local_idx; i < uniforms.head_size_vec; i += workgroup_size_x) { + let out_offset = head_idx * uniforms.num_total_seq_length_tile * uniforms.head_size_vec + (workgroup_idx % uniforms.num_total_seq_length_tile) * uniforms.head_size_vec + i; + out_split_vx[out_offset] = tile_output[i]; + } +)MAIN_FN"; + + return Status::OK(); +} + +Status ComputeFlashAttentionDecodeSplitVxScore(onnxruntime::webgpu::ComputeContext& context, + const Tensor* metadata, + const Tensor* qk, + Tensor* out_split_vx, + Tensor* present_value, + const WebgpuAttentionParameters& parameters, + uint32_t num_total_seq_length_tile, + uint32_t tile_size) { + const int components = 4; + int head_size_vec = parameters.v_head_size_ / components; + FlashAttentionDecodeSplitVxProgram program{"FlashAttentionDecodeSplitVx", tile_size, head_size_vec}; + program.AddInputs({{metadata, ProgramTensorMetadataDependency::TypeAndRank, 2}, + {qk, ProgramTensorMetadataDependency::TypeAndRank}, + {present_value, ProgramTensorMetadataDependency::TypeAndRank, components}}); + program.AddOutputs({{out_split_vx, ProgramTensorMetadataDependency::TypeAndRank, components}}); // [B, N, split_k, head_size] + program.SetDispatchGroupSize(parameters.num_heads_ * num_total_seq_length_tile) + .CacheHint(tile_size, head_size_vec) + .SetWorkgroupSize(64) + .AddUniformVariables({{static_cast(parameters.total_sequence_length_)}, + {static_cast(head_size_vec)}, + {static_cast(parameters.is_gqa_ ? parameters.seqlen_present_kv_cache_ : parameters.total_sequence_length_)}, + {static_cast(parameters.n_reps)}, + num_total_seq_length_tile}); + + return context.RunProgram(program); +} + +Status FlashAttentionDecodeVxReduceProgram::GenerateShaderCode(ShaderHelper& shader) const { + shader.AddInput("input", ShaderUsage::UseUniform); + shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); + + // Inputs are splits of the GQA output, split into num_total_seq_length_tiles rows. + // This shader needs to add these splits across the row dimension to arrive at the final result. The column is head size wide. + // The reduction achieves maximum parallelization by splitting this task first into tile_size columns that each workgroup is responsible for. + // Then within each workgroup the task of summation over the num_total_seq_length_tile for the tile_size columns is further split in two ways. + // First across the row dimension to have WORKGROUP_SIZE/TILE_SIZE parallel computations of summation of TILE_SIZE rows. + // Then across the column dimension where each thread is responsible for 1 column of the TILE_SIZE columns the workgroup is resposible for. + shader.AdditionalImplementation() << "const TILE_SIZE = " << tile_size_ << ";\n"; + shader.AdditionalImplementation() << R"HELPER_FN( +var tile_input: array, TILE_SIZE>; + )HELPER_FN"; + + shader.MainFunctionBody() << R"MAIN_FN( + let head_size_offset = (workgroup_idx % uniforms.num_head_size_tile) * TILE_SIZE; + let head_idx = u32(workgroup_idx / uniforms.num_head_size_tile); + let in_offset = head_idx * uniforms.num_total_seq_length_tile * uniforms.head_size_vec; + var value = output_value_t(0); + let local_row = u32(local_idx / TILE_SIZE); + let local_col = local_idx % TILE_SIZE; + + if (head_size_offset + local_col < uniforms.head_size_vec) { + for (var r = 0u; r < uniforms.num_total_seq_length_tile; r += TILE_SIZE) { + if (r + local_row < uniforms.num_total_seq_length_tile) { + value += input[in_offset + (r + local_row) * uniforms.head_size_vec + head_size_offset + local_col]; + } + } + } + + tile_input[local_row][local_col] = value; + workgroupBarrier(); + + if (local_idx < TILE_SIZE && head_size_offset + local_idx < uniforms.head_size_vec) { + value = output_value_t(0); + for (var i = 0u; i < TILE_SIZE; i++) { + value += tile_input[i][local_idx]; + } + let output_id = head_idx * uniforms.head_size_vec + head_size_offset + local_idx; + output[output_id] = value; + } +)MAIN_FN"; + + return Status::OK(); +} + +Status ComputeFlashAttentionDecodeVxReduce(onnxruntime::webgpu::ComputeContext& context, + const Tensor* out_split_vx, + Tensor* output, + const WebgpuAttentionParameters& parameters, + uint32_t num_total_seq_length_tile) { + const int components = 4; + constexpr int tile_size = 8; + int tile_head_size = tile_size * components; + FlashAttentionDecodeVxReduceProgram program{"FlashAttentionDecodeVxReduce", tile_size}; + program.AddInputs({{out_split_vx, ProgramTensorMetadataDependency::TypeAndRank, components}}); + program.AddOutputs({{output, ProgramTensorMetadataDependency::TypeAndRank, components}}); + const uint32_t num_head_size_tile = static_cast((parameters.v_head_size_ + tile_head_size - 1) / tile_head_size); + program.SetDispatchGroupSize(parameters.num_heads_ * num_head_size_tile) + .CacheHint(tile_size) + .SetWorkgroupSize(tile_size * tile_size) + .AddUniformVariables({{static_cast(parameters.v_head_size_ / components)}, + num_total_seq_length_tile, + {num_head_size_tile}}); + + return context.RunProgram(program); +} + +Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, const Tensor* attention_bias, + Tensor* output, const Tensor* past_key, Tensor* present_key, const Tensor* past_value, Tensor* present_value, + const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context) { + ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value)); + + if (parameters.sequence_length_ > 1) { + const uint32_t tile_size = 64; + bool has_attention_bias = attention_bias != nullptr; + FlashAttentionProgram program{"FlashAttention", has_attention_bias, parameters.head_size_, parameters.num_heads_}; + program.AddInputs({{Q, ProgramTensorMetadataDependency::TypeAndRank, 4}, + {present_key, ProgramTensorMetadataDependency::TypeAndRank, 4}, + {present_value, ProgramTensorMetadataDependency::TypeAndRank, 4}}); + if (has_attention_bias) { + program.AddInputs({{attention_bias, ProgramTensorMetadataDependency::TypeAndRank}}); + } + program.AddOutputs({{output, ProgramTensorMetadataDependency::TypeAndRank, 4}}); + const float alpha = parameters.scale_ == 0.0f ? 1.f / sqrt(static_cast(parameters.head_size_)) + : parameters.scale_; + std::string cache_hint = std::to_string(has_attention_bias) + + std::to_string(parameters.head_size_) + + std::to_string(parameters.num_heads_); + const uint32_t num_seq_tile = (parameters.sequence_length_ + tile_size - 1) / tile_size; + program.SetDispatchGroupSize(parameters.num_heads_ * num_seq_tile) + .SetWorkgroupSize(tile_size) + .CacheHint(cache_hint) + .AddUniformVariables({{static_cast(parameters.sequence_length_)}, + {static_cast(parameters.total_sequence_length_)}, + {static_cast(parameters.past_present_share_buffer_ ? parameters.past_sequence_length_ : parameters.total_sequence_length_)}, + {static_cast(parameters.total_sequence_length_ - parameters.kv_sequence_length_)}, + {static_cast(parameters.is_gqa_ ? 1 : 0)}, + {static_cast(parameters.n_reps)}, + {alpha}, + {num_seq_tile}}); + + return context.RunProgram(program); + } + + const TensorShapeVector qk_dims({parameters.batch_size_, parameters.num_heads_, + parameters.sequence_length_, parameters.total_sequence_length_}); + const TensorShape qk_shape(qk_dims); + Tensor qk = context.CreateGPUTensor(Q->DataType(), qk_shape); + constexpr uint32_t tile_size = 64; + const uint32_t num_total_seq_length_tile = (parameters.total_sequence_length_ + tile_size - 1) / tile_size; + // The metadata is used to store the max and sum of each tile. + const TensorShapeVector metadata_dims({parameters.batch_size_, parameters.num_heads_, + num_total_seq_length_tile, 2}); + const TensorShape metadata_shape(metadata_dims); + Tensor metadata = context.CreateGPUTensor(Q->DataType(), metadata_shape); + ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeQKT(context, Q, attention_bias, &qk, present_key, &metadata, + parameters, num_total_seq_length_tile, tile_size)); + + const TensorShapeVector out_split_vx_dims({parameters.batch_size_, parameters.num_heads_, num_total_seq_length_tile, parameters.head_size_}); + const TensorShape out_split_vx_shape(out_split_vx_dims); + Tensor out_split_vx = context.CreateGPUTensor(Q->DataType(), out_split_vx_shape); + ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeSplitVxScore(context, &metadata, &qk, &out_split_vx, present_value, parameters, num_total_seq_length_tile, tile_size)); + ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeVxReduce(context, &out_split_vx, output, parameters, num_total_seq_length_tile)); + + return Status::OK(); +} + bool CanApplyFlashAttention(const Tensor* bias, const Tensor* present_key, const Tensor* present_value, const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context) { return parameters.batch_size_ == 1 && !parameters.is_packed_qkv_ && parameters.head_size_ == parameters.v_head_size_ && bias == nullptr && - parameters.sequence_length_ > 1 && context.HasFeature(wgpu::FeatureName::Subgroups) && present_key != nullptr && present_value != nullptr && present_key->SizeInBytes() > 0 && present_value->SizeInBytes() > 0 && parameters.head_size_ % 4 == 0; diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index 8931403641a81..c066d6249c8b2 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -23,7 +23,7 @@ class CopyKVCacheProgram final : public Program { Status GenerateShaderCode(ShaderHelper& sh) const override; - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"valid_present_size", ProgramUniformVariableDataType::Uint32}, + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"copy_size", ProgramUniformVariableDataType::Uint32}, {"past_sequence_length", ProgramUniformVariableDataType::Uint32}); private: @@ -61,6 +61,62 @@ class FlashAttentionProgram final : public Program { int qkv_num_heads_; }; +class FlashAttentionDecodeQKTProgram final : public Program { + public: + FlashAttentionDecodeQKTProgram(const std::string& kernel_name, + bool has_attention_bias, uint32_t tile_size) + : Program{kernel_name}, has_attention_bias_(has_attention_bias), tile_size_(tile_size) { + } + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"head_size_vec", ProgramUniformVariableDataType::Uint32}, + {"total_sequence_length", ProgramUniformVariableDataType::Uint32}, + {"alpha", ProgramUniformVariableDataType::Float32}, + {"present_sequence_length", ProgramUniformVariableDataType::Uint32}, + {"n_reps", ProgramUniformVariableDataType::Uint32}, + {"num_total_seq_length_tile", ProgramUniformVariableDataType::Uint32}); + + private: + bool has_attention_bias_; + uint32_t tile_size_; +}; + +class FlashAttentionDecodeSplitVxProgram final : public Program { + public: + FlashAttentionDecodeSplitVxProgram(const std::string& kernel_name, uint32_t tile_size, int head_size_vec) + : Program{kernel_name}, tile_size_(tile_size), head_size_vec_(head_size_vec) { + } + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"total_sequence_length", ProgramUniformVariableDataType::Uint32}, + {"head_size_vec", ProgramUniformVariableDataType::Uint32}, + {"present_sequence_length", ProgramUniformVariableDataType::Uint32}, + {"n_reps", ProgramUniformVariableDataType::Uint32}, + {"num_total_seq_length_tile", ProgramUniformVariableDataType::Uint32}); + + private: + uint32_t tile_size_; + int head_size_vec_; +}; + +class FlashAttentionDecodeVxReduceProgram final : public Program { + public: + FlashAttentionDecodeVxReduceProgram(const std::string& kernel_name, uint32_t tile_size) + : Program{kernel_name}, tile_size_(tile_size) { + } + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"head_size_vec", ProgramUniformVariableDataType::Uint32}, + {"num_total_seq_length_tile", ProgramUniformVariableDataType::Uint32}, + {"num_head_size_tile", ProgramUniformVariableDataType::Uint32}); + + private: + uint32_t tile_size_; +}; + Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, const Tensor* attention_bias, Tensor* output, const Tensor* past_key, Tensor* present_key, const Tensor* past_value, Tensor* present_value, const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context); diff --git a/onnxruntime/contrib_ops/webgpu/bert/skip_layer_norm.cc b/onnxruntime/contrib_ops/webgpu/bert/skip_layer_norm.cc index 61f701f7911a7..2126022f8b547 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/skip_layer_norm.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/skip_layer_norm.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/string_macros.h" #include "core/providers/webgpu/webgpu_utils.h" #include "core/providers/webgpu/webgpu_supported_types.h" #include "contrib_ops/webgpu/webgpu_contrib_kernels.h" @@ -12,7 +13,7 @@ namespace contrib { namespace webgpu { Status SkipLayerNormProgram::GenerateShaderCode(ShaderHelper& shader) const { - const auto& x = shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); + const auto& x = shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); shader.AddInput("skip", ShaderUsage::UseUniform); shader.AddInput("gamma", ShaderUsage::UseUniform); if (hasBeta_) { @@ -26,57 +27,112 @@ Status SkipLayerNormProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddOutput("input_skip_bias_sum", ShaderUsage::UseUniform); } - int components = x.NumComponents(); - - std::string bias = (hasBias_) ? " + bias[offset1d + i] " : ""; std::string simpl1 = (simplified_) ? "" : "- mean * mean "; - std::string simpl2 = (simplified_) ? "" : "- element_t(mean) "; - std::string beta = (hasBeta_) ? " + beta[offset1d + i] " : ""; - std::string input_skip_bias_sum = (has_input_skip_bias_sum_) ? "input_skip_bias_sum[offset + i] = value;\n" : ""; - - shader.AdditionalImplementation() - << "alias element_t = " << (is_fp16_ ? "f16;\n" : "f32;\n") - << "alias f32_val_t = " << (components == 4 ? "vec4" : (components == 2 ? "vec2" : "f32")) << ";\n" - << "var sum_shared : array;\n" - << "var sum_squared_shared : array;\n"; - - shader.MainFunctionBody() - << "let ix = local_idx;\n" - << "let iy = global_idx / workgroup_size_x;\n" - << "let hidden_size_vectorized: u32 = uniforms.hidden_size / uniforms.components;\n" - << "var stride = hidden_size_vectorized / workgroup_size_x;\n" - << "let offset = ix * stride + iy * hidden_size_vectorized;\n" - << "let offset1d = stride * ix;\n" - << "if (ix == workgroup_size_x - 1) {\n" - << " stride = hidden_size_vectorized - stride * ix;\n" - << "}\n" - << "for (var i: u32 = 0; i < stride; i++) {\n" - << " let skip_value = skip[offset + i];\n" - << " let input_value = x[offset + i];\n" - << " let value = input_value + skip_value" << bias << ";\n" - << " output[offset + i] = value;\n" - << input_skip_bias_sum - << " let f32_value = f32_val_t(value);\n" - << " sum_shared[ix] += f32_value;\n" - << " sum_squared_shared[ix] += f32_value * f32_value;\n" - << "}\n" - << "workgroupBarrier();\n" - << "var reduce_size : u32 = workgroup_size_x;\n" - << "for (var curr_size = reduce_size >> 1; curr_size > 0; curr_size = reduce_size >> 1) {\n" - << " reduce_size = curr_size + (reduce_size & 1);\n" - << " if (ix < curr_size) {\n" - << " sum_shared[ix] += sum_shared[ix + reduce_size];\n" - << " sum_squared_shared[ix] += sum_squared_shared[ix + reduce_size];\n" - << " }\n" - << " workgroupBarrier();\n" - << "}\n" - << "let sum = sum_shared[0];\n" - << "let square_sum = sum_squared_shared[0];\n" - << "let mean = " << SumVector("sum", components) << " / f32(uniforms.hidden_size);\n" - << "let inv_std_dev = inverseSqrt(" << SumVector("square_sum", components) << " / f32(uniforms.hidden_size) " << simpl1 << "+ uniforms.epsilon);\n" - << "for (var i: u32 = 0; i < stride; i++) {\n" - << " output[offset + i] = (output[offset + i] " << simpl2 << ") * element_t(inv_std_dev) * gamma[offset1d + i]" << beta << ";\n" - << "};\n"; + std::string simpl2 = (simplified_) ? "" : "- x_element_t(mean) "; + if (split_hidden_dim_) { + shader.AdditionalImplementation() + << "var sum_shared : array;\n" + << "var sum_squared_shared : array;\n"; + + SS(input_skip_bias_sum_ss, 512); + if (has_input_skip_bias_sum_) { + input_skip_bias_sum_ss + << " let workgroup_half_idx = uniforms.hidden_size / (workgroup_size_x * 4);\n" + << " if (workgroup_idx >= workgroup_half_idx) {\n" + << " offset = (workgroup_idx - workgroup_half_idx) * workgroup_size_x + local_idx;\n" + << " let skip_value = skip[offset];\n" + << " let input_value = x[offset];\n" + << " let value = input_value + skip_value" << (hasBias_ ? " + bias[offset]" : "") << ";\n" + << " input_skip_bias_sum[offset] = value;\n" + << " return;\n" + << " }\n"; + } + + shader.MainFunctionBody() + << " var offset: u32 = 0;\n" + << (has_input_skip_bias_sum_ ? SS_GET(input_skip_bias_sum_ss) : "") + << " var sum_vec4 = vec4(0);\n" + << " var sum_squared_vec4 = vec4(0);\n" + << " var cur_input_skip_bias_sum = x_value_t(0);\n" + << " for (var i: u32 = 0; i < uniforms.hidden_size / (workgroup_size_x * 4); i++) {\n" + << " let input_offset = i * workgroup_size_x + local_idx;\n" + << " let skip_value = skip[input_offset];\n" + << " let input_value = x[input_offset];\n" + << " let value = input_value + skip_value" << (hasBias_ ? " + bias[input_offset]" : "") << ";\n" + << " if (i == workgroup_idx) {\n" + << " cur_input_skip_bias_sum = value;\n" + << " }\n" + << " let f32_value = vec4(value);\n" + << " sum_vec4 += f32_value;\n" + << " sum_squared_vec4 += f32_value * f32_value;\n" + << " }\n" + << " var sum = " << SumVector("sum_vec4", 4) << ";\n" + << " var sum_squared = " << SumVector("sum_squared_vec4", 4) << ";\n" + << " sum_shared[local_idx] = sum;\n" + << " sum_squared_shared[local_idx] = sum_squared;\n" + << " workgroupBarrier();\n" + << " var reduce_size : u32 = workgroup_size_x;\n" + << " for (var curr_size = reduce_size >> 1; curr_size > 0; curr_size = reduce_size >> 1) {\n" + << " reduce_size = curr_size + (reduce_size & 1);\n" + << " if (local_idx < curr_size) {\n" + << " sum_shared[local_idx] += sum_shared[local_idx + reduce_size];\n" + << " sum_squared_shared[local_idx] += sum_squared_shared[local_idx + reduce_size];\n" + << " }\n" + << " workgroupBarrier();\n" + << " }\n" + << " let mean = sum_shared[0] / f32(uniforms.hidden_size);\n" + << " let inv_std_dev = inverseSqrt(sum_squared_shared[0] / f32(uniforms.hidden_size) " << simpl1 << "+ uniforms.epsilon);\n" + << " offset = workgroup_idx * workgroup_size_x + local_idx;\n" + << " output[offset] = ((cur_input_skip_bias_sum " << simpl2 << ") * x_element_t(inv_std_dev) * gamma[offset]" << (hasBeta_ ? " + beta[offset] " : "") << ");\n"; + } else { + int components = x.NumComponents(); + std::string bias = (hasBias_) ? " + bias[offset1d + i] " : ""; + std::string beta = (hasBeta_) ? " + beta[offset1d + i] " : ""; + std::string input_skip_bias_sum = (has_input_skip_bias_sum_) ? "input_skip_bias_sum[offset + i] = value;\n" : ""; + + shader.AdditionalImplementation() + << "alias f32_val_t = " << (components == 4 ? "vec4" : (components == 2 ? "vec2" : "f32")) << ";\n" + << "var sum_shared : array;\n" + << "var sum_squared_shared : array;\n"; + + shader.MainFunctionBody() + << "let ix = local_idx;\n" + << "let iy = global_idx / workgroup_size_x;\n" + << "let hidden_size_vectorized: u32 = uniforms.hidden_size / uniforms.components;\n" + << "var stride = hidden_size_vectorized / workgroup_size_x;\n" + << "let offset = ix * stride + iy * hidden_size_vectorized;\n" + << "let offset1d = stride * ix;\n" + << "if (ix == workgroup_size_x - 1) {\n" + << " stride = hidden_size_vectorized - stride * ix;\n" + << "}\n" + << "for (var i: u32 = 0; i < stride; i++) {\n" + << " let skip_value = skip[offset + i];\n" + << " let input_value = x[offset + i];\n" + << " let value = input_value + skip_value" << bias << ";\n" + << " output[offset + i] = value;\n" + << input_skip_bias_sum + << " let f32_value = f32_val_t(value);\n" + << " sum_shared[ix] += f32_value;\n" + << " sum_squared_shared[ix] += f32_value * f32_value;\n" + << "}\n" + << "workgroupBarrier();\n" + << "var reduce_size : u32 = workgroup_size_x;\n" + << "for (var curr_size = reduce_size >> 1; curr_size > 0; curr_size = reduce_size >> 1) {\n" + << " reduce_size = curr_size + (reduce_size & 1);\n" + << " if (ix < curr_size) {\n" + << " sum_shared[ix] += sum_shared[ix + reduce_size];\n" + << " sum_squared_shared[ix] += sum_squared_shared[ix + reduce_size];\n" + << " }\n" + << " workgroupBarrier();\n" + << "}\n" + << "let sum = sum_shared[0];\n" + << "let square_sum = sum_squared_shared[0];\n" + << "let mean = " << SumVector("sum", components) << " / f32(uniforms.hidden_size);\n" + << "let inv_std_dev = inverseSqrt(" << SumVector("square_sum", components) << " / f32(uniforms.hidden_size) " << simpl1 << "+ uniforms.epsilon);\n" + << "for (var i: u32 = 0; i < stride; i++) {\n" + << " output[offset + i] = (output[offset + i] " << simpl2 << ") * x_element_t(inv_std_dev) * gamma[offset1d + i]" << beta << ";\n" + << "};\n"; + } return Status::OK(); } @@ -100,14 +156,15 @@ Status SkipLayerNorm::ComputeInternal(onnxruntime::webgpu::ComputeCo return Status::OK(); } - const bool is_fp16 = x->GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16; const uint32_t hidden_size = onnxruntime::narrow(x_shape[x_shape.NumDimensions() - 1]); const int components = GetMaxComponents(hidden_size); const bool has_input_skip_bias_sum = input_skip_bias_sum != nullptr; + const uint32_t norm_count = onnxruntime::narrow(x_shape.SizeToDimension(x_shape.NumDimensions() - 1)); + const bool split_hidden_dim = hidden_size % 512 == 0 && norm_count == 1; - SkipLayerNormProgram program{beta != nullptr, bias != nullptr, epsilon_, hidden_size, has_input_skip_bias_sum, is_fp16, simplified}; + SkipLayerNormProgram program{beta != nullptr, bias != nullptr, epsilon_, hidden_size, has_input_skip_bias_sum, simplified, split_hidden_dim}; program - .CacheHint(simplified, has_input_skip_bias_sum) + .CacheHint(simplified, has_input_skip_bias_sum, split_hidden_dim) .AddInputs({{x, ProgramTensorMetadataDependency::Type, components}}) .AddInputs({{skip, ProgramTensorMetadataDependency::Type, components}}) .AddInputs({{gamma, ProgramTensorMetadataDependency::Type, components}}) @@ -123,6 +180,13 @@ Status SkipLayerNorm::ComputeInternal(onnxruntime::webgpu::ComputeCo {static_cast(epsilon_)}, }); + if (split_hidden_dim) { + const uint32_t workgroup_size_x = 128; + const uint32_t dispatch_size_x = (has_input_skip_bias_sum ? 2 : 1) * hidden_size / (workgroup_size_x * components); + program.SetDispatchGroupSize(dispatch_size_x, 1, 1) + .SetWorkgroupSize(workgroup_size_x); + } + if (beta != nullptr) { program.AddInput({beta, ProgramTensorMetadataDependency::Type, components}); } diff --git a/onnxruntime/contrib_ops/webgpu/bert/skip_layer_norm.h b/onnxruntime/contrib_ops/webgpu/bert/skip_layer_norm.h index 03de1a4b568b9..73f02f0ad8ec0 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/skip_layer_norm.h +++ b/onnxruntime/contrib_ops/webgpu/bert/skip_layer_norm.h @@ -15,7 +15,7 @@ using onnxruntime::webgpu::ComputeContext; class SkipLayerNormProgram final : public Program { public: - SkipLayerNormProgram(bool hasBeta, bool hasBias, float epsilon, uint32_t hidden_size, bool has_input_skip_bias_sum, bool is_fp16, bool simplified) : Program{"SkipLayerNorm"} { + SkipLayerNormProgram(bool hasBeta, bool hasBias, float epsilon, uint32_t hidden_size, bool has_input_skip_bias_sum, bool simplified, bool split_hidden_dim) : Program{"SkipLayerNorm"} { epsilon_ = epsilon; hasBeta_ = hasBeta; hasBias_ = hasBias; @@ -23,7 +23,7 @@ class SkipLayerNormProgram final : public Program { hidden_size_ = hidden_size; has_input_skip_bias_sum_ = has_input_skip_bias_sum; simplified_ = simplified; - is_fp16_ = is_fp16; + split_hidden_dim_ = split_hidden_dim; } Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -39,8 +39,8 @@ class SkipLayerNormProgram final : public Program { float epsilon_; uint32_t hidden_size_; bool has_input_skip_bias_sum_; - bool is_fp16_; bool simplified_; + bool split_hidden_dim_; }; template diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 7532f8eeb5a9b..4782e479753a2 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -1455,6 +1455,9 @@ MlasConvDepthwiseFloat_CHW( #endif #elif defined(MLAS_TARGET_WASM_SIMD) #define MLAS_WASM_SIMD_INTRINSICS +#if defined(MLAS_TARGET_WASM_RELAXED_SIMD) +#define MLAS_WASM_RELAXED_SIMD_INTRINSICS +#endif #elif defined(MLAS_TARGET_LARCH64) #define MLAS_LSX_INTRINSICS #endif @@ -2265,6 +2268,8 @@ MlasMaximumFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2) #elif defined(MLAS_VSX_INTRINSICS) // Don't use vec_max to avoid undefined behavior if NAN return vec_sel(Vector2, Vector1, vec_cmpgt(Vector1, Vector2)); +#elif defined(MLAS_WASM_RELAXED_SIMD_INTRINSICS) + return wasm_f32x4_relaxed_max(Vector1, Vector2); #elif defined(MLAS_WASM_SIMD_INTRINSICS) return wasm_f32x4_max(Vector1, Vector2); #elif defined(MLAS_LSX_INTRINSICS) @@ -2285,6 +2290,8 @@ MlasMinimumFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2) #elif defined(MLAS_VSX_INTRINSICS) // Don't use vec_min to avoid undefined behavior if NAN return vec_sel(Vector2, Vector1, vec_cmpgt(Vector2, Vector1)); +#elif defined(MLAS_WASM_RELAXED_SIMD_INTRINSICS) + return wasm_f32x4_relaxed_min(Vector1, Vector2); #elif defined(MLAS_WASM_SIMD_INTRINSICS) return wasm_f32x4_min(Vector1, Vector2); #elif defined(MLAS_LSX_INTRINSICS) diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc index bdc7297b08fe8..02d2bf22b8144 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.cc @@ -31,10 +31,9 @@ Status BaseOpBuilder::AddToModelBuilder(QnnModelWrapper& qnn_model_wrapper, // Inputs & output handling mostly same for most of the Ops, just node attributes are different ORT_RETURN_IF_ERROR(ProcessInputs(qnn_model_wrapper, node_unit, logger, input_names, do_op_validation)); - + ORT_RETURN_IF_ERROR(ProcessInt64Tensors(qnn_model_wrapper, node_unit, input_names)); ORT_RETURN_IF_ERROR(ProcessAttributesAndOutputs(qnn_model_wrapper, node_unit, std::move(input_names), logger, do_op_validation)); - return Status::OK(); } @@ -132,6 +131,48 @@ Status BaseOpBuilder::AddZeroBiasInput(QnnModelWrapper& qnn_model_wrapper, return Status::OK(); } +Status BaseOpBuilder::ProcessInt64Tensors(QnnModelWrapper& qnn_model_wrapper, + const NodeUnit& node_unit, + std::vector& input_names) const { + if (input_names.size() < 1) { + return Status::OK(); + } + for (size_t i = 0; i < input_names.size(); i++) { + auto& input_tensorwrapper = qnn_model_wrapper.GetQnnTensorWrapper(input_names[i]); + // Insert cast to int32 if input dtype is int64 + if (input_tensorwrapper.GetTensorDataType() == QNN_DATATYPE_INT_64) { + const Qnn_TensorType_t tensor_type = QNN_TENSOR_TYPE_NATIVE; + const std::string cast_output_name = input_names[i] + "_cast_int32"; + if (!qnn_model_wrapper.IsQnnTensorWrapperExist(cast_output_name)) { + Qnn_DataType_t qnn_data_type = QNN_DATATYPE_INT_32; + const auto& input_i = node_unit.Inputs()[i]; + std::vector output_shape; + ORT_RETURN_IF_NOT(qnn_model_wrapper.GetOnnxShape(input_i.node_arg, output_shape), + "QNN EP: Cannot get input shape for ", input_i.node_arg.Name().c_str()); + QnnTensorWrapper output_tensorwrapper(cast_output_name, + tensor_type, + qnn_data_type, + QnnQuantParamsWrapper(), + std::move(output_shape)); + ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(output_tensorwrapper)), + "Failed to add output tensor for QNN Cast node."); + + ORT_RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(cast_output_name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + QNN_OP_CAST, + {input_names[i]}, + {cast_output_name}, + {}, + false), + "Failed to create QNN Cast node."); + } + + input_names[i] = cast_output_name; + } + } + return Status::OK(); +} + Status BaseOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wrapper, const NodeUnit& node_unit, std::vector&& input_names, @@ -140,7 +181,6 @@ Status BaseOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wra if (input_names.size() < 1) { return Status::OK(); } - ORT_RETURN_IF_ERROR(ProcessOutputs(qnn_model_wrapper, node_unit, std::move(input_names), {}, logger, do_op_validation, GetQnnOpType(node_unit.OpType()))); return Status::OK(); @@ -179,7 +219,35 @@ Status BaseOpBuilder::ProcessOutputs(QnnModelWrapper& qnn_model_wrapper, Qnn_DataType_t supported_qnn_data_type = GetSupportedOutputDataType(output_i, output_info.qnn_data_type); bool is_graph_output = qnn_model_wrapper.IsGraphOutput(output_name); - if (supported_qnn_data_type != output_info.qnn_data_type && is_graph_output && !do_op_validation) { + + // Check if we need to add a cast node for int64 + bool needs_int64_cast = false; + if (is_graph_output) { + for (const auto& input_name : input_names) { + if (input_name.find("_cast_int32") != std::string::npos) { + needs_int64_cast = true; + break; + } + } + } + + if (needs_int64_cast) { + std::string cast_node_name = output_name + "_cast_int64"; + std::string cast_input_name = output_name + "_cast_int64_aux"; + QnnQuantParamsWrapper quant_params = output_info.quant_param.Copy(); + std::vector cast_output_shape = output_info.shape; + + // Create the cast input tensor wrapper + QnnTensorWrapper cast_input_tensorwrapper(cast_input_name, + QNN_TENSOR_TYPE_NATIVE, + supported_qnn_data_type, + output_info.quant_param.Copy(), + std::move(cast_output_shape)); + ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cast_input_tensorwrapper)), "Failed to add tensor."); + output_names.push_back(cast_input_name); + // Store the cast node information for later addition + cast_node_info_vec.push_back({cast_node_name, cast_input_name, output_name}); + } else if (supported_qnn_data_type != output_info.qnn_data_type && is_graph_output && !do_op_validation) { std::string cast_node_name = output_name + "_ort_qnn_ep_cast"; std::string cast_input_name = output_name + "_ort_qnn_ep_aux"; std::vector cast_output_shape = output_info.shape; diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h index df9d0de8e0e3e..37060fcd9ba93 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/base_op_builder.h @@ -76,6 +76,10 @@ class BaseOpBuilder : public IOpBuilder { std::vector& input_names, bool do_op_validation = false) const ORT_MUST_USE_RESULT; + Status ProcessInt64Tensors(QnnModelWrapper& qnn_model_wrapper, + const NodeUnit& node_unit, + std::vector& input_names) const ORT_MUST_USE_RESULT; + virtual Status ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wrapper, const NodeUnit& node_unit, std::vector&& input_names, diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/cast_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/cast_op_builder.cc index b7f056ae65260..27b0c86827531 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/cast_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/cast_op_builder.cc @@ -102,6 +102,9 @@ Status CastOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wra const bool is_graph_output = qnn_model_wrapper.IsGraphOutput(output_name); const Qnn_TensorType_t tensor_type = is_graph_output ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE; + if (qnn_data_type == QNN_DATATYPE_INT_64 && tensor_type == QNN_TENSOR_TYPE_NATIVE) { + qnn_data_type = QNN_DATATYPE_INT_32; + } QnnTensorWrapper output_tensorwrapper(output_name, tensor_type, qnn_data_type, diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/gather_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/gather_op_builder.cc index d25ec3f333bf1..f2992196f7811 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/gather_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/gather_op_builder.cc @@ -257,7 +257,57 @@ Status GatherOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_w bool is_graph_output = qnn_model_wrapper.IsGraphOutput(output_name); bool reshape_required = (qnn_output_shape.size() != target_output_shape.size()); - std::string gather_output_name = output_name + (reshape_required ? "_ort_qnn_ep_reshape" : ""); + + struct CastNodeInfo { + std::string node_name; + std::string input_name; + std::string output_name; + }; + std::vector cast_node_info_vec; + + // Get the output info for the gather output tensor + TensorInfo output_info = {}; + ORT_RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(gather_output, output_info)); + + // Check if we need to add a cast node for int64 + bool needs_int64_cast = false; + if (is_graph_output) { + for (const auto& input_name : input_names) { + if (input_name.find("_cast_int32") != std::string::npos) { + needs_int64_cast = true; + break; + } + } + } + + // If a cast to int64 is needed, add the cast node + if (needs_int64_cast) { + std::string cast_node_name = output_name + "_cast_int64"; + std::string cast_input_name = output_name + "_cast_int64_aux"; + std::string cast_output_name = output_name; + + // Create the cast input tensor wrapper + QnnTensorWrapper cast_input_tensorwrapper(cast_input_name, + QNN_TENSOR_TYPE_NATIVE, + output_info.qnn_data_type, + output_info.quant_param.Copy(), + std::move(qnn_output_shape)); + + ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cast_input_tensorwrapper)), "Failed to add tensor."); + cast_node_info_vec.push_back({cast_node_name, cast_input_name, cast_output_name}); + Qnn_TensorType_t cast_tensor_type = is_graph_output ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE; + QnnTensorWrapper cast_output(output_name, cast_tensor_type, qnn_data_type, std::move(quantize_param), + std::move(target_output_shape)); + ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cast_output)), "Failed to add tensor."); + } + + std::string gather_output_name = output_name; + if (reshape_required) { + gather_output_name += "_ort_qnn_ep_reshape"; + } else if (needs_int64_cast) { + gather_output_name += "_cast_int64_aux"; + } + Qnn_TensorType_t tensor_type = (!reshape_required && is_graph_output) ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE; QnnTensorWrapper gather_output_wrapper(gather_output_name, tensor_type, qnn_data_type, quantize_param.Copy(), std::move(qnn_output_shape)); @@ -279,16 +329,35 @@ Status GatherOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_w std::move(target_output_shape)); ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(reshape_output)), "Failed to add tensor."); const static std::string qnn_node_type = "Reshape"; + std::string node_output_name = output_name; + + if (needs_int64_cast) { + // If needs_int64 is true, the output name should be the input name of the cast node + node_output_name = output_name + "_cast_int64_aux"; + } ORT_RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(output_name, QNN_OP_PACKAGE_NAME_QTI_AISW, qnn_node_type, {gather_output_name}, - {output_name}, + {node_output_name}, {}, do_op_validation), "Failed to add node."); } + if (needs_int64_cast) { + for (const auto& cast_node_info : cast_node_info_vec) { + // Insert cast node. + ORT_RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(cast_node_info.node_name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + "Cast", + {cast_node_info.input_name}, + {cast_node_info.output_name}, + {}), + " Failed to add Cast node"); + } + } + return Status::OK(); } diff --git a/onnxruntime/core/providers/qnn/builder/opbuilder/transpose_op_builder.cc b/onnxruntime/core/providers/qnn/builder/opbuilder/transpose_op_builder.cc index bcd8a6d0f78f6..5b26c20cef825 100644 --- a/onnxruntime/core/providers/qnn/builder/opbuilder/transpose_op_builder.cc +++ b/onnxruntime/core/providers/qnn/builder/opbuilder/transpose_op_builder.cc @@ -81,12 +81,51 @@ Status TransposeOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_mode bool is_graph_output = qnn_model_wrapper.IsGraphOutput(output_name); Qnn_TensorType_t tensor_type = is_graph_output ? QNN_TENSOR_TYPE_APP_READ : QNN_TENSOR_TYPE_NATIVE; + struct CastNodeInfo { + std::string node_name; + std::string input_name; + std::string output_name; + }; + std::vector cast_node_info_vec; + + // Check if we need to add a cast node for int64 + bool needs_int64_cast = false; + if (is_graph_output) { + for (const auto& input_name : input_names) { + if (input_name.find("_cast_int32") != std::string::npos) { + needs_int64_cast = true; + break; + } + } + } + + const auto& transpose_output = node_unit.Outputs()[0]; + // Get the output info for the gather output tensor + TensorInfo output_info = {}; + ORT_RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(transpose_output, output_info)); std::vector output_shape; ORT_RETURN_IF_NOT(qnn_model_wrapper.GetOnnxShape(node_unit.Outputs()[0].node_arg, output_shape), "Cannot get shape"); const QnnTensorWrapper& input_tensor_wrapper = qnn_model_wrapper.GetQnnTensorWrapper(input_names[0]); + // If a cast to int64 is needed, add the cast node + if (needs_int64_cast) { + std::string cast_node_name = output_name + "_cast_int64"; + std::string cast_input_name = output_name + "_cast_int64_aux"; + std::string cast_output_name = output_name; + + // Create the cast input tensor wrapper + QnnTensorWrapper cast_input_tensorwrapper(cast_input_name, + QNN_TENSOR_TYPE_NATIVE, + output_info.qnn_data_type, + output_info.quant_param.Copy(), + std::move(output_shape)); + + ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cast_input_tensorwrapper)), "Failed to add tensor."); + cast_node_info_vec.push_back({cast_node_name, cast_input_name, cast_output_name}); + } + // Transpose output uses same data type and quantization parameter with input // 1. In QDQ model, the optimization may create scenario like Q -> Transpose -> DQ, Transpose is single node // Input tensor is created by previous node which is quantized tensor, @@ -110,6 +149,18 @@ Status TransposeOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_mode do_op_validation), "Failed to add node."); + if (needs_int64_cast) { + for (const auto& cast_node_info : cast_node_info_vec) { + // Insert cast node. + ORT_RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(cast_node_info.node_name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + "Cast", + {cast_node_info.input_name}, + {cast_node_info.output_name}, + {}), + " Failed to add Cast node"); + } + } return Status::OK(); } diff --git a/onnxruntime/core/providers/qnn/builder/qnn_def.h b/onnxruntime/core/providers/qnn/builder/qnn_def.h index 4088670bc6217..6351779b23f99 100644 --- a/onnxruntime/core/providers/qnn/builder/qnn_def.h +++ b/onnxruntime/core/providers/qnn/builder/qnn_def.h @@ -155,6 +155,24 @@ class QnnTensorWrapper { dimensions_(std::move(shape)), client_buf_(std::move(client_buf)), quant_params_(quantize_params) { + if (data_type == QNN_DATATYPE_INT_64) { + // QNN doesn't support int64_t, so we cast to int32_t. + if (tensor_type == QNN_TENSOR_TYPE_NATIVE) { + data_type = QNN_DATATYPE_INT_32; + } + if (client_buf_.size()) { + const size_t num_elems = client_buf_.size() / sizeof(int64_t); + std::vector cast_data; + cast_data.resize(num_elems * sizeof(int32_t)); + gsl::span origin_values{reinterpret_cast(client_buf_.data()), num_elems}; + gsl::span new_values(reinterpret_cast(cast_data.data()), num_elems); + for (size_t i = 0; i < num_elems; i++) { + new_values[i] = static_cast(origin_values[i]); + } + data_type = QNN_DATATYPE_INT_32; + client_buf_ = std::move(cast_data); + } + } SetQnnTensorType(qnn_tensor_, tensor_type); SetQnnTensorName(qnn_tensor_, tensor_name_.c_str()); SetQnnTensorDataType(qnn_tensor_, data_type); diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 2c443cdc33458..ded135bf50ec8 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -2296,8 +2296,9 @@ SubGraphCollection_t TensorrtExecutionProvider::GetSupportedList(SubGraphCollect auto network_flags = 0; #if NV_TENSORRT_MAJOR > 8 network_flags |= fp16_enable_ || int8_enable_ ? 0 : 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kSTRONGLY_TYPED); -#endif +#else network_flags |= 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); +#endif auto trt_network = std::unique_ptr(trt_builder->createNetworkV2(network_flags)); auto trt_parser = tensorrt_ptr::unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); @@ -2907,8 +2908,9 @@ Status TensorrtExecutionProvider::CreateNodeComputeInfoFromGraph(const GraphView auto network_flags = 0; #if NV_TENSORRT_MAJOR > 8 network_flags |= fp16_enable_ || int8_enable_ ? 0 : 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kSTRONGLY_TYPED); -#endif +#else network_flags |= 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); +#endif auto trt_network = std::unique_ptr(trt_builder->createNetworkV2(network_flags)); auto trt_config = std::unique_ptr(trt_builder->createBuilderConfig()); auto trt_parser = tensorrt_ptr::unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); diff --git a/onnxruntime/core/providers/webgpu/nn/activation_util.h b/onnxruntime/core/providers/webgpu/nn/activation_util.h index 1c9fd93e35384..3a0f5d4900a17 100644 --- a/onnxruntime/core/providers/webgpu/nn/activation_util.h +++ b/onnxruntime/core/providers/webgpu/nn/activation_util.h @@ -4,6 +4,7 @@ #pragma once #include +#include namespace onnxruntime { namespace webgpu { @@ -12,4 +13,4 @@ extern std::string TypeSnippet(uint32_t component, std::string data_type); extern std::string BiasSnippet(bool has_bias); } // namespace webgpu -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/program.cc b/onnxruntime/core/providers/webgpu/program.cc index 73291e1e93ff1..89f547481b6e4 100644 --- a/onnxruntime/core/providers/webgpu/program.cc +++ b/onnxruntime/core/providers/webgpu/program.cc @@ -102,6 +102,9 @@ constexpr std::string_view ProgramVariableDataTypeName[] = { "u8x4", // Uint8x4 "u8x8", // Uint8x8 "u8x16", // Uint8x16 + "i8x4", // Int8x4 + "i8x8", // Int8x8 + "i8x16", // Int8x16 }; std::ostream& operator<<(std::ostream& os, ProgramVariableDataType type) { os << ProgramVariableDataTypeName[std::underlying_type::type(type)]; @@ -129,6 +132,7 @@ int NumberOfComponents(ProgramVariableDataType type) { case ProgramVariableDataType::Float16x4: case ProgramVariableDataType::Boolx4: case ProgramVariableDataType::Uint8x4: + case ProgramVariableDataType::Int8x4: return 4; case ProgramVariableDataType::Uint8x8: return 8; @@ -142,6 +146,10 @@ int NumberOfComponents(ProgramVariableDataType type) { ProgramVariableDataType ToProgramVariableDataType(int32_t element_type, int component /* = 1 */) { if (component == 1) { switch (element_type) { + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: + return ProgramVariableDataType::Uint8x4; // shader needs to be aware that only 1 value is valid + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: + return ProgramVariableDataType::Int8x4; // shader needs to be aware that only 1 value is valid case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: return ProgramVariableDataType::Float32; case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: @@ -174,6 +182,8 @@ ProgramVariableDataType ToProgramVariableDataType(int32_t element_type, int comp switch (element_type) { case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: return ProgramVariableDataType::Uint8x4; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: + return ProgramVariableDataType::Int8x4; case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: return ProgramVariableDataType::Float32x4; case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: diff --git a/onnxruntime/core/providers/webgpu/program.h b/onnxruntime/core/providers/webgpu/program.h index ea7d8ae5471af..3b0acfa7d0d35 100644 --- a/onnxruntime/core/providers/webgpu/program.h +++ b/onnxruntime/core/providers/webgpu/program.h @@ -197,7 +197,10 @@ enum class ProgramVariableDataType { Boolx4, Uint8x4, Uint8x8, - Uint8x16 + Uint8x16, + Int8x4, + Int8x8, + Int8x16, }; #ifndef NDEBUG std::ostream& operator<<(std::ostream& os, ProgramVariableDataType); diff --git a/onnxruntime/core/providers/webgpu/quantization/quantize_linear.cc b/onnxruntime/core/providers/webgpu/quantization/quantize_linear.cc new file mode 100644 index 0000000000000..866b1debf6dc8 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/quantization/quantize_linear.cc @@ -0,0 +1,221 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include "core/util/math.h" +#include "core/providers/webgpu/quantization/quantize_linear.h" +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" +#include "core/providers/webgpu/webgpu_utils.h" + +namespace onnxruntime { +namespace webgpu { + +Status DequantizeLinearProgram::GenerateShaderCode(ShaderHelper& shader) const { + const auto& x = shader.AddInput("input", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseElementTypeAlias); + const auto& scale = shader.AddInput("scale", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias); + const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseShapeAndStride | ShaderUsage::UseValueTypeAlias); + + shader.MainFunctionBody() + << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size") + << "let output_indices = " << output.OffsetToIndices("global_idx") << ";\n"; + + // Get x input + if (packed_) { + std::string unpack = (signed_) ? "unpack4xI8(x)" : "unpack4xU8(x)"; + if (output.NumComponents() == 1) { + shader.MainFunctionBody() + << "let x = " << x.GetByOffset("global_idx / 4") << ";\n" + << "let x_vec = " << unpack << ";\n" + << "let x_value = x_vec[global_idx % 4];\n"; + } else { + shader.MainFunctionBody() + << "let x = " << x.GetByOffset("global_idx") << ";\n" + << "let x_vec = " << unpack << ";\n" + << "let x_value = x_vec;\n"; + } + } else { + shader.MainFunctionBody() + << "let x_value = " << x.GetByOffset("global_idx") << ";\n"; + } + + // Get scaler + if (per_layer_) { + // scale input is a scalar () + shader.MainFunctionBody() + << "let scale_value = " << scale.GetByOffset("0") << ";\n"; + } else if (per_axis_) { + shader.MainFunctionBody() + << "let scale_index = " << output.IndicesGet("output_indices", "uniforms.axis") << ";\n" + << "let scale_value = " << scale.GetByOffset("scale_index") << ";\n"; + } else { + // Block quantization. Scale input rank is same as input/output rank. + shader.MainFunctionBody() + << "var scale_indices: scale_indices_t = output_indices;\n" + << "let index = " << scale.IndicesGet("scale_indices", "uniforms.axis") << "/ uniforms.block_size;\n" + << scale.IndicesSet("scale_indices", "uniforms.axis", "index") << ";\n" + << "let scale_value = " << scale.GetByIndices("scale_indices") << ";\n"; + } + + // Get zero-point + if (has_zeropoint_) { + const auto& zero_point = shader.AddInput("zero_point", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias); + + std::string unpack = (signed_) ? "unpack4xI8(zero_point_input)" : "unpack4xU8(zero_point_input)"; + if (per_layer_) { + // zero-point input is a scalar + if (packed_) { + shader.MainFunctionBody() + << "let zero_point_input = " << zero_point.GetByOffset("0") << ";\n" + << "let zero_point_vec = " << unpack << ";\n" + << "let zero_point_value = zero_point_vec[0];\n"; + } else { + shader.MainFunctionBody() + << "let zero_point_value = " << zero_point.GetByOffset("0") << ";\n"; + } + } else if (per_axis_) { + // zero-point input is a 1D tensor + if (packed_) { + shader.MainFunctionBody() + << "let zero_point_index = " << output.IndicesGet("output_indices", "uniforms.axis") << ";\n" + << "let zero_point_input = " << zero_point.GetByOffset("zero_point_index / 4") << ";\n" + << "let zero_point_vec = " << unpack << ";\n" + << "let zero_point_value = zero_point_vec[zero_point_index % 4];\n"; + } else { + shader.MainFunctionBody() + << "let zero_point_index = " << output.IndicesGet("output_indices", "uniforms.axis") << ";\n" + << "let zero_point_value = " << zero_point.GetByOffset("zero_point_index") << ";\n"; + } + } else { + // BlockedQuantization. The zero-point input shape is same as the input shape except along axis. + if (packed_) { + shader.MainFunctionBody() + << "let zero_point_offset = " << scale.GetByIndices("scale_indices") << ";\n" + << "let zero_point_input = " << zero_point.GetByOffset("zero_point_offset / 4") << ";\n" + << "let zero_point_vec = " << unpack << ";\n" + << "let zero_point_value = zero_point_vec[zero_point_offset % 4];\n"; + } else { + shader.MainFunctionBody() + << "let zero_point_value = " << zero_point.GetByIndices("scale_indices") << ";\n"; + } + } + } else { + shader.MainFunctionBody() + << "let zero_point_value = input_element_t(0);\n"; + } + + // compute and write output + shader.MainFunctionBody() + << output.SetByOffset("global_idx", "(output_value_t(x_value) - scale_value_t(zero_point_value)) * scale_value"); + + return Status::OK(); +} + +Status DequantizeLinear::ComputeInternal(ComputeContext& context) const { + const auto* x = context.Input(0); + const auto* x_scale = context.Input(1); + const auto* x_zeropoint = context.Input(2); + const auto x_shape = x->Shape(); + int64_t x_size = x_shape.Size(); + auto* output_tensor = context.Output(0, x_shape); + int64_t x_scale_rank = x_scale->Shape().NumDimensions(); + + bool packed = x->GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8 || x->GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8; + bool is_signed = x->GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8; + int64_t axis = (axis_ >= 0) ? axis_ : axis_ + x_shape.NumDimensions(); + + int max_components = GetMaxComponents(x_size); + if (max_components != 4) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "DequantizeLinear: components must be 4, but got ", max_components); + } + + // scaler - single scaler for all elements + bool per_layer = x_scale_rank == 0 || (x_scale_rank == 1 && x_scale->Shape()[0] == 1); + + // 1D tensor - 1 scaler for per axis + bool per_axis = per_layer == false && x_scale_rank == 1; + + bool use_components = per_layer && (!packed || max_components == 4); + int components = use_components ? max_components : 1; + int input_component = use_components && !packed ? max_components : 1; + + DequantizeLinearProgram program{packed, is_signed, per_layer, per_axis, x_zeropoint != nullptr}; + + program + .AddInputs({{x, ProgramTensorMetadataDependency::TypeAndRank, input_component}}) + .AddInputs({{x_scale, ProgramTensorMetadataDependency::TypeAndRank}}) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, components}) + .SetDispatchGroupSize((x_size / components + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) + .AddUniformVariables({{static_cast(axis)}}) + .AddUniformVariables({{static_cast(block_size_)}}) + .AddUniformVariables({{static_cast(x_size / components)}}) + .CacheHint(std::to_string(axis), std::to_string(is_signed), std::to_string(per_layer), std::to_string(per_axis), std::to_string(block_size_)); + + if (x_zeropoint != nullptr) { + program.AddInputs({{x_zeropoint, ProgramTensorMetadataDependency::TypeAndRank}}); + } + + return context.RunProgram(program); +} + +namespace { +const std::vector& DequantizeLinearConstraints() { + static std::vector types{ + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}; + return types; +} +} // namespace + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + DequantizeLinear, + kOnnxDomain, + 10, 12, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DequantizeLinearConstraints()), + DequantizeLinear); + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + DequantizeLinear, + kOnnxDomain, + 13, 18, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DequantizeLinearConstraints()), + DequantizeLinear); + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + DequantizeLinear, + kOnnxDomain, + 19, 20, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T1", DequantizeLinearConstraints()) + .TypeConstraint("T2", WebGpuSupportedFloatTypes()), + DequantizeLinear); + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + DequantizeLinear, + kOnnxDomain, + 21, 22, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T1", DequantizeLinearConstraints()) + .TypeConstraint("T2", WebGpuSupportedFloatTypes()), + DequantizeLinear); + +ONNX_OPERATOR_KERNEL_EX( + DequantizeLinear, + kOnnxDomain, + 23, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T1", DequantizeLinearConstraints()) + .TypeConstraint("T2", WebGpuSupportedFloatTypes()), + DequantizeLinear); + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/quantization/quantize_linear.h b/onnxruntime/core/providers/webgpu/quantization/quantize_linear.h new file mode 100644 index 0000000000000..95614998017e9 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/quantization/quantize_linear.h @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/webgpu/webgpu_kernel.h" + +namespace onnxruntime { +namespace webgpu { + +class DequantizeLinearProgram final : public Program { + public: + DequantizeLinearProgram(const bool packed, const bool issigned, const bool per_layer, + const bool per_axis, bool has_zeropoint) : Program{"DequantizeLinear"}, + packed_{packed}, + signed_{issigned}, + per_layer_{per_layer}, + per_axis_{per_axis}, + has_zeropoint_{has_zeropoint} {} + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"axis", ProgramUniformVariableDataType::Uint32}, + {"block_size", ProgramUniformVariableDataType::Uint32}, + {"output_size", ProgramUniformVariableDataType::Uint32}); + + private: + bool packed_; + bool signed_; + bool per_layer_; + bool per_axis_; + bool has_zeropoint_; +}; + +class DequantizeLinear final : public WebGpuKernel { + public: + DequantizeLinear(const OpKernelInfo& info) : WebGpuKernel(info) { + axis_ = info.GetAttrOrDefault("axis", 1); + block_size_ = info.GetAttrOrDefault("block_size", 0); + output_dtype_ = info.GetAttrOrDefault("output_dtype", 0); + } + + Status ComputeInternal(ComputeContext& context) const override; + + private: + int64_t axis_; + int64_t block_size_; + int64_t output_dtype_; +}; + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/shader_helper.cc b/onnxruntime/core/providers/webgpu/shader_helper.cc index db14cb88d1963..bac360c4c270e 100644 --- a/onnxruntime/core/providers/webgpu/shader_helper.cc +++ b/onnxruntime/core/providers/webgpu/shader_helper.cc @@ -168,6 +168,12 @@ Status ValidateVariableDataType(int32_t element_type, ProgramVariableDataType va var_type == ProgramVariableDataType::Uint8x16, "Unexpected program variable type ", int(var_type), " for uint8 tensor"); break; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: + ORT_RETURN_IF_NOT(var_type == ProgramVariableDataType::Int8x4 || + var_type == ProgramVariableDataType::Int8x8 || + var_type == ProgramVariableDataType::Int8x16, + "Unexpected program variable type ", int(var_type), " for int8 tensor"); + break; default: ORT_RETURN_IF(true, "Unsupported data type: ", element_type); // todo: add int4/uint4 diff --git a/onnxruntime/core/providers/webgpu/shader_variable.cc b/onnxruntime/core/providers/webgpu/shader_variable.cc index f8e1e0b3b8d2b..502d03c2c2dd8 100644 --- a/onnxruntime/core/providers/webgpu/shader_variable.cc +++ b/onnxruntime/core/providers/webgpu/shader_variable.cc @@ -32,6 +32,7 @@ constexpr static const std::string_view STORAGE_TYPE_ARRAY[] = { "u32", // Uint8x4 "vec2", // Uint8x8 "vec4", // Uint8x16 + "u32", // Int8x4 }; constexpr static const auto STORAGE_TYPE = details::_to_std_array(STORAGE_TYPE_ARRAY); @@ -54,6 +55,7 @@ constexpr static const std::string_view VALUE_TYPE_ARRAY[] = { "u32", // Uint8x4 (u32 as 4 elements of uint8) "vec2", // Uint8x8 (vec2 as 2x4 elements of uint8) "vec4", // Uint8x16 (vec4 as 4x4 elements of uint8) + "i32", // Int8x4 }; constexpr static const auto VALUE_TYPE = details::_to_std_array(VALUE_TYPE_ARRAY); @@ -76,6 +78,9 @@ constexpr static const std::string_view ELEMENT_TYPE_ARRAY[] = { "u32", // Uint8x4 "u32", // Uint8x8 "u32", // Uint8x16 + "i32", // Int8x4 + "i32", // Int8x8 + "i32", // Int8x16 }; constexpr static const auto ELEMENT_TYPE = details::_to_std_array(ELEMENT_TYPE_ARRAY); diff --git a/onnxruntime/core/providers/webgpu/tensor/pad.cc b/onnxruntime/core/providers/webgpu/tensor/pad.cc index f24578a145aae..5adf873be7897 100644 --- a/onnxruntime/core/providers/webgpu/tensor/pad.cc +++ b/onnxruntime/core/providers/webgpu/tensor/pad.cc @@ -5,6 +5,7 @@ #include #include "core/util/math.h" +#include "core/providers/webgpu/string_macros.h" #include "core/providers/webgpu/tensor/pad.h" #include "core/providers/webgpu/shader_helper.h" #include "core/providers/webgpu/webgpu_supported_types.h" @@ -38,38 +39,47 @@ Status PadProgram::GenerateShaderCode(ShaderHelper& shader) const { std::string lower_pads_str = GetElementAt("uniforms.lower_pads", "dim", rank); std::string data_shape_str = "i32(" + GetElementAt("uniforms.data_shape", "dim", rank) + ")"; std::string data_stride_str = rank == 1 ? "" : " * " + GetElementAt("uniforms.data_stride", "dim", rank - 1); - std::string begin_axis_statement = "in_coord = "; - std::string end_axis_statement = "in_coord = "; - std::string in_axis_statement = "in_coord = " + output_indices_str + " - " + lower_pads_str + ";\n"; + SS(axis_body_ss, 1024); switch (mode_) { case Mode::Constant: - begin_axis_statement = "use_pad_value = true;\n"; - end_axis_statement = "use_pad_value = true;\n"; + axis_body_ss << " if (" << output_indices_str << " < " << lower_pads_str << " || " << output_indices_str << " >= " << lower_pads_str << " + " << data_shape_str << ") {\n" + << " use_pad_value = true;\n"; break; case Mode::Edge: - begin_axis_statement += "0;\n"; - end_axis_statement += data_shape_str + " - 1;\n"; + axis_body_ss << " if (" << output_indices_str << " < " << lower_pads_str << ") {\n" + << " in_coord = 0;\n" + << " } else if (" << output_indices_str << " >= " << lower_pads_str << " + " << data_shape_str << ") {\n" + << " in_coord = " << data_shape_str + " - 1;\n"; break; case Mode::Reflect: - begin_axis_statement += lower_pads_str + " - " + output_indices_str + ";\n"; - end_axis_statement += data_shape_str + " - 2 - (" + output_indices_str + - " - (" + lower_pads_str + " + " + data_shape_str + "));\n"; + axis_body_ss << " if (" << output_indices_str << " < " << lower_pads_str << " || " << output_indices_str << " >= " << lower_pads_str << " + " << data_shape_str << ") {\n" + << " in_coord = " << output_indices_str << " - " << lower_pads_str << ";\n" + << " if (in_coord < 0) {\n" + << " in_coord = -in_coord;\n" + << " }\n" + << " {\n" + << " let _2n_1 = 2 * (" << data_shape_str << " - 1);\n" + << " in_coord = in_coord % _2n_1;\n" + << " if(in_coord >= " << data_shape_str << ") {\n" + << " in_coord = _2n_1 - in_coord;\n" + << " }\n" + << " }\n"; break; case Mode::Wrap: - begin_axis_statement += data_shape_str + " + " + output_indices_str + " - " + lower_pads_str + ";\n"; - end_axis_statement += output_indices_str + " - " + lower_pads_str + " - " + data_shape_str + ";\n"; + axis_body_ss << " if (" << output_indices_str << " < " << lower_pads_str << ") {\n" + << " in_coord = " << data_shape_str << " + " << output_indices_str << " - " << lower_pads_str + ";\n" + << " } else if (" << output_indices_str << " >= " << lower_pads_str << " + " << data_shape_str << ") {\n" + << " in_coord = " << output_indices_str << " - " << lower_pads_str << " - " << data_shape_str << ";\n"; break; default: return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported mode type: ", static_cast(mode_)); } + axis_body_ss << " } else {\n" + << " " << "in_coord = " << output_indices_str << " - " << lower_pads_str << ";\n" + << " }\n"; shader.MainFunctionBody() << " for (var dim = 0; dim < " << rank << " && !use_pad_value; dim++) {\n" - << " if (" << output_indices_str << " < " << lower_pads_str << ") {\n" - << " " << begin_axis_statement << " }\n" - << " else if (" << output_indices_str << " >= " << lower_pads_str << " + " << data_shape_str << ") {\n" - << " " << end_axis_statement << " }\n" - << " else {\n" - << " " << in_axis_statement << " }\n" + << SS_GET(axis_body_ss) << " input_index += select(u32(in_coord)" << data_stride_str << ", u32(in_coord), dim == " << rank - 1 << ");\n" << " }\n" << " " << constant_value_str diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 2987d3905fe54..17f4fa1bd44b3 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -409,31 +409,19 @@ Status WebGpuContext::Run(ComputeContext& context, const ProgramBase& program) { WriteTimestamp(num_pending_dispatches_ * 2); - uint32_t entry_index = 0; - std::vector bind_group_entries; + std::vector bind_buffers; + bind_buffers.reserve(inputs.size() + outputs.size() + (uniform_buffer ? 1 : 0)); for (const auto& input : inputs) { - bind_group_entries.push_back({nullptr, entry_index++, reinterpret_cast(const_cast(input.tensor->DataRaw()))}); + bind_buffers.push_back(reinterpret_cast(const_cast(input.tensor->DataRaw()))); } for (const auto& output : outputs) { - bind_group_entries.push_back({nullptr, entry_index++, reinterpret_cast(output.tensor->MutableDataRaw())}); + bind_buffers.push_back(reinterpret_cast(output.tensor->MutableDataRaw())); } if (uniform_buffer) { - bind_group_entries.push_back({nullptr, entry_index++, uniform_buffer}); + bind_buffers.push_back(uniform_buffer); } - wgpu::BindGroupDescriptor bind_group_desc{}; - bind_group_desc.layout = program_artifact->compute_pipeline.GetBindGroupLayout(0); - bind_group_desc.entryCount = bind_group_entries.size(); - bind_group_desc.entries = bind_group_entries.data(); - bind_group_desc.label = program_artifact->name.c_str(); - - auto bind_group = Device().CreateBindGroup(&bind_group_desc); - - // TODO support graph capture - - compute_pass_encoder.SetPipeline(program_artifact->compute_pipeline); - compute_pass_encoder.SetBindGroup(0, bind_group); - compute_pass_encoder.DispatchWorkgroups(x, y, z); + LaunchComputePipeline(compute_pass_encoder, bind_buffers, *program_artifact, x, y, z); if (uniform_buffer) { buffer_mgr_->Release(uniform_buffer); @@ -708,6 +696,35 @@ void WebGpuContext::OnRunEnd() { #endif // ENABLE_PIX_FOR_WEBGPU_EP } +void WebGpuContext::LaunchComputePipeline(const wgpu::ComputePassEncoder& compute_pass_encoder, + const std::vector& bind_buffers, + const ProgramArtifact& program_artifact, + uint32_t x, uint32_t y, uint32_t z) { + uint32_t entry_index = 0; + std::vector bind_group_entries; + for (WGPUBuffer buffer : bind_buffers) { + bind_group_entries.push_back({nullptr, entry_index++, buffer, 0, WGPU_WHOLE_SIZE, nullptr, nullptr}); + } + + WGPUBindGroupLayout bind_group_layout = program_artifact.compute_pipeline.GetBindGroupLayout(0).MoveToCHandle(); + WGPUBindGroupDescriptor bind_group_desc{}; + bind_group_desc.layout = bind_group_layout; + bind_group_desc.entryCount = bind_group_entries.size(); + bind_group_desc.entries = bind_group_entries.data(); + bind_group_desc.label = {program_artifact.name.data(), program_artifact.name.length()}; + + auto bind_group = wgpuDeviceCreateBindGroup(Device().Get(), &bind_group_desc); + + // TODO support graph capture + + compute_pass_encoder.SetPipeline(program_artifact.compute_pipeline); + wgpuComputePassEncoderSetBindGroup(compute_pass_encoder.Get(), 0, bind_group, 0, nullptr); + compute_pass_encoder.DispatchWorkgroups(x, y, z); + + wgpuBindGroupRelease(bind_group); + wgpuBindGroupLayoutRelease(bind_group_layout); +} + std::unordered_map WebGpuContextFactory::contexts_; std::mutex WebGpuContextFactory::mutex_; std::once_flag WebGpuContextFactory::init_default_flag_; diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index 8ebb122103177..68005d8afec16 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -156,6 +156,11 @@ class WebGpuContext final { : instance_{instance}, device_{device}, validation_mode_{validation_mode}, query_type_{TimestampQueryType::None}, preserve_device_{preserve_device} {} ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(WebGpuContext); + void LaunchComputePipeline(const wgpu::ComputePassEncoder& compute_pass_encoder, + const std::vector& bind_buffers, + const ProgramArtifact& program_artifact, + uint32_t x, uint32_t y, uint32_t z); + std::vector GetEnabledAdapterToggles() const; std::vector GetEnabledDeviceToggles() const; std::vector GetDisabledDeviceToggles() const; diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index eb65e998c81c5..1d4fea09dd4a6 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -391,18 +391,11 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kMSInternalNHWCD class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 13, CumSum); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 14, CumSum); -class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 10, 12, uint8_t, DequantizeLinear); -class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 10, 12, int8_t, DequantizeLinear); -class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 10, 12, int32_t, DequantizeLinear); -class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, 18, uint8_t, DequantizeLinear); -class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, 18, int8_t, DequantizeLinear); -class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, 18, int32_t, DequantizeLinear); -class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 19, 20, uint8_t, DequantizeLinear); -class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 19, 20, int8_t, DequantizeLinear); -class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 19, 20, int32_t, DequantizeLinear); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 21, uint8_t, DequantizeLinear); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 21, int8_t, DequantizeLinear); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 21, int32_t, DequantizeLinear); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 10, 12, DequantizeLinear); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, 18, DequantizeLinear); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 19, 20, DequantizeLinear); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 21, 22, DequantizeLinear); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 23, DequantizeLinear); std::unique_ptr RegisterKernels() { auto kernel_registry = std::make_unique(); @@ -723,20 +716,15 @@ std::unique_ptr RegisterKernels() { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + + KERNEL_CREATE_INFO_VERSIONED(10, 12, DequantizeLinear), + KERNEL_CREATE_INFO_VERSIONED(13, 18, DequantizeLinear), + KERNEL_CREATE_INFO_VERSIONED(19, 20, DequantizeLinear), + KERNEL_CREATE_INFO_VERSIONED(21, 22, DequantizeLinear), + KERNEL_CREATE_INFO(23, DequantizeLinear), + BuildKernelCreateInfo, BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/core/providers/webnn/builders/helper.cc b/onnxruntime/core/providers/webnn/builders/helper.cc index e0d82b26c2174..79875cbe6f76c 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.cc +++ b/onnxruntime/core/providers/webnn/builders/helper.cc @@ -121,15 +121,15 @@ std::unordered_set GetSupportedNodes(const GraphViewer& graph_viewe return supported_nodes; } -bool AreInputDataTypesSame(const std::string_view op_type, - gsl::span input_types, - const logging::Logger& logger) { - for (size_t i = 1; i < input_types.size(); i++) { - if (input_types[0] != input_types[i]) { +bool AreDataTypesSame(const std::string_view op_type, + gsl::span data_types, + const logging::Logger& logger) { + for (size_t i = 1; i < data_types.size(); i++) { + if (data_types[0] != data_types[i]) { LOGS(logger, VERBOSE) << "[" << op_type - << "] Input data types should be the same, but [" - << input_types[0] << "] does not match " - << input_types[i] << "]."; + << "] data types should be the same, but [" + << data_types[0] << "] does not match " + << data_types[i] << "]."; return false; } } @@ -272,22 +272,23 @@ bool IsMLTensorSupported() { return is_supported; } -// Convert int8 to uint4/int4 (stored as uint8) -uint8_t PackInt8ToUint8AsNibble(int8_t value, const int32_t& data_type) { - uint8_t result = 0; +// Convert int8 to uint4/int4 (stored as uint8), used for creating WebNN Constant +// with same value in both high and low nibbles for uint4/int4 data type. +uint8_t PackInt8ToUint8DoubledNibbles(int8_t value, const int32_t& data_type) { if (data_type == ONNX_NAMESPACE::TensorProto_DataType_UINT4) { if (value < 0 || value > 15) { ORT_THROW("Value cannot be safely converted to uint4."); } - result |= (static_cast(value) << 4); } else { if (value < -8 || value > 7) { ORT_THROW("Value cannot be safely converted to int4."); } - result |= (value << 4); } - return result; + // Explicit conversion + truncate to lower 4 bits + const uint8_t result = static_cast(value) & 0x0F; + // Duplicate the 4-bit value to both high and low nibbles + return (result << 4) | result; } // Convert float32 to float16 (stored as uint16) diff --git a/onnxruntime/core/providers/webnn/builders/helper.h b/onnxruntime/core/providers/webnn/builders/helper.h index 95c4b79053a1f..4fd1ed1d2b299 100644 --- a/onnxruntime/core/providers/webnn/builders/helper.h +++ b/onnxruntime/core/providers/webnn/builders/helper.h @@ -199,7 +199,11 @@ std::unordered_set GetSupportedNodes(const GraphViewer& graph_viewe // Some ONNX ops are supported by decomposed WebNN ops. const std::map> decomposed_op_map = { + {"GroupQueryAttention", + {"add", "cast", "concat", "constant", "cumulativeSum", "div", "expand", "lesser", "matmul", "reshape", "scatterND", + "softmax", "transpose", "where"}}, {"LRN", {"add", "averagePool2d", "div", "mul", "pad", "pow", "transpose"}}, + {"MatMulNBits", {"add", "dequantizeLinear", "matmul", "reshape", "transpose"}}, {"RotaryEmbedding", {"add", "concat", "gather", "mul", "reshape", "split"}}, {"SimplifiedLayerNormalization", {"add", "div", "mul", "pow", "reduceMean", "sqrt"}}, {"SkipSimplifiedLayerNormalization", {"add", "div", "mul", "pow", "reduceMean", "sqrt"}}, @@ -360,9 +364,9 @@ const std::map onnx_to_w {ONNX_NAMESPACE::TensorProto_DataType_UINT64, "uint64"}, }; -bool AreInputDataTypesSame(const std::string_view op_type, - gsl::span input_types, - const logging::Logger& logger); +bool AreDataTypesSame(const std::string_view op_type, + gsl::span input_types, + const logging::Logger& logger); bool IsSupportedDataType(const int32_t onnx_data_type, const emscripten::val& webnn_supported_data_types); bool IsDataTypeSupportedByOp(const std::string_view onnx_op_type, const int32_t onnx_data_type, @@ -386,7 +390,7 @@ bool SetWebnnDataType(emscripten::val& desc, const int32_t data_type); bool IsMLTensorSupported(); -uint8_t PackInt8ToUint8AsNibble(int8_t value, const int32_t& data_type); +uint8_t PackInt8ToUint8DoubledNibbles(int8_t value, const int32_t& data_type); uint16_t PackFloat32ToUint16AsFloat16(float value); } // namespace webnn diff --git a/onnxruntime/core/providers/webnn/builders/impl/attention_helper.h b/onnxruntime/core/providers/webnn/builders/impl/attention_helper.h new file mode 100644 index 0000000000000..3cac8a03a700c --- /dev/null +++ b/onnxruntime/core/providers/webnn/builders/impl/attention_helper.h @@ -0,0 +1,72 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Copyright (c) Intel Corporation. All rights reserved. +// Licensed under the MIT License. + +namespace onnxruntime { +namespace webnn { +/* + ScaledDotProductAttention Subgraph: The basis for MultiHeadAttention and GroupQueryAttention + inputs: query, key, value, scale, attention mask, and reshape_output_shape (for reshape) + Abbreviatios: B is batch_size, S is query sequence_length, kv_S is key/value sequence length, + N is number of attention heads, H is head size, W is hidden_size + + query key + | | + +---matmul---+ scale + | | + +-----div-----+ attn_mask + | | + +-----add-----+ value + | | + +------matmul-----+ + | + (0,2,1,3) transpose B,H,S,N -> B,S,H,N + | + reshape B,S,H,N -> B,S,W + | + output +*/ +emscripten::val ScaledDotProductAttention(ModelBuilder& model_builder, const Node& node, const logging::Logger& logger, + emscripten::val query, emscripten::val key, emscripten::val value, + emscripten::val scale, emscripten::val attn_mask, + std::vector reshape_output_shape) { + emscripten::val common_options = emscripten::val::object(); + // B,H,S,N * B,H,kv_S,N = B,H,S,kv_S + common_options.set("label", node.Name() + "_/Attention/qkv/matmul_1"); + emscripten::val matmul_output = + model_builder.GetBuilder().call("matmul", query, key, common_options); + + common_options.set("label", node.Name() + "_/Attention/qkv/div"); + emscripten::val div_output = + model_builder.GetBuilder().call("mul", matmul_output, scale, common_options); + + emscripten::val softmax_input = div_output; + if (attn_mask != emscripten::val::undefined()) { + common_options.set("label", node.Name() + "_/Attention/attn_mask/softmax_input"); + softmax_input = model_builder.GetBuilder().call("add", div_output, attn_mask, common_options); + } + + common_options.set("label", node.Name() + "_/Attention/attn_mask/softmax_input"); + int32_t softmax_axis = 3; + emscripten::val softmax_output = + model_builder.GetBuilder().call("softmax", softmax_input, softmax_axis, common_options); + + // B,H,S,kv_S * B,H,kv_S,N = B,H,S,N + common_options.set("label", node.Name() + "_/Attention/qkv/matmul_2"); + emscripten::val attn_output = + model_builder.GetBuilder().call("matmul", softmax_output, value, common_options); + + emscripten::val options = emscripten::val::object(); + options.set("permutation", emscripten::val::array(std::vector({0, 2, 1, 3}))); + options.set("label", node.Name() + "_/Attention/qkv/transpose"); + attn_output = model_builder.GetBuilder().call("transpose", attn_output, options); + + common_options.set("label", node.Name() + "_/Attention/qkv/reshape"); + attn_output = model_builder.GetBuilder().call( + "reshape", attn_output, emscripten::val::array(reshape_output_shape), common_options); + + return attn_output; +} + +} // namespace webnn +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webnn/builders/impl/binary_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/binary_op_builder.cc index 29d02690d17c8..130fa068219a2 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/binary_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/binary_op_builder.cc @@ -69,7 +69,7 @@ bool BinaryOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const Node& nod return false; std::array input_types{input0_type, input1_type}; - if (!AreInputDataTypesSame(op_type, input_types, logger)) { + if (!AreDataTypesSame(op_type, input_types, logger)) { return false; } diff --git a/onnxruntime/core/providers/webnn/builders/impl/concat_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/concat_op_builder.cc index f5b78bf4bc16b..8589237617745 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/concat_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/concat_op_builder.cc @@ -70,7 +70,7 @@ bool ConcatOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const Node& nod } std::array input_types{input0_type, input_type}; - if (!AreInputDataTypesSame(op_type, input_types, logger)) { + if (!AreDataTypesSame(op_type, input_types, logger)) { return false; } } diff --git a/onnxruntime/core/providers/webnn/builders/impl/conv_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/conv_op_builder.cc index 436324e087321..1924c3cb5e698 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/conv_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/conv_op_builder.cc @@ -406,7 +406,7 @@ bool ConvOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const Node& node, if (has_input3) { input_types.push_back(input3_type); } - if (!AreInputDataTypesSame(op_type, input_types, logger)) { + if (!AreDataTypesSame(op_type, input_types, logger)) { return false; } diff --git a/onnxruntime/core/providers/webnn/builders/impl/gemm_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/gemm_op_builder.cc index fbf3ac1df2bc2..d4be0f1bee18e 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/gemm_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/gemm_op_builder.cc @@ -237,7 +237,7 @@ bool GemmOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const Node& node, if (has_input3) { input_types.push_back(input3_type); } - if (!AreInputDataTypesSame(op_type, input_types, logger)) { + if (!AreDataTypesSame(op_type, input_types, logger)) { return false; } diff --git a/onnxruntime/core/providers/webnn/builders/impl/gqa_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/gqa_op_builder.cc new file mode 100644 index 0000000000000..e402104251527 --- /dev/null +++ b/onnxruntime/core/providers/webnn/builders/impl/gqa_op_builder.cc @@ -0,0 +1,540 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Copyright (c) Intel Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/shared/utils/utils.h" +#include "core/providers/webnn/builders/helper.h" +#include "core/providers/webnn/builders/model_builder.h" +#include "core/providers/webnn/builders/op_builder_factory.h" +#include "cmath" + +#include "base_op_builder.h" +#include "attention_helper.h" + +namespace onnxruntime { +namespace webnn { + +class GroupQueryAttentionOpBuilder : public BaseOpBuilder { + public: + void AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const override; + // Add operator related. + private: + Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, + const logging::Logger& logger) const override ORT_MUST_USE_RESULT; + + // Operator support related. + private: + bool IsOpSupportedImpl(const GraphViewer&, const Node& node, const WebnnDeviceType /* device_type */, + const logging::Logger& logger) const override; + bool HasSupportedInputsImpl(const GraphViewer&, const Node& node, const emscripten::val& wnn_limits, + const logging::Logger& logger) const override; + bool HasSupportedOutputsImpl(const Node& node, const emscripten::val& wnn_limits, + const logging::Logger& logger) const override; +}; + +void GroupQueryAttentionOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const { + // We check the value of input total_sequence_length in function IsOpSupportedImpl, + // and it should be an initializer and does not participate in Op calculation. + const auto input_name = node.InputDefs()[6]->Name(); + model_builder.AddInitializerToSkip(input_name); + model_builder.AddInputToSkip(input_name); +} + +std::vector generate_indices(int32_t batch_size, int32_t kv_num_heads, int32_t sequence_length) { + std::vector indices; + for (int32_t i = 0; i < sequence_length; ++i) { + for (int32_t j = 0; j < batch_size * kv_num_heads; ++j) { + indices.push_back(j / kv_num_heads); + indices.push_back(j % kv_num_heads); + } + } + return indices; +} + +std::vector repeat_sequence(int32_t sequence_length, int32_t kv_num_heads, int32_t batch_size) { + std::vector repeated; + for (int32_t i = 0; i < sequence_length; ++i) { + for (int32_t j = 0; j < batch_size * kv_num_heads; ++j) { + repeated.push_back(i); + } + } + return repeated; +} + +/** GroupQueryAttention SubGraph. + Abbreviatios: B is batch_size, S is sequence_length, W is hidden_size, P is past_sequence_length + N is number of attention heads, kv_N is number of attention heads for kv, H is head size + G is group size, and G=N/kv_N, W=N*H, h=Sqrt(H). + GQA inputs: query, key, value, past_key, past_value, seqlens_k, total_sequence_length + Notes: cos_cache, sin_cache inputs are not supported. If the data type of the inputs (qkv and past kv) is float16, + we cast them to float32 to ensure data precision. + + query key value + | | | + Reshape Reshape Reshape (B,S,H,N) seqlens_k + | | | / | + | | past_value | (scatter_indices*) | + q_Transpose | \ | / | + (0,2,1,3) | past_key ScatterND-----------------------|------> present_value + \ | / | | +present_key<--\----ScatterND Expand(G) (attention_bias, one/finfo_min mask*) + \ | | / + | Expand(G) | / + | | | / + | k_Transpose | / + | (0,1,3,2) | / + | | | / + +---------------------------------------+ + | ScaledDotProductAttention | + +---------------------------------------+ + | + output +*/ + +Status GroupQueryAttentionOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, + const logging::Logger& logger) const { + const auto& input_defs = node.InputDefs(); + + emscripten::val query_input = model_builder.GetOperand(input_defs[0]->Name()); + emscripten::val key_input = model_builder.GetOperand(input_defs[1]->Name()); + emscripten::val value_input = model_builder.GetOperand(input_defs[2]->Name()); + emscripten::val past_key_input = model_builder.GetOperand(input_defs[3]->Name()); + emscripten::val past_value_input = model_builder.GetOperand(input_defs[4]->Name()); + emscripten::val seqlens_k_input = model_builder.GetOperand(input_defs[5]->Name()); + + std::vector input_q_shape, input_past_k_shape; + ORT_RETURN_IF_NOT(GetShape(*input_defs[0], input_q_shape, logger), "Cannot get query shape"); + ORT_RETURN_IF_NOT(GetShape(*input_defs[3], input_past_k_shape, logger), "Cannot get past_key shape"); + + NodeAttrHelper helper(node); + const uint32_t kv_num_heads = helper.Get("kv_num_heads", 0); + const uint32_t num_heads = helper.Get("num_heads", 0); + + const uint32_t batch_size = SafeInt(input_q_shape[0]); + const uint32_t qkv_sequence_length = SafeInt(input_q_shape[1]); + const uint32_t qkv_hidden_size = SafeInt(input_q_shape[2]); + const uint32_t head_size = SafeInt(qkv_hidden_size / num_heads); + const uint32_t past_sequence_length = SafeInt(input_past_k_shape[2]); + const uint32_t group_size = SafeInt(num_heads / kv_num_heads); + + const float scale_value = helper.Get("scale", 1 / sqrt(static_cast(head_size))); + + const std::vector reshape_output_shape = {batch_size, qkv_sequence_length, qkv_hidden_size}; + const std::vector scatter_indices_shape = {batch_size, qkv_sequence_length, kv_num_heads, 3}; + const std::vector reshape_tensor_shape = {batch_size, qkv_sequence_length, num_heads, head_size}; + + emscripten::val common_options = emscripten::val::object(); + emscripten::val common_desc = emscripten::val::object(); + + int32_t q_type = 0; + ORT_RETURN_IF_NOT(GetType(*input_defs[0], q_type, logger), "Could not get input data type."); + + // Check whether inputs' data type is fp16, if so, we should cast them to fp32 to ensure the calculation precision. + if (q_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) { + common_options.set("label", node.Name() + "_/GQA/preprocess/cast/query_input"); + query_input = model_builder.GetBuilder().call("cast", query_input, emscripten::val("float32"), + common_options); + + common_options.set("label", node.Name() + "_/GQA/preprocess/cast/key_input"); + key_input = + model_builder.GetBuilder().call("cast", key_input, emscripten::val("float32"), common_options); + + common_options.set("label", node.Name() + "_/GQA/preprocess/cast/value_input"); + value_input = model_builder.GetBuilder().call("cast", value_input, emscripten::val("float32"), + common_options); + + common_options.set("label", node.Name() + "_/GQA/preprocess/cast/past_key_input"); + past_key_input = model_builder.GetBuilder().call("cast", past_key_input, + emscripten::val("float32"), common_options); + + common_options.set("label", node.Name() + "_/GQA/preprocess/cast/past_value_input"); + past_value_input = model_builder.GetBuilder().call("cast", past_value_input, + emscripten::val("float32"), common_options); + } + + // Reshape and transpose the input "query" + common_options.set("label", node.Name() + "_/GQA/query/reshape"); + emscripten::val reshaped_query = model_builder.GetBuilder().call( + "reshape", query_input, emscripten::val::array(reshape_tensor_shape), common_options); + + emscripten::val transpose_options = emscripten::val::object(); + transpose_options.set("permutation", emscripten::val::array(std::vector({0, 2, 1, 3}))); + transpose_options.set("label", node.Name() + "_/GQA/query/transpose"); + emscripten::val new_query = + model_builder.GetBuilder().call("transpose", reshaped_query, transpose_options); + + // Reshape the inputs "key" and "value" for scatterND + std::vector reshape_kv_shape = {batch_size, qkv_sequence_length, kv_num_heads, head_size}; + common_options.set("label", node.Name() + "_/GQA/key/reshape_1"); + emscripten::val key_for_scatter = model_builder.GetBuilder().call( + "reshape", key_input, emscripten::val::array(reshape_kv_shape), common_options); + + common_options.set("label", node.Name() + "_/GQA/value/reshape_1"); + emscripten::val value_for_scatter = model_builder.GetBuilder().call( + "reshape", value_input, emscripten::val::array(reshape_kv_shape), common_options); + + /* Calculate scatter_indices for kv's scatterND + if_prefill (0/1 constant) + | + scatter_indices_left_constant scatter_indices_right_constant 0 ---> Where <--- Cast <---seqlens_k + | | | + | Add <--------------------------- scatter_pos* + | | + +------------------+-------------------+ + | + scatter_indices + */ + // Prepare the constant materials for scatter indices + auto left = generate_indices(batch_size, kv_num_heads, qkv_sequence_length); + auto right = repeat_sequence(qkv_sequence_length, kv_num_heads, batch_size); + + std::string left_name = "webnn_GQA_left_constant_of_scatter_indices_" + std::to_string(batch_size) + "_" + + std::to_string(qkv_sequence_length) + "_" + std::to_string(kv_num_heads) + "_2"; + emscripten::val left_constant = model_builder.CreateOrGetConstant( + ONNX_NAMESPACE::TensorProto_DataType_INT32, left_name, left, + std::vector({batch_size * qkv_sequence_length * kv_num_heads, 2})); + + std::string right_name = "webnn_GQA_right_constant_of_scatter_indices_" + std::to_string(batch_size) + "_" + + std::to_string(qkv_sequence_length) + "_" + std::to_string(kv_num_heads) + "_1"; + emscripten::val right_constant = model_builder.CreateOrGetConstant( + ONNX_NAMESPACE::TensorProto_DataType_INT32, right_name, right, + std::vector({batch_size * qkv_sequence_length * kv_num_heads, 1})); + + // The prefilling and decoding stages require different index construction for ScatterND operations. + // Similar to other EPs like CPU and DirectML, when qkv_sequence_length > 1, the key and value are scattered to the + // beginning of kv cache. + std::vector first_condition({(qkv_sequence_length > 1)}); + std::string condition_name = "webnn_GQA_condition_constant_for_where_1"; + emscripten::val condition_constant = model_builder.CreateOrGetConstant( + ONNX_NAMESPACE::TensorProto_DataType_UINT8, condition_name, first_condition, std::vector({1})); + + emscripten::val value_zero_constant = + model_builder.CreateOrGetConstant(ONNX_NAMESPACE::TensorProto_DataType_INT32, 0, {1}); + + // Use concat and reshape to achieve scatter_indices + common_options.set("label", node.Name() + "_/GQA/scatter/where"); + emscripten::val scatter_pos = model_builder.GetBuilder().call( + "where", condition_constant, value_zero_constant, seqlens_k_input, common_options); + + common_options.set("label", node.Name() + "_/GQA/right_constant/add"); + right_constant = model_builder.GetBuilder().call("add", right_constant, scatter_pos, common_options); + + common_options.set("label", node.Name() + "_/GQA/concat_for_pre_scatter_indices"); + std::vector inputs({left_constant, right_constant}); + uint32_t axis = 1; + emscripten::val pre_scatter_indices = + model_builder.GetBuilder().call("concat", emscripten::val::array(inputs), axis, common_options); + + common_options.set("label", node.Name() + "_/GQA/pre_scatter_indices/reshape"); + emscripten::val scatter_indices = model_builder.GetBuilder().call( + "reshape", pre_scatter_indices, emscripten::val::array(scatter_indices_shape), common_options); + + // scatterND for present_key and present_value + common_options.set("label", node.Name() + "_/GQA/present_key/ScatterND"); + emscripten::val present_key = model_builder.GetBuilder().call( + "scatterND", past_key_input, scatter_indices, key_for_scatter, common_options); + + common_options.set("label", node.Name() + "_/GQA/present_value/ScatterND"); + emscripten::val present_value = model_builder.GetBuilder().call( + "scatterND", past_value_input, scatter_indices, value_for_scatter, common_options); + + emscripten::val true_present_key; + emscripten::val true_present_value; + if (group_size != 1) { + // Broadcast key and value for group query by reshape, expand and reshape. + // present kv shape (B,kv_N,P,H) -> (B,kv_N,1,P,H) -> (B,kv_N,N/kv_N,P,H) -> (B,N,P,H) broadcasted kv shape + const std::vector group_broadcast_tensor_shape_1 = {batch_size, kv_num_heads, 1, past_sequence_length, + head_size}; + const std::vector group_broadcast_tensor_shape_2 = {batch_size, kv_num_heads, group_size, + past_sequence_length, head_size}; + const std::vector group_broadcast_tensor_shape_3 = {batch_size, num_heads, past_sequence_length, + head_size}; + common_options.set("label", node.Name() + "_/GQA/true_present_key/reshape_1"); + true_present_key = model_builder.GetBuilder().call( + "reshape", present_key, emscripten::val::array(group_broadcast_tensor_shape_1), common_options); + common_options.set("label", node.Name() + "_/GQA/true_present_key/expand"); + true_present_key = model_builder.GetBuilder().call( + "expand", true_present_key, emscripten::val::array(group_broadcast_tensor_shape_2), common_options); + common_options.set("label", node.Name() + "_/GQA/true_present_key/reshape_2"); + true_present_key = model_builder.GetBuilder().call( + "reshape", true_present_key, emscripten::val::array(group_broadcast_tensor_shape_3), common_options); + + common_options.set("label", node.Name() + "_/GQA/true_present_value/reshape_1"); + true_present_value = model_builder.GetBuilder().call( + "reshape", present_value, emscripten::val::array(group_broadcast_tensor_shape_1), common_options); + common_options.set("label", node.Name() + "_/GQA/true_present_value/expand"); + true_present_value = model_builder.GetBuilder().call( + "expand", true_present_value, emscripten::val::array(group_broadcast_tensor_shape_2), common_options); + common_options.set("label", node.Name() + "_/GQA/true_present_value/reshape_2"); + true_present_value = model_builder.GetBuilder().call( + "reshape", true_present_value, emscripten::val::array(group_broadcast_tensor_shape_3), common_options); + } else { // no need for broadcast + true_present_key = present_key; + true_present_value = present_value; + } + + // Transpose key for matrix multiplication + transpose_options.set("permutation", emscripten::val::array(std::vector({0, 1, 3, 2}))); + transpose_options.set("label", node.Name() + "_/GQA/present_key/transpose"); + true_present_key = model_builder.GetBuilder().call("transpose", true_present_key, transpose_options); + + emscripten::val scale_constant = + model_builder.CreateOrGetConstant(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, scale_value, {1}); + + /* Calculate attention_bias for masking softmax + ones_array (shape=B,N,S,P) range_of_qkv_sequence_length_constant (0,1,2,...) (shape=S) + | | + CumSum (axis=3, exclusive=true, reversed=false) Add <--- scatter_pos + | | + | Expand (shape=P,S) + | | + +-------------------------------> Lesser <---------------------Transpose (1,0) + | + 1 ---> Where <--- finfo_min (minimum value of FP32) + | + attention_bias + */ + const std::vector mask_shape_ones_shape(batch_size * num_heads * qkv_sequence_length * past_sequence_length, + 1); + std::string mask_shape_ones_shape_name = "webnn_GQA_left_constant_of_scatter_indices_" + std::to_string(batch_size) + + "_" + std::to_string(num_heads) + "_" + std::to_string(qkv_sequence_length) + + "_" + std::to_string(past_sequence_length); + emscripten::val mask_shape_ones_shape_constant = model_builder.CreateOrGetConstant( + ONNX_NAMESPACE::TensorProto_DataType_INT32, mask_shape_ones_shape_name, mask_shape_ones_shape, + std::vector({batch_size, num_heads, qkv_sequence_length, past_sequence_length})); + + emscripten::val cumsum_options = emscripten::val::object(); + cumsum_options.set("label", node.Name() + "_range_of_mask_shape"); + cumsum_options.set("exclusive", true); + cumsum_options.set("reversed", false); + emscripten::val neq_left = model_builder.GetBuilder().call( + "cumulativeSum", mask_shape_ones_shape_constant, gsl::narrow(3), cumsum_options); + + std::vector reshape_pre_neq_right = {past_sequence_length, qkv_sequence_length}; + std::vector pre_neq_right_data_range(qkv_sequence_length); + std::iota(pre_neq_right_data_range.begin(), pre_neq_right_data_range.end(), 1); + + std::string pre_neq_right_data_range_name = + "webnn_GQA_left_constant_of_scatter_indices_" + std::to_string(qkv_sequence_length); + emscripten::val pre_neq_right_data_range_constant = model_builder.CreateOrGetConstant( + ONNX_NAMESPACE::TensorProto_DataType_INT32, pre_neq_right_data_range_name, pre_neq_right_data_range, + std::vector({qkv_sequence_length})); + + common_options.set("label", node.Name() + "_/GQA/attn_mask/add"); + emscripten::val pre_neq_right = model_builder.GetBuilder().call( + "add", pre_neq_right_data_range_constant, scatter_pos, common_options); + + common_options.set("label", node.Name() + "_/GQA/expand_neq_right"); + emscripten::val expanded_neq_right = model_builder.GetBuilder().call( + "expand", pre_neq_right, emscripten::val::array(reshape_pre_neq_right), common_options); + + transpose_options.set("permutation", emscripten::val::array(std::vector({1, 0}))); + transpose_options.set("label", node.Name() + "_/GQA/neq_right/transpose"); + emscripten::val neq_right = + model_builder.GetBuilder().call("transpose", expanded_neq_right, transpose_options); + + common_options.set("label", node.Name() + "_/GQA/attn_mask/condition"); + emscripten::val condition = + model_builder.GetBuilder().call("lesser", neq_left, neq_right, common_options); + + emscripten::val value_one_constant = + model_builder.CreateOrGetConstant(ONNX_NAMESPACE::TensorProto_DataType_FLOAT, 1, {1}); + + // finfo_min: the minimum value of float32 + emscripten::val finfo_min_constant = model_builder.CreateOrGetConstant( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT, -3.4028234663852886e+38, {1}); + + common_options.set("label", node.Name() + "_/GQA/attn_mask/where"); + emscripten::val attn_mask = model_builder.GetBuilder().call("where", condition, value_one_constant, + finfo_min_constant, common_options); + + // Execute ScaledDotProductAttention + emscripten::val output = + ScaledDotProductAttention(model_builder, node, logger, new_query, true_present_key, true_present_value, + scale_constant, attn_mask, reshape_output_shape); + + if (q_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) { + common_options.set("label", node.Name() + "_/GQA/postprocess/cast/output"); + output = + model_builder.GetBuilder().call("cast", output, emscripten::val("float16"), common_options); + + common_options.set("label", node.Name() + "_/GQA/postprocess/cast/present_key"); + present_key = model_builder.GetBuilder().call("cast", present_key, emscripten::val("float16"), + common_options); + + common_options.set("label", node.Name() + "_/GQA/postprocess/cast/present_value"); + present_value = model_builder.GetBuilder().call("cast", present_value, emscripten::val("float16"), + common_options); + } + + model_builder.AddOperand(node.OutputDefs()[0]->Name(), std::move(output)); + model_builder.AddOperand(node.OutputDefs()[1]->Name(), std::move(present_key)); + model_builder.AddOperand(node.OutputDefs()[2]->Name(), std::move(present_value)); + + return Status::OK(); +} + +// Operator support related. + +bool GroupQueryAttentionOpBuilder::IsOpSupportedImpl(const GraphViewer& graph_viewer, const Node& node, + const WebnnDeviceType /* device_type */, + const logging::Logger& logger) const { + const auto& input_defs = node.InputDefs(); + const auto& op_type = node.OpType(); + NodeAttrHelper helper(node); + + const auto& total_sequence_length_name = input_defs[6]->Name(); + const auto* total_sequence_length_initializer = graph_viewer.GetConstantInitializer(total_sequence_length_name); + if (!total_sequence_length_initializer) { + LOGS(logger, VERBOSE) << "total_sequence_length must be constant"; + return false; + } + + const auto total_sequence_length_tensor = *total_sequence_length_initializer; + emscripten::val total_sequence_length = emscripten::val::undefined(); + if (!ReadScalarTensorData(total_sequence_length_tensor, total_sequence_length, logger)) { + return false; + } + + std::vector query_shape; + if (!GetShape(*input_defs[0], query_shape, logger)) { + LOGS(logger, VERBOSE) << "Cannot get query shape."; + return false; + } + const auto sequence_length = query_shape[1]; + + std::vector past_key_shape; + if (!GetShape(*input_defs[3], past_key_shape, logger)) { + LOGS(logger, VERBOSE) << "Cannot get past_key shape."; + return false; + } + const auto past_sequence_length = past_key_shape[2]; + + // WebNN EP only supports past_sequence_length of past kv equals to present_sequence_length of present kv + // According to CPU EP, present_sequence_length = max(past_sequence_length,total_sequence_length) + // For prefilling stage (the first prompt), it requires sequence_length == total_sequence_length. + if (sequence_length != 1) { + if (sequence_length != total_sequence_length.as()) { + LOGS(logger, VERBOSE) << op_type << " sequence_length != total_sequence_length. Not first prompt."; + return false; + } + } else { // For decoding stage, it requires past_sequence_length == total_sequence_length. + if (past_sequence_length != total_sequence_length.as()) { + LOGS(logger, VERBOSE) << op_type << " past_sequence_length != total_sequence_length."; + return false; + } + } + + const auto& output_defs = node.OutputDefs(); + if (output_defs.size() != 3) { + LOGS(logger, VERBOSE) << op_type << " output count must be three."; + return false; + } + + return true; +} + +bool GroupQueryAttentionOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const Node& node, + const emscripten::val& wnn_limits, + const logging::Logger& logger) const { + const auto& input_defs = node.InputDefs(); + const auto& op_type = node.OpType(); + + for (int i = 0; i < 9; i++) { + if (i < 7) { + if (!TensorExists(input_defs, i)) { + LOGS(logger, VERBOSE) << op_type << " requires input " << i; + return false; + } + } else { // cos_cache and sin_cache are not supported + if (TensorExists(input_defs, i)) { + LOGS(logger, VERBOSE) << op_type << " does not support input " << i; + return false; + } + } + } + + int32_t q_type = 0; + int32_t k_type = 0; + int32_t v_type = 0; + int32_t past_k_type = 0; + int32_t past_v_type = 0; + int32_t seqlens_k_type = 0; + int32_t total_sequence_length_type = 0; + if (!GetType(*input_defs[0], q_type, logger) || !GetType(*input_defs[1], k_type, logger) || + !GetType(*input_defs[2], v_type, logger) || !GetType(*input_defs[3], past_k_type, logger) || + !GetType(*input_defs[4], past_v_type, logger) || !GetType(*input_defs[5], seqlens_k_type, logger) || + !GetType(*input_defs[6], total_sequence_length_type, logger)) { + return false; + } + + std::array input_types{q_type, k_type, v_type, past_k_type, past_v_type}; + if (!AreDataTypesSame(op_type, input_types, logger)) { + return false; + } + + if (q_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT && q_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) { + return false; + } + + if (seqlens_k_type != ONNX_NAMESPACE::TensorProto_DataType_INT32 && + total_sequence_length_type != ONNX_NAMESPACE::TensorProto_DataType_INT32) { + return false; + } + + std::vector input_q_shape, input_k_shape, input_v_shape, input_past_k_shape, input_past_v_shape; + if (!GetShape(*input_defs[0], input_q_shape, logger) || !GetShape(*input_defs[1], input_k_shape, logger) || + !GetShape(*input_defs[2], input_v_shape, logger) || !GetShape(*input_defs[3], input_past_k_shape, logger) || + !GetShape(*input_defs[4], input_past_v_shape, logger)) { + return false; + } + const auto q_rank = input_q_shape.size(); + const auto k_rank = input_k_shape.size(); + const auto v_rank = input_v_shape.size(); + const auto past_k_rank = input_past_k_shape.size(); + const auto past_v_rank = input_past_v_shape.size(); + if (q_rank != 3 || k_rank != 3 || v_rank != 3) { // The qkv shape should be BSW + LOGS(logger, VERBOSE) << op_type << " qkv shape is not BSW."; + return false; + } + + if (past_k_rank != 4 || past_v_rank != 4) { // The past qkv shape should be BNSH + LOGS(logger, VERBOSE) << op_type << " past qkv shape is not BNSH."; + return false; + } + + return true; +} + +bool GroupQueryAttentionOpBuilder::HasSupportedOutputsImpl(const Node& node, const emscripten::val& wnn_limits, + const logging::Logger& logger) const { + const auto& output_defs = node.OutputDefs(); + const std::string_view op_type = node.OpType(); + int32_t output_type = 0; + int32_t present_k_type = 0; + int32_t present_v_type = 0; + if (!GetType(*output_defs[0], output_type, logger) || !GetType(*output_defs[1], present_k_type, logger) || + !GetType(*output_defs[2], present_v_type, logger)) { + return false; + } + + std::array output_types{output_type, present_k_type, present_v_type}; + if (!AreDataTypesSame(op_type, output_types, logger)) { + return false; + } + + // GQA allows float16, bfloat16 and float32, but WebNN only supports float16 and float32. + if (output_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT && + output_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) { + return false; + } + return true; +} + +void CreateGroupQueryAttentionOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) { + op_registrations.builders.push_back(std::make_unique()); + op_registrations.op_builder_map.emplace(op_type, op_registrations.builders.back().get()); +} + +} // namespace webnn +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webnn/builders/impl/gru_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/gru_op_builder.cc index 403bc8af8ac1f..68abc8ce834f9 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/gru_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/gru_op_builder.cc @@ -216,7 +216,7 @@ bool GruOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const Node& node, if (has_input_initial_h) { input_types.push_back(input_initial_h_type); } - if (!AreInputDataTypesSame(op_type, input_types, logger)) { + if (!AreDataTypesSame(op_type, input_types, logger)) { return false; } diff --git a/onnxruntime/core/providers/webnn/builders/impl/logical_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/logical_op_builder.cc index 7c6de428d0934..42940083cad8e 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/logical_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/logical_op_builder.cc @@ -86,7 +86,7 @@ bool LogicalOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const Node& no if (!GetType(*input_defs[1], input1_type, logger)) return false; std::array input_types{input0_type, input1_type}; - if (!AreInputDataTypesSame(op_type, input_types, logger)) { + if (!AreDataTypesSame(op_type, input_types, logger)) { return false; } } diff --git a/onnxruntime/core/providers/webnn/builders/impl/lstm_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/lstm_op_builder.cc index c49f360c11737..7b5757ecc0faa 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/lstm_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/lstm_op_builder.cc @@ -239,7 +239,7 @@ bool LstmOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const Node& node, if (has_input7) { input_types.push_back(input7_type); } - if (!AreInputDataTypesSame(op_type, input_types, logger)) { + if (!AreDataTypesSame(op_type, input_types, logger)) { return false; } diff --git a/onnxruntime/core/providers/webnn/builders/impl/matMulNBits_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/matMulNBits_op_builder.cc new file mode 100644 index 0000000000000..521aa4a4bfc5a --- /dev/null +++ b/onnxruntime/core/providers/webnn/builders/impl/matMulNBits_op_builder.cc @@ -0,0 +1,255 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Copyright (c) Intel Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/providers/shared/utils/utils.h" +#include "core/providers/webnn/builders/helper.h" +#include "core/providers/webnn/builders/model_builder.h" +#include "core/providers/webnn/builders/op_builder_factory.h" + +#include "base_op_builder.h" + +namespace onnxruntime { +namespace webnn { + +class MatMulNBitsBuilder : public BaseOpBuilder { + // Add operator related. + public: + void AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const override; + + private: + Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, + const logging::Logger& logger) const override ORT_MUST_USE_RESULT; + + // Operator support related. + private: + bool IsOpSupportedImpl(const GraphViewer& graph_viewer, const Node& node, + const WebnnDeviceType /* device_type */, const logging::Logger& logger) const override; + bool HasSupportedInputsImpl(const GraphViewer&, const Node& node, + const emscripten::val& wnn_limits, const logging::Logger& logger) const override; + bool HasSupportedOutputsImpl(const Node& node, const emscripten::val& wnn_limits, + const logging::Logger& logger) const override; +}; + +void MatMulNBitsBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const { + // Inputs B and zero_points (if present) must be initializers. If they are of type uint8, + // they should be stored as uint4 constants in WebNN. Therefore, we skip them here and + // delay their registration as WebNN constants. + const auto& input_defs = node.InputDefs(); + if (input_defs[1]->TypeAsProto()->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_UINT8) { + model_builder.AddInitializerToSkip(input_defs[1]->Name()); // B + if (TensorExists(input_defs, 3)) { + model_builder.AddInitializerToSkip(input_defs[3]->Name()); // zero_points + } + } +} + +// WebNN doesn't provide a dedicated op for MatMulNBits, it can be simply decomposed by +// DequantizeLinear + Transpose + MatMul. Given that the CPU EP currently only supports +// 4-bit quantization, we only handle 4-bit quantization here. +// +// To align with WebNN's dequantizeLinear op contraints, the following transformations are +// required for MatMulNBits inputs: +// 1. B: must be a constant initializer and registered as a 'uint4' WebNN constant with shape +// [N, n_blocks_per_col, blob_size * 2]. +// 2. scales: reshape it to [N, n_blocks_per_col, 1]. +// 3. zero_points: it has the same shape as reshaped scales. If it presents, it must be a +// constant initializer and registered as a 'uint4' WebNN constant. +// Otherwise, it must be registered as a 'uint4' WebNN constant with default value 8. +Status MatMulNBitsBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, + const Node& node, + const logging::Logger& logger) const { + const auto& input_defs = node.InputDefs(); + const auto& output_defs = node.OutputDefs(); + const auto& initializers = model_builder.GetInitializerTensors(); + + emscripten::val input = model_builder.GetOperand(input_defs[0]->Name()); + emscripten::val scales = model_builder.GetOperand(input_defs[2]->Name()); + + std::vector B_shape; // [N, n_blocks_per_col, blob_size] + ORT_RETURN_IF_NOT(GetShape(*input_defs[1], B_shape, logger), "Cannot get B shape"); + + NodeAttrHelper helper(node); + const uint32_t K = helper.Get("K", 0); + const uint32_t N = helper.Get("N", 0); + const uint32_t n_blocks_per_col = SafeInt(B_shape[1]); + const uint32_t double_blob_size = SafeInt(B_shape[2] * 2); + + // Prepare DequantizeLinear's x input + // Input B is an initializer with data type 'uint8', we need to register it as 'uint4' WebNN constant + const std::vector x_shape{N, n_blocks_per_col, double_blob_size}; + emscripten::val x_shape_array = emscripten::val::array(x_shape); + emscripten::val x_desc = emscripten::val::object(); + x_desc.set("dataType", emscripten::val("uint4")); + x_desc.set("shape", x_shape_array); + x_desc.set("dimensions", x_shape_array); + emscripten::val dq_x = emscripten::val::undefined(); + const auto B_tensor = *initializers.at(input_defs[1]->Name()); + ORT_RETURN_IF_ERROR(model_builder.RegisterConstant(B_tensor, dq_x, x_desc, logger)); + + // Prepare DequantizeLinear's x_scale input + // DequantizeLinear's x_scale should be [N, n_blocks_per_col, 1], reshape scales to [N, n_blocks_per_col, 1] + emscripten::val options = emscripten::val::object(); + options.set("label", node.Name() + "_reshape_scales"); + const std::vector x_scale_shape{N, n_blocks_per_col, 1}; + emscripten::val x_scale_shape_array = emscripten::val::array(x_scale_shape); + emscripten::val x_scale = + model_builder.GetBuilder().call("reshape", scales, x_scale_shape_array, options); + + // Prepare DequantizeLinear's x_zero_point input + // x_zero_point has the same shape as x_scale + const bool has_zero_points = TensorExists(input_defs, 3); + emscripten::val x_zero_point = emscripten::val::undefined(); + if (has_zero_points) { + // zero_points is an initializer with data type 'uint8', we need to register it as 'uint4' WebNN constant + const auto zero_points_tensor = *initializers.at(input_defs[3]->Name()); + emscripten::val zero_points_desc = emscripten::val::object(); + zero_points_desc.set("dataType", emscripten::val("uint4")); + zero_points_desc.set("shape", x_scale_shape_array); + zero_points_desc.set("dimensions", x_scale_shape_array); + ORT_RETURN_IF_ERROR(model_builder.RegisterConstant(zero_points_tensor, x_zero_point, zero_points_desc, logger)); + } else { + // zero_points' default value is 8, referred from CPU EP + const int8_t default_zero_point = 8; + x_zero_point = model_builder.CreateOrGetConstant(ONNX_NAMESPACE::TensorProto_DataType_UINT4, + default_zero_point, + x_scale_shape); + } + + // DequantizeLinear + options.set("label", node.Name() + "_dequantizeLinear"); + emscripten::val dq = + model_builder.GetBuilder().call("dequantizeLinear", dq_x, x_scale, x_zero_point, options); + + // Reshape DequantizeLinear to [N, K] + options.set("label", node.Name() + "_reshape_dequantizeLinear"); + const std::vector new_dq_shape{N, K}; + emscripten::val new_dq_shape_array = emscripten::val::array(new_dq_shape); + emscripten::val dq_reshaped = + model_builder.GetBuilder().call("reshape", dq, new_dq_shape_array, options); + + // Transpose reshaped DequantizeLinear to [K, N] + options.set("label", node.Name() + "_transpose_dequantizeLinear"); + emscripten::val dq_transposed = model_builder.GetBuilder().call("transpose", dq_reshaped, options); + + // MatMul + options.set("label", node.Name() + "_matmul"); + emscripten::val output = model_builder.GetBuilder().call("matmul", input, dq_transposed, options); + + // Add output with bias if present + if (TensorExists(input_defs, 5)) { + emscripten::val bias = model_builder.GetOperand(input_defs[5]->Name()); + options.set("label", node.Name() + "_add_bias"); + output = model_builder.GetBuilder().call("add", output, bias, options); + } + + model_builder.AddOperand(output_defs[0]->Name(), std::move(output)); + + return Status::OK(); +} + +bool MatMulNBitsBuilder::IsOpSupportedImpl(const GraphViewer& graph_viewer, + const Node& node, + const WebnnDeviceType /* device_type */, + const logging::Logger& logger) const { + const auto& name = node.Name(); + const auto& input_defs = node.InputDefs(); + std::vector input_shape; + if (!GetShape(*input_defs[0], input_shape, logger)) { + return false; + } + + // Inputs B and zero_points (if present) must be initializers + if (!graph_viewer.GetConstantInitializer(input_defs[1]->Name())) { // B + LOGS(logger, VERBOSE) << "Input B of MatMulNBits [" << name << "] must be known as initializer"; + return false; + } + if (TensorExists(input_defs, 3) && !graph_viewer.GetConstantInitializer(input_defs[3]->Name())) { // zero_points + LOGS(logger, VERBOSE) << "Input zero_points of MatMulNBits [" << name << "] must be known as initializer"; + return false; + } + + // WebNN doesn't support g_idx input + if (TensorExists(input_defs, 4)) { // g_idx + LOGS(logger, VERBOSE) << "Input g_idx of MatMulNBits [" << name << "] is not supported"; + return false; + } + + NodeAttrHelper helper(node); + if (helper.Get("bits", 4) != 4) { + LOGS(logger, VERBOSE) << "Only 4-bit quantization is supported for MatMulNBits, additional bits support is planned"; + } + + return true; +} + +bool MatMulNBitsBuilder::HasSupportedInputsImpl(const GraphViewer&, + const Node& node, const emscripten::val& wnn_limits, + const logging::Logger& logger) const { + const auto& input_defs = node.InputDefs(); + const std::string_view op_type = node.OpType(); + + int32_t A_type = 0; + int32_t B_type = 0; + int32_t scales_type = 0; + int32_t zero_points_type = 0; + if (!GetType(*input_defs[0], A_type, logger) || + !GetType(*input_defs[1], B_type, logger) || + !GetType(*input_defs[2], scales_type, logger)) { + return false; + } + + const bool has_zero_points = TensorExists(input_defs, 3); + if (has_zero_points && !GetType(*input_defs[3], zero_points_type, logger)) { + return false; + } + + InlinedVector input_types = {A_type, scales_type}; + if (!AreDataTypesSame(op_type, input_types, logger)) { + return false; + } + + if (A_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT && A_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) { + LOGS(logger, VERBOSE) << "WebNN only supports float32 or float16 data type for input A of MatMulNBits"; + return false; + } + if (B_type != ONNX_NAMESPACE::TensorProto_DataType_UINT8) { + LOGS(logger, VERBOSE) << "WebNN only supports uint8 data type for input B of MatMulNBits"; + return false; + } + if (has_zero_points && zero_points_type != ONNX_NAMESPACE::TensorProto_DataType_UINT8) { + LOGS(logger, VERBOSE) << "WebNN only supports uint8 data type for input zero_points of MatMulNBits"; + return false; + } + + // We only support 4-bit quantization, which is represented as the uint4 data type in WebNN. + // Ensure that uint4 is supported. + return IsDataTypeSupportedByOp("DequantizeLinear", ONNX_NAMESPACE::TensorProto_DataType_UINT4, + wnn_limits, "input", "x", logger); +} + +bool MatMulNBitsBuilder::HasSupportedOutputsImpl(const Node& node, const emscripten::val& wnn_limits, + const logging::Logger& logger) const { + const auto& output_defs = node.OutputDefs(); + + int32_t output_type = 0; + if (!GetType(*output_defs[0], output_type, logger)) { + return false; + } + + if (output_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT && + output_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) { + LOGS(logger, VERBOSE) << "WebNN only supports float32 or float16 data type for output of MatMulNBits"; + return false; + } + + return true; +} + +void CreateMatMulNBitsOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) { + op_registrations.builders.push_back(std::make_unique()); + op_registrations.op_builder_map.emplace(op_type, op_registrations.builders.back().get()); +} + +} // namespace webnn +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webnn/builders/impl/max_min_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/max_min_op_builder.cc index 7ec4ff640132c..4e4014e3553ea 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/max_min_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/max_min_op_builder.cc @@ -103,7 +103,7 @@ bool MaxMinOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const Node& nod } std::array input_types{input0_type, input_type}; - if (!AreInputDataTypesSame(op_type, input_types, logger)) { + if (!AreDataTypesSame(op_type, input_types, logger)) { return false; } } diff --git a/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc index 704c6a65624d8..e59c9137dbfc5 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/normalization_op_builder.cc @@ -302,7 +302,7 @@ bool NormalizationOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const No if (has_input4) { input_types.push_back(input4_type); } - if (!AreInputDataTypesSame(op_type, input_types, logger)) { + if (!AreDataTypesSame(op_type, input_types, logger)) { return false; } diff --git a/onnxruntime/core/providers/webnn/builders/impl/rotaryEmbedding_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/rotaryEmbedding_op_builder.cc index bdd0c97b7b81c..c8e28818829e5 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/rotaryEmbedding_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/rotaryEmbedding_op_builder.cc @@ -343,7 +343,7 @@ bool RotaryEmbeddingOpBuilder::HasSupportedInputsImpl(const GraphViewer&, } std::array input_types{input_type, cos_cache_type, sin_cache_type}; - if (!AreInputDataTypesSame(op_type, input_types, logger)) { + if (!AreDataTypesSame(op_type, input_types, logger)) { return false; } diff --git a/onnxruntime/core/providers/webnn/builders/impl/ternary_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/ternary_op_builder.cc index f6c1744ca7a3e..7a7f64b1ec96d 100644 --- a/onnxruntime/core/providers/webnn/builders/impl/ternary_op_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/impl/ternary_op_builder.cc @@ -62,7 +62,7 @@ bool TernaryOpBuilder::HasSupportedInputsImpl(const GraphViewer&, const Node& no // ONNX's condition data type is bool which is same as WebNN. // Only need to check X, Y data types. std::array input_types{input1_type, input2_type}; - if (!AreInputDataTypesSame(op_type, input_types, logger)) { + if (!AreDataTypesSame(op_type, input_types, logger)) { return false; } diff --git a/onnxruntime/core/providers/webnn/builders/model_builder.cc b/onnxruntime/core/providers/webnn/builders/model_builder.cc index 399cc5faf6273..e42a33e858a07 100644 --- a/onnxruntime/core/providers/webnn/builders/model_builder.cc +++ b/onnxruntime/core/providers/webnn/builders/model_builder.cc @@ -95,6 +95,114 @@ void ModelBuilder::PreprocessInitializers() { } } +Status ModelBuilder::RegisterConstant(const onnx::TensorProto& tensor, emscripten::val& operand, + emscripten::val& desc, const logging::Logger& logger) { + emscripten::val wnn_builder = GetBuilder(); + const auto data_type = tensor.data_type(); + + // A flag to indicate if we should convert int64 to int32. + const bool should_convert_int64_to_int32 = !IsInt64Supported() && + data_type == ONNX_NAMESPACE::TensorProto_DataType_INT64; + + if (utils::HasExternalData(tensor)) { + // Create WebNN Constant from external data. + std::basic_string external_file_path; + onnxruntime::FileOffsetType data_offset; + SafeInt tensor_byte_size; + ORT_RETURN_IF_ERROR(utils::GetExternalDataInfo( + tensor, graph_viewer_.ModelPath(), external_file_path, data_offset, tensor_byte_size)); + + auto webnnRegisterMLConstant = emscripten::val::module_property("webnnRegisterMLConstant"); + operand = webnnRegisterMLConstant(emscripten::val(external_file_path), + static_cast(data_offset), + static_cast(tensor_byte_size), + wnn_builder, + desc, + should_convert_int64_to_int32); + } else { + std::byte* tensor_ptr = nullptr; + std::vector unpacked_tensor; + emscripten::val view = emscripten::val::undefined(); + + if (tensor.has_raw_data()) { + tensor_ptr = reinterpret_cast(const_cast(tensor.raw_data().c_str())); + } else { + ORT_RETURN_IF_ERROR(onnxruntime::utils::UnpackInitializerData(tensor, unpacked_tensor)); + tensor_ptr = reinterpret_cast(unpacked_tensor.data()); + } + + const auto& shape = tensor.dims(); + auto num_elements = SafeInt(Product(shape)); + if (data_type == ONNX_NAMESPACE::TensorProto_DataType_INT4 || + data_type == ONNX_NAMESPACE::TensorProto_DataType_UINT4) { + // For WebNN int4 and uint4 tensors are stored in Uint8Array, + // so we need to adjust the number of elements. + num_elements = (static_cast(num_elements) + 1) / 2; + } + switch (data_type) { + case ONNX_NAMESPACE::TensorProto_DataType_BOOL: + case ONNX_NAMESPACE::TensorProto_DataType_INT4: + case ONNX_NAMESPACE::TensorProto_DataType_UINT4: + case ONNX_NAMESPACE::TensorProto_DataType_UINT8: + view = emscripten::val{emscripten::typed_memory_view(num_elements, + reinterpret_cast(tensor_ptr))}; + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT8: + view = emscripten::val{emscripten::typed_memory_view(num_elements, + reinterpret_cast(tensor_ptr))}; + break; + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: + view = emscripten::val{emscripten::typed_memory_view(num_elements, + reinterpret_cast(tensor_ptr))}; + break; + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: + view = emscripten::val{emscripten::typed_memory_view(num_elements, + reinterpret_cast(tensor_ptr))}; + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT32: + view = emscripten::val{emscripten::typed_memory_view(num_elements, + reinterpret_cast(tensor_ptr))}; + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + view = emscripten::val{emscripten::typed_memory_view(num_elements, + reinterpret_cast(tensor_ptr))}; + break; + case ONNX_NAMESPACE::TensorProto_DataType_UINT32: + view = emscripten::val{emscripten::typed_memory_view(num_elements, + reinterpret_cast(tensor_ptr))}; + break; + case ONNX_NAMESPACE::TensorProto_DataType_UINT64: + view = emscripten::val{emscripten::typed_memory_view(num_elements, + reinterpret_cast(tensor_ptr))}; + break; + default: + break; + } + + // If int64 is not supported, convert int64 to int32. + std::vector int32_data; + if (should_convert_int64_to_int32) { + try { + int32_data = GetNarrowedIntfromInt64( + gsl::span(reinterpret_cast(tensor_ptr), num_elements)); + LOGS(logger, VERBOSE) << "Initializer '" << tensor.name() << "' is converted from int64 to int32."; + } catch (const std::exception& e) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, e.what()); + } + view = emscripten::val{emscripten::typed_memory_view(num_elements, int32_data.data())}; + + desc.set("dataType", emscripten::val("int32")); + } + + // Wasm memory growth will cause all array buffers reallocation, which will be treated as detached + // buffers in JS side. Simply create a copy to fix it. + view = view.call("slice"); + operand = wnn_builder.call("constant", desc, view["buffer"]); + } + + return Status::OK(); +} + Status ModelBuilder::RegisterInitializers() { for (const auto& pair : GetInitializerTensors()) { const auto& tensor = *pair.second; @@ -119,109 +227,11 @@ Status ModelBuilder::RegisterInitializers() { // in WebNN EP for a while to support older Chromium versions. desc.set("dimensions", emscripten::val::array(dims)); desc.set("shape", emscripten::val::array(dims)); - auto data_type = tensor.data_type(); + const auto data_type = tensor.data_type(); emscripten::val operand = emscripten::val::object(); if (IsSupportedDataType(data_type, wnn_limits_["constant"]["dataTypes"])) { ORT_RETURN_IF_NOT(SetWebnnDataType(desc, data_type), "Unsupported data type"); - auto num_elements = SafeInt(Product(shape)); - emscripten::val view = emscripten::val::undefined(); - std::byte* tensor_ptr = nullptr; - - // A flag to indicate if we should convert int64 to int32. - const bool should_convert_int64_to_int32 = !is_int64_supported_ && - data_type == ONNX_NAMESPACE::TensorProto_DataType_INT64; - - if (utils::HasExternalData(tensor)) { - // Create WebNN Constant from external data. - std::basic_string external_file_path; - onnxruntime::FileOffsetType data_offset; - SafeInt tensor_byte_size; - ORT_RETURN_IF_ERROR(utils::GetExternalDataInfo( - tensor, graph_viewer_.ModelPath(), external_file_path, data_offset, tensor_byte_size)); - - auto webnnRegisterMLConstant = emscripten::val::module_property("webnnRegisterMLConstant"); - operand = webnnRegisterMLConstant(emscripten::val(external_file_path), - static_cast(data_offset), - static_cast(tensor_byte_size), - wnn_builder_, - desc, - should_convert_int64_to_int32); - } else { - if (tensor.has_raw_data()) { - tensor_ptr = reinterpret_cast(const_cast(tensor.raw_data().c_str())); - } else { - // Store temporary unpacked_tensor. - unpacked_tensors_.push_back({}); - std::vector& unpacked_tensor = unpacked_tensors_.back(); - ORT_RETURN_IF_ERROR(onnxruntime::utils::UnpackInitializerData(tensor, unpacked_tensor)); - tensor_ptr = reinterpret_cast(unpacked_tensor.data()); - } - if (data_type == ONNX_NAMESPACE::TensorProto_DataType_INT4 || - data_type == ONNX_NAMESPACE::TensorProto_DataType_UINT4) { - // For WebNN int4 and uint4 tensors are stored in Uint8Array, - // so we need to adjust the number of elements. - num_elements = (static_cast(num_elements) + 1) / 2; - } - switch (data_type) { - case ONNX_NAMESPACE::TensorProto_DataType_BOOL: - case ONNX_NAMESPACE::TensorProto_DataType_INT4: - case ONNX_NAMESPACE::TensorProto_DataType_UINT4: - case ONNX_NAMESPACE::TensorProto_DataType_UINT8: - view = emscripten::val{emscripten::typed_memory_view(num_elements, - reinterpret_cast(tensor_ptr))}; - break; - case ONNX_NAMESPACE::TensorProto_DataType_INT8: - view = emscripten::val{emscripten::typed_memory_view(num_elements, - reinterpret_cast(tensor_ptr))}; - break; - case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: - view = emscripten::val{emscripten::typed_memory_view(num_elements, - reinterpret_cast(tensor_ptr))}; - break; - case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: - view = emscripten::val{emscripten::typed_memory_view(num_elements, - reinterpret_cast(tensor_ptr))}; - break; - case ONNX_NAMESPACE::TensorProto_DataType_INT32: - view = emscripten::val{emscripten::typed_memory_view(num_elements, - reinterpret_cast(tensor_ptr))}; - break; - case ONNX_NAMESPACE::TensorProto_DataType_INT64: - view = emscripten::val{emscripten::typed_memory_view(num_elements, - reinterpret_cast(tensor_ptr))}; - break; - case ONNX_NAMESPACE::TensorProto_DataType_UINT32: - view = emscripten::val{emscripten::typed_memory_view(num_elements, - reinterpret_cast(tensor_ptr))}; - break; - case ONNX_NAMESPACE::TensorProto_DataType_UINT64: - view = emscripten::val{emscripten::typed_memory_view(num_elements, - reinterpret_cast(tensor_ptr))}; - break; - default: - break; - } - - // If int64 is not supported, convert int64 to int32. - std::vector int32_data; - if (should_convert_int64_to_int32) { - try { - int32_data = GetNarrowedIntfromInt64( - gsl::span(reinterpret_cast(tensor_ptr), num_elements)); - LOGS(logger_, VERBOSE) << "Initializer '" << name << "' is converted from int64 to int32."; - } catch (const std::exception& e) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, e.what()); - } - view = emscripten::val{emscripten::typed_memory_view(num_elements, int32_data.data())}; - - desc.set("dataType", emscripten::val("int32")); - } - - // Wasm memory grow will cause all array buffers reallocation, which will be treated as detached - // buffers in JS side. Simply create a copy to fix it. - view = view.call("slice"); - operand = wnn_builder_.call("constant", desc, view["buffer"]); - } + ORT_RETURN_IF_ERROR(RegisterConstant(tensor, operand, desc, logger_)); } else { // TODO: support other type. return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, @@ -375,7 +385,7 @@ Status ModelBuilder::AddOperandFromPersistMemoryBuffer( desc.set("dimensions", emscripten::val::array(shape)); desc.set("shape", emscripten::val::array(shape)); emscripten::val operand = emscripten::val::object(); - // Wasm memory grow will cause all array buffers reallocation, which will be treated as detached + // Wasm memory growth will cause all array buffers reallocation, which will be treated as detached // buffers in JS side. Simply create a copy to fix it. view = view.call("slice"); operand = wnn_builder_.call("constant", desc, view["buffer"]); diff --git a/onnxruntime/core/providers/webnn/builders/model_builder.h b/onnxruntime/core/providers/webnn/builders/model_builder.h index f45e38935651f..4635121090a08 100644 --- a/onnxruntime/core/providers/webnn/builders/model_builder.h +++ b/onnxruntime/core/providers/webnn/builders/model_builder.h @@ -22,9 +22,9 @@ class IOpBuilder; class ModelBuilder { public: - ModelBuilder(const GraphViewer& graph_viewer, const logging::Logger& logger, - const emscripten::val& context, const DataLayout preferred_layout, - const WebnnDeviceType wnn_device_type, const emscripten::val& wnn_limits); + ModelBuilder(const GraphViewer& graph_viewer, const logging::Logger& logger, const emscripten::val& context, + const DataLayout preferred_layout, const WebnnDeviceType wnn_device_type, + const emscripten::val& wnn_limits); ~ModelBuilder() = default; Status Compile(std::unique_ptr& model) ORT_MUST_USE_RESULT; @@ -43,17 +43,23 @@ class ModelBuilder { void AddOperand(const std::string& name, const emscripten::val& operand); + // Register a WebNN constant operand using the provided tensor and descriptor information. + Status RegisterConstant(const onnx::TensorProto& tensor, emscripten::val& operand, + emscripten::val& desc, const logging::Logger& logger); template const emscripten::val& CreateOrGetConstant(const int32_t& data_type, T value, const std::vector& shape = {}); + template + const emscripten::val& CreateOrGetConstant(const int32_t& data_type, std::string name, std::vector value, + const std::vector& shape = {}); + // Use the buffers to persist WebNN allocated data like transposed weight. // It ensures the validity during inference session. std::vector> mem_persist_buffers_; // Add a constant operand (allocate persist buffer and move the ownership to mem_persist_buffers_). - Status AddOperandFromPersistMemoryBuffer( - const std::string& name, const void* buffer, - const size_t size, const std::vector shape, const int32_t data_type); + Status AddOperandFromPersistMemoryBuffer(const std::string& name, const void* buffer, const size_t size, + const std::vector shape, const int32_t data_type); DataLayout GetPreferredLayout() const { return preferred_layout_; } @@ -83,7 +89,6 @@ class ModelBuilder { InlinedHashMap wnn_operands_; std::vector input_names_; std::vector output_names_; - std::vector> unpacked_tensors_; InlinedHashMap input_output_info_; @@ -143,8 +148,10 @@ const emscripten::val& ModelBuilder::CreateOrGetConstant(const int32_t& data_typ name = name_stream.str(); } - // If the operand does not exist, create it. - if (wnn_operands_.find(name) == wnn_operands_.end()) { + auto result = wnn_operands_.find(name); + if (result != wnn_operands_.end()) { + return result->second; + } else { // If the operand does not exist, create it. emscripten::val desc = emscripten::val::object(); desc.set("shape", dims); desc.set("dimensions", dims); @@ -161,7 +168,7 @@ const emscripten::val& ModelBuilder::CreateOrGetConstant(const int32_t& data_typ num_elements = (num_elements + 1) / 2; buffer = emscripten::val::global("Uint8Array").new_(num_elements); if (value) { - buffer.call("fill", emscripten::val(PackInt8ToUint8AsNibble(value, data_type))); + buffer.call("fill", emscripten::val(PackInt8ToUint8DoubledNibbles(value, data_type))); } break; case ONNX_NAMESPACE::TensorProto_DataType_BOOL: @@ -178,9 +185,8 @@ const emscripten::val& ModelBuilder::CreateOrGetConstant(const int32_t& data_typ } break; case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: - buffer = is_float16array_available_ - ? emscripten::val::global("Float16Array").new_(num_elements) - : emscripten::val::global("Uint16Array").new_(num_elements); + buffer = is_float16array_available_ ? emscripten::val::global("Float16Array").new_(num_elements) + : emscripten::val::global("Uint16Array").new_(num_elements); if (value) { buffer.call("fill", emscripten::val(is_float16array_available_ ? value : PackFloat32ToUint16AsFloat16(value))); @@ -221,11 +227,67 @@ const emscripten::val& ModelBuilder::CreateOrGetConstant(const int32_t& data_typ } const emscripten::val constant = wnn_builder_.call("constant", desc, buffer); - wnn_operands_.insert(std::make_pair(name, constant)); + auto insert_result = wnn_operands_.insert(std::make_pair(name, constant)); + return insert_result.first->second; } - - return wnn_operands_.at(name); } +// Create or retrieve a WebNN constant MLOperand with the specified name, value, data type, and shape. +template +const emscripten::val& ModelBuilder::CreateOrGetConstant(const int32_t& data_type, std::string name, + std::vector value, const std::vector& shape) { + emscripten::val dims = emscripten::val::array(); + if (!shape.empty()) { + dims = emscripten::val::array(shape); + } + + auto result = wnn_operands_.find(name); + if (result != wnn_operands_.end()) { + return result->second; + } else { // If the operand does not exist, create it. + emscripten::val desc = emscripten::val::object(); + desc.set("shape", dims); + desc.set("dimensions", dims); + emscripten::val buffer = emscripten::val::undefined(); + if (!SetWebnnDataType(desc, data_type)) { + ORT_THROW("Unsupported data type: " + std::to_string(data_type)); + } + switch (data_type) { + case ONNX_NAMESPACE::TensorProto_DataType_BOOL: + case ONNX_NAMESPACE::TensorProto_DataType_UINT8: + buffer = emscripten::val::global("Uint8Array").new_(emscripten::val::array(value)); + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT8: + buffer = emscripten::val::global("Int8Array").new_(emscripten::val::array(value)); + break; + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: + buffer = is_float16array_available_ + ? emscripten::val::global("Float16Array").new_(emscripten::val::array(value)) + : emscripten::val::global("Uint16Array").new_(emscripten::val::array(value)); + break; + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: + buffer = emscripten::val::global("Float32Array").new_(emscripten::val::array(value)); + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT32: + buffer = emscripten::val::global("Int32Array").new_(emscripten::val::array(value)); + break; + case ONNX_NAMESPACE::TensorProto_DataType_UINT32: + buffer = emscripten::val::global("Uint32Array").new_(emscripten::val::array(value)); + break; + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + buffer = emscripten::val::global("BigInt64Array").new_(emscripten::val::array(value)); + break; + case ONNX_NAMESPACE::TensorProto_DataType_UINT64: + buffer = emscripten::val::global("BigUint64Array").new_(emscripten::val::array(value)); + break; + default: + break; + } + + const emscripten::val constant = wnn_builder_.call("constant", desc, buffer); + auto insert_result = wnn_operands_.insert(std::make_pair(name, constant)); + return insert_result.first->second; + } +} } // namespace webnn } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc b/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc index ee21a33091078..fd7ed74b9d1d6 100644 --- a/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc +++ b/onnxruntime/core/providers/webnn/builders/op_builder_factory.cc @@ -115,6 +115,10 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() { CreateGatherNDOpBuilder("GatherND", op_registrations); } + { // GroupQueryAttention + CreateGroupQueryAttentionOpBuilder("GroupQueryAttention", op_registrations); + } + { // Flatten CreateFlattenOpBuilder("Flatten", op_registrations); } @@ -149,6 +153,10 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() { CreateLstmOpBuilder("LSTM", op_registrations); } + { // MatMulNBits + CreateMatMulNBitsOpBuilder("MatMulNBits", op_registrations); + } + { // Max/Min CreateMaxMinOpBuilder("Max", op_registrations); CreateMaxMinOpBuilder("Min", op_registrations); diff --git a/onnxruntime/core/providers/webnn/builders/op_builder_factory.h b/onnxruntime/core/providers/webnn/builders/op_builder_factory.h index 1c4a7b32f842b..a7d74340c8f37 100644 --- a/onnxruntime/core/providers/webnn/builders/op_builder_factory.h +++ b/onnxruntime/core/providers/webnn/builders/op_builder_factory.h @@ -36,10 +36,12 @@ void CreateGatherOpBuilder(const std::string& op_type, OpBuilderRegistrations& o void CreateGatherElementsOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateGatherNDOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateGemmOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); +void CreateGroupQueryAttentionOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateGruOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateLogicalOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateLRNOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateLstmOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); +void CreateMatMulNBitsOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateMaxMinOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateNormalizationOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreatePadOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); diff --git a/onnxruntime/core/session/provider_registration.cc b/onnxruntime/core/session/provider_registration.cc index 7fb518cdc05ca..e578eb0dacd2d 100644 --- a/onnxruntime/core/session/provider_registration.cc +++ b/onnxruntime/core/session/provider_registration.cc @@ -169,7 +169,7 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider, } else { ORT_UNUSED_PARAMETER(options); status = OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, - "Unknown provider name. Currently supported values are 'OPENVINO', 'SNPE', 'XNNPACK', 'QNN', 'WEBNN' ,'CoreML', and 'AZURE'"); + "Unknown provider name. Currently supported values are 'DML', 'QNN', 'OpenVINO', 'SNPE', 'XNNPACK', 'WEBNN', 'WebGPU', 'AZURE', 'JS', 'VitisAI', and 'CoreML'"); } return status; diff --git a/onnxruntime/python/tools/tensorrt/perf/parse_post_perf.py b/onnxruntime/python/tools/tensorrt/perf/parse_post_perf.py new file mode 100644 index 0000000000000..76194e4ea96bc --- /dev/null +++ b/onnxruntime/python/tools/tensorrt/perf/parse_post_perf.py @@ -0,0 +1,183 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +""" +Parse and Post customized perf data to DB +""" + +import argparse +import csv +import datetime +import logging +import sys + +import pandas as pd +from azure.kusto.data import KustoConnectionStringBuilder +from azure.kusto.data.data_format import DataFormat +from azure.kusto.ingest import IngestionProperties, QueuedIngestClient, ReportLevel + +# Configure logging +LOG_FORMAT = "%(asctime)s - %(levelname)s - %(message)s" +logging.basicConfig(level=logging.INFO, format=LOG_FORMAT) + + +def parse_mobile_perf(log_data: str, model: str, device_id: str, ep: str, commit_id: str, csv_filename: str): + """ + Parse log data and save metrics to a CSV file. + + Args: + log_data (str): The log data to parse. + csv_filename (str): The filename to save the parsed metrics. + """ + metrics = { + "Model": model, + "DeviceId": device_id, + "Ep": ep, + "CommitId": commit_id, + "TTFTAvgTimeSec": None, + "TTFTAvgTokenPerSec": None, + "TokenGenerationAvgTimeSec": None, + "TokenGenerationAvgTokenPerSec": None, + "TokenSamplingAvgTimeSec": None, + "TokenSamplingAvgTokenPerSec": None, + "E2EGenerationAvgTimeSec": None, + "PeakMemoryMB": None, + } + + current_section = None + + for line in log_data.split("\n").strip(): + if "Prompt processing" in line: + current_section = "TTFT" + elif "Token generation" in line: + current_section = "TokenGeneration" + elif "Token sampling" in line: + current_section = "TokenSampling" + elif "E2E generation" in line: + current_section = "E2EGeneration" + elif "Peak working set size" in line: + current_section = "PeakMemory" + + if line.startswith("avg (us):"): + value = float(line.split(":")[1].strip()) / 1_000_000 + if current_section == "TTFT": + metrics["TTFTAvgTimeSec"] = value + elif current_section == "TokenGeneration": + metrics["TokenGenerationAvgTimeSec"] = value + elif current_section == "TokenSampling": + metrics["TokenSamplingAvgTimeSec"] = value + + elif line.startswith("avg (tokens/s):"): + value = float(line.split(":")[1].strip()) + if current_section == "TTFT": + metrics["TTFTAvgTokenPerSec"] = value + elif current_section == "TokenGeneration": + metrics["TokenGenerationAvgTokenPerSec"] = value + elif current_section == "TokenSampling": + metrics["TokenSamplingAvgTokenPerSec"] = value + + elif line.startswith("avg (ms):"): + value = float(line.split(":")[1].strip()) / 1_000 + if current_section == "E2EGeneration": + metrics["E2EGenerationAvgTimeSec"] = value + + elif line.startswith("Peak working set size (bytes):"): + value = float(line.split(":")[1].strip()) / (1024 * 1024) + metrics["PeakMemoryMB"] = value + + for key in metrics.items(): + if metrics[key] is not None and isinstance(metrics[key], (int, float)): + metrics[key] = f"{metrics[key]:.8f}" + + save_to_csv(metrics, csv_filename) + + +def save_to_csv(metrics: dict, csv_filename: str) -> None: + try: + with open(csv_filename, mode="w", newline="") as file: + writer = csv.writer(file) + writer.writerow(metrics.keys()) + writer.writerow(metrics.values()) + logging.info(f"Metrics saved to {csv_filename}") + except OSError as e: + logging.error(f"Failed to save metrics to {csv_filename}: {e}, abort.") + sys.exit(1) + + +def post_to_db(csv_file: str, kusto_table: str, kusto_conn: str, kusto_db: str): + """ + Post data to Kusto DB. + + Args: + csv_file (str): The path to csv file. + kusto_table (str): The Kusto table name. + kusto_conn (str): The Kusto connection string. + kusto_db (str): The Kusto database name. + """ + try: + table = pd.read_csv(csv_file) + upload_time = datetime.datetime.now(tz=datetime.timezone.utc).replace(microsecond=0).isoformat() + table["UploadTime"] = upload_time + + kcsb_ingest = KustoConnectionStringBuilder.with_az_cli_authentication(kusto_conn) + ingest_client = QueuedIngestClient(kcsb_ingest) + + ingestion_props = IngestionProperties( + database=kusto_db, + table=kusto_table, + data_format=DataFormat.CSV, + report_level=ReportLevel.FailuresAndSuccesses, + ) + + ingest_client.ingest_from_dataframe(table, ingestion_properties=ingestion_props) + logging.info(f"Data uploaded to Kusto table {kusto_table} in database {kusto_db}") + except Exception as e: + logging.error(f"Failed to upload data to Kusto DB: {e}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Parse and post perf data to DB") + parser.add_argument("--kusto-table", required=True, help="Kusto table name") + parser.add_argument("--kusto-conn", required=True, help="Kusto connection string") + parser.add_argument("--kusto-db", required=True, help="Kusto database name") + # Post mobile perf data + parser.add_argument("--parse-mobile-perf", action="store_true", help="Parse mobile perf data and post to DB") + parser.add_argument("--log-file", help="Path to log file containing performance data") + parser.add_argument("--model", help="Testing model") + parser.add_argument("--device-id", help="The local mobile device id") + parser.add_argument("--commit-id", help="The ORT commit id") + parser.add_argument("--ep", help="The execution provider running on devices") + parser.add_argument("--output-csv", default="data.csv", help="CSV file to save parsed metrics") + # Post csv data + parser.add_argument("--upload-csv", help="CSV file to upload to DB") + + args = parser.parse_args() + + if args.parse_mobile_perf: + for arg in [args.log_file, args.model, args.device_id, args.ep, args.commit_id]: + if arg is None: + raise ValueError(f"Missing required parameter {arg} for parsing mobile perf data") + log_data = "" + try: + with open(args.log_file) as f: + log_data = f.read() + logging.info(f"Read mobile perf log from {args.log_file}") + except OSError as e: + logging.error(f"Failed to read log file {args.log_file}: {e}") + sys.exit(1) + # Parse the mobile perf data for further upload + parse_mobile_perf(log_data, args.model, args.device_id, args.ep, args.commit_id, args.output_csv) + post_to_db(args.output_csv, args.kusto_table, args.kusto_conn, args.kusto_db) + elif args.upload_csv: + # Upload an existing CSV to the database + # Need to make sure schema is correct + try: + post_to_db(args.upload_csv, args.kusto_table, args.kusto_conn, args.kusto_db) + except Exception as e: + logging.error(f"Failed to upload CSV {args.upload_csv} to DB: {e}") + sys.exit(1) + else: + parser.print_help() + logging.error("Either --parse-mobile-perf or --upload-csv must be specified") + sys.exit(1) diff --git a/onnxruntime/test/providers/cpu/nn/conv_fp16_test.cc b/onnxruntime/test/providers/cpu/nn/conv_fp16_test.cc index 1404071928e09..049f70d660f1e 100644 --- a/onnxruntime/test/providers/cpu/nn/conv_fp16_test.cc +++ b/onnxruntime/test/providers/cpu/nn/conv_fp16_test.cc @@ -968,7 +968,7 @@ TEST(ConvFp16Test, ConvDimWithZero) { vector{1, 1}, // kernel_shape vector{0, 0, 0, 0}, // pads vector{1, 1}, // strides - {} // excluded EPs + {kWebGpuExecutionProvider} // excluded EPs }; vector X; diff --git a/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc b/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc index 51aae0cfd4adf..4e7a6356a5129 100644 --- a/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/quantize_linear_test.cc @@ -156,7 +156,8 @@ TEST(DequantizeLinearOpTest, Scalar) { test.AddInput("x_zero_point", {}, {-10}); test.AddOutput("y", {}, {220.0f}); // Disable Tensorrt EP due to error:node1_quantize_scale_node: out of bounds channel axis 1. Number of input dimensions is 0. - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); + // Disable WebGPU EP due to error: needs at least component size 4 + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); } // dequantize with scalar data @@ -167,7 +168,8 @@ TEST(DequantizeLinearOpMLFloat16Test, Scalar) { test.AddInput("x_zero_point", {}, {-10}); test.AddOutput("y", {}, {MLFloat16(220.0f)}); // Disable Tensorrt EP due to error:node1_quantize_scale_node: out of bounds channel axis 1. Number of input dimensions is 0. - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); + // Disable WebGPU EP due to error: needs at least component size 4 + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); } // dequantize without zero point @@ -176,6 +178,45 @@ TEST(DequantizeLinearOpTest, Without_Zero_Point) { test.AddInput("x", {}, {100}); test.AddInput("x_scale", {}, {2.0f}); test.AddOutput("y", {}, {200.0f}); + // No DQ allowed without corresponding Q. Skip since TRT10 + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kWebGpuExecutionProvider}); +} + +// dequantize without zero point int8 (testing 8 elements for webgpu) +TEST(DequantizeLinearOpTest, No_Zero_Point_int8) { + OpTester test("DequantizeLinear", 10); + test.AddInput("x", {1, 8}, {-10, 50, 100, 120, -9, 49, 99, 119}); + test.AddInput("x_scale", {}, {2.0f}); + test.AddOutput("y", {1, 8}, {-20.0f, 100.0f, 200.0f, 240.0f, -18.0f, 98.0f, 198.0f, 238.0f}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); // No DQ allowed without corresponding Q. Skip since TRT10 +} + +// dequantize without zero point uint8 (testing 8 elements for webgpu) +TEST(DequantizeLinearOpTest, No_Zero_Point_uint8) { + OpTester test("DequantizeLinear", 10); + test.AddInput("x", {1, 8}, {10, 50, 100, 180, 9, 49, 99, 179}); + test.AddInput("x_scale", {}, {2.0f}); + test.AddOutput("y", {1, 8}, {20.0f, 100.0f, 200.0f, 360.0f, 18.0f, 98.0f, 198.0f, 358.0f}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); // No DQ allowed without corresponding Q. Skip since TRT10 +} + +// dequantize zero point int8 (testing 8 elements for webgpu) +TEST(DequantizeLinearOpTest, Zero_Point_int8) { + OpTester test("DequantizeLinear", 10); + test.AddInput("x", {1, 8}, {-10, 50, 100, 120, -9, 49, 99, 119}); + test.AddInput("x_scale", {}, {2.0f}); + test.AddInput("zero_point", {}, {-10}); + test.AddOutput("y", {1, 8}, {0.0f, 120.0f, 220.0f, 260.0f, 2.0f, 118.0f, 218.0f, 258.0f}); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); // No DQ allowed without corresponding Q. Skip since TRT10 +} + +// dequantize zero point uint8 (testing 8 elements for webgpu) +TEST(DequantizeLinearOpTest, Zero_Point_uint8) { + OpTester test("DequantizeLinear", 10); + test.AddInput("x", {1, 8}, {10, 50, 100, 180, 9, 49, 99, 119}); + test.AddInput("x_scale", {}, {2.0f}); + test.AddInput("zero_point", {}, {10}); + test.AddOutput("y", {1, 8}, {0.0f, 80.0f, 180.0f, 340.0f, -2.0f, 78.0f, 178.0f, 218.0f}); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); // No DQ allowed without corresponding Q. Skip since TRT10 } diff --git a/onnxruntime/test/providers/qnn/reshape_expand_op_test.cc b/onnxruntime/test/providers/qnn/reshape_expand_op_test.cc index b978e3a0790db..a5f19881e3d2b 100644 --- a/onnxruntime/test/providers/qnn/reshape_expand_op_test.cc +++ b/onnxruntime/test/providers/qnn/reshape_expand_op_test.cc @@ -238,6 +238,17 @@ TEST_F(QnnHTPBackendTests, Reshape_4D_int32) { 19); // Opset } +// Test that int64 Reshape runs on HTP backend. +TEST_F(QnnHTPBackendTests, Reshape_4D_int64) { + std::vector input_data = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}; + RunReshapeExpandTestOnHTP("Reshape", + TestInputDef({1, 3, 2, 2}, false, input_data), + TestInputDef({3}, true, {1, 1, 12}), + {}, // Attributes + ExpectedEPNodeAssignment::All, + 19); // Opset +} + // Test QDQ Reshape with a shape value of 0 (copy dimension from input) TEST_F(QnnHTPBackendTests, Reshape_4D_0MeansCopy) { RunQDQReshapeExpandTestOnHTP("Reshape", diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index 0fea6acc110f0..e8625e77e9a63 100644 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -4,7 +4,6 @@ # Licensed under the MIT License. from __future__ import annotations -import argparse import contextlib import json import os @@ -14,7 +13,6 @@ import shutil import subprocess import sys -import warnings from pathlib import Path @@ -30,9 +28,8 @@ def version_to_tuple(version: str) -> tuple: REPO_DIR = os.path.normpath(os.path.join(SCRIPT_DIR, "..", "..")) sys.path.insert(0, os.path.join(REPO_DIR, "tools", "python")) - - import util.android as android # noqa: E402 +from build_args import parse_arguments # noqa: E402 from util import ( # noqa: E402 generate_android_triplets, generate_linux_triplets, @@ -76,774 +73,9 @@ def _check_python_version(): ) -def _str_to_bool(s): - """Convert string to bool (in argparse context).""" - if s.lower() not in ["true", "false"]: - raise ValueError(f"Need bool; got {s!r}") - return {"true": True, "false": False}[s.lower()] - - _check_python_version() -def _openvino_verify_device_type(device_read): - choices = ["CPU", "GPU", "NPU"] - - choices1 = [ - "CPU_NO_PARTITION", - "GPU_NO_PARTITION", - "NPU_NO_PARTITION", - "NPU_NO_CPU_FALLBACK", - ] - status_hetero = True - res = False - if device_read in choices: - res = True - elif device_read in choices1: - res = True - elif device_read.startswith(("HETERO:", "MULTI:", "AUTO:")): - res = True - comma_separated_devices = device_read.split(":") - comma_separated_devices = comma_separated_devices[1].split(",") - if len(comma_separated_devices) < 2: - print("At least two devices required in Hetero/Multi/Auto Mode") - status_hetero = False - dev_options = ["CPU", "GPU", "NPU"] - for dev in comma_separated_devices: - if dev not in dev_options: - status_hetero = False - break - - def invalid_hetero_build(): - print("\nIf trying to build Hetero/Multi/Auto, specify the supported devices along with it.\n") - print("specify the keyword HETERO or MULTI or AUTO followed by the devices ") - print("in the order of priority you want to build\n") - print("The different hardware devices that can be added in HETERO or MULTI or AUTO") - print("are ['CPU','GPU','NPU'] \n") - print("An example of how to specify the hetero build type. Ex: HETERO:GPU,CPU \n") - print("An example of how to specify the MULTI build type. Ex: MULTI:GPU,CPU \n") - print("An example of how to specify the AUTO build type. Ex: AUTO:GPU,CPU \n") - sys.exit("Wrong Build Type selected") - - if res is False: - print("\nYou have selected wrong configuration for the build.") - print("pick the build type for specific Hardware Device from following options: ", choices) - print("(or) from the following options with graph partitioning disabled: ", choices1) - print("\n") - if not (device_read.startswith(("HETERO", "MULTI", "AUTO"))): - invalid_hetero_build() - sys.exit("Wrong Build Type selected") - - if status_hetero is False: - invalid_hetero_build() - - return device_read - - -def _qnn_verify_library_kind(library_kind): - choices = ["shared_lib", "static_lib"] - if library_kind not in choices: - print("\nYou have specified an invalid library kind for QNN EP.") - print(f"The invalid library kind was: {library_kind}") - print("Provide a library kind from the following options: ", choices) - print(f"Example: --use_qnn {choices[0]}") - sys.exit("Incorrect build configuration") - return library_kind - - -def parse_arguments(): - class Parser(argparse.ArgumentParser): - # override argument file line parsing behavior - allow multiple arguments per line and handle quotes - def convert_arg_line_to_args(self, arg_line): - return shlex.split(arg_line) - - parser = Parser( - description="ONNXRuntime CI build driver.", - usage=""" - Default behavior is --update --build --test for native architecture builds. - Default behavior is --update --build for cross-compiled builds. - - The Update phase will update git submodules, and run cmake to generate makefiles. - The Build phase will build all projects. - The Test phase will run all unit tests, and optionally the ONNX tests. - - Use the individual flags to only run the specified stages. - """, - # files containing arguments can be specified on the command line with "@" and the arguments within - # will be included at that point - fromfile_prefix_chars="@", - ) - # Main arguments - parser.add_argument("--build_dir", required=True, help="Path to the build directory.") - parser.add_argument( - "--config", - nargs="+", - default=["Debug"], - choices=["Debug", "MinSizeRel", "Release", "RelWithDebInfo"], - help="Configuration(s) to build.", - ) - parser.add_argument("--update", action="store_true", help="Update makefiles.") - parser.add_argument("--build", action="store_true", help="Build.") - parser.add_argument( - "--clean", action="store_true", help="Run 'cmake --build --target clean' for the selected config/s." - ) - parser.add_argument( - "--parallel", - nargs="?", - const="0", - default="1", - type=int, - help="Use parallel build. The optional value specifies the maximum number of parallel jobs. " - "If the optional value is 0 or unspecified, it is interpreted as the number of CPUs.", - ) - parser.add_argument( - "--nvcc_threads", - nargs="?", - default=-1, - type=int, - help="Maximum number of NVCC threads in each parallel job." - "If the value is unspecified, it will be computed based on available memory and number of parallel jobs.", - ) - - parser.add_argument("--test", action="store_true", help="Run unit tests.") - parser.add_argument("--skip_tests", action="store_true", help="Skip all tests.") - parser.add_argument( - "--compile_no_warning_as_error", - action="store_true", - help="Preventing warnings from being treated as errors on compile.", - ) - # Training options - parser.add_argument("--enable_nvtx_profile", action="store_true", help="Enable NVTX profile in ORT.") - parser.add_argument("--enable_memory_profile", action="store_true", help="Enable memory profile in ORT.") - parser.add_argument( - "--enable_training", - action="store_true", - help="Enable full training functionality in ORT. Includes ORTModule and ORT Training APIs.", - ) - parser.add_argument("--enable_training_apis", action="store_true", help="Enable ort training apis.") - parser.add_argument("--enable_training_ops", action="store_true", help="Enable training ops in inference graph.") - - parser.add_argument("--enable_nccl", action="store_true", help="Enable Nccl.") - parser.add_argument("--mpi_home", help="Path to MPI installation dir") - parser.add_argument("--nccl_home", help="Path to NCCL installation dir") - parser.add_argument( - "--use_mpi", nargs="?", default=False, const=True, type=_str_to_bool, help="Disabled by default." - ) - - # enable ONNX tests - parser.add_argument( - "--enable_onnx_tests", - action="store_true", - help="""When running the Test phase, run onnx_test_running against - available test data directories.""", - ) - parser.add_argument("--path_to_protoc_exe", help="Path to protoc exe.") - parser.add_argument("--fuzz_testing", action="store_true", help="Enable Fuzz testing of the onnxruntime.") - parser.add_argument( - "--enable_symbolic_shape_infer_tests", - action="store_true", - help="""When running the Test phase, run symbolic shape inference against - available test data directories.""", - ) - - # generate documentation - parser.add_argument( - "--gen_doc", - nargs="?", - const="yes", - type=str, - help="Generate documentation listing standard ONNX operators and types implemented by " - "various execution providers and contrib operator schemas. Must be used for inference builds, only!" - "Use `--gen_doc validate` to validate these match the current contents in /docs.", - ) - - parser.add_argument("--gen-api-doc", action="store_true", help="Generate API documentation for PyTorch frontend") - - # CUDA related - parser.add_argument("--use_cuda", action="store_true", help="Enable CUDA.") - parser.add_argument( - "--cuda_version", help="The version of CUDA toolkit to use. Auto-detect if not specified. e.g. 9.0" - ) - parser.add_argument( - "--cuda_home", - help="Path to CUDA home." - "Read from CUDA_HOME environment variable if --use_cuda is true and " - "--cuda_home is not specified.", - ) - parser.add_argument( - "--cudnn_home", - help="Path to CUDNN home. " - "Read from CUDNN_HOME environment variable if --use_cuda is true and " - "--cudnn_home is not specified.", - ) - parser.add_argument("--enable_cuda_line_info", action="store_true", help="Enable CUDA line info.") - - parser.add_argument( - "--enable_cuda_nhwc_ops", action="store_true", help="Deprecated; default to enable CUDA NHWC ops in build." - ) - - parser.add_argument("--disable_cuda_nhwc_ops", action="store_true", help="Disable CUDA NHWC ops in build.") - parser.add_argument("--enable_cuda_minimal_build", action="store_true", help="Enable CUDA minimal build.") - - # Python bindings - parser.add_argument("--enable_pybind", action="store_true", help="Enable Python Bindings.") - parser.add_argument("--build_wheel", action="store_true", help="Build Python Wheel.") - parser.add_argument( - "--wheel_name_suffix", - help="Suffix to append to created wheel names. This value is currently only used for nightly builds.", - ) - parser.add_argument("--skip-keras-test", action="store_true", help="Skip tests with Keras if keras is installed") - - # C-Sharp bindings - parser.add_argument( - "--build_csharp", - action="store_true", - help="Build C#.Net DLL and NuGet package. This should be only used in CI pipelines. " - "For building C# bindings and packaging them into nuget package use --build_nuget arg.", - ) - - parser.add_argument( - "--build_nuget", - action="store_true", - help="Build C#.Net DLL and NuGet package on the local machine. " - "Currently only Windows and Linux platforms are supported.", - ) - - parser.add_argument( - "--msbuild_extra_options", - nargs="+", - action="append", - help="Extra properties to pass to msbuild during build. " - "These are just msbuild /p: options without the leading /p:.", - ) - - # Java bindings - parser.add_argument("--build_java", action="store_true", help="Build Java bindings.") - - # Node.js binding - parser.add_argument("--build_nodejs", action="store_true", help="Build Node.js binding and NPM package.") - - # Objective-C binding - parser.add_argument("--build_objc", action="store_true", help="Build Objective-C binding.") - - # Build a shared lib - parser.add_argument("--build_shared_lib", action="store_true", help="Build a shared library for the ONNXRuntime.") - - # Build a shared lib - parser.add_argument( - "--build_apple_framework", action="store_true", help="Build a macOS/iOS framework for the ONNXRuntime." - ) - - # Build options - parser.add_argument( - "--cmake_extra_defines", - nargs="+", - action="append", - help="Extra definitions to pass to CMake during build system " - "generation. These are just CMake -D options without the leading -D.", - ) - parser.add_argument("--target", help="Build a specific target, e.g. winml_dll") - # This flag is needed when : - # 1. The OS is 64 bits Windows - # 2. And the target binary is for 32 bits Windows - # 3. And the python used for running this script is 64 bits. - # But if you can get a 32 bits python, the build will run better and you won't need this flag. - parser.add_argument( - "--x86", - action="store_true", - help="[cross-compiling] Create Windows x86 makefiles. Requires --update and no existing cache " - "CMake setup. Delete CMakeCache.txt if needed", - ) - parser.add_argument( - "--rv64", - action="store_true", - help="[cross-compiling] Create riscv64 makefiles. Requires --update and no existing cache " - "CMake setup. Delete CMakeCache.txt if needed", - ) - parser.add_argument( - "--arm", - action="store_true", - help="[cross-compiling] Create ARM makefiles. Requires --update and no existing cache " - "CMake setup. Delete CMakeCache.txt if needed", - ) - parser.add_argument( - "--arm64", - action="store_true", - help="[cross-compiling] Create ARM64 makefiles. Requires --update and no existing cache " - "CMake setup. Delete CMakeCache.txt if needed", - ) - parser.add_argument( - "--arm64ec", - action="store_true", - help="[cross-compiling] Create ARM64EC makefiles. Requires --update and no existing cache " - "CMake setup. Delete CMakeCache.txt if needed", - ) - parser.add_argument( - "--buildasx", - action="store_true", - help="[cross-compiling] Create ARM64X Binary.", - ) - parser.add_argument( - "--riscv_toolchain_root", - type=str, - default="", - help="Path to RISC-V toolchain root dir. e.g. --riscv_toolchain_root=$HOME/riscv-tools/", - ) - parser.add_argument( - "--riscv_qemu_path", - type=str, - default="", - help="Path to RISC-V qemu. e.g. --riscv_qemu_path=$HOME/qemu-dir/qemu-riscv64", - ) - # https://gitlab.kitware.com/cmake/cmake/-/issues/25192 - parser.add_argument( - "--msvc_toolset", - help="MSVC toolset to use. e.g. 14.11. It doesn't work if the version number is in the range of [14.36, 14.39]", - ) - parser.add_argument("--windows_sdk_version", help="Windows SDK version to use. e.g. 10.0.19041.0") - parser.add_argument("--android", action="store_true", help="Build for Android") - parser.add_argument( - "--android_abi", - default="arm64-v8a", - choices=["armeabi-v7a", "arm64-v8a", "x86", "x86_64"], - help="Specify the target Android Application Binary Interface (ABI)", - ) - parser.add_argument("--android_api", type=int, default=27, help="Android API Level, e.g. 21") - parser.add_argument( - "--android_sdk_path", type=str, default=os.environ.get("ANDROID_HOME", ""), help="Path to the Android SDK" - ) - parser.add_argument( - "--android_ndk_path", type=str, default=os.environ.get("ANDROID_NDK_HOME", ""), help="Path to the Android NDK" - ) - parser.add_argument( - "--android_cpp_shared", - action="store_true", - help="Build with shared libc++ instead of the default static libc++.", - ) - parser.add_argument("--android_run_emulator", action="store_true", help="Start up an Android emulator if needed.") - - parser.add_argument("--use_gdk", action="store_true", help="Build with the GDK toolchain.") - parser.add_argument( - "--gdk_edition", - default=os.path.normpath(os.environ.get("GameDKLatest", "")).split(os.sep)[-1], # noqa: SIM112 - help="Build with a specific GDK edition. Defaults to the latest installed.", - ) - parser.add_argument("--gdk_platform", default="Scarlett", help="Sets the GDK target platform.") - parser.add_argument("--enable_wasm_memory64", action="store_true", help="Enable WebAssembly 64bit support") - platform_group = parser.add_mutually_exclusive_group() - platform_group.add_argument("--ios", action="store_true", help="build for ios") - platform_group.add_argument("--visionos", action="store_true", help="build for visionOS") - platform_group.add_argument("--tvos", action="store_true", help="build for tvOS") - platform_group.add_argument( - "--macos", - choices=["MacOSX", "Catalyst"], - help="Specify the target platform for macOS build. Only specify this argument when --build_apple_framework is present.", - ) - - parser.add_argument( - "--apple_sysroot", default="", help="Specify the location name of the macOS platform SDK to be used" - ) - parser.add_argument( - "--ios_toolchain_file", - default="", - help="Path to ios toolchain file, or cmake/onnxruntime_ios.toolchain.cmake will be used", - ) - parser.add_argument( - "--visionos_toolchain_file", - default="", - help="Path to visionos toolchain file, or cmake/onnxruntime_visionos.toolchain.cmake will be used", - ) - parser.add_argument( - "--tvos_toolchain_file", - default="", - help="Path to tvos toolchain file, or cmake/onnxruntime_tvos.toolchain.cmake will be used", - ) - parser.add_argument( - "--xcode_code_signing_team_id", default="", help="The development team ID used for code signing in Xcode" - ) - parser.add_argument( - "--xcode_code_signing_identity", default="", help="The development identity used for code signing in Xcode" - ) - parser.add_argument( - "--use_xcode", - action="store_const", - const="Xcode", - dest="cmake_generator", - help="Use Xcode as cmake generator, this is only supported on MacOS. (non Catalyst build). Equivalent to '--cmake_generator Xcode'.", - ) - parser.add_argument( - "--osx_arch", - default="arm64" if platform.machine() == "arm64" else "x86_64", - choices=["arm64", "arm64e", "x86_64"], - help="Specify the Target specific architectures for macOS and iOS, This is only supported on MacOS", - ) - parser.add_argument( - "--apple_deploy_target", - type=str, - help="Specify the minimum version of the target platform (e.g. macOS or iOS)This is only supported on MacOS", - ) - # A 32-bit progress doesn't have enough memory to run all the tests in onnxruntime_test_all. - # Mimalloc is incompatible with address sanitizer. - # Address sanitizer itself is also a memory leak checker, so when it is enabled we should disable_memleak_checker. - parser.add_argument( - "--enable_address_sanitizer", action="store_true", help="Enable address sanitizer. Windows/Linux/MacOS only." - ) - # The following flag is mostly designed to be used in ONNX Runtime's Azure DevOps/Github build pipelines. Its main purpose is to make the built binaries pass BinSkim scan. - parser.add_argument("--use_binskim_compliant_compile_flags", action="store_true", help="Use preset compile flags.") - parser.add_argument( - "--disable_memleak_checker", - action="store_true", - help="Disable memory leak checker from Windows build. By default it is enabled in Windows Debug build. This option is Windows only.", - ) - # Dependency search with vcpkg - parser.add_argument( - "--use_vcpkg", - action="store_true", - help="Use vcpkg to search dependencies. Requires CMAKE_TOOLCHAIN_FILE for vcpkg.cmake", - ) - parser.add_argument( - "--use_vcpkg_ms_internal_asset_cache", - action="store_true", - help="Microsoft internal option", - ) - - # WebAssembly build - parser.add_argument("--build_wasm", action="store_true", help="Build for WebAssembly") - parser.add_argument("--build_wasm_static_lib", action="store_true", help="Build for WebAssembly static library") - parser.add_argument("--emsdk_version", default="4.0.4", help="Specify version of emsdk") - - parser.add_argument("--enable_wasm_simd", action="store_true", help="Enable WebAssembly SIMD") - parser.add_argument("--enable_wasm_relaxed_simd", action="store_true", help="Enable WebAssembly Relaxed SIMD") - parser.add_argument("--enable_wasm_threads", action="store_true", help="Enable WebAssembly multi-threads support") - - parser.add_argument( - "--disable_wasm_exception_catching", action="store_true", help="Disable exception catching in WebAssembly." - ) - parser.add_argument( - "--enable_wasm_api_exception_catching", action="store_true", help="Catch exceptions at top level api." - ) - parser.add_argument( - "--enable_wasm_exception_throwing_override", - action="store_true", - help="Enable exception throwing in WebAssembly, this will override default disabling exception throwing " - "behavior when disable exceptions.", - ) - parser.add_argument("--wasm_run_tests_in_browser", action="store_true", help="Run WebAssembly tests in browser") - - parser.add_argument( - "--enable_wasm_profiling", action="store_true", help="Enable WebAssembly profiling and preserve function names" - ) - parser.add_argument( - "--enable_wasm_debug_info", action="store_true", help="Build WebAssembly with DWARF format debug info" - ) - - parser.add_argument("--wasm_malloc", help="Specify memory allocator for WebAssembly") - - parser.add_argument( - "--emscripten_settings", - nargs="+", - action="append", - help="Extra emscripten settings to pass to emcc using '-s =' during build.", - ) - - # Enable onnxruntime-extensions - parser.add_argument( - "--use_extensions", - action="store_true", - help="Enable custom operators in onnxruntime-extensions, use git submodule onnxruntime-extensions " - "in path cmake/external/onnxruntime-extensions by default.", - ) - parser.add_argument( - "--extensions_overridden_path", - type=str, - help="Path to pre-pulled onnxruntime-extensions, will override default onnxruntime-extensions path.", - ) - - # Arguments needed by CI - parser.add_argument("--cmake_path", default="cmake", help="Path to the CMake program.") - parser.add_argument( - "--ctest_path", - default="ctest", - help="Path to the CTest program. It can be an empty string. If it is empty, " - "we will use this script driving the test programs directly.", - ) - parser.add_argument( - "--skip_submodule_sync", - action="store_true", - help="Don't do a 'git submodule update'. Makes the Update phase faster.", - ) - parser.add_argument("--use_mimalloc", action="store_true", help="Use mimalloc allocator") - parser.add_argument("--use_dnnl", action="store_true", help="Build with DNNL.") - parser.add_argument( - "--dnnl_gpu_runtime", action="store", default="", type=str.lower, help="e.g. --dnnl_gpu_runtime ocl" - ) - parser.add_argument( - "--dnnl_opencl_root", - action="store", - default="", - help="Path to OpenCL SDK. " - 'e.g. --dnnl_opencl_root "C:/Program Files (x86)/IntelSWTools/sw_dev_tools/OpenCL/sdk"', - ) - parser.add_argument( - "--use_openvino", - nargs="?", - const="CPU", - type=_openvino_verify_device_type, - help="Build with OpenVINO for specific hardware.", - ) - parser.add_argument( - "--dnnl_aarch64_runtime", action="store", default="", type=str.lower, help="e.g. --dnnl_aarch64_runtime acl" - ) - parser.add_argument( - "--dnnl_acl_root", - action="store", - default="", - help='Path to ACL ROOT DIR. e.g. --dnnl_acl_root "$HOME/ComputeLibrary/"', - ) - parser.add_argument("--use_coreml", action="store_true", help="Build with CoreML support.") - parser.add_argument("--use_webnn", action="store_true", help="Build with WebNN support.") - parser.add_argument("--use_snpe", action="store_true", help="Build with SNPE support.") - parser.add_argument("--snpe_root", help="Path to SNPE SDK root.") - parser.add_argument("--use_nnapi", action="store_true", help="Build with NNAPI support.") - parser.add_argument("--use_vsinpu", action="store_true", help="Build with VSINPU support.") - parser.add_argument( - "--nnapi_min_api", type=int, help="Minimum Android API level to enable NNAPI, should be no less than 27" - ) - parser.add_argument("--use_jsep", action="store_true", help="Build with JavaScript kernels.") - parser.add_argument("--use_webgpu", action="store_true", help="Build with WebGPU support.") - parser.add_argument("--use_external_dawn", action="store_true", help="Treat Dawn as an external dependency.") - parser.add_argument( - "--use_qnn", - nargs="?", - const="shared_lib", # If provide --use_qnn without an arg, defaults to a shared library. - type=_qnn_verify_library_kind, - help="Build with QNN support. Specify 'shared_lib' or 'static_lib' to build QNN EP " - "as a shared or static library, respectively.", - ) - parser.add_argument("--qnn_home", help="Path to QNN SDK dir.") - parser.add_argument("--use_rknpu", action="store_true", help="Build with RKNPU.") - parser.add_argument("--use_preinstalled_eigen", action="store_true", help="Use pre-installed Eigen.") - parser.add_argument("--eigen_path", help="Path to pre-installed Eigen.") - parser.add_argument("--enable_msinternal", action="store_true", help="Enable for Microsoft internal builds only.") - parser.add_argument("--use_vitisai", action="store_true", help="Build with Vitis-AI") - parser.add_argument("--use_tensorrt", action="store_true", help="Build with TensorRT") - parser.add_argument( - "--use_tensorrt_builtin_parser", action="store_true", default=True, help="Use TensorRT builtin parser" - ) - parser.add_argument("--use_tensorrt_oss_parser", action="store_true", help="Use TensorRT OSS parser") - parser.add_argument("--tensorrt_home", help="Path to TensorRT installation dir") - parser.add_argument("--test_all_timeout", default="10800", help="Set timeout for onnxruntime_test_all") - parser.add_argument("--use_migraphx", action="store_true", help="Build with MIGraphX") - parser.add_argument("--migraphx_home", help="Path to MIGraphX installation dir") - parser.add_argument("--use_full_protobuf", action="store_true", help="Use the full protobuf library") - parser.add_argument("--enable_pix_capture", action="store_true", help="Enable Pix Support.") - - parser.add_argument( - "--skip_onnx_tests", - action="store_true", - help="Explicitly disable all onnx related tests. Note: Use --skip_tests to skip all tests.", - ) - parser.add_argument("--skip_winml_tests", action="store_true", help="Explicitly disable all WinML related tests") - parser.add_argument("--skip_nodejs_tests", action="store_true", help="Explicitly disable all Node.js binding tests") - - parser.add_argument( - "--enable_msvc_static_runtime", action="store_true", help="Enable static linking of MSVC runtimes." - ) - parser.add_argument( - "--cmake_generator", - choices=[ - "MinGW Makefiles", - "Ninja", - "NMake Makefiles", - "NMake Makefiles JOM", - "Unix Makefiles", - "Visual Studio 17 2022", - "Xcode", - ], - default=None, - help="Specify the generator that CMake invokes.", - ) - parser.add_argument("--use_dml", action="store_true", help="Build with DirectML.") - parser.add_argument( - "--dml_path", - type=str, - default="", - help="Path to a custom DirectML installation (must have bin/, lib/, and include/ subdirectories).", - ) - parser.add_argument("--use_winml", action="store_true", help="Build with WinML.") - parser.add_argument( - "--winml_root_namespace_override", type=str, help="Specify the namespace that WinML builds into." - ) - parser.add_argument( - "--dml_external_project", action="store_true", help="Build with DirectML as an external project." - ) - parser.add_argument( - "--use_telemetry", action="store_true", help="Only official builds can set this flag to enable telemetry." - ) - parser.add_argument("--enable_wcos", action="store_true", help="Build for Windows Core OS.") - # Do not enable LTO when the compiler is MSVC and the flag for generating debug symbols is set to /Z7 and training - # is also enabled. Because both LTO and /Zi could significantly increase *.obj/*.lib files' size, and on Windows - # there is a 4GB per file limit(ERROR LNK1248). We may solve the issue by splitting the big static libs to smaller - # ones. Before the refactoring work is done, we should avoid enabling LTO and ccache at the same time because ccache - # needs /Z7. - parser.add_argument("--enable_lto", action="store_true", help="Enable Link Time Optimization") - parser.add_argument("--enable_transformers_tool_test", action="store_true", help="Enable transformers tool test") - parser.add_argument( - "--use_acl", - action="store_true", - help="Build with ACL for ARM architectures.", - ) - parser.add_argument("--acl_home", help="Path to ACL home dir") - parser.add_argument("--acl_libs", help="Path to ACL libraries") - parser.add_argument("--use_armnn", action="store_true", help="Enable ArmNN Execution Provider.") - parser.add_argument( - "--armnn_relu", action="store_true", help="Use the Relu operator implementation from the ArmNN EP." - ) - parser.add_argument( - "--armnn_bn", action="store_true", help="Use the Batch Normalization operator implementation from the ArmNN EP." - ) - parser.add_argument("--armnn_home", help="Path to ArmNN home dir") - parser.add_argument("--armnn_libs", help="Path to ArmNN libraries") - parser.add_argument("--build_micro_benchmarks", action="store_true", help="Build ONNXRuntime micro-benchmarks.") - parser.add_argument("--no_kleidiai", action="store_true", help="Disable KleidiAI integration on Arm platforms.") - - # options to reduce binary size - parser.add_argument( - "--minimal_build", - default=None, - nargs="*", - type=str.lower, - help="Create a build that only supports ORT format models. " - "See https://onnxruntime.ai/docs/tutorials/mobile/ for more information. " - "RTTI is automatically disabled in a minimal build. " - "To enable execution providers that compile kernels at runtime (e.g. NNAPI) pass 'extended' " - "as a parameter. e.g. '--minimal_build extended'. " - "To enable support for custom operators pass 'custom_ops' as a parameter. " - "e.g. '--minimal_build custom_ops'. This can be combined with an 'extended' build by passing " - "'--minimal_build extended custom_ops'", - ) - - parser.add_argument( - "--include_ops_by_config", - type=str, - help="Include ops from config file. See /docs/Reduced_Operator_Kernel_build.md for more information.", - ) - parser.add_argument( - "--enable_reduced_operator_type_support", - action="store_true", - help="If --include_ops_by_config is specified, and the configuration file has type reduction " - "information, limit the types individual operators support where possible to further " - "reduce the build size. " - "See /docs/Reduced_Operator_Kernel_build.md for more information.", - ) - - parser.add_argument("--disable_contrib_ops", action="store_true", help="Disable contrib ops (reduces binary size)") - parser.add_argument( - "--disable_ml_ops", action="store_true", help="Disable traditional ML ops (reduces binary size)" - ) - # Please note in our CMakeLists.txt this is already default on. But in this file we reverse it to default OFF. - parser.add_argument("--disable_rtti", action="store_true", help="Disable RTTI (reduces binary size)") - parser.add_argument( - "--disable_types", - nargs="+", - default=[], - choices=["float8", "optional", "sparsetensor"], - help="Disable selected data types (reduces binary size)", - ) - parser.add_argument( - "--disable_exceptions", - action="store_true", - help="Disable exceptions to reduce binary size. Requires --minimal_build.", - ) - - parser.add_argument("--rocm_version", help="The version of ROCM stack to use. ") - parser.add_argument("--use_rocm", action="store_true", help="Build with ROCm") - parser.add_argument("--rocm_home", help="Path to ROCm installation dir") - - # Code coverage - parser.add_argument( - "--code_coverage", action="store_true", help="Generate code coverage when targeting Android (only)." - ) - - # lazy tensor support. - parser.add_argument( - "--enable_lazy_tensor", action="store_true", help="Enable use ORT as backend in Pytorch LazyTensor." - ) - - parser.add_argument("--ms_experimental", action="store_true", help="Build microsoft experimental operators.") - - parser.add_argument( - "--enable_external_custom_op_schemas", - action="store_true", - help="Enable registering user defined custom operation schemas at shared library load time.\ - This feature is only supported/available on Ubuntu.", - ) - - parser.add_argument( - "--external_graph_transformer_path", type=str, help="path to the external graph transformer dir." - ) - - parser.add_argument( - "--enable_cuda_profiling", - action="store_true", - help="enable cuda kernel profiling, \ - cupti library must be added to PATH beforehand.", - ) - parser.add_argument("--use_cann", action="store_true", help="Build with CANN") - parser.add_argument("--cann_home", help="Path to CANN installation dir") - - parser.add_argument( - "--enable_rocm_profiling", - action="store_true", - help="enable rocm kernel profiling.", - ) - - parser.add_argument("--use_xnnpack", action="store_true", help="Enable xnnpack EP.") - parser.add_argument("--use_avx512", action="store_true", help="Enable AVX512 instructions") - parser.add_argument("--use_azure", action="store_true", help="Enable azure EP.") - - parser.add_argument("--use_cache", action="store_true", help="Use compiler cache in CI") - - parser.add_argument("--use_triton_kernel", action="store_true", help="Use triton compiled kernels") - parser.add_argument("--use_lock_free_queue", action="store_true", help="Use lock-free task queue for threadpool.") - - parser.add_argument( - "--enable_generic_interface", - action="store_true", - help="build ORT shared library and compatible bridge with primary EPs(tensorRT, OpenVino, Qnn, vitisai) but not tests", - ) - - if not is_windows(): - parser.add_argument( - "--allow_running_as_root", - action="store_true", - help="Allow build to be run as root user. This is not allowed by default.", - ) - - args = parser.parse_args() - if args.android_sdk_path: - args.android_sdk_path = os.path.normpath(args.android_sdk_path) - if args.android_ndk_path: - args.android_ndk_path = os.path.normpath(args.android_ndk_path) - - if args.enable_wasm_api_exception_catching: - # if we catch on api level, we don't want to catch all - args.disable_wasm_exception_catching = True - if not args.disable_wasm_exception_catching or args.enable_wasm_api_exception_catching: - # doesn't make sense to catch if no one throws - args.enable_wasm_exception_throwing_override = True - - if args.cmake_generator is None and is_windows(): - args.cmake_generator = "Ninja" if args.build_wasm else "Visual Studio 17 2022" - - if args.enable_cuda_nhwc_ops: - warnings.warn( - "The argument '--enable_cuda_nhwc_ops' is deprecated and is default to True. ", DeprecationWarning - ) - - return args - - def is_reduced_ops_build(args): return args.include_ops_by_config is not None @@ -953,7 +185,7 @@ def use_dev_mode(args): return False if args.use_armnn: return False - if (args.ios or args.visionos or args.tvos) and is_macOS(): + if is_macOS() and (args.ios or args.visionos or args.tvos): return False SYSTEM_COLLECTIONURI = os.getenv("SYSTEM_COLLECTIONURI") # noqa: N806 if SYSTEM_COLLECTIONURI and SYSTEM_COLLECTIONURI != "https://dev.azure.com/onnxruntime/": @@ -1112,7 +344,6 @@ def generate_build_tree( cuda_home, cudnn_home, rocm_home, - mpi_home, nccl_home, tensorrt_home, migraphx_home, @@ -1140,7 +371,49 @@ def generate_build_tree( disable_float8_types = args.android or ("float8" in types_to_disable) disable_optional_type = "optional" in types_to_disable disable_sparse_tensors = "sparsetensor" in types_to_disable + if is_windows(): + cmake_args += [ + "-Donnxruntime_USE_DML=" + ("ON" if args.use_dml else "OFF"), + "-Donnxruntime_USE_WINML=" + ("ON" if args.use_winml else "OFF"), + "-Donnxruntime_USE_TELEMETRY=" + ("ON" if args.use_telemetry else "OFF"), + "-Donnxruntime_ENABLE_PIX_FOR_WEBGPU_EP=" + ("ON" if args.enable_pix_capture else "OFF"), + ] + if args.winml_root_namespace_override: + cmake_args.append("-Donnxruntime_WINML_NAMESPACE_OVERRIDE=" + args.winml_root_namespace_override) + if args.disable_memleak_checker or args.enable_address_sanitizer: + cmake_args.append("-Donnxruntime_ENABLE_MEMLEAK_CHECKER=OFF") + else: + cmake_args.append("-Donnxruntime_ENABLE_MEMLEAK_CHECKER=ON") + + if args.use_winml: + cmake_args.append("-Donnxruntime_BUILD_WINML_TESTS=" + ("OFF" if args.skip_winml_tests else "ON")) + if args.dml_path: + cmake_args += [ + "-Donnxruntime_USE_CUSTOM_DIRECTML=ON", + "-Ddml_INCLUDE_DIR=" + os.path.join(args.dml_path, "include"), + "-Ddml_LIB_DIR=" + os.path.join(args.dml_path, "lib"), + ] + if args.dml_external_project: + cmake_args += [ + "-Donnxruntime_USE_CUSTOM_DIRECTML=ON", + "-Ddml_EXTERNAL_PROJECT=ON", + ] + + if args.use_gdk: + cmake_args += [ + "-DCMAKE_TOOLCHAIN_FILE=" + os.path.join(source_dir, "cmake", "gdk_toolchain.cmake"), + "-DGDK_EDITION=" + args.gdk_edition, + "-DGDK_PLATFORM=" + args.gdk_platform, + "-Donnxruntime_BUILD_UNIT_TESTS=OFF", # gtest doesn't build for GDK + ] + if args.use_dml and not (args.dml_path or args.dml_external_project): + raise BuildError("You must set dml_path or dml_external_project when building with the GDK.") + elif not is_macOS(): + cmake_args.append( + "-Donnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS=" + + ("ON" if args.enable_external_custom_op_schemas else "OFF") + ) cmake_args += [ "-Donnxruntime_RUN_ONNX_TESTS=" + ("ON" if args.enable_onnx_tests else "OFF"), "-Donnxruntime_GENERATE_TEST_REPORTS=ON", @@ -1187,10 +460,7 @@ def generate_build_tree( else "OFF" ), "-Donnxruntime_REDUCED_OPS_BUILD=" + ("ON" if is_reduced_ops_build(args) else "OFF"), - "-Donnxruntime_USE_DML=" + ("ON" if args.use_dml else "OFF"), - "-Donnxruntime_USE_WINML=" + ("ON" if args.use_winml else "OFF"), "-Donnxruntime_BUILD_MS_EXPERIMENTAL_OPS=" + ("ON" if args.ms_experimental else "OFF"), - "-Donnxruntime_USE_TELEMETRY=" + ("ON" if args.use_telemetry else "OFF"), "-Donnxruntime_ENABLE_LTO=" + ("ON" if args.enable_lto else "OFF"), "-Donnxruntime_USE_ACL=" + ("ON" if args.use_acl else "OFF"), "-Donnxruntime_USE_ARMNN=" + ("ON" if args.use_armnn else "OFF"), @@ -1198,7 +468,6 @@ def generate_build_tree( "-Donnxruntime_ARMNN_BN_USE_CPU=" + ("OFF" if args.armnn_bn else "ON"), "-Donnxruntime_USE_JSEP=" + ("ON" if args.use_jsep else "OFF"), "-Donnxruntime_USE_WEBGPU=" + ("ON" if args.use_webgpu else "OFF"), - "-Donnxruntime_ENABLE_PIX_FOR_WEBGPU_EP=" + ("ON" if args.enable_pix_capture else "OFF"), "-Donnxruntime_USE_EXTERNAL_DAWN=" + ("ON" if args.use_external_dawn else "OFF"), # Training related flags "-Donnxruntime_ENABLE_NVTX_PROFILE=" + ("ON" if args.enable_nvtx_profile else "OFF"), @@ -1211,7 +480,6 @@ def generate_build_tree( "-Donnxruntime_BUILD_BENCHMARKS=" + ("ON" if args.build_micro_benchmarks else "OFF"), "-Donnxruntime_USE_ROCM=" + ("ON" if args.use_rocm else "OFF"), "-Donnxruntime_GCOV_COVERAGE=" + ("ON" if args.code_coverage else "OFF"), - "-Donnxruntime_USE_MPI=" + ("ON" if args.use_mpi else "OFF"), "-Donnxruntime_ENABLE_MEMORY_PROFILE=" + ("ON" if args.enable_memory_profile else "OFF"), "-Donnxruntime_ENABLE_CUDA_LINE_NUMBER_INFO=" + ("ON" if args.enable_cuda_line_info else "OFF"), "-Donnxruntime_USE_CUDA_NHWC_OPS=" + ("ON" if args.use_cuda and not args.disable_cuda_nhwc_ops else "OFF"), @@ -1228,8 +496,6 @@ def generate_build_tree( "-Donnxruntime_ENABLE_WEBASSEMBLY_DEBUG_INFO=" + ("ON" if args.enable_wasm_debug_info else "OFF"), "-Donnxruntime_ENABLE_WEBASSEMBLY_PROFILING=" + ("ON" if args.enable_wasm_profiling else "OFF"), "-Donnxruntime_ENABLE_LAZY_TENSOR=" + ("ON" if args.enable_lazy_tensor else "OFF"), - "-Donnxruntime_ENABLE_EXTERNAL_CUSTOM_OP_SCHEMAS=" - + ("ON" if args.enable_external_custom_op_schemas else "OFF"), "-Donnxruntime_ENABLE_CUDA_PROFILING=" + ("ON" if args.enable_cuda_profiling else "OFF"), "-Donnxruntime_ENABLE_ROCM_PROFILING=" + ("ON" if args.enable_rocm_profiling else "OFF"), "-Donnxruntime_USE_XNNPACK=" + ("ON" if args.use_xnnpack else "OFF"), @@ -1414,7 +680,8 @@ def generate_build_tree( add_default_definition(cmake_extra_defines, "VCPKG_TARGET_TRIPLET", triplet) # By default on Windows we currently support only cross compiling for ARM/ARM64 - if (args.arm64 or args.arm64ec or args.arm) and platform.architecture()[0] != "AMD64": + if is_windows() and (args.arm64 or args.arm64ec or args.arm) and platform.architecture()[0] != "AMD64": + # The onnxruntime_CROSS_COMPILING flag is deprecated. Prefer to use CMAKE_CROSSCOMPILING. add_default_definition(cmake_extra_defines, "onnxruntime_CROSS_COMPILING", "ON") if args.use_extensions: add_default_definition(cmake_extra_defines, "OPENCV_SKIP_SYSTEM_PROCESSOR_DETECTION", "ON") @@ -1429,8 +696,7 @@ def generate_build_tree( cmake_args.append("-DCMAKE_HIP_COMPILER_LAUNCHER=ccache") if args.external_graph_transformer_path: cmake_args.append("-Donnxruntime_EXTERNAL_TRANSFORMER_SRC_PATH=" + args.external_graph_transformer_path) - if args.use_winml: - cmake_args.append("-Donnxruntime_BUILD_WINML_TESTS=" + ("OFF" if args.skip_winml_tests else "ON")) + if args.use_dnnl: cmake_args.append("-Donnxruntime_DNNL_GPU_RUNTIME=" + args.dnnl_gpu_runtime) cmake_args.append("-Donnxruntime_DNNL_OPENCL_ROOT=" + args.dnnl_opencl_root) @@ -1493,14 +759,6 @@ def generate_build_tree( if armnn_libs and os.path.exists(armnn_libs): cmake_args += ["-Donnxruntime_ARMNN_LIBS=" + armnn_libs] - if mpi_home and os.path.exists(mpi_home): - if args.use_mpi: - cmake_args += ["-Donnxruntime_MPI_HOME=" + mpi_home] - else: - log.warning( - "mpi_home is supplied but use_mpi is set to false. Build will continue without linking MPI libraries." - ) - if nccl_home and os.path.exists(nccl_home): cmake_args += ["-Donnxruntime_NCCL_HOME=" + nccl_home] @@ -1513,8 +771,6 @@ def generate_build_tree( if cann_home and os.path.exists(cann_home): cmake_args += ["-Donnxruntime_CANN_HOME=" + cann_home] - if args.winml_root_namespace_override: - cmake_args += ["-Donnxruntime_WINML_NAMESPACE_OVERRIDE=" + args.winml_root_namespace_override] if args.use_openvino: cmake_args += [ "-Donnxruntime_USE_OPENVINO=ON", @@ -1539,9 +795,6 @@ def generate_build_tree( nvml_stub_path = cuda_home + "/lib64/stubs" cmake_args += ["-DCUDA_CUDA_LIBRARY=" + nvml_stub_path] - if args.use_preinstalled_eigen: - cmake_args += ["-Donnxruntime_USE_PREINSTALLED_EIGEN=ON", "-Deigen_SOURCE_PATH=" + args.eigen_path] - if args.nnapi_min_api: cmake_args += ["-Donnxruntime_NNAPI_MIN_API=" + str(args.nnapi_min_api)] @@ -1569,29 +822,6 @@ def generate_build_tree( if args.android_cpp_shared: cmake_args += ["-DANDROID_STL=c++_shared"] - if args.dml_path: - cmake_args += [ - "-Donnxruntime_USE_CUSTOM_DIRECTML=ON", - "-Ddml_INCLUDE_DIR=" + os.path.join(args.dml_path, "include"), - "-Ddml_LIB_DIR=" + os.path.join(args.dml_path, "lib"), - ] - - if args.dml_external_project: - cmake_args += [ - "-Donnxruntime_USE_CUSTOM_DIRECTML=ON", - "-Ddml_EXTERNAL_PROJECT=ON", - ] - - if args.use_gdk: - cmake_args += [ - "-DCMAKE_TOOLCHAIN_FILE=" + os.path.join(source_dir, "cmake", "gdk_toolchain.cmake"), - "-DGDK_EDITION=" + args.gdk_edition, - "-DGDK_PLATFORM=" + args.gdk_platform, - "-Donnxruntime_BUILD_UNIT_TESTS=OFF", # gtest doesn't build for GDK - ] - if args.use_dml and not (args.dml_path or args.dml_external_project): - raise BuildError("You must set dml_path or dml_external_project when building with the GDK.") - if is_macOS() and not args.android: add_default_definition(cmake_extra_defines, "CMAKE_OSX_ARCHITECTURES", args.osx_arch) if args.apple_deploy_target: @@ -1628,20 +858,22 @@ def generate_build_tree( # if args.use_jsep and args.use_webgpu: # raise BuildError("JSEP (--use_jsep) and WebGPU (--use_webgpu) cannot be enabled at the same time.") - if args.use_external_dawn and not args.use_webgpu: - raise BuildError("External Dawn (--use_external_dawn) must be enabled with WebGPU (--use_webgpu).") + if not args.use_webgpu: + if args.use_external_dawn: + raise BuildError("External Dawn (--use_external_dawn) must be enabled with WebGPU (--use_webgpu).") - if args.enable_pix_capture and (not args.use_webgpu or not is_windows()): - raise BuildError( - "Enable PIX Capture (--enable_pix_capture) must be enabled with WebGPU (--use_webgpu) on Windows" - ) + if is_windows(): + if args.enable_pix_capture: + raise BuildError( + "Enable PIX Capture (--enable_pix_capture) must be enabled with WebGPU (--use_webgpu) on Windows" + ) if args.use_snpe: cmake_args += ["-Donnxruntime_USE_SNPE=ON"] cmake_args += ["-Donnxruntime_USE_KLEIDIAI=" + ("OFF" if args.no_kleidiai else "ON")] - if args.macos or args.ios or args.visionos or args.tvos: + if is_macOS() and (args.macos or args.ios or args.visionos or args.tvos): # Note: Xcode CMake generator doesn't have a good support for Mac Catalyst yet. if args.macos == "Catalyst" and args.cmake_generator == "Xcode": raise BuildError("Xcode CMake generator ('--cmake_generator Xcode') doesn't support Mac Catalyst build.") @@ -1797,7 +1029,7 @@ def generate_build_tree( add_default_definition(cmake_extra_defines, "onnxruntime_USE_LOCK_FREE_QUEUE", "ON") if is_windows(): - if not args.ios and not args.android and not args.build_wasm: + if not args.android and not args.build_wasm: if args.use_cache: add_default_definition( cmake_extra_defines, @@ -1862,7 +1094,7 @@ def generate_build_tree( cxxflags = None ldflags = None cudaflags = [] - if is_windows() and not args.ios and not args.android and not args.build_wasm: + if is_windows() and not args.android and not args.build_wasm: njobs = number_of_parallel_jobs(args) if args.use_cuda: cudaflags.append("-allow-unsupported-compiler") @@ -1873,11 +1105,7 @@ def generate_build_tree( cflags += [f"/MP{njobs}"] # Setup default values for cflags/cxxflags/ldflags. # The values set here are purely for security and compliance purposes. ONNX Runtime should work fine without these flags. - if ( - (args.use_binskim_compliant_compile_flags or args.enable_address_sanitizer) - and not args.ios - and not args.android - ): + if (args.use_binskim_compliant_compile_flags or args.enable_address_sanitizer) and not args.android: if is_windows() and not args.build_wasm: cflags += ["/guard:cf", "/DWIN32", "/D_WINDOWS"] if not args.use_gdk: @@ -2375,7 +1603,7 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs): if args.android: run_android_tests(args, source_dir, build_dir, config, cwd) continue - elif args.ios: + elif is_macOS() and args.ios: run_ios_tests(args, source_dir, config, cwd) continue dll_path_list = [] @@ -2937,7 +2165,7 @@ def main(): if args.android and args.use_vcpkg and args.android_ndk_path is not None and os.path.exists(args.android_ndk_path): os.environ["ANDROID_NDK_HOME"] = args.android_ndk_path - if not is_windows(): + if not is_windows() and not is_macOS(): if not args.allow_running_as_root: is_root_user = os.geteuid() == 0 if is_root_user: @@ -2946,34 +2174,15 @@ def main(): ) cmake_extra_defines = normalize_arg_list(args.cmake_extra_defines) - cross_compiling = args.arm or args.arm64 or args.arm64ec or args.android - - if args.enable_address_sanitizer: - # Disable ONNX Runtime's builtin memory checker - args.disable_memleak_checker = True # When this flag is enabled, it is possible ONNXRuntime shared library is build separately, expecting some compatible EP - # shared lib being build in a seperate process. So we skip the testing if none of the primary EPs are built with ONNXRuntime + # shared lib being build in a separate process. So we skip the testing if none of the primary EPs are built with ONNXRuntime # shared lib if args.enable_generic_interface and not ( args.use_tensorrt or args.use_openvino or args.use_vitisai or (args.use_qnn and args.use_qnn != "static_lib") ): args.test = False - # If there was no explicit argument saying what to do, default - # to update, build and test (for native builds). - if not (args.update or args.clean or args.build or args.test or args.gen_doc): - log.debug("Defaulting to running update, build [and test for native builds].") - args.update = True - args.build = True - if cross_compiling: - args.test = args.android_abi == "x86_64" or args.android_abi == "arm64-v8a" - else: - args.test = True - - if args.skip_tests: - args.test = False - if args.use_tensorrt: args.use_cuda = True @@ -3032,15 +2241,12 @@ def main(): if args.code_coverage and not args.android: raise BuildError("Using --code_coverage requires --android") - if args.gen_api_doc and len(args.config) != 1: - raise BuildError("Using --get-api-doc requires a single build config") - # Disabling unit tests for GPU on nuget creation if args.use_openvino and args.use_openvino != "CPU" and args.build_nuget: args.test = False # GDK builds don't support testing - if args.use_gdk: + if is_windows() and args.use_gdk: args.test = False # enable_training is a higher level flag that enables all training functionality. @@ -3065,7 +2271,6 @@ def main(): if args.use_cuda: cuda_home, cudnn_home = setup_cuda_vars(args) - mpi_home = args.mpi_home nccl_home = args.nccl_home snpe_root = args.snpe_root @@ -3216,9 +2421,6 @@ def main(): if args.use_rocm and args.rocm_version is None: args.rocm_version = "" - if args.enable_external_custom_op_schemas and not is_linux(): - raise BuildError("Registering external custom op schemas is only supported on Linux.") - generate_build_tree( cmake_path, source_dir, @@ -3226,7 +2428,6 @@ def main(): cuda_home, cudnn_home, rocm_home, - mpi_home, nccl_home, tensorrt_home, migraphx_home, @@ -3247,8 +2448,9 @@ def main(): if args.clean: clean_targets(cmake_path, build_dir, configs) - # if using DML, perform initial nuget package restore - setup_dml_build(args, cmake_path, build_dir, configs) + if is_windows(): + # if using DML, perform initial nuget package restore + setup_dml_build(args, cmake_path, build_dir, configs) if args.build: if args.parallel < 0: @@ -3318,9 +2520,9 @@ def main(): args.use_openvino, args.use_tensorrt, args.use_dnnl, - args.use_winml, + getattr(args, "use_winml", False), args.use_qnn, - args.use_dml, + getattr(args, "use_dml", False), args.enable_training_apis, normalize_arg_list(args.msbuild_extra_options), ) @@ -3347,14 +2549,6 @@ def main(): # documentation generation as a separate task post-build generate_documentation(source_dir, build_dir, configs, args.gen_doc == "validate") - if args.gen_api_doc and (args.build or args.test): - print("Generating Python doc for ORTModule...") - docbuild_dir = os.path.join(source_dir, "tools", "doc") - run_subprocess( - ["bash", "builddoc.sh", os.path.dirname(sys.executable), source_dir, build_dir, args.config[0]], - cwd=docbuild_dir, - ) - log.info("Build complete") diff --git a/tools/ci_build/build_args.py b/tools/ci_build/build_args.py new file mode 100644 index 0000000000000..a54500c176a87 --- /dev/null +++ b/tools/ci_build/build_args.py @@ -0,0 +1,941 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License.import argparse +import argparse +import os +import platform +import shlex +import sys +import warnings + +from util import ( + is_macOS, + is_windows, +) + + +def _str_to_bool(s: str) -> bool: + """Convert string to bool (in argparse context) using match/case.""" + match s.lower(): + case "true": + return True + case "false": + return False + case _: + raise ValueError(f"Invalid boolean value: {s!r}. Use 'true' or 'false'.") + return False + + +# --- Argument Verification Helpers --- +def _qnn_verify_library_kind(library_kind: str) -> str: + """Verifies the library kind for the QNN Execution Provider.""" + choices = ["shared_lib", "static_lib"] + if library_kind not in choices: + print("\nYou have specified an invalid library kind for QNN EP.") + print(f"The invalid library kind was: {library_kind}") + print("Provide a library kind from the following options: ", choices) + print(f"Example: --use_qnn {choices[0]}") + sys.exit("Incorrect build configuration") + return library_kind + + +def _openvino_verify_device_type(device_read: str) -> str: + """Verifies the device type string for the OpenVINO Execution Provider.""" + choices = ["CPU", "GPU", "NPU"] + choices1 = [ + "CPU_NO_PARTITION", + "GPU_NO_PARTITION", + "NPU_NO_PARTITION", + "NPU_NO_CPU_FALLBACK", + ] + status_hetero = True + res = False + if device_read in choices: + res = True + elif device_read in choices1: + res = True + elif device_read.startswith(("HETERO:", "MULTI:", "AUTO:")): + res = True + parts = device_read.split(":") + if len(parts) < 2 or not parts[1]: + print("Hetero/Multi/Auto mode requires devices to be specified after the colon.") + status_hetero = False + else: + comma_separated_devices = parts[1].split(",") + if len(comma_separated_devices) < 2: + print("At least two devices required in Hetero/Multi/Auto Mode") + status_hetero = False + dev_options = ["CPU", "GPU", "NPU"] + for dev in comma_separated_devices: + if dev not in dev_options: + status_hetero = False + print(f"Invalid device '{dev}' found in Hetero/Multi/Auto specification.") + break + + def invalid_hetero_build() -> None: + print("\nIf trying to build Hetero/Multi/Auto, specify the supported devices along with it.\n") + print("Specify the keyword HETERO or MULTI or AUTO followed by a colon and comma-separated devices ") + print("in the order of priority you want to build (e.g., HETERO:GPU,CPU).\n") + print("The different hardware devices that can be added are ['CPU','GPU','NPU'] \n") + print("An example of how to specify the hetero build type: --use_openvino HETERO:GPU,CPU \n") + print("An example of how to specify the MULTI build type: --use_openvino MULTI:GPU,CPU \n") + print("An example of how to specify the AUTO build type: --use_openvino AUTO:GPU,CPU \n") + sys.exit("Wrong Build Type selected") + + if res is False: + print("\nYou have selected wrong configuration for the build.") + print("Pick the build type for specific Hardware Device from following options: ", choices) + print("(or) from the following options with graph partitioning disabled: ", choices1) + print("\n") + if not (device_read.startswith(("HETERO", "MULTI", "AUTO"))): + invalid_hetero_build() # Will exit + sys.exit("Wrong Build Type selected") # Should not be reached if invalid_hetero_build exits + + if status_hetero is False: + invalid_hetero_build() # Will exit + + return device_read + + +# --- Argument Grouping Functions --- + + +def add_core_build_args(parser: argparse.ArgumentParser) -> None: + """Adds core build process arguments.""" + parser.add_argument("--build_dir", required=True, help="Path to the build directory.") + parser.add_argument( + "--config", + nargs="+", + default=["Debug"], + choices=["Debug", "MinSizeRel", "Release", "RelWithDebInfo"], + help="Configuration(s) to build.", + ) + parser.add_argument("--update", action="store_true", help="Update makefiles.") + parser.add_argument("--build", action="store_true", help="Build.") + parser.add_argument( + "--clean", action="store_true", help="Run 'cmake --build --target clean' for the selected config/s." + ) + parser.add_argument( + "--parallel", + nargs="?", + const="0", + default="1", + type=int, + help="Use parallel build. Optional value specifies max jobs (0=num CPUs).", + ) + parser.add_argument("--target", help="Build a specific CMake target (e.g., winml_dll).") + parser.add_argument( + "--compile_no_warning_as_error", + action="store_true", + help="Prevent warnings from being treated as errors during compile. Only works for cmake targets that honor the COMPILE_WARNING_AS_ERROR property", + ) + parser.add_argument("--build_shared_lib", action="store_true", help="Build a shared library for ONNXRuntime.") + parser.add_argument( + "--build_apple_framework", action="store_true", help="Build a macOS/iOS framework for ONNXRuntime." + ) + parser.add_argument("--enable_lto", action="store_true", help="Enable Link Time Optimization (LTO).") + parser.add_argument("--use_cache", action="store_true", help="Use ccache in CI") + parser.add_argument( + "--use_binskim_compliant_compile_flags", + action="store_true", + help="[MS Internal] Use preset compile flags for BinSkim compliance.", + ) + + +def add_cmake_build_config_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments related to CMake and general build system configuration.""" + parser.add_argument( + "--cmake_extra_defines", + nargs="+", + action="append", + help="Extra CMake definitions (-D=). Provide as =.", + ) + parser.add_argument("--cmake_path", default="cmake", help="Path to the CMake executable.") + parser.add_argument( + "--cmake_generator", + choices=[ + "MinGW Makefiles", + "Ninja", + "NMake Makefiles", + "NMake Makefiles JOM", + "Unix Makefiles", + "Visual Studio 17 2022", + "Xcode", + ], + default=None, # Will be set later based on OS and WASM + help="Specify the generator for CMake.", + ) + parser.add_argument( + "--use_vcpkg", action="store_true", help="Use vcpkg for dependencies (requires CMAKE_TOOLCHAIN_FILE)." + ) + parser.add_argument( + "--use_vcpkg_ms_internal_asset_cache", action="store_true", help="[MS Internal] Use internal vcpkg asset cache." + ) + parser.add_argument("--skip_submodule_sync", action="store_true", help="Skip 'git submodule update'.") + + +def add_testing_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments related to running tests.""" + parser.add_argument("--test", action="store_true", help="Run unit tests.") + parser.add_argument("--skip_tests", action="store_true", help="Skip all tests.") + parser.add_argument( + "--ctest_path", + default="ctest", + help="Path to CTest. Empty string uses script to drive tests.", + ) + parser.add_argument( + "--enable_onnx_tests", + action="store_true", + help="Run onnx_test_runner against test data. Only used in ONNX Runtime's CI pipelines", + ) + parser.add_argument("--path_to_protoc_exe", help="Path to protoc executable.") + parser.add_argument("--fuzz_testing", action="store_true", help="Enable Fuzz testing.") + parser.add_argument( + "--enable_symbolic_shape_infer_tests", + action="store_true", + help="Run symbolic shape inference tests.", + ) + parser.add_argument("--skip_onnx_tests", action="store_true", help="Explicitly disable ONNX related tests.") + parser.add_argument("--skip_winml_tests", action="store_true", help="Explicitly disable WinML related tests.") + parser.add_argument("--skip_nodejs_tests", action="store_true", help="Explicitly disable Node.js binding tests.") + parser.add_argument("--test_all_timeout", default="10800", help="Timeout for onnxruntime_test_all (seconds).") + parser.add_argument("--enable_transformers_tool_test", action="store_true", help="Enable transformers tool test.") + parser.add_argument("--build_micro_benchmarks", action="store_true", help="Build ONNXRuntime micro-benchmarks.") + parser.add_argument("--code_coverage", action="store_true", help="Generate code coverage report (Android only).") + + +def add_training_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments related to ONNX Runtime Training.""" + parser.add_argument( + "--enable_training", + action="store_true", + help="Enable full ORT Training (ORTModule, Training APIs).", + ) + parser.add_argument("--enable_training_apis", action="store_true", help="Enable ORT Training APIs.") + parser.add_argument("--enable_training_ops", action="store_true", help="Enable training ops in inference graph.") + parser.add_argument("--enable_nccl", action="store_true", help="Enable NCCL for distributed training.") + parser.add_argument("--nccl_home", help="Path to NCCL installation directory.") + + +def add_general_profiling_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments related to general (non-EP specific) profiling.""" + parser.add_argument("--enable_memory_profile", action="store_true", help="Enable memory profiling.") + + +def add_debugging_sanitizer_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments related to debugging, sanitizers, and compliance.""" + parser.add_argument( + "--enable_address_sanitizer", action="store_true", help="Enable Address Sanitizer (ASan) (Linux/macOS/Windows)." + ) + + +def add_documentation_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments related to documentation generation.""" + parser.add_argument( + "--gen_doc", + nargs="?", + const="yes", + type=str, + help="Generate operator/type docs. Use '--gen_doc validate' to check against /docs.", + ) + + +def add_cross_compile_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for cross-compiling to non-Windows target CPU architectures.""" + parser.add_argument( + "--rv64", + action="store_true", + help="[cross-compiling] Target RISC-V 64-bit.", + ) + parser.add_argument( + "--riscv_toolchain_root", + type=str, + default="", + help="Path to RISC-V toolchain root.", + ) + parser.add_argument( + "--riscv_qemu_path", + type=str, + default="", + help="Path to RISC-V qemu executable.", + ) + + +def add_android_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for Android platform builds.""" + parser.add_argument("--android", action="store_true", help="Build for Android.") + parser.add_argument( + "--android_abi", + default="arm64-v8a", + choices=["armeabi-v7a", "arm64-v8a", "x86", "x86_64"], + help="Target Android ABI.", + ) + parser.add_argument("--android_api", type=int, default=27, help="Android API Level (e.g., 21).") + parser.add_argument( + "--android_sdk_path", type=str, default=os.environ.get("ANDROID_HOME", ""), help="Path to Android SDK." + ) + parser.add_argument( + "--android_ndk_path", type=str, default=os.environ.get("ANDROID_NDK_HOME", ""), help="Path to Android NDK." + ) + parser.add_argument( + "--android_cpp_shared", + action="store_true", + help="Link shared libc++ instead of static (default).", + ) + parser.add_argument( + "--android_run_emulator", action="store_true", help="Start an Android emulator if needed for tests." + ) + + +def add_apple_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for Apple platform builds (iOS, macOS, visionOS, tvOS).""" + platform_group = parser.add_mutually_exclusive_group() + platform_group.add_argument("--ios", action="store_true", help="Build for iOS.") + platform_group.add_argument("--visionos", action="store_true", help="Build for visionOS.") + platform_group.add_argument("--tvos", action="store_true", help="Build for tvOS.") + platform_group.add_argument( + "--macos", + choices=["MacOSX", "Catalyst"], + help="Target platform for macOS build (requires --build_apple_framework).", + ) + + parser.add_argument("--apple_sysroot", default="", help="Specify the macOS platform SDK location name.") + parser.add_argument( + "--ios_toolchain_file", + default="", + help="Path to iOS CMake toolchain file (defaults to cmake/onnxruntime_ios.toolchain.cmake).", + ) + parser.add_argument( + "--visionos_toolchain_file", + default="", + help="Path to visionOS CMake toolchain file (defaults to cmake/onnxruntime_visionos.toolchain.cmake).", + ) + parser.add_argument( + "--tvos_toolchain_file", + default="", + help="Path to tvOS CMake toolchain file (defaults to cmake/onnxruntime_tvos.toolchain.cmake).", + ) + parser.add_argument("--xcode_code_signing_team_id", default="", help="Development team ID for Xcode code signing.") + parser.add_argument( + "--xcode_code_signing_identity", default="", help="Development identity for Xcode code signing." + ) + parser.add_argument( + "--use_xcode", + action="store_const", + const="Xcode", + dest="cmake_generator", # Overwrites the general cmake_generator if specified + help="Use Xcode CMake generator (macOS only, non-Catalyst).", + ) + parser.add_argument( + "--osx_arch", + default="arm64" if platform.machine() == "arm64" else "x86_64", + choices=["arm64", "arm64e", "x86_64"], + help="Target architecture for macOS/iOS (macOS host only).", + ) + parser.add_argument( + "--apple_deploy_target", + type=str, + help="Minimum deployment target version (e.g., 11.0 for macOS).", + ) + + +def add_webassembly_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for WebAssembly (WASM) platform builds.""" + parser.add_argument("--build_wasm", action="store_true", help="Build for WebAssembly.") + parser.add_argument("--build_wasm_static_lib", action="store_true", help="Build WebAssembly static library.") + parser.add_argument("--emsdk_version", default="4.0.4", help="Specify version of emsdk.") + parser.add_argument("--enable_wasm_simd", action="store_true", help="Enable WebAssembly SIMD.") + parser.add_argument("--enable_wasm_relaxed_simd", action="store_true", help="Enable WebAssembly Relaxed SIMD.") + parser.add_argument("--enable_wasm_threads", action="store_true", help="Enable WebAssembly multi-threading.") + parser.add_argument("--enable_wasm_memory64", action="store_true", help="Enable WebAssembly 64-bit memory.") + parser.add_argument("--disable_wasm_exception_catching", action="store_true", help="Disable exception catching.") + parser.add_argument( + "--enable_wasm_api_exception_catching", + action="store_true", + help="Catch exceptions only at top-level API calls.", + ) + parser.add_argument( + "--enable_wasm_exception_throwing_override", + action="store_true", + help="Override default behavior to allow throwing exceptions even when catching is generally disabled.", + ) + parser.add_argument("--wasm_run_tests_in_browser", action="store_true", help="Run WASM tests in a browser.") + parser.add_argument( + "--enable_wasm_profiling", action="store_true", help="Enable WASM profiling and preserve function names." + ) + parser.add_argument("--enable_wasm_debug_info", action="store_true", help="Build WASM with DWARF debug info.") + parser.add_argument("--wasm_malloc", help="Specify memory allocator for WebAssembly (e.g., dlmalloc).") + parser.add_argument( + "--emscripten_settings", + nargs="+", + action="append", + help="Extra emscripten settings (-s =). Provide as =.", + ) + + +def add_gdk_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for GDK (Xbox) platform builds.""" + parser.add_argument("--use_gdk", action="store_true", help="Build with the GDK toolchain.") + default_gdk_edition = "" + gdk_latest_env = os.environ.get("GameDKLatest", "") # noqa: SIM112 + if gdk_latest_env: + try: + default_gdk_edition = os.path.basename(os.path.normpath(gdk_latest_env)) + except Exception as e: + warnings.warn(f"Failed to determine GDK edition from GameDKLatest env var: {e}") + parser.add_argument( + "--gdk_edition", + default=default_gdk_edition, + help="Specific GDK edition to build with (defaults to latest installed via GameDKLatest env var).", + ) + parser.add_argument("--gdk_platform", default="Scarlett", help="GDK target platform (e.g., Scarlett, XboxOne).") + + +def add_windows_specific_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments specific to Windows builds or Windows cross-compilation.""" + # Build tools / config + parser.add_argument("--msvc_toolset", help="MSVC toolset version (e.g., 14.11). Must be >=14.40") + parser.add_argument("--windows_sdk_version", help="Windows SDK version (e.g., 10.0.19041.0).") + parser.add_argument("--enable_msvc_static_runtime", action="store_true", help="Statically link MSVC runtimes.") + parser.add_argument("--use_telemetry", action="store_true", help="Enable telemetry (official builds only).") + + # Cross-compilation targets hosted on Windows + parser.add_argument( + "--x86", + action="store_true", + help="[Windows cross-compiling] Target Windows x86.", + ) + parser.add_argument( + "--arm", + action="store_true", + help="[Windows cross-compiling] Target Windows ARM.", + ) + parser.add_argument( + "--arm64", + action="store_true", + help="[Windows cross-compiling] Target Windows ARM64.", + ) + parser.add_argument( + "--arm64ec", + action="store_true", + help="[Windows cross-compiling] Target Windows ARM64EC.", + ) + parser.add_argument( + "--buildasx", + action="store_true", + help="[Windows cross-compiling] Create ARM64X Binary.", + ) + + parser.add_argument( + "--disable_memleak_checker", + action="store_true", + help="Disable memory leak checker (enabled by default in Debug builds).", + ) + parser.add_argument( + "--enable_pix_capture", action="store_true", help="Enable Pix support for GPU debugging (requires D3D12)." + ) + + parser.add_argument( + "--enable_wcos", + action="store_true", + help="Build for Windows Core OS. Link to Windows umbrella libraries instead of kernel32.lib.", + ) + + add_gdk_args(parser) + + # --- WinML --- + winml_group = parser.add_argument_group("WinML API (Windows)") + winml_group.add_argument( + "--use_winml", action="store_true", help="Enable WinML API (Windows). Requires --enable_wcos." + ) + winml_group.add_argument( + "--winml_root_namespace_override", type=str, help="Override the namespace WinML builds into." + ) + + +def add_linux_specific_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments specific to Linux builds.""" + parser.add_argument( + "--allow_running_as_root", + action="store_true", + help="Allow build script to run as root (disallowed by default).", + ) + parser.add_argument( + "--enable_external_custom_op_schemas", + action="store_true", + help="Enable loading custom op schemas from external shared libraries (Ubuntu only).", + ) + + +def add_dependency_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments related to external dependencies.""" + parser.add_argument("--use_full_protobuf", action="store_true", help="Use the full (non-lite) protobuf library.") + parser.add_argument("--use_mimalloc", action="store_true", help="Use mimalloc memory allocator.") + parser.add_argument( + "--external_graph_transformer_path", type=str, help="Path to external graph transformer directory." + ) + + +def add_extension_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments related to ONNX Runtime Extensions.""" + parser.add_argument( + "--use_extensions", + action="store_true", + help="Enable ONNX Runtime Extensions (uses submodule by default).", + ) + parser.add_argument( + "--extensions_overridden_path", + type=str, + help="Path to an external ONNX Runtime Extensions source directory.", + ) + + +def add_size_reduction_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for reducing the binary size.""" + parser.add_argument( + "--minimal_build", + default=None, + nargs="*", + type=str.lower, + help="Create a minimal build supporting only ORT format models. " + "Options: 'extended' (runtime kernel compilation), 'custom_ops'. " + "e.g., '--minimal_build extended custom_ops'. RTTI disabled automatically.", + ) + parser.add_argument( + "--include_ops_by_config", + type=str, + help="Include only ops specified in the config file (see docs/Reduced_Operator_Kernel_build.md).", + ) + parser.add_argument( + "--enable_reduced_operator_type_support", + action="store_true", + help="Further reduce size by limiting operator data types based on --include_ops_by_config file.", + ) + parser.add_argument("--disable_contrib_ops", action="store_true", help="Disable contrib operators.") + parser.add_argument("--disable_ml_ops", action="store_true", help="Disable traditional ML operators.") + parser.add_argument("--disable_rtti", action="store_true", help="Disable Run-Time Type Information (RTTI).") + parser.add_argument( + "--disable_types", + nargs="+", + default=[], + choices=["float8", "optional", "sparsetensor"], + help="Disable selected data types.", + ) + parser.add_argument( + "--disable_exceptions", + action="store_true", + help="Disable exceptions (requires --minimal_build).", + ) + + +def add_python_binding_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for Python bindings.""" + parser.add_argument("--enable_pybind", action="store_true", help="Enable Python bindings.") + parser.add_argument("--build_wheel", action="store_true", help="Build Python wheel package.") + parser.add_argument( + "--wheel_name_suffix", + help="Suffix for wheel name (used for nightly builds).", + ) + parser.add_argument("--skip-keras-test", action="store_true", help="Skip Keras-related tests.") + + +def add_csharp_binding_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for C# bindings.""" + parser.add_argument( + "--build_csharp", + action="store_true", + help="Build C# DLL and NuGet package (CI usage). Use --build_nuget for local build.", + ) + parser.add_argument( + "--build_nuget", + action="store_true", + help="Build C# DLL and NuGet package locally (Windows/Linux only).", + ) + parser.add_argument( + "--msbuild_extra_options", + nargs="+", + action="append", + help="Extra MSBuild properties (/p:key=value). Provide as key=value.", + ) + + +def add_java_binding_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for Java bindings.""" + parser.add_argument("--build_java", action="store_true", help="Build Java bindings.") + + +def add_nodejs_binding_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for Node.js bindings.""" + parser.add_argument("--build_nodejs", action="store_true", help="Build Node.js binding and NPM package.") + # Note: --skip_nodejs_tests is handled in add_testing_args + + +def add_objc_binding_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for Objective-C bindings.""" + parser.add_argument("--build_objc", action="store_true", help="Build Objective-C binding.") + + +def add_execution_provider_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for enabling various Execution Providers (EPs).""" + # --- CUDA --- + cuda_group = parser.add_argument_group("CUDA Execution Provider") + cuda_group.add_argument("--use_cuda", action="store_true", help="Enable CUDA EP.") + cuda_group.add_argument("--cuda_version", help="CUDA toolkit version (e.g., 11.8). Auto-detect if omitted.") + cuda_group.add_argument("--cuda_home", help="Path to CUDA toolkit (uses CUDA_HOME env var if unset).") + cuda_group.add_argument("--cudnn_home", help="Path to cuDNN (uses CUDNN_HOME env var if unset).") + cuda_group.add_argument("--enable_cuda_line_info", action="store_true", help="Enable CUDA line info for debugging.") + cuda_group.add_argument("--enable_cuda_nhwc_ops", action="store_true", help="[Deprecated] Default enabled.") + cuda_group.add_argument("--disable_cuda_nhwc_ops", action="store_true", help="Disable CUDA NHWC layout ops.") + cuda_group.add_argument("--enable_cuda_minimal_build", action="store_true", help="Enable CUDA minimal build.") + cuda_group.add_argument( + "--nvcc_threads", + nargs="?", + default=-1, # -1 signifies auto-detect based on jobs/memory + type=int, + help="Max NVCC threads per parallel job (-1=auto).", + ) + # CUDA-specific profiling + cuda_group.add_argument( + "--enable_nvtx_profile", action="store_true", help="Enable NVTX profile markers for CUDA EP." + ) + cuda_group.add_argument( + "--enable_cuda_profiling", + action="store_true", + help="Enable CUDA kernel profiling (requires CUPTI in PATH).", + ) + + # --- ROCm --- + rocm_group = parser.add_argument_group("ROCm Execution Provider") + rocm_group.add_argument("--use_rocm", action="store_true", help="Enable ROCm EP.") + rocm_group.add_argument("--rocm_version", help="ROCm stack version.") + rocm_group.add_argument("--rocm_home", help="Path to ROCm installation directory.") + # ROCm-specific profiling + rocm_group.add_argument( + "--enable_rocm_profiling", + action="store_true", + help="Enable ROCm kernel profiling.", + ) + + # --- DNNL (formerly MKL-DNN / oneDNN) --- + dnnl_group = parser.add_argument_group("DNNL Execution Provider") + dnnl_group.add_argument("--use_dnnl", action="store_true", help="Enable DNNL EP.") + dnnl_group.add_argument( + "--dnnl_gpu_runtime", + action="store", + default="", + type=str.lower, + choices=["ocl", ""], + help="DNNL GPU backend (e.g., ocl for OpenCL).", + ) + dnnl_group.add_argument("--dnnl_opencl_root", action="store", default="", help="Path to OpenCL SDK (for DNNL GPU).") + dnnl_group.add_argument( + "--dnnl_aarch64_runtime", + action="store", + default="", + type=str.lower, + choices=["acl", ""], + help="DNNL AArch64 backend (e.g., acl for Arm Compute Library).", + ) + dnnl_group.add_argument( + "--dnnl_acl_root", action="store", default="", help="Path to Arm Compute Library (ACL) root." + ) + + # --- OpenVINO --- + openvino_group = parser.add_argument_group("OpenVINO Execution Provider") + openvino_group.add_argument( + "--use_openvino", + nargs="?", + const="CPU", # Default device if only flag is present + type=_openvino_verify_device_type, + help="Enable OpenVINO EP for specific hardware (e.g., CPU, GPU, NPU, HETERO:GPU,CPU).", + ) + + # --- TensorRT --- + trt_group = parser.add_argument_group("TensorRT Execution Provider") + trt_group.add_argument("--use_tensorrt", action="store_true", help="Enable TensorRT EP.") + trt_group.add_argument( + "--use_tensorrt_builtin_parser", + action="store_true", + default=True, + help="Use TensorRT internal ONNX parser (default).", + ) + trt_group.add_argument("--use_tensorrt_oss_parser", action="store_true", help="Use TensorRT OSS ONNX parser.") + trt_group.add_argument("--tensorrt_home", help="Path to TensorRT installation directory.") + + # --- DirectML --- + dml_group = parser.add_argument_group("DirectML Execution Provider (Windows)") + dml_group.add_argument("--use_dml", action="store_true", help="Enable DirectML EP (Windows).") + dml_group.add_argument("--dml_path", type=str, default="", help="Path to custom DirectML SDK.") + dml_group.add_argument("--dml_external_project", action="store_true", help="Build DirectML as an external project.") + + # --- NNAPI --- + nnapi_group = parser.add_argument_group("NNAPI Execution Provider (Android)") + nnapi_group.add_argument("--use_nnapi", action="store_true", help="Enable NNAPI EP (Android).") + nnapi_group.add_argument("--nnapi_min_api", type=int, help="Minimum Android API level for NNAPI (>= 27).") + + # --- CoreML --- + coreml_group = parser.add_argument_group("CoreML Execution Provider (Apple)") + coreml_group.add_argument("--use_coreml", action="store_true", help="Enable CoreML EP (Apple platforms).") + + # --- QNN --- + qnn_group = parser.add_argument_group("QNN Execution Provider (Qualcomm)") + qnn_group.add_argument( + "--use_qnn", + nargs="?", + const="shared_lib", # Default linkage if only flag is present + type=_qnn_verify_library_kind, + help="Enable QNN EP. Optionally specify 'shared_lib' (default) or 'static_lib'.", + ) + qnn_group.add_argument("--qnn_home", help="Path to QNN SDK directory.") + + # --- SNPE --- + snpe_group = parser.add_argument_group("SNPE Execution Provider (Qualcomm)") + snpe_group.add_argument("--use_snpe", action="store_true", help="Enable SNPE EP.") + snpe_group.add_argument("--snpe_root", help="Path to SNPE SDK root directory.") + + # --- Vitis-AI --- + vitis_group = parser.add_argument_group("Vitis-AI Execution Provider (Xilinx)") + vitis_group.add_argument("--use_vitisai", action="store_true", help="Enable Vitis-AI EP.") + + # --- ArmNN --- + armnn_group = parser.add_argument_group("ArmNN Execution Provider") + armnn_group.add_argument("--use_armnn", action="store_true", help="Enable ArmNN EP.") + armnn_group.add_argument("--armnn_relu", action="store_true", help="Use ArmNN Relu implementation.") + armnn_group.add_argument("--armnn_bn", action="store_true", help="Use ArmNN BatchNormalization implementation.") + armnn_group.add_argument("--armnn_home", help="Path to ArmNN home directory.") + armnn_group.add_argument("--armnn_libs", help="Path to ArmNN libraries directory.") + + # --- ACL (Arm Compute Library) --- + acl_group = parser.add_argument_group("ACL Execution Provider") + acl_group.add_argument("--use_acl", action="store_true", help="Enable ACL EP (ARM architectures).") + acl_group.add_argument("--acl_home", help="Path to ACL home directory.") + acl_group.add_argument("--acl_libs", help="Path to ACL libraries directory.") + acl_group.add_argument( + "--no_kleidiai", action="store_true", help="Disable KleidiAI integration (used with ACL/ArmNN)." + ) + + # --- RKNPU --- + rknpu_group = parser.add_argument_group("RKNPU Execution Provider") + rknpu_group.add_argument("--use_rknpu", action="store_true", help="Enable RKNPU EP.") + + # --- CANN (Huawei Ascend) --- + cann_group = parser.add_argument_group("CANN Execution Provider") + cann_group.add_argument("--use_cann", action="store_true", help="Enable CANN EP.") + cann_group.add_argument("--cann_home", help="Path to CANN installation directory.") + + # --- MIGraphX (AMD) --- + migx_group = parser.add_argument_group("MIGraphX Execution Provider") + migx_group.add_argument("--use_migraphx", action="store_true", help="Enable MIGraphX EP.") + migx_group.add_argument("--migraphx_home", help="Path to MIGraphX installation directory.") + + # --- WebNN --- + webnn_group = parser.add_argument_group("WebNN Execution Provider") + webnn_group.add_argument("--use_webnn", action="store_true", help="Enable WebNN EP.") + + # --- JSEP (JavaScript EP for WASM) --- + jsep_group = parser.add_argument_group("JSEP Execution Provider (WebAssembly)") + jsep_group.add_argument("--use_jsep", action="store_true", help="Enable JavaScript EP (used with WebAssembly).") + + # --- WebGPU --- + webgpu_group = parser.add_argument_group("WebGPU Execution Provider") + webgpu_group.add_argument("--use_webgpu", action="store_true", help="Enable WebGPU EP.") + webgpu_group.add_argument( + "--use_external_dawn", action="store_true", help="Use external Dawn dependency for WebGPU." + ) + + # --- XNNPACK --- + xnn_group = parser.add_argument_group("XNNPACK Execution Provider") + xnn_group.add_argument("--use_xnnpack", action="store_true", help="Enable XNNPACK EP.") + + # --- VSINPU (VeriSilicon NPU) --- + vsi_group = parser.add_argument_group("VSINPU Execution Provider") + vsi_group.add_argument("--use_vsinpu", action="store_true", help="Enable VSINPU EP.") + + # --- Azure --- + azure_group = parser.add_argument_group("Azure Execution Provider") + azure_group.add_argument("--use_azure", action="store_true", help="Enable Azure EP.") + + +def add_other_feature_args(parser: argparse.ArgumentParser) -> None: + """Adds arguments for other miscellaneous features.""" + parser.add_argument("--enable_lazy_tensor", action="store_true", help="Enable ORT backend for PyTorch LazyTensor.") + parser.add_argument("--ms_experimental", action="store_true", help="Build Microsoft experimental operators.") + parser.add_argument( + "--enable_msinternal", action="store_true", help="[MS Internal] Enable Microsoft internal build features." + ) + parser.add_argument( + "--use_triton_kernel", action="store_true", help="Use Triton compiled kernels (requires Triton)." + ) + parser.add_argument("--use_lock_free_queue", action="store_true", help="Use lock-free task queue for threadpool.") + parser.add_argument( + "--enable_generic_interface", + action="store_true", + help="Build ORT shared lib with compatible bridge for primary EPs (TRT, OV, QNN, VitisAI), excludes tests.", + ) + + +def is_cross_compiling(args: argparse.Namespace) -> bool: + return any( + [ + # Check existence before accessing for conditionally added args + getattr(args, "x86", False), + getattr(args, "arm", False), + getattr(args, "arm64", False), + getattr(args, "arm64ec", False), + args.rv64, # General cross-compile arg + args.android, + # Check existence for macOS/Apple specific args + getattr(args, "ios", False), + getattr(args, "visionos", False), + getattr(args, "tvos", False), + args.build_wasm, + getattr(args, "use_gdk", False), # GDK args added conditionally + ] + ) + + +# --- Main Argument Parsing Function --- +def parse_arguments() -> argparse.Namespace: + """Parses command line arguments for the ONNX Runtime build.""" + + class Parser(argparse.ArgumentParser): + # override argument file line parsing behavior - allow multiple arguments per line and handle quotes + def convert_arg_line_to_args(self, arg_line: str) -> list[str]: # Use list[str] for Python 3.9+ + return shlex.split(arg_line) + + parser = Parser( + description="ONNXRuntime CI build driver.", + usage=""" + Default behavior is --update --build --test for native architecture builds. + Default behavior is --update --build for cross-compiled builds. + + The Update phase will update git submodules, and run cmake to generate makefiles. + The Build phase will build all projects. + The Test phase will run all unit tests, and optionally the ONNX tests. + + Use the individual flags (--update, --build, --test) to only run specific stages. + """, + fromfile_prefix_chars="@", # Allow args from file (@filename) + ) + + # Add arguments by category + add_core_build_args(parser) + add_cmake_build_config_args(parser) + add_testing_args(parser) + add_training_args(parser) + add_general_profiling_args(parser) + add_debugging_sanitizer_args(parser) + add_documentation_args(parser) + add_cross_compile_args(parser) # Non-Windows cross-compile args + add_android_args(parser) + add_webassembly_args(parser) + add_dependency_args(parser) + add_extension_args(parser) + add_size_reduction_args(parser) + + # Language Bindings + add_python_binding_args(parser) + add_csharp_binding_args(parser) + add_java_binding_args(parser) + add_nodejs_binding_args(parser) + add_objc_binding_args(parser) + + # Execution Providers (now includes EP-specific profiling args) + add_execution_provider_args(parser) + + # Other Features + add_other_feature_args(parser) + + # Platform specific args (now includes Windows cross-compile targets & specific config/debug args) + if is_windows(): + add_windows_specific_args(parser) + elif is_macOS(): + add_apple_args(parser) + else: # Assuming Linux or other non-Windows, non-macOS Unix-like + add_linux_specific_args(parser) + + # --- Parse Arguments --- + args: argparse.Namespace = parser.parse_args() + + # --- Post-processing and Defaults --- + + # Normalize paths + if args.android_sdk_path: + args.android_sdk_path = os.path.normpath(args.android_sdk_path) + if args.android_ndk_path: + args.android_ndk_path = os.path.normpath(args.android_ndk_path) + + # Handle WASM exception logic + if args.enable_wasm_api_exception_catching: + args.disable_wasm_exception_catching = True # Catching at API level implies disabling broader catching + if not args.disable_wasm_exception_catching or args.enable_wasm_api_exception_catching: + # Doesn't make sense to catch if nothing throws + args.enable_wasm_exception_throwing_override = True + + # Set default CMake generator if not specified + # Check if cmake_generator attribute exists (it might if --use_xcode was used) + # before checking if it's None. + if not hasattr(args, "cmake_generator") or args.cmake_generator is None: + if is_windows(): + # Default to Ninja for WASM on Windows for potential speedup, VS otherwise + args.cmake_generator = "Ninja" if args.build_wasm else "Visual Studio 17 2022" + # else: Linux/macOS default (usually Makefiles or Ninja) is handled by CMake itself + + # Handle deprecated args + if hasattr(args, "enable_cuda_nhwc_ops") and args.enable_cuda_nhwc_ops: + warnings.warn("The argument '--enable_cuda_nhwc_ops' is deprecated and enabled by default.", DeprecationWarning) + + # Default behavior (update/build/test) if no action flags are specified + # Determine if it's a cross-compiled build (approximated by checking common cross-compile flags) + if not (args.update or args.build or args.test or args.clean or args.gen_doc): + args.update = True + args.build = True + # Only default to running tests for native builds if tests aren't explicitly skipped + if not is_cross_compiling(args) and not args.skip_tests: + args.test = True + elif is_cross_compiling(args): + print( + "Cross-compiling build detected: Defaulting to --update --build. Specify --test explicitly to run tests." + ) + + # Validation: Minimal build requires disabling exceptions + if args.disable_exceptions and args.minimal_build is None: + parser.error("--disable_exceptions requires --minimal_build to be specified.") + if is_windows(): + if getattr(args, "use_winml", False) and not getattr(args, "enable_wcos", False): + parser.error("--use_winml requires --enable_wcos to be specified.") + if hasattr(args, "msvc_toolset") and args.msvc_toolset: + try: + # Extract major.minor version parts (e.g., "14.36") + version_parts = args.msvc_toolset.split(".") + if len(version_parts) >= 2: + major = int(version_parts[0]) + minor = int(version_parts[1]) + # Check known problematic range based on previous script comments/help text + # Refined check: >= 14.36 and <= 14.39 + # Help text now says >= 14.40 is required, so check < 14.40 + if major == 14 and minor < 40: + # You could make this an error or just a warning + # parser.error(f"MSVC toolset version {args.msvc_toolset} is not supported. Use 14.40 or higher.") + warnings.warn( + f"Specified MSVC toolset version {args.msvc_toolset} might have compatibility issues. Version 14.40 or higher is recommended." + ) + + except (ValueError, IndexError): + warnings.warn( + f"Could not parse MSVC toolset version: {args.msvc_toolset}. Skipping compatibility check." + ) + + elif is_macOS(): + if getattr(args, "build_apple_framework", False) and not any( + [ + getattr(args, "ios", False), + getattr(args, "macos", None), + getattr(args, "visionos", False), + getattr(args, "tvos", False), + ] + ): + parser.error("--build_apple_framework requires --ios, --macos, --visionos, or --tvos to be specified.") + + if getattr(args, "macos", None) and not getattr(args, "build_apple_framework", False): + parser.error("--macos target requires --build_apple_framework.") + return args diff --git a/tools/ci_build/github/azure-pipelines/linux-migraphx-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/linux-migraphx-ci-pipeline.yml index 7f54eeeb5d90e..de5df97d37d3d 100644 --- a/tools/ci_build/github/azure-pipelines/linux-migraphx-ci-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/linux-migraphx-ci-pipeline.yml @@ -102,7 +102,6 @@ jobs: CMAKE_HIP_COMPILER=/opt/rocm/llvm/bin/clang++ \ onnxruntime_BUILD_KERNEL_EXPLORER=OFF \ onnxruntime_USE_COMPOSABLE_KERNEL=OFF \ - --mpi_home /opt/ompi \ --use_migraphx \ --rocm_version=$(RocmVersion) \ --rocm_home /opt/rocm \ diff --git a/tools/ci_build/github/azure-pipelines/linux-rocm-ci-pipeline.yml b/tools/ci_build/github/azure-pipelines/linux-rocm-ci-pipeline.yml index 40b7752fbbf4a..af5a8d1decb6e 100644 --- a/tools/ci_build/github/azure-pipelines/linux-rocm-ci-pipeline.yml +++ b/tools/ci_build/github/azure-pipelines/linux-rocm-ci-pipeline.yml @@ -101,7 +101,6 @@ jobs: CMAKE_HIP_COMPILER=/opt/rocm/llvm/bin/clang++ \ onnxruntime_BUILD_KERNEL_EXPLORER=ON \ CMAKE_HIP_ARCHITECTURES=gfx90a \ - --mpi_home /opt/ompi \ --use_rocm \ --rocm_version=$(RocmVersion) \ --rocm_home /opt/rocm \ diff --git a/tools/ci_build/github/azure-pipelines/templates/android-java-api-aar-test.yml b/tools/ci_build/github/azure-pipelines/templates/android-java-api-aar-test.yml index f5c0cfa68c0c7..1d35c9461768d 100644 --- a/tools/ci_build/github/azure-pipelines/templates/android-java-api-aar-test.yml +++ b/tools/ci_build/github/azure-pipelines/templates/android-java-api-aar-test.yml @@ -19,6 +19,11 @@ parameters: type: string default: '2.32.0.250228' +- name: enableWebGpu + displayName: Enable WebGPU test + type: boolean + default: true + jobs: - job: Final_AAR_Testing_Android pool: @@ -59,7 +64,21 @@ jobs: workingDirectory: $(Build.BinariesDirectory) # skip emulator tests for qnn package as there are no arm64-v8a emulators and no qnn libraries for x86 - - ${{ if not(contains(parameters.packageName, 'qnn')) }}: + # skip emulator tests for webgpu package as software rendering with swiftshader does not work without host GPU + - script: | + # Do not output ##vso[] commands with `set -x` or they may be parsed again and include a trailing quote. + set -e + if [[ "${{ parameters.packageName }}" == *"qnn"* || "${{ parameters.enableWebGpu }}" == "true" ]]; then + echo "##vso[task.setvariable variable=useEmulator]false" + echo "WebGPU is enabled or it's a QNN package. Emulator tests will be skipped." + else + echo "##vso[task.setvariable variable=useEmulator]true" + echo "Emulator tests will be executed." + fi + + displayName: Determine if emulator tests should be skipped + + - ${{ if eq(variables['useEmulator'], 'true')}}: - template: use-android-emulator.yml parameters: create: true @@ -77,6 +96,7 @@ jobs: stop: true - ${{ else }}: + - script: | # QNN SDK version string, expected format: 2.28.0.241029 # Extract the first three parts of the version string to get the Maven package version (e.g., 2.28.0) @@ -84,12 +104,17 @@ jobs: echo "QnnMavenPackageVersion: $QnnMavenPackageVersion" echo "##vso[task.setvariable variable=QnnMavenPackageVersion]$QnnMavenPackageVersion" displayName: Trim QNN SDK version to major.minor.patch + condition: contains('${{ parameters.packageName }}', 'qnn') - script: | - set -e -x - # build apks for qnn package as they are not built in the emulator test step - $(Build.SourcesDirectory)/java/gradlew --no-daemon clean assembleDebug assembleAndroidTest -DqnnVersion=$(QnnMavenPackageVersion) --stacktrace - displayName: Build QNN APK + set -e -x + # build apks for qnn package and webgpu as they are not built in the emulator test step + if [[ "${{ parameters.packageName }}" == *"qnn"* ]]; then + $(Build.SourcesDirectory)/java/gradlew --no-daemon clean assembleDebug assembleAndroidTest -DqnnVersion=$(QnnMavenPackageVersion) --stacktrace + else + $(Build.SourcesDirectory)/java/gradlew --no-daemon clean assembleDebug assembleAndroidTest --stacktrace + fi + displayName: Build APK for QNN and WebGPU package workingDirectory: $(Build.BinariesDirectory)/android_test/android # we run e2e tests on one older device (Pixel 3) and one newer device (Galaxy 23) diff --git a/tools/ci_build/github/pai/orttraining-ci.yml b/tools/ci_build/github/pai/orttraining-ci.yml deleted file mode 100644 index e504977353013..0000000000000 --- a/tools/ci_build/github/pai/orttraining-ci.yml +++ /dev/null @@ -1,47 +0,0 @@ -protocolVersion: 2 -name: "@@job_name@@" -type: job -jobRetryCount: 0 -prerequisites: - - type: dockerimage - uri: onnxruntimeregistry.azurecr.io/internal/azureml/onnxruntimepaibuild:rocm3.7-paiagent - name: docker_image_0 - auth: - username: "@@docker_user@@" - password: <% $secrets.docker_password_0 %> - registryuri: onnxruntimeregistry.azurecr.io -taskRoles: - taskrole: - instances: 1 - completion: - minFailedInstances: 1 - minSucceededInstances: -1 - taskRetryCount: 0 - dockerImage: docker_image_0 - resourcePerInstance: - gpu: 1 - cpu: 8 - memoryMB: 16384 - commands: - - >- - git clone https://github.com/microsoft/onnxruntime.git && - cd onnxruntime && - git checkout @@commit@@ && - python tools/ci_build/build.py - --config RelWithDebInfo - --build_dir ./build - --build_wheel - --enable_training - --use_rocm --rocm_home /opt/rocm - --mpi_home /usr/mpi/gcc/openmpi-4.0.4rc3 - --nccl_home /opt/rocm - --update - --build --parallel 8 - --skip_tests && - cd ./build/RelWithDebInfo && - LC_ALL=en_US.UTF-8 LANG=en_US.UTF-8 - ../../tools/ci_build/github/pai/pai_test_launcher.sh -secrets: - docker_password_0: "@@docker_password@@" -defaults: - virtualCluster: rocm diff --git a/tools/ci_build/requirements/pybind/requirements.txt b/tools/ci_build/requirements/pybind/requirements.txt index 6c31267d8bc84..c26f80ce443ad 100644 --- a/tools/ci_build/requirements/pybind/requirements.txt +++ b/tools/ci_build/requirements/pybind/requirements.txt @@ -1,6 +1,6 @@ setuptools flatbuffers -wheel +wheel==0.45.1 pytest numpy>=1.19.0 sympy>=1.10 diff --git a/tools/doc/builddoc.bat b/tools/doc/builddoc.bat deleted file mode 100644 index 75429e716ba40..0000000000000 --- a/tools/doc/builddoc.bat +++ /dev/null @@ -1,5 +0,0 @@ -@echo off -rem This script must be executed from this folder. -pip install -r ../../docs/python/requirements.txt -python -m sphinx -j2 -v -T -b html -d ../../build/docs/_doctrees/html ../../docs/python ../../build/docs/html -python -u rename_folders.py ../../build/docs/html diff --git a/tools/doc/builddoc.sh b/tools/doc/builddoc.sh deleted file mode 100755 index 88d4fc5d409b1..0000000000000 --- a/tools/doc/builddoc.sh +++ /dev/null @@ -1,21 +0,0 @@ -# This script must be executed from this folder. - -# $1 python path -# $2 source folder -# $3 build folder -# $4 build config - -# Fail the document generation if anything goes wrong in the process -set -e -x - -# Install doc generation tools -$1/python -m pip install -r $2/docs/python/requirements.txt - -# Fake onnxruntime installation -export PYTHONPATH=$3/$4:$PYTHONPATH - -# Remove old docs -rm -rf $3/docs/ - -$1/python -m sphinx -v -T -b html -d $3/docs/_doctrees/html $2/docs/python/ $3/docs/html -$1/python -u $2/tools/doc/rename_folders.py $3/docs/html diff --git a/tools/doc/rename_folders.py b/tools/doc/rename_folders.py deleted file mode 100644 index 587755d101ce2..0000000000000 --- a/tools/doc/rename_folders.py +++ /dev/null @@ -1,85 +0,0 @@ -""" -Github publishes the markdown documentation with jekyll enabled. -This extension does not publish any folder starting with `_`. -These folders need to be renamed. -""" - -import os -import re - - -def rename_folder(root): - """ - Renames all folder starting with `_`. - Returns the list of renamed folders. - """ - found = [] - for r, dirs, _files in os.walk(root): - for name in dirs: - if name.startswith("_"): - found.append((r, name)) - renamed = [] - for r, name in found: - into = name.lstrip("_") - renamed.append((r, name, into)) - full_src = os.path.join(r, name) - full_into = os.path.join(r, into) - if os.path.exists(full_into): - raise RuntimeError("%r already exists, previous documentation should be removed.") - print(f"rename {full_src!r}") - os.rename(full_src, full_into) - - return renamed - - -def replace_files(root, renamed): - subs = {r[1]: r[2] for r in renamed} - reg = re.compile('(\\"[a-zA-Z0-9\\.\\/\\?\\:@\\-_=#]+\\.([a-zA-Z]){2,6}([a-zA-Z0-9\\.\\&\\/\\?\\:@\\-_=#])*\\")') - - for r, _dirs, files in os.walk(root): - for name in files: - if os.path.splitext(name)[-1] != ".html": - continue - full = os.path.join(r, name) - with open(full, encoding="utf-8") as f: - content = f.read() - find = reg.findall(content) - repl = [] - for f in find: - if f[0].startswith("http"): - continue - for k, v in subs.items(): - if k == v: - raise ValueError(f"{k!r} == {v!r}") - if (f'"{k}') in f[0]: - repl.append((f[0], f[0].replace(f'"{k}', f'"{v}'))) - if (f"/{k}") in f[0]: - repl.append((f[0], f[0].replace(f"/{k}", f"/{v}"))) - if len(repl) == 0: - continue - print(f"update {full!r}") - for k, v in repl: - content = content.replace(k, v) - with open(full, "w", encoding="utf-8") as f: - f.write(content) - - -if __name__ == "__main__": - import sys - - if len(sys.argv) > 1: - root = sys.argv[-1] - else: - root = "../../build/docs/html" - print(f"look into {root!r}") - ren = rename_folder(root) - if len(ren) == 0: - ren = [ - ("", "_static", "static"), - ("", "_images", "images"), - ("", "_downloads", "downloads"), - ("", "_sources", "sources"), - ("", "_modules", "modules"), - ] - replace_files(root, ren) - print("done.")