diff --git a/.github/workflows/windows_cuda.yml b/.github/workflows/windows_cuda.yml index 7eaf7e06b2df5..12b0fe749a542 100644 --- a/.github/workflows/windows_cuda.yml +++ b/.github/workflows/windows_cuda.yml @@ -116,13 +116,34 @@ jobs: exit $lastExitCode } # Execute the build process - python.exe ${{ github.workspace }}\tools\ci_build\build.py --update --build --config RelWithDebInfo --build_dir build --skip_submodule_sync --build_csharp --parallel --nvcc_threads 4 --flash_nvcc_threads 4 --use_binskim_compliant_compile_flags --cmake_generator "Visual Studio 17 2022" --build_shared_lib --build_wheel --build_java --use_cuda --cuda_home="$env:RUNNER_TEMP\v12.8" --enable_cuda_profiling --use_vcpkg --use_vcpkg_ms_internal_asset_cache --enable_transformers_tool_test --cmake_extra_defines onnxruntime_QUICK_BUILD=ON --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 --cmake_extra_defines onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON + python.exe ${{ github.workspace }}\tools\ci_build\build.py --update --build --config RelWithDebInfo --build_dir build --skip_submodule_sync --build_csharp --parallel --nvcc_threads 4 --flash_nvcc_threads 4 --use_binskim_compliant_compile_flags --cmake_generator "Visual Studio 17 2022" --build_shared_lib --build_wheel --build_java --use_cuda --cuda_home="$env:RUNNER_TEMP\v12.8" --enable_cuda_profiling --use_vcpkg --use_vcpkg_ms_internal_asset_cache --enable_transformers_tool_test --cmake_extra_defines onnxruntime_QUICK_BUILD=ON --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 --cmake_extra_defines onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON onnxruntime_USE_CUDA_DIRECTSTORAGE=ON if ($lastExitCode -ne 0) { exit $lastExitCode } - # Clean up the output directory before uploading artifacts $outputDir = "${{ runner.temp }}\build\RelWithDebInfo" + # DirectStorage resolves dstoragecore.dll relative to the process executable. + $testDir = Join-Path $outputDir "RelWithDebInfo" + $sdkManifest = Join-Path $outputDir "directstorage-source-dir.txt" + if (!(Test-Path -LiteralPath $sdkManifest -PathType Leaf)) { + throw "Missing DirectStorage SDK source manifest: $sdkManifest" + } + if (!(Test-Path -LiteralPath (Join-Path $testDir "onnxruntime_provider_test.exe") -PathType Leaf)) { + throw "Missing CUDA provider test executable in $testDir" + } + $sdkDir = (Get-Content -LiteralPath $sdkManifest -Raw -ErrorAction Stop).Trim() + if ([string]::IsNullOrWhiteSpace($sdkDir)) { + throw "Empty DirectStorage SDK source manifest: $sdkManifest" + } + foreach ($dll in @("dstorage.dll", "dstoragecore.dll")) { + $source = Join-Path $sdkDir "native\bin\x64\$dll" + if (!(Test-Path -LiteralPath $source -PathType Leaf)) { + throw "Missing DirectStorage runtime: $source" + } + Copy-Item -LiteralPath $source -Destination $testDir -ErrorAction Stop + } + + # Clean up the output directory before uploading artifacts Write-Host "Cleaning up files from $outputDir..." Remove-Item -Path "$outputDir\onnxruntime" -Recurse -Force -ErrorAction SilentlyContinue @@ -243,7 +264,7 @@ jobs: exit $lastExitCode } - python.exe ${{ github.workspace }}\tools\ci_build\build.py --test --config RelWithDebInfo --build_dir build --skip_submodule_sync --build_csharp --parallel --nvcc_threads 4 --flash_nvcc_threads 4 --use_binskim_compliant_compile_flags --cmake_generator "Visual Studio 17 2022" --build_shared_lib --build_wheel --build_java --use_cuda --cuda_home="$env:RUNNER_TEMP\v12.8" --enable_cuda_profiling --use_vcpkg --use_vcpkg_ms_internal_asset_cache --enable_transformers_tool_test --cmake_extra_defines onnxruntime_QUICK_BUILD=ON --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 --cmake_extra_defines onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON + python.exe ${{ github.workspace }}\tools\ci_build\build.py --test --config RelWithDebInfo --build_dir build --skip_submodule_sync --build_csharp --parallel --nvcc_threads 4 --flash_nvcc_threads 4 --use_binskim_compliant_compile_flags --cmake_generator "Visual Studio 17 2022" --build_shared_lib --build_wheel --build_java --use_cuda --cuda_home="$env:RUNNER_TEMP\v12.8" --enable_cuda_profiling --use_vcpkg --use_vcpkg_ms_internal_asset_cache --enable_transformers_tool_test --cmake_extra_defines onnxruntime_QUICK_BUILD=ON --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=86 --cmake_extra_defines onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON onnxruntime_USE_CUDA_DIRECTSTORAGE=ON if ($lastExitCode -ne 0) { exit $lastExitCode } diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 21f1ee7c3c4eb..57fe08645dab1 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -70,13 +70,12 @@ option(onnxruntime_ENABLE_PYTHON "Enable python bindings" OFF) option(onnxruntime_ENABLE_MEMLEAK_CHECKER "Experimental: Enable memory leak checker in Windows debug build" OFF) option(onnxruntime_ENABLE_CONVSYMKERNELAVX2_SAT_CHECKER "Experimental: Enable ConvSymKernelAvx2 assembly saturation checker in build" OFF) option(onnxruntime_USE_CUDA "Build with CUDA support" OFF) -# Enable ONNX Runtime CUDA EP's internal unit tests that directly access the EP's internal functions instead of through -# OpKernels. When the option is ON, we will have two copies of GTest library in the same process. It is not a typical -# use. If you hit any problem with that, please do not report it to GTest. Turn OFF the following build option instead. -cmake_dependent_option(onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS "Build with CUDA unit tests" OFF "onnxruntime_USE_CUDA;onnxruntime_BUILD_UNIT_TESTS" OFF) - +cmake_dependent_option(onnxruntime_USE_CUDA_DIRECTSTORAGE "Build Microsoft DirectStorage CUDA loading support" OFF "onnxruntime_USE_CUDA;WIN32" OFF) cmake_dependent_option(onnxruntime_USE_CUDA_NHWC_OPS "Build CUDA with NHWC op support" ON "onnxruntime_USE_CUDA" OFF) cmake_dependent_option(onnxruntime_BUILD_CUDA_EP_AS_PLUGIN "Build CUDA EP as a separate plugin shared library instead of the legacy in-tree provider" OFF "onnxruntime_USE_CUDA" OFF) +if(onnxruntime_USE_CUDA_DIRECTSTORAGE AND onnxruntime_BUILD_CUDA_EP_AS_PLUGIN) + message(FATAL_ERROR "onnxruntime_USE_CUDA_DIRECTSTORAGE is not supported with onnxruntime_BUILD_CUDA_EP_AS_PLUGIN.") +endif() option(onnxruntime_BUILD_CUDA_QUANT_PREPROCESS "Build CUDA weight-packing module onnxruntime_cuda_quant_preprocess.so" OFF) option(onnxruntime_CUDA_MINIMAL "Build CUDA without any operations apart from memcpy ops. Useful for a very minimal TRT build" OFF) option(onnxruntime_ENABLE_CUDA_LINE_NUMBER_INFO "When building with CUDA support, generate device code line number information." OFF) @@ -97,6 +96,12 @@ option(onnxruntime_USE_ARM_NEON_NCHWC "Build with ARM Neon NCHWc kernels in MLAS option(onnxruntime_USE_KLEIDIAI "Build with KleidiAI integration in MLAS" OFF) option(onnxruntime_USE_QMX_KLEIDIAI_COEXIST "Build with QMX and Arm KLEIDIAI libraries" OFF) option(onnxruntime_BUILD_UNIT_TESTS "Build ONNXRuntime unit tests" ON) +# Declare both prerequisites before evaluating this dependent option on a fresh configure. +# Enable ONNX Runtime CUDA EP's internal unit tests that directly access the EP's internal functions instead of through +# OpKernels. When the option is ON, we will have two copies of GTest library in the same process. It is not a typical +# use. If you hit any problem with that, please do not report it to GTest. Turn OFF the following build option instead. +cmake_dependent_option(onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS "Build with CUDA unit tests" OFF "onnxruntime_USE_CUDA;onnxruntime_BUILD_UNIT_TESTS" OFF) + # Materialize the ONNX node-test corpus from ONNX's Python generators into the build tree at # configure/build time, instead of depending on the on-disk corpus shipped in the ONNX source # archive. This detaches ORT from ONNX PR #7959 (which deletes onnx/backend/test/data/node). @@ -1545,6 +1550,24 @@ if (onnxruntime_USE_CUDA) endif() find_package(CUDAToolkit REQUIRED) + if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND NOT onnxruntime_MINIMAL_BUILD AND NOT onnxruntime_CUDA_MINIMAL) + include(CheckCXXSourceCompiles) + include(CMakePushCheckState) + cmake_push_check_state(RESET) + set(CMAKE_REQUIRED_INCLUDES ${CUDAToolkit_INCLUDE_DIRS}) + check_cxx_source_compiles(" + #include + using SetBoolParameterFn = decltype(&cuFileSetParameterBool); + int main() { + return CUFILE_PARAM_USE_PCIP2PDMA == CUFILE_PARAM_PROPERTIES_ALLOW_COMPAT_MODE; + }" onnxruntime_CUFILE_CONFIG_API_SUPPORTED) + cmake_pop_check_state() + if(onnxruntime_CUFILE_CONFIG_API_SUPPORTED) + set_property(SOURCE "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_external_data_loader_gds.cc" + APPEND PROPERTY COMPILE_DEFINITIONS ORT_CUDA_GDS_AVAILABLE) + endif() + endif() + if(MSVC AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9 AND CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 13.0) foreach(_cuda_include_dir IN LISTS CUDAToolkit_INCLUDE_DIRS) diff --git a/cmake/deps.txt b/cmake/deps.txt index dd582d7651e16..fc1c2a55d8cdf 100644 --- a/cmake/deps.txt +++ b/cmake/deps.txt @@ -58,6 +58,7 @@ cutlass;https://github.com/NVIDIA/cutlass/archive/refs/tags/v4.7.0.zip;51d4f1ba4 deep_gemm;https://github.com/deepseek-ai/DeepGEMM/archive/559d79fb6994a58b8a15b4b93bf13ccc16edf247.tar.gz;76a0076386991cac8e5d32c3e7e74d9bb8102115 extensions;https://github.com/microsoft/onnxruntime-extensions/archive/c24b7bab0c12f53da76d0c31b03b9f0f8ec8f3b4.zip;239063aee4946a9af147b473a4c3da78ba7413b4 directx_headers;https://github.com/microsoft/DirectX-Headers/archive/refs/tags/v1.613.1.zip;47653509a3371eabb156360f42faf582f314bf2e +directstorage;https://www.nuget.org/api/v2/package/Microsoft.Direct3D.DirectStorage/1.2.3;be08f099a75c54997a753f444224af96ed0ee3d9 cudnn_frontend;https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v1.27.0.zip;1e4c9a464d3437e388ab0163f3be068dba783c08 dawn;https://github.com/google/dawn/archive/refs/tags/v20260916.214021.zip;514e914f23a0787213c89eed879f74b5e275690d dawn_agility_sdk;https://www.nuget.org/api/v2/package/Microsoft.Direct3D.D3D12/1.721.3-preview;fa5f5fc8d0c8c209cfbb530be634960d20595841 diff --git a/cmake/external/directstorage.cmake b/cmake/external/directstorage.cmake new file mode 100644 index 0000000000000..d5fcbf58b2f91 --- /dev/null +++ b/cmake/external/directstorage.cmake @@ -0,0 +1,16 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +include_guard(GLOBAL) +onnxruntime_fetchcontent_declare( + directstorage + URL ${DEP_URL_directstorage} + URL_HASH SHA1=${DEP_SHA1_directstorage} + DOWNLOAD_NAME directstorage.zip +) +onnxruntime_fetchcontent_makeavailable(directstorage) + +if(onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS) + # CI deploys the SDK runtime only beside test executables, including when the SDK source is overridden. + file(WRITE "${CMAKE_BINARY_DIR}/directstorage-source-dir.txt" "${directstorage_SOURCE_DIR}\n") +endif() diff --git a/cmake/onnxruntime_providers_cuda.cmake b/cmake/onnxruntime_providers_cuda.cmake index f0e37e8aa4c50..0af481707c674 100644 --- a/cmake/onnxruntime_providers_cuda.cmake +++ b/cmake/onnxruntime_providers_cuda.cmake @@ -24,6 +24,12 @@ endif() # Exclude plugin directory if it was picked up by GLOB_RECURSE list(FILTER onnxruntime_providers_cuda_cc_srcs EXCLUDE REGEX "core/providers/cuda/plugin/.*") + if(onnxruntime_MINIMAL_BUILD OR onnxruntime_CUDA_MINIMAL) + list(REMOVE_ITEM onnxruntime_providers_cuda_cc_srcs + "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_external_data_loader_directstorage.cc" + "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_external_data_loader_gds.cc" + ) + endif() # Remove pch files list(REMOVE_ITEM onnxruntime_providers_cuda_cc_srcs @@ -236,7 +242,18 @@ # config_cuda_provider_shared_module can be used to config onnxruntime_providers_cuda_obj, onnxruntime_providers_cuda & onnxruntime_providers_cuda_ut. # This function guarantees that all 3 targets have the same configurations. + if(onnxruntime_USE_CUDA_DIRECTSTORAGE AND + NOT onnxruntime_MINIMAL_BUILD AND NOT onnxruntime_CUDA_MINIMAL) + include(external/directstorage.cmake) + endif() + function(config_cuda_provider_shared_module target) + if(onnxruntime_USE_CUDA_DIRECTSTORAGE AND + NOT onnxruntime_MINIMAL_BUILD AND NOT onnxruntime_CUDA_MINIMAL) + target_compile_definitions(${target} PRIVATE ORT_CUDA_DIRECTSTORAGE_AVAILABLE) + target_include_directories(${target} PRIVATE "${directstorage_SOURCE_DIR}/native/include") + target_link_libraries(${target} PRIVATE d3d12 dxgi) + endif() if (onnxruntime_REDUCED_OPS_BUILD) add_op_reduction_include_dirs(${target}) endif() diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 1d01a4482e678..46543348f8895 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1075,6 +1075,7 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND onnxruntime_BUILD_CUDA_EP_AS_P NOT onnxruntime_MINIMAL_BUILD AND NOT onnxruntime_REDUCED_OPS_BUILD) set(onnxruntime_test_providers_cuda_plugin_internal_test_src "${TEST_SRC_DIR}/providers/cuda/test_cases/allocator_cuda_test.cc" + "${TEST_SRC_DIR}/providers/cuda/test_cases/cuda_external_data_loader_gds_test.cc" "${TEST_SRC_DIR}/providers/cuda/test_cases/cuda_utils_test.cc" "${TEST_SRC_DIR}/providers/cuda/test_cases/group_query_attention_workspace_header_test.cc" "${TEST_SRC_DIR}/providers/cuda/test_cases/packed_attention_workspace_header_test.cc" diff --git a/docs/CUDA_cuDNN_Optional_Design.md b/docs/CUDA_cuDNN_Optional_Design.md index 1d761d31b72b0..e3d80549b29c9 100644 --- a/docs/CUDA_cuDNN_Optional_Design.md +++ b/docs/CUDA_cuDNN_Optional_Design.md @@ -292,9 +292,10 @@ Implementation details: `CUDAExecutionProviderInfo::FromProviderOptions(...)`. - Emit it from `CUDAExecutionProviderInfo::ToProviderOptions(...)`. - Include it in `std::hash` because it changes the EP behavior. -- Do **not** add a field to `OrtCUDAProviderOptionsV2` for Phase 1. That struct is public C - ABI surface; string-key provider options are sufficient and can be set through existing - provider-options APIs. +- Phase 1 keeps this policy in `CUDAExecutionProviderInfo`. `OrtCUDAProviderOptionsV2` is opaque in the + public C API: callers obtain it through `CreateCUDAProviderOptions` and configure it through string keys. + Its definition in `include/onnxruntime/core/providers/cuda/cuda_provider_options.h` is internal and may be extended for new + options. This differs from the publicly defined `OrtCUDAProviderOptions`, whose layout must remain stable. - Add an EP helper such as `CUDAExecutionProvider::IsCudnnEnabled()` or `CudaKernel::IsCudnnEnabled()` so kernels can distinguish: - cuDNN disabled by user (`enable_cudnn=0`), and diff --git a/docs/FAQ.md b/docs/FAQ.md index 1664fee5953e4..a8387edebbb86 100644 --- a/docs/FAQ.md +++ b/docs/FAQ.md @@ -7,7 +7,8 @@ The default CUDA build supports 3 standard quantization operators: QuantizeLinea ## How can I reduce model loading time for large CPU or CUDA models? See [Accelerate model loading](Model_Loading_Performance.md) for parallel CPU weight prepacking and CUDA external-data -loading through pinned host buffers, including the session and execution provider options that control them. +loading through GPUDirect Storage, Microsoft DirectStorage, or pinned host buffers, including the session and execution provider options that +control them. ## How do I change the severity level of the default logger to something other than the default (WARNING)? Setting the severity level to VERBOSE is most useful when debugging errors. diff --git a/docs/Model_Loading_Performance.md b/docs/Model_Loading_Performance.md index 014f728dda2e5..7d0a3ac6d7762 100644 --- a/docs/Model_Loading_Performance.md +++ b/docs/Model_Loading_Performance.md @@ -5,10 +5,10 @@ ONNX Runtime provides two independent loading optimizations for models with larg | Target | Mechanism | Configuration | Default | |---|---|---|---| | CPU prepacking | Run eligible CPU kernel `PrePack()` calls concurrently | `session.prepack.enable_parallel` | Disabled (`"0"`) | -| CUDA external initializers | Read external data through reusable pinned buffers while copying to the GPU | CUDA EP option `external_data_loader_reading_threads` | 4 readers | +| CUDA external initializers | Read external data with GPUDirect Storage, Microsoft DirectStorage, or reusable pinned buffers | CUDA EP options `external_data_loader_use_gds`, `external_data_loader_use_directstorage`, and `external_data_loader_reading_threads` | Direct storage disabled; 4 readers | -The CPU option is a session configuration entry. The CUDA option is an execution provider option passed when the -CUDA EP is appended to `SessionOptions`. +The CPU option is a session configuration entry. The CUDA options are execution provider options passed when the CUDA +EP is appended to `SessionOptions`. ## Parallel CPU weight prepacking @@ -46,11 +46,98 @@ session = ort.InferenceSession( The best thread count depends on available CPU cores, memory bandwidth, storage, and concurrent workloads. Setting `intra_op_num_threads` to `1` keeps prepacking sequential even if the configuration entry is enabled. -## Pinned-buffer loading for CUDA external initializers +## CUDA external initializer loading Models saved with [external data](https://onnx.ai/onnx/repo-docs/ExternalData.html) normally load weights through -pageable CPU memory before copying them to the GPU. The CUDA execution provider can instead load external -initializers through two reusable 64 MiB pinned host buffers: +pageable CPU memory before copying them to the GPU. The CUDA execution provider can instead use NVIDIA GPUDirect +Storage (GDS) on Linux, Microsoft DirectStorage on Windows, or reusable pinned host buffers. + +### GPUDirect Storage + +GDS loads external initializers without staging file data in CPU memory. It is opt-in. If it cannot be initialized or +cannot read an external-data file, ONNX Runtime logs a warning and uses the configured pinned/pageable host-memory +fallback for the rest of the session. + +With GDS enabled, ONNX Runtime opens each external-data file with `O_DIRECT` and uses `libcufile` to read 64 MiB +blocks into a reusable, cuFile-registered CUDA buffer. Each block is then copied device-to-device into the +initializer allocation owned by the CUDA arena: + +```text +external-data file -> registered CUDA staging buffer -> CUDA arena initializer + cuFileRead device-to-device copy +``` + +The reusable staging buffer bounds additional GPU memory usage to 64 MiB per CUDA external-data loader. Each +device-to-device copy completes before that buffer is reused. String and Boolean initializers retain the existing +loading path because they require host-side conversion. Loaders share a process-wide cuFile driver. Its final release +and subsequent initialization are serialized, so a new loader cannot configure or reopen the driver until the +previous driver has finished closing. + +GDS requires: + +- Linux and `cufile.h` with `cuFileSetParameterBool`, `CUFILE_PARAM_USE_PCIP2PDMA`, and + `CUFILE_PARAM_PROPERTIES_ALLOW_COMPAT_MODE`; +- `libcufile.so` at runtime; +- either `nvidia-fs` or a recent open NVIDIA kernel module that supports PCI P2PDMA; +- a supported storage/filesystem and PCIe topology; and +- external weights stored in a file that can be opened with `O_DIRECT`. + +ONNX Runtime loads `libcufile` dynamically, so enabling the option does not add a mandatory runtime dependency for +users who keep GDS disabled. It requests PCI P2PDMA, which can provide GDS without `nvidia-fs` on supported recent +kernels, GPUs, and storage devices. It also disables cuFile compatibility mode: if the storage stack cannot provide +a native GDS path, ONNX Runtime uses its configured host-memory fallback instead of cuFile's internal POSIX fallback. +GDS is attempted only for external-data ranges whose offset and length are both 4 KiB aligned. An unaligned +initializer uses the configured fallback without disabling GDS for later aligned initializers. + +The build checks for the required cuFile configuration API. Older CUDA toolkits without it remain supported, but +enabling GDS in those builds logs a warning and uses the configured host-memory fallback. + +### Microsoft DirectStorage (Windows) + +The built-in CUDA execution provider also supports Microsoft's DirectStorage API through D3D12/CUDA +interoperability. Build with `--cmake_extra_defines onnxruntime_USE_CUDA_DIRECTSTORAGE=ON` in addition to the +usual CUDA build options. This option is not supported with `onnxruntime_BUILD_CUDA_EP_AS_PLUGIN=ON`; +CMake rejects that combination rather than silently compiling an unavailable backend. +This opt-in build downloads the pinned DirectStorage SDK headers; it does not +introduce a link-time dependency on `dstorage.dll`. Deploy the SDK's matching x64 `dstorage.dll` and +`dstoragecore.dll` beside the application executable, following Microsoft's +[DirectStorage deployment guidance](https://github.com/microsoft/DirectStorage/blob/main/Docs/DeveloperGuidance.md#sdk-path). +For Python, the application executable is `python.exe`, not the CUDA provider DLL. +The application must make `dstorage.dll` discoverable through the Windows application/system/user DLL search +directories; the current working directory is not searched. + +Set the CUDA EP option `external_data_loader_use_directstorage` to `"1"` to enable this path. +It requires a Windows D3D12-capable NVIDIA adapter that supports CUDA external memory and fence import. +The D3D12 adapter is selected by the CUDA device's LUID, not by assuming that both APIs enumerate GPUs +in the same order. Linked D3D12 adapters are not supported. + +```text +external-data file -> DirectStorage -> shared D3D12 GPU buffer -> CUDA arena initializer + internal staging CUDA external memory device-to-device copy +``` + +The loader reuses a 32 MiB shared GPU buffer and a DirectStorage queue. A shared D3D12 fence establishes +completion and visibility to CUDA; request errors are checked before copying bytes into the initializer. +CUDA copies complete before the next DirectStorage write reuses the buffer. The additional GPU buffer +does not include DirectStorage's own internal staging allocations. This implementation loads uncompressed +ONNX external data, not GDeflate-compressed weights. + +**DirectStorage is not NVIDIA GPUDirect Storage:** Microsoft's uncompressed data flow can use system-memory +and upload-heap staging. Successful DirectStorage loading does not establish zero-copy storage-to-VRAM DMA +or prove that Windows BypassIO was used. Compare actual timings rather than assuming the API is faster. +See Microsoft's [uncompressed data flow documentation](https://github.com/microsoft/DirectStorage/blob/main/Docs/DeveloperGuidance.md#uncompressed-data-flow). + +DirectStorage accepts unaligned offsets and lengths. Because its API opens files by path, the loader compares +the opened file's identity with the handle already validated by ONNX Runtime before submitting reads. +Initialization or read failures produce a warning and disable DirectStorage for the remainder of that loader's +lifetime, using the configured host-memory fallback. Boolean tensors retain the host conversion path. +If both direct-storage options are enabled, DirectStorage is tried first, followed by GDS, then the host path; +normally enable only the option appropriate for the operating system. Unsupported builds/platforms report +unavailability and use the fallback instead of silently claiming DirectStorage support. + +### Pinned-buffer host loading + +The CUDA execution provider can load external initializers through two reusable 64 MiB pinned host buffers: ```text external-data file -> pinned buffer 0/1 -> CUDA initializer allocation @@ -61,29 +148,39 @@ The buffers alternate so that reading the next block can overlap the host-to-dev Each buffer is synchronized before reuse. Initializers are loaded one at a time through the shared staging resources, which bounds pinned host memory use at 128 MiB per loader. -The `external_data_loader_reading_threads` CUDA provider option controls how each block is filled: +The following CUDA execution provider options control the primary and fallback paths: -| Value | Behavior | -|---:|---| -| `0` | Disable the CUDA external-data loader and use the framework's pageable-memory path | -| `1` | Use pinned buffers with synchronous reads on the calling thread | -| `2` to `64` | Use that many independent CPU read tasks per block | +| Option | Values | Default | Purpose | +|---|---|---:|---| +| `external_data_loader_use_gds` | `0` or `1` | `0` | Try GDS before another external-data loading path | +| `external_data_loader_use_directstorage` | `0` or `1` | `0` | Try Microsoft DirectStorage through D3D12/CUDA on Windows | +| `external_data_loader_reading_threads` | `0` to `64` | `4` | Configure the pinned-buffer fallback; `0` disables it | -The default is `4`, so the pinned-buffer loader is enabled without additional configuration. Parallel reads are used -for external tensors of at least 16 MiB; smaller tensors use one read. The optimal reader count depends on the storage -device and filesystem. If pinned buffers or CUDA streams cannot be created, loading falls back to a pageable buffer. -Models with weights embedded in the ONNX file do not use this external-data path. +Keep `external_data_loader_reading_threads` greater than zero when enabling either direct-storage option to retain +the pinned-buffer backup. +`1` uses synchronous reads into pinned memory. Values from `2` through `64` use that many parallel read tasks per +64 MiB pinned buffer. The default is `4`, so the pinned-buffer loader is enabled without additional configuration. +Parallel reads are used for external tensors of at least 16 MiB; smaller tensors use one read. The optimal reader count +depends on the storage device and filesystem. If pinned buffers or CUDA streams cannot be created, loading falls back +to a pageable buffer. If the value is `0`, the pageable path is used when direct storage is disabled or unavailable. Models with +weights embedded in the ONNX file do not use this external-data path. Configure the CUDA EP in Python: ```python +import sys + import onnxruntime as ort session_options = ort.SessionOptions() providers = [ ( "CUDAExecutionProvider", - {"external_data_loader_reading_threads": "4"}, + { + "external_data_loader_use_gds": "0" if sys.platform == "win32" else "1", + "external_data_loader_use_directstorage": "1" if sys.platform == "win32" else "0", + "external_data_loader_reading_threads": "4", + }, ), "CPUExecutionProvider", ] @@ -103,7 +200,14 @@ session_options.SetIntraOpNumThreads(8); session_options.AddConfigEntry("session.prepack.enable_parallel", "1"); Ort::CUDAProviderOptions cuda_options; -cuda_options.Update({{"external_data_loader_reading_threads", "4"}}); +cuda_options.Update({ +#ifdef _WIN32 + {"external_data_loader_use_directstorage", "1"}, +#else + {"external_data_loader_use_gds", "1"}, +#endif + {"external_data_loader_reading_threads", "4"}, +}); session_options.AppendExecutionProvider_CUDA_V2(*cuda_options); Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "model_loading"); @@ -113,3 +217,80 @@ Ort::Session session(env, ORT_TSTR("model.onnx"), session_options); The CPU prepacking option only parallelizes CPU EP kernels. The CUDA provider option only changes how external initializers assigned to CUDA memory are staged and copied; CPU and other execution providers retain their existing loading paths. They can be enabled together for models partitioned between CPU and CUDA. + +### Tests + +`CudaExternalDataLoaderTest.*Gds*` exercises the real loader with GDS enabled, including aligned and unaligned ranges, +multiple buffers, repeated loads, and different host-memory fallback configurations. These tests require a CUDA GPU, +but use the configured fallback when native GDS is unavailable. A passing result alone does not prove native GDS usage +or performance. `CudaGdsDriverTest.*` separately checks shared-driver lifetime synchronization without GPU hardware. + +`CApiTest.CUDAProviderOptions*Gds*` checks string-based configuration, invalid values, and serialization round trips. + +`CudaExternalDataLoaderTest.*DirectStorage*` covers option validation, aligned and unaligned reads, repeated +loads, multiple buffers, and configured fallbacks. `NativeDirectStorageWithoutFallback` calls the DirectStorage +backend directly, verifies loaded bytes, rejects out-of-range and mismatched-file requests, and cannot succeed +by using the host fallback. It explicitly skips when D3D12/CUDA/DirectStorage initialization is unavailable. +`CApiTest.CUDAProviderOptionsDirectStorageRoundTrip` covers the string-based option. + +Successful loads emit INFO records of the form `CUDA external data loader: path= bytes=`, +where `` is `pageable`, `pinned`, `gds`, or `directstorage`. Enable the **default** logger's INFO severity +(`ort.set_default_logger_severity(1)` in Python) as well as the session logger when collecting these records. +They identify the path actually used, including fallbacks; requested provider options alone are not proof. + +### Comparing all three loading paths + +`onnxruntime/test/python/transformers/benchmark_cuda_model_loading.py` compares pageable CPU staging +(the loader-disabled baseline), pinned buffers, and the platform's direct-storage API, all targeting the +same CUDA device. Use the Python package from the build being evaluated, not an installed older wheel. +On Windows, enable the DirectStorage build option and deploy the SDK runtime as described above. + +For example, generate identical aligned external weights totaling 1 GiB and collect five measurements per path: + +```powershell +python onnxruntime\test\python\transformers\benchmark_cuda_model_loading.py ` + --generate-model .\cuda-loading-1gib --weight-count 16 --weight-dim 4096 ` + --threads 1 --repetitions 5 --output .\cuda-loading-1gib.json +``` + +Use `--weight-count 64` and a different output directory for 4 GiB of weights, provided enough GPU memory +is available. Generation refuses to overwrite an existing fixture. To benchmark a real model instead, supply +`--model model.onnx --inputs inputs.npz --expected-outputs expected.npz`; the NPZ keys must match tensor names, +and reference outputs must be computed independently. + +Each sample uses a fresh process. The timed interval covers session construction, including graph initialization +and completed weight transfers, but not imports, fixture generation, or output verification. A blocking inference +and comparison against independent reference values happen afterward. Reported GiB/s is therefore **effective +end-to-end initialization throughput**, not raw disk bandwidth. The JSON report includes timing distributions, +GPU/configuration metadata, actual loaded bytes by path, and separately classified fallback or mixed-path samples. +A direct-storage fallback is never counted as a successful direct-storage measurement. + +Caches remain OS-managed: fixture creation and warmups may warm the filesystem cache, and a fresh process does +not imply a cold disk. The optional POSIX per-file eviction hint is also not proof of cold-cache operation. +The script never performs privileged or global cache flushing. Record storage/filesystem details alongside results, +and do not compare Windows DirectStorage and Linux GDS numbers as if they came from an identical software stack. + +#### Windows measurements (2026-09-23) + +Measured on Windows 11 (build 26200), an NVIDIA RTX 4060 Laptop GPU (8 GiB, WDDM, driver 591.55), +and a local WD `SDCPNRZ-2T00-1124-WD` 2 TB SSD. The source build used MSVC 2022, CUDA 13.0.2, +DirectStorage 1.2.3, and CPython 3.13.14. It was a Release/quick build restricted to SM89 and +MatMul/Gather registrations, with contrib support enabled. The fixture used FP32 4096-by-4096 weights, +one intra-op thread, four pinned-buffer readers, disabled graph optimization/prepacking, and disabled TF32. + +Each row summarizes five fresh-process samples after one warmup per path, with OS-managed (potentially warm) +caches. The complete weight-byte counts were confirmed from actual-path logs and every sample's outputs +matched independent references. DirectStorage rows used the DirectStorage API, not the pinned fallback. + +| External weights | Loading path | Median initialization | Effective throughput | +|---|---|---:|---:| +| 1 GiB | CPU pageable | 1.517 s | 0.659 GiB/s | +| 1 GiB | Pinned buffers | 1.105 s | 0.905 GiB/s | +| 1 GiB | Microsoft DirectStorage | 1.853 s | 0.540 GiB/s | +| 4 GiB | CPU pageable | 6.201 s | 0.645 GiB/s | +| 4 GiB | Pinned buffers | 2.292 s | 1.745 GiB/s | +| 4 GiB | Microsoft DirectStorage | 7.870 s | 0.508 GiB/s | + +On this machine and workload, pinned buffers outperform both pageable loading and the current uncompressed +DirectStorage implementation. These are complete session-initialization measurements, not isolated I/O timings, +and are not evidence of cold-cache disk bandwidth or Linux GDS performance. diff --git a/include/onnxruntime/core/providers/cuda/cuda_provider_options.h b/include/onnxruntime/core/providers/cuda/cuda_provider_options.h index 0afbb13739205..36433d0405234 100644 --- a/include/onnxruntime/core/providers/cuda/cuda_provider_options.h +++ b/include/onnxruntime/core/providers/cuda/cuda_provider_options.h @@ -42,5 +42,7 @@ struct OrtCUDAProviderOptionsV2 { int use_tf32 = 1; // use TF32 int fuse_conv_bias = 0; // Enable CUDNN Frontend kernel fusing, results in JIT compiles int sdpa_kernel = 0; // Scaled Dot Product Attention kernel option - size_t external_data_loader_reading_threads = 4; // Number of CPU read tasks per external-data staging buffer. 0 disables the loader; 1 disables parallel reads. + size_t external_data_loader_reading_threads = 4; // Number of CPU read tasks per external-data staging buffer. 0 disables pinned-buffer fallback; 1 disables parallel reads. + int external_data_loader_use_gds = 0; // Try GPUDirect Storage before the configured pinned/pageable host-memory fallback. + int external_data_loader_use_directstorage = 0; // Try Microsoft DirectStorage on Windows before the configured host-memory fallback. }; diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index db988873539ad..bc99373e8bf92 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -142,6 +142,7 @@ static common::Status DeserializeTensorProto(const Env& env, const std::basic_st // Bool external initializers are copied verbatim and may carry bytes outside the canonical // {0, 1} set. The CPU staging tensor above can be backed by a read-only mmap, so normalize into // a writable CPU copy before copying to the device (see utils::NormalizeBoolTensorIfNeeded). +#if defined(ORT_MINIMAL_BUILD) if (cpu_staging_tensor.IsDataType()) { Tensor normalized_cpu_tensor; ORT_RETURN_IF_ERROR(AllocateTensorOnDeviceOrMemory(/* use_device_allocator_for_initializers =*/true, @@ -154,6 +155,25 @@ static common::Status DeserializeTensorProto(const Env& env, const std::basic_st return CopyTensorFromCPUToDevice(data_transfer_mgr, deserialized_value.Get(), std::move(tensor), ort_value); +#else + if (cpu_staging_tensor.IsDataType()) { + Tensor normalized_cpu_tensor; + ORT_RETURN_IF_ERROR(AllocateTensorOnDeviceOrMemory(/* use_device_allocator_for_initializers =*/true, + tensor_shape, type, + default_cpu_alloc, normalized_cpu_tensor)); + utils::MakeCpuTensorCopy(cpu_staging_tensor, normalized_cpu_tensor); + utils::NormalizeBoolTensorIfNeeded(normalized_cpu_tensor); + ORT_RETURN_IF_ERROR( + CopyTensorFromCPUToDevice(data_transfer_mgr, normalized_cpu_tensor, std::move(tensor), ort_value)); + } else { + ORT_RETURN_IF_ERROR(CopyTensorFromCPUToDevice(data_transfer_mgr, deserialized_value.Get(), + std::move(tensor), ort_value)); + } + if (device.Type() == OrtDevice::GPU && device.Vendor() == OrtDevice::VendorIds::NVIDIA) { + LOGS_DEFAULT(INFO) << "CUDA external data loader: path=pageable bytes=" << cpu_staging_tensor.SizeInBytes(); + } + return Status::OK(); +#endif } } else { if (device == default_cpu_device) { diff --git a/onnxruntime/core/platform/env.h b/onnxruntime/core/platform/env.h index 8e0f6669a9dbc..aa191f103dce7 100644 --- a/onnxruntime/core/platform/env.h +++ b/onnxruntime/core/platform/env.h @@ -137,6 +137,24 @@ class RandomAccessFile { ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(RandomAccessFile); }; +#if !defined(ORT_MINIMAL_BUILD) +class PosixFileDescriptorProvider { + public: + virtual ~PosixFileDescriptorProvider() = default; + + // The descriptor remains owned by the provider and is valid only for its lifetime. + virtual int GetFileDescriptor() const = 0; +}; + +class WindowsFileHandleProvider { + public: + virtual ~WindowsFileHandleProvider() = default; + + // The handle remains owned by the provider and is valid only for its lifetime. + virtual void* GetFileHandle() const = 0; +}; +#endif + /// \brief An interface used by the onnxruntime implementation to /// access operating system functionality like the filesystem etc. /// diff --git a/onnxruntime/core/platform/posix/env.cc b/onnxruntime/core/platform/posix/env.cc index 43b2c4b9a73ae..8755acc613e6c 100644 --- a/onnxruntime/core/platform/posix/env.cc +++ b/onnxruntime/core/platform/posix/env.cc @@ -130,7 +130,12 @@ common::Status GetFileLength(int fd, size_t& file_size) { return common::Status::OK(); } -class PosixRandomAccessFile final : public RandomAccessFile { +class PosixRandomAccessFile final : public RandomAccessFile +#if !defined(ORT_MINIMAL_BUILD) + , + public PosixFileDescriptorProvider +#endif +{ public: PosixRandomAccessFile(ScopedFileDescriptor descriptor, std::string path) : descriptor_(std::move(descriptor)), path_(std::move(path)) {} @@ -163,6 +168,12 @@ class PosixRandomAccessFile final : public RandomAccessFile { return common::Status::OK(); } +#if !defined(ORT_MINIMAL_BUILD) + int GetFileDescriptor() const override { + return descriptor_.Get(); + } +#endif + private: ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(PosixRandomAccessFile); ScopedFileDescriptor descriptor_; diff --git a/onnxruntime/core/platform/windows/env.cc b/onnxruntime/core/platform/windows/env.cc index 33f8b3e20994d..250f5a116e8d9 100644 --- a/onnxruntime/core/platform/windows/env.cc +++ b/onnxruntime/core/platform/windows/env.cc @@ -360,11 +360,20 @@ common::Status WindowsEnv::GetFileLength(int fd, /*out*/ size_t& file_size) cons namespace { -class WindowsRandomAccessFile final : public RandomAccessFile { +class WindowsRandomAccessFile final : public RandomAccessFile +#if !defined(ORT_MINIMAL_BUILD) + , + public WindowsFileHandleProvider +#endif +{ public: explicit WindowsRandomAccessFile(wil::unique_hfile file_handle) : file_handle_(std::move(file_handle)) {} ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(WindowsRandomAccessFile); +#if !defined(ORT_MINIMAL_BUILD) + void* GetFileHandle() const override { return file_handle_.get(); } +#endif + Status GetLength(size_t& length) const override { LARGE_INTEGER file_size{}; if (!GetFileSizeEx(file_handle_.get(), &file_size)) { diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index e6645a0d4cb77..be973b5a8d578 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -3433,12 +3433,15 @@ std::unique_ptr CUDAExecutionProvider::GetDataTransf } std::unique_ptr CUDAExecutionProvider::GetExternalDataLoader() const { - if (info_.external_data_loader_reading_threads == 0) { + if (info_.external_data_loader_reading_threads == 0 && + !info_.external_data_loader_use_gds && + !info_.external_data_loader_use_directstorage) { return nullptr; } return std::make_unique( - info_.device_id, info_.external_data_loader_reading_threads); + info_.device_id, info_.external_data_loader_reading_threads, + info_.external_data_loader_use_gds, info_.external_data_loader_use_directstorage); } std::vector> diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc index 14899aee419ec..9a961e35b0680 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.cc @@ -39,6 +39,8 @@ constexpr const char* kUseTF32 = "use_tf32"; constexpr const char* kFuseConvBias = "fuse_conv_bias"; constexpr const char* kSdpaKernel = "sdpa_kernel"; constexpr const char* kExternalDataLoaderReadingThreads = "external_data_loader_reading_threads"; +constexpr const char* kExternalDataLoaderUseGds = "external_data_loader_use_gds"; +constexpr const char* kExternalDataLoaderUseDirectStorage = "external_data_loader_use_directstorage"; } // namespace provider_option_names } // namespace cuda @@ -146,6 +148,12 @@ CUDAExecutionProviderInfo CUDAExecutionProviderInfo::FromProviderOptions(const P OrtCUDAProviderOptionsV2::kMaxExternalDataLoaderReadingThreadCount, "."); return Status::OK(); }) + .AddAssignmentToReference( + cuda::provider_option_names::kExternalDataLoaderUseGds, + info.external_data_loader_use_gds) + .AddAssignmentToReference( + cuda::provider_option_names::kExternalDataLoaderUseDirectStorage, + info.external_data_loader_use_directstorage) .AddValueParser( cuda::provider_option_names::kTunableOpEnable, [&info](const std::string& value_str) -> Status { @@ -203,6 +211,10 @@ ProviderOptions CUDAExecutionProviderInfo::ToProviderOptions(const CUDAExecution {cuda::provider_option_names::kFuseConvBias, MakeStringWithClassicLocale(info.fuse_conv_bias)}, {cuda::provider_option_names::kExternalDataLoaderReadingThreads, MakeStringWithClassicLocale(info.external_data_loader_reading_threads)}, + {cuda::provider_option_names::kExternalDataLoaderUseGds, + MakeStringWithClassicLocale(info.external_data_loader_use_gds)}, + {cuda::provider_option_names::kExternalDataLoaderUseDirectStorage, + MakeStringWithClassicLocale(info.external_data_loader_use_directstorage)}, }; return options; @@ -230,6 +242,10 @@ ProviderOptions CUDAExecutionProviderInfo::ToProviderOptions(const OrtCUDAProvid {cuda::provider_option_names::kSdpaKernel, MakeStringWithClassicLocale(info.sdpa_kernel)}, {cuda::provider_option_names::kExternalDataLoaderReadingThreads, MakeStringWithClassicLocale(info.external_data_loader_reading_threads)}, + {cuda::provider_option_names::kExternalDataLoaderUseGds, + MakeStringWithClassicLocale(info.external_data_loader_use_gds)}, + {cuda::provider_option_names::kExternalDataLoaderUseDirectStorage, + MakeStringWithClassicLocale(info.external_data_loader_use_directstorage)}, }; return options; diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h index b4d4cb758d2f9..2a6826e926b84 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider_info.h @@ -84,9 +84,11 @@ struct CUDAExecutionProviderInfo { int sdpa_kernel{0}; - // 0 disables the custom external-data loader and retains the framework's existing path. + // 0 disables the pinned-buffer loader and uses the pageable fallback if direct storage is unavailable. // 1 uses the pinned-buffer loader with synchronous reads. 2..64 use that many parallel read tasks per block. size_t external_data_loader_reading_threads{4}; + bool external_data_loader_use_gds{false}; + bool external_data_loader_use_directstorage{false}; static CUDAExecutionProviderInfo FromProviderOptions(const ProviderOptions& options); static ProviderOptions ToProviderOptions(const CUDAExecutionProviderInfo& info); @@ -122,6 +124,8 @@ struct std::hash<::onnxruntime::CUDAExecutionProviderInfo> { onnxruntime::HashCombine(info.sdpa_kernel, value); onnxruntime::HashCombine(info.enable_cudnn, value); onnxruntime::HashCombine(info.external_data_loader_reading_threads, value); + onnxruntime::HashCombine(info.external_data_loader_use_gds, value); + onnxruntime::HashCombine(info.external_data_loader_use_directstorage, value); // Memory pointers onnxruntime::HashCombine(reinterpret_cast(info.user_compute_stream), value); diff --git a/onnxruntime/core/providers/cuda/cuda_external_data_loader.cc b/onnxruntime/core/providers/cuda/cuda_external_data_loader.cc index a587c252c3b34..34468355a2307 100644 --- a/onnxruntime/core/providers/cuda/cuda_external_data_loader.cc +++ b/onnxruntime/core/providers/cuda/cuda_external_data_loader.cc @@ -116,13 +116,21 @@ common::Status LoadWithPageableBuffer(const RandomAccessFile& file, FileOffsetTy } // namespace -ExternalDataLoader::ExternalDataLoader(int device_id, size_t reading_thread_count, - AllocatePinnedBufferFn allocate_pinned_buffer, - CreateStreamFn create_stream) +ExternalDataLoader::ExternalDataLoader(int device_id, size_t reading_thread_count, bool use_gds, + bool use_directstorage) : device_id_(device_id), - reading_thread_count_(reading_thread_count), - allocate_pinned_buffer_(allocate_pinned_buffer), - create_stream_(create_stream) {} + reading_thread_count_(reading_thread_count) +#if !defined(ORT_MINIMAL_BUILD) && !defined(USE_CUDA_MINIMAL) + , + use_gds_(use_gds), + use_directstorage_(use_directstorage) +#endif +{ +#if defined(ORT_MINIMAL_BUILD) || defined(USE_CUDA_MINIMAL) + ORT_UNUSED_PARAMETER(use_gds); + ORT_UNUSED_PARAMETER(use_directstorage); +#endif +} ExternalDataLoader::~ExternalDataLoader() { reader_pool_.reset(); @@ -143,13 +151,13 @@ common::Status ExternalDataLoader::EnsureResources() const { } for (size_t i = 0; i < buffers_.size(); ++i) { - auto status = CUDA_CALL(allocate_pinned_buffer_(&buffers_[i], kExternalDataLoaderBufferSize)); + auto status = CUDA_CALL(cudaMallocHost(&buffers_[i], kExternalDataLoaderBufferSize)); if (!status.IsOK()) { ReleaseResources(); return status; } - status = CUDA_CALL(create_stream_(&streams_[i], cudaStreamNonBlocking)); + status = CUDA_CALL(cudaStreamCreateWithFlags(&streams_[i], cudaStreamNonBlocking)); if (!status.IsOK()) { ReleaseResources(); return status; @@ -166,6 +174,11 @@ void ExternalDataLoader::ReleaseResources() const noexcept { previous_device != device_id_ && cudaSetDevice(device_id_) == cudaSuccess; +#if !defined(ORT_MINIMAL_BUILD) && !defined(USE_CUDA_MINIMAL) + gds_loader_.reset(); + directstorage_loader_.reset(); +#endif + for (auto& stream : streams_) { if (stream != nullptr) { ORT_IGNORE_RETURN_VALUE(CUDA_CALL(cudaStreamSynchronize(stream))); @@ -210,11 +223,83 @@ common::Status ExternalDataLoader::LoadTensor(const Env& env, std::lock_guard lock(mutex_); CudaDeviceGuard device_guard; ORT_RETURN_IF_ERROR(device_guard.SetDevice(device_id_)); + +#if !defined(ORT_MINIMAL_BUILD) && !defined(USE_CUDA_MINIMAL) + if (use_directstorage_ && !directstorage_disabled_ && length != 0 && + std::endian::native == std::endian::little && !tensor.IsDataType()) { + Status status = Status::OK(); + if (!directstorage_loader_) { + status = DirectStorageLoader::Create(device_id_, directstorage_loader_); + } + if (status.IsOK()) { + void* handle = nullptr; +#if !defined(ORT_NO_RTTI) + const auto* provider = dynamic_cast(file.get()); + if (provider != nullptr) { + handle = provider->GetFileHandle(); + } +#endif + status = directstorage_loader_->Load(data_file_path, handle, data_offset, length, tensor); + } + if (status.IsOK()) { + LOGS_DEFAULT(INFO) << "CUDA external data loader: path=directstorage bytes=" << length; + return Status::OK(); + } + directstorage_disabled_ = true; + directstorage_loader_.reset(); + LOGS_DEFAULT(WARNING) << "Microsoft DirectStorage could not load external data; falling back to another CUDA " + << "external-data loading path. " << status.ErrorMessage(); + } + + const bool gds_range_is_aligned = + data_offset % static_cast(kGdsIoAlignment) == 0 && + length % kGdsIoAlignment == 0; + if (use_gds_ && !gds_disabled_ && gds_range_is_aligned && + std::endian::native == std::endian::little && + !tensor.IsDataType()) { + Status gds_status = Status::OK(); + if (!gds_loader_) { + gds_status = GdsLoader::Create(device_id_, gds_loader_); + } + if (gds_status.IsOK()) { +#if defined(ORT_NO_RTTI) + constexpr int file_descriptor = -1; +#else + const auto* descriptor_provider = dynamic_cast(file.get()); + const int file_descriptor = + descriptor_provider == nullptr ? -1 : descriptor_provider->GetFileDescriptor(); +#endif + gds_status = gds_loader_->Load( + file_descriptor, data_offset, length, tensor); + } + if (gds_status.IsOK()) { + LOGS_DEFAULT(INFO) << "CUDA external data loader: path=gds bytes=" << length; + return Status::OK(); + } + + gds_disabled_ = true; + gds_loader_.reset(); + LOGS_DEFAULT(WARNING) << "GPUDirect Storage could not load external data; falling back to the CUDA " + << (reading_thread_count_ == 0 ? "pageable-buffer" : "pinned-buffer") + << " loader. " + << gds_status.ErrorMessage(); + } +#endif + + if (reading_thread_count_ == 0) { + ORT_RETURN_IF_ERROR(LoadWithPageableBuffer(*file, data_offset, length, tensor, 1, reader_pool_)); + LOGS_DEFAULT(INFO) << "CUDA external data loader: path=pageable bytes=" << length; + return Status::OK(); + } + const auto resource_status = EnsureResources(); if (!resource_status.IsOK()) { - // TODO: Remember setup failures during initialization and report the first CUDA error - // so later initializers do not repeatedly retry unavailable pinned buffers or streams. - return LoadWithPageableBuffer(*file, data_offset, length, tensor, reading_thread_count_, reader_pool_); + LOGS_DEFAULT(WARNING) << "CUDA pinned-buffer setup failed; falling back to pageable memory. " + << resource_status.ErrorMessage(); + ORT_RETURN_IF_ERROR( + LoadWithPageableBuffer(*file, data_offset, length, tensor, reading_thread_count_, reader_pool_)); + LOGS_DEFAULT(INFO) << "CUDA external data loader: path=pageable bytes=" << length; + return Status::OK(); } auto* destination = static_cast(tensor.MutableDataRaw()); @@ -278,7 +363,9 @@ common::Status ExternalDataLoader::LoadTensor(const Env& env, offset += chunk_size; } - return synchronize_streams(); + ORT_RETURN_IF_ERROR(synchronize_streams()); + LOGS_DEFAULT(INFO) << "CUDA external data loader: path=pinned bytes=" << length; + return Status::OK(); } } // namespace cuda diff --git a/onnxruntime/core/providers/cuda/cuda_external_data_loader.h b/onnxruntime/core/providers/cuda/cuda_external_data_loader.h index bcd795021a71a..5cffc09642578 100644 --- a/onnxruntime/core/providers/cuda/cuda_external_data_loader.h +++ b/onnxruntime/core/providers/cuda/cuda_external_data_loader.h @@ -8,6 +8,10 @@ #include #include "core/framework/external_data_loader.h" +#if !defined(ORT_MINIMAL_BUILD) && !defined(USE_CUDA_MINIMAL) +#include "core/providers/cuda/cuda_external_data_loader_gds.h" +#include "core/providers/cuda/cuda_external_data_loader_directstorage.h" +#endif #include "cuda_pch.h" namespace onnxruntime { @@ -64,12 +68,8 @@ class ExternalDataLoaderThreadPool; */ class ExternalDataLoader final : public IExternalDataLoader { public: - using AllocatePinnedBufferFn = cudaError_t (*)(void**, size_t); - using CreateStreamFn = cudaError_t (*)(cudaStream_t*, unsigned int); - - ExternalDataLoader(int device_id, size_t reading_thread_count, - AllocatePinnedBufferFn allocate_pinned_buffer = cudaMallocHost, - CreateStreamFn create_stream = cudaStreamCreateWithFlags); + ExternalDataLoader(int device_id, size_t reading_thread_count, bool use_gds = false, + bool use_directstorage = false); ~ExternalDataLoader() override; bool CanLoad(const OrtMemoryInfo& target_memory_info) const override; @@ -89,8 +89,14 @@ class ExternalDataLoader final : public IExternalDataLoader { mutable std::array buffers_{}; mutable std::array streams_{}; const size_t reading_thread_count_; - const AllocatePinnedBufferFn allocate_pinned_buffer_; - const CreateStreamFn create_stream_; +#if !defined(ORT_MINIMAL_BUILD) && !defined(USE_CUDA_MINIMAL) + const bool use_gds_; + mutable bool gds_disabled_{false}; + mutable std::unique_ptr gds_loader_; + const bool use_directstorage_; + mutable bool directstorage_disabled_{false}; + mutable std::unique_ptr directstorage_loader_; +#endif mutable std::unique_ptr reader_pool_; }; diff --git a/onnxruntime/core/providers/cuda/cuda_external_data_loader_directstorage.cc b/onnxruntime/core/providers/cuda/cuda_external_data_loader_directstorage.cc new file mode 100644 index 0000000000000..88b9155702ed4 --- /dev/null +++ b/onnxruntime/core/providers/cuda/cuda_external_data_loader_directstorage.cc @@ -0,0 +1,255 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// provider_api.h must be first to set SHARED_PROVIDER +#include "core/providers/shared_library/provider_api.h" + +#include "core/providers/cuda/cuda_external_data_loader_directstorage.h" + +#if defined(ORT_CUDA_DIRECTSTORAGE_AVAILABLE) +#include +#include +#include +#include +#include +#include + +#include "core/common/safeint.h" +#include "core/providers/cuda/cuda_common.h" +#endif + +namespace onnxruntime { +namespace cuda { +namespace { + +#if defined(ORT_CUDA_DIRECTSTORAGE_AVAILABLE) + +using Microsoft::WRL::ComPtr; +constexpr size_t kDirectStorageBufferSize = 32 * 1024 * 1024; + +common::Status CheckHResult(HRESULT result, const char* operation) { + ORT_RETURN_IF(FAILED(result), operation, " failed, HRESULT=", static_cast(result)); + return Status::OK(); +} + +class DirectStorageLibrary { + public: + DirectStorageLibrary() = default; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(DirectStorageLibrary); + ~DirectStorageLibrary() { + if (handle != nullptr) { + FreeLibrary(handle); + } + } + HMODULE handle{nullptr}; +}; + +class WindowsDirectStorageLoader final : public DirectStorageLoader { + public: + ~WindowsDirectStorageLoader() override { + if (stream_ != nullptr) { + ORT_IGNORE_RETURN_VALUE(CUDA_CALL(cudaStreamSynchronize(stream_))); + } + if (queue_) { + queue_->Close(); + } + if (stream_ != nullptr) { + ORT_IGNORE_RETURN_VALUE(CUDA_CALL(cudaStreamDestroy(stream_))); + } + if (cuda_fence_ != nullptr) { + ORT_IGNORE_RETURN_VALUE(CUDA_CALL(cudaDestroyExternalSemaphore(cuda_fence_))); + } + if (buffer_ != nullptr) { + ORT_IGNORE_RETURN_VALUE(CUDA_CALL(cudaFree(buffer_))); + } + if (cuda_memory_ != nullptr) { + ORT_IGNORE_RETURN_VALUE(CUDA_CALL(cudaDestroyExternalMemory(cuda_memory_))); + } + } + + static common::Status Create(int device_id, std::unique_ptr& loader) { + auto candidate = std::unique_ptr(new WindowsDirectStorageLoader()); + ORT_RETURN_IF_ERROR(candidate->Initialize(device_id)); + loader = std::move(candidate); + return Status::OK(); + } + + common::Status Load(const std::filesystem::path& path, void* validated_file_handle, + int64_t data_offset, size_t data_length, Tensor& tensor) override { + ORT_RETURN_IF(validated_file_handle == nullptr, + "Microsoft DirectStorage requires the validated Windows file handle."); + BY_HANDLE_FILE_INFORMATION original{}; + ORT_RETURN_IF_NOT(GetFileInformationByHandle(validated_file_handle, &original), + "GetFileInformationByHandle failed: ", GetLastError()); + ComPtr file; + ORT_RETURN_IF_ERROR(CheckHResult(factory_->OpenFile(path.c_str(), IID_PPV_ARGS(&file)), + "DirectStorage OpenFile")); + bool pending = false; + auto drain_on_error = gsl::finally([&]() { + if (pending) { + queue_->Close(); + } + }); + BY_HANDLE_FILE_INFORMATION opened{}; + ORT_RETURN_IF_ERROR(CheckHResult(file->GetFileInformation(&opened), "DirectStorage GetFileInformation")); + // DirectStorage opens by path. Never read a replacement for the file validated by the caller. + ORT_RETURN_IF(original.dwVolumeSerialNumber != opened.dwVolumeSerialNumber || + original.nFileIndexHigh != opened.nFileIndexHigh || + original.nFileIndexLow != opened.nFileIndexLow, + "External-data file changed before DirectStorage opened it."); + const uint64_t file_size = (static_cast(opened.nFileSizeHigh) << 32) | opened.nFileSizeLow; + ORT_RETURN_IF(data_offset < 0 || static_cast(data_offset) > file_size || + data_length > file_size - static_cast(data_offset), + "DirectStorage external-data range is outside the file."); + + auto* destination = static_cast(tensor.MutableDataRaw()); + for (size_t offset = 0; offset < data_length;) { + const auto chunk = static_cast(std::min(kDirectStorageBufferSize, data_length - offset)); + DSTORAGE_REQUEST request{}; + request.Options.SourceType = DSTORAGE_REQUEST_SOURCE_FILE; + request.Options.DestinationType = DSTORAGE_REQUEST_DESTINATION_BUFFER; + request.Source.File.Source = file.Get(); + request.Source.File.Offset = SafeInt(data_offset) + offset; + request.Source.File.Size = chunk; + request.UncompressedSize = chunk; + request.Destination.Buffer.Resource = resource_.Get(); + request.Destination.Buffer.Size = chunk; + queue_->EnqueueRequest(&request); + queue_->EnqueueStatus(status_.Get(), 0); + queue_->EnqueueSignal(fence_.Get(), ++fence_value_); + queue_->Submit(); + pending = true; + + cudaExternalSemaphoreWaitParams wait{}; + wait.params.fence.value = fence_value_; + CUDA_RETURN_IF_ERROR(cudaWaitExternalSemaphoresAsync(&cuda_fence_, &wait, 1, stream_)); + CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream_)); + pending = false; + ORT_RETURN_IF_ERROR(CheckHResult(status_->GetHResult(0), "DirectStorage read")); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(destination + offset, buffer_, chunk, + cudaMemcpyDeviceToDevice, stream_)); + // Complete CUDA reads before DirectStorage writes the shared buffer again. + CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream_)); + offset += chunk; + } + return Status::OK(); + } + + private: + WindowsDirectStorageLoader() = default; + + common::Status Initialize(int device_id) { + CUDA_RETURN_IF_ERROR(cudaSetDevice(device_id)); + cudaDeviceProp properties{}; + CUDA_RETURN_IF_ERROR(cudaGetDeviceProperties(&properties, device_id)); + const unsigned int node_mask = properties.luidDeviceNodeMask; + ORT_RETURN_IF(node_mask == 0 || (node_mask & (node_mask - 1)) != 0, + "DirectStorage requires a single CUDA device node."); + LUID adapter_luid{}; + static_assert(sizeof(properties.luid) == sizeof(adapter_luid)); + std::memcpy(&adapter_luid, properties.luid, sizeof(adapter_luid)); + ComPtr dxgi; + ORT_RETURN_IF_ERROR(CheckHResult(CreateDXGIFactory1(IID_PPV_ARGS(&dxgi)), "CreateDXGIFactory1")); + ComPtr adapter; + ORT_RETURN_IF_ERROR(CheckHResult(dxgi->EnumAdapterByLuid(adapter_luid, IID_PPV_ARGS(&adapter)), + "Find CUDA DXGI adapter")); + ORT_RETURN_IF_ERROR(CheckHResult( + D3D12CreateDevice(adapter.Get(), D3D_FEATURE_LEVEL_11_0, IID_PPV_ARGS(&device_)), "D3D12CreateDevice")); + ORT_RETURN_IF(device_->GetNodeCount() != 1, "DirectStorage does not support linked D3D12 adapters."); + + library_.handle = LoadLibraryExW(L"dstorage.dll", nullptr, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS); + ORT_RETURN_IF(library_.handle == nullptr, "Microsoft DirectStorage is unavailable: dstorage.dll load failed: ", + GetLastError()); + const auto get_factory = + reinterpret_cast(GetProcAddress(library_.handle, "DStorageGetFactory")); + ORT_RETURN_IF(get_factory == nullptr, "DStorageGetFactory is unavailable: ", GetLastError()); + ORT_RETURN_IF_ERROR(CheckHResult(get_factory(IID_PPV_ARGS(&factory_)), "DStorageGetFactory")); + // The default DirectStorage staging size is 32 MiB. Do not reconfigure its process-wide factory. + DSTORAGE_QUEUE_DESC queue_desc{}; + queue_desc.SourceType = DSTORAGE_REQUEST_SOURCE_FILE; + queue_desc.Capacity = DSTORAGE_MIN_QUEUE_CAPACITY; + queue_desc.Priority = DSTORAGE_PRIORITY_NORMAL; + queue_desc.Device = device_.Get(); + ORT_RETURN_IF_ERROR(CheckHResult(factory_->CreateQueue(&queue_desc, IID_PPV_ARGS(&queue_)), + "DirectStorage CreateQueue")); + ORT_RETURN_IF_ERROR(CheckHResult(factory_->CreateStatusArray(1, "ORT external data", IID_PPV_ARGS(&status_)), + "DirectStorage CreateStatusArray")); + + D3D12_HEAP_PROPERTIES heap{}; + heap.Type = D3D12_HEAP_TYPE_DEFAULT; + heap.CreationNodeMask = node_mask; + heap.VisibleNodeMask = node_mask; + D3D12_RESOURCE_DESC desc{}; + desc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER; + desc.Width = kDirectStorageBufferSize; + desc.Height = 1; + desc.DepthOrArraySize = 1; + desc.MipLevels = 1; + desc.SampleDesc.Count = 1; + desc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; + ORT_RETURN_IF_ERROR(CheckHResult( + device_->CreateCommittedResource(&heap, D3D12_HEAP_FLAG_SHARED, &desc, + D3D12_RESOURCE_STATE_COMMON, nullptr, IID_PPV_ARGS(&resource_)), + "Create shared DirectStorage buffer")); + HANDLE shared_memory = nullptr; + ORT_RETURN_IF_ERROR(CheckHResult( + device_->CreateSharedHandle(resource_.Get(), nullptr, GENERIC_ALL, nullptr, &shared_memory), + "Share DirectStorage buffer")); + auto close_memory = gsl::finally([&]() { CloseHandle(shared_memory); }); + cudaExternalMemoryHandleDesc memory_desc{}; + memory_desc.type = cudaExternalMemoryHandleTypeD3D12Resource; + memory_desc.handle.win32.handle = shared_memory; + memory_desc.size = device_->GetResourceAllocationInfo(node_mask, 1, &desc).SizeInBytes; + memory_desc.flags = cudaExternalMemoryDedicated; + CUDA_RETURN_IF_ERROR(cudaImportExternalMemory(&cuda_memory_, &memory_desc)); + cudaExternalMemoryBufferDesc buffer_desc{}; + buffer_desc.size = kDirectStorageBufferSize; + CUDA_RETURN_IF_ERROR(cudaExternalMemoryGetMappedBuffer(&buffer_, cuda_memory_, &buffer_desc)); + + ORT_RETURN_IF_ERROR(CheckHResult( + device_->CreateFence(0, D3D12_FENCE_FLAG_SHARED, IID_PPV_ARGS(&fence_)), "Create DirectStorage fence")); + HANDLE shared_fence = nullptr; + ORT_RETURN_IF_ERROR(CheckHResult( + device_->CreateSharedHandle(fence_.Get(), nullptr, GENERIC_ALL, nullptr, &shared_fence), + "Share DirectStorage fence")); + auto close_fence = gsl::finally([&]() { CloseHandle(shared_fence); }); + cudaExternalSemaphoreHandleDesc fence_desc{}; + fence_desc.type = cudaExternalSemaphoreHandleTypeD3D12Fence; + fence_desc.handle.win32.handle = shared_fence; + CUDA_RETURN_IF_ERROR(cudaImportExternalSemaphore(&cuda_fence_, &fence_desc)); + CUDA_RETURN_IF_ERROR(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking)); + return Status::OK(); + } + + DirectStorageLibrary library_; + ComPtr device_; + ComPtr factory_; + ComPtr resource_; + ComPtr fence_; + ComPtr status_; + ComPtr queue_; + cudaExternalMemory_t cuda_memory_{nullptr}; + cudaExternalSemaphore_t cuda_fence_{nullptr}; + cudaStream_t stream_{nullptr}; + void* buffer_{nullptr}; + uint64_t fence_value_{0}; +}; + +#endif + +} // namespace + +common::Status DirectStorageLoader::Create(int device_id, std::unique_ptr& loader) { +#if defined(ORT_CUDA_DIRECTSTORAGE_AVAILABLE) + return WindowsDirectStorageLoader::Create(device_id, loader); +#else + ORT_UNUSED_PARAMETER(device_id); + ORT_UNUSED_PARAMETER(loader); + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "Microsoft DirectStorage requires a Windows CUDA build with " + "onnxruntime_USE_CUDA_DIRECTSTORAGE=ON."); +#endif +} + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cuda_external_data_loader_directstorage.h b/onnxruntime/core/providers/cuda/cuda_external_data_loader_directstorage.h new file mode 100644 index 0000000000000..9821a46427adb --- /dev/null +++ b/onnxruntime/core/providers/cuda/cuda_external_data_loader_directstorage.h @@ -0,0 +1,36 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include + +#include "core/common/common.h" +#include "core/common/status.h" + +namespace onnxruntime { +#ifndef SHARED_PROVIDER +class Tensor; +#endif + +namespace cuda { + +class DirectStorageLoader { + public: + virtual ~DirectStorageLoader() = default; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(DirectStorageLoader); + + virtual common::Status Load(const std::filesystem::path& path, void* validated_file_handle, + int64_t data_offset, size_t data_length, Tensor& tensor) = 0; + + static common::Status Create(int device_id, std::unique_ptr& loader); + + protected: + DirectStorageLoader() = default; +}; + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cuda_external_data_loader_gds.cc b/onnxruntime/core/providers/cuda/cuda_external_data_loader_gds.cc new file mode 100644 index 0000000000000..9ce52309a7be8 --- /dev/null +++ b/onnxruntime/core/providers/cuda/cuda_external_data_loader_gds.cc @@ -0,0 +1,252 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// provider_api.h must be first to set SHARED_PROVIDER +#include "core/providers/shared_library/provider_api.h" + +#include "core/providers/cuda/cuda_external_data_loader_gds.h" + +#include +#include +#include +#include + +#include "core/common/common.h" +#include "core/common/safeint.h" +#include "core/providers/cuda/cuda_common.h" + +#if defined(ORT_CUDA_GDS_AVAILABLE) +#include +#include +#include +#include +#endif + +namespace onnxruntime { +namespace cuda { +namespace { + +#if defined(ORT_CUDA_GDS_AVAILABLE) + +constexpr size_t kGdsBufferSize = 64 * 1024 * 1024; + +template +common::Status LoadSymbol(void* library, const char* name, T& function) { + dlerror(); + function = reinterpret_cast(dlsym(library, name)); + const char* error = dlerror(); + ORT_RETURN_IF(function == nullptr || error != nullptr, + "Failed to load ", name, " from libcufile: ", + error == nullptr ? "symbol not found" : error); + return Status::OK(); +} + +common::Status CheckCuFileStatus(CUfileError_t status, std::string_view operation) { + ORT_RETURN_IF(status.err != CU_FILE_SUCCESS, operation, " failed: ", + cufileop_status_error(status.err), " (", static_cast(status.err), ")"); + return Status::OK(); +} + +class CuFileDriver { + public: + using DriverOpenFn = decltype(&cuFileDriverOpen); + using DriverCloseFn = CUfileError_t (*)(); + using HandleRegisterFn = decltype(&cuFileHandleRegister); + using HandleDeregisterFn = decltype(&cuFileHandleDeregister); + using BufferRegisterFn = decltype(&cuFileBufRegister); + using BufferDeregisterFn = decltype(&cuFileBufDeregister); + using SetBoolParameterFn = decltype(&cuFileSetParameterBool); + using ReadFn = decltype(&cuFileRead); + + CuFileDriver() = default; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CuFileDriver); + + ~CuFileDriver() { + if (driver_initialized_) { + ORT_IGNORE_RETURN_VALUE(driver_close_()); + } + if (library_ != nullptr) { + ORT_IGNORE_RETURN_VALUE(dlclose(library_)); + } + } + + CUfileError_t RegisterHandle(CUfileHandle_t* handle, CUfileDescr_t* descriptor) const { + return handle_register_(handle, descriptor); + } + + void DeregisterHandle(CUfileHandle_t handle) const { + handle_deregister_(handle); + } + + CUfileError_t RegisterBuffer(const void* buffer, size_t length) const { + return buffer_register_(buffer, length, 0); + } + + CUfileError_t DeregisterBuffer(const void* buffer) const { + return buffer_deregister_(buffer); + } + + ssize_t Read(CUfileHandle_t handle, void* buffer, size_t length, + off_t file_offset, off_t buffer_offset) const { + return read_(handle, buffer, length, file_offset, buffer_offset); + } + + common::Status Initialize() { + library_ = dlopen("libcufile.so.0", RTLD_NOW | RTLD_LOCAL); + if (library_ == nullptr) { + library_ = dlopen("libcufile.so", RTLD_NOW | RTLD_LOCAL); + } + const char* library_error = dlerror(); + ORT_RETURN_IF(library_ == nullptr, "GPUDirect Storage is unavailable: ", + library_error == nullptr ? "libcufile could not be loaded" : library_error); + + ORT_RETURN_IF_ERROR(LoadSymbol(library_, "cuFileDriverOpen", driver_open_)); + auto close_status = LoadSymbol(library_, "cuFileDriverClose_v2", driver_close_); + if (!close_status.IsOK()) { + ORT_RETURN_IF_ERROR(LoadSymbol(library_, "cuFileDriverClose", driver_close_)); + } + ORT_RETURN_IF_ERROR(LoadSymbol(library_, "cuFileHandleRegister", handle_register_)); + ORT_RETURN_IF_ERROR(LoadSymbol(library_, "cuFileHandleDeregister", handle_deregister_)); + ORT_RETURN_IF_ERROR(LoadSymbol(library_, "cuFileBufRegister", buffer_register_)); + ORT_RETURN_IF_ERROR(LoadSymbol(library_, "cuFileBufDeregister", buffer_deregister_)); + ORT_RETURN_IF_ERROR(LoadSymbol(library_, "cuFileSetParameterBool", set_bool_parameter_)); + ORT_RETURN_IF_ERROR(LoadSymbol(library_, "cuFileRead", read_)); + + ORT_RETURN_IF_ERROR(CheckCuFileStatus( + set_bool_parameter_(CUFILE_PARAM_USE_PCIP2PDMA, true), + "Enabling cuFile PCI P2PDMA")); + ORT_RETURN_IF_ERROR(CheckCuFileStatus( + set_bool_parameter_(CUFILE_PARAM_PROPERTIES_ALLOW_COMPAT_MODE, false), + "Disabling cuFile compatibility mode")); + ORT_RETURN_IF_ERROR(CheckCuFileStatus(driver_open_(), "cuFileDriverOpen")); + driver_initialized_ = true; + return Status::OK(); + } + + private: + void* library_{nullptr}; + bool driver_initialized_{false}; + DriverOpenFn driver_open_{nullptr}; + DriverCloseFn driver_close_{nullptr}; + HandleRegisterFn handle_register_{nullptr}; + HandleDeregisterFn handle_deregister_{nullptr}; + BufferRegisterFn buffer_register_{nullptr}; + BufferDeregisterFn buffer_deregister_{nullptr}; + SetBoolParameterFn set_bool_parameter_{nullptr}; + ReadFn read_{nullptr}; +}; + +class LinuxGdsLoader final : public GdsLoader { + public: + ~LinuxGdsLoader() override { + if (gds_buffer_registered_) { + ORT_IGNORE_RETURN_VALUE(driver_->DeregisterBuffer(gds_buffer_)); + } + if (gds_buffer_ != nullptr) { + ORT_IGNORE_RETURN_VALUE(CUDA_CALL(cudaFree(gds_buffer_))); + } + } + + static common::Status Create(int device_id, std::unique_ptr& loader) { + auto candidate = std::unique_ptr(new LinuxGdsLoader()); + ORT_RETURN_IF_ERROR(candidate->Initialize(device_id)); + loader = std::move(candidate); + return Status::OK(); + } + + common::Status Load(int file_descriptor, + int64_t data_offset, + size_t data_length, + Tensor& tensor) const override { + ORT_RETURN_IF(file_descriptor < 0, + "GPUDirect Storage requires an open POSIX file descriptor."); + + const int direct_descriptor = fcntl(file_descriptor, F_DUPFD_CLOEXEC, 0); + ORT_RETURN_IF(direct_descriptor < 0, "Failed to duplicate external-data file descriptor: ", + std::strerror(errno)); + auto close_file = gsl::finally([direct_descriptor]() { + ORT_IGNORE_RETURN_VALUE(close(direct_descriptor)); + }); + + const int original_flags = fcntl(direct_descriptor, F_GETFL); + ORT_RETURN_IF(original_flags < 0, "Failed to query external-data file flags: ", + std::strerror(errno)); + ORT_RETURN_IF(fcntl(direct_descriptor, F_SETFL, original_flags | O_DIRECT) < 0, + "Failed to enable O_DIRECT for GPUDirect Storage: ", std::strerror(errno)); + auto restore_flags = gsl::finally([direct_descriptor, original_flags]() { + ORT_IGNORE_RETURN_VALUE(fcntl(direct_descriptor, F_SETFL, original_flags)); + }); + + CUfileDescr_t descriptor{}; + descriptor.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD; + descriptor.handle.fd = direct_descriptor; + CUfileHandle_t file_handle = nullptr; + ORT_RETURN_IF_ERROR(CheckCuFileStatus(driver_->RegisterHandle(&file_handle, &descriptor), + "cuFileHandleRegister")); + auto deregister_file = gsl::finally([&]() { driver_->DeregisterHandle(file_handle); }); + + auto* destination = static_cast(tensor.MutableDataRaw()); + for (size_t offset = 0; offset < data_length;) { + const size_t chunk_size = std::min(kGdsBufferSize, data_length - offset); + const auto file_offset = SafeInt(data_offset) + offset; + const ssize_t bytes_read = driver_->Read(file_handle, gds_buffer_, chunk_size, file_offset, 0); + if (bytes_read != static_cast(chunk_size)) { + if (bytes_read == -1) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "cuFileRead failed: ", std::strerror(errno)); + } + if (bytes_read < 0) { + const auto cu_file_error = static_cast(-bytes_read); + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "cuFileRead failed: ", + cufileop_status_error(cu_file_error), + " (", static_cast(cu_file_error), ")"); + } + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "cuFileRead returned ", bytes_read, + " bytes; expected ", chunk_size, "."); + } + + CUDA_RETURN_IF_ERROR( + cudaMemcpy(destination + offset, gds_buffer_, chunk_size, cudaMemcpyDeviceToDevice)); + CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(nullptr)); + offset += chunk_size; + } + + return Status::OK(); + } + + private: + common::Status Initialize(int device_id) { + ORT_RETURN_IF_ERROR(driver_.Acquire()); + + CUDA_RETURN_IF_ERROR(cudaSetDevice(device_id)); + CUDA_RETURN_IF_ERROR(cudaMalloc(&gds_buffer_, kGdsBufferSize)); + ORT_RETURN_IF_ERROR(CheckCuFileStatus( + driver_->RegisterBuffer(gds_buffer_, kGdsBufferSize), "cuFileBufRegister")); + gds_buffer_registered_ = true; + return Status::OK(); + } + + LinuxGdsLoader() = default; + + GdsDriverHandle driver_; + void* gds_buffer_{nullptr}; + bool gds_buffer_registered_{false}; +}; + +#endif + +} // namespace + +common::Status GdsLoader::Create(int device_id, std::unique_ptr& loader) { +#if defined(ORT_CUDA_GDS_AVAILABLE) + return LinuxGdsLoader::Create(device_id, loader); +#else + ORT_UNUSED_PARAMETER(device_id); + ORT_UNUSED_PARAMETER(loader); + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "GPUDirect Storage requires Linux and cuFile headers with cuFileSetParameterBool, " + "CUFILE_PARAM_USE_PCIP2PDMA, and CUFILE_PARAM_PROPERTIES_ALLOW_COMPAT_MODE."); +#endif +} + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cuda_external_data_loader_gds.h b/onnxruntime/core/providers/cuda/cuda_external_data_loader_gds.h new file mode 100644 index 0000000000000..2b15fdd3a0cd6 --- /dev/null +++ b/onnxruntime/core/providers/cuda/cuda_external_data_loader_gds.h @@ -0,0 +1,84 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include +#include + +#include "core/common/common.h" +#include "core/common/status.h" + +namespace onnxruntime { +#ifndef SHARED_PROVIDER +class Tensor; +#endif + +namespace cuda { + +inline constexpr size_t kGdsIoAlignment = 4096; + +template +class GdsDriverHandle { + public: + GdsDriverHandle() : state_(SharedState()) {} + + ~GdsDriverHandle() { + // Lock before the final strong reference expires, not inside Driver's destructor. + std::lock_guard lock(state_->mutex); + driver_.reset(); + } + + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(GdsDriverHandle); + + template + common::Status Acquire(Args&&... args) { + std::lock_guard lock(state_->mutex); + if (driver_) { + return common::Status::OK(); + } + + driver_ = state_->driver.lock(); + if (!driver_) { + auto candidate = std::make_shared(std::forward(args)...); + ORT_RETURN_IF_ERROR(candidate->Initialize()); + state_->driver = candidate; + driver_ = std::move(candidate); + } + return common::Status::OK(); + } + + Driver* operator->() const { return driver_.get(); } + + private: + struct State { + std::mutex mutex; + std::weak_ptr driver; + }; + + static std::shared_ptr SharedState() { + static auto state = std::make_shared(); + return state; + } + + std::shared_ptr state_; + std::shared_ptr driver_; +}; + +class GdsLoader { + public: + virtual ~GdsLoader() = default; + + virtual common::Status Load(int file_descriptor, + int64_t data_offset, + size_t data_length, + Tensor& tensor) const = 0; + + static common::Status Create(int device_id, std::unique_ptr& loader); +}; + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index 00519584950b8..ffee8f37c18eb 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -210,6 +210,12 @@ struct CUDA_Provider : Provider { OrtCUDAProviderOptionsV2::kMaxExternalDataLoaderReadingThreadCount, "external_data_loader_reading_threads must be between 0 and ", OrtCUDAProviderOptionsV2::kMaxExternalDataLoaderReadingThreadCount, "."); + ORT_ENFORCE(params->external_data_loader_use_gds == 0 || + params->external_data_loader_use_gds == 1, + "external_data_loader_use_gds must be 0 or 1."); + ORT_ENFORCE(params->external_data_loader_use_directstorage == 0 || + params->external_data_loader_use_directstorage == 1, + "external_data_loader_use_directstorage must be 0 or 1."); // Calling a function like ::cudaDeviceSynchronize will cause CUDA to ensure there is binary code for the current GPU architecture // Ideally this will be already part of the binary, but if not, CUDA will JIT it during this call. This can take a very long time @@ -251,6 +257,8 @@ struct CUDA_Provider : Provider { info.use_tf32 = params->use_tf32 != 0; info.sdpa_kernel = params->sdpa_kernel; info.external_data_loader_reading_threads = params->external_data_loader_reading_threads; + info.external_data_loader_use_gds = params->external_data_loader_use_gds != 0; + info.external_data_loader_use_directstorage = params->external_data_loader_use_directstorage != 0; return std::make_shared(info); } @@ -288,6 +296,10 @@ struct CUDA_Provider : Provider { cuda_options.fuse_conv_bias = internal_options.fuse_conv_bias; cuda_options.external_data_loader_reading_threads = internal_options.external_data_loader_reading_threads; + cuda_options.external_data_loader_use_gds = + internal_options.external_data_loader_use_gds; + cuda_options.external_data_loader_use_directstorage = + internal_options.external_data_loader_use_directstorage; } ProviderOptions GetProviderOptions(const void* provider_options) override { diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index 8ece3b74db06a..44555b8f83ec7 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -2949,6 +2949,15 @@ ORT_API(void, OrtApis::ReleaseTensorRTProviderOptions, _Frees_ptr_opt_ OrtTensor ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider_CUDA_V2, _In_ OrtSessionOptions* options, _In_ const OrtCUDAProviderOptionsV2* cuda_options) { API_IMPL_BEGIN +#if !defined(ORT_MINIMAL_BUILD) + if (cuda_options->external_data_loader_use_directstorage != 0 && + cuda_options->external_data_loader_use_directstorage != 1) { + const auto message = onnxruntime::MakeString( + "external_data_loader_use_directstorage got ", + cuda_options->external_data_loader_use_directstorage, + "; must be 0 or 1."); + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, message.c_str()); + } if (cuda_options->external_data_loader_reading_threads > OrtCUDAProviderOptionsV2::kMaxExternalDataLoaderReadingThreadCount) { const auto message = onnxruntime::MakeString( @@ -2958,6 +2967,17 @@ ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendExecutionProvider_CUDA_V2, _In_ ORT_INVALID_ARGUMENT, message.c_str()); } + if (cuda_options->external_data_loader_use_gds != 0 && + cuda_options->external_data_loader_use_gds != 1) { + const auto message = onnxruntime::MakeString( + "external_data_loader_use_gds got ", + cuda_options->external_data_loader_use_gds, + "; must be 0 or 1."); + return OrtApis::CreateStatus( + ORT_INVALID_ARGUMENT, + message.c_str()); + } +#endif auto factory = onnxruntime::CudaProviderFactoryCreator::Create(cuda_options); if (!factory) { diff --git a/onnxruntime/test/providers/cuda/test_cases/cuda_external_data_loader_gds_test.cc b/onnxruntime/test/providers/cuda/test_cases/cuda_external_data_loader_gds_test.cc new file mode 100644 index 0000000000000..75c9978ec9892 --- /dev/null +++ b/onnxruntime/test/providers/cuda/test_cases/cuda_external_data_loader_gds_test.cc @@ -0,0 +1,137 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include + +#include "core/providers/cuda/cuda_external_data_loader_gds.h" +#include "gtest/gtest.h" + +namespace onnxruntime { +namespace test { +namespace { + +struct DriverState { + std::atomic initialize_attempts{0}; + std::atomic opens{0}; + std::atomic closes{0}; + std::atomic active{0}; + bool fail_initialization{false}; + bool block_first_close{false}; + std::latch close_started{1}; + std::latch allow_close{1}; +}; + +class TestDriver { + public: + explicit TestDriver(DriverState& state) : state_(state) {} + + ~TestDriver() { + if (initialized_) { + if (state_.block_first_close && state_.closes.load() == 0) { + state_.close_started.count_down(); + state_.allow_close.wait(); + } + --state_.active; + ++state_.closes; + } + } + + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(TestDriver); + + Status Initialize() { + ++state_.initialize_attempts; + ORT_RETURN_IF(state_.fail_initialization, "Driver initialization failed"); + ORT_RETURN_IF(state_.active.exchange(1) != 0, "Previous driver is still open"); + ++state_.opens; + initialized_ = true; + return Status::OK(); + } + + private: + DriverState& state_; + bool initialized_{false}; +}; + +using DriverHandle = cuda::GdsDriverHandle; + +TEST(CudaGdsDriverTest, SharesDriverUntilLastHandleIsReleased) { + DriverState state; + { + auto first = std::make_unique(); + auto second = std::make_unique(); + ASSERT_TRUE(first->Acquire(state).IsOK()); + ASSERT_TRUE(first->Acquire(state).IsOK()); + ASSERT_TRUE(second->Acquire(state).IsOK()); + EXPECT_EQ(first->operator->(), second->operator->()); + EXPECT_EQ(state.initialize_attempts.load(), 1U); + + first.reset(); + EXPECT_EQ(state.closes.load(), 0U); + second.reset(); + EXPECT_EQ(state.closes.load(), 1U); + EXPECT_EQ(state.active.load(), 0U); + + DriverHandle next; + ASSERT_TRUE(next.Acquire(state).IsOK()); + EXPECT_EQ(state.opens.load(), 2U); + } + EXPECT_EQ(state.closes.load(), 2U); +} + +TEST(CudaGdsDriverTest, FailedInitializationCanBeRetried) { + DriverState state; + state.fail_initialization = true; + { + DriverHandle handle; + EXPECT_FALSE(handle.Acquire(state).IsOK()); + EXPECT_EQ(state.closes.load(), 0U); + EXPECT_EQ(state.active.load(), 0U); + + state.fail_initialization = false; + ASSERT_TRUE(handle.Acquire(state).IsOK()); + EXPECT_EQ(state.initialize_attempts.load(), 2U); + EXPECT_EQ(state.opens.load(), 1U); + } + EXPECT_EQ(state.closes.load(), 1U); +} + +TEST(CudaGdsDriverTest, ReacquireWaitsForFinalClose) { + DriverState state; + state.block_first_close = true; + auto first = std::make_unique(); + ASSERT_TRUE(first->Acquire(state).IsOK()); + + std::jthread releasing([handle = std::move(first)]() mutable { handle.reset(); }); + state.close_started.wait(); + + std::latch acquire_started{1}; + std::promise acquired; + auto result = acquired.get_future(); + std::jthread acquiring([&]() { + DriverHandle next; + acquire_started.count_down(); + acquired.set_value(next.Acquire(state)); + }); + acquire_started.wait(); + + const auto waiting = result.wait_for(std::chrono::milliseconds(100)); + state.allow_close.count_down(); + releasing.join(); + acquiring.join(); + + EXPECT_EQ(waiting, std::future_status::timeout); + EXPECT_TRUE(result.get().IsOK()); + EXPECT_EQ(state.initialize_attempts.load(), 2U); + EXPECT_EQ(state.opens.load(), 2U); + EXPECT_EQ(state.closes.load(), 2U); + EXPECT_EQ(state.active.load(), 0U); +} + +} // namespace +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/providers/cuda/test_cases/cuda_external_data_loader_test.cc b/onnxruntime/test/providers/cuda/test_cases/cuda_external_data_loader_test.cc index 9c688e2788561..16e8f540a0f35 100644 --- a/onnxruntime/test/providers/cuda/test_cases/cuda_external_data_loader_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/cuda_external_data_loader_test.cc @@ -42,13 +42,15 @@ uint8_t TestValue(size_t index, size_t load = 0) { void CreateExternalDataFile(size_t length, PathString& path, gsl::span suffix = {}, - size_t load = 0) { + size_t load = 0, + size_t prefix_size = kFilePrefixSize) { FILE* file = nullptr; path = ORT_TSTR("cuda_external_data_loader_XXXXXX"); CreateTestFile(file, path); std::vector chunk(1024 * 1024); - ASSERT_EQ(kFilePrefixSize, fwrite(chunk.data(), 1, kFilePrefixSize, file)); + ASSERT_LE(prefix_size, chunk.size()); + ASSERT_EQ(prefix_size, fwrite(chunk.data(), 1, prefix_size, file)); for (size_t offset = 0; offset < length;) { const size_t chunk_size = std::min(chunk.size(), length - offset); for (size_t i = 0; i < chunk_size; ++i) { @@ -63,11 +65,15 @@ void CreateExternalDataFile(size_t length, PathString& path, EXPECT_EQ(0, fclose(file)); } -void VerifyLoad(size_t length, size_t load_count = 1, size_t reading_thread_count = 4) { +void VerifyLoad(size_t length, size_t load_count = 1, size_t reading_thread_count = 4, + bool use_gds = false, size_t prefix_size = kFilePrefixSize, + bool use_directstorage = false) { OrtCUDAProviderOptionsV2 provider_options{}; provider_options.do_copy_in_default_stream = true; provider_options.use_tf32 = false; provider_options.external_data_loader_reading_threads = reading_thread_count; + provider_options.external_data_loader_use_gds = use_gds; + provider_options.external_data_loader_use_directstorage = use_directstorage; auto execution_provider = CudaExecutionProviderWithOptions(&provider_options); ASSERT_NE(execution_provider, nullptr); auto loader = execution_provider->GetExternalDataLoader(); @@ -82,11 +88,11 @@ void VerifyLoad(size_t length, size_t load_count = 1, size_t reading_thread_coun for (size_t load = 0; load < load_count; ++load) { PathString path; - CreateExternalDataFile(length, path, {}, load); + CreateExternalDataFile(length, path, {}, load, prefix_size); ScopedFileDeleter file_deleter{path}; ASSERT_EQ(cudaSuccess, cudaMemset(tensor.MutableDataRaw(), 0xa5, length)); ASSERT_EQ(cudaSuccess, cudaStreamSynchronize(nullptr)); - ASSERT_STATUS_OK(loader->LoadTensor(Env::Default(), path, kFilePrefixSize, length, tensor)); + ASSERT_STATUS_OK(loader->LoadTensor(Env::Default(), path, prefix_size, length, tensor)); std::vector output(length); ASSERT_EQ(cudaSuccess, cudaMemcpy(output.data(), tensor.DataRaw(), length, cudaMemcpyDeviceToHost)); @@ -193,6 +199,17 @@ TEST(CudaExternalDataLoaderTest, DisablesLoaderWhenConfiguredWithZeroReadingThre EXPECT_EQ(execution_provider->GetExternalDataLoader(), nullptr); } +TEST(CudaExternalDataLoaderTest, EnablesLoaderForGdsWithZeroFallbackReadingThreads) { + OrtCUDAProviderOptionsV2 provider_options{}; + provider_options.do_copy_in_default_stream = true; + provider_options.use_tf32 = false; + provider_options.external_data_loader_reading_threads = 0; + provider_options.external_data_loader_use_gds = 1; + auto execution_provider = CudaExecutionProviderWithOptions(&provider_options); + ASSERT_NE(execution_provider, nullptr); + EXPECT_NE(execution_provider->GetExternalDataLoader(), nullptr); +} + TEST(CudaExternalDataLoaderTest, RejectsTooManyReadingThreadsFromStructOptions) { OrtCUDAProviderOptionsV2 provider_options{}; provider_options.external_data_loader_reading_threads = @@ -213,19 +230,110 @@ TEST(CudaExternalDataLoaderTest, RejectsTooManyReadingThreadsFromStructOptions) EXPECT_EQ(CudaExecutionProviderWithOptions(&provider_options), nullptr); } +TEST(CudaExternalDataLoaderTest, RejectsInvalidGdsOptionFromStructOptions) { + OrtCUDAProviderOptionsV2 provider_options{}; + provider_options.external_data_loader_use_gds = 2; + Ort::SessionOptions session_options; + try { + session_options.AppendExecutionProvider_CUDA_V2(provider_options); + FAIL() << "Expected an invalid external_data_loader_use_gds value to be rejected."; + } catch (const Ort::Exception& ex) { + EXPECT_THAT(ex.what(), testing::HasSubstr("external_data_loader_use_gds")); + EXPECT_THAT(ex.what(), testing::HasSubstr("got 2")); + } + + EXPECT_EQ(CudaExecutionProviderWithOptions(&provider_options), nullptr); +} + TEST(CudaExternalDataLoaderTest, ReusesAlternatingBuffersAcrossRepeatedLoads) { VerifyLoad(2 * cuda::kExternalDataLoaderBufferSize + 1, 2); } -cudaError_t FailPinnedBufferAllocation(void**, size_t) { - return cudaErrorMemoryAllocation; +TEST(CudaExternalDataLoaderTest, LoadsAlignedDataWithGdsEnabled) { + for (const size_t reading_thread_count : {0, 1, 4}) { + SCOPED_TRACE(reading_thread_count); + VerifyLoad(cuda::kGdsIoAlignment, 2, reading_thread_count, true, cuda::kGdsIoAlignment); + } +} + +TEST(CudaExternalDataLoaderTest, LoadsMultipleBuffersWithGdsEnabled) { + VerifyLoad(2 * cuda::kExternalDataLoaderBufferSize + cuda::kGdsIoAlignment, + 2, 4, true, cuda::kGdsIoAlignment); +} + +TEST(CudaExternalDataLoaderTest, LoadsUnalignedDataWithGdsEnabled) { + VerifyLoad(cuda::kExternalDataLoaderParallelReadThreshold + 1, 2, 4, true); +} + +TEST(CudaExternalDataLoaderTest, LoadsWithDirectStorageOrConfiguredFallback) { + for (const size_t readers : {0, 1, 4}) { + SCOPED_TRACE(readers); + VerifyLoad(cuda::kGdsIoAlignment, 2, readers, false, cuda::kGdsIoAlignment, true); + VerifyLoad(cuda::kExternalDataLoaderParallelReadThreshold + 1, 2, readers, false, kFilePrefixSize, true); + } +} + +TEST(CudaExternalDataLoaderTest, LoadsMultipleBuffersWithDirectStorageEnabled) { + VerifyLoad(2 * cuda::kExternalDataLoaderBufferSize + 1, 2, 4, false, kFilePrefixSize, true); +} + +TEST(CudaExternalDataLoaderTest, RejectsInvalidDirectStorageOptionFromStructOptions) { + OrtCUDAProviderOptionsV2 provider_options{}; + provider_options.external_data_loader_use_directstorage = 2; + Ort::SessionOptions session_options; + Ort::Status status(Ort::GetApi().SessionOptionsAppendExecutionProvider_CUDA_V2(session_options, &provider_options)); + ASSERT_FALSE(status.IsOK()); + EXPECT_THAT(status.GetErrorMessage(), testing::HasSubstr("external_data_loader_use_directstorage")); + EXPECT_THAT(status.GetErrorMessage(), testing::HasSubstr("got 2")); + EXPECT_EQ(CudaExecutionProviderWithOptions(&provider_options), nullptr); } -cudaError_t FailStreamCreation(cudaStream_t*, unsigned int) { - return cudaErrorInitializationError; +#if defined(ORT_CUDA_DIRECTSTORAGE_AVAILABLE) && !defined(ORT_NO_RTTI) +TEST(CudaExternalDataLoaderTest, NativeDirectStorageWithoutFallback) { + std::unique_ptr loader; + const auto setup_status = cuda::DirectStorageLoader::Create(0, loader); + if (!setup_status.IsOK()) { + GTEST_SKIP() << "DirectStorage/D3D12/CUDA interoperability unavailable: " << setup_status.ErrorMessage(); + } + auto execution_provider = DefaultCudaExecutionProvider(); + ASSERT_NE(execution_provider, nullptr); + auto allocators = execution_provider->CreatePreferredAllocators(); + const auto allocator = std::find_if(allocators.begin(), allocators.end(), [](const AllocatorPtr& candidate) { + return candidate->Info().device.Type() == OrtDevice::GPU && + candidate->Info().mem_type == OrtMemTypeDefault; + }); + ASSERT_NE(allocator, allocators.end()); + constexpr size_t length = cuda::kExternalDataLoaderBufferSize + 17; + Tensor tensor(DataTypeImpl::GetType(), TensorShape({length}), *allocator); + for (size_t load = 0; load < 2; ++load) { + PathString path; + ASSERT_NO_FATAL_FAILURE(CreateExternalDataFile(length, path, {}, load)); + ScopedFileDeleter deleter{path}; + std::unique_ptr file; + ASSERT_STATUS_OK(Env::Default().OpenRandomAccessFile(path.c_str(), file)); + const auto* handles = dynamic_cast(file.get()); + ASSERT_NE(handles, nullptr); + ASSERT_STATUS_OK(loader->Load(path, handles->GetFileHandle(), kFilePrefixSize, length, tensor)); + std::vector output(length); + ASSERT_EQ(cudaSuccess, cudaMemcpy(output.data(), tensor.DataRaw(), length, cudaMemcpyDeviceToHost)); + for (size_t i = 0; i < length; ++i) { + ASSERT_EQ(TestValue(i, load), output[i]) << "byte=" << i << " load=" << load; + } + EXPECT_FALSE(loader->Load(path, handles->GetFileHandle(), -1, length, tensor).IsOK()); + EXPECT_FALSE(loader->Load(path, handles->GetFileHandle(), kFilePrefixSize + 1, length, tensor).IsOK()); + EXPECT_FALSE(loader->Load(path, nullptr, kFilePrefixSize, length, tensor).IsOK()); + + PathString other_path; + ASSERT_NO_FATAL_FAILURE(CreateExternalDataFile(length, other_path, {}, load + 1)); + ScopedFileDeleter other_deleter{other_path}; + const auto replaced = loader->Load(other_path, handles->GetFileHandle(), kFilePrefixSize, length, tensor); + EXPECT_FALSE(replaced.IsOK()); + EXPECT_THAT(replaced.ErrorMessage(), testing::HasSubstr("file changed")); + } } +#endif -TEST(CudaExternalDataLoaderTest, NormalizesBoolWithPinnedAndPageableFallback) { +TEST(CudaExternalDataLoaderTest, NormalizesBoolWithPinnedAndPageableLoading) { const std::array input{0, 1, 2, 255}; const std::array expected{0, 1, 1, 1}; PathString path; @@ -241,22 +349,12 @@ TEST(CudaExternalDataLoaderTest, NormalizesBoolWithPinnedAndPageableFallback) { }); ASSERT_NE(allocator, allocators.end()); - for (int failure_mode = 0; failure_mode < 3; ++failure_mode) { - SCOPED_TRACE(failure_mode); - std::unique_ptr loader; - if (failure_mode == 1) { - loader = std::make_unique( - 0, 4, FailPinnedBufferAllocation); - } else if (failure_mode == 2) { - loader = std::make_unique( - 0, 4, static_cast(cudaMallocHost), - FailStreamCreation); - } else { - loader = std::make_unique(0, 4); - } + for (const size_t reading_thread_count : {0, 1, 4}) { + SCOPED_TRACE(reading_thread_count); + cuda::ExternalDataLoader loader(0, reading_thread_count); Tensor tensor(DataTypeImpl::GetType(), TensorShape({static_cast(input.size())}), *allocator); - ASSERT_STATUS_OK(loader->LoadTensor( + ASSERT_STATUS_OK(loader.LoadTensor( Env::Default(), path, kFilePrefixSize, input.size(), tensor)); std::array output{}; ASSERT_EQ(cudaSuccess, cudaMemcpy( diff --git a/onnxruntime/test/python/transformers/benchmark_cuda_model_loading.py b/onnxruntime/test/python/transformers/benchmark_cuda_model_loading.py index 8b16d7e469368..6c916673edc6d 100644 --- a/onnxruntime/test/python/transformers/benchmark_cuda_model_loading.py +++ b/onnxruntime/test/python/transformers/benchmark_cuda_model_loading.py @@ -3,25 +3,60 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""Measure CUDA InferenceSession creation for a model with external data. +"""Compare pageable, pinned and direct-storage external weights on one CUDA device. -Example: - python benchmark_cuda_model_loading.py \ - --model /path/to/model.onnx \ - --external-data /path/to/model.onnx.data \ - --evict-file-cache +Generate a deterministic, 4-KiB-aligned fixture and benchmark it (no PyTorch needed): + python benchmark_cuda_model_loading.py --generate-model benchmark_data --repetitions 5 -Run this script in a fresh process for every sample. Cache eviction is applied -only to the paths passed on the command line and remains an operating-system -hint rather than a guarantee. +Reuse a model with independently computed reference outputs (NPZ keys are tensor names): + python benchmark_cuda_model_loading.py --model model.onnx --inputs inputs.npz \ + --expected-outputs expected.npz --output results.json + +Every warmup and measured sample runs in a fresh process. Only InferenceSession +construction is timed; imports, input loading, and correctness checks are excluded. +Dependencies are imported lazily so CLI help and reporting helpers need no ORT build. +Initialization completes weight transfers; a blocking inference and NumPy comparison +then verify the result. Loader INFO logs, not requested provider options, determine +the observed path. Direct storage means GDS on Linux or DirectStorage on Windows. +Uncompressed DirectStorage uses host/upload staging with a shared D3D12/CUDA +destination, not zero-host-copy NVMe-to-VRAM DMA. GDS requests native I/O with cuFile +compatibility mode disabled; the physical path remains platform-dependent. +ORT logs prove API use, not the physical storage-to-GPU DMA route. + +Caches remain OS-managed (warmups and fixture creation can warm them). The optional +legacy --evict-file-cache is only a per-file, unprivileged POSIX hint, never proof +of cold storage. No privileged or global cache flushing is performed. """ import argparse import json +import math import os +import platform +import random +import re +import statistics +import subprocess +import sys import time +from pathlib import Path -import onnxruntime +LOADER_LOG = re.compile(r"CUDA external data loader: path=(pageable|pinned|gds|directstorage) bytes=(\d+)\b") +UTF16LE_ASCII_LOG = re.compile(r"(?:[\t\r\n\x20-\x7e]\x00){2,}") +RESULT_PREFIX = "CUDA_LOADING_RESULT=" +PATHS = ("pageable", "pinned", "direct") +PATH_DESCRIPTIONS = { + "pageable": "CPU pageable buffer followed by H2D copy", + "pinned": "CPU pinned staging buffers followed by H2D copy", + "gds": ( + "NVIDIA cuFile API; native GDS requested with compatibility mode disabled, " + "physical storage-to-GPU path remains platform-dependent" + ), + "directstorage": ( + "Microsoft DirectStorage API; uncompressed host/upload staging with a shared D3D12/CUDA destination, " + "not zero-host-copy NVMe-to-VRAM DMA" + ), +} def reading_thread_count(value): @@ -31,81 +66,437 @@ def reading_thread_count(value): return count +def positive_int(value): + count = int(value) + if count <= 0: + raise argparse.ArgumentTypeError("must be positive") + return count + + +def nonnegative_int(value): + count = int(value) + if count < 0: + raise argparse.ArgumentTypeError("must be nonnegative") + return count + + def evict_file_cache(path): if not hasattr(os, "posix_fadvise"): raise RuntimeError("file-cache eviction requires os.posix_fadvise") - with open(path, "rb") as file: os.posix_fadvise(file.fileno(), 0, 0, os.POSIX_FADV_DONTNEED) -def main(): - parser = argparse.ArgumentParser(description="Benchmark CUDA model loading") - parser.add_argument("--model", required=True, help="Path to the ONNX model") - parser.add_argument( - "--external-data", - action="append", - default=[], - help="External-data file to evict with --evict-file-cache; may be specified more than once", - ) - parser.add_argument("--device-id", type=int, default=0, help="CUDA device ID") - parser.add_argument("--threads", type=int, default=96, help="Intra-op thread count") - parser.add_argument( - "--reading-threads", - type=reading_thread_count, - help="Override parallel reads per CUDA pinned staging buffer; 0 disables the loader (runtime default: 4)", - ) - parser.add_argument( - "--evict-file-cache", - action="store_true", - help="Advise the OS to evict the model and external-data files before creating the session", - ) - args = parser.parse_args() +def direct_path(platform_name): + if platform_name.startswith("linux"): + return "gds" + if platform_name == "win32": + return "directstorage" + return None + - paths = [args.model, *args.external_data] +def provider_options(path, device_id, reading_threads, platform_name): + options = { + "device_id": device_id, + "external_data_loader_reading_threads": 0 if path == "pageable" else reading_threads, + "external_data_loader_use_gds": 0, + "external_data_loader_use_directstorage": 0, + "use_tf32": 0, + } + if path == "direct": + api_path = direct_path(platform_name) + if api_path is None: + raise ValueError(f"Direct storage is unsupported on {platform_name}") + options[f"external_data_loader_use_{api_path}"] = 1 + return options + + +def parse_loader_logs(stderr): + """Parse mixed Python text and Windows native UTF-16LE ASCII log segments. + + Only paired ASCII/NUL runs are normalized, not arbitrary NUL characters. + The caller retains the original stderr as evidence in the sample. + """ + stderr = UTF16LE_ASCII_LOG.sub(lambda match: match[0][::2], stderr) + observed = {} + for path, size in LOADER_LOG.findall(stderr): + observed[path] = observed.get(path, 0) + int(size) + warnings = [ + line + for line in stderr.splitlines() + if "CUDA external data loader" in line + and ("fallback" in line.lower() or "falling back" in line.lower() or "[W:" in line) + ] + return observed, warnings + + +def classify_path(expected_path, observed, expected_bytes): + """Require complete byte accounting before claiming the requested path ran.""" + observed = {path: size for path, size in observed.items() if size > 0} + if not observed: + return "unverified" + if sum(observed.values()) != expected_bytes: + return "incomplete" + if len(observed) > 1: + return "mixed" + if expected_path not in observed: + return "fallback" + return "verified" + + +def percentile(values, fraction): + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = math.floor(position) + upper = math.ceil(position) + return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower) + + +def summarize_samples(samples, expected_bytes): + """Keep verified, mixed, and fallback measurements in separate distributions.""" + summaries = [] + groups = {} + for sample in samples: + key = (sample["requested_path"], sample["status"], tuple(sorted(sample["observed_bytes"].items()))) + groups.setdefault(key, []).append(sample) + for (requested, status, observed), group in groups.items(): + seconds = [sample["seconds"] for sample in group] + median = statistics.median(seconds) + summaries.append( + { + "requested_path": requested, + "status": status, + "observed_bytes": dict(observed), + "observed_path_descriptions": {path: PATH_DESCRIPTIONS[path] for path, _ in observed}, + "throughput_definition": ( + "External weight bytes divided by end-to-end InferenceSession initialization time; " + "not storage bandwidth" + ), + "count": len(group), + "seconds": { + "min": min(seconds), + "median": median, + "mean": statistics.mean(seconds), + "stdev": statistics.stdev(seconds) if len(seconds) > 1 else 0.0, + "p90": percentile(seconds, 0.9), + "p95": percentile(seconds, 0.95), + "max": max(seconds), + }, + # End-to-end initialization throughput, not a device or disk bandwidth measurement. + "effective_gib_per_second": { + "median": statistics.median(expected_bytes / (2**30) / value for value in seconds), + "min": expected_bytes / (2**30) / max(seconds), + "max": expected_bytes / (2**30) / min(seconds), + }, + } + ) + return summaries + + +def generate_model(directory, weight_count, dimension, seed): + import numpy as np # noqa: PLC0415 + import onnx # noqa: PLC0415 + + if dimension % 32: + raise ValueError("--weight-dim must be divisible by 32 for 4-KiB-aligned lengths") + directory = Path(directory).resolve() + directory.mkdir(parents=True, exist_ok=True) + paths = [directory / name for name in ("model.onnx", "weights.bin", "inputs.npz", "expected.npz")] for path in paths: - if not os.path.isfile(path): - raise FileNotFoundError(path) + if path.exists(): + raise FileExistsError(f"Refusing to overwrite {path}; use --model to reuse an existing fixture") + rng = np.random.default_rng(seed) + initializers, nodes, outputs, expected = [], [], [], {} + with paths[1].open("wb") as weights_file: + for index in range(weight_count): + name, output = f"weight_{index}", f"output_{index}" + # Binary fractions and a ones input yield stable sums without TF32 ambiguity. + weight = rng.integers(-8, 9, size=(dimension, dimension), dtype=np.int32).astype(" round_trip_options; + for (std::string entry; std::getline(stream, entry, ';');) { + const auto separator = entry.find('='); + ASSERT_NE(separator, std::string::npos); + round_trip_options.emplace(entry.substr(0, separator), entry.substr(separator + 1)); + } + ASSERT_EQ(round_trip_options.at(key), value); + Ort::CUDAProviderOptions restored; + restored.Update(round_trip_options); + EXPECT_EQ((*restored).external_data_loader_use_directstorage, value[0] - '0'); + EXPECT_EQ(restored.GetCUDAProviderOptionsAsString(), serialized); + Ort::SessionOptions session_options; + EXPECT_NO_THROW(session_options.AppendExecutionProvider_CUDA_V2(*options)); + } + options.Update({{key, "1"}}); + for (const char* value : {"2", "-1", "invalid", ""}) { + const char* keys[] = {key}; + const char* values[] = {value}; + Ort::Status status(Ort::GetApi().UpdateCUDAProviderOptions(options, keys, values, 1)); + ASSERT_FALSE(status.IsOK()); + EXPECT_THAT(status.GetErrorMessage(), + testing::HasSubstr(value[0] == '\0' ? "key/value cannot be empty" : key)); + EXPECT_EQ((*options).external_data_loader_use_directstorage, 1); + } +} + +TEST(CApiTest, CUDAProviderOptionsGdsRoundTrip) { + constexpr const char* key = "external_data_loader_use_gds"; + Ort::CUDAProviderOptions cuda_options; + cuda_options.Update({}); + EXPECT_EQ((*cuda_options).external_data_loader_use_gds, 0); + + for (const char* value : {"0", "1", "0"}) { + SCOPED_TRACE(value); + cuda_options.Update({{key, value}}); + EXPECT_EQ((*cuda_options).external_data_loader_use_gds, value[0] - '0'); + + const auto serialized = cuda_options.GetCUDAProviderOptionsAsString(); + std::istringstream stream(serialized); + std::unordered_map round_trip_options; + for (std::string entry; std::getline(stream, entry, ';');) { + const auto separator = entry.find('='); + ASSERT_NE(separator, std::string::npos) << entry; + ASSERT_TRUE(round_trip_options.emplace( + entry.substr(0, separator), entry.substr(separator + 1)) + .second); + } + ASSERT_EQ(round_trip_options.at(key), value); + + Ort::CUDAProviderOptions restored; + restored.Update(round_trip_options); + EXPECT_EQ((*restored).external_data_loader_use_gds, value[0] - '0'); + EXPECT_EQ(restored.GetCUDAProviderOptionsAsString(), serialized); + } +} + +#ifndef ORT_NO_EXCEPTIONS +TEST(CApiTest, CUDAProviderOptionsRejectInvalidGdsValue) { + const char* keys[] = {"external_data_loader_use_gds"}; + for (const char* value : {"2", "-1", "invalid", ""}) { + SCOPED_TRACE(value); + Ort::CUDAProviderOptions cuda_options; + cuda_options.Update({{keys[0], "1"}}); + const char* values[] = {value}; + Ort::Status status(Ort::GetApi().UpdateCUDAProviderOptions(cuda_options, keys, values, 1)); + ASSERT_FALSE(status.IsOK()); + const char* expected_error = value[0] == '\0' ? "key/value cannot be empty" : keys[0]; + EXPECT_THAT(status.GetErrorMessage(), testing::HasSubstr(expected_error)); + EXPECT_EQ((*cuda_options).external_data_loader_use_gds, 1); + } +} +#endif + // This test uses CreateCUDAProviderOptions/UpdateCUDAProviderOptions/UpdateCUDAProviderOptionsWithValue APIs to configure and create a CUDA Execution Provider instance TEST(CApiTest, TestConfigureCUDAProviderOptions) { Ort::CUDAProviderOptions cuda_options;