From dbe7655d8b0be637881a0631060c611366d96987 Mon Sep 17 00:00:00 2001 From: Gopalakrishnan Nallasamy Date: Wed, 8 Jul 2026 15:51:11 -0700 Subject: [PATCH 01/11] Avoid small MatMul batch parameter heap allocations (#29085) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description Reduce per-call heap allocation overhead in the CPU `MatMul` kernel for the small batched-GEMM case. `MatMul::Compute()` builds an array of MLAS batch-parameter structs before calling `MlasGemmBatch` / `MlasSBGemmBatch`. Previously this always used `std::vector`, so even the common single-GEMM path performed one heap allocation per `Compute()` call. This change uses stack-backed `InlinedVector` storage when `helper.OutputOffsets().size() <= 2`, and keeps the existing `std::vector` path for larger batches: - `MLAS_SGEMM_DATA_PARAMS` and `MLAS_SBGEMM_DATA_PARAMS` (aarch64/Linux fast-math path) both use `InlinedVector<..., 2>` for `max_len <= 2`. - Larger batches fall back to `std::vector`, preserving prior behavior (and avoiding the `InlinedVector` overflow path, which I measured to regress for larger batches). The struct is ~64 bytes and trivially copyable, so the net effect is removing exactly **one `malloc`/`free` pair per `Compute()` call** in the small-batch case. ### Motivation and Context Many real CPU MatMul shapes flatten to `helper.OutputOffsets().size() == 1`. `MatMulComputeHelper` special-cases an activation × 2D-weight matrix into a single GEMM (offsets `{0}`), so shapes like the following all hit `max_len == 1`: ```text [M, K] x [K, N] [B, S, K] x [K, N] [B, H, S, K] x [K, N] ``` These appear in transformer projections, MLP layers, and classifier heads. Removing a heap allocation from that path is a small, low-risk latency improvement. ### Performance Measurements Measured on Windows (MSVC, `RelWithDebInfo`) with `onnxruntime_perf_test`, using 256-node **square** MatMul chains (`[N,N]·[N,N]`, which flatten to `max_len == 1` — the case this PR optimizes). 5 trials each; medians shown. `per-node saving` = `(base_p50 − opt_p50) / 256`. | N | base avg (ms) | opt avg (ms) | base p50 (µs) | opt p50 (µs) | per-node saving | |---:|---:|---:|---:|---:|---:| | 1 | 0.178246 | 0.177635 | 159.7 | 153.4 | 24.6 ns | | 8 | 0.189659 | 0.182407 | 167.7 | 158.4 | 36.3 ns | | 16 | 0.209381 | 0.194601 | 178.4 | 168.4 | 39.1 ns | | 32 | 0.307178 | 0.293141 | 279.4 | 269.7 | 37.9 ns | | 64 | 1.389432 | 1.362902 | 1371.2 | 1333.6 | 146.9 ns | | 128 | 3.403018 | 3.145609 | 3305.9 | 3123.3 | 713.3 ns | **Interpretation (honest scope):** - The benefit is a **fixed ~36–39 ns per `Compute()` call**, clearly and consistently visible at N=8/16/32 — exactly one `malloc`/`free` pair removed. - Because the saving is fixed, its true contribution shrinks as the GEMM grows: ~5–6% at N≤16, ~3.5% at N=32, but **well under 1% for N≥64**. - The larger apparent deltas at N=64/128 (147 ns / 713 ns per node) exceed what a fixed allocation can account for and are dominated by threaded-GEMM run-to-run variance (the baseline/optimized ranges overlap), so they should **not** be read as a real ~5% win. In short: this is a targeted micro-optimization that helps **tiny-GEMM-heavy** CPU workloads; for mainstream MatMul sizes the effect is below the measurement noise floor. It is intentionally low-risk rather than a broad throughput win. ### Validation - Built `onnxruntime_perf_test` separately for the baseline (`std::vector`) and optimized (`InlinedVector`) code for an apples-to-apples comparison. - Built `onnxruntime_provider_test` with the final change. - Focused MatMul tests pass: ```text onnxruntime_provider_test.exe --gtest_filter=*MatMul* [==========] 430 tests from 16 test suites ran. [ PASSED ] 430 tests. YOU HAVE 1 DISABLED TEST ``` --------- Co-authored-by: Gopalakrishnan Nallasamy --- onnxruntime/core/providers/cpu/math/matmul.cc | 79 ++++++++++++------- 1 file changed, 51 insertions(+), 28 deletions(-) diff --git a/onnxruntime/core/providers/cpu/math/matmul.cc b/onnxruntime/core/providers/cpu/math/matmul.cc index 6fa3e0d9a4827..3117001eae5b1 100644 --- a/onnxruntime/core/providers/cpu/math/matmul.cc +++ b/onnxruntime/core/providers/cpu/math/matmul.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/providers/cpu/math/matmul.h" +#include "core/common/inlined_containers_fwd.h" #include "core/providers/cpu/math/gemm_matmul_common.h" #include "core/providers/cpu/math/matmul_helper.h" #include "core/util/math.h" @@ -315,6 +316,10 @@ Status MatMul::Compute(OpKernelContext* ctx) const { const size_t K = static_cast(helper.K()); const size_t lda = helper.Lda(trans_a); const size_t ldb = helper.Ldb(trans_b); + // When Abseil is enabled (the default), small batches use stack-backed InlinedVector + // storage to avoid a per-Compute() heap allocation; larger batches use std::vector. + // (Under DISABLE_ABSEIL, InlinedVector is std::vector, so this is a no-op.) + constexpr size_t kInlineBatchCutoff = 2; #if defined(__aarch64__) && defined(__linux__) const bool can_use_fastmath_sbgemm = CanUseFastMathModeSBGemm(N, K); if (packed_b_) { @@ -329,39 +334,57 @@ Status MatMul::Compute(OpKernelContext* ctx) const { } if (can_use_fastmath_sbgemm) { - std::vector data(max_len); - for (size_t i = 0; i < max_len; i++) { - data[i].BIsfp32 = !(bool(packed_b_)); - data[i].AIsfp32 = true; - data[i].A = a_data + helper.LeftOffsets()[i]; - data[i].lda = lda; - data[i].B = data[i].BIsfp32 ? b_data + helper.RightOffsets()[i] : (float*)packed_b_.get(); - data[i].ldb = ldb; - data[i].C = y_data + helper.OutputOffsets()[i]; - data[i].ldc = N; - data[i].Bias = nullptr; - data[i].OutputProcessor = nullptr; - data[i].BIsPacked = static_cast(packed_b_); + auto gemm_batch = [&](auto& data) { + for (size_t i = 0; i < max_len; i++) { + data[i].BIsfp32 = !(bool(packed_b_)); + data[i].AIsfp32 = true; + data[i].A = a_data + helper.LeftOffsets()[i]; + data[i].lda = lda; + data[i].B = data[i].BIsfp32 ? b_data + helper.RightOffsets()[i] : (float*)packed_b_.get(); + data[i].ldb = ldb; + data[i].C = y_data + helper.OutputOffsets()[i]; + data[i].ldc = N; + data[i].Bias = nullptr; + data[i].OutputProcessor = nullptr; + data[i].BIsPacked = static_cast(packed_b_); + } + MlasSBGemmBatch(trans_a ? CblasTrans : CblasNoTrans, trans_b ? CblasTrans : CblasNoTrans, + M, N, K, max_len, data.data(), thread_pool, &mlas_backend_kernel_selector_config_); + }; + + if (max_len <= kInlineBatchCutoff) { + InlinedVector data(max_len); + gemm_batch(data); + } else { + std::vector data(max_len); + gemm_batch(data); } - MlasSBGemmBatch(trans_a ? CblasTrans : CblasNoTrans, trans_b ? CblasTrans : CblasNoTrans, - M, N, K, max_len, data.data(), thread_pool, &mlas_backend_kernel_selector_config_); } else #endif { - std::vector data(max_len); - for (size_t i = 0; i < max_len; i++) { - data[i].BIsPacked = bool(packed_b_); - data[i].A = a_data + helper.LeftOffsets()[i]; - data[i].lda = lda; - data[i].B = data[i].BIsPacked ? (float*)packed_b_.get() : b_data + helper.RightOffsets()[i]; - data[i].ldb = ldb; - data[i].C = y_data + helper.OutputOffsets()[i]; - data[i].ldc = N; - data[i].alpha = alpha_attr_; - data[i].beta = 0.0f; + auto gemm_batch = [&](auto& data) { + for (size_t i = 0; i < max_len; i++) { + data[i].BIsPacked = bool(packed_b_); + data[i].A = a_data + helper.LeftOffsets()[i]; + data[i].lda = lda; + data[i].B = data[i].BIsPacked ? (float*)packed_b_.get() : b_data + helper.RightOffsets()[i]; + data[i].ldb = ldb; + data[i].C = y_data + helper.OutputOffsets()[i]; + data[i].ldc = N; + data[i].alpha = alpha_attr_; + data[i].beta = 0.0f; + } + MlasGemmBatch(trans_a ? CblasTrans : CblasNoTrans, trans_b ? CblasTrans : CblasNoTrans, + M, N, K, data.data(), max_len, thread_pool, &mlas_backend_kernel_selector_config_); + }; + + if (max_len <= kInlineBatchCutoff) { + InlinedVector data(max_len); + gemm_batch(data); + } else { + std::vector data(max_len); + gemm_batch(data); } - MlasGemmBatch(trans_a ? CblasTrans : CblasNoTrans, trans_b ? CblasTrans : CblasNoTrans, - M, N, K, data.data(), max_len, thread_pool, &mlas_backend_kernel_selector_config_); } return Status::OK(); } From 1c34fd5c28308e573680fc4324914d9b9c460169 Mon Sep 17 00:00:00 2001 From: Gopalakrishnan Nallasamy Date: Wed, 8 Jul 2026 16:25:07 -0700 Subject: [PATCH 02/11] Add IOBinding test for preallocated CPU output buffer reuse (#29566) ### Description Adds a CPU IOBinding test that makes the no-copy/in-place contract for preallocated outputs explicit. The new test: - builds the existing in-memory MatMul model used by `io_binding_test.cc`; - binds CPU inputs; - allocates a CPU output `OrtValue` and records its raw tensor buffer pointer; - binds that `OrtValue` as output `Y`; - runs the session with `IOBinding`; - verifies the returned output still points to the exact bound buffer and validates the numeric result. This guards the intended behavior that a preallocated CPU output is written in-place rather than ORT allocating a separate output buffer and copying into/returning that. ### Motivation and Context Preallocated I/O binding is one of the local zero-copy/copy-elision paths that can be validated without a GPU or multi-EP setup. Existing tests checked values; this adds a direct buffer-identity assertion so regressions that accidentally allocate/copy to a different buffer are caught. ### Testing Built `onnxruntime_provider_test` (Release, Windows) and ran: ```powershell .\onnxruntime_provider_test.exe --gtest_filter="InferenceSessionTests.TestBindCpu:InferenceSessionTests.TestBindCpuPreallocatedOutputUsesBoundBuffer" ``` Result: ```text [==========] 2 tests from 1 test suite ran. [ PASSED ] 2 tests. ``` Co-authored-by: Gopalakrishnan Nallasamy --- onnxruntime/test/providers/io_binding_test.cc | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/onnxruntime/test/providers/io_binding_test.cc b/onnxruntime/test/providers/io_binding_test.cc index 12dce160ea8c9..44c2faf231007 100644 --- a/onnxruntime/test/providers/io_binding_test.cc +++ b/onnxruntime/test/providers/io_binding_test.cc @@ -284,6 +284,56 @@ TEST(InferenceSessionTests, TestBindCpu) { false /* don't preallocate output */); } +TEST(InferenceSessionTests, TestBindCpuPreallocatedOutputUsesBoundBuffer) { + SessionOptions so; + so.session_logid = "InferenceSessionTests.TestBindCpuPreallocatedOutputUsesBoundBuffer"; + InferenceSession session_object{so, GetEnvironment()}; + + std::unique_ptr p_model; + CreateMatMulModel(p_model, kCpuExecutionProvider); + + std::string serialized_model; + p_model->ToProto().SerializeToString(&serialized_model); + std::stringstream model_stream(serialized_model); + ASSERT_STATUS_OK(session_object.Load(model_stream)); + ASSERT_STATUS_OK(session_object.Initialize()); + + std::unique_ptr io_binding; + ASSERT_STATUS_OK(session_object.NewIOBinding(&io_binding)); + + auto cpu_alloc = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; + const std::vector dims_mul_x_a = {3, 4}; + const std::vector dims_mul_x_b = {4, 3}; + const std::vector values_mul_x = {0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, + 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 11.0f}; + + OrtValue input_ml_value_a; + CreateMLValue(cpu_alloc, dims_mul_x_a, values_mul_x, &input_ml_value_a); + ASSERT_STATUS_OK(io_binding->BindInput("A", input_ml_value_a)); + + OrtValue input_ml_value_b; + CreateMLValue(cpu_alloc, dims_mul_x_b, values_mul_x, &input_ml_value_b); + ASSERT_STATUS_OK(io_binding->BindInput("B", input_ml_value_b)); + + const std::vector expected_output_dims = {3, 3}; + OrtValue output_ml_value; + AllocateMLValue(cpu_alloc, expected_output_dims, &output_ml_value); + const void* const bound_output_buffer = output_ml_value.Get().DataRaw(); + ASSERT_STATUS_OK(io_binding->BindOutput("Y", output_ml_value)); + + RunOptions run_options; + run_options.run_tag = so.session_logid; + ASSERT_STATUS_OK(session_object.Run(run_options, *io_binding)); + + const auto& outputs = io_binding->GetOutputs(); + ASSERT_EQ(outputs.size(), 1U); + EXPECT_EQ(outputs[0].Get().DataRaw(), bound_output_buffer) + << "Preallocated CPU output should be written in-place without allocating or copying to a new buffer."; + + const std::vector expected_values_mul_y = {42, 48, 54, 114, 136, 158, 186, 224, 262}; + VerifySingleOutput(outputs, expected_output_dims, expected_values_mul_y); +} + TEST(InferenceSessionTests, TestIOBindingReuse) { SessionOptions so; InferenceSession session_object(so, GetEnvironment()); From c57e0e50ad068905a4140d361b0b3fd8c251e540 Mon Sep 17 00:00:00 2001 From: Edward Chen <18449977+edgchen1@users.noreply.github.com> Date: Wed, 8 Jul 2026 16:26:34 -0700 Subject: [PATCH 03/11] Avoid throwing C++ exceptions from WebGPU EP Dawn callbacks (#29625) ### Description Mark every Dawn callback (RequestAdapter, RequestDevice, SetUncapturedErrorCallback, SetDeviceLostCallback, MapAsync, PopErrorScope, CreateComputePipelineAsync) noexcept and avoid throwing from inside them. A C++ exception unwinding out of a Dawn callback runs while Dawn holds internal locks and is undefined behavior, and can leave the EventManager mutex locked so a later instance release self-deadlocks. The MapAsync callbacks in BufferManager::Download and WebGpuContext::CollectProfilingData now write the status and message into a MapAsyncResult struct and the outcome is checked after Wait returns (ORT_THROW_IF_ERROR + ORT_ENFORCE). PopErrorScope builds and returns an onnxruntime::Status for both a failed pop and a validation error instead of enforcing inside the callback. ### Motivation and Context Fix crashes caused by exceptions crossing the WebGPU callback boundary. Follow up to #29591 with similar changes for other callbacks. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../core/providers/webgpu/buffer_manager.cc | 25 +++++- .../core/providers/webgpu/program_manager.cc | 7 +- .../core/providers/webgpu/webgpu_context.cc | 81 ++++++++++++------- 3 files changed, 81 insertions(+), 32 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.cc b/onnxruntime/core/providers/webgpu/buffer_manager.cc index 462fd002da800..2467d8e7d81df 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.cc +++ b/onnxruntime/core/providers/webgpu/buffer_manager.cc @@ -599,9 +599,28 @@ void BufferManager::Download(WGPUBuffer src, void* dst, size_t size) const { // TODO: revise wait in whole project - ORT_ENFORCE(context_.Wait(staging_buffer.MapAsync(wgpu::MapMode::Read, 0, buffer_size, wgpu::CallbackMode::WaitAnyOnly, [](wgpu::MapAsyncStatus status, wgpu::StringView message) { - ORT_ENFORCE(status == wgpu::MapAsyncStatus::Success, "Failed to download data from buffer: ", std::string_view{message}); - })) == Status::OK()); + struct MapAsyncResult { + wgpu::MapAsyncStatus status{}; + std::string message{}; + } map_async_result; + + ORT_THROW_IF_ERROR(context_.Wait( + staging_buffer.MapAsync( + wgpu::MapMode::Read, 0, buffer_size, wgpu::CallbackMode::WaitAnyOnly, + // Note: Don't throw from a Dawn callback. + [](wgpu::MapAsyncStatus status, wgpu::StringView message, + MapAsyncResult* result) noexcept { + result->status = status; + if (auto message_sv = static_cast(message); + !message_sv.empty()) { + result->message = std::string{message_sv}; + } + }, + &map_async_result))); + + ORT_ENFORCE(map_async_result.status == wgpu::MapAsyncStatus::Success, + "Failed to download data from buffer. wgpu::MapAsyncStatus value: ", + static_cast(map_async_result.status), ", message: ", map_async_result.message); auto mapped_data = staging_buffer.GetConstMappedRange(); memcpy(dst, mapped_data, size); diff --git a/onnxruntime/core/providers/webgpu/program_manager.cc b/onnxruntime/core/providers/webgpu/program_manager.cc index 136e7d503f59f..71cab39f239f8 100644 --- a/onnxruntime/core/providers/webgpu/program_manager.cc +++ b/onnxruntime/core/providers/webgpu/program_manager.cc @@ -224,11 +224,14 @@ Status ProgramManager::Build(const ProgramBase& program, device.CreateComputePipelineAsync( &pipeline_descriptor, wgpu::CallbackMode::WaitAnyOnly, - [](wgpu::CreatePipelineAsyncStatus status, wgpu::ComputePipeline pipeline, wgpu::StringView message, CreateComputePipelineContext* context) { + // Note: Don't throw from a Dawn callback. + [](wgpu::CreatePipelineAsyncStatus status, wgpu::ComputePipeline pipeline, wgpu::StringView message, + CreateComputePipelineContext* context) noexcept { if (status == wgpu::CreatePipelineAsyncStatus::Success) { context->pipeline = std::move(pipeline); } else { - context->status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to create a WebGPU compute pipeline: ", std::string_view{message}); + context->status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to create a WebGPU compute pipeline: ", + std::string_view{message}); } }, &create_pipeline_context))); diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 141c4e19be078..f8485c5f6ac1f 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -73,7 +73,7 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { &req_adapter_options, wgpu::CallbackMode::WaitAnyOnly, [](wgpu::RequestAdapterStatus status, wgpu::Adapter adapter, wgpu::StringView message, - RequestAdapterResult* result) { + RequestAdapterResult* result) noexcept { result->status = status; if (status == wgpu::RequestAdapterStatus::Success) { result->adapter = std::move(adapter); @@ -113,20 +113,24 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { device_desc.requiredLimits = &required_limits; // TODO: revise temporary error handling - device_desc.SetUncapturedErrorCallback([](const wgpu::Device& /*device*/, wgpu::ErrorType type, wgpu::StringView message) { - if (logging::LoggingManager::HasDefaultLogger()) { - LOGS_DEFAULT(ERROR) << "WebGPU device error(" << int(type) << "): " << std::string_view{message}; - } - }); + device_desc.SetUncapturedErrorCallback( + // Note: Don't throw from a Dawn callback. + [](const wgpu::Device& /*device*/, wgpu::ErrorType type, + wgpu::StringView message) noexcept { + if (logging::LoggingManager::HasDefaultLogger()) { + LOGS_DEFAULT(ERROR) << "WebGPU device error(" << int(type) << "): " << std::string_view{message}; + } + }); // TODO: revise temporary device lost handling - device_desc.SetDeviceLostCallback(wgpu::CallbackMode::AllowSpontaneous, [](const wgpu::Device& /*device*/, wgpu::DeviceLostReason reason, wgpu::StringView message) { - if (logging::LoggingManager::HasDefaultLogger()) { - LOGS_DEFAULT(INFO) << "WebGPU device lost (" << int(reason) << "): " << std::string_view{message}; - } - }); + device_desc.SetDeviceLostCallback( + wgpu::CallbackMode::AllowSpontaneous, + // Note: Don't throw from a Dawn callback. + [](const wgpu::Device& /*device*/, wgpu::DeviceLostReason reason, wgpu::StringView message) noexcept { + if (logging::LoggingManager::HasDefaultLogger()) { + LOGS_DEFAULT(INFO) << "WebGPU device lost (" << int(reason) << "): " << std::string_view{message}; + } + }); - // Capture device request result without throwing inside the Dawn callback (same - // reasoning as the adapter callback above). struct RequestDeviceResult { wgpu::RequestDeviceStatus status = wgpu::RequestDeviceStatus::Error; wgpu::Device device; @@ -136,8 +140,11 @@ void WebGpuContext::Initialize(const WebGpuContextConfig& config) { ORT_ENFORCE(wgpu::WaitStatus::Success == instance_.WaitAny(adapter.RequestDevice( &device_desc, wgpu::CallbackMode::WaitAnyOnly, - [](wgpu::RequestDeviceStatus status, wgpu::Device device, wgpu::StringView message, - RequestDeviceResult* result) { + // Note: Don't throw from a Dawn callback. + [](wgpu::RequestDeviceStatus status, + wgpu::Device device, + wgpu::StringView message, + RequestDeviceResult* result) noexcept { result->status = status; if (status == wgpu::RequestDeviceStatus::Success) { result->device = std::move(device); @@ -713,13 +720,30 @@ void WebGpuContext::CollectProfilingData(profiling::Events& events) { const auto& pending_kernels = pending_query.kernels; const auto& query_read_buffer = pending_query.query_buffer; - ORT_ENFORCE(Wait(query_read_buffer.MapAsync(wgpu::MapMode::Read, - 0, - static_cast(query_read_buffer.GetSize()), - wgpu::CallbackMode::WaitAnyOnly, - [](wgpu::MapAsyncStatus status, wgpu::StringView message) { - ORT_ENFORCE(status == wgpu::MapAsyncStatus::Success, "Failed to download data from buffer: ", std::string_view{message}); - })) == Status::OK()); + struct MapAsyncResult { + wgpu::MapAsyncStatus status{}; + std::string message{}; + } map_async_result; + + ORT_THROW_IF_ERROR(Wait(query_read_buffer.MapAsync( + wgpu::MapMode::Read, + 0, + static_cast(query_read_buffer.GetSize()), + wgpu::CallbackMode::WaitAnyOnly, + // Note: Don't throw from a Dawn callback. + [](wgpu::MapAsyncStatus status, wgpu::StringView message, MapAsyncResult* result) noexcept { + result->status = status; + if (auto message_sv = static_cast(message); + !message_sv.empty()) { + result->message = std::string{message_sv}; + } + }, + &map_async_result))); + + ORT_ENFORCE(map_async_result.status == wgpu::MapAsyncStatus::Success, + "Failed to download data from buffer. wgpu::MapAsyncStatus value: ", + static_cast(map_async_result.status), ", message: ", map_async_result.message); + auto mapped_data = static_cast(query_read_buffer.GetConstMappedRange()); for (size_t i = 0; i < pending_kernels.size(); i++) { @@ -793,12 +817,15 @@ Status WebGpuContext::PopErrorScope() { Status status{}; ORT_RETURN_IF_ERROR(Wait(device_.PopErrorScope( wgpu::CallbackMode::WaitAnyOnly, - [](wgpu::PopErrorScopeStatus pop_status, wgpu::ErrorType error_type, char const* message, Status* status) { - ORT_ENFORCE(pop_status == wgpu::PopErrorScopeStatus::Success, "Instance dropped."); - if (error_type == wgpu::ErrorType::NoError) { - return; + // Note: Don't throw from a Dawn callback. + [](wgpu::PopErrorScopeStatus pop_status, wgpu::ErrorType error_type, char const* message, + Status* status) noexcept { + if (pop_status != wgpu::PopErrorScopeStatus::Success) { + *status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to pop WebGPU error scope. status=", + static_cast(pop_status)); + } else if (error_type != wgpu::ErrorType::NoError) { + *status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "WebGPU validation failed. ", message); } - *status = ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "WebGPU validation failed. ", message); }, &status))); return status; From f631f0d93bf94cb3031953375319dea09196a232 Mon Sep 17 00:00:00 2001 From: Gopalakrishnan Nallasamy Date: Wed, 8 Jul 2026 18:07:31 -0700 Subject: [PATCH 04/11] Guard GridSample CUDA coordinate conversion (#29581) ### Description Guards the CUDA GridSample coordinate conversion paths for non-finite and extreme grid values, matching the existing CPU behavior for float grids. The CUDA path now sanitizes coordinates before integer conversion, uses wider intermediate indices where needed, and clamps reflected indices before sampling. Per review feedback, `GsReflect` performs the `isfinite` check before computing the reflection range (`x_max - x_min`), so a non-finite coordinate returns early without the extra subtraction. The same reorder is applied to the CPU `GsReflect` (`core/providers/cpu/tensor/grid_sample.cc`) to keep the CPU and CUDA implementations in lockstep; behavior is unchanged. ### Tests - `.\.venv\Scripts\python.exe tools\ci_build\build.py --config RelWithDebInfo --build --parallel --target onnxruntime_provider_test --build_dir build\Windows` - `.\build\Windows\RelWithDebInfo\RelWithDebInfo\onnxruntime_provider_test.exe --gtest_filter=*Grid*` - 180 tests passed - `.\.venv\Scripts\clang-format.exe` on touched C++ files Note: the local build is CPU-only (`onnxruntime_USE_CUDA=OFF`), so CUDA compilation/runtime coverage will come from CUDA CI. ### Test coverage The hardening is applied at shared choke-points: coordinate sanitization runs before interpolation-mode dispatch, and the `int64_t` index widening and reflected-index clamps are single shared branches. So one representative case per {mode, padding, dimensionality} exercises the hardened path rather than the full cross-product. Regression tests use constant-valued images, so the expected output is well-defined regardless of which (now sanitized/clamped) index each adversarial coordinate resolves to; this also sidesteps a pre-existing CPU-vs-CUDA reflected-index difference (double-reflect + round-half-to-even). The new `GridSampleCudaHardeningTest` cases run on CPU and CUDA (and CUDA-NHWC when `ENABLE_CUDA_NHWC_OPS` is enabled); the CUDA kernel is registered for `float` only, so `double` is exercised through the shared templated CPU path. --------- Co-authored-by: Gopalakrishnan Nallasamy --- .../core/providers/cpu/tensor/grid_sample.cc | 12 +- .../providers/cuda/tensor/grid_sample_impl.cu | 91 +++++-- .../cpu/tensor/grid_sample_test_custom.cc | 238 ++++++++++++++++++ 3 files changed, 320 insertions(+), 21 deletions(-) diff --git a/onnxruntime/core/providers/cpu/tensor/grid_sample.cc b/onnxruntime/core/providers/cpu/tensor/grid_sample.cc index 34135ca62e4a8..e80ecd8fee7fe 100644 --- a/onnxruntime/core/providers/cpu/tensor/grid_sample.cc +++ b/onnxruntime/core/providers/cpu/tensor/grid_sample.cc @@ -64,10 +64,16 @@ template T GsReflect(T x, T x_min, T x_max) { T dx = {}; T fx = static_cast(x); + // Guard against NaN or Inf first, before computing the range, so a non-finite coordinate + // returns early without an unnecessary subtraction and can never reach the float->int cast + // (undefined behavior) below. + if (!std::isfinite(fx)) { + return x_min; + } T range = x_max - x_min; - // Guard against NaN, Inf, or non-positive range (e.g. dim==1 with align_corners=true) - // which would otherwise produce wild indices via division by zero or UB float->int casts. - if (!std::isfinite(fx) || !(range > T{0})) { + // Guard against a non-positive range (e.g. dim==1 with align_corners=true) which would + // otherwise produce wild indices via division by zero. + if (!(range > T{0})) { return x_min; } if (fx < x_min) { diff --git a/onnxruntime/core/providers/cuda/tensor/grid_sample_impl.cu b/onnxruntime/core/providers/cuda/tensor/grid_sample_impl.cu index 0e7d947741924..51258e6e8f9cd 100755 --- a/onnxruntime/core/providers/cuda/tensor/grid_sample_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/grid_sample_impl.cu @@ -25,10 +25,23 @@ template __device__ T GsReflect(T x, float x_min, float x_max) { float fx = static_cast(x); float dx = {}; + // Guard against NaN or Inf first, before computing the range, so a non-finite coordinate + // returns early without an unnecessary subtraction and can never reach the float->int cast + // (undefined behavior) below. + if (!isfinite(fx)) { + return static_cast(x_min); + } float range = x_max - x_min; + // Guard against a non-positive range (e.g. dim==1 with align_corners=true) which would + // otherwise produce wild indices via division by zero. + if (!(range > 0.0f)) { + return static_cast(x_min); + } if (fx < x_min) { dx = x_min - fx; - int n = static_cast(dx / range); + // Use int64_t rather than int: for extreme (but finite) inputs, dx / range can exceed + // INT_MAX, making a float->int cast undefined behavior. + int64_t n = static_cast(dx / range); float r = dx - n * range; if (n % 2 == 0) { fx = x_min + r; @@ -37,7 +50,7 @@ __device__ T GsReflect(T x, float x_min, float x_max) { } } else if (fx > x_max) { dx = fx - x_max; - int n = static_cast(dx / range); + int64_t n = static_cast(dx / range); float r = dx - n * range; if (n % 2 == 0) { fx = x_max - r; @@ -49,6 +62,15 @@ __device__ T GsReflect(T x, float x_min, float x_max) { return static_cast(fx); } +// Returns true when v is finite and its magnitude is small enough that converting it to +// int64_t via floor / nearbyint is well-defined. 2^62 is far below INT64_MAX (~9.22e18) +// and leaves ample margin for any realistic image dimension. +template +__device__ inline bool IsSafeForInt64Conversion(T v) { + constexpr T kSafeBound = static_cast(int64_t{1} << 62); + return isfinite(v) && v >= -kSafeBound && v <= kSafeBound; +} + template __device__ T PixelAtGrid(const T* input_data, int64_t bIdx, int64_t cIdx, int64_t y, int64_t x, int64_t padding_mode, int64_t N, int64_t C, int64_t H, int64_t W, float border[4]) { @@ -71,6 +93,11 @@ __device__ T PixelAtGrid(const T* input_data, int64_t bIdx, int64_t cIdx, int64_ } else { // Reflection x = (int64_t)GsReflect(x, border[0], border[2]); y = (int64_t)GsReflect(y, border[1], border[3]); + // Safety clamp: GsReflect is computed in floating point and cast back to int64_t. + // Extreme grid coordinates can overflow that cast, so clamp the resulting indices + // back into the image range before indexing. + x = max((int64_t)0, min((int64_t)W - 1, x)); + y = max((int64_t)0, min((int64_t)H - 1, y)); pixel = input_data[PixelOffset(x, y)]; } return pixel; @@ -168,6 +195,16 @@ __global__ void _GridSampleKernel( y_max = float(H_in) - 1.0f; } float border[] = {x_min, y_min, x_max, y_max}; // l-t-r-b + + // Sanitize coordinates that are non-finite or whose magnitude is too large for a safe + // float->int64 conversion. Substituting the in-range lower border keeps the subsequent + // floor / nearbyint casts well-defined for every padding mode. + if (!IsSafeForInt64Conversion(grid_x_imgSpace)) { + grid_x_imgSpace = x_min; + } + if (!IsSafeForInt64Conversion(grid_y_imgSpace)) { + grid_y_imgSpace = y_min; + } if (grid_x_imgSpace < x_min || grid_x_imgSpace > x_max || grid_y_imgSpace < y_min || grid_y_imgSpace > y_max) { // out of bound if (padding_mode == 1) { // border @@ -181,10 +218,10 @@ __global__ void _GridSampleKernel( } if (mode == 0) { // bilinear - int x1 = floor(grid_x_imgSpace); - int y1 = floor(grid_y_imgSpace); - int x2 = x1 + 1; - int y2 = y1 + 1; + int64_t x1 = static_cast(floor(grid_x_imgSpace)); + int64_t y1 = static_cast(floor(grid_y_imgSpace)); + int64_t x2 = x1 + 1; + int64_t y2 = y1 + 1; T w_lt = 0.0f; T w_rt = 0.0f; T w_lb = 0.0f; @@ -209,8 +246,8 @@ __global__ void _GridSampleKernel( return; } if (mode == 1) { // nearest - int x_n = grid_x_imgSpace; - int y_n = grid_y_imgSpace; + int64_t x_n = static_cast(grid_x_imgSpace); + int64_t y_n = static_cast(grid_y_imgSpace); output_data[outIdx] = PixelAtGrid(input_data, BIdx, cIdx, y_n, x_n, padding_mode, N, C, H_in, W_in, border); return; @@ -286,7 +323,12 @@ __device__ T PixelAtGrid3D(const T* input_data, int64_t bIdx, int64_t cIdx, int6 z = (int64_t)GsReflect(z, border[0], border[3]); y = (int64_t)GsReflect(y, border[1], border[4]); x = (int64_t)GsReflect(x, border[2], border[5]); - + // Safety clamp: GsReflect is computed in floating point and cast back to int64_t. + // Extreme grid coordinates can overflow that cast, so clamp the resulting indices + // back into the volume range before indexing. + z = max((int64_t)0, min((int64_t)D - 1, z)); + y = max((int64_t)0, min((int64_t)H - 1, y)); + x = max((int64_t)0, min((int64_t)W - 1, x)); pixel = input_data[PixelOffset3D(z, y, x)]; } return pixel; @@ -381,6 +423,19 @@ __global__ void _GridSampleKernel3D( } float border[] = {z_min, y_min, x_min, z_max, y_max, x_max}; // zmin,ymin,xmin,zmax,ymax,xmax + + // Sanitize coordinates that are non-finite or whose magnitude is too large for a safe + // float->int64 conversion. Substituting the in-range lower border keeps the subsequent + // floor / nearbyint casts well-defined for every padding mode. + if (!IsSafeForInt64Conversion(grid_x_volSpace)) { + grid_x_volSpace = x_min; + } + if (!IsSafeForInt64Conversion(grid_y_volSpace)) { + grid_y_volSpace = y_min; + } + if (!IsSafeForInt64Conversion(grid_z_volSpace)) { + grid_z_volSpace = z_min; + } if (grid_z_volSpace < z_min || grid_z_volSpace > z_max || grid_y_volSpace < y_min || grid_y_volSpace > y_max || grid_x_volSpace < x_min || grid_x_volSpace > x_max) { // out of bound @@ -397,12 +452,12 @@ __global__ void _GridSampleKernel3D( } if (mode == 0) { // bilinear - int z1 = floor(grid_z_volSpace); - int y1 = floor(grid_y_volSpace); - int x1 = floor(grid_x_volSpace); - int z2 = z1 + 1; - int y2 = y1 + 1; - int x2 = x1 + 1; + int64_t z1 = static_cast(floor(grid_z_volSpace)); + int64_t y1 = static_cast(floor(grid_y_volSpace)); + int64_t x1 = static_cast(floor(grid_x_volSpace)); + int64_t z2 = z1 + 1; + int64_t y2 = y1 + 1; + int64_t x2 = x1 + 1; // Weights T w_lt_front = 0.0f; @@ -450,9 +505,9 @@ __global__ void _GridSampleKernel3D( return; } if (mode == 1) { // nearest - int x_n = grid_x_volSpace; - int y_n = grid_y_volSpace; - int z_n = grid_z_volSpace; + int64_t x_n = static_cast(grid_x_volSpace); + int64_t y_n = static_cast(grid_y_volSpace); + int64_t z_n = static_cast(grid_z_volSpace); output_data[outIdx] = PixelAtGrid3D(input_data, BIdx, cIdx, z_n, y_n, x_n, padding_mode, N, C, D_in, H_in, W_in, border); diff --git a/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc b/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc index dd23b76d8275d..fe4f551046474 100644 --- a/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc +++ b/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc @@ -63,6 +63,30 @@ void RunCpuOnly(T& test) { RunTests(test, GetCpuExecutionProviders()); } +// Execution providers targeted by the float->int64 cast hardening in the CUDA GridSample +// kernels (onnxruntime/core/providers/cuda/tensor/grid_sample_impl.cu). CoreML and WebGPU +// are intentionally excluded here: their integer-conversion semantics differ from the +// CPU/CUDA "substitute the in-range border" behavior on NaN/Inf/extreme coordinates, so +// their outputs for these adversarial inputs are not expected to match. +std::vector> GetCpuAndCudaExecutionProviders() { + std::vector> execution_providers; + execution_providers.emplace_back(DefaultCpuExecutionProvider()); + +#ifdef USE_CUDA + execution_providers.emplace_back(DefaultCudaExecutionProvider()); +#ifdef ENABLE_CUDA_NHWC_OPS + execution_providers.push_back(DefaultCudaNHWCExecutionProvider()); +#endif +#endif + + return execution_providers; +} + +template +void RunCpuAndCuda(T& test) { + RunTests(test, GetCpuAndCudaExecutionProviders()); +} + } // namespace template @@ -578,5 +602,219 @@ TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_reflection_mixed RunCpuOnly(test); } +// ----------------------------------------------------------------------------------------- +// CUDA hardening regression tests. +// +// These mirror the CPU-only cases above but also run on the CUDA execution provider (when +// built) to cover the equivalent float->int64 cast hardening in +// onnxruntime/core/providers/cuda/tensor/grid_sample_impl.cu. They are float-only because +// the CUDA GridSample kernel is registered for float only, and they use constant-valued +// images so the expected output is well-defined regardless of which in-range pixel each +// (now clamped) index resolves to. +// ----------------------------------------------------------------------------------------- + +TEST(GridSampleCudaHardeningTest, nearest_reflection_extreme_and_nan_coords) { + // Exercises the CUDA _GridSampleKernel reflection path: GsReflect now guards NaN/Inf and + // uses int64_t for the wrap count (1e10 exceeds INT_MAX), and PixelAtGrid clamps the + // reflected index back into range. + OpTester test("GridSample", 20); + test.AddAttribute("mode", std::string("nearest")); + test.AddAttribute("padding_mode", std::string("reflection")); + test.AddAttribute("align_corners", int64_t{0}); + + std::initializer_list X_shape{1, 1, 2, 2}; + std::initializer_list X_data{1.0f, 1.0f, 1.0f, 1.0f}; + + std::initializer_list Grid_shape{1, 1, 2, 2}; + std::initializer_list Grid_data{ + 1.0e+10f, 1.0e+10f, + std::numeric_limits::quiet_NaN(), std::numeric_limits::quiet_NaN()}; + + std::initializer_list Y_shape{1, 1, 1, 2}; + + test.AddInput("X", X_shape, X_data); + test.AddInput("Grid", Grid_shape, Grid_data); + test.AddOutput("Y", Y_shape, {1.0f, 1.0f}); + RunCpuAndCuda(test); +} + +TEST(GridSampleCudaHardeningTest, bilinear_reflection_infinity_coords) { + // CUDA bilinear reflection path: the GsReflect finite guard and coordinate sanitization + // keep the floor(...) cast well-defined for +/-Inf inputs. + OpTester test("GridSample", 20); + test.AddAttribute("mode", std::string("linear")); + test.AddAttribute("padding_mode", std::string("reflection")); + test.AddAttribute("align_corners", int64_t{0}); + + std::initializer_list X_shape{1, 1, 2, 2}; + std::initializer_list X_data{1.0f, 1.0f, 1.0f, 1.0f}; + + std::initializer_list Grid_shape{1, 1, 2, 2}; + const float inf_val = std::numeric_limits::infinity(); + std::initializer_list Grid_data{ + inf_val, -inf_val, + -inf_val, inf_val}; + + std::initializer_list Y_shape{1, 1, 1, 2}; + + test.AddInput("X", X_shape, X_data); + test.AddInput("Grid", Grid_shape, Grid_data); + test.AddOutput("Y", Y_shape, {1.0f, 1.0f}); + RunCpuAndCuda(test); +} + +TEST(GridSampleCudaHardeningTest, bilinear_border_nan_inf_extreme_coords) { + // CUDA bilinear border path: coordinate sanitization keeps the floor(...) cast + // well-defined before PixelAtGrid clamps border-padding samples into range. + OpTester test("GridSample", 20); + test.AddAttribute("mode", std::string("linear")); + test.AddAttribute("padding_mode", std::string("border")); + test.AddAttribute("align_corners", int64_t{0}); + + std::initializer_list X_shape{1, 1, 2, 2}; + std::initializer_list X_data{1.0f, 1.0f, 1.0f, 1.0f}; + + std::initializer_list Grid_shape{1, 1, 3, 2}; + std::initializer_list Grid_data{ + std::numeric_limits::quiet_NaN(), std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), -std::numeric_limits::infinity(), + 1.0e+20f, -1.0e+20f}; + + std::initializer_list Y_shape{1, 1, 1, 3}; + + test.AddInput("X", X_shape, X_data); + test.AddInput("Grid", Grid_shape, Grid_data); + test.AddOutput("Y", Y_shape, {1.0f, 1.0f, 1.0f}); + RunCpuAndCuda(test); +} + +TEST(GridSampleCudaHardeningTest, bilinear_zeros_nan_inf_extreme_coords) { + // Zeros padding: the CUDA kernel's IsSafeForInt64Conversion sanitization substitutes the + // lower border (-0.5) for NaN/Inf/extreme coordinates before the floor cast. Only the + // bottom-right neighbor (pixel(0, 0) = 1.0) is in bounds with weight 0.5 * 0.5 = 0.25. + OpTester test("GridSample", 20); + test.AddAttribute("mode", std::string("linear")); + test.AddAttribute("padding_mode", std::string("zeros")); + test.AddAttribute("align_corners", int64_t{0}); + + std::initializer_list X_shape{1, 1, 2, 2}; + std::initializer_list X_data{1.0f, 1.0f, 1.0f, 1.0f}; + + std::initializer_list Grid_shape{1, 1, 3, 2}; + std::initializer_list Grid_data{ + std::numeric_limits::quiet_NaN(), std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), -std::numeric_limits::infinity(), + 1.0e+20f, -1.0e+20f}; + + std::initializer_list Y_shape{1, 1, 1, 3}; + + test.AddInput("X", X_shape, X_data); + test.AddInput("Grid", Grid_shape, Grid_data); + test.AddOutput("Y", Y_shape, {0.25f, 0.25f, 0.25f}); + RunCpuAndCuda(test); +} + +TEST(GridSampleCudaHardeningTest, cubic_reflection_extreme_coords) { + // CUDA bicubic path (4-D only): static_cast(std::floor(...)) - 1 is now + // protected by the coordinate sanitization. Constant-valued image => 1.0. + OpTester test("GridSample", 20); + test.AddAttribute("mode", std::string("cubic")); + test.AddAttribute("padding_mode", std::string("reflection")); + test.AddAttribute("align_corners", int64_t{0}); + + std::initializer_list X_shape{1, 1, 4, 4}; + std::initializer_list X_data{ + 1.0f, 1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f, 1.0f}; + + std::initializer_list Grid_shape{1, 1, 1, 2}; + std::initializer_list Grid_data{1.0e+10f, -1.0e+10f}; + + std::initializer_list Y_shape{1, 1, 1, 1}; + + test.AddInput("X", X_shape, X_data); + test.AddInput("Grid", Grid_shape, Grid_data); + test.AddOutput("Y", Y_shape, {1.0f}); + RunCpuAndCuda(test); +} + +TEST(GridSampleCudaHardeningTest, nearest_reflection_dim1_align_corners) { + // 1x1 image with align_corners=1 makes the reflection range zero (x_min == x_max), so + // GsReflect must not divide by zero. The single input pixel is the only valid sample. + OpTester test("GridSample", 20); + test.AddAttribute("mode", std::string("nearest")); + test.AddAttribute("padding_mode", std::string("reflection")); + test.AddAttribute("align_corners", int64_t{1}); + + std::initializer_list X_shape{1, 1, 1, 1}; + std::initializer_list X_data{7.0f}; + + std::initializer_list Grid_shape{1, 1, 2, 2}; + std::initializer_list Grid_data{ + 0.5f, 0.5f, + -0.5f, -0.5f}; + + std::initializer_list Y_shape{1, 1, 1, 2}; + + test.AddInput("X", X_shape, X_data); + test.AddInput("Grid", Grid_shape, Grid_data); + test.AddOutput("Y", Y_shape, {7.0f, 7.0f}); + RunCpuAndCuda(test); +} + +TEST(GridSampleCudaHardeningTest, nearest_reflection_5D_extreme_coords) { + // CUDA _GridSampleKernel3D reflection path: PixelAtGrid3D received the same safety clamp + // as PixelAtGrid. Constant-valued volume => 1.0. + OpTester test("GridSample", 20); + test.AddAttribute("mode", std::string("nearest")); + test.AddAttribute("padding_mode", std::string("reflection")); + test.AddAttribute("align_corners", int64_t{0}); + + std::initializer_list X_shape{1, 1, 2, 2, 2}; + std::initializer_list X_data{ + 1.0f, 1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f, 1.0f}; + + std::initializer_list Grid_shape{1, 1, 1, 2, 3}; + std::initializer_list Grid_data{ + 1.0e+10f, 1.0e+10f, 1.0e+10f, + -1.0e+10f, -1.0e+10f, -1.0e+10f}; + + std::initializer_list Y_shape{1, 1, 1, 1, 2}; + + test.AddInput("X", X_shape, X_data); + test.AddInput("Grid", Grid_shape, Grid_data); + test.AddOutput("Y", Y_shape, {1.0f, 1.0f}); + RunCpuAndCuda(test); +} + +TEST(GridSampleCudaHardeningTest, bilinear_reflection_5D_extreme_coords) { + // CUDA 3D trilinear path: the x/y/z indices are computed via int64_t floor casts after + // sanitization. Constant-valued volume => 1.0. + OpTester test("GridSample", 20); + test.AddAttribute("mode", std::string("linear")); + test.AddAttribute("padding_mode", std::string("reflection")); + test.AddAttribute("align_corners", int64_t{0}); + + std::initializer_list X_shape{1, 1, 2, 2, 2}; + std::initializer_list X_data{ + 1.0f, 1.0f, 1.0f, 1.0f, + 1.0f, 1.0f, 1.0f, 1.0f}; + + std::initializer_list Grid_shape{1, 1, 1, 2, 3}; + std::initializer_list Grid_data{ + 1.0e+20f, 1.0e+20f, 1.0e+20f, + -1.0e+20f, -1.0e+20f, -1.0e+20f}; + + std::initializer_list Y_shape{1, 1, 1, 1, 2}; + + test.AddInput("X", X_shape, X_data); + test.AddInput("Grid", Grid_shape, Grid_data); + test.AddOutput("Y", Y_shape, {1.0f, 1.0f}); + RunCpuAndCuda(test); +} + } // namespace test } // namespace onnxruntime From 4d8886f833de548797bd41ea4e159bf5d90628d0 Mon Sep 17 00:00:00 2001 From: Ankit Maheshkar Date: Thu, 9 Jul 2026 07:18:44 +0530 Subject: [PATCH 05/11] [ORT] onnxruntime_perf_test: Add multi-shape profiling support to perftest tool (--data_shape) (#29571) ### Description This PR adds a `--data_shape` flag to `onnxruntime_perf_test` app that enables profiling multiple input shapes within a single session. The model is compiled once and then run with each specified shape group in round-robin order, providing per-shape latency statistics (Avg, Min, Max, P50, P90, P95, P99). This is useful for evaluating dynamic-shape model performance across different input dimensions without needing to recompile or restart the session for each shape. Usage: #With -I (generate random default inputs of specified shapes): `onnxruntime_perf_test.exe -v -e cpu -m times -r 10 -I --data_shape "data:[1,3,60,10][1,3,60,20]" "model.onnx"` #Without -I (select matching test data folders): `onnxruntime_perf_test.exe -v -e cpu -m times -r 10 --data_shape "data:[1,3,60,10][1,3,60,20]" "model.onnx"` ### Motivation and Context Addresses the feature request: #28628 --------- Co-authored-by: n1harika --- .../test/perftest/command_args_parser.cc | 19 ++ onnxruntime/test/perftest/ort_test_session.cc | 162 +++++++++++++---- onnxruntime/test/perftest/ort_test_session.h | 14 ++ .../test/perftest/performance_runner.cc | 135 +++++++++++++- .../test/perftest/performance_runner.h | 17 +- onnxruntime/test/perftest/strings_helper.cc | 165 ++++++++++++++++++ onnxruntime/test/perftest/strings_helper.h | 13 +- .../test/perftest/test_configuration.h | 2 + onnxruntime/test/perftest/test_session.h | 1 + 9 files changed, 482 insertions(+), 46 deletions(-) diff --git a/onnxruntime/test/perftest/command_args_parser.cc b/onnxruntime/test/perftest/command_args_parser.cc index 3f878ac796ee3..c8698140a8731 100644 --- a/onnxruntime/test/perftest/command_args_parser.cc +++ b/onnxruntime/test/perftest/command_args_parser.cc @@ -214,6 +214,15 @@ ABSL_FLAG(bool, compile_ep_context, DefaultPerformanceTestConfig().run_config.co ABSL_FLAG(std::string, compile_model_path, "model_ctx.onnx", "The compiled model path for saving EP context model. Overwrites if already exists"); ABSL_FLAG(bool, compile_binary_embed, DefaultPerformanceTestConfig().run_config.compile_binary_embed, "Embed binary blob within EP context node"); ABSL_FLAG(bool, compile_only, DefaultPerformanceTestConfig().run_config.compile_only, "Only compile EP context model without running it"); +ABSL_FLAG(std::string, data_shape, "", + "Specifies input shapes for multi-shape profiling within a single session.\n" + "The model is compiled once and run with each shape group in round-robin order.\n" + "[Usage]: --data_shape \"input_name:[d0,d1,...][d0,d1,...] ...\"\n" + "[Example]: --data_shape \"input:[1,3,224,224][1,3,448,448][1,3,112,112]\"\n" + "With -I: generates random input of the specified shapes.\n" + "Without -I: selects test data folders whose tensor shapes match.\n" + "All inputs must have the same number of shape groups.\n" + "All dimension values must be positive integers."); ABSL_FLAG(bool, h, false, "Print program usage."); namespace onnxruntime { @@ -481,6 +490,16 @@ bool CommandLineParser::ParseArguments(PerformanceTestConfig& test_config, int a // -I test_config.run_config.generate_model_input_binding = absl::GetFlag(FLAGS_I); + // --data_shape + { + const auto& data_shape_str = absl::GetFlag(FLAGS_data_shape); + if (!data_shape_str.empty()) { + if (!ParseDataShapeGroups(data_shape_str, test_config.run_config.data_shape_groups)) { + return false; + } + } + } + // -d if (absl::GetFlag(FLAGS_d) < 0) return false; test_config.run_config.cudnn_conv_algo = absl::GetFlag(FLAGS_d); diff --git a/onnxruntime/test/perftest/ort_test_session.cc b/onnxruntime/test/perftest/ort_test_session.cc index 679be75fa586e..ae9ed2a504320 100644 --- a/onnxruntime/test/perftest/ort_test_session.cc +++ b/onnxruntime/test/perftest/ort_test_session.cc @@ -42,9 +42,15 @@ namespace onnxruntime { namespace perftest { RunTiming OnnxRuntimeTestSession::Run() { - // Randomly pick one OrtValueArray from test_inputs_. (NOT ThreadSafe) - const std::uniform_int_distribution::param_type p(0, static_cast(test_inputs_.size() - 1)); - const size_t id = static_cast(dist_(rand_engine_, p)); + // Select input set: round-robin for multi-shape mode, random otherwise. + size_t id; + if (use_round_robin_ && test_inputs_.size() > 1) { + id = round_robin_counter_.fetch_add(1, std::memory_order_relaxed) % test_inputs_.size(); + } else { + // Random selection (not thread-safe). + const std::uniform_int_distribution::param_type p(0, static_cast(test_inputs_.size() - 1)); + id = static_cast(dist_(rand_engine_, p)); + } auto& input = test_inputs_.at(id); auto start = std::chrono::high_resolution_clock::now(); @@ -88,6 +94,7 @@ RunTiming OnnxRuntimeTestSession::Run() { timing.submit_timing = std::chrono::high_resolution_clock::now() - start; timing.total_timing = timing.submit_timing; } + timing.test_input_index = id; return timing; } @@ -1083,54 +1090,139 @@ static void InitializeTensorWithSeed(int32_t seed, Ort::Value& tensor) { #undef CASE_FOR_TYPE } +void OnnxRuntimeTestSession::CreateAndStoreGeneratedInput(size_t test_data_id, size_t input_idx, + const std::vector& dims, + ONNXTensorElementDataType element_type, int32_t seed) { + if (device_memory_name_ != CUDA) { + Ort::Value input_tensor = Ort::Value::CreateTensor(allocator_, (const int64_t*)dims.data(), + dims.size(), element_type); + InitializeTensorWithSeed(seed, input_tensor); + PreLoadTestData(test_data_id, input_idx, std::move(input_tensor)); + } +#if defined(USE_CUDA) || defined(USE_TENSORRT) || defined(USE_NV) + else { + Ort::AllocatorWithDefaultOptions default_allocator; + Ort::Value default_tensor = Ort::Value::CreateTensor(default_allocator, (const int64_t*)dims.data(), + dims.size(), element_type); + InitializeTensorWithSeed(seed, default_tensor); + + const void* default_ptr = default_tensor.GetTensorRawData(); + size_t total_bytes = default_tensor.GetTensorSizeInBytes(); + + Ort::Value cuda_tensor = Ort::Value::CreateTensor(allocator_, dims.data(), + dims.size(), element_type); + void* cuda_ptr = cuda_tensor.GetTensorMutableData(); + + cudaError_t cuda_err = cudaMemcpy(cuda_ptr, default_ptr, total_bytes, cudaMemcpyHostToDevice); + if (cuda_err != cudaSuccess) { + ORT_THROW("Failed to copy tensor data from CPU to CUDA device. CUDA Error: ", cudaGetErrorString(cuda_err)); + } + PreLoadTestData(test_data_id, input_idx, std::move(cuda_tensor)); + } +#endif +} + bool OnnxRuntimeTestSession::PopulateGeneratedInputTestData(int32_t seed) { - Ort::AllocatorWithDefaultOptions default_allocator; - // iterate over all input nodes for (size_t i = 0; i < static_cast(input_length_); i++) { Ort::TypeInfo type_info = session_.GetInputTypeInfo(i); - if (type_info.GetONNXType() == ONNX_TYPE_TENSOR) { - auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); - std::vector input_node_dim = tensor_info.GetShape(); + if (type_info.GetONNXType() != ONNX_TYPE_TENSOR) { + continue; + } - // free dimensions are treated as 1 if not overridden - auto transform_fcn = [](int64_t input) { return (input == -1) ? -input : input; }; - std::transform(input_node_dim.begin(), input_node_dim.end(), input_node_dim.begin(), transform_fcn); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector input_node_dim = tensor_info.GetShape(); - if (device_memory_name_ != CUDA) { - Ort::Value input_tensor = Ort::Value::CreateTensor(allocator_, (const int64_t*)input_node_dim.data(), - input_node_dim.size(), tensor_info.GetElementType()); - InitializeTensorWithSeed(seed, input_tensor); - PreLoadTestData(0, i, std::move(input_tensor)); - } -// Create tensor on CPU, initialize and copy to CUDA tensor -#if defined(USE_CUDA) || defined(USE_TENSORRT) || defined(USE_NV) - else { - Ort::Value default_tensor = Ort::Value::CreateTensor(default_allocator, (const int64_t*)input_node_dim.data(), - input_node_dim.size(), tensor_info.GetElementType()); - InitializeTensorWithSeed(seed, default_tensor); + // free dimensions are treated as 1 if not overridden + auto transform_fcn = [](int64_t input) { return (input == -1) ? -input : input; }; + std::transform(input_node_dim.begin(), input_node_dim.end(), input_node_dim.begin(), transform_fcn); - // Get pointer to CPU tensor data - const void* default_ptr = default_tensor.GetTensorRawData(); + CreateAndStoreGeneratedInput(0, i, input_node_dim, tensor_info.GetElementType(), seed); + } + return true; +} - size_t total_bytes = default_tensor.GetTensorSizeInBytes(); +bool OnnxRuntimeTestSession::PopulateGeneratedMultiShapeInputTestData( + int32_t seed, + const std::map>>& data_shape_groups) { + // Validate that all input names in data_shape_groups exist in the model + for (const auto& [name, groups] : data_shape_groups) { + bool found = false; + for (int i = 0; i < input_length_; i++) { + if (input_names_str_[i] == name) { + found = true; + break; + } + } + if (!found) { + std::cerr << "Error: --data_shape specifies unknown input '" << name << "'." << std::endl; + return false; + } + } - Ort::Value cuda_tensor = Ort::Value::CreateTensor(allocator_, input_node_dim.data(), - input_node_dim.size(), tensor_info.GetElementType()); + const size_t num_groups = data_shape_groups.begin()->second.size(); - void* cuda_ptr = cuda_tensor.GetTensorMutableData(); + for (size_t g = 0; g < num_groups; g++) { + for (size_t i = 0; i < static_cast(input_length_); i++) { + Ort::TypeInfo type_info = session_.GetInputTypeInfo(i); + if (type_info.GetONNXType() != ONNX_TYPE_TENSOR) { + continue; + } - // Copy the initialized data from CPU to GPU - cudaError_t cuda_err = cudaMemcpy(cuda_ptr, default_ptr, total_bytes, cudaMemcpyHostToDevice); - if (cuda_err != cudaSuccess) { - ORT_THROW("Failed to copy tensor data from CPU to CUDA device. CUDA Error: ", cudaGetErrorString(cuda_err)); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + std::vector input_node_dim; + + // Use user-specified shape if available, otherwise fall back to model metadata + auto it = data_shape_groups.find(input_names_str_[i]); + if (it != data_shape_groups.end()) { + input_node_dim = it->second[g]; + const auto model_shape = tensor_info.GetShape(); + if (!model_shape.empty() && input_node_dim.size() != model_shape.size()) { + std::cerr << "Error: --data_shape rank mismatch for input '" << input_names_str_[i] + << "': expected " << model_shape.size() << " dims but got " << input_node_dim.size() << "." << std::endl; + return false; + } + } else { + input_node_dim = tensor_info.GetShape(); + bool has_dynamic_dim = std::any_of(input_node_dim.begin(), input_node_dim.end(), + [](int64_t d) { return d == -1; }); + auto transform_fcn = [](int64_t input) { return (input == -1) ? -input : input; }; + std::transform(input_node_dim.begin(), input_node_dim.end(), input_node_dim.begin(), transform_fcn); + if (g == 0 && has_dynamic_dim) { + std::cerr << "Warning: input '" << input_names_str_[i] + << "' not specified in --data_shape; using inferred shape ["; + for (size_t d = 0; d < input_node_dim.size(); d++) { + if (d > 0) std::cerr << ","; + std::cerr << input_node_dim[d]; + } + std::cerr << "] (dynamic dims defaulted to 1)." << std::endl; } - PreLoadTestData(0, i, std::move(cuda_tensor)); } -#endif + + CreateAndStoreGeneratedInput(g, i, input_node_dim, tensor_info.GetElementType(), seed); } } + use_round_robin_ = true; return true; } + +std::vector OnnxRuntimeTestSession::GetLoadedInputShape(size_t test_data_id, size_t input_id) const { + const auto& v = test_inputs_.at(test_data_id).at(input_id); + if (!v.IsTensor()) { + ORT_THROW("--data_shape only supports tensor inputs; input_id=", input_id, " in test_data_id=", test_data_id, + " is not a tensor."); + } + return v.GetTensorTypeAndShapeInfo().GetShape(); +} + +void OnnxRuntimeTestSession::SelectTestDataSets(const std::vector& selected_ids) { + std::vector> filtered; + filtered.reserve(selected_ids.size()); + for (size_t id : selected_ids) { + filtered.push_back(std::move(test_inputs_.at(id))); + } + test_inputs_ = std::move(filtered); +} + OnnxRuntimeTestSession::~OnnxRuntimeTestSession() { #ifdef USE_CUDA if (device_memory_name_ == CUDA && stream_ != nullptr) { diff --git a/onnxruntime/test/perftest/ort_test_session.h b/onnxruntime/test/perftest/ort_test_session.h index 0a2c683428bf2..422e76ba0e72d 100644 --- a/onnxruntime/test/perftest/ort_test_session.h +++ b/onnxruntime/test/perftest/ort_test_session.h @@ -2,6 +2,7 @@ // Licensed under the MIT License. #pragma once +#include #include #include #include "test_configuration.h" @@ -31,6 +32,13 @@ class OnnxRuntimeTestSession : public TestSession { } bool PopulateGeneratedInputTestData(int32_t seed); + bool PopulateGeneratedMultiShapeInputTestData( + int32_t seed, + const std::map>>& data_shape_groups); + + std::vector GetLoadedInputShape(size_t test_data_id, size_t input_id) const; + void SelectTestDataSets(const std::vector& selected_ids); + void SetUseRoundRobin(bool v) { use_round_robin_ = v; } ~OnnxRuntimeTestSession(); @@ -39,6 +47,10 @@ class OnnxRuntimeTestSession : public TestSession { ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(OnnxRuntimeTestSession); private: + void CreateAndStoreGeneratedInput(size_t test_data_id, size_t input_idx, + const std::vector& dims, + ONNXTensorElementDataType element_type, int32_t seed); + Ort::Session session_{nullptr}; std::mt19937 rand_engine_; std::uniform_int_distribution dist_; @@ -60,6 +72,8 @@ class OnnxRuntimeTestSession : public TestSession { std::string device_memory_name_; // Device memory type name to use from the list in allocator.h const std::unordered_map& run_config_entries_; bool has_dynamic_output_shapes_ = false; + std::atomic round_robin_counter_{0}; + bool use_round_robin_{false}; #if defined(USE_CUDA) || defined(USE_TENSORRT) || defined(USE_NV) cudaStream_t stream_; // Device stream if required by IO bindings #endif diff --git a/onnxruntime/test/perftest/performance_runner.cc b/onnxruntime/test/perftest/performance_runner.cc index a14e6437c0e01..91334ae00f81f 100644 --- a/onnxruntime/test/perftest/performance_runner.cc +++ b/onnxruntime/test/perftest/performance_runner.cc @@ -10,6 +10,7 @@ #include #include "TestCase.h" +#include "strings_helper.h" #include "utils.h" #include "ort_test_session.h" using onnxruntime::Status; @@ -54,7 +55,8 @@ Eigen::ThreadPoolInterface* GetDefaultThreadPool(const onnxruntime::Env& env) { namespace onnxruntime { namespace perftest { -void PerformanceResult::DumpToFile(const std::basic_string& path, bool f_include_statistics) const { +void PerformanceResult::DumpToFile(const std::basic_string& path, bool f_include_statistics, + const std::map>>& shape_groups) const { bool have_file = !path.empty(); std::ofstream outfile; @@ -109,6 +111,48 @@ void PerformanceResult::DumpToFile(const std::basic_string& path, boo output_stats(std::cout); } + + // Per-shape statistics (when --data_shape is used) + if (!per_shape_time_costs_total.empty() && f_include_statistics) { + auto output_per_shape = [&](std::ostream& ostream) { + ostream << "\nLatency per shape group:" << std::endl; + for (size_t g = 0; g < per_shape_time_costs_total.size(); g++) { + const auto& shape_costs = per_shape_time_costs_total[g]; + + std::string label = FormatShapeGroup(shape_groups, g); + ostream << " " << (g + 1) << ". " << (label.empty() ? "shape group" : label) << std::endl; + + if (shape_costs.empty()) { + ostream << " (no data)" << std::endl; + continue; + } + + std::vector sorted = shape_costs; + std::sort(sorted.begin(), sorted.end()); + size_t count = sorted.size(); + size_t s50 = static_cast(count * 0.5); + size_t s90 = static_cast(count * 0.9); + size_t s95 = static_cast(count * 0.95); + size_t s99 = static_cast(count * 0.99); + double avg = std::accumulate(sorted.begin(), sorted.end(), 0.0) / static_cast(count); + + ostream << " Iterations: " << count << "\n" + << " Average Latency: " << avg * 1000.0 << " ms\n" + << " Min Latency: " << sorted[0] * 1000.0 << " ms\n" + << " Max Latency: " << sorted[count - 1] * 1000.0 << " ms\n" + << " P50 Latency: " << sorted[s50] * 1000.0 << " ms\n" + << " P90 Latency: " << sorted[s90] * 1000.0 << " ms\n" + << " P95 Latency: " << sorted[s95] * 1000.0 << " ms\n" + << " P99 Latency: " << sorted[s99] * 1000.0 << " ms" << std::endl; + } + }; + + if (have_file) { + output_per_shape(outfile); + } + + output_per_shape(std::cout); + } } void PerformanceRunner::LogSessionCreationTime() { @@ -121,9 +165,15 @@ Status PerformanceRunner::Run() { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "failed to initialize."); } - // warm up + // warm up — run one iteration per shape group to ensure all shapes are warmed + size_t warmup_count = 1; + if (!performance_test_config_.run_config.data_shape_groups.empty()) { + warmup_count = performance_test_config_.run_config.data_shape_groups.begin()->second.size(); + } initial_inference_result_.start = std::chrono::high_resolution_clock::now(); - ORT_RETURN_IF_ERROR(RunOneIteration()); + for (size_t w = 0; w < warmup_count; w++) { + ORT_RETURN_IF_ERROR(RunOneIteration()); + } initial_inference_result_.end = std::chrono::high_resolution_clock::now(); // TODO: start profiling @@ -306,9 +356,20 @@ bool PerformanceRunner::Initialize() { test_case_ = CreateOnnxTestCase(narrow_model_name, std::move(test_model_info_), 0.0, 0.0); if (performance_test_config_.run_config.generate_model_input_binding) { - return static_cast( - session_.get()) - ->PopulateGeneratedInputTestData(performance_test_config_.run_config.random_seed_for_input_data); + auto* ort_session = static_cast(session_.get()); + if (!performance_test_config_.run_config.data_shape_groups.empty()) { + if (!ort_session->PopulateGeneratedMultiShapeInputTestData( + performance_test_config_.run_config.random_seed_for_input_data, + performance_test_config_.run_config.data_shape_groups)) { + return false; + } + // Pre-size per-shape timing vectors + size_t num_groups = performance_test_config_.run_config.data_shape_groups.begin()->second.size(); + performance_result_.per_shape_time_costs_total.resize(num_groups); + return true; + } + return ort_session->PopulateGeneratedInputTestData( + performance_test_config_.run_config.random_seed_for_input_data); } // TODO: Place input tensor on cpu memory if dnnl provider type to avoid CopyTensor logic in CopyInputAcrossDevices @@ -317,11 +378,10 @@ bool PerformanceRunner::Initialize() { std::cout << "there is no test data for model " << test_case_->GetTestCaseName() << std::endl; return false; } + int input_count = test_model_info->GetInputCount(); for (size_t test_data_id = 0; test_data_id != test_data_count; ++test_data_id) { std::unordered_map feeds; test_case_->LoadTestData(test_data_id /* id */, b_, feeds, true); - // Discard the names in feeds - int input_count = test_model_info->GetInputCount(); for (int i = 0; i != input_count; ++i) { auto iter = feeds.find(test_model_info->GetInputName(i)); if (iter == feeds.end()) { @@ -333,6 +393,65 @@ bool PerformanceRunner::Initialize() { } } + // When --data_shape is specified without -I, select only the test data sets + // whose tensor shapes match the requested shape groups. + const auto& data_shape_groups = performance_test_config_.run_config.data_shape_groups; + if (!data_shape_groups.empty()) { + auto* ort_session = static_cast(session_.get()); + size_t num_groups = data_shape_groups.begin()->second.size(); + + // Build input name -> index map + std::unordered_map input_name_to_idx; + for (int i = 0; i < input_count; ++i) { + input_name_to_idx[test_model_info->GetInputName(i)] = static_cast(i); + } + + // Validate that all names in data_shape_groups exist in the model + for (const auto& [name, groups] : data_shape_groups) { + if (input_name_to_idx.find(name) == input_name_to_idx.end()) { + std::cerr << "Error: --data_shape specifies unknown input '" << name << "'." << std::endl; + return false; + } + } + + // For each shape group, find the matching test data set + std::vector selected_ids; + for (size_t g = 0; g < num_groups; ++g) { + bool found = false; + for (size_t test_data_id = 0; test_data_id < test_data_count; ++test_data_id) { + bool match = true; + for (const auto& [input_name, shape_list] : data_shape_groups) { + size_t input_idx = input_name_to_idx[input_name]; + auto loaded_shape = ort_session->GetLoadedInputShape(test_data_id, input_idx); + if (loaded_shape != shape_list[g]) { + match = false; + break; + } + } + if (match) { + if (std::find(selected_ids.begin(), selected_ids.end(), test_data_id) != selected_ids.end()) { + std::cerr << "Error: --data_shape shape group " << (g + 1) + << " matches a test data set already selected. " + << "Duplicate shape groups are not supported when selecting from test data." << std::endl; + return false; + } + selected_ids.push_back(test_data_id); + found = true; + break; + } + } + if (!found) { + std::cerr << "No test data found matching shape group " << (g + 1) << ": " + << FormatShapeGroup(data_shape_groups, g) << "." << std::endl; + return false; + } + } + + ort_session->SelectTestDataSets(selected_ids); + ort_session->SetUseRoundRobin(true); + performance_result_.per_shape_time_costs_total.resize(num_groups); + } + return true; } diff --git a/onnxruntime/test/perftest/performance_runner.h b/onnxruntime/test/perftest/performance_runner.h index f67ab9f5c2287..34b699573ddc4 100644 --- a/onnxruntime/test/perftest/performance_runner.h +++ b/onnxruntime/test/perftest/performance_runner.h @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -35,9 +36,12 @@ struct PerformanceResult { double total_time_cost{0}; std::vector time_costs_submission; std::vector time_costs_total; + // Per-shape-group timing (only populated when --data_shape is used) + std::vector> per_shape_time_costs_total; std::string model_name; - void DumpToFile(const std::basic_string& path, bool f_include_statistics = false) const; + void DumpToFile(const std::basic_string& path, bool f_include_statistics = false, + const std::map>>& shape_groups = {}) const; }; class PerformanceRunner { @@ -53,7 +57,8 @@ class PerformanceRunner { inline void SerializeResult() const { performance_result_.DumpToFile(performance_test_config_.model_info.result_file_path, - performance_test_config_.run_config.f_dump_statistics); + performance_test_config_.run_config.f_dump_statistics, + performance_test_config_.run_config.data_shape_groups); } ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(PerformanceRunner); @@ -80,6 +85,14 @@ class PerformanceRunner { performance_result_.time_costs_total.emplace_back(duration_seconds.total_timing.count()); // if the test did not run async the device timing will be equal to the CPU timing performance_result_.total_time_cost += duration_seconds.total_timing.count(); + // Record per-shape timing when multi-shape mode is active + if (!performance_result_.per_shape_time_costs_total.empty()) { + size_t shape_idx = duration_seconds.test_input_index; + if (shape_idx < performance_result_.per_shape_time_costs_total.size()) { + performance_result_.per_shape_time_costs_total[shape_idx].emplace_back( + duration_seconds.total_timing.count()); + } + } if (performance_test_config_.run_config.f_verbose) { std::cout << "iteration:" << performance_result_.time_costs_submission.size() << "," << "time_cost_submission:" << performance_result_.time_costs_submission.back() << "," diff --git a/onnxruntime/test/perftest/strings_helper.cc b/onnxruntime/test/perftest/strings_helper.cc index d9fd2a2a55c09..bb57b66303c0b 100644 --- a/onnxruntime/test/perftest/strings_helper.cc +++ b/onnxruntime/test/perftest/strings_helper.cc @@ -3,8 +3,10 @@ // SPDX-FileCopyrightText: Copyright 2024 Arm Limited and/or its affiliates // Licensed under the MIT License. +#include #include #include +#include #include "strings_helper.h" #include "core/common/common.h" @@ -154,5 +156,168 @@ void ParseEpDeviceFilterKeyValuePairs(const std::string& input, std::vector>>& data_shape_groups) { + // Parse format: "input_name:[d0,d1,...][d0,d1,...] input_name2:[d0,d1][d0,d1]" + // Split on whitespace that is outside brackets to allow spaces inside shape specs. + std::vector input_specs; + std::string current; + int bracket_depth = 0; + for (char c : input) { + if (c == '[') { + bracket_depth++; + current += c; + } else if (c == ']') { + bracket_depth--; + current += c; + } else if ((c == ' ' || c == '\t') && bracket_depth == 0) { + if (!current.empty()) { + input_specs.push_back(current); + current.clear(); + } + } else { + current += c; + } + } + if (!current.empty()) { + input_specs.push_back(current); + } + + if (bracket_depth != 0) { + std::cerr << "Error parsing --data_shape: mismatched brackets (unbalanced '[' and ']')." << std::endl; + return false; + } + + for (const auto& input_spec : input_specs) { + // Split on first ':' to get input_name and shapes string + size_t colon_pos = input_spec.find(':'); + if (colon_pos == std::string::npos || colon_pos == 0) { + std::cerr << "Error parsing --data_shape: expected 'name:[shape]...' format, got '" << input_spec << "'." << std::endl; + return false; + } + + std::string input_name = input_spec.substr(0, colon_pos); + std::string shapes_str = input_spec.substr(colon_pos + 1); + + if (shapes_str.empty()) { + std::cerr << "Error parsing --data_shape: no shape groups provided for input '" << input_name << "'." << std::endl; + return false; + } + + std::vector> shape_groups; + + // Extract bracket-delimited shape groups: [d0,d1,...][d0,d1,...] + size_t pos = 0; + while (pos < shapes_str.size()) { + while (pos < shapes_str.size() && (shapes_str[pos] == ' ' || shapes_str[pos] == '\t')) { + ++pos; + } + if (pos >= shapes_str.size()) break; + if (shapes_str[pos] != '[') { + std::cerr << "Error parsing --data_shape: expected '[' at position " << pos + << " in shapes for input '" << input_name << "'." << std::endl; + return false; + } + + size_t close_bracket = shapes_str.find(']', pos); + if (close_bracket == std::string::npos) { + std::cerr << "Error parsing --data_shape: unmatched bracket for input '" << input_name << "'." << std::endl; + return false; + } + + std::string dims_str = shapes_str.substr(pos + 1, close_bracket - pos - 1); + if (dims_str.empty()) { + std::cerr << "Error parsing --data_shape: empty shape group for input '" << input_name << "'." << std::endl; + return false; + } + + // Parse comma-separated dimensions + std::vector dims; + std::stringstream dims_ss(dims_str); + std::string dim_token; + while (std::getline(dims_ss, dim_token, ',')) { + // Trim whitespace from dimension token + size_t start = dim_token.find_first_not_of(" \t"); + size_t end = dim_token.find_last_not_of(" \t"); + if (start == std::string::npos) { + std::cerr << "Error parsing --data_shape: empty dimension token for input '" << input_name << "'." << std::endl; + return false; + } + dim_token = dim_token.substr(start, end - start + 1); + + int64_t dim_val = 0; + const char* begin = dim_token.data(); + const char* end_ptr = begin + dim_token.size(); + auto [ptr, ec] = std::from_chars(begin, end_ptr, dim_val); + if (ec != std::errc{} || ptr != end_ptr) { + std::cerr << "Error parsing --data_shape: invalid dimension value '" + << dim_token << "' for input '" << input_name << "'." << std::endl; + return false; + } + if (dim_val <= 0) { + std::cerr << "Error parsing --data_shape: dimensions must be positive integers, got '" + << dim_token << "' for input '" << input_name << "'." << std::endl; + return false; + } + dims.push_back(dim_val); + } + + shape_groups.push_back(std::move(dims)); + pos = close_bracket + 1; + } + + if (shape_groups.empty()) { + std::cerr << "Error parsing --data_shape: no shape groups found for input '" << input_name << "'." << std::endl; + return false; + } + + if (data_shape_groups.count(input_name) > 0) { + std::cerr << "Error parsing --data_shape: duplicate input name '" << input_name << "'." << std::endl; + return false; + } + + data_shape_groups[input_name] = std::move(shape_groups); + } + + if (data_shape_groups.empty()) { + std::cerr << "Error parsing --data_shape: no input specifications found." << std::endl; + return false; + } + + // Validate all inputs have the same number of shape groups + if (data_shape_groups.size() > 1) { + size_t expected_count = data_shape_groups.begin()->second.size(); + for (const auto& [name, groups] : data_shape_groups) { + if (groups.size() != expected_count) { + std::cerr << "Error parsing --data_shape: all inputs must have the same number of shape groups. " + << "Input '" << data_shape_groups.begin()->first << "' has " << expected_count + << " groups but input '" << name << "' has " << groups.size() << "." << std::endl; + return false; + } + } + } + + return true; +} + +std::string FormatShapeGroup(const std::map>>& groups, size_t g) { + if (groups.empty() || g >= groups.begin()->second.size()) { + return ""; + } + std::string result; + for (const auto& [name, shapes] : groups) { + if (!result.empty()) result += ", "; + result += name + ":["; + const auto& dims = shapes[g]; + for (size_t d = 0; d < dims.size(); d++) { + if (d > 0) result += ","; + result += std::to_string(dims[d]); + } + result += "]"; + } + return result; +} + } // namespace perftest } // namespace onnxruntime diff --git a/onnxruntime/test/perftest/strings_helper.h b/onnxruntime/test/perftest/strings_helper.h index d6c6f6112ab6c..38ad40cf1a92a 100644 --- a/onnxruntime/test/perftest/strings_helper.h +++ b/onnxruntime/test/perftest/strings_helper.h @@ -2,8 +2,13 @@ // Copyright (c) 2023 NVIDIA Corporation. // SPDX-FileCopyrightText: Copyright 2024 Arm Limited and/or its affiliates // Licensed under the MIT License. -#include + +#pragma once + +#include #include +#include +#include #include #include #include @@ -26,5 +31,11 @@ void ParseEpOptions(const std::string& input, std::vector& result); void ParseEpDeviceFilterKeyValuePairs(const std::string& input, std::vector>& result); + +bool ParseDataShapeGroups(const std::string& input, + std::map>>& data_shape_groups); + +std::string FormatShapeGroup(const std::map>>& groups, size_t g); + } // namespace perftest } // namespace onnxruntime diff --git a/onnxruntime/test/perftest/test_configuration.h b/onnxruntime/test/perftest/test_configuration.h index dea0b196f6bcb..b4e525902954a 100644 --- a/onnxruntime/test/perftest/test_configuration.h +++ b/onnxruntime/test/perftest/test_configuration.h @@ -8,6 +8,7 @@ #include #include #include +#include #include "core/graph/constants.h" #include "core/framework/session_options.h" @@ -67,6 +68,7 @@ struct RunConfig { std::unordered_map run_config_entries; std::map free_dim_name_overrides; std::map free_dim_denotation_overrides; + std::map>> data_shape_groups; std::string intra_op_thread_affinities; bool disable_spinning = false; bool disable_spinning_between_run = false; diff --git a/onnxruntime/test/perftest/test_session.h b/onnxruntime/test/perftest/test_session.h index 71cb6062a93f4..d2b7cee97efb8 100644 --- a/onnxruntime/test/perftest/test_session.h +++ b/onnxruntime/test/perftest/test_session.h @@ -12,6 +12,7 @@ namespace perftest { struct RunTiming { std::chrono::duration submit_timing = std::chrono::seconds(0); std::chrono::duration total_timing = std::chrono::seconds(0); + size_t test_input_index = 0; // which test_inputs_ entry was used }; class TestSession { From 1ad2a1f2ec4d019704dcee269b46ecd984489398 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 9 Jul 2026 13:51:29 +0800 Subject: [PATCH 06/11] webgpu: Bump FlashAttentionDecodeQKV workgroup to 128, tile_size_k_vec to 32 (#29586) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Bump `FlashAttentionDecodeQKVProgram` workgroup size from 64 to 128. - Bump `tile_size_k_vec` from 8 to 32. - `sub_tile_count` is derived from `WorkgroupSizeX() / tile_size_k_vec`, so it goes from 8 to 4 automatically. The shader template expresses all loops in terms of these parameters, so no shader-side changes are needed. ## Motivation `FlashAttentionDecodeQKVProgram` was dispatching with a smaller workgroup / tile_size_k_vec than `MatMulNBitsProgram` on the same class of hardware. Mirroring the MatMulNBits dispatch shape improves occupancy and reduces the per-call cost of the fused QK^T + softmax + V multiply step during token generation. Measured on Qwen3-1.7B-graph-prune (WebGPU EP, D3D12, decode): | Metric | Before | After | |---|---|---| | `GroupQueryAttention\|FlashAttentionDecodeQKV` per-call GPU time | 0.83 ms | 0.47 ms | | Token generation throughput | 257 tps | 282 tps | ## Test plan - [x] Built ORT with WebGPU EP (Release, D3D12) on Windows. - [x] `onnxruntime_provider_test.exe --gtest_filter="GroupQueryAttentionTest.*"` — 52 passed, 12 CUDA-only tests skipped. All WebGPU FlashAttention paths pass, including `BatchedRightPaddedRotaryPrefillFlashAttention_WebGPU`, `BatchedRightPaddedRotaryPrefillFlashAttentionLargeSpread_WebGPU`, and `WebGPU_SharedKV_*` variants. - [x] End-to-end correctness on Qwen3-1.7B-graph-prune (decode). - [x] Profiled Qwen3-1.7B-graph-prune with the WebGPU profiling pipeline; confirmed the per-op time improvement above. --- .../contrib_ops/webgpu/bert/flash_attention.cc | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index ce4018956f5c0..b2979a95ce22e 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -304,7 +304,11 @@ Status FlashAttentionDecodeQKVProgram::GenerateShaderCode(ShaderHelper& shader) const auto& out_split_vx = shader.AddOutput("out_split_vx", ShaderUsage::UseUniform); const auto& metadata = shader.AddOutput("metadata", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); - const uint32_t tile_size_k_vec = 8; + // Wider K tiling (32 vec4) with a 128-thread workgroup is used for decode (m_tile == 1) to + // mirror MatMulNBits and improve GPU time. For prefill (m_tile > 1) the shared-memory + // arrays that scale with tile_size_k_vec and m_tile would exceed the 32 KB workgroup + // storage limit on some adapters, so keep the original 8 vec4 / 64-thread shape there. + const uint32_t tile_size_k_vec = (m_tile_ == 1u) ? 32u : 8u; const uint32_t sub_tile_count = WorkgroupSizeX() / tile_size_k_vec; return WGSL_TEMPLATE_APPLY(shader, "bert/flash_attention_decode_qkv.wgsl.template", WGSL_TEMPLATE_PARAMETER(compressed_head_size_u32, compressed_head_size_u32_), @@ -376,7 +380,11 @@ Status ComputeFlashAttentionDecodeQKV(onnxruntime::webgpu::ComputeContext& conte } else { program.SetDispatchGroupSize(parameters.batch_size_ * parameters.num_heads_ * ((parameters.sequence_length_ + m_tile - 1) / m_tile) * num_total_seq_length_tile); } - program.SetWorkgroupSize(64) + // Workgroup size mirrors the tile_size_k_vec choice inside the program's shader (see + // FlashAttentionDecodeQKVProgram::GenerateShaderCode): 128 threads with 32 vec4 K tiles + // for decode, 64 threads with 8 vec4 K tiles for prefill. + const uint32_t workgroup_size = (m_tile == 1u) ? 128u : 64u; + program.SetWorkgroupSize(workgroup_size) .CacheHint(tile_size, head_size_vec, has_attention_bias, use_indirect_dispatch, q_BNSH, is_unidirectional, m_tile, use_seqlen_k, turbo_quant, compressed_head_size_u32) .AddUniformVariables({{static_cast(vectorized_head_size)}, {static_cast(parameters.total_sequence_length_)}, From 632de62365cc45f1bf7381a809f87319a75ffdd0 Mon Sep 17 00:00:00 2001 From: Jianhui Dai Date: Fri, 10 Jul 2026 00:41:23 +0800 Subject: [PATCH 07/11] [webgpu] Enable im2col-matmul for newer Xe platforms; limit to fp16 (#29505) ### Description - Add xe-2hpg, xe-3lpg, and xe-3lpg-xs to the execution allowlist. - Gate im2col-matmul on fp16 as the kernel is tuned for it. - Fall back to default convolution for fp32 and other data types. ### Motivation and Context See above. --- onnxruntime/core/providers/webgpu/nn/conv.cc | 6 ++++-- .../core/providers/webgpu/nn/im2col_matmul.cc | 14 ++++++++++++-- .../core/providers/webgpu/nn/im2col_matmul.h | 3 ++- 3 files changed, 18 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/nn/conv.cc b/onnxruntime/core/providers/webgpu/nn/conv.cc index c2a8896b84a7e..d980f819a264f 100644 --- a/onnxruntime/core/providers/webgpu/nn/conv.cc +++ b/onnxruntime/core/providers/webgpu/nn/conv.cc @@ -164,7 +164,8 @@ Status Conv::ComputeInternal(ComputeContext& context is_channels_last, activation_.activation_kind_ != ActivationKind::None, kernel_shape, - onnxruntime::narrow(conv_attrs_.group))) { + onnxruntime::narrow(conv_attrs_.group), + kernel->DataType())) { return ApplyIm2ColMatMulProgram(context, is_channels_last, dilations, @@ -334,7 +335,8 @@ Status Conv::PrePackInternal(ComputeContextBase& con // kernel directly from context.Input(1), ignoring prepacked weights. // Skip prepacking when this path will be used at runtime. if (CanApplyIm2ColMatMulProgram(context, is_channels_last, activation_.activation_kind_ != ActivationKind::None, - kernel_shape, onnxruntime::narrow(conv_attrs_.group))) { + kernel_shape, onnxruntime::narrow(conv_attrs_.group), + tensor.DataType())) { return Status::OK(); } diff --git a/onnxruntime/core/providers/webgpu/nn/im2col_matmul.cc b/onnxruntime/core/providers/webgpu/nn/im2col_matmul.cc index a1f385aee7626..8c33539253055 100644 --- a/onnxruntime/core/providers/webgpu/nn/im2col_matmul.cc +++ b/onnxruntime/core/providers/webgpu/nn/im2col_matmul.cc @@ -43,7 +43,10 @@ bool IsDeviceSupported(const ComputeContextBase& context) { const wgpu::AdapterInfo& adapter_info = context.AdapterInfo(); if (adapter_info.vendor == std::string_view("intel")) { - if (adapter_info.architecture == std::string_view("xe-2lpg")) { + if (adapter_info.architecture == std::string_view("xe-2lpg") || + adapter_info.architecture == std::string_view("xe-2hpg") || + adapter_info.architecture == std::string_view("xe-3lpg") || + adapter_info.architecture == std::string_view("xe-3lpg-xs")) { return true; } } @@ -167,11 +170,18 @@ bool CanApplyIm2ColMatMulProgram(ComputeContextBase& context, const bool is_channels_last, const bool is_fused, const TensorShape weight_shape, - const uint32_t group) { + const uint32_t group, + const MLDataType data_type) { if (!IsDeviceSupported(context)) { return false; } + // The im2col-matmul kernel is performance-tuned for fp16. Use the default + // conv path for fp32. + if (data_type != DataTypeImpl::GetType()) { + return false; + } + // TODO: Support !is_channels_last // TODO: Support fuse // TODO: Support group conv diff --git a/onnxruntime/core/providers/webgpu/nn/im2col_matmul.h b/onnxruntime/core/providers/webgpu/nn/im2col_matmul.h index fc01785442310..de46689bda921 100644 --- a/onnxruntime/core/providers/webgpu/nn/im2col_matmul.h +++ b/onnxruntime/core/providers/webgpu/nn/im2col_matmul.h @@ -65,7 +65,8 @@ bool CanApplyIm2ColMatMulProgram(ComputeContextBase& context, const bool is_channels_last, const bool is_fused, const TensorShape kernel_shape, - const uint32_t group); + const uint32_t group, + const MLDataType data_type); Status ApplyIm2ColMatMulProgram(ComputeContext& context, const bool is_channels_last, From 780d714226f250104503be378106b165b51e7db7 Mon Sep 17 00:00:00 2001 From: Gopalakrishnan Nallasamy Date: Thu, 9 Jul 2026 11:15:30 -0700 Subject: [PATCH 08/11] Add error-path tests for InferenceSession Run/Load/Initialize validation (#29555) ### Description Adds error-path unit tests for `InferenceSession` validation / early-return branches, which were under-covered (only failure/`kExpectFailure`-style assertions were a small fraction of the suite). All nine target previously-untested branches and are verified passing against a local Release build. | Test | Branch exercised | Asserted message | |---|---|---| | `RunBeforeInitializeReturnsError` | `Run` when not initialized | `Session not initialized` | | `InitializeBeforeLoadReturnsError` | `Initialize` with no model loaded | `Model was not loaded` | | `RunWithInvalidOutputNameReturnsError` | `ValidateOutputs` unknown output | `Invalid output name` | | `RunWithWrongInputTypeReturnsError` | `CheckTypes` (input) | `Unexpected input data type` | | `RunWithWrongInputRankReturnsError` | `CheckShapes` (input rank) | `Invalid rank for input` | | `RunWithMismatchedFeedCountReturnsError` | feed/name count mismatch | `feed names has ...` | | `RunWithWrongOutputTypeReturnsError` | `CheckTypes` (output) | `Unexpected output data type` | | `LoadMalformedModelFromArrayReturnsError` | malformed model bytes | graceful failure | | `LoadNonexistentModelReturnsError` | missing model file | graceful failure | ### Motivation and Context Error/validation paths in `InferenceSession::Run` / `Initialize` / `Load` were thinly tested relative to the numeric happy-path bulk. These tests lock in the failure contracts (both the non-OK `Status` and the message substring, via `ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR`) so a regression that changes the behavior or the message surfaces clearly. ### Notes - `"Invalid input name"` is already covered by `TestOptionalInputs`, so it is intentionally not duplicated. - Test-only change (`onnxruntime/test/framework/inference_session_test.cc`); no product/runtime code is touched. --------- Co-authored-by: Gopalakrishnan Nallasamy --- .../test/framework/inference_session_test.cc | 165 ++++++++++++++++++ 1 file changed, 165 insertions(+) diff --git a/onnxruntime/test/framework/inference_session_test.cc b/onnxruntime/test/framework/inference_session_test.cc index e478a23770afd..0da7d1d4786ca 100644 --- a/onnxruntime/test/framework/inference_session_test.cc +++ b/onnxruntime/test/framework/inference_session_test.cc @@ -353,6 +353,171 @@ TEST(InferenceSessionTests, RequestLoadCancellation) { } } +// Error-path coverage for InferenceSession validation / early-return branches. +// "Invalid input name" is already covered by TestOptionalInputs; the tests below cover the +// remaining gaps in Run/Initialize/Load: Run-before-Initialize, Initialize-before-Load, an +// invalid output name, input type/rank and output type mismatches, feed name/value count +// mismatch, and load failures (malformed model bytes and a nonexistent model path). +TEST(InferenceSessionTests, RunBeforeInitializeReturnsError) { + SessionOptions so; + so.session_logid = "InferenceSessionTests.RunBeforeInitializeReturnsError"; + InferenceSession session_object{so, GetEnvironment()}; + ASSERT_STATUS_OK(session_object.Load(MODEL_URI)); + // Intentionally do NOT call Initialize(). + + std::vector dims_x = {3, 2}; + std::vector values_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + OrtValue ml_value; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], dims_x, values_x, &ml_value); + NameMLValMap feeds; + feeds.insert(std::make_pair("X", ml_value)); + std::vector output_names{"Y"}; + std::vector fetches; + + RunOptions run_options; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(session_object.Run(run_options, feeds, output_names, &fetches), + "not initialized"); +} + +TEST(InferenceSessionTests, RunWithInvalidOutputNameReturnsError) { + SessionOptions so; + so.session_logid = "InferenceSessionTests.RunWithInvalidOutputNameReturnsError"; + InferenceSession session_object{so, GetEnvironment()}; + ASSERT_STATUS_OK(session_object.Load(MODEL_URI)); + ASSERT_STATUS_OK(session_object.Initialize()); + + std::vector dims_x = {3, 2}; + std::vector values_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + OrtValue ml_value; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], dims_x, values_x, &ml_value); + NameMLValMap feeds; + feeds.insert(std::make_pair("X", ml_value)); + std::vector output_names{"not_a_real_output"}; + std::vector fetches; + + RunOptions run_options; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(session_object.Run(run_options, feeds, output_names, &fetches), + "Invalid output name"); +} + +TEST(InferenceSessionTests, LoadMalformedModelFromArrayReturnsError) { + SessionOptions so; + so.session_logid = "InferenceSessionTests.LoadMalformedModelFromArrayReturnsError"; + InferenceSession session_object{so, GetEnvironment()}; + + // Bytes that are not a valid ModelProto: Load must fail gracefully (return an error, not crash). + const std::string garbage = "this is definitely not a valid onnx model proto"; + ASSERT_FALSE(session_object.Load(garbage.data(), static_cast(garbage.size())).IsOK()); +} + +TEST(InferenceSessionTests, RunWithWrongInputTypeReturnsError) { + SessionOptions so; + so.session_logid = "InferenceSessionTests.RunWithWrongInputTypeReturnsError"; + InferenceSession session_object{so, GetEnvironment()}; + ASSERT_STATUS_OK(session_object.Load(MODEL_URI)); + ASSERT_STATUS_OK(session_object.Initialize()); + + // Model input "X" is float; feed int32 to trigger the element-type check (CheckTypes). + std::vector dims_x = {3, 2}; + std::vector values_x = {1, 2, 3, 4, 5, 6}; + OrtValue ml_value; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], dims_x, values_x, &ml_value); + NameMLValMap feeds; + feeds.insert(std::make_pair("X", ml_value)); + std::vector output_names{"Y"}; + std::vector fetches; + + RunOptions run_options; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(session_object.Run(run_options, feeds, output_names, &fetches), + "Unexpected input data type"); +} + +TEST(InferenceSessionTests, RunWithWrongInputRankReturnsError) { + SessionOptions so; + so.session_logid = "InferenceSessionTests.RunWithWrongInputRankReturnsError"; + InferenceSession session_object{so, GetEnvironment()}; + ASSERT_STATUS_OK(session_object.Load(MODEL_URI)); + ASSERT_STATUS_OK(session_object.Initialize()); + + // Model input "X" has rank 2; feed a rank-1 tensor to trigger the rank check (CheckShapes). + std::vector bad_dims = {6}; + std::vector values_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + OrtValue ml_value; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], bad_dims, values_x, &ml_value); + NameMLValMap feeds; + feeds.insert(std::make_pair("X", ml_value)); + std::vector output_names{"Y"}; + std::vector fetches; + + RunOptions run_options; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(session_object.Run(run_options, feeds, output_names, &fetches), + "Invalid rank for input"); +} + +TEST(InferenceSessionTests, RunWithMismatchedFeedCountReturnsError) { + SessionOptions so; + so.session_logid = "InferenceSessionTests.RunWithMismatchedFeedCountReturnsError"; + InferenceSession session_object{so, GetEnvironment()}; + ASSERT_STATUS_OK(session_object.Load(MODEL_URI)); + ASSERT_STATUS_OK(session_object.Initialize()); + + std::vector dims_x = {3, 2}; + std::vector values_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + OrtValue ml_value; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], dims_x, values_x, &ml_value); + + // Two feed names but only one feed value -> count-mismatch branch in ValidateInputsOutputs. + std::vector feed_names{"X", "X"}; + std::vector feeds{ml_value}; + std::vector output_names{"Y"}; + std::vector fetches; + + RunOptions run_options; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(session_object.Run(run_options, feed_names, feeds, output_names, &fetches), + "feed names has"); +} + +TEST(InferenceSessionTests, InitializeBeforeLoadReturnsError) { + SessionOptions so; + so.session_logid = "InferenceSessionTests.InitializeBeforeLoadReturnsError"; + InferenceSession session_object{so, GetEnvironment()}; + // Initialize() without a prior successful Load(). + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(session_object.Initialize(), "Model was not loaded"); +} + +TEST(InferenceSessionTests, LoadNonexistentModelReturnsError) { + SessionOptions so; + so.session_logid = "InferenceSessionTests.LoadNonexistentModelReturnsError"; + InferenceSession session_object{so, GetEnvironment()}; + // Loading a path that does not exist must fail gracefully (return an error, not crash). + ASSERT_FALSE(session_object.Load(ORT_TSTR("testdata/this_model_does_not_exist.onnx")).IsOK()); +} + +TEST(InferenceSessionTests, RunWithWrongOutputTypeReturnsError) { + SessionOptions so; + so.session_logid = "InferenceSessionTests.RunWithWrongOutputTypeReturnsError"; + InferenceSession session_object{so, GetEnvironment()}; + ASSERT_STATUS_OK(session_object.Load(MODEL_URI)); + ASSERT_STATUS_OK(session_object.Initialize()); + + std::vector dims_x = {3, 2}; + std::vector values_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; + OrtValue x_value; + CreateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], dims_x, values_x, &x_value); + NameMLValMap feeds; + feeds.insert(std::make_pair("X", x_value)); + + // Pre-allocate the fetch for float output "Y" as int32 to trigger the output type check. + std::vector output_names{"Y"}; + std::vector fetches; + fetches.resize(1); + AllocateMLValue(TestCPUExecutionProvider()->CreatePreferredAllocators()[0], dims_x, &fetches[0]); + + RunOptions run_options; + ASSERT_STATUS_NOT_OK_AND_HAS_SUBSTR(session_object.Run(run_options, feeds, output_names, &fetches), + "Unexpected output data type"); +} + TEST(InferenceSessionTests, CheckRunLogger) { if constexpr (!SessionOptions::DEFAULT_USE_PER_SESSION_THREADS) { GTEST_SKIP() << "Skipping the test"; From 73260b2f6192f4816004c9bf97f66090f1a814c6 Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Thu, 9 Jul 2026 11:44:46 -0700 Subject: [PATCH 09/11] Fix JSEP pooling output shape to honor ceil_mode (#29627) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary The ORT-Web JSEP pooling output-shape helper (`PoolConvUtil.computePoolOutputShape` → `computeShapeHelper` → `adjustPadAndReturnShape`) was **floor-only**: it had no `ceilMode` parameter and silently ignored `ceil_mode`. Under `ceil_mode=1` this allocated the output tensor one element too small along each affected spatial axis, producing a wrong output shape (and downstream shape mismatches) — independent of any pooling divisor concern. The same floor-only helper is duplicated in the legacy `js/web/lib/onnxjs/util.ts`; both copies are fixed. ### Fix - Thread a `ceilMode` parameter (default `0`, i.e. floor) through `computePoolOutputShape` → `computeShapeHelper` → `adjustPadAndReturnShape` in **both** `js/web/lib/wasm/jsep/util.ts` and `js/web/lib/onnxjs/util.ts`. - Route the NOTSET / VALID / SAME_UPPER / SAME_LOWER branches through a new `computeOutputSize()` helper that produces results identical to the C++ reference `PoolAttributes::ComputeOutputSize` (`onnxruntime/core/providers/cpu/nn/pool_attributes.h`), including the `ceil_mode` **"shrink the last window if it starts entirely in the trailing padding"** rule (ref: onnx/onnx#5741). - The floor path (`ceilMode` default) is algebraically unchanged, so Conv and `auto_pad` pad-adjustment are unaffected. - `js/web/lib/wasm/jsep/webgpu/ops/pool.ts` now passes `attributes.ceilMode` into the shape computation. ### Scope: SHAPE-only This PR fixes the **output-shape** computation only. End-to-end `ceil_mode` execution remains gated by the existing `throw` guards in `parseAveragePoolAttributes` / `parseMaxPoolAttributes`, because the WebGPU pooling kernel does not yet implement `ceil_mode` trailing-padding handling (and, for AveragePool, the `count_include_pad` divisor). Removing those throws + adding kernel support is a **tracked follow-up**; the now-correct shape path is exercised directly by the added unit tests until then. Comments at the throw sites and in the test header document this. ### Tests Adds `js/web/test/unittests/pool-output-shape.ts` (registered in `unittests/index.ts`), asserting output shapes for **both** implementations (jsep + onnxjs) against the C++ CPU reference ground truth: - `test_maxpool_2d_ceil` → `[1,1,2,2]`, `AveragePool_10_ceil1_2d` → `[1,1,2,3]` - AvgPool `ceil` 1D `[1,2,4]` / 2D `[1,1,3,3]` / 3D `[1,1,2,2,2]` (matching the CPU `pool_op_test.cc` cases) - `ceil_mode` with dilation, `VALID` / `SAME_UPPER` / `SAME_LOWER` auto_pad - a shrink-rule discriminator (naive `ceil()` would give 3; correct = 2) - floor-mode no-regression case ### Related - PyTorch context: pytorch/pytorch#183528 - Related ORT PR: #16752 ### Post-review fixups - **SAME_UPPER/SAME_LOWER integer division (external-review Major).** The `legacyTargetSize = (inSize + stride - 1) / stride` computation now uses `Math.floor(...)` to match C++ `pool_attributes.h` `ComputeSizePadDilations`, which uses integer division. This fixes a latent float-division divergence that mis-rounded the auto_pad pad distribution. The correction applies to **all** ceil modes (not only `ceil_mode=1`), aligning JSEP/onnxjs with the CPU/CUDA EPs. The floor-path (`ceil_mode=0`) **output shape is unchanged** — only the SAME_* pad split is corrected — so nothing previously-correct regresses. Added `ceil_mode=0` SAME_UPPER/SAME_LOWER non-divisible regression tests that assert the corrected pad out-param (SAME_UPPER → `[0,1]`, SAME_LOWER → `[1,0]`). - **Throw-message / TODO clarity.** The retained `ceil_mode` `throw` statements and the pool op TODO banner were reworded to make explicit that the output **shape** is now computed correctly, while `ceil_mode` **kernel execution** (padding/divisor) remains the pending WebGPU follow-up. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- js/web/lib/onnxjs/util.ts | 59 ++++- js/web/lib/wasm/jsep/util.ts | 56 ++++- js/web/lib/wasm/jsep/webgpu/ops/pool.ts | 25 +- js/web/test/unittests/index.ts | 2 + js/web/test/unittests/pool-output-shape.ts | 278 +++++++++++++++++++++ 5 files changed, 407 insertions(+), 13 deletions(-) create mode 100644 js/web/test/unittests/pool-output-shape.ts diff --git a/js/web/lib/onnxjs/util.ts b/js/web/lib/onnxjs/util.ts index a2e048680a6ef..0ce035528bdbb 100644 --- a/js/web/lib/onnxjs/util.ts +++ b/js/web/lib/onnxjs/util.ts @@ -1260,6 +1260,8 @@ export class PoolConvUtil { * @param pads Padding for the beginning and ending along each axis. * @param autoPad DEPRECATED attribute supported for legacy models. Specifies how to implicitly calculate pads in each * dimension. Can take values NOTSET, SAME_UPPER, SAME_LOWER, or VALID. + * @param ceilMode When set to 1, use ceil() instead of floor() to compute the output spatial size (and apply the + * "shrink the last window if it starts entirely in padding" rule). Defaults to 0 (floor). */ static computePoolOutputShape( isGlobalOperator: boolean, @@ -1269,6 +1271,7 @@ export class PoolConvUtil { kernelShape: number[], pads: number[], autoPad?: string, + ceilMode = 0, ): number[] { if (inputDims.length <= 0) { throw new Error('input shape must be of size greater than 0'); @@ -1286,6 +1289,7 @@ export class PoolConvUtil { kernelShape, pads, autoPad, + ceilMode, ); return outputDims; } @@ -1332,6 +1336,7 @@ export class PoolConvUtil { kernelShape: readonly number[], pads: number[], autoPad?: string, + ceilMode = 0, ) { if (isGlobalOperator) { for (let dim = 0; dim < inputDims.length - 2; dim++) { @@ -1349,12 +1354,44 @@ export class PoolConvUtil { dim, dim + inputDims.length - 2, autoPad, + ceilMode, ), ); } } } + // Computes the output spatial size for a single dimension. + // Produces results identical to the C++ PoolAttributes::ComputeOutputSize + // (onnxruntime/core/providers/cpu/nn/pool_attributes.h), including the ceil_mode + // "shrink the last window if it starts entirely in the trailing padding" rule. The JS + // signature takes a pre-computed `numerator` (equal to `inSize + padHead + padTail - dkernel`, + // matching the C++ `in_size + pad_head + pad_tail - dilation * (kernel - 1) - 1`) instead of + // the raw pooling attributes, but the computed output size is the same. + // Keep in sync with the onnxjs/jsep copy. + // NOTE: In this onnxjs copy the ceilMode path exists for shape-test parity with the jsep copy; + // the onnxjs WebGL pooling caller (backends/webgl/ops/pool.ts) intentionally does NOT pass + // ceilMode (legacy path still throws on ceil_mode != 0), so it always uses the floor default. + private static computeOutputSize( + numerator: number, + stride: number, + inSize: number, + padHead: number, + ceilMode: number, + ): number { + let outSize = Math.floor(numerator / stride) + 1; + // Match C++ `ceil_mode == 1` exactly so out-of-spec ceil_mode values do not diverge. + if (ceilMode === 1) { + outSize = Math.ceil(numerator / stride) + 1; + // Ensure the last pooling window starts inside the image (ref: https://github.com/onnx/onnx/pull/5741). + // inSize and padHead are needed here to reconstruct the last window's start position. + if ((outSize - 1) * stride >= inSize + padHead) { + outSize -= 1; + } + } + return outSize; + } + // helper for computeShapeHelper() and adjustPadsBasedOnAutoPad() // adjusts pad value for given 'autoPad' string and computes output shape along a particular dimension private static adjustPadAndReturnShape( @@ -1366,6 +1403,7 @@ export class PoolConvUtil { padHeadIndex: number, padTailIndex: number, autoPad?: string, + ceilMode = 0, ): number { const dkernel = dilation * (kernel - 1) + 1; if (autoPad && autoPad !== 'NOTSET') { @@ -1373,23 +1411,36 @@ export class PoolConvUtil { case 'VALID': pads[padHeadIndex] = 0; pads[padTailIndex] = 0; - return Math.floor((inSize - dkernel) / stride + 1); + return PoolConvUtil.computeOutputSize(inSize - dkernel, stride, inSize, 0, ceilMode); case 'SAME_LOWER': case 'SAME_UPPER': if (dilation !== 1) { throw new Error('Dilation not supported for SAME_UPPER or SAME_LOWER'); } else { - const legacyTargetSize = (inSize + stride - 1) / stride; + // Integer division to match C++ pool_attributes.h ComputeSizePadDilations; float division mis-rounds SAME_* pads. + const legacyTargetSize = Math.floor((inSize + stride - 1) / stride); const padNeeded = (legacyTargetSize - 1) * stride + kernel - inSize; pads[padHeadIndex] = autoPad === 'SAME_LOWER' ? Math.floor((padNeeded + 1) / 2) : Math.floor(padNeeded / 2); pads[padTailIndex] = padNeeded - pads[padHeadIndex]; - return Math.floor((inSize + padNeeded - kernel) / stride + 1); + return PoolConvUtil.computeOutputSize( + inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel, + stride, + inSize, + pads[padHeadIndex], + ceilMode, + ); } default: throw new Error('Unsupported AutoPad type'); } } else { - return Math.floor((inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel) / stride + 1); + return PoolConvUtil.computeOutputSize( + inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel, + stride, + inSize, + pads[padHeadIndex], + ceilMode, + ); } } } diff --git a/js/web/lib/wasm/jsep/util.ts b/js/web/lib/wasm/jsep/util.ts index 27a718ac2fcd6..38ed670413cfa 100644 --- a/js/web/lib/wasm/jsep/util.ts +++ b/js/web/lib/wasm/jsep/util.ts @@ -366,6 +366,8 @@ export class PoolConvUtil { * @param pads Padding for the beginning and ending along each axis. * @param autoPad DEPRECATED attribute supported for legacy models. Specifies how to implicitly calculate pads in each * dimension. Can take values NOTSET, SAME_UPPER, SAME_LOWER, or VALID. + * @param ceilMode When set to 1, use ceil() instead of floor() to compute the output spatial size (and apply the + * "shrink the last window if it starts entirely in padding" rule). Defaults to 0 (floor). */ static computePoolOutputShape( isGlobalOperator: boolean, @@ -375,6 +377,7 @@ export class PoolConvUtil { kernelShape: number[], pads: number[], autoPad?: string, + ceilMode = 0, ): number[] { if (inputDims.length <= 0) { throw new Error('input shape must be of size greater than 0'); @@ -392,6 +395,7 @@ export class PoolConvUtil { kernelShape, pads, autoPad, + ceilMode, ); return outputDims; } @@ -438,6 +442,7 @@ export class PoolConvUtil { kernelShape: readonly number[], pads: number[], autoPad?: string, + ceilMode = 0, ) { if (isGlobalOperator) { for (let dim = 0; dim < inputDims.length - 2; dim++) { @@ -455,12 +460,41 @@ export class PoolConvUtil { dim, dim + inputDims.length - 2, autoPad, + ceilMode, ), ); } } } + // Computes the output spatial size for a single dimension. + // Produces results identical to the C++ PoolAttributes::ComputeOutputSize + // (onnxruntime/core/providers/cpu/nn/pool_attributes.h), including the ceil_mode + // "shrink the last window if it starts entirely in the trailing padding" rule. The JS + // signature takes a pre-computed `numerator` (equal to `inSize + padHead + padTail - dkernel`, + // matching the C++ `in_size + pad_head + pad_tail - dilation * (kernel - 1) - 1`) instead of + // the raw pooling attributes, but the computed output size is the same. + // Keep in sync with the onnxjs/jsep copy. + private static computeOutputSize( + numerator: number, + stride: number, + inSize: number, + padHead: number, + ceilMode: number, + ): number { + let outSize = Math.floor(numerator / stride) + 1; + // Match C++ `ceil_mode == 1` exactly so out-of-spec ceil_mode values do not diverge. + if (ceilMode === 1) { + outSize = Math.ceil(numerator / stride) + 1; + // Ensure the last pooling window starts inside the image (ref: https://github.com/onnx/onnx/pull/5741). + // inSize and padHead are needed here to reconstruct the last window's start position. + if ((outSize - 1) * stride >= inSize + padHead) { + outSize -= 1; + } + } + return outSize; + } + // helper for computeShapeHelper() and adjustPadsBasedOnAutoPad() // adjusts pad value for given 'autoPad' string and computes output shape along a particular dimension private static adjustPadAndReturnShape( @@ -472,6 +506,7 @@ export class PoolConvUtil { padHeadIndex: number, padTailIndex: number, autoPad?: string, + ceilMode = 0, ): number { const dkernel = dilation * (kernel - 1) + 1; if (autoPad && autoPad !== 'NOTSET') { @@ -479,23 +514,36 @@ export class PoolConvUtil { case 'VALID': pads[padHeadIndex] = 0; pads[padTailIndex] = 0; - return Math.floor((inSize - dkernel) / stride + 1); + return PoolConvUtil.computeOutputSize(inSize - dkernel, stride, inSize, 0, ceilMode); case 'SAME_LOWER': case 'SAME_UPPER': if (dilation !== 1) { throw new Error('Dilation not supported for SAME_UPPER or SAME_LOWER'); } else { - const legacyTargetSize = (inSize + stride - 1) / stride; + // Integer division to match C++ pool_attributes.h ComputeSizePadDilations; float division mis-rounds SAME_* pads. + const legacyTargetSize = Math.floor((inSize + stride - 1) / stride); const padNeeded = (legacyTargetSize - 1) * stride + kernel - inSize; pads[padHeadIndex] = autoPad === 'SAME_LOWER' ? Math.floor((padNeeded + 1) / 2) : Math.floor(padNeeded / 2); pads[padTailIndex] = padNeeded - pads[padHeadIndex]; - return Math.floor((inSize + padNeeded - kernel) / stride + 1); + return PoolConvUtil.computeOutputSize( + inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel, + stride, + inSize, + pads[padHeadIndex], + ceilMode, + ); } default: throw new Error('Unsupported AutoPad type'); } } else { - return Math.floor((inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel) / stride + 1); + return PoolConvUtil.computeOutputSize( + inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel, + stride, + inSize, + pads[padHeadIndex], + ceilMode, + ); } } } diff --git a/js/web/lib/wasm/jsep/webgpu/ops/pool.ts b/js/web/lib/wasm/jsep/webgpu/ops/pool.ts index 8b2438e45d6b4..de6c08dfa31af 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/pool.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/pool.ts @@ -20,7 +20,8 @@ import { } from './common'; // TODO: support: -// - ceil_mode "test_maxpool_2d_ceil" +// - ceil_mode kernel execution "test_maxpool_2d_ceil" (output SHAPE already honors ceil_mode +// via PoolConvUtil.computePoolOutputShape; the WebGPU kernel padding/divisor handling is pending) // - storage_order "test_maxpool_with_argmax_2d_precomputed_strides" // - [MaxPool] dilations "test_maxpool_2d_dilations" // - [MaxPool] output[1] "test_maxpool_with_argmax_2d_precomputed_pads" @@ -56,6 +57,7 @@ const getAdjustedPoolAttributesAndOutputShape = ): const countIncludePad = (attributes.count_include_pad as number) === 0 ? false : true; const attr = parsePoolCommonAttributes(attributes); - // TODO: support attribute 'ceil_mode' + // ceil_mode is honored by the output-shape math (PoolConvUtil.computePoolOutputShape now + // threads ceilMode and matches the C++ reference), but end-to-end execution stays gated + // here: the WebGPU pooling kernel must also apply the ceil_mode trailing-padding handling + // (and, for AveragePool, the count_include_pad divisor) before this throw can be removed. + // Tracked follow-up: remove this guard + add kernel ceil_mode padding support. if (attr.ceilMode !== 0) { - throw new Error('using ceil() in shape computation is not yet supported for AveragePool'); + throw new Error( + 'ceil_mode output-shape is computed, but ceil_mode kernel execution (padding/divisor) is not yet implemented in the WebGPU AveragePool kernel', + ); } const averagePoolAttributes = { countIncludePad, ...attr, cacheKey: '' }; return { ...averagePoolAttributes, cacheKey: createAveragePoolShaderKeyFromAttributes(averagePoolAttributes) }; @@ -491,12 +499,19 @@ export const parseMaxPoolAttributes = (attributes: Record): Max const dilations = attributes.dilations as [number, number]; const attr = parsePoolCommonAttributes(attributes); - // TODO: support attribute 'ceil_mode' and 'storage_order' + // TODO: support attribute 'storage_order' if (storageOrder !== 0) { throw new Error('column major storage order is not yet supported for MaxPool'); } + // ceil_mode is honored by the output-shape math (PoolConvUtil.computePoolOutputShape now + // threads ceilMode and matches the C++ reference), but end-to-end execution stays gated + // here: the WebGPU pooling kernel must also apply the ceil_mode trailing-padding handling + // before this throw can be removed. Tracked follow-up: remove this guard + add kernel + // ceil_mode padding support. if (attr.ceilMode !== 0) { - throw new Error('using ceil() in shape computation is not yet supported for MaxPool'); + throw new Error( + 'ceil_mode output-shape is computed, but ceil_mode kernel execution (padding) is not yet implemented in the WebGPU MaxPool kernel', + ); } const maxPoolAttributes = { storageOrder, dilations, ...attr, cacheKey: '' }; return { ...maxPoolAttributes, cacheKey: createMaxPoolShaderKeyFromAttributes(maxPoolAttributes) }; diff --git a/js/web/test/unittests/index.ts b/js/web/test/unittests/index.ts index b68681f4977c0..0f1349dda4795 100644 --- a/js/web/test/unittests/index.ts +++ b/js/web/test/unittests/index.ts @@ -13,4 +13,6 @@ if (typeof window !== 'undefined') { require('./backends/wasm/test-model-metadata'); +require('./pool-output-shape'); + require('./opset'); diff --git a/js/web/test/unittests/pool-output-shape.ts b/js/web/test/unittests/pool-output-shape.ts new file mode 100644 index 0000000000000..89ddf5ae91fc1 --- /dev/null +++ b/js/web/test/unittests/pool-output-shape.ts @@ -0,0 +1,278 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// These tests target PoolConvUtil.computePoolOutputShape directly (the pure shape function), +// not the end-to-end pooling op. That is intentional: while the shape math now honors +// ceil_mode, execution is still gated by the ceil_mode throw in parseAveragePool/MaxPoolAttributes +// (js/web/lib/wasm/jsep/webgpu/ops/pool.ts), because the WebGPU kernel does not yet implement +// ceil_mode trailing-padding handling. Removing that throw + adding kernel support is a tracked +// follow-up; until then the shape path is only reachable via this unit test. + +import { expect } from 'chai'; + +import { PoolConvUtil as JsepPoolConvUtil } from '../../lib/wasm/jsep/util'; +import { PoolConvUtil as OnnxjsPoolConvUtil } from '../../lib/onnxjs/util'; + +// Ground-truth output shapes are taken from the C++ CPU reference +// (onnxruntime/core/providers/cpu/nn/pool_attributes.h ComputeOutputSize) and the +// matching CPU pool_op_test.cc cases, so all pooling targets share one oracle and +// cannot silently re-diverge on ceil_mode. +interface PoolShapeCase { + readonly name: string; + readonly isGlobalOperator: boolean; + readonly inputDims: number[]; + readonly strides: number[]; + readonly dilations: number[]; + readonly kernelShape: number[]; + readonly pads: number[]; + readonly autoPad?: string; + readonly ceilMode: number; + readonly expectedShape: number[]; + // When set, asserts the auto_pad pad distribution written back to the pads out-param. + readonly expectedPads?: number[]; +} + +const poolShapeCases: PoolShapeCase[] = [ + // floor mode must stay identical to the previous behavior (no regression). + { + name: 'NOTSET floor mode is unchanged (test_maxpool_2d_default analogue)', + isGlobalOperator: false, + inputDims: [1, 1, 4, 4], + strides: [2, 2], + dilations: [1, 1], + kernelShape: [3, 3], + pads: [0, 0, 0, 0], + ceilMode: 0, + expectedShape: [1, 1, 1, 1], + }, + // test_maxpool_2d_ceil: floor would give [1,1] but ceil_mode gives [2,2]. + { + name: 'test_maxpool_2d_ceil -> [1,1,2,2]', + isGlobalOperator: false, + inputDims: [1, 1, 4, 4], + strides: [2, 2], + dilations: [1, 1], + kernelShape: [3, 3], + pads: [0, 0, 0, 0], + ceilMode: 1, + expectedShape: [1, 1, 2, 2], + }, + // AveragePool_10_ceil1_2d: asymmetric strides [3,1] -> [2,3]. + { + name: 'AveragePool_10_ceil1_2d -> [1,1,2,3]', + isGlobalOperator: false, + inputDims: [1, 1, 4, 4], + strides: [3, 1], + dilations: [1, 1], + kernelShape: [2, 2], + pads: [0, 0, 0, 0], + ceilMode: 1, + expectedShape: [1, 1, 2, 3], + }, + // AveragePool_18/19_ceil_count_include_pad_1d: PyTorch #183528 repro shape. + { + name: 'AveragePool ceil 1d (pads 3,3 kernel 7 stride 3) -> [1,2,4]', + isGlobalOperator: false, + inputDims: [1, 2, 9], + strides: [3], + dilations: [1], + kernelShape: [7], + pads: [3, 3], + ceilMode: 1, + expectedShape: [1, 2, 4], + }, + // AveragePool_18_ceil_count_include_pad_2d. + { + name: 'AveragePool ceil 2d (pads 1 kernel 3 stride 2) -> [1,1,3,3]', + isGlobalOperator: false, + inputDims: [1, 1, 4, 4], + strides: [2, 2], + dilations: [1, 1], + kernelShape: [3, 3], + pads: [1, 1, 1, 1], + ceilMode: 1, + expectedShape: [1, 1, 3, 3], + }, + // AveragePool_18_ceil_count_include_pad_3d. + { + name: 'AveragePool ceil 3d (pads 1 kernel 3 stride 2) -> [1,1,2,2,2]', + isGlobalOperator: false, + inputDims: [1, 1, 3, 3, 3], + strides: [2, 2, 2], + dilations: [1, 1, 1], + kernelShape: [3, 3, 3], + pads: [1, 1, 1, 1, 1, 1], + ceilMode: 1, + expectedShape: [1, 1, 2, 2, 2], + }, + // Shrink-rule discriminator: naive ceil() would give 3, but the last window would start + // entirely in the trailing padding, so C++ ComputeOutputSize shrinks it back to 2. + { + name: 'shrink rule: last window starting in padding is dropped -> [1,1,2]', + isGlobalOperator: false, + inputDims: [1, 1, 3], + strides: [2], + dilations: [1], + kernelShape: [2], + pads: [1, 1], + ceilMode: 1, + expectedShape: [1, 1, 2], + }, + // ceil_mode with dilation (AveragePool_19_dilation_2d shape). + { + name: 'ceil 2d with dilation 2 -> [1,1,2,2]', + isGlobalOperator: false, + inputDims: [1, 1, 4, 4], + strides: [1, 1], + dilations: [2, 2], + kernelShape: [2, 2], + pads: [0, 0, 0, 0], + ceilMode: 1, + expectedShape: [1, 1, 2, 2], + }, + // VALID auto_pad honors ceil_mode. + { + name: 'VALID auto_pad + ceil_mode -> [1,1,2,2]', + isGlobalOperator: false, + inputDims: [1, 1, 4, 4], + strides: [2, 2], + dilations: [1, 1], + kernelShape: [3, 3], + pads: [0, 0, 0, 0], + autoPad: 'VALID', + ceilMode: 1, + expectedShape: [1, 1, 2, 2], + }, + // SAME_UPPER auto_pad honors ceil_mode (exercises the recomputed-pad branch). + // Reviewer-verified: in=5, stride=2, kernel=3, SAME_UPPER, ceil -> 3. + { + name: 'SAME_UPPER auto_pad + ceil_mode -> [1,1,3]', + isGlobalOperator: false, + inputDims: [1, 1, 5], + strides: [2], + dilations: [1], + kernelShape: [3], + pads: [0, 0], + autoPad: 'SAME_UPPER', + ceilMode: 1, + expectedShape: [1, 1, 3], + }, + // SAME_LOWER auto_pad honors ceil_mode (same numeric case, mirrored padding split). + { + name: 'SAME_LOWER auto_pad + ceil_mode -> [1,1,3]', + isGlobalOperator: false, + inputDims: [1, 1, 5], + strides: [2], + dilations: [1], + kernelShape: [3], + pads: [0, 0], + autoPad: 'SAME_LOWER', + ceilMode: 1, + expectedShape: [1, 1, 3], + }, + // Non-divisible SAME_UPPER + ceil_mode regression guard: legacyTargetSize must use C++ + // integer division. With float division this yielded [1,1,2] instead of the correct [1,1,1]. + // (in=2, stride=2, kernel=3, SAME_UPPER, ceil -> 1, matching C++ ComputeSizePadDilations.) + { + name: 'SAME_UPPER auto_pad + ceil_mode non-divisible -> [1,1,1]', + isGlobalOperator: false, + inputDims: [1, 1, 2], + strides: [2], + dilations: [1], + kernelShape: [3], + pads: [0, 0], + autoPad: 'SAME_UPPER', + ceilMode: 1, + expectedShape: [1, 1, 1], + }, + // SAME_LOWER symmetric counterpart of the non-divisible regression guard. + { + name: 'SAME_LOWER auto_pad + ceil_mode non-divisible -> [1,1,1]', + isGlobalOperator: false, + inputDims: [1, 1, 2], + strides: [2], + dilations: [1], + kernelShape: [3], + pads: [0, 0], + autoPad: 'SAME_LOWER', + ceilMode: 1, + expectedShape: [1, 1, 1], + }, + // ceilMode=0 SAME_* regression guards: the Math.floor(legacyTargetSize) fix also corrects the + // auto_pad pad DISTRIBUTION on the floor path (output shape is unchanged, so only asserting the + // shape would miss it). Old float division gave pads (1,1); C++ integer division gives (0,1) for + // SAME_UPPER and (1,0) for SAME_LOWER at in=2, stride=2, kernel=3. + { + name: 'SAME_UPPER auto_pad ceilMode=0 non-divisible corrects pad split -> pads [0,1]', + isGlobalOperator: false, + inputDims: [1, 1, 2], + strides: [2], + dilations: [1], + kernelShape: [3], + pads: [0, 0], + autoPad: 'SAME_UPPER', + ceilMode: 0, + expectedShape: [1, 1, 1], + expectedPads: [0, 1], + }, + { + name: 'SAME_LOWER auto_pad ceilMode=0 non-divisible corrects pad split -> pads [1,0]', + isGlobalOperator: false, + inputDims: [1, 1, 2], + strides: [2], + dilations: [1], + kernelShape: [3], + pads: [0, 0], + autoPad: 'SAME_LOWER', + ceilMode: 0, + expectedShape: [1, 1, 1], + expectedPads: [1, 0], + }, +]; + +function runPoolConvUtil( + computePoolOutputShape: typeof JsepPoolConvUtil.computePoolOutputShape, + testCase: PoolShapeCase, +): { shape: number[]; pads: number[] } { + // strides/dilations/kernelShape/pads are copied because computePoolOutputShape mutates pads + // in the auto_pad branches. The mutated pads copy is returned so tests can assert the + // computed auto_pad pad distribution, not just the output shape. + const pads = testCase.pads.slice(); + const shape = computePoolOutputShape( + testCase.isGlobalOperator, + testCase.inputDims, + testCase.strides.slice(), + testCase.dilations.slice(), + testCase.kernelShape.slice(), + pads, + testCase.autoPad, + testCase.ceilMode, + ); + return { shape, pads }; +} + +describe('PoolConvUtil.computePoolOutputShape ceil_mode', () => { + for (const testCase of poolShapeCases) { + it(`[jsep] ${testCase.name}`, () => { + const { shape, pads } = runPoolConvUtil(JsepPoolConvUtil.computePoolOutputShape, testCase); + expect(shape).to.deep.equal(testCase.expectedShape); + if (testCase.expectedPads) { + expect(pads).to.deep.equal(testCase.expectedPads); + } + }); + it(`[onnxjs] ${testCase.name}`, () => { + const { shape, pads } = runPoolConvUtil(OnnxjsPoolConvUtil.computePoolOutputShape, testCase); + expect(shape).to.deep.equal(testCase.expectedShape); + if (testCase.expectedPads) { + expect(pads).to.deep.equal(testCase.expectedPads); + } + }); + } + + it('defaults to floor when ceilMode is omitted', () => { + const jsep = JsepPoolConvUtil.computePoolOutputShape(false, [1, 1, 4, 4], [2, 2], [1, 1], [3, 3], [0, 0, 0, 0]); + const onnxjs = OnnxjsPoolConvUtil.computePoolOutputShape(false, [1, 1, 4, 4], [2, 2], [1, 1], [3, 3], [0, 0, 0, 0]); + expect(jsep).to.deep.equal([1, 1, 1, 1]); + expect(onnxjs).to.deep.equal([1, 1, 1, 1]); + }); +}); From 930b764dcd40abd7c8cb704db7439bb46109bf30 Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Thu, 9 Jul 2026 13:34:31 -0700 Subject: [PATCH 10/11] Hoist CropAndResize index offset to a named int64_t (#29623) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Description Follow-up to the merged #29605, addressing a review nit from @apsonawane on the base-offset index arithmetic in `CropAndResizeForward`. The suggestion was that the `static_cast(...)` around the per-channel base offset looked redundant. It turns out the cast could **not** simply be removed: the offset expression is a `SafeInt`, and `pointer + SafeInt` is ambiguous — `SafeInt` exposes many implicit conversion operators (`char`, `short`, `int`, …), so the operand for pointer arithmetic cannot be resolved unambiguously (and `SafeInt`'s own `operator+(U, SafeInt)` tries to treat the pointer as an integer). The cast was therefore load-bearing. To honor the readability intent without the awkward inline cast, this change hoists the offset into a named `int64_t` local in both the bilinear and nearest branches: ```cpp const int64_t bottom_data_offset = (SafeInt(roi_batch_ind) * channels + c) * height * width; const T* offset_bottom_data = bottom_data + bottom_data_offset; ``` Initializing an `int64_t` directly from the `SafeInt` expression is unambiguous (it selects `operator int64_t()`), so this compiles cleanly, keeps the `SafeInt` overflow checking on the offset computation, and reads more clearly than the inline cast. **No behavior change.** ### Motivation and Context Improves readability of the index arithmetic introduced in #29605 while keeping the overflow-checked offset computation intact. ### Tests Covered by the existing `CropAndResizeTest` suite — all 10 tests pass: ``` onnxruntime_provider_test --gtest_filter='CropAndResizeTest.*' ``` Signed-off-by: Tita Wang Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cpu/crop_and_resize.cc | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/crop_and_resize.cc b/onnxruntime/contrib_ops/cpu/crop_and_resize.cc index 47f164fde8322..5da7d5e3c016a 100644 --- a/onnxruntime/contrib_ops/cpu/crop_and_resize.cc +++ b/onnxruntime/contrib_ops/cpu/crop_and_resize.cc @@ -149,9 +149,8 @@ void CropAndResizeForward(const TensorShape& output_shape, for (auto c = 0; c < channels; c++) { int64_t index_n_c = index_n + c * pooled_width * pooled_height; int64_t index = index_n_c + ph * pooled_width + pw; - const T* offset_bottom_data = - bottom_data + - static_cast((SafeInt(roi_batch_ind) * channels + c) * height * width); + const int64_t bottom_data_offset = (SafeInt(roi_batch_ind) * channels + c) * height * width; + const T* offset_bottom_data = bottom_data + bottom_data_offset; const float top_left(static_cast(offset_bottom_data[top_left_index])); const float top_right(static_cast(offset_bottom_data[top_right_index])); const float bottom_left(static_cast(offset_bottom_data[bottom_left_index])); @@ -169,9 +168,8 @@ void CropAndResizeForward(const TensorShape& output_shape, for (auto c = 0; c < channels; c++) { int64_t index_n_c = index_n + c * pooled_width * pooled_height; int64_t index = index_n_c + ph * pooled_width + pw; - const T* offset_bottom_data = - bottom_data + - static_cast((SafeInt(roi_batch_ind) * channels + c) * height * width); + const int64_t bottom_data_offset = (SafeInt(roi_batch_ind) * channels + c) * height * width; + const T* offset_bottom_data = bottom_data + bottom_data_offset; top_data[index] = static_cast(offset_bottom_data[closest_index]); } } From f4aa2b4407a4f4d6dfb0707e7cb320690d4d707b Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Thu, 9 Jul 2026 14:53:58 -0700 Subject: [PATCH 11/11] Validate initializer data length and guard element-count computation in contrib Range shape inference (#29265) ### Summary The `com.microsoft` Range operator's shape-inference helper read a fixed number of bytes from an initializer's `raw_data` without first checking the buffer length, and the element-count computation could pass non-finite or out-of-range values to an `int64` cast. This change adds the missing validations and aligns the shape-inference and CPU kernel paths. ### Changes - `GetFirstElement` now checks `raw_data` length is at least the element size before reading, and reads via `std::memcpy` into an aligned local. - `CalcRangeDim` and the CPU kernel `ComputeRange` now reject non-finite computed counts, handle non-positive counts before the `int64` cast, and reject counts that are not representable as `int64` (`>= 2^63`). Both paths use identical messages and semantics. - The output dimension for empty/backward ranges is clamped to 0 in shape inference to match the kernel. ### Tests Added contrib Range model-load regression tests in `range_test.cc` covering: - truncated `raw_data` for `start`, `limit`, and `delta` (double, plus float and int64 element types), - zero delta, - a finite-but-too-large element count, - an exact-size success boundary, - backward-range zero-dimension inference. Tests assert on `Status` (safe for no-exception builds) and are guarded by `#ifndef DISABLE_CONTRIB_OPS` (throwing cases additionally by `!defined(ORT_NO_EXCEPTIONS)`). 21/21 `RangeTest` cases pass locally. ### Follow-up `int16` inputs are currently supported only via `raw_data` (there is no `get_data` specialization for the non-raw-data path); this is left as a separate follow-up. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../graph/contrib_ops/range_schema_defs.cc | 35 +- .../core/providers/cpu/generator/range.cc | 28 +- .../core/providers/cuda/generator/range.cc | 29 +- .../providers/cuda/generator/range_impl.cu | 27 +- .../providers/cuda/generator/range_impl.h | 2 +- .../providers/cpu/generator/range_test.cc | 339 ++++++++++++++++++ 6 files changed, 437 insertions(+), 23 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc b/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc index c2d853af86723..19e39e48a330d 100644 --- a/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc @@ -7,6 +7,7 @@ #include "core/graph/constants.h" #include "core/graph/op.h" #include +#include #include namespace onnxruntime { @@ -42,6 +43,13 @@ int64_t get_data(const TensorProto* shapeInitializer) { fail_shape_inference("Can not get shape initializer data!"); } +template <> +int16_t get_data(const TensorProto* shapeInitializer) { + // ONNX packs INT16 initializer values into the int32_data repeated field. + if (shapeInitializer->int32_data_size() > 0) return static_cast(shapeInitializer->int32_data(0)); + fail_shape_inference("Can not get shape initializer data!"); +} + template <> float get_data(const TensorProto* shapeInitializer) { if (shapeInitializer->float_data_size() > 0) return shapeInitializer->float_data(0); @@ -60,7 +68,13 @@ static T GetFirstElement(const TensorProto* shapeInitializer) { if (utils::HasRawData(*shapeInitializer)) { const std::string& bytes = shapeInitializer->raw_data(); - return *reinterpret_cast(bytes.c_str()); + if (bytes.size() < sizeof(T)) { + fail_shape_inference("Range: raw_data size is smaller than the element size for the given data type."); + } + // std::string data is only char-aligned, so copy the bytes into a properly aligned value. + T value; + std::memcpy(&value, bytes.data(), sizeof(T)); + return value; } return get_data(shapeInitializer); } @@ -75,7 +89,24 @@ static int64_t CalcRangeDim(const TensorProto* startShapeInitializer, if (delta == 0) { fail_shape_inference("delta in Range operator can not be zero!"); } - return static_cast(ceil((1.0 * (limit - start)) / delta)); + // Mirror the CPU kernel (core/providers/cpu/generator/range.cc ComputeRange) exactly so the + // two paths stay byte-consistent: clamp empty or backward ranges to 0 and promote the + // operands to double before the subtraction. Keep this expression identical to the kernel. + double count = ceil((static_cast(limit) - static_cast(start)) / static_cast(delta)); + if (!std::isfinite(count)) { + fail_shape_inference("Range: the computed number of elements is not a finite value."); + } + // Empty or backward ranges clamp to 0; handle the non-positive case before the cast so a + // large-magnitude negative count can never reach (and overflow) the int64 conversion below. + if (count <= 0) { + return 0; + } + // static_cast(INT64_MAX) rounds up to 2^63 (9223372036854775808.0), which is not + // representable as int64_t, so reject any count at or above that boundary before the cast. + if (count >= 9223372036854775808.0) { + fail_shape_inference("Range: the computed number of elements exceeds the supported range."); + } + return static_cast(count); } static int64_t CalcResultDim(const TensorProto* startShapeInitializer, diff --git a/onnxruntime/core/providers/cpu/generator/range.cc b/onnxruntime/core/providers/cpu/generator/range.cc index fcd4abcf72f87..f1f42e10f5541 100644 --- a/onnxruntime/core/providers/cpu/generator/range.cc +++ b/onnxruntime/core/providers/cpu/generator/range.cc @@ -6,10 +6,7 @@ #include #include "core/providers/op_kernel_type_control.h" -// TODO: fix the warnings -#if defined(_MSC_VER) && !defined(__clang__) -#pragma warning(disable : 26451) -#endif + namespace onnxruntime { namespace op_kernel_type_control { @@ -80,9 +77,26 @@ static Status ComputeRange( if (delta == T{0}) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "delta in Range operator can not be zero!"); } - int64_t n = static_cast(ceil((1.0 * (limit - start)) / delta)); - if (n <= 0) - n = 0; + // Compute the element count in double, mirroring the shape-inference path + // (core/graph/contrib_ops/range_schema_defs.cc CalcRangeDim) exactly. The operands are + // promoted to double before the subtraction so integral inputs cannot overflow in T. + double count = ceil((static_cast(limit) - static_cast(start)) / static_cast(delta)); + if (!std::isfinite(count)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Range: the computed number of elements is not a finite value."); + } + // Empty or backward ranges clamp to 0; handle the non-positive case before the cast so a + // large-magnitude negative count can never reach (and overflow) the int64 conversion. + int64_t n = 0; + if (count > 0) { + // static_cast(INT64_MAX) rounds up to 2^63 (9223372036854775808.0), which is not + // representable as int64_t, so reject any count at or above that boundary before the cast. + if (count >= 9223372036854775808.0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Range: the computed number of elements exceeds the supported range."); + } + n = static_cast(count); + } TensorShape shape = {n}; T* y = ctx->Output(0, shape)->MutableData(); for (int64_t i = 0; i < n; ++i) { diff --git a/onnxruntime/core/providers/cuda/generator/range.cc b/onnxruntime/core/providers/cuda/generator/range.cc index f9952df6ece42..9d1e33c728412 100644 --- a/onnxruntime/core/providers/cuda/generator/range.cc +++ b/onnxruntime/core/providers/cuda/generator/range.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "core/providers/shared_library/provider_api.h" #include "core/providers/cuda/cuda_common.h" #include "range.h" @@ -87,11 +89,28 @@ static Status ComputeRange(cudaStream_t stream, OpKernelContext* ctx) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "delta in Range operator can not be zero!"); } - double num = (static_cast(limit) - static_cast(start)) / static_cast(delta); - int count = static_cast(ceil(num)); - if (count <= 0) - count = 0; - TensorShape shape = {static_cast(count)}; + // Compute the element count in double, mirroring the CPU kernel's guard structure and error + // messages (core/providers/cpu/generator/range.cc ComputeRange) and the shape-inference path + // (core/graph/contrib_ops/range_schema_defs.cc CalcRangeDim). The operands are + // promoted to double before the subtraction so integral inputs cannot overflow in T. + double num = ceil((static_cast(limit) - static_cast(start)) / static_cast(delta)); + if (!std::isfinite(num)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Range: the computed number of elements is not a finite value."); + } + // Empty or backward ranges clamp to 0; handle the non-positive case before the cast so a + // large-magnitude negative count can never reach (and overflow) the int64 conversion. + int64_t count = 0; + if (num > 0) { + // static_cast(INT64_MAX) rounds up to 2^63 (9223372036854775808.0), which is not + // representable as int64_t, so reject any count at or above that boundary before the cast. + if (num >= 9223372036854775808.0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Range: the computed number of elements exceeds the supported range."); + } + count = static_cast(num); + } + TensorShape shape = {count}; T* y = ctx->Output(0, shape)->MutableData(); if (count > 0) { diff --git a/onnxruntime/core/providers/cuda/generator/range_impl.cu b/onnxruntime/core/providers/cuda/generator/range_impl.cu index 5cc7522a16c2a..a278b40ea8f60 100644 --- a/onnxruntime/core/providers/cuda/generator/range_impl.cu +++ b/onnxruntime/core/providers/cuda/generator/range_impl.cu @@ -13,23 +13,34 @@ namespace onnxruntime { namespace cuda { template -__global__ void RangeKernel(const T start, const T delta, const int count, T* output) { - int index = blockIdx.x * blockDim.x + threadIdx.x; - if (index < count) { - output[index] = start + delta * index; +__global__ void RangeKernel(const T start, const T delta, const int64_t count, T* output) { + // Use a 64-bit grid-stride loop so counts larger than the launch grid (and larger than + // INT_MAX) are handled correctly without index truncation or grid-dimension overflow. + int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = index; i < count; i += stride) { + output[i] = start + delta * static_cast(i); } } template -Status RangeImpl(cudaStream_t stream, const T start, const T delta, const int count, T* output) { - constexpr int block_size = 256; - int grid_size = (count + block_size - 1) / block_size; +Status RangeImpl(cudaStream_t stream, const T start, const T delta, const int64_t count, T* output) { + if (count <= 0) { + return Status::OK(); + } + constexpr int block_size = GridDim::maxThreadsPerBlock; + int64_t num_blocks = (count + block_size - 1) / block_size; + // CUDA limits the x-dimension of the launch grid to 2^31 - 1 blocks. Cap the grid to that + // maximum; the grid-stride loop in RangeKernel covers any remaining elements when the count + // requires more blocks than can be launched at once. + constexpr int64_t kMaxGridDimX = 2147483647; + int grid_size = static_cast(num_blocks < kMaxGridDimX ? num_blocks : kMaxGridDimX); RangeKernel<<>>(start, delta, count, output); return CUDA_CALL(cudaGetLastError()); } #define SPECIALIZED_IMPL(T) \ - template Status RangeImpl(cudaStream_t stream, const T start, const T delta, const int count, T* output); + template Status RangeImpl(cudaStream_t stream, const T start, const T delta, const int64_t count, T* output); SPECIALIZED_IMPL(int16_t) SPECIALIZED_IMPL(int32_t) diff --git a/onnxruntime/core/providers/cuda/generator/range_impl.h b/onnxruntime/core/providers/cuda/generator/range_impl.h index 0dd2973d03385..be95a3077a921 100644 --- a/onnxruntime/core/providers/cuda/generator/range_impl.h +++ b/onnxruntime/core/providers/cuda/generator/range_impl.h @@ -8,7 +8,7 @@ namespace onnxruntime { namespace cuda { template -Status RangeImpl(cudaStream_t stream, const T start, const T delta, const int count, T* output); +Status RangeImpl(cudaStream_t stream, const T start, const T delta, const int64_t count, T* output); } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/generator/range_test.cc b/onnxruntime/test/providers/cpu/generator/range_test.cc index 04664628ecdd1..8560c06db81a1 100644 --- a/onnxruntime/test/providers/cpu/generator/range_test.cc +++ b/onnxruntime/test/providers/cpu/generator/range_test.cc @@ -1,8 +1,14 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + +#include "core/graph/constants.h" +#include "core/session/inference_session.h" #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "test/test_environment.h" +#include "default_providers.h" namespace onnxruntime { namespace test { @@ -88,5 +94,338 @@ TEST(RangeTest, AlmostSameStartAndLimitHighDelta) { RunTest(2.0f, 2.01f, 1000000.0f, {1}, {2.0f}); } +#ifndef DISABLE_CONTRIB_OPS + +namespace { +// Describes a single Range input supplied as a graph initializer: its declared dimensions and +// either the exact raw_data bytes or a typed int32_data payload. Tests use this to craft +// well-formed, truncated (e.g. dims=[0] with empty raw_data), and typed-field initializers for +// the contrib Range schema. When int32_data is non-empty it is stored in the typed field instead +// of raw_data (ONNX packs sub-32-bit int types such as INT16 into int32_data). +struct RangeInputSpec { + std::vector dims; + std::string raw_data; + std::vector int32_data; +}; + +// Serializes a scalar value to its raw_data byte representation. +template +std::string ToRawData(T value) { + std::string bytes(sizeof(T), '\0'); + std::memcpy(bytes.data(), &value, sizeof(T)); + return bytes; +} + +// A correctly-sized scalar initializer (no dims) holding a single value. +template +RangeInputSpec ScalarInput(T value) { + return RangeInputSpec{{}, ToRawData(value), {}}; +} + +// A zero-element initializer (dims=[0]) whose raw_data is empty. The declared size matches +// the (empty) payload, so initializer size validation accepts it; shape inference, however, +// still attempts to read the first element. +RangeInputSpec EmptyInput() { + return RangeInputSpec{{0}, std::string{}, {}}; +} + +// A scalar initializer whose single value is carried in the typed int32_data field. ONNX packs +// INT16 (and other sub-32-bit int) initializer values into int32_data, so this exercises the +// typed-field shape-inference path rather than the raw_data path. +RangeInputSpec Int16TypedInput(int16_t value) { + return RangeInputSpec{{}, std::string{}, {static_cast(value)}}; +} + +void AddInitializer(ONNX_NAMESPACE::GraphProto* graph, const std::string& name, int data_type, + const RangeInputSpec& spec) { + auto* initializer = graph->add_initializer(); + initializer->set_name(name); + initializer->set_data_type(data_type); + for (int64_t dim : spec.dims) { + initializer->add_dims(dim); + } + if (!spec.int32_data.empty()) { + for (int32_t value : spec.int32_data) { + initializer->add_int32_data(value); + } + } else { + initializer->set_raw_data(spec.raw_data); + } +} + +// Builds a single com.microsoft Range node model whose start/limit/delta inputs are supplied +// as graph initializers of the given element type, then loads and initializes it in a session. +common::Status BuildAndInitializeContribRangeModel(int data_type, + const RangeInputSpec& start, + const RangeInputSpec& limit, + const RangeInputSpec& delta, + InferenceSession& session_object) { + ONNX_NAMESPACE::ModelProto model; + model.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); + model.add_opset_import()->set_version(11); + auto* ms_opset = model.add_opset_import(); + ms_opset->set_domain(kMSDomain); + ms_opset->set_version(1); + + auto* graph = model.mutable_graph(); + graph->set_name("ContribRangeGraph"); + + AddInitializer(graph, "start", data_type, start); + AddInitializer(graph, "limit", data_type, limit); + AddInitializer(graph, "delta", data_type, delta); + + auto* output = graph->add_output(); + output->set_name("Y"); + output->mutable_type()->mutable_tensor_type()->set_elem_type(data_type); + + auto* range_node = graph->add_node(); + range_node->set_name("Range"); + range_node->set_op_type("Range"); + range_node->set_domain(kMSDomain); + range_node->add_input("start"); + range_node->add_input("limit"); + range_node->add_input("delta"); + range_node->add_output("Y"); + + std::string serialized_model; + if (!model.SerializeToString(&serialized_model)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to serialize test model."); + } + + std::stringstream model_stream(serialized_model); + ORT_RETURN_IF_ERROR(session_object.Load(model_stream)); + return session_object.Initialize(); +} + +// Returns the inferred dim_value of output Y after a successful Initialize, or -1 if absent. +int64_t GetInferredOutputDim(const InferenceSession& session_object) { + const auto outputs = session_object.GetModelOutputs(); + if (!outputs.first.IsOK() || outputs.second->size() != 1u) { + return -1; + } + const auto* shape = (*outputs.second)[0]->Shape(); + if (shape == nullptr || shape->dim_size() != 1 || !shape->dim(0).has_dim_value()) { + return -1; + } + return shape->dim(0).dim_value(); +} +} // namespace + +// Verifies that shape inference clamps an empty/backward range to a zero-sized dimension, +// matching the CPU kernel behavior, instead of emitting a negative dimension value. +TEST(RangeTest, ContribOp_BackwardRange_InfersZeroDim) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_BackwardRange_InfersZeroDim"; + InferenceSession session_object{so, GetEnvironment()}; + // start > limit with a positive delta yields an empty range. + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(5.0), ScalarInput(0.0), + ScalarInput(1.0), session_object); + ASSERT_STATUS_OK(status); + ASSERT_EQ(GetInferredOutputDim(session_object), 0); +} + +// Verifies that a correctly-sized raw_data initializer (exactly sizeof(T) bytes) loads +// successfully and produces the expected inferred dimension, so the length check does not +// over-reject valid models. +TEST(RangeTest, ContribOp_ExactSizeRawData_LoadsSuccessfully) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_ExactSizeRawData_LoadsSuccessfully"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(5.0), + ScalarInput(1.0), session_object); + ASSERT_STATUS_OK(status); + ASSERT_EQ(GetInferredOutputDim(session_object), 5); +} + +// Verifies that shape inference handles an INT16 Range whose initializers store their values in +// the typed int32_data field (the ONNX packing for INT16) rather than raw_data. This exercises +// the int16 typed-field path in get_data: start=0, limit=5, delta=1 -> 5 elements. +TEST(RangeTest, ContribOp_Int16TypedField_InfersDim) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_Int16TypedField_InfersDim"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_INT16, Int16TypedInput(0), Int16TypedInput(5), + Int16TypedInput(1), session_object); + ASSERT_STATUS_OK(status); + ASSERT_EQ(GetInferredOutputDim(session_object), 5); +} + +// The following two tests run the contrib Range CPU kernel with runtime (non-constant) inputs. +// Because the inputs are not initializers, shape inference cannot fold the output length, so +// the element-count guards in ComputeRange are exercised at execution time rather than at load. +// They are pinned to the CPU execution provider because the contrib Range kernel is CPU-only. + +// Verifies that the kernel rejects an element count that exceeds the int64 range. +TEST(RangeTest, ContribOp_Kernel_CountExceedsInt64Range_FailsAtRuntime) { + OpTester test("Range", 1, kMSDomain); + test.AddInput("start", {}, {0.0}); + test.AddInput("limit", {}, {1e19}); + test.AddInput("delta", {}, {1.0}); + test.AddOutput("output", {0}, {}); + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, + "Range: the computed number of elements exceeds the supported range.", {}, nullptr, + &execution_providers); +} + +// Verifies that the kernel rejects a non-finite computed element count (the difference of the +// two inputs overflows to infinity). +TEST(RangeTest, ContribOp_Kernel_NonFiniteCount_FailsAtRuntime) { + OpTester test("Range", 1, kMSDomain); + test.AddInput("start", {}, {-1e308}); + test.AddInput("limit", {}, {1e308}); + test.AddInput("delta", {}, {1.0}); + test.AddOutput("output", {0}, {}); + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, + "Range: the computed number of elements is not a finite value.", {}, nullptr, + &execution_providers); +} + +// The following tests exercise failure paths that surface via fail_shape_inference, which +// reports the error by throwing an inference error. They are excluded from no-exception +// builds where such a throw would abort rather than yield a failure Status. +#if !defined(ORT_NO_EXCEPTIONS) + +// Verifies that Range shape inference rejects an initializer whose raw_data holds fewer bytes +// than its declared data type requires (here: a zero-element initializer with empty raw_data, +// which initializer size validation accepts), returning a clean failure status instead of +// reading more bytes than the initializer declares. One test per input position. +TEST(RangeTest, ContribOp_TruncatedStartRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedStartRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, EmptyInput(), ScalarInput(5.0), + ScalarInput(1.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + +TEST(RangeTest, ContribOp_TruncatedLimitRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedLimitRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), EmptyInput(), + ScalarInput(1.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + +TEST(RangeTest, ContribOp_TruncatedDeltaRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedDeltaRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(5.0), + EmptyInput(), session_object); + ASSERT_FALSE(status.IsOK()); +} + +// Exercises the length check for different sizeof(T) template instantiations (float, int64). +TEST(RangeTest, ContribOp_TruncatedFloatRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedFloatRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT, EmptyInput(), ScalarInput(5.0f), + ScalarInput(1.0f), session_object); + ASSERT_FALSE(status.IsOK()); +} + +TEST(RangeTest, ContribOp_TruncatedInt64RawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedInt64RawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_INT64, EmptyInput(), ScalarInput(5), + ScalarInput(1), session_object); + ASSERT_FALSE(status.IsOK()); +} + +// Verifies that a zero delta is rejected cleanly by shape inference. +TEST(RangeTest, ContribOp_ZeroDelta_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_ZeroDelta_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(10.0), + ScalarInput(0.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + +// Verifies that a finite element count that exceeds the int64 range is rejected cleanly, +// rather than reaching an out-of-range conversion (start=0, limit=1e19, delta=1). +TEST(RangeTest, ContribOp_CountExceedsInt64Range_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_CountExceedsInt64Range_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(1e19), + ScalarInput(1.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + +#endif // !defined(ORT_NO_EXCEPTIONS) + +#endif // DISABLE_CONTRIB_OPS + +#ifdef USE_CUDA +// The following tests exercise the standard onnx-domain Range CUDA kernel with runtime +// (non-constant) inputs so the element-count guards in the CUDA ComputeRange are evaluated at +// execution time. They are pinned to the CUDA execution provider and are skipped when a CUDA +// provider is not available in the current build/runtime. The expected failure messages match +// the CPU kernel exactly so both paths stay consistent. + +// Verifies that the CUDA kernel rejects an element count that exceeds the int64 range. +TEST(RangeTest, CudaKernel_CountExceedsInt64Range_FailsAtRuntime) { + auto cuda_provider = DefaultCudaExecutionProvider(); + if (cuda_provider == nullptr) { + GTEST_SKIP() << "CUDA execution provider is not available."; + } + OpTester test("Range", 11); + test.AddInput("start", {}, {0.0}); + test.AddInput("limit", {}, {1e19}); + test.AddInput("delta", {}, {1.0}); + test.AddOutput("output", {0}, {}); + std::vector> execution_providers; + execution_providers.push_back(std::move(cuda_provider)); + test.Run(OpTester::ExpectResult::kExpectFailure, + "Range: the computed number of elements exceeds the supported range.", {}, nullptr, + &execution_providers); +} + +// Verifies that the CUDA kernel rejects a non-finite computed element count (the difference of +// the two inputs overflows to infinity). +TEST(RangeTest, CudaKernel_NonFiniteCount_FailsAtRuntime) { + auto cuda_provider = DefaultCudaExecutionProvider(); + if (cuda_provider == nullptr) { + GTEST_SKIP() << "CUDA execution provider is not available."; + } + OpTester test("Range", 11); + test.AddInput("start", {}, {-1e308}); + test.AddInput("limit", {}, {1e308}); + test.AddInput("delta", {}, {1.0}); + test.AddOutput("output", {0}, {}); + std::vector> execution_providers; + execution_providers.push_back(std::move(cuda_provider)); + test.Run(OpTester::ExpectResult::kExpectFailure, + "Range: the computed number of elements is not a finite value.", {}, nullptr, + &execution_providers); +} + +// Note on large-count coverage: a test that actually materializes a valid count > INT_MAX to +// prove the int64 launch-path widening is intentionally omitted. OpTester requires a host-side +// reference output tensor of the same element count, and a value above INT_MAX would need a +// multi-GB host allocation (and matching device memory), which is impractical and OOM-prone on +// CI. The int64 widening (count, grid sizing, and the grid-stride kernel index) is instead +// covered by code review; the two guards above ensure non-finite and >= 2^63 counts are rejected +// before any launch. A device-level large-count test can be added as a separate, opt-in benchmark. +#endif // USE_CUDA + } // namespace test } // namespace onnxruntime