From c87fbb0de6cd80008c8d52a94306eca11ae28800 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Wed, 24 Jun 2026 22:34:05 +0000 Subject: [PATCH 1/2] contrib: validate MaxpoolWithMask kernel rank matches input spatial rank Residual hardening for MaxpoolWithMask::Compute: the 1D/2D/3D dispatch is selected by kernel_shape.size() and reads x_shape[2..4] / output_dims[2..4] accordingly, but the only rank guard was x_shape.NumDimensions() >= 3. A kernel_shape rank that exceeds the input spatial rank (e.g. rank-3 X with a 2D kernel_shape) caused out-of-bounds reads of x_shape[3]/x_shape[4] and output_dims[3]/output_dims[4]. ONNX shape inference catches this for shaped ONNX models, but the kernel must validate at Compute time for defense in depth (ORT-format / shapeless models bypass shape inference). Add ORT_RETURN_IF_NOT(kernel rank == input spatial rank) alongside the existing mask/input guards, and a kExpectFailure regression test that bypasses ONNX shape inference via AddShapeToTensorData(false) to exercise the Compute-time guard directly. Agent-signed-off: Developer (0b207188) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cpu/maxpool_with_mask.h | 5 +++ .../test/contrib_ops/maxpool_mask_test.cc | 33 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h index 7dfb2770a6979..7b7eda12d6ab2 100644 --- a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h +++ b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h @@ -216,6 +216,11 @@ class MaxpoolWithMask : public OpKernel, public PoolBase { "Mask and input spatial dimensions mismatch at dimension ", i, ": mask=", m_shape[i], " input=", x_shape[i]); } + // The pooling kernel rank drives the 1D/2D/3D dispatch below, which reads x_shape[2..4] and + // output_dims[2..4]. Require it to match the input spatial rank so those reads stay in bounds. + ORT_RETURN_IF_NOT(pool_attrs_.kernel_shape.size() == x_shape.NumDimensions() - 2, + "Pooling kernel rank must equal input spatial rank. Got kernel rank: ", + pool_attrs_.kernel_shape.size(), " input spatial rank: ", x_shape.NumDimensions() - 2); TensorShapeVector pads = pool_attrs_.pads; TensorShapeVector kernel_shape = pool_attrs_.kernel_shape; diff --git a/onnxruntime/test/contrib_ops/maxpool_mask_test.cc b/onnxruntime/test/contrib_ops/maxpool_mask_test.cc index ed65700ccd336..7da97302d0f19 100644 --- a/onnxruntime/test/contrib_ops/maxpool_mask_test.cc +++ b/onnxruntime/test/contrib_ops/maxpool_mask_test.cc @@ -161,5 +161,38 @@ TEST(ContribOpTest, MaxPoolWithMask_MaskEmptyBatchDim) { "Mask N and C dimensions must be greater than 0"); } +TEST(ContribOpTest, MaxPoolWithMask_KernelRankMismatch) { + OpTester test("MaxpoolWithMask", 1, onnxruntime::kMSDomain); + + // AddShapeToTensorData(false) omits input shape from the graph so ONNX shape inference is bypassed + // (convPoolShapeInference returns early when hasInputShape is false). This lets the model pass + // Graph::Resolve() and reach Compute() where the kernel-rank guard fires. + test.AddShapeToTensorData(false); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{1, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0}); + // 2D kernel_shape, but X has only one spatial dimension (rank 3). + test.AddAttribute("kernel_shape", std::vector{8, 8}); + + // Input X has shape {1, 1, 8} (rank 3 => one spatial dim) + std::vector x_dims = {1, 1, 8}; + std::vector x_vals(8, 1.0f); + + // Mask M matches X shape so the earlier spatial/rank guards pass and we reach the kernel-rank guard. + std::vector m_dims = {1, 1, 8}; + std::vector m_vals(8, 1); + + // Placeholder output shape and values (not validated since we expect failure) + std::vector expected_dims = {1, 1, 1, 1}; + std::vector expected_vals = {1.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddInput("M", m_dims, m_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(BaseTester::ExpectResult::kExpectFailure, + "Pooling kernel rank must equal input spatial rank"); +} + } // namespace test } // namespace onnxruntime From b13dc903449938577740fc3470df9f13139a6466 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Wed, 24 Jun 2026 22:40:33 +0000 Subject: [PATCH 2/2] contrib: reject unsupported MaxpoolWithMask pooling rank and DRY rank check Follow-up polish on the kernel-rank guard. The 1D/2D/3D dispatch switch has a default case that already returns INVALID_ARGUMENT, so an oversized kernel rank was never a silent/uninitialized-output bug, but the default produced a vague 'Unsupported pooling size :' message late in Compute after allocating the output. Add an explicit early guard requiring the input spatial rank (== pooling kernel rank) to be in {1, 2, 3}, with a clear message, alongside the existing guards and before output allocation. Extract the repeated x_shape.NumDimensions() - 2 into a single input_spatial_rank local (defined after the rank >= 3 guard so it cannot underflow) and reuse it. Add a MaxPoolWithMask_KernelRankTooLarge regression test (rank-6 X + 4D kernel_shape) that passes the equality guard but trips the new supported-rank guard, asserting the specific message substring. Agent-signed-off: Developer (0b207188) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cpu/maxpool_with_mask.h | 9 ++++-- .../test/contrib_ops/maxpool_mask_test.cc | 32 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h index 7b7eda12d6ab2..4c7f2268310ff 100644 --- a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h +++ b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h @@ -216,11 +216,16 @@ class MaxpoolWithMask : public OpKernel, public PoolBase { "Mask and input spatial dimensions mismatch at dimension ", i, ": mask=", m_shape[i], " input=", x_shape[i]); } + // x_shape.NumDimensions() >= 3 is guaranteed above, so this subtraction cannot underflow. + const size_t input_spatial_rank = x_shape.NumDimensions() - 2; // The pooling kernel rank drives the 1D/2D/3D dispatch below, which reads x_shape[2..4] and // output_dims[2..4]. Require it to match the input spatial rank so those reads stay in bounds. - ORT_RETURN_IF_NOT(pool_attrs_.kernel_shape.size() == x_shape.NumDimensions() - 2, + ORT_RETURN_IF_NOT(pool_attrs_.kernel_shape.size() == input_spatial_rank, "Pooling kernel rank must equal input spatial rank. Got kernel rank: ", - pool_attrs_.kernel_shape.size(), " input spatial rank: ", x_shape.NumDimensions() - 2); + pool_attrs_.kernel_shape.size(), " input spatial rank: ", input_spatial_rank); + // Only 1D/2D/3D pooling is implemented by the dispatch below; a larger rank would match no case. + ORT_RETURN_IF_NOT(input_spatial_rank >= 1 && input_spatial_rank <= 3, + "Only 1D, 2D, and 3D pooling are supported. Got input spatial rank: ", input_spatial_rank); TensorShapeVector pads = pool_attrs_.pads; TensorShapeVector kernel_shape = pool_attrs_.kernel_shape; diff --git a/onnxruntime/test/contrib_ops/maxpool_mask_test.cc b/onnxruntime/test/contrib_ops/maxpool_mask_test.cc index 7da97302d0f19..af095a141dfcd 100644 --- a/onnxruntime/test/contrib_ops/maxpool_mask_test.cc +++ b/onnxruntime/test/contrib_ops/maxpool_mask_test.cc @@ -194,5 +194,37 @@ TEST(ContribOpTest, MaxPoolWithMask_KernelRankMismatch) { "Pooling kernel rank must equal input spatial rank"); } +TEST(ContribOpTest, MaxPoolWithMask_KernelRankTooLarge) { + OpTester test("MaxpoolWithMask", 1, onnxruntime::kMSDomain); + + // Bypass ONNX shape inference (see MaxPoolWithMask_KernelRankMismatch) so the model reaches Compute(). + test.AddShapeToTensorData(false); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{1, 1, 1, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0, 0, 0, 0, 0}); + // 4D kernel_shape with a matching rank-6 X (4 spatial dims): passes the kernel-rank equality guard + // but exceeds the supported 1D/2D/3D pooling ranks. + test.AddAttribute("kernel_shape", std::vector{2, 2, 2, 2}); + + // Input X has shape {1, 1, 2, 2, 2, 2} (rank 6 => four spatial dims) + std::vector x_dims = {1, 1, 2, 2, 2, 2}; + std::vector x_vals(16, 1.0f); + + // Mask M matches X so the earlier guards and the kernel-rank equality guard all pass. + std::vector m_dims = {1, 1, 2, 2, 2, 2}; + std::vector m_vals(16, 1); + + // Placeholder output shape and values (not validated since we expect failure) + std::vector expected_dims = {1, 1, 1, 1}; + std::vector expected_vals = {1.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddInput("M", m_dims, m_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(BaseTester::ExpectResult::kExpectFailure, + "Only 1D, 2D, and 3D pooling are supported"); +} + } // namespace test } // namespace onnxruntime