diff --git a/onnxruntime/core/optimizer/nchwc_transformer.cc b/onnxruntime/core/optimizer/nchwc_transformer.cc index a971a058f43b7..f60412e19aee0 100644 --- a/onnxruntime/core/optimizer/nchwc_transformer.cc +++ b/onnxruntime/core/optimizer/nchwc_transformer.cc @@ -533,6 +533,31 @@ void NchwcTransformerImpl::TransformPool(Node& node) { return; } + // Bail out for AveragePool with ceil_mode==1 && count_include_pad==1. The default float CPU + // AveragePool kernel was fixed (PR #29629) to divide the trailing ceil_mode window by its + // clamped size instead of the full kernel size. NchwcAveragePool routes to + // MlasAveragePoolingIncludePad, which still uses the full-kernel-size divisor and would + // silently reintroduce the wrong-average bug for optimized NCHWc graphs. Leaving this combo + // unconverted keeps it on the fixed CPU path; every other AveragePool (and all MaxPool / + // global pooling) conversion is unaffected, so there is no perf impact on the common case. + if (node.OpType() == "AveragePool") { + const NodeAttributes& attrs = node.GetAttributes(); + const auto get_int_attr = [&attrs](const char* name) -> int64_t { + const auto it = attrs.find(name); + return (it != attrs.end() && it->second.type() == ONNX_NAMESPACE::AttributeProto_AttributeType_INT) + ? it->second.i() + : 0; + }; + // Gate count_include_pad as (!= 0) to match how PoolBase and the float (pool.cc) / fp16 + // kernels treat it: any nonzero value enables include-pad. Using == 1 here would let an + // out-of-spec count_include_pad=2 model escape the bail-out, convert to NchwcAveragePool, + // and silently hit the buggy full-kernel divisor in the optimized graph. ceil_mode stays + // == 1 because the kernels gate the ceil-window fix on exactly ceil_mode == 1. + if (get_int_attr("ceil_mode") == 1 && get_int_attr("count_include_pad") != 0) { + return; + } + } + const size_t nchwc_block_size = MlasNchwcGetBlockSize(); const auto* input_type = input_defs[0]->TypeAsProto(); diff --git a/onnxruntime/core/providers/cpu/fp16/fp16_pool.cc b/onnxruntime/core/providers/cpu/fp16/fp16_pool.cc index 8be90a249f711..20e957494c591 100644 --- a/onnxruntime/core/providers/cpu/fp16/fp16_pool.cc +++ b/onnxruntime/core/providers/cpu/fp16/fp16_pool.cc @@ -7,6 +7,7 @@ #include "core/framework/op_kernel.h" #include "core/util/math.h" +#include "core/providers/cpu/nn/pool.h" #include "core/providers/cpu/nn/pool_attributes.h" #include "core/platform/threadpool.h" #include "core/common/safeint.h" @@ -36,6 +37,22 @@ class PoolFp16 : public OpKernel { PoolAttributes pool_attrs_; bool is_max_pool_; // either max pool or average pool bool channels_last_; + + private: + // Correct reference fallback for ceil_mode==1 && count_include_pad==1 AveragePool. The MLAS + // fp16 im2col path divides by the full kernel_size and cannot drop the ceil_mode phantom tail + // cells. Rather than duplicate the pooling loop, this casts the tensor MLFloat16->float, + // delegates to the shared float ComputeAveragePoolReferenceCompute (cpu/nn/pool.cc), then casts + // float->MLFloat16 with a single round at store. Accumulation happens in float inside the shared + // helper, so the result is identical to accumulate-in-float + round-once. Handles both NCHW and + // channels-last (NHWC) layouts (transposing to NCHW for the shared helper) and 1D/2D/3D. Zero + // MLAS edits. + Status ComputeAveragePoolFp16Reference(OpKernelContext* context, + const Tensor* X, + const TensorShapeVector& output_dims, + const TensorShapeVector& pads, + int64_t N, + int64_t C) const; }; Status PoolFp16::Compute(OpKernelContext* context) const { @@ -115,6 +132,17 @@ Status PoolFp16::Compute(OpKernelContext* context) const { output_dims.push_back(C); } + // The MLAS fp16 NHWC avg-pool (MlasNhwcAvgPool -> mlas/lib/pooling_fp16.cpp) divides every + // window by the full kernel_size and cannot drop the ceil_mode phantom tail cells, giving a + // too-small average for ceil_mode==1 && count_include_pad. Route only that combo to a correct + // float-accumulating reference loop; every other case keeps the fast MLAS im2col path. This + // mirrors the float fix in pool.cc (Pool::Compute guard) so the fp16 and + // float AveragePool paths stay dtype-consistent. + if (!is_max_pool_ && pool_attrs_.ceil_mode == 1 && pool_attrs_.count_include_pad && + !pool_attrs_.global_pooling) { + return ComputeAveragePoolFp16Reference(context, X, output_dims, pads, N, C); + } + const bool need_padding = !is_max_pool_ && pool_attrs_.count_include_pad; std::vector padding_data; if (need_padding) { @@ -217,9 +245,94 @@ Status PoolFp16::Compute(OpKernelContext* context) const { return Status::OK(); } -// -// Operator definitions -// +// fp16 AveragePool reference fallback for ceil_mode + count_include_pad. Casts the input +// MLFloat16->float (transposing NHWC->NCHW when channels_last_), runs the shared float pooling +// core (ComputeAveragePoolReferenceCompute in cpu/nn/pool.cc), then casts float->MLFloat16 into +// the output (transposing NCHW->NHWC back when channels_last_). A single round happens at the +// float->fp16 store; all accumulation is in float inside the shared core, so this is numerically +// equivalent to accumulate-in-float + round-once and stays in lockstep with the float path's +// clamped-divisor semantics. Reuses the float loop per xadupre's review request. +Status PoolFp16::ComputeAveragePoolFp16Reference(OpKernelContext* context, + const Tensor* X, + const TensorShapeVector& output_dims, + const TensorShapeVector& pads, + int64_t N, + int64_t C) const { + const TensorShape& input_shape = X->Shape(); + const size_t spatial_dims = output_dims.size() - 2; + const size_t spatial_dim_start = channels_last_ ? 1 : 2; + + // Build NCHW spatial extents plus flat spatial sizes for input and output. + int64_t input_image_size = 1; + int64_t output_image_size = 1; + TensorShapeVector x_nchw_dims({N, C}); + TensorShapeVector out_nchw_dims({N, C}); + for (size_t d = 0; d < spatial_dims; ++d) { + const int64_t in_d = input_shape[d + spatial_dim_start]; + const int64_t out_d = output_dims[d + spatial_dim_start]; + x_nchw_dims.push_back(in_d); + out_nchw_dims.push_back(out_d); + input_image_size *= in_d; + output_image_size *= out_d; + } + const TensorShape x_nchw_shape(x_nchw_dims); + + AllocatorPtr alloc; + ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc)); + + const size_t x_float_count = static_cast(N) * static_cast(C) * + static_cast(input_image_size); + const size_t y_float_count = static_cast(N) * static_cast(C) * + static_cast(output_image_size); + auto x_float = IAllocator::MakeUniquePtr(alloc, x_float_count); + auto y_float = IAllocator::MakeUniquePtr(alloc, y_float_count); + + // Cast MLFloat16 -> float into an NCHW buffer (transpose from NHWC when channels_last_). + const auto* Xdata = X->Data(); + float* x_ptr = x_float.get(); + if (channels_last_) { + for (int64_t n = 0; n < N; ++n) { + for (int64_t c = 0; c < C; ++c) { + for (int64_t s = 0; s < input_image_size; ++s) { + x_ptr[(n * C + c) * input_image_size + s] = + Xdata[(n * input_image_size + s) * C + c].ToFloat(); + } + } + } + } else { + for (size_t i = 0; i < x_float_count; ++i) { + x_ptr[i] = Xdata[i].ToFloat(); + } + } + + // Run the shared float pooling core. The guard in Compute() excludes global_pooling, so + // pool_attrs_ carries fully-populated kernel_shape/strides/dilations that match this call. + concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); + ORT_RETURN_IF_ERROR(ComputeAveragePoolReferenceCompute(x_ptr, y_float.get(), x_nchw_shape, + out_nchw_dims, pads, pool_attrs_, tp)); + + // Cast float -> MLFloat16 into the output (transpose back to NHWC when channels_last_). + // MLFloat16(float) is the single, final rounding step. + auto* Y = context->Output(0, output_dims); + auto* Ydata = Y->MutableData(); + const float* y_ptr = y_float.get(); + if (channels_last_) { + for (int64_t n = 0; n < N; ++n) { + for (int64_t c = 0; c < C; ++c) { + for (int64_t s = 0; s < output_image_size; ++s) { + Ydata[(n * output_image_size + s) * C + c] = + MLFloat16(y_ptr[(n * C + c) * output_image_size + s]); + } + } + } + } else { + for (size_t i = 0; i < y_float_count; ++i) { + Ydata[i] = MLFloat16(y_ptr[i]); + } + } + + return Status::OK(); +} ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL( MaxPool, 8, 11, MLFloat16, diff --git a/onnxruntime/core/providers/cpu/nn/pool.cc b/onnxruntime/core/providers/cpu/nn/pool.cc index d6b9ed693432b..2fe135e117436 100644 --- a/onnxruntime/core/providers/cpu/nn/pool.cc +++ b/onnxruntime/core/providers/cpu/nn/pool.cc @@ -41,6 +41,117 @@ inline static void RunLoop(concurrency::ThreadPool* tp, std::ptrdiff_t total_cha concurrency::ThreadPool::TryParallelFor(tp, total_channels, task.Cost(), task); } +// Core (non-MLAS) average-pooling loop, operating on caller-provided NCHW float buffers and +// pre-computed output_dims/pads. Correct for ceil_mode + count_include_pad because the +// AveragePool{1,2,3}DTask functors clamp the window end to input+tail_pad, excluding the +// ceil-mode phantom cells that MLAS's full-kernel divisor would count. +// +// This is the single source of truth for the reference average, shared by: +// - the float production path (AveragePoolV19 and the Pool opset 7-18 +// ceil_mode+count_include_pad fallback), via the ComputeAveragePoolReference wrapper below; +// - the fp16 fallback (PoolFp16::ComputeAveragePoolFp16Reference), which casts the tensor +// MLFloat16->float, calls this, then casts float->MLFloat16 (single round at store). +// Accumulation is in float regardless of caller dtype, so the fp16 path is numerically +// equivalent to accumulate-in-float + round-once. +Status ComputeAveragePoolReferenceCompute(const float* X_data, float* Y_data, + const TensorShape& x_shape, + gsl::span output_dims, + gsl::span pads, + const PoolAttributes& pool_attrs, + concurrency::ThreadPool* tp) { + const auto& kernel_shape = pool_attrs.kernel_shape; + + const int64_t channels = x_shape[1]; + const int64_t height = x_shape[2]; + const int64_t width = kernel_shape.size() > 1 ? x_shape[3] : 1; + const int64_t depth = kernel_shape.size() > 2 ? x_shape[4] : 1; + const int64_t pooled_height = output_dims[2]; + const int64_t pooled_width = kernel_shape.size() > 1 ? output_dims[3] : 1; + const int64_t pooled_depth = kernel_shape.size() > 2 ? output_dims[4] : 1; + const int64_t total_channels = x_shape[0] * channels; + + const int64_t stride_h = pool_attrs.strides[0]; + const int64_t stride_w = kernel_shape.size() > 1 ? pool_attrs.strides[1] : 1; + const int64_t stride_d = kernel_shape.size() > 2 ? pool_attrs.strides[2] : 1; + + // p (LpPool p-norm) is unused by AveragePool functors; pass 0. + constexpr int64_t p = 0; + + switch (kernel_shape.size()) { + case 1: { + int64_t x_step = height; + int64_t y_step = pooled_height; + const int64_t dilation_h = pool_attrs.dilations[0]; + RunLoop>(tp, onnxruntime::narrow(total_channels), + {X_data, Y_data, x_step, y_step, dilation_h, pooled_height, stride_h, + height, kernel_shape, pads, pool_attrs.count_include_pad, p}); + break; + } + case 2: { + int64_t x_step = height * width; + int64_t y_step = pooled_height * pooled_width; + const int64_t dilation_h = pool_attrs.dilations[0]; + const int64_t dilation_w = pool_attrs.dilations[1]; + RunLoop>( + tp, onnxruntime::narrow(total_channels), + {X_data, Y_data, x_step, y_step, dilation_h, dilation_w, pooled_height, pooled_width, stride_h, + stride_w, height, width, kernel_shape, pads, pool_attrs.count_include_pad, p}); + break; + } + case 3: { + int64_t x_step = height * width * depth; + int64_t y_step = pooled_height * pooled_width * pooled_depth; + const int64_t dilation_h = pool_attrs.dilations[0]; + const int64_t dilation_w = pool_attrs.dilations[1]; + const int64_t dilation_d = pool_attrs.dilations[2]; + RunLoop>(tp, onnxruntime::narrow(total_channels), + {X_data, Y_data, x_step, y_step, dilation_h, dilation_w, dilation_d, + pooled_height, pooled_width, pooled_depth, stride_h, stride_w, stride_d, + height, width, depth, kernel_shape, pads, pool_attrs.count_include_pad, p}); + break; + } + default: + return Status(ONNXRUNTIME, INVALID_ARGUMENT, + "Unsupported kernel dimension : " + std::to_string(kernel_shape.size())); + } + return Status::OK(); +} + +// Context wrapper for the float production path (AveragePoolV19, opset >= 19; and the +// Pool opset 7-18 ceil_mode+count_include_pad fallback). Pulls X/Y from +// the kernel context, mirrors PoolBase's rank checks, computes output_dims/pads, then runs the +// shared float compute. Behavior-preserving: identical geometry + functor dispatch as before. +template +Status ComputeAveragePoolReference(OpKernelContext* context, + const PoolAttributes& pool_attrs, + concurrency::ThreadPool* tp) { + ORT_ENFORCE(!pool_attrs.global_pooling, + "ComputeAveragePoolReference requires non-global pooling; " + "strides/dilations are not populated for global pooling."); + + const auto* X = context->Input(0); + const TensorShape& x_shape = X->Shape(); + ORT_RETURN_IF_NOT(x_shape.NumDimensions() >= 3, "Input dimension cannot be less than 3."); + + // Mirror the rank checks PoolBase::Compute performs before SetOutputSize. This function is + // reached directly from the Pool::Compute ceil_mode+count_include_pad + // guard (and from AveragePoolV19), bypassing PoolBase::Compute, so without these a + // rank-mismatched model would reach SetOutputSize/InferOutputSize and read out of bounds. + const size_t pooling_dims = x_shape.NumDimensions() - 2; + if (pooling_dims > 3) { + return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported pooling size."); + } + ORT_RETURN_IF_NOT(pooling_dims == pool_attrs.kernel_shape.size(), + "kernel_shape num_dims is not compatible with X num_dims."); + + auto pads = pool_attrs.pads; + auto output_dims = pool_attrs.SetOutputSize(x_shape, x_shape[1], &pads); + Tensor* Y = context->Output(0, output_dims); + + return ComputeAveragePoolReferenceCompute(X->Data(), Y->MutableData(), x_shape, + output_dims, pads, pool_attrs, tp); +} + template Status Pool::Compute(OpKernelContext* context) const { concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); @@ -152,6 +263,15 @@ Status Pool>::Compute(OpKernelContext* context) co template <> Status Pool::Compute(OpKernelContext* context) const { + // MLAS's include-pad path divides by the full kernel size and cannot drop the + // ceil_mode phantom tail cells, giving a wrong average for + // ceil_mode=1 && count_include_pad. Route only that combo to the reference loop, + // which clamps the window to input+tail_pad. Everything else keeps the MLAS fast path. + // Global pooling is excluded because it has no ceil/pads (so it is already correct) and + // its strides[] are not populated, which ComputeAveragePoolReference requires. + if (pool_attrs_.ceil_mode == 1 && pool_attrs_.count_include_pad && !pool_attrs_.global_pooling) { + return ComputeAveragePoolReference(context, pool_attrs_, context->GetOperatorThreadPool()); + } return PoolBase::Compute(context, pool_attrs_.count_include_pad ? MlasAveragePoolingIncludePad : MlasAveragePoolingExcludePad); } @@ -251,77 +371,7 @@ Status MaxPoolV8::ComputeImpl(OpKernelContext* context) const { template Status AveragePoolV19::Compute(OpKernelContext* context) const { - concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); - bool need_dilation = false; - for (auto n : pool_attrs_.dilations) { - need_dilation |= n > 1; - } - - const auto* X = context->Input(0); - const TensorShape& x_shape = X->Shape(); - - ORT_RETURN_IF_NOT(x_shape.NumDimensions() >= 3, "Input dimension cannot be less than 3."); - - auto pads = pool_attrs_.pads; - auto kernel_shape = pool_attrs_.kernel_shape; - - auto output_dims = pool_attrs_.SetOutputSize(x_shape, x_shape[1], &pads); - Tensor* Y = context->Output(0, output_dims); - - const auto* X_data = X->Data(); - auto* Y_data = Y->MutableData(); - - // The main loop - int64_t channels = x_shape[1]; - int64_t height = x_shape[2]; - int64_t width = kernel_shape.size() > 1 ? x_shape[3] : 1; - int64_t depth = kernel_shape.size() > 2 ? x_shape[4] : 1; - int64_t pooled_height = output_dims[2]; - int64_t pooled_width = kernel_shape.size() > 1 ? output_dims[3] : 1; - int64_t pooled_depth = kernel_shape.size() > 2 ? output_dims[4] : 1; - const int64_t total_channels = x_shape[0] * channels; - - switch (kernel_shape.size()) { - case 1: { - int64_t x_step = height; - int64_t y_step = pooled_height; - const int64_t dilation_h = pool_attrs_.dilations[0]; - - RunLoop>(tp, onnxruntime::narrow(total_channels), - {X_data, Y_data, x_step, y_step, dilation_h, pooled_height, stride_h(), - height, kernel_shape, pads, pool_attrs_.count_include_pad, p_}); - break; - } - - case 2: { - int64_t x_step = height * width; - int64_t y_step = pooled_height * pooled_width; - const int64_t dilation_h = pool_attrs_.dilations[0]; - const int64_t dilation_w = pool_attrs_.dilations[1]; - RunLoop>( - tp, onnxruntime::narrow(total_channels), - {X_data, Y_data, x_step, y_step, dilation_h, dilation_w, pooled_height, pooled_width, stride_h(), - stride_w(), height, width, kernel_shape, pads, pool_attrs_.count_include_pad, p_}); - break; - } - case 3: { - int64_t x_step = height * width * depth; - int64_t y_step = pooled_height * pooled_width * pooled_depth; - const int64_t dilation_h = pool_attrs_.dilations[0]; - const int64_t dilation_w = pool_attrs_.dilations[1]; - const int64_t dilation_d = pool_attrs_.dilations[2]; - RunLoop>(tp, onnxruntime::narrow(total_channels), - {X_data, Y_data, x_step, y_step, - dilation_h, dilation_w, dilation_d, pooled_height, pooled_width, - pooled_depth, stride_h(), stride_w(), stride_d(), height, - width, depth, kernel_shape, pads, pool_attrs_.count_include_pad, p_}); - break; - } - default: - return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported kernel dimension : " + std::to_string(kernel_shape.size())); - } - - return Status::OK(); + return ComputeAveragePoolReference(context, pool_attrs_, context->GetOperatorThreadPool()); } template diff --git a/onnxruntime/core/providers/cpu/nn/pool.h b/onnxruntime/core/providers/cpu/nn/pool.h index 7796b28b8369b..7ceb3588b4aa6 100644 --- a/onnxruntime/core/providers/cpu/nn/pool.h +++ b/onnxruntime/core/providers/cpu/nn/pool.h @@ -7,6 +7,22 @@ namespace onnxruntime { +namespace concurrency { +class ThreadPool; +} + +// Core float average-pooling reference loop (defined in pool.cc). Operates on caller-provided +// NCHW float buffers plus pre-computed output_dims/pads. Shared single source of truth for the +// float production path and the fp16 fallback (which casts MLFloat16<->float around it). The +// window end is clamped to input+tail_pad so ceil_mode + count_include_pad divides by the real +// (non-phantom) window. Accumulation is always in float. +Status ComputeAveragePoolReferenceCompute(const float* X_data, float* Y_data, + const TensorShape& x_shape, + gsl::span output_dims, + gsl::span pads, + const PoolAttributes& pool_attrs, + concurrency::ThreadPool* tp); + template class Pool : public OpKernel, public PoolBase { public: @@ -34,9 +50,6 @@ class AveragePoolV19 : public OpKernel, public PoolBase { } Status Compute(OpKernelContext* context) const override; - - private: - int64_t p_; }; // For maxpool v8 and beyond diff --git a/onnxruntime/test/optimizer/nchwc_optimizer_test.cc b/onnxruntime/test/optimizer/nchwc_optimizer_test.cc index a07f173950051..c41235138d10a 100644 --- a/onnxruntime/test/optimizer/nchwc_optimizer_test.cc +++ b/onnxruntime/test/optimizer/nchwc_optimizer_test.cc @@ -498,6 +498,42 @@ TEST(NchwcOptimizerTests, ConvAveragePool) { test_case(true); } +// PR #29629 (M1): AveragePool with ceil_mode==1 && count_include_pad==1 must NOT be converted to +// NchwcAveragePool. NchwcAveragePool routes to MlasAveragePoolingIncludePad, which divides by the +// full kernel size and cannot drop the ceil_mode phantom tail cells — reintroducing the wrong +// average that the default float CPU kernel was fixed for. The transformer bails out for this +// combo, leaving the node as the ONNX AveragePool on the fixed CPU path. The Level2-vs-Level3 +// output comparison inside NchwcOptimizerTester additionally guards correctness: if the node were +// (incorrectly) converted, the buggy NCHWc divisor would diverge from the fixed unconverted kernel +// and fail the comparison. +TEST(NchwcOptimizerTests, ConvAveragePoolCeilCountIncludePadNotConverted) { + auto build_test_case = [&](NchwcTestHelper& helper) { + auto* input_arg = helper.MakeInput({1, 48, 34, 34}); + auto* conv_output_arg = helper.MakeIntermediate(); + auto* output_arg = helper.MakeOutput(); + + helper.AddConvNode(input_arg, conv_output_arg, {128, 48, 5, 5}); + + // kernel 4 / stride 4 over the 30x30 conv output: ceil_mode adds a trailing window that + // overruns the input, so the buggy full-kernel divisor and the fixed clamped divisor differ. + auto& pool_node = helper.AddNode("AveragePool", {conv_output_arg}, {output_arg}); + pool_node.AddAttribute("kernel_shape", std::vector{4, 4}); + pool_node.AddAttribute("strides", std::vector{4, 4}); + pool_node.AddAttribute("ceil_mode", static_cast(1)); + pool_node.AddAttribute("count_include_pad", static_cast(1)); + }; + + auto check_nchwc_graph = [&](InferenceSessionWrapper& session) { + auto op_to_count = CountOpsInGraph(session.GetGraph()); + // Conv still converts; the AveragePool stays on the fixed ONNX CPU path (not NCHWc). + EXPECT_EQ(op_to_count["com.microsoft.nchwc.Conv"], 1); + EXPECT_EQ(op_to_count["com.microsoft.nchwc.AveragePool"], 0); + EXPECT_EQ(op_to_count["AveragePool"], 1); + }; + + NchwcOptimizerTester(build_test_case, check_nchwc_graph); +} + TEST(NchwcOptimizerTests, ConvGlobalPool) { auto test_case = [&](const std::string& op_type) { auto build_test_case = [&](NchwcTestHelper& helper) { diff --git a/onnxruntime/test/providers/cpu/nn/pool_fp16_op_test.cc b/onnxruntime/test/providers/cpu/nn/pool_fp16_op_test.cc index 9be733e22f2e6..ec599a6509e35 100644 --- a/onnxruntime/test/providers/cpu/nn/pool_fp16_op_test.cc +++ b/onnxruntime/test/providers/cpu/nn/pool_fp16_op_test.cc @@ -3,6 +3,9 @@ #include "core/mlas/inc/mlas.h" +// NOTE: MLAS_F16VEC_INTRINSICS_SUPPORTED is defined only on ARM64 (non-Apple), so on x86 this +// whole file (and the fp16 PoolFp16 kernel it exercises) compiles out. These tests therefore run +// only on the ARM64 / CoreML / XNNPACK / WebGPU CI legs, never on an x86 host build. #if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) || defined(USE_COREML) || defined(USE_XNNPACK) || defined(USE_WEBGPU) #include "core/providers/cpu/nn/pool.h" @@ -583,6 +586,131 @@ TEST(PoolFp16Test, AveragePool_10_ceil1_2d) { {kTensorrtExecutionProvider, kAclExecutionProvider, kWebGpuExecutionProvider}); } +// fp16 analog of the float AveragePool_18_ceil_count_include_pad_1d bug (PyTorch #183528). +// k=7 s=3 pad=(3,3) ceil_mode=1 count_include_pad=1 over {1,2,9}: the trailing ceil window ends +// past input+pad_tail, so the MLAS fp16 full-kernel divisor gave a too-small average. The new +// ComputeAveragePoolFp16Reference clamps the window end and matches the float reference. +// Expected: out0 = (1+2+9)/7 = 12/7, out1 = (1+2+9)/6 = 2. +TEST(PoolFp16Test, AveragePool_Ceil_CountIncludePad_1d) { + if (DefaultDmlExecutionProvider().get() != nullptr) { + GTEST_SKIP() << "Skipping because of the following error: MLOperatorAuthorImpl.cpp(2100): The parameter is incorrect."; + } + + OpTester test("AveragePool", 11); + 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 = {MLFloat16(1.f), MLFloat16(2.f), MLFloat16(9.f)}; + std::vector x_dims = {1, 1, 3}; + std::vector expected_dims = {1, 1, 2}; + std::vector expected_vals = {MLFloat16(12.0f / 7.0f), MLFloat16(2.0f)}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kAclExecutionProvider, kWebGpuExecutionProvider}); +} + +// fp16 2D analog: PyTorch #183528 repro, x = arange(1,17).reshape(1,1,4,4), k=3 s=2 pad=1, +// ceil_mode=1, count_include_pad=1. The trailing row/col ceil windows extend past input+pad_tail; +// the reference divides by the clamped include-pad window, matching the float fix. +TEST(PoolFp16Test, AveragePool_Ceil_CountIncludePad_2d) { + if (DefaultDmlExecutionProvider().get() != nullptr) { + GTEST_SKIP() << "Skipping because of the following error: MLOperatorAuthorImpl.cpp(2100): The parameter is incorrect."; + } + + OpTester test("AveragePool", 11); + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{2, 2}); + test.AddAttribute("pads", std::vector{1, 1, 1, 1}); + test.AddAttribute("kernel_shape", std::vector{3, 3}); + test.AddAttribute("ceil_mode", (int64_t)1); + test.AddAttribute("count_include_pad", (int64_t)1); + + std::vector x_vals; + for (int i = 1; i <= 16; ++i) { + x_vals.push_back(MLFloat16(static_cast(i))); + } + std::vector x_dims = {1, 1, 4, 4}; + std::vector expected_dims = {1, 1, 3, 3}; + std::vector expected_vals = { + MLFloat16(14.0f / 9.0f), MLFloat16(30.0f / 9.0f), MLFloat16(2.0f), + MLFloat16(38.0f / 6.0f), MLFloat16(11.0f), MLFloat16(6.0f), + MLFloat16(4.5f), MLFloat16(7.5f), MLFloat16(4.0f)}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kAclExecutionProvider, kWebGpuExecutionProvider}); +} + +// No-regression control: ceil_mode=1 but count_include_pad=0 keeps the standard MLAS path +// (guard requires count_include_pad==1), so the divisor already counts only in-bounds cells. +TEST(PoolFp16Test, AveragePool_Ceil_CountExcludePad_2d) { + if (DefaultDmlExecutionProvider().get() != nullptr) { + GTEST_SKIP() << "Skipping because of the following error: MLOperatorAuthorImpl.cpp(2100): The parameter is incorrect."; + } + + OpTester test("AveragePool", 11); + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{2, 2}); + test.AddAttribute("pads", std::vector{1, 1, 1, 1}); + test.AddAttribute("kernel_shape", std::vector{3, 3}); + test.AddAttribute("ceil_mode", (int64_t)1); + test.AddAttribute("count_include_pad", (int64_t)0); + + std::vector x_vals; + for (int i = 1; i <= 16; ++i) { + x_vals.push_back(MLFloat16(static_cast(i))); + } + std::vector x_dims = {1, 1, 4, 4}; + std::vector expected_dims = {1, 1, 3, 3}; + std::vector expected_vals = { + MLFloat16(3.5f), MLFloat16(5.0f), MLFloat16(6.0f), + MLFloat16(9.5f), MLFloat16(11.0f), MLFloat16(12.0f), + MLFloat16(13.5f), MLFloat16(15.0f), MLFloat16(16.0f)}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kAclExecutionProvider, kWebGpuExecutionProvider}); +} + +// No-regression control: count_include_pad=1 but ceil_mode=0 keeps the standard MLAS path +// (guard requires ceil_mode==1), so floor-mode output stays correct. +TEST(PoolFp16Test, AveragePool_Floor_CountIncludePad_2d) { + if (DefaultDmlExecutionProvider().get() != nullptr) { + GTEST_SKIP() << "Skipping because of the following error: MLOperatorAuthorImpl.cpp(2100): The parameter is incorrect."; + } + + OpTester test("AveragePool", 11); + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{2, 2}); + test.AddAttribute("pads", std::vector{1, 1, 1, 1}); + test.AddAttribute("kernel_shape", std::vector{3, 3}); + test.AddAttribute("ceil_mode", (int64_t)0); + test.AddAttribute("count_include_pad", (int64_t)1); + + std::vector x_vals; + for (int i = 1; i <= 16; ++i) { + x_vals.push_back(MLFloat16(static_cast(i))); + } + std::vector x_dims = {1, 1, 4, 4}; + std::vector expected_dims = {1, 1, 2, 2}; + std::vector expected_vals = { + MLFloat16(14.0f / 9.0f), MLFloat16(30.0f / 9.0f), + MLFloat16(38.0f / 6.0f), MLFloat16(11.0f)}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", + {kTensorrtExecutionProvider, kAclExecutionProvider, kWebGpuExecutionProvider}); +} + TEST(PoolFp16Test, GlobalAveragePool) { OpTester test("GlobalAveragePool"); diff --git a/onnxruntime/test/providers/cpu/nn/pool_op_test.cc b/onnxruntime/test/providers/cpu/nn/pool_op_test.cc index 95c0c9d35c7f0..99a9a2773e516 100644 --- a/onnxruntime/test/providers/cpu/nn/pool_op_test.cc +++ b/onnxruntime/test/providers/cpu/nn/pool_op_test.cc @@ -9,6 +9,26 @@ namespace onnxruntime { namespace test { +// Execution providers excluded from the four opset-18 AveragePool ceil_mode + +// count_include_pad parity tests below. Shared by all four tests to prevent scoping +// drift. Two rationales: +// - Most EPs (kTensorrt, kNvTensorRTRTX, kAcl, kOpenVINO, kDml, kWebGpu, kDnnl, +// kCoreML, kQnn) do not implement the clamped-window divisor (PyTorch #183528), +// so they return the pre-fix full-kernel-size average and would fail these tests. +// - kCuda / kCudaNHWC DO support the semantics, but these opset-18 cases are +// CPU-reference gate tests (the CUDA path has its own parity tests) and cuDNN-NHWC +// can flap on the 2D case, so they are excluded here too. +// +// NOTE: do not confuse this with kPoolingEpsExcludedFromCeilCipTests (defined ~L1150), which +// has the OPPOSITE kCuda membership. This set is the CPU-reference GATE and therefore INCLUDES +// kCuda/kCudaNHWC in the exclusion list (CPU is the oracle here); that other set EXCLUDES +// kCuda/kCudaNHWC because there CUDA is the tested target. Pick the one matching your intent. +static const std::unordered_set kPoolingEpsExcludedFromCeilCountIncludePadTests = { + kCudaExecutionProvider, kCudaNHWCExecutionProvider, kTensorrtExecutionProvider, + kNvTensorRTRTXExecutionProvider, kAclExecutionProvider, kOpenVINOExecutionProvider, + kDmlExecutionProvider, kWebGpuExecutionProvider, kDnnlExecutionProvider, + kCoreMLExecutionProvider, kQnnExecutionProvider}; + template class PoolTest : public ::testing::Test { }; @@ -1130,6 +1150,11 @@ TEST(PoolTest, AveragePool_19_ceil_count_include_pad_1d) { // 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. +// +// NOTE: do not confuse this with kPoolingEpsExcludedFromCeilCountIncludePadTests (top of file, +// ~L21), which has the OPPOSITE kCuda membership. That set is a CPU-reference GATE and INCLUDES +// kCuda/kCudaNHWC in its exclusions; this set EXCLUDES them because here CUDA/CUDA-NHWC ARE the +// tested targets. Opposite kCuda intent — pick the one matching your test. // --------------------------------------------------------------------------- const std::unordered_set kPoolingEpsExcludedFromCeilCipTests = { kTensorrtExecutionProvider, kNvTensorRTRTXExecutionProvider, kDnnlExecutionProvider, @@ -1378,6 +1403,112 @@ TEST(PoolTest, AveragePool_CUDA_same_lower_asymmetric_1d) { // type-checker rejects such a graph at load. The fp16 case above already exercises the // accumulate-in-float path. +// (a)-gate regression test for the CPU/MLAS AvgPool ceil_mode + count_include_pad bug +// (PyTorch #183528). This is the opset-18 clone of AveragePool_19_ceil_count_include_pad_1d: +// same X and same expected_vals, but at opset 18 the float path routes through MLAS (which +// divided by the full kernel size and produced a wrong average) instead of the v19 reference +// loop. GPU / other EPs are excluded so the test is green the moment the CPU fix lands; the +// CUDA leg is tracked separately as the (b) probe. +TEST(PoolTest, AveragePool_18_ceil_count_include_pad_1d) { + OpTester test("AveragePool", 18); + + 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 = {2.0903f, 4.6493f, 1.6320f, -3.2051f, 4.6975f, 4.7296f, 3.3653f, -1.5815f, -2.3832f, 0.9628f, -1.5899f, -2.6820f, 5.7529f, 7.7346f, -0.8910f, -2.0151f, 0.1313f, -0.5374f}; + std::vector x_dims = {1, 2, 9}; + std::vector expected_dims = {1, 2, 4}; + std::vector expected_vals = {0.73807144f, 2.5655572f, 0.8032287f, -0.09990001f, 0.34911433f, 1.0389f, 1.4536142f, -0.40353334f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCountIncludePadTests); +} + +// 2D opset-18 case for the same bug. Input is the PyTorch #183528 repro: +// x = arange(1, 17).reshape(1, 1, 4, 4), kernel=3, stride=2, pad=1, ceil_mode=1, +// count_include_pad=1. The ceil-mode trailing window ends past input+pad_tail, so MLAS's +// full-kernel divisor gave a wrong average; the reference loop divides by the clamped +// window (in-bounds + real pad cells only). +TEST(PoolTest, AveragePool_18_ceil_count_include_pad_2d) { + OpTester test("AveragePool", 18); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{2, 2}); + test.AddAttribute("pads", std::vector{1, 1, 1, 1}); + test.AddAttribute("kernel_shape", std::vector{3, 3}); + 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, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 16.0f}; + std::vector x_dims = {1, 1, 4, 4}; + std::vector expected_dims = {1, 1, 3, 3}; + std::vector expected_vals = {1.5555556f, 3.3333333f, 2.0f, + 6.3333335f, 11.0f, 6.0f, + 4.5f, 7.5f, 4.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCountIncludePadTests); +} + +// 3D opset-18 case for the same bug, exercising the AveragePool3DTask path. +TEST(PoolTest, AveragePool_18_ceil_count_include_pad_3d) { + OpTester test("AveragePool", 18); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{2, 2, 2}); + test.AddAttribute("pads", std::vector{1, 1, 1, 1, 1, 1}); + test.AddAttribute("kernel_shape", std::vector{3, 3, 3}); + test.AddAttribute("ceil_mode", (int64_t)1); + test.AddAttribute("count_include_pad", (int64_t)1); + + std::vector x_vals(27); + for (int i = 0; i < 27; ++i) { + x_vals[i] = static_cast(i + 1); + } + std::vector x_dims = {1, 1, 3, 3, 3}; + std::vector expected_dims = {1, 1, 2, 2, 2}; + // Ground truth from the CPU v19 reference loop (window clamped to input + real pad). + std::vector expected_vals = {2.2222223f, 2.5185184f, 3.1111112f, 3.4074075f, + 4.888889f, 5.185185f, 5.7777777f, 6.074074f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCountIncludePadTests); +} + +// No-regression guard: with count_include_pad=0 the divisor already counts only in-bounds +// cells, so this combo stays on the MLAS fast path and must remain correct. +TEST(PoolTest, AveragePool_18_ceil_count_exclude_pad_2d) { + OpTester test("AveragePool", 18); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{2, 2}); + test.AddAttribute("pads", std::vector{1, 1, 1, 1}); + test.AddAttribute("kernel_shape", std::vector{3, 3}); + 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, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 16.0f}; + std::vector x_dims = {1, 1, 4, 4}; + std::vector expected_dims = {1, 1, 3, 3}; + // count_include_pad=0: each output divides by the number of in-bounds cells only. + std::vector expected_vals = {3.5f, 5.0f, 6.0f, + 9.5f, 11.0f, 12.0f, + 13.5f, 15.0f, 16.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", kPoolingEpsExcludedFromCeilCountIncludePadTests); +} + TEST(PoolTest, GlobalAveragePool) { OpTester test("GlobalAveragePool");