diff --git a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h index 7dfb2770a6979..4c7f2268310ff 100644 --- a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h +++ b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h @@ -216,6 +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() == input_spatial_rank, + "Pooling kernel rank must equal input spatial rank. Got kernel rank: ", + 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 ed65700ccd336..af095a141dfcd 100644 --- a/onnxruntime/test/contrib_ops/maxpool_mask_test.cc +++ b/onnxruntime/test/contrib_ops/maxpool_mask_test.cc @@ -161,5 +161,70 @@ 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"); +} + +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