From f79df447ebe067ca49799da7cf680ab0868ec7ff Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Fri, 10 Jul 2026 10:48:00 -0700 Subject: [PATCH 1/2] Fix CUDA AveragePool wrong results with asymmetric padding and dilation (#29631) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary Fixes wrong `AveragePool` output on the **CUDA** EP whenever cuDNN's pooling descriptor cannot represent the requested pooling — i.e. **asymmetric padding** or **dilation > 1**. This is the CUDA-EP counterpart of the CPU fix in #29629 and, together with it, the JSEP shape fix in #29627, closes out the CUDA leg of pytorch/pytorch#183528. ### Root cause `CudnnPoolingDescriptor::Set` copies only the **begin** pads (`pads[0..rank)`) into the cuDNN pooling descriptor, which stores a *single symmetric pad value per axis* and applies it to both sides. The ONNX **end** pads (`pads[rank..2*rank)`) are silently dropped. As a result, **all** asymmetric-pad AveragePool on CUDA is wrong: - explicit asymmetric `pads`, - `auto_pad = SAME_UPPER` / `SAME_LOWER` (which produce naturally asymmetric pads), - `ceil_mode = 1` boundary windows. Separately, the cuDNN pooling descriptor has **no dilation parameter at all**, so any `dilations > 1` is also silently ignored — even with symmetric pads. Example divergence from the CPU reference (QA probe): 1D `pad(0,3)`, k7, s3, `ceil_mode=1`, `count_include_pad=1` gave CUDA `[4, 6.5, 8]` vs. correct `[4, 5.571, 4]`; 2D `pad(0,0,3,3)` gave `71.0` vs. correct `17.75`. ### Fix Add a custom CUDA average-pool kernel (`avg_pool_impl.cu` / `.h`, modeled on `max_pool_with_index.cu`) that honors **per-side pads** and **dilation**, and computes the `count_include_pad` divisor exactly like the CPU v19 reference functor (`AveragePool{1,2,3}DTask`): - include-pad divisor clamps the window end to `input + pad_tail` (drops the ceil-mode phantom cells), dividing by `∏ (1 + (end - start - 1)/dilation)`; - exclude-pad divisor counts only in-bounds cells. In `Pool::ComputeInternal`, a cheap dispatch guard routes to the custom kernel **only** when the pooling is non-global **and** (`asymmetric pads` **or** `!default_dilations`). Every symmetric, non-dilated case — the overwhelmingly common path, including **all** `GlobalAveragePool` — stays on the existing cuDNN fast path, so there is **zero perf regression** on the common path. `GlobalAveragePool` is excluded explicitly because `PoolAttributes` leaves its `kernel_shape`/`strides`/`dilations` unpopulated. `MaxPool` is unaffected (it ignores pad cells; `MaxPool<8>` already routes dilation to its own custom kernel via the same `!default_dilations` check we mirror here). Covers fp32, fp64, fp16, and bf16 (fp16/bf16 accumulate in float). ### Tests Added CUDA-**un-excluded** parity tests in `pool_op_test.cc` — the CUDA leg actually runs and must match the CPU reference oracle: - asymmetric tail-pad 1D/2D, include- and exclude-pad; - `auto_pad=SAME_UPPER` and `SAME_LOWER` (naturally asymmetric); - **symmetric-pad + dilation>1** (wrong on cuDNN, correct via the kernel — locks in the dilation guard); - a symmetric-pad regression case (must stay on cuDNN and remain correct); - fp16; - a `MaxPool` asymmetric-pad case confirming MaxPool is unaffected. The `ceil_mode + count_include_pad` cases use opset 19 so the CPU leg runs the already-correct v19 reference functor and validates the CUDA kernel independently of the separate opset-7..18 CPU MLAS fix (#29629); CUDA routing is opset-independent. Full `PoolTest` suite passes on an A100 (53 passed / 2 DML-skipped / 0 failures). ### Related - CPU (a): #29629 — opset 7-18 MLAS `ceil_mode + count_include_pad` divisor fix. - JSEP (c): #29627 — JSEP pooling output shape honoring `ceil_mode`. - pytorch/pytorch#183528. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/providers/cuda/nn/avg_pool_impl.cu | 247 +++++++++++++++ .../core/providers/cuda/nn/avg_pool_impl.h | 39 +++ onnxruntime/core/providers/cuda/nn/pool.cc | 32 ++ .../test/providers/cpu/nn/pool_op_test.cc | 295 +++++++++++++++++- 4 files changed, 607 insertions(+), 6 deletions(-) create mode 100644 onnxruntime/core/providers/cuda/nn/avg_pool_impl.cu create mode 100644 onnxruntime/core/providers/cuda/nn/avg_pool_impl.h diff --git a/onnxruntime/core/providers/cuda/nn/avg_pool_impl.cu b/onnxruntime/core/providers/cuda/nn/avg_pool_impl.cu new file mode 100644 index 0000000000000..028f3f2bde9a4 --- /dev/null +++ b/onnxruntime/core/providers/cuda/nn/avg_pool_impl.cu @@ -0,0 +1,247 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "avg_pool_impl.h" + +#include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/shared_inc/fast_divmod.h" +#include "core/providers/cuda/shared_inc/cuda_utils.h" + +namespace onnxruntime { +namespace cuda { + +// Accumulate half in float for precision; keep native type otherwise. +template +struct AveragePoolAccumulator { + using type = T; +}; +template <> +struct AveragePoolAccumulator { + using type = float; +}; +template <> +struct AveragePoolAccumulator { + using type = float; +}; + +template +__global__ void AveragePoolWithPadKernel( + int64_t channels, + int64_t height, + int64_t width, + int64_t depth, + int64_t pooled_height, + int64_t pooled_width, + int64_t pooled_depth, + int64_t kernel_h, + int64_t kernel_w, + int64_t kernel_d, + int64_t stride_h, + int64_t stride_w, + int64_t stride_d, + int64_t pad_h_head, + int64_t pad_w_head, + int64_t pad_d_head, + int64_t pad_h_tail, + int64_t pad_w_tail, + int64_t pad_d_tail, + int64_t dilation_h, + int64_t dilation_w, + int64_t dilation_d, + fast_divmod fdm_c, + fast_divmod fdm_h, + fast_divmod fdm_w, + fast_divmod fdm_d, + bool count_include_pad, + const T* p_input, + int64_t output_size, + T* p_output) { + int id = blockIdx.x * blockDim.x + threadIdx.x; + if (id >= output_size) return; + + auto compute_offset = + [height, width, depth, channels](int n_index, int c_index, int h_index, int w_index, int d_index) -> int64_t { + if constexpr (Layout == LAYOUT_NCHW) { + return (((n_index * channels + c_index) * height + h_index) * width + w_index) * depth + d_index; + } else if constexpr (Layout == LAYOUT_NHWC) { + return (((n_index * height + h_index) * width + w_index) * depth + d_index) * channels + c_index; + } + }; + + int d_index, w_index, h_index, c_index, n_index, id_tmp; + if constexpr (Layout == LAYOUT_NCHW) { + fdm_d.divmod(id, id_tmp, d_index); + fdm_w.divmod(id_tmp, id_tmp, w_index); + fdm_h.divmod(id_tmp, id_tmp, h_index); + fdm_c.divmod(id_tmp, n_index, c_index); + } else if constexpr (Layout == LAYOUT_NHWC) { + fdm_c.divmod(id, id_tmp, c_index); + fdm_d.divmod(id_tmp, id_tmp, d_index); + fdm_w.divmod(id_tmp, id_tmp, w_index); + fdm_h.divmod(id_tmp, n_index, h_index); + } + + // Window bounds mirror the CPU AveragePool{1,2,3}DTask reference exactly. + int64_t h_start = h_index * stride_h - pad_h_head; + int64_t w_start = w_index * stride_w - pad_w_head; + int64_t d_start = d_index * stride_d - pad_d_head; + + int64_t h_end = _Min(h_start + kernel_h * dilation_h, height + pad_h_tail); + int64_t w_end = _Min(w_start + kernel_w * dilation_w, width + pad_w_tail); + int64_t d_end = _Min(d_start + kernel_d * dilation_d, depth + pad_d_tail); + + using AccT = typename AveragePoolAccumulator::type; + AccT acc = static_cast(0); + int64_t counted = 0; + + int64_t offset = compute_offset(n_index, c_index, 0, 0, 0); + const T* p_slice = p_input + offset; + for (int64_t h = h_start; h < h_end; h += dilation_h) { + if (h < 0 || h >= height) continue; + for (int64_t w = w_start; w < w_end; w += dilation_w) { + if (w < 0 || w >= width) continue; + for (int64_t d = d_start; d < d_end; d += dilation_d) { + if (d < 0 || d >= depth) continue; + acc += static_cast(p_slice[compute_offset(0, 0, h, w, d)]); + ++counted; + } + } + } + + AccT result = static_cast(0); + if (counted > 0) { + if (count_include_pad) { + int64_t divisor = (1 + (h_end - h_start - 1) / dilation_h) * + (1 + (w_end - w_start - 1) / dilation_w) * + (1 + (d_end - d_start - 1) / dilation_d); + result = acc / static_cast(divisor); + } else { + result = acc / static_cast(counted); + } + } + p_output[id] = static_cast(result); +} + +template +void AveragePoolWithPad( + cudaStream_t stream, + const TensorShape& input_shape, + const TensorShape& output_shape, + const gsl::span& kernel_shape, + const gsl::span& stride_shape, + const gsl::span& pads, + const gsl::span& dilations, + bool count_include_pad, + const T* p_input, + T* p_output) { + int64_t channels, height, width, depth; + int64_t pooled_height, pooled_width, pooled_depth; + if constexpr (Layout == LAYOUT_NCHW) { + channels = input_shape[1]; + height = input_shape[2]; + width = kernel_shape.size() > 1 ? input_shape[3] : 1; + depth = kernel_shape.size() > 2 ? input_shape[4] : 1; + + pooled_height = output_shape[2]; + pooled_width = kernel_shape.size() > 1 ? output_shape[3] : 1; + pooled_depth = kernel_shape.size() > 2 ? output_shape[4] : 1; + } else if constexpr (Layout == LAYOUT_NHWC) { + height = input_shape[1]; + width = kernel_shape.size() > 1 ? input_shape[2] : 1; + depth = kernel_shape.size() > 2 ? input_shape[3] : 1; + channels = input_shape[input_shape.NumDimensions() - 1]; + + pooled_height = output_shape[1]; + pooled_width = kernel_shape.size() > 1 ? output_shape[2] : 1; + pooled_depth = kernel_shape.size() > 2 ? output_shape[3] : 1; + } + + const int64_t rank = static_cast(kernel_shape.size()); + int64_t kernel_h = kernel_shape[0]; + int64_t kernel_w = rank > 1 ? kernel_shape[1] : 1; + int64_t kernel_d = rank > 2 ? kernel_shape[2] : 1; + int64_t stride_h = stride_shape[0]; + int64_t stride_w = rank > 1 ? stride_shape[1] : 1; + int64_t stride_d = rank > 2 ? stride_shape[2] : 1; + + // pads: [x1_begin,...,xN_begin, x1_end,...,xN_end]. Begin at [i], end at [rank + i]. + int64_t pad_h_head = pads[0]; + int64_t pad_w_head = rank > 1 ? pads[1] : 0; + int64_t pad_d_head = rank > 2 ? pads[2] : 0; + int64_t pad_h_tail = pads[rank + 0]; + int64_t pad_w_tail = rank > 1 ? pads[rank + 1] : 0; + int64_t pad_d_tail = rank > 2 ? pads[rank + 2] : 0; + + int64_t dilation_h = dilations[0]; + int64_t dilation_w = rank > 1 ? dilations[1] : 1; + int64_t dilation_d = rank > 2 ? dilations[2] : 1; + + int64_t output_size = output_shape.Size(); + if (output_size == 0) return; + + fast_divmod fdm_c(static_cast(channels)); + fast_divmod fdm_h(static_cast(pooled_height)); + fast_divmod fdm_w(static_cast(pooled_width)); + fast_divmod fdm_d(static_cast(pooled_depth)); + + int blocksPerGrid = (int)((output_size + GridDim::maxThreadsPerBlock - 1) / GridDim::maxThreadsPerBlock); + AveragePoolWithPadKernel<<>>( + channels, + height, + width, + depth, + pooled_height, + pooled_width, + pooled_depth, + kernel_h, + kernel_w, + kernel_d, + stride_h, + stride_w, + stride_d, + pad_h_head, + pad_w_head, + pad_d_head, + pad_h_tail, + pad_w_tail, + pad_d_tail, + dilation_h, + dilation_w, + dilation_d, + fdm_c, + fdm_h, + fdm_w, + fdm_d, + count_include_pad, + p_input, + output_size, + p_output); +} + +#define INSTANTIATE_AVERAGEPOOLWITHPAD(T, Layout) \ + template void AveragePoolWithPad( \ + cudaStream_t stream, \ + const TensorShape& input_shape, \ + const TensorShape& output_shape, \ + const gsl::span& kernel_shape, \ + const gsl::span& stride_shape, \ + const gsl::span& pads, \ + const gsl::span& dilations, \ + bool count_include_pad, \ + const T* p_input, \ + T* p_output); + +INSTANTIATE_AVERAGEPOOLWITHPAD(float, LAYOUT_NCHW) +INSTANTIATE_AVERAGEPOOLWITHPAD(double, LAYOUT_NCHW) +INSTANTIATE_AVERAGEPOOLWITHPAD(half, LAYOUT_NCHW) +INSTANTIATE_AVERAGEPOOLWITHPAD(BFloat16, LAYOUT_NCHW) + +#ifdef ENABLE_CUDA_NHWC_OPS +INSTANTIATE_AVERAGEPOOLWITHPAD(float, LAYOUT_NHWC) +INSTANTIATE_AVERAGEPOOLWITHPAD(double, LAYOUT_NHWC) +INSTANTIATE_AVERAGEPOOLWITHPAD(half, LAYOUT_NHWC) +INSTANTIATE_AVERAGEPOOLWITHPAD(BFloat16, LAYOUT_NHWC) +#endif + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/nn/avg_pool_impl.h b/onnxruntime/core/providers/cuda/nn/avg_pool_impl.h new file mode 100644 index 0000000000000..9cf51a5442f9f --- /dev/null +++ b/onnxruntime/core/providers/cuda/nn/avg_pool_impl.h @@ -0,0 +1,39 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/framework/tensor_shape.h" + +namespace onnxruntime { +namespace cuda { + +// Custom average-pooling CUDA kernel that honors per-side (asymmetric) padding. +// +// cuDNN's pooling descriptor stores a single symmetric pad value per axis, so it cannot +// represent ONNX asymmetric padding (pad_begin != pad_end), which arises from explicit +// asymmetric `pads` or `auto_pad = SAME_UPPER/SAME_LOWER` resolving to asymmetric pads. This +// kernel is the CUDA fallback for that case and mirrors the CPU reference functor +// (AveragePool{1,2,3}DTask) exactly: +// start = out_idx * stride - pad_begin +// end = min(start + kernel * dilation, in_size + pad_end) +// sum over cells [start, end) with dilation step that are in [0, in_size) +// count_include_pad == 1: divisor = product of (1 + (end - start - 1) / dilation) +// count_include_pad == 0: divisor = number of summed in-bounds cells +// +// `pads` is the full ONNX layout [x1_begin,...,xN_begin, x1_end,...,xN_end] (2 * rank). +template +void AveragePoolWithPad( + cudaStream_t stream, + const TensorShape& input_shape, + const TensorShape& output_shape, + const gsl::span& kernel_shape, + const gsl::span& stride_shape, + const gsl::span& pads, + const gsl::span& dilations, + bool count_include_pad, + const T* p_input, + T* p_output); + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/nn/pool.cc b/onnxruntime/core/providers/cuda/nn/pool.cc index 89c6047769cf0..b955e68268a2f 100644 --- a/onnxruntime/core/providers/cuda/nn/pool.cc +++ b/onnxruntime/core/providers/cuda/nn/pool.cc @@ -6,6 +6,7 @@ #include "core/providers/cuda/nn/pool.h" #include "core/providers/cuda/cudnn_common.h" #include "core/providers/cuda/nn/max_pool_with_index.h" +#include "core/providers/cuda/nn/avg_pool_impl.h" #include "core/providers/cuda/math/unary_elementwise_ops_impl.h" using namespace onnxruntime::common; @@ -205,6 +206,37 @@ Status Pool::ComputeInternal(OpKernelContext* context) cons auto x_data = reinterpret_cast(X->Data()); auto y_data = reinterpret_cast(Y->MutableData()); + // cuDNN's pooling descriptor cannot represent two ONNX features: + // (1) Asymmetric padding: it stores a single symmetric pad value per axis and applies it to + // both sides, so it silently drops ONNX end pads when pad_begin != pad_end (explicit + // asymmetric pads, or auto_pad=SAME_UPPER/SAME_LOWER resolving to asymmetric pads). + // (2) Dilation: the pooling descriptor has no dilation parameter at all, so any dilation > 1 + // is silently ignored. + // Either case produces wrong sums and divisors on cuDNN, so route it to the custom kernel, + // which honors per-side pads AND dilation and matches the CPU reference divisor exactly. + // Symmetric, non-dilated pooling (the common case, including all global pooling) keeps the + // fast cuDNN path unchanged, so there is zero perf regression. (MaxPool<8> guards dilation the + // same way via !default_dilations.) Global pooling is always symmetric and never dilated, and + // its kernel_shape/strides/dilations are left unpopulated (PoolAttributes returns early), so it + // is excluded here and stays on cuDNN. + if constexpr (PoolType::type == onnxruntime::PoolType::kAveragePool) { + if (!pool_attrs_.global_pooling) { + const size_t spatial_rank = kernel_shape.size(); + bool asymmetric_pads = false; + for (size_t i = 0; i < spatial_rank; ++i) { + if (pads[i] != pads[spatial_rank + i]) { + asymmetric_pads = true; + break; + } + } + if (asymmetric_pads || !pool_attrs_.default_dilations) { + AveragePoolWithPad(Stream(context), x_shape, y_shape, kernel_shape, strides, pads, + pool_attrs_.dilations, pool_attrs_.count_include_pad, x_data, y_data); + return Status::OK(); + } + } + } + TensorShapeVector x_dims_cudnn(x_dims.begin(), x_dims.end()); TensorShapeVector y_dims_cudnn(y_dims); if (kernel_shape.size() < 2) { diff --git a/onnxruntime/test/providers/cpu/nn/pool_op_test.cc b/onnxruntime/test/providers/cpu/nn/pool_op_test.cc index e7f870fc43d7f..95c0c9d35c7f0 100644 --- a/onnxruntime/test/providers/cpu/nn/pool_op_test.cc +++ b/onnxruntime/test/providers/cpu/nn/pool_op_test.cc @@ -959,10 +959,13 @@ TEST(PoolTest, AveragePool_CountIncludePad_AsymmetricPads) { test.AddInput("X", x_dims, x_vals); test.AddOutput("Y", expected_dims, expected_vals); - // This test targets the CPU fix only. Exclude EPs whose external libraries - // (cuDNN, CoreML, etc.) also produce wrong results for this case. + // The CUDA custom AveragePoolWithPad kernel now honors per-side (asymmetric) pads, so the + // CUDA (NCHW) leg is un-excluded here to lock in that fix. kCudaNHWCExecutionProvider is + // excluded here only to avoid redundant coverage: the asymmetric NHWC-CUDA decode branch is + // now exercised (and passing) by the 1D/2D AveragePool_CUDA_* parity tests below. The remaining + // exclusions are EPs whose external libraries (CoreML, etc.) still produce wrong results here. test.Run(OpTester::ExpectResult::kExpectSuccess, "", - {kCudaExecutionProvider, kCudaNHWCExecutionProvider, + {kCudaNHWCExecutionProvider, kTensorrtExecutionProvider, kAclExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider, kCoreMLExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider}); @@ -994,10 +997,14 @@ TEST(PoolTest, AveragePool3D_CountIncludePad_AsymmetricPads) { 0.5f, 0.25f}; test.AddInput("X", x3d_dims, x3d_vals); test.AddOutput("Y", expected3d_dims, expected3d_vals); - // This test targets the CPU fix only. Exclude EPs whose external libraries - // (cuDNN, CoreML, etc.) also produce wrong results for this case. + // The CUDA custom AveragePoolWithPad kernel now honors per-side (asymmetric) pads in 3D, so the + // CUDA (NCHW) leg is un-excluded here to lock in that fix. kCudaNHWCExecutionProvider stays + // excluded here: the asymmetric NHWC-CUDA decode branch is now exercised (and passing) by the + // 1D/2D AveragePool_CUDA_* parity tests below, but 3D (NDHWC) NHWC pooling is not among them, so + // it remains a follow-up. The remaining exclusions are EPs whose external libraries (CoreML, + // etc.) still produce wrong results here. test.Run(OpTester::ExpectResult::kExpectSuccess, "", - {kCudaExecutionProvider, kCudaNHWCExecutionProvider, + {kCudaNHWCExecutionProvider, kTensorrtExecutionProvider, kAclExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider, kCoreMLExecutionProvider, kQnnExecutionProvider, kDmlExecutionProvider}); @@ -1095,6 +1102,282 @@ TEST(PoolTest, AveragePool_19_ceil_count_include_pad_1d) { {kTensorrtExecutionProvider, kAclExecutionProvider, kOpenVINOExecutionProvider, kDmlExecutionProvider}); } +// --------------------------------------------------------------------------- +// CUDA AveragePool asymmetric-padding parity tests. +// +// cuDNN's pooling descriptor stores one symmetric pad per axis, so it silently drops the +// ONNX end pad when pad_begin != pad_end, producing wrong averages on CUDA (e.g. for a 1D +// pad of (0,3) cuDNN yields [4, 6.5, 8] while the CPU reference yields [4, 5.571, 4], and a +// 2D pad of (0,0,3,3) diverges by up to 53.25). The custom AveragePoolWithPad CUDA kernel +// fixes this. These cases keep the CUDA EP UN-excluded so the CUDA leg actually runs and must +// match the CPU reference oracle. Expected values are the CPU reference outputs. +// +// ceil_mode + count_include_pad cases use opset 19 so the CPU leg runs the already-correct v19 +// reference functor and validates the CUDA kernel independently of the separate CPU opset-7..18 +// MLAS fix (PR #29629); the CUDA routing is opset-independent, so this still exercises the fix. +// +// Exception: the fp16 case below runs CUDA-only (CPU has no fp16 AveragePool kernel on x64 and +// the Arm64 NEON fp16 pooling kernel mishandles the ceil_mode + count_include_pad divisor), so +// the "CPU leg runs the v19 reference oracle" statement above does not apply to it — see its own +// comment for how it validates the CUDA half accumulate-in-float path without a CPU oracle. +// +// The float cases share a single exclusion set (kPoolingEpsExcludedFromCeilCipTests) so the list cannot +// drift test-to-test. It names every EP whose pooling does NOT implement ONNX's asymmetric-pad / +// dilated / ceil_mode + count_include_pad clamped-divisor semantics (they would produce wrong +// values and cannot serve as an oracle). The CPU EP (correct v19 reference) stays un-excluded as +// the float oracle, and the CUDA + CUDA-NHWC EPs stay un-excluded as the tested targets — the +// asymmetric NHWC-CUDA path is intentionally exercised here (USE_CUDA_NHWC_OPS defaults ON) and +// passes. kWebGpuExecutionProvider is listed defensively: it auto-skips in a CUDA-only build +// (DefaultWebGpuExecutionProvider returns nullptr), but naming it keeps a future WebGPU build leg +// from re-triggering the CI failure this list fixes. +// --------------------------------------------------------------------------- +const std::unordered_set kPoolingEpsExcludedFromCeilCipTests = { + kTensorrtExecutionProvider, kNvTensorRTRTXExecutionProvider, kDnnlExecutionProvider, + kOpenVINOExecutionProvider, kAclExecutionProvider, kCoreMLExecutionProvider, + kQnnExecutionProvider, kDmlExecutionProvider, kWebGpuExecutionProvider}; + +TEST(PoolTest, AveragePool_CUDA_asymmetric_tail_pad_1d) { + OpTester test("AveragePool", 19); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{3}); + test.AddAttribute("pads", std::vector{0, 3}); + test.AddAttribute("kernel_shape", std::vector{7}); + test.AddAttribute("ceil_mode", (int64_t)1); + test.AddAttribute("count_include_pad", (int64_t)1); + + std::vector x_vals = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f}; + std::vector x_dims = {1, 1, 9}; + std::vector expected_dims = {1, 1, 3}; + std::vector expected_vals = {4.0f, 5.5714283f, 4.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCipTests); +} + +TEST(PoolTest, AveragePool_CUDA_asymmetric_tail_pad_1d_exclude_pad) { + OpTester test("AveragePool", 18); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{3}); + test.AddAttribute("pads", std::vector{0, 3}); + test.AddAttribute("kernel_shape", std::vector{7}); + test.AddAttribute("ceil_mode", (int64_t)1); + test.AddAttribute("count_include_pad", (int64_t)0); + + std::vector x_vals = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f}; + std::vector x_dims = {1, 1, 9}; + std::vector expected_dims = {1, 1, 3}; + // exclude-pad divides by in-bounds cells only. + std::vector expected_vals = {4.0f, 6.5f, 8.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCipTests); +} + +TEST(PoolTest, AveragePool_CUDA_asymmetric_tail_pad_2d) { + OpTester test("AveragePool", 19); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{3, 3}); + test.AddAttribute("pads", std::vector{0, 0, 3, 3}); + test.AddAttribute("kernel_shape", std::vector{7, 7}); + test.AddAttribute("ceil_mode", (int64_t)1); + test.AddAttribute("count_include_pad", (int64_t)1); + + std::vector x_vals(81); + for (int i = 0; i < 81; ++i) { + x_vals[i] = static_cast(i + 1); + } + std::vector x_dims = {1, 1, 9, 9}; + std::vector expected_dims = {1, 1, 3, 3}; + std::vector expected_vals = {31.0f, 28.714287f, 17.5f, + 45.857143f, 41.142857f, 24.642858f, + 33.5f, 29.785715f, 17.75f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCipTests); +} + +TEST(PoolTest, AveragePool_CUDA_asymmetric_tail_pad_2d_exclude_pad) { + OpTester test("AveragePool", 18); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{3, 3}); + test.AddAttribute("pads", std::vector{0, 0, 3, 3}); + test.AddAttribute("kernel_shape", std::vector{7, 7}); + test.AddAttribute("ceil_mode", (int64_t)1); + test.AddAttribute("count_include_pad", (int64_t)0); + + std::vector x_vals(81); + for (int i = 0; i < 81; ++i) { + x_vals[i] = static_cast(i + 1); + } + std::vector x_dims = {1, 1, 9, 9}; + std::vector expected_dims = {1, 1, 3, 3}; + std::vector expected_vals = {31.0f, 33.5f, 35.0f, + 53.5f, 56.0f, 57.5f, + 67.0f, 69.5f, 71.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCipTests); +} + +// auto_pad=SAME_UPPER produces naturally asymmetric pads (here pad(0,1)); proves the latent +// SAME-pad bug on CUDA is also fixed by the same kernel. +TEST(PoolTest, AveragePool_CUDA_same_upper_asymmetric_1d) { + OpTester test("AveragePool", 18); + + test.AddAttribute("auto_pad", "SAME_UPPER"); + test.AddAttribute("strides", std::vector{2}); + test.AddAttribute("kernel_shape", std::vector{3}); + test.AddAttribute("count_include_pad", (int64_t)1); + + std::vector x_vals = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f}; + std::vector x_dims = {1, 1, 10}; + std::vector expected_dims = {1, 1, 5}; + std::vector expected_vals = {2.0f, 4.0f, 6.0f, 8.0f, 6.3333335f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCipTests); +} + +// Regression guard: symmetric pads must STAY on the fast cuDNN path and remain correct. +TEST(PoolTest, AveragePool_CUDA_symmetric_pad_regression_1d) { + OpTester test("AveragePool", 19); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{3}); + test.AddAttribute("pads", std::vector{3, 3}); + test.AddAttribute("kernel_shape", std::vector{7}); + test.AddAttribute("ceil_mode", (int64_t)1); + test.AddAttribute("count_include_pad", (int64_t)1); + + std::vector x_vals = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f}; + std::vector x_dims = {1, 1, 9}; + std::vector expected_dims = {1, 1, 4}; + std::vector expected_vals = {1.4285715f, 4.0f, 5.5714283f, 4.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCipTests); +} + +// MaxPool asymmetric-pad probe/regression on CUDA (verify-then-decide per design). MaxPool +// ignores pad cells (no divisor); asymmetric tail pad only changes output size, computed +// correctly upstream. CUDA un-excluded to confirm parity with the CPU reference. +TEST(PoolTest, MaxPool_CUDA_asymmetric_tail_pad_1d) { + OpTester test("MaxPool", 12); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{3}); + test.AddAttribute("pads", std::vector{0, 3}); + test.AddAttribute("kernel_shape", std::vector{7}); + test.AddAttribute("ceil_mode", (int64_t)1); + + std::vector x_vals = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f}; + std::vector x_dims = {1, 1, 9}; + std::vector expected_dims = {1, 1, 3}; + std::vector expected_vals = {7.0f, 9.0f, 9.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCipTests); +} + +// fp16 asymmetric-pad AveragePool. Runs on the CUDA EP ONLY (via an explicit provider list): +// the CPU AveragePool has no fp16 kernel on x64, and the Arm64 NEON fp16 pooling kernel does +// not honor the ceil_mode + count_include_pad divisor rule (a separate, pre-existing CPU +// limitation), so it cannot serve as the fp16 oracle. This test validates that the CUDA +// AveragePoolWithPad kernel's half accumulate-in-float path matches the reference values. +TEST(PoolTest, AveragePool_CUDA_asymmetric_tail_pad_1d_fp16) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + return; + } + + OpTester test("AveragePool", 19); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{3}); + test.AddAttribute("pads", std::vector{0, 3}); + test.AddAttribute("kernel_shape", std::vector{7}); + test.AddAttribute("ceil_mode", (int64_t)1); + test.AddAttribute("count_include_pad", (int64_t)1); + + std::vector x_vals = {MLFloat16(1.0f), MLFloat16(2.0f), MLFloat16(3.0f), + MLFloat16(4.0f), MLFloat16(5.0f), MLFloat16(6.0f), + MLFloat16(7.0f), MLFloat16(8.0f), MLFloat16(9.0f)}; + std::vector x_dims = {1, 1, 9}; + std::vector expected_dims = {1, 1, 3}; + std::vector expected_vals = {MLFloat16(4.0f), MLFloat16(5.5714283f), MLFloat16(4.0f)}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.SetOutputTolerance(0.005f); + + std::vector> execution_providers; + execution_providers.push_back(std::move(cuda_ep)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Symmetric pads BUT dilation > 1. cuDNN's pooling descriptor has no dilation parameter, so the +// old (asymmetric-pads-only) guard let this fall through to cuDNN, which silently ignored the +// dilation and produced the wrong result. The dilation guard (!default_dilations) now routes this +// to the custom kernel. opset 19 so the CPU AveragePoolV19 reference (which honors dilation) also +// runs and must match. Expected values come from that CPU reference. +TEST(PoolTest, AveragePool_CUDA_symmetric_pad_dilation_1d) { + OpTester test("AveragePool", 19); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{1}); + test.AddAttribute("pads", std::vector{2, 2}); + test.AddAttribute("kernel_shape", std::vector{3}); + test.AddAttribute("dilations", std::vector{2}); + test.AddAttribute("count_include_pad", (int64_t)1); + + std::vector x_vals = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f}; + std::vector x_dims = {1, 1, 9}; + std::vector expected_dims = {1, 1, 9}; + std::vector expected_vals = {1.3333334f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, + 4.6666665f, 5.3333335f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCipTests); +} + +// auto_pad=SAME_LOWER produces naturally asymmetric pads with the extra pad on the LOW side +// (here pad(1,0)); companion to the SAME_UPPER case. opset 19 so the CPU reference also runs. +TEST(PoolTest, AveragePool_CUDA_same_lower_asymmetric_1d) { + OpTester test("AveragePool", 19); + + test.AddAttribute("auto_pad", "SAME_LOWER"); + test.AddAttribute("strides", std::vector{2}); + test.AddAttribute("kernel_shape", std::vector{3}); + test.AddAttribute("count_include_pad", (int64_t)1); + + std::vector x_vals = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f}; + std::vector x_dims = {1, 1, 10}; + std::vector expected_dims = {1, 1, 5}; + std::vector expected_vals = {1.0f, 3.0f, 5.0f, 7.0f, 9.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCipTests); +} + +// bf16 AveragePool is intentionally NOT tested here: although the CUDA kernel instantiates +// BFloat16 (compile-checked) and AveragePoolWithPad accumulates it in float like fp16, the ONNX +// AveragePool schema type constraint does not include tensor(bfloat16), so OpTester's model +// type-checker rejects such a graph at load. The fp16 case above already exercises the +// accumulate-in-float path. + TEST(PoolTest, GlobalAveragePool) { OpTester test("GlobalAveragePool"); From 6054e36155b2f816a1786a22ef8fe6f0bda6105e Mon Sep 17 00:00:00 2001 From: neilmsft Date: Fri, 10 Jul 2026 13:29:11 -0700 Subject: [PATCH 2/2] Fix over-copy of packed sub-byte tensors in OrtApi::GetValue (#29157) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary OrtApi::GetValue on a sequence of tensors copied elements using element_type->Size() * element_count bytes. For packed sub-byte types (INT4/UINT4, two elements per byte) this is ~2× the real storage size, causing a heap over-read of the source and overflow of the destination. ### Fix In PopulateTensorWithData (onnxruntime/core/session/onnxruntime_c_api.cc), copy the tensor's actual packed size via Tensor::SizeInBytes() instead of element_type->Size() * num_elems, and drop the now-unused elem_size parameter. SizeInBytes() is packing-aware: identical for ≥1-byte types (no behavior change) and correct for sub-byte types. ### Testing - Added CApiTest.CreateGetSeqSubByteTensors: sequences of INT4/UINT4 tensors with an odd length (7 elements → 4 packed bytes), verifying GetValue() round-trips correctly. - Full onnxruntime_shared_lib_test suite passes (185/185); other paths unchanged. - Confirmed under AddressSanitizer: the old formula triggers a heap-buffer-overflow, the new one is clean. --- .../core/session/onnxruntime_c_api.h | 10 ++- .../core/session/onnxruntime_cxx_api.h | 13 +++- .../core/optimizer/constant_folding.cc | 60 +++++++++++++--- onnxruntime/core/session/onnxruntime_c_api.cc | 36 +++++----- .../test/optimizer/graph_transform_test.cc | 71 +++++++++++++++++++ .../test/shared_lib/test_nontensor_types.cc | 47 ++++++++++++ 6 files changed, 205 insertions(+), 32 deletions(-) diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index 34a15076c8a7f..5e752a7a9fcf6 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -5767,9 +5767,13 @@ struct OrtApi { /** \brief Compute total size in bytes of the tensor data contained in an OrtValue. * - * Returns the total number of bytes used to store the tensor data. For numeric tensors, - * this is sizeof(element_type) * total_element_count. OrtValues that are not tensors or - * that are tensors that contain strings will cause an error to be returned. + * Returns the total number of bytes used to store the tensor data. For numeric tensors of a + * type that occupies at least one byte per element, this is sizeof(element_type) * + * total_element_count. For packed sub-byte types (e.g. int4/uint4, which store multiple + * elements per byte) it is the actual packed storage size, which is smaller than + * sizeof(element_type) * total_element_count. Use this value (not the element count) when + * copying or bounds-checking the raw tensor buffer. OrtValues that are not tensors or that are + * tensors that contain strings will cause an error to be returned. * * \param[in] ort_value OrtValue instance containing a tensor * \param[out] size The total size of the tensor data in bytes diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index 55a4e36167e86..f15b19a56af2e 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -2064,7 +2064,12 @@ struct TensorTypeAndShapeInfoImpl : Base { using B::B; ONNXTensorElementDataType GetElementType() const; ///< Wraps OrtApi::GetTensorElementType - size_t GetElementCount() const; ///< Wraps OrtApi::GetTensorShapeElementCount + + /// Wraps OrtApi::GetTensorShapeElementCount. + /// Returns the number of logical elements in the tensor (the product of its shape dimensions). + /// Use Ort::Value::GetTensorSizeInBytes() when sizing or bounds-checking the raw buffer returned + /// by GetTensorRawData()/GetTensorData\(). + size_t GetElementCount() const; size_t GetDimensionsCount() const; ///< Wraps OrtApi::GetDimensionsCount @@ -2345,7 +2350,11 @@ struct ConstValueImpl : Base { /// /// Returns the total size of the tensor data in bytes. Throws an exception if the OrtValue /// does not contain a tensor or if it contains a tensor that contains strings. - /// For numeric tensors, this is sizeof(element_type) * total_element_count. + /// For numeric tensors of a type that occupies at least one byte per element, this is + /// sizeof(element_type) * total_element_count. For packed sub-byte types (e.g. int4/uint4) + /// it is the actual packed storage size, which is smaller than the element count returned by + /// GetTensorTypeAndShapeInfo().GetElementCount(). Use this value (not the element count) when + /// copying or bounds-checking the raw buffer from GetTensorRawData()/GetTensorData\(). /// /// The total size of the tensor data in bytes size_t GetTensorSizeInBytes() const; ///< Wraps OrtApi::GetTensorSizeInBytes diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index 84bd36938c9c2..092d9c7720520 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -10,6 +10,7 @@ #include "core/graph/graph_utils.h" #include "core/optimizer/optimizer_execution_frame.h" #include "core/framework/op_kernel.h" +#include "core/framework/tensor.h" #include "core/framework/tensorprotoutils.h" #include "core/session/onnxruntime_session_options_config_keys.h" #include "core/common/safeint.h" @@ -148,6 +149,15 @@ static Status ConstantFoldIfNode(Graph& graph, Node& if_node, const logging::Log static constexpr int64_t kDefaultConstantFoldingMaxOutputSizeInBytes = 1024 * 1024 * 1024; static size_t GetElementSizeForConstantFolding(ONNX_NAMESPACE::TensorProto_DataType elem_type) { + // Complex types are excluded: ORT's tensor type system (DataTypeImpl::TensorTypeFromONNXEnum) + // does not support them, so we cannot size them via Tensor::CalculateTensorStorageSize(). + // Returning 0 makes callers treat the size as unknown instead of attempting an unsupported + // conversion. + if (elem_type == ONNX_NAMESPACE::TensorProto_DataType_COMPLEX64 || + elem_type == ONNX_NAMESPACE::TensorProto_DataType_COMPLEX128) { + return 0; + } + const size_t element_size = utils::GetElementSizeOfTensor(elem_type); if (element_size != 0) { return element_size; @@ -157,6 +167,28 @@ static size_t GetElementSizeForConstantFolding(ONNX_NAMESPACE::TensorProto_DataT return elem_type == ONNX_NAMESPACE::TensorProto_DataType_STRING ? sizeof(std::string) : 0; } +// Computes the packed storage size in bytes for `num_elements` elements of `elem_type`. +// Delegates to Tensor::CalculateTensorStorageSize() so the sub-byte packing math (e.g. int4/uint4 +// pack 2 elements per byte, int2/uint2 pack 4) lives in a single place and does not have to be kept +// in sync here. Returns -1 if the element type is not a recognized tensor element type or the +// computed size cannot be represented. +static int64_t ComputeTensorSizeInBytesForConstantFolding(ONNX_NAMESPACE::TensorProto_DataType elem_type, + int64_t num_elements) { + if (GetElementSizeForConstantFolding(elem_type) == 0) { + return -1; // Unknown element type. + } + + const MLDataType elt_type = DataTypeImpl::TensorTypeFromONNXEnum(elem_type)->GetElementType(); + size_t storage_size = 0; + const Status status = + Tensor::CalculateTensorStorageSize(elt_type, TensorShape({num_elements}), /*alignment*/ 0, storage_size); + if (!status.IsOK() || storage_size > static_cast(std::numeric_limits::max())) { + return -1; + } + + return static_cast(storage_size); +} + static int64_t EstimateTensorElementCount(const ONNX_NAMESPACE::TensorShapeProto& shape) { SafeInt num_elements = 1; for (int i = 0; i < shape.dim_size(); ++i) { @@ -197,7 +229,7 @@ static int64_t EstimateTensorSizeInBytes(const NodeArg& node_arg) { return -1; } - return SafeInt(num_elements) * static_cast(element_size); + return ComputeTensorSizeInBytesForConstantFolding(elem_type, num_elements); } static int64_t EstimateUniqueOutputSizeInBytes(const Node& node) { @@ -233,8 +265,17 @@ static int64_t EstimateUniqueOutputSizeInBytes(const Node& node) { continue; } - const size_t element_size = output_idx == 0 ? input_element_size : sizeof(int64_t); - total_size += SafeInt(input_num_elements) * static_cast(element_size); + if (output_idx == 0) { + // Output 0 has the same (possibly packed sub-byte) element type as the input. + const int64_t output0_size = ComputeTensorSizeInBytesForConstantFolding(input_elem_type, input_num_elements); + if (output0_size < 0) { + return -1; // Output 0 size is unknown; treat the whole node's output size as unknown. + } + total_size += output0_size; + } else { + // The remaining Unique outputs are int64 index/count tensors. + total_size += SafeInt(input_num_elements) * sizeof(int64_t); + } } return total_size; @@ -280,21 +321,20 @@ static int64_t EstimateConstantOfShapeOutputSizeInBytes(const Node& node, const num_elements *= dim; } - // Determine the element size of the output. The ONNX spec for ConstantOfShape defaults the + // Determine the element type of the output. The ONNX spec for ConstantOfShape defaults the // element type to float when the optional 'value' attribute is absent. - size_t element_size = sizeof(float); + auto elem_type = ONNX_NAMESPACE::TensorProto_DataType_FLOAT; const auto& attrs = node.GetAttributes(); auto it = attrs.find("value"); if (it != attrs.end() && it->second.type() == ONNX_NAMESPACE::AttributeProto::TENSOR) { - const auto elem_type = static_cast( + const auto value_elem_type = static_cast( it->second.t().data_type()); - const size_t es = GetElementSizeForConstantFolding(elem_type); - if (es != 0) { - element_size = es; + if (GetElementSizeForConstantFolding(value_elem_type) != 0) { + elem_type = value_elem_type; } } - return num_elements * static_cast(element_size); + return ComputeTensorSizeInBytesForConstantFolding(elem_type, num_elements); } // Estimate the total output size in bytes for a node using shape inference results. diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index e55af70915a39..34369788ec445 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -2018,17 +2018,24 @@ static ORT_STATUS_PTR OrtGetValueImplSeqOfMap(const OrtValue* p_ml_value, int in } #endif -ORT_STATUS_PTR PopulateTensorWithData(Tensor& tensor, bool is_string, _In_ const void* data_elem, size_t num_elems, - size_t elem_size) { - auto len = narrow(tensor.Shape().Size()); - if (num_elems < len) { - return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "input array is too short"); - } - if (!is_string) { - memcpy(tensor.MutableDataRaw(), data_elem, elem_size * num_elems); +// Copies the contents of `data_elem` into `tensor`. +// +// Both the element type and the number of elements to copy are derived from `tensor`: +// - For numeric tensors, tensor.SizeInBytes() bytes are copied. This is packing-aware, +// so sub-byte types (e.g. int4/uint4) copy only their packed storage size rather than +// one byte per logical element. +// - For string tensors, exactly tensor.Shape().Size() std::string elements are copied. +// +// Assumption: `data_elem` points to a buffer that holds at least as many elements as +// `tensor`'s shape, with a matching element type. All current callers satisfy this because +// they size the tensor from the same source data. +ORT_STATUS_PTR PopulateTensorWithData(Tensor& tensor, _In_ const void* data_elem) { + if (!tensor.IsDataTypeString()) { + memcpy(tensor.MutableDataRaw(), data_elem, tensor.SizeInBytes()); } else { + const auto len = narrow(tensor.Shape().Size()); const std::string* strings = reinterpret_cast(data_elem); - auto str_span = gsl::make_span(strings, num_elems); + auto str_span = gsl::make_span(strings, len); auto* dst = tensor.MutableData(); std::copy(str_span.begin(), str_span.end(), dst); } @@ -2036,10 +2043,9 @@ ORT_STATUS_PTR PopulateTensorWithData(Tensor& tensor, bool is_string, _In_ const } ORT_STATUS_PTR CreateTensorAndPopulate(MLDataType element_type, const int64_t* shape, size_t shape_len, - const void* data, size_t num_elements, _Inout_ OrtAllocator* allocator, OrtValue& result) { + const void* data, _Inout_ OrtAllocator* allocator, OrtValue& result) { ORT_API_RETURN_IF_ERROR(CreateTensorImpl(element_type, shape, shape_len, allocator, result)); - ORT_API_RETURN_IF_ERROR(PopulateTensorWithData(*result.GetMutable(), utils::IsDataTypeString(element_type), - data, num_elements, element_type->Size())); + ORT_API_RETURN_IF_ERROR(PopulateTensorWithData(*result.GetMutable(), data)); return nullptr; } @@ -2057,7 +2063,6 @@ static ORT_STATUS_PTR OrtGetValueImplSeqOfTensors(_In_ const OrtValue* p_ml_valu auto result = std::make_unique(); ORT_API_RETURN_IF_ERROR(c_api_internal::CreateTensorAndPopulate(one_tensor.DataType(), tensor_shape.GetDims().data(), tensor_shape.NumDimensions(), one_tensor.DataRaw(), - narrow(one_tensor.Shape().Size()), allocator, *result)); *out = result.release(); return nullptr; @@ -2105,7 +2110,6 @@ static ORT_STATUS_PTR OrtGetValueImplMapHelper(_In_ const OrtValue* p_ml_value, std::vector vec_keys; std::vector vec_vals; const void* data_ptr; - size_t data_size; MLDataType element_type; switch (index) { case 0: { // user is requesting keys @@ -2113,20 +2117,18 @@ static ORT_STATUS_PTR OrtGetValueImplMapHelper(_In_ const OrtValue* p_ml_value, vec_keys.reserve(static_cast(num_kv_pairs)); std::transform(data.cbegin(), data.cend(), std::back_inserter(vec_keys), [](const auto& k) { return k.first; }); data_ptr = vec_keys.data(); - data_size = vec_keys.size(); } break; case 1: { // user is requesting values element_type = DataTypeImpl::TensorTypeFromONNXEnum(GetONNXTensorElementDataType())->GetElementType(); vec_vals.reserve(static_cast(num_kv_pairs)); std::transform(data.cbegin(), data.cend(), std::back_inserter(vec_vals), [](const auto& k) { return k.second; }); data_ptr = vec_vals.data(); - data_size = vec_vals.size(); } break; default: return OrtApis::CreateStatus(ORT_FAIL, "Invalid index requested for map type."); } ORT_API_RETURN_IF_ERROR(c_api_internal::CreateTensorAndPopulate(element_type, dims.data(), dims.size(), data_ptr, - data_size, allocator, *result)); + allocator, *result)); *out = result.release(); return nullptr; } diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 1f07bf3204894..3327b29983d1a 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -19,6 +19,7 @@ #include "core/common/span_utils.h" #include "core/framework/data_types.h" +#include "core/framework/int4.h" #include "core/framework/ort_value.h" #include "core/graph/graph_utils.h" #include "core/graph/graph_viewer.h" @@ -1753,6 +1754,76 @@ TEST_F(GraphTransformationTests, ConstantFoldingConstantOfShapeBlockedWhenOutput pre_graph_checker, post_graph_checker)); } +// The constant-folding output-size estimate must be packing-aware for sub-byte types. +// A QuantizeLinear node that produces a packed int4 output stores 2 elements per byte, so its +// real storage size is ceil(num_elements / 2) bytes. The pre-execution size estimate previously +// used num_elements * 1 byte, over-counting by ~2x and needlessly blocking folding of sub-byte +// outputs that actually fit within the configured limit. +TEST_F(GraphTransformationTests, ConstantFoldingSubByteOutputSizeIsPacked) { + // 20000 int4 elements -> 10000 packed bytes of real storage (the old estimate was 20000 bytes). + constexpr int64_t kNumElements = 20000; + + auto build_model = [&](ModelTestBuilder& builder) { + auto* input_data = builder.MakeInitializer({kNumElements}, -1.0f, 1.0f); + auto* output_arg = builder.MakeOutput(); + builder.AddQuantizeLinearNode(input_data, 1.0f, Int4x2(0, 0), output_arg); + }; + + // Case 1: limit (15000 bytes) sits between the packed size (10000) and the old over-estimate + // (20000). With the packing-aware estimate the node SHOULD be folded. Under the old estimate + // it would have been (incorrectly) skipped. + { + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["QuantizeLinear"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // Real packed size (10000 bytes) is within the 15000 byte limit, so folding proceeds. + TEST_RETURN_IF_NOT(op_to_count["QuantizeLinear"] == 0); + return Status::OK(); + }; + + std::unique_ptr e = std::make_unique(CPUExecutionProviderInfo()); + ConfigOptions config_options; + ASSERT_STATUS_OK(config_options.AddConfigEntry( + kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes, "15000")); + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 21, *logger_, + std::make_unique(*e.get(), false, config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); + } + + // Case 2: limit (5000 bytes) is below even the packed size (10000), so folding is blocked. + { + auto pre_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + TEST_RETURN_IF_NOT(op_to_count["QuantizeLinear"] == 1); + return Status::OK(); + }; + + auto post_graph_checker = [](Graph& graph) -> Status { + auto op_to_count = CountOpsInGraph(graph); + // Real packed size (10000 bytes) exceeds the 5000 byte limit, so the node is not folded. + TEST_RETURN_IF_NOT(op_to_count["QuantizeLinear"] == 1); + return Status::OK(); + }; + + std::unique_ptr e = std::make_unique(CPUExecutionProviderInfo()); + ConfigOptions config_options; + ASSERT_STATUS_OK(config_options.AddConfigEntry( + kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes, "5000")); + + ASSERT_STATUS_OK(TestGraphTransformer(build_model, 21, *logger_, + std::make_unique(*e.get(), false, config_options), + TransformerLevel::Level1, 1, + pre_graph_checker, post_graph_checker)); + } +} + // Test that small constant folding still works with the size limit. TEST_F(GraphTransformationTests, ConstantFoldingSmallOutputAllowed) { // Build a model with a small Expand: scalar -> [4, 4] = 16 * 4 = 64 bytes. diff --git a/onnxruntime/test/shared_lib/test_nontensor_types.cc b/onnxruntime/test/shared_lib/test_nontensor_types.cc index dc361e6a65510..a1a7e0c1df8ce 100644 --- a/onnxruntime/test/shared_lib/test_nontensor_types.cc +++ b/onnxruntime/test/shared_lib/test_nontensor_types.cc @@ -5,6 +5,7 @@ #include #include #include +#include #include "core/common/common.h" #include "core/session/onnxruntime_cxx_api.h" @@ -278,6 +279,52 @@ TEST(CApiTest, CreateGetSeqStringTensors) { ASSERT_EQ(string_set, std::set(std::begin(string_input_data), std::end(string_input_data))); } +// Test - GetValue() on a sequence of packed sub-byte tensors +// (int4/uint4) must copy only the packed storage bytes. +TEST(CApiTest, CreateGetSeqSubByteTensors) { + auto default_allocator = std::make_unique(); + Ort::MemoryInfo info("Cpu", OrtDeviceAllocator, 0, OrtMemTypeDefault); + + auto run_for_type = [&](ONNXTensorElementDataType elem_type, std::array packed) { + const std::vector dims{7}; // 7 4-bit elements -> 4 packed bytes + constexpr int N = 2; + + std::vector in; + for (int i = 0; i < N; ++i) { + Ort::Value tensor = Ort::Value::CreateTensor(info, packed.data(), packed.size(), + dims.data(), dims.size(), elem_type); + in.push_back(std::move(tensor)); + } + + Ort::Value seq_ort = Ort::Value::CreateSequence(in); + + for (int idx = 0; idx < N; ++idx) { + Ort::Value out = seq_ort.GetValue(idx, default_allocator.get()); + + auto type_info = out.GetTypeInfo(); + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + ASSERT_EQ(tensor_info.GetElementType(), elem_type); + ASSERT_EQ(tensor_info.GetShape(), dims); + + // `out` is a fresh tensor that GetValue() allocated and copied into, so its data is a + // distinct buffer from the input `packed` bytes (the input tensors alias `packed`). + // Comparing them therefore validates the copy rather than reading the same memory twice. + const size_t out_bytes = out.GetTensorSizeInBytes(); + ASSERT_EQ(out_bytes, packed.size()); + const auto* ret = static_cast(out.GetTensorRawData()); + ASSERT_NE(static_cast(ret), static_cast(packed.data())); + for (size_t i = 0; i < out_bytes; ++i) { + ASSERT_EQ(ret[i], packed[i]); + } + } + }; + + // {0, 1, 2, 3, -8, 7, 6, pad_0} + run_for_type(ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_INT4, {0x10, 0x32, 0x78, 0x06}); + // {0, 1, 2, 3, 4, 5, 15, pad_0} + run_for_type(ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT4, {0x10, 0x32, 0x54, 0x0F}); +} + TEST(CApiTest, TypeInfoSequence) { // Creation auto default_allocator = std::make_unique();