diff --git a/onnxruntime/core/providers/cpu/tensor/col2im.cc b/onnxruntime/core/providers/cpu/tensor/col2im.cc index 26a2e458ffa45..7262ad17e1d3e 100644 --- a/onnxruntime/core/providers/cpu/tensor/col2im.cc +++ b/onnxruntime/core/providers/cpu/tensor/col2im.cc @@ -21,12 +21,22 @@ Status Col2Im::Compute(OpKernelContext* context) const { const auto* image_shape = context->Input(1); const auto* kernel_shape = context->Input(2); + ORT_RETURN_IF_NOT(image_shape->Shape().NumDimensions() == 1, + "'image_shape' input must be a 1-D tensor."); + ORT_RETURN_IF_NOT(kernel_shape->Shape().NumDimensions() == 1, + "'block_shape' input must be a 1-D tensor."); + ORT_RETURN_IF_NOT(image_shape->Shape().Size() == kernel_shape->Shape().Size(), + "'image_shape' and 'block_shape' must have the same number of elements."); + size_t image_dim_number = onnxruntime::narrow(image_shape->Shape().Size()); + ORT_RETURN_IF_NOT(image_dim_number > 0, "'image_shape' must have at least one element."); + TensorShapeVector dilations; if (dilations_.empty()) { dilations.resize(image_dim_number, 1); } else { - ORT_ENFORCE(dilations_.size() == image_dim_number, "size of 'dilations' attribute, if provided, should equal to the number of image dimmensions."); + ORT_RETURN_IF_NOT(dilations_.size() == image_dim_number, + "size of 'dilations' attribute, if provided, should equal to the number of image dimensions."); dilations = dilations_; } @@ -34,7 +44,8 @@ Status Col2Im::Compute(OpKernelContext* context) const { if (pads_.empty()) { pads.resize(image_dim_number * 2, 0); } else { - ORT_ENFORCE(pads_.size() == 2 * image_dim_number, "size of 'pads' attribute, if provided, should equal to twice the number of image dimmensions."); + ORT_RETURN_IF_NOT(pads_.size() == 2 * image_dim_number, + "size of 'pads' attribute, if provided, should equal to twice the number of image dimensions."); pads = pads_; } @@ -42,33 +53,57 @@ Status Col2Im::Compute(OpKernelContext* context) const { if (strides_.empty()) { strides.resize(image_dim_number, 1); } else { - ORT_ENFORCE(strides_.size() == image_dim_number, "size of 'strides' attribute, if provided, should equal to the number of image dimmensions."); + ORT_RETURN_IF_NOT(strides_.size() == image_dim_number, + "size of 'strides' attribute, if provided, should equal to the number of image dimensions."); strides = strides_; } SafeInt image_shape_size = 1; SafeInt kernel_shape_size = 1; + SafeInt expected_col_blocks = 1; TensorShapeVector adjusted_kernel_shape_dims; + TensorShapeVector sliding_block_shape_dims; auto image_dims = image_shape->Data(); auto kernel_dims = kernel_shape->Data(); for (size_t i = 0; i < image_dim_number; ++i) { + ORT_RETURN_IF_NOT(image_dims[i] > 0, "All 'image_shape' values must be positive."); + ORT_RETURN_IF_NOT(kernel_dims[i] > 0, "All 'block_shape' values must be positive."); + ORT_RETURN_IF_NOT(strides[i] > 0, "All stride values must be positive."); + ORT_RETURN_IF_NOT(dilations[i] > 0, "All dilation values must be positive."); image_shape_size *= image_dims[i]; kernel_shape_size *= kernel_dims[i]; - adjusted_kernel_shape_dims.push_back(SafeInt(dilations[i]) * (kernel_dims[i] - 1) + 1); + const int64_t adjusted_kernel = SafeInt(dilations[i]) * (kernel_dims[i] - 1) + 1; + adjusted_kernel_shape_dims.push_back(adjusted_kernel); + const int64_t padded_extent = SafeInt(image_dims[i]) + pads[i] + pads[i + image_dim_number]; + ORT_RETURN_IF_NOT(padded_extent >= adjusted_kernel, + "Padded image extent is smaller than the dilated kernel for spatial dimension ", i, "."); + const int64_t sliding_blocks = (padded_extent - adjusted_kernel) / strides[i] + 1; + sliding_block_shape_dims.push_back(sliding_blocks); + expected_col_blocks *= sliding_blocks; } - ORT_ENFORCE(kernel_shape_size > 0, "kernel_shape_size must be positive"); + ORT_RETURN_IF_NOT(kernel_shape_size > 0, "kernel_shape_size must be positive"); + TensorShape col_shape = col_tensor->Shape(); + ORT_RETURN_IF_NOT(col_shape.NumDimensions() == 3, + "'input' tensor must be 3-D with shape (N, C * prod(block_shape), L)."); + ORT_RETURN_IF_NOT(col_shape[1] > 0 && col_shape[1] % static_cast(kernel_shape_size) == 0, + "'input' dim[1] (", col_shape[1], + ") must be a positive multiple of prod(block_shape) (", static_cast(kernel_shape_size), ")."); + ORT_RETURN_IF_NOT(col_shape[2] == static_cast(expected_col_blocks), + "'input' dim[2] (", col_shape[2], + ") does not match the number of sliding blocks (", static_cast(expected_col_blocks), + ") implied by 'image_shape', 'block_shape', 'pads', 'strides', and 'dilations'."); + const auto N = col_shape[0]; const int64_t C = col_shape[1] / static_cast(kernel_shape_size); const int64_t col_stride = SafeInt(C) * image_shape_size; TensorShape adjusted_kernel_shape(adjusted_kernel_shape_dims); const int64_t col_data_stride = col_shape.SizeFromDimension(1); - TensorShapeVector batched_image_shape_dims, adjusted_image_shape_dims; + TensorShapeVector batched_image_shape_dims; batched_image_shape_dims.insert(batched_image_shape_dims.begin(), {N, C}); for (size_t i = 0; i < image_dim_number; ++i) { batched_image_shape_dims.push_back(image_dims[i]); - adjusted_image_shape_dims.push_back(image_dims[i] - adjusted_kernel_shape[i] + 1); } TensorShape batched_image_shape(batched_image_shape_dims); T* image_data = context->Output(0, batched_image_shape)->template MutableData(); @@ -97,7 +132,7 @@ Status Col2Im::Compute(OpKernelContext* context) const { math::Col2imNd( col_data + image_id * col_data_stride, image_dims, - adjusted_image_shape_dims.data(), + sliding_block_shape_dims.data(), kernel_shape_size * C, image_shape_size * C, adjusted_kernel_shape.GetDims().data(), diff --git a/onnxruntime/test/providers/cpu/tensor/col2im_test.cc b/onnxruntime/test/providers/cpu/tensor/col2im_test.cc index 3a4539024e5a9..05e862866db23 100644 --- a/onnxruntime/test/providers/cpu/tensor/col2im_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/col2im_test.cc @@ -165,5 +165,117 @@ TEST(Col2ImOpTest, Simple5dNCHWD) { test.Run(); } +TEST(Col2ImOpTest, WithStrides5dNCHWD) { + OpTester test("Col2Im", 18); + + test.AddAttribute("strides", std::vector{2, 2, 2}); + test.AddAttribute("dilations", std::vector{1, 1, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0, 0, 0}); + + std::vector output(64, 0.0f); + output[0] = 1.0f; + output[2] = 2.0f; + output[8] = 3.0f; + output[10] = 4.0f; + output[32] = 5.0f; + output[34] = 6.0f; + output[40] = 7.0f; + output[42] = 8.0f; + + test.AddInput("input", {1, 1, 8}, + std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); + test.AddInput("image_shape", {3}, std::vector{4, 4, 4}); + test.AddInput("block_shape", {3}, std::vector{1, 1, 1}); + test.AddOutput("output", {1, 1, 4, 4, 4}, output); + test.Run(); +} + +// Regression test for a heap buffer over-read in Col2Im when 'image_shape' implies more +// sliding-block positions than the column tensor actually contains. The kernel must reject +// the inputs instead of reading past the column allocation. +TEST(Col2ImOpTest, ImageShapeLargerThanColumnTensor) { + OpTester test("Col2Im", 18); + + test.AddAttribute("strides", std::vector{1, 1}); + test.AddAttribute("dilations", std::vector{1, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0}); + + // col holds only 4 spatial positions but image_shape implies 4*4=16. + test.AddInput("input", {1, 1, 4}, std::vector{1.0f, 2.0f, 3.0f, 4.0f}); + test.AddInput("image_shape", {2}, std::vector{4, 4}); + test.AddInput("block_shape", {2}, std::vector{1, 1}); + + test.AddOutput("output", {1, 1, 4, 4}, std::vector(16, 0.0f)); + // Restrict to the CPU kernel: other execution providers (e.g. DML) reject the + // inconsistent shapes during graph partitioning with a different message. + test.Run(OpTester::ExpectResult::kExpectFailure, "does not match the number of sliding blocks", + {kDmlExecutionProvider}); +} + +TEST(Col2ImOpTest, StridesMustBePositive) { + OpTester test("Col2Im", 18); + + test.AddAttribute("strides", std::vector{0, 1}); + test.AddAttribute("dilations", std::vector{1, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0}); + + test.AddInput("input", {1, 1, 1}, std::vector{1.0f}); + test.AddInput("image_shape", {2}, std::vector{1, 1}); + test.AddInput("block_shape", {2}, std::vector{1, 1}); + test.AddOutput("output", {1, 1, 1, 1}, std::vector{0.0f}); + test.Run(OpTester::ExpectResult::kExpectFailure, "All stride values must be positive", + {kDmlExecutionProvider}); +} + +TEST(Col2ImOpTest, DilationsMustBePositive) { + OpTester test("Col2Im", 18); + + test.AddAttribute("strides", std::vector{1, 1}); + test.AddAttribute("dilations", std::vector{0, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0}); + + test.AddInput("input", {1, 2, 1}, std::vector{1.0f, 2.0f}); + test.AddInput("image_shape", {2}, std::vector{1, 1}); + test.AddInput("block_shape", {2}, std::vector{2, 1}); + test.AddOutput("output", {1, 1, 1, 1}, std::vector{0.0f}); + test.Run(OpTester::ExpectResult::kExpectFailure, "All dilation values must be positive", + {kDmlExecutionProvider}); +} + +// The dilated kernel must fit within the padded image. With dilations={1,2} and block_shape={1,2} +// the second dimension's dilated kernel extent is 3, larger than the padded image extent of 1, so +// the kernel must reject the inputs instead of computing a negative/zero sliding-block count. +TEST(Col2ImOpTest, PaddedImageSmallerThanDilatedKernel) { + OpTester test("Col2Im", 18); + + test.AddAttribute("strides", std::vector{1, 1}); + test.AddAttribute("dilations", std::vector{1, 2}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0}); + + test.AddInput("input", {1, 2, 1}, std::vector{1.0f, 2.0f}); + test.AddInput("image_shape", {2}, std::vector{1, 1}); + test.AddInput("block_shape", {2}, std::vector{1, 2}); + test.AddOutput("output", {1, 1, 1, 1}, std::vector{0.0f}); + test.Run(OpTester::ExpectResult::kExpectFailure, + "Padded image extent is smaller than the dilated kernel", {kDmlExecutionProvider}); +} + +// The column tensor's channel dimension (dim[1]) must be a positive multiple of prod(block_shape). +// Here prod(block_shape) = 4 but dim[1] = 3, so the kernel must reject the inputs. +TEST(Col2ImOpTest, ColumnChannelsNotMultipleOfBlock) { + OpTester test("Col2Im", 18); + + test.AddAttribute("strides", std::vector{1, 1}); + test.AddAttribute("dilations", std::vector{1, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0}); + + test.AddInput("input", {1, 3, 1}, std::vector{1.0f, 2.0f, 3.0f}); + test.AddInput("image_shape", {2}, std::vector{2, 2}); + test.AddInput("block_shape", {2}, std::vector{2, 2}); + test.AddOutput("output", {1, 1, 2, 2}, std::vector(4, 0.0f)); + test.Run(OpTester::ExpectResult::kExpectFailure, "must be a positive multiple of prod(block_shape)", + {kDmlExecutionProvider}); +} + } // namespace test } // namespace onnxruntime