Skip to content

Validate Col2Im inputs to prevent heap over-read #28706

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Gopalakrishnan Nallasamy (GopalakrishnanN) merged 8 commits into from
Jun 6, 2026
Merged

Validate Col2Im inputs to prevent heap over-read #28706

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
b532d69
Validate Col2Im inputs to prevent heap over-read
May 28, 2026
0df6250
Use validated sliding block shape in Col2ImNd
May 29, 2026
917a152
Restrict Col2Im negative test to CPU EP
May 29, 2026
3ec71ab
Validate Col2Im strides and dilations
Jun 4, 2026
798c67f
Add Col2Im negative tests for dilated-kernel and column-channel guards
Jun 4, 2026
62a51b2
Align Col2Im strides/dilations validation messages with constructor
Jun 4, 2026
0061d8d
Fix transformer conversion test module loading in staged CI
Jun 5, 2026
efea5e9
Revert "Fix transformer conversion test module loading in staged CI"
Jun 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 43 additions & 8 deletions onnxruntime/core/providers/cpu/tensor/col2im.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,54 +21,89 @@ Status Col2Im<T>::Compute(OpKernelContext* context) const {
const auto* image_shape = context->Input<Tensor>(1);
const auto* kernel_shape = context->Input<Tensor>(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<size_t>(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_;
}

TensorShapeVector pads;
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_;
}

TensorShapeVector strides;
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<int64_t> image_shape_size = 1;
SafeInt<int64_t> kernel_shape_size = 1;
SafeInt<int64_t> expected_col_blocks = 1;
TensorShapeVector adjusted_kernel_shape_dims;
TensorShapeVector sliding_block_shape_dims;
auto image_dims = image_shape->Data<int64_t>();
auto kernel_dims = kernel_shape->Data<int64_t>();
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<int64_t>(dilations[i]) * (kernel_dims[i] - 1) + 1);
const int64_t adjusted_kernel = SafeInt<int64_t>(dilations[i]) * (kernel_dims[i] - 1) + 1;
adjusted_kernel_shape_dims.push_back(adjusted_kernel);
const int64_t padded_extent = SafeInt<int64_t>(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;
Comment thread
GopalakrishnanN marked this conversation as resolved.
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<int64_t>(kernel_shape_size) == 0,
"'input' dim[1] (", col_shape[1],
") must be a positive multiple of prod(block_shape) (", static_cast<int64_t>(kernel_shape_size), ").");
ORT_RETURN_IF_NOT(col_shape[2] == static_cast<int64_t>(expected_col_blocks),
"'input' dim[2] (", col_shape[2],
") does not match the number of sliding blocks (", static_cast<int64_t>(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<int64_t>(kernel_shape_size);
const int64_t col_stride = SafeInt<int64_t>(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<T>();
Expand Down Expand Up @@ -97,7 +132,7 @@ Status Col2Im<T>::Compute(OpKernelContext* context) const {
math::Col2imNd<T, CPUMathUtil, StorageOrder::NCHW>(
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(),
Expand Down
112 changes: 112 additions & 0 deletions onnxruntime/test/providers/cpu/tensor/col2im_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,5 +165,117 @@ TEST(Col2ImOpTest, Simple5dNCHWD) {
test.Run();
}

TEST(Col2ImOpTest, WithStrides5dNCHWD) {
OpTester test("Col2Im", 18);

test.AddAttribute("strides", std::vector<int64_t>{2, 2, 2});
test.AddAttribute("dilations", std::vector<int64_t>{1, 1, 1});
test.AddAttribute("pads", std::vector<int64_t>{0, 0, 0, 0, 0, 0});

std::vector<float> 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<float>("input", {1, 1, 8},
std::vector<float>{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f});
test.AddInput<int64_t>("image_shape", {3}, std::vector<int64_t>{4, 4, 4});
test.AddInput<int64_t>("block_shape", {3}, std::vector<int64_t>{1, 1, 1});
test.AddOutput<float>("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<int64_t>{1, 1});
test.AddAttribute("dilations", std::vector<int64_t>{1, 1});
test.AddAttribute("pads", std::vector<int64_t>{0, 0, 0, 0});

// col holds only 4 spatial positions but image_shape implies 4*4=16.
test.AddInput<float>("input", {1, 1, 4}, std::vector<float>{1.0f, 2.0f, 3.0f, 4.0f});
test.AddInput<int64_t>("image_shape", {2}, std::vector<int64_t>{4, 4});
test.AddInput<int64_t>("block_shape", {2}, std::vector<int64_t>{1, 1});

test.AddOutput<float>("output", {1, 1, 4, 4}, std::vector<float>(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<int64_t>{0, 1});
test.AddAttribute("dilations", std::vector<int64_t>{1, 1});
test.AddAttribute("pads", std::vector<int64_t>{0, 0, 0, 0});

test.AddInput<float>("input", {1, 1, 1}, std::vector<float>{1.0f});
test.AddInput<int64_t>("image_shape", {2}, std::vector<int64_t>{1, 1});
test.AddInput<int64_t>("block_shape", {2}, std::vector<int64_t>{1, 1});
test.AddOutput<float>("output", {1, 1, 1, 1}, std::vector<float>{0.0f});
test.Run(OpTester::ExpectResult::kExpectFailure, "All stride values must be positive",
Comment thread
GopalakrishnanN marked this conversation as resolved.
{kDmlExecutionProvider});
}

TEST(Col2ImOpTest, DilationsMustBePositive) {
OpTester test("Col2Im", 18);

test.AddAttribute("strides", std::vector<int64_t>{1, 1});
test.AddAttribute("dilations", std::vector<int64_t>{0, 1});
test.AddAttribute("pads", std::vector<int64_t>{0, 0, 0, 0});

test.AddInput<float>("input", {1, 2, 1}, std::vector<float>{1.0f, 2.0f});
test.AddInput<int64_t>("image_shape", {2}, std::vector<int64_t>{1, 1});
test.AddInput<int64_t>("block_shape", {2}, std::vector<int64_t>{2, 1});
test.AddOutput<float>("output", {1, 1, 1, 1}, std::vector<float>{0.0f});
test.Run(OpTester::ExpectResult::kExpectFailure, "All dilation values must be positive",
Comment thread
GopalakrishnanN marked this conversation as resolved.
{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<int64_t>{1, 1});
test.AddAttribute("dilations", std::vector<int64_t>{1, 2});
test.AddAttribute("pads", std::vector<int64_t>{0, 0, 0, 0});

test.AddInput<float>("input", {1, 2, 1}, std::vector<float>{1.0f, 2.0f});
test.AddInput<int64_t>("image_shape", {2}, std::vector<int64_t>{1, 1});
test.AddInput<int64_t>("block_shape", {2}, std::vector<int64_t>{1, 2});
test.AddOutput<float>("output", {1, 1, 1, 1}, std::vector<float>{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<int64_t>{1, 1});
test.AddAttribute("dilations", std::vector<int64_t>{1, 1});
test.AddAttribute("pads", std::vector<int64_t>{0, 0, 0, 0});

test.AddInput<float>("input", {1, 3, 1}, std::vector<float>{1.0f, 2.0f, 3.0f});
test.AddInput<int64_t>("image_shape", {2}, std::vector<int64_t>{2, 2});
test.AddInput<int64_t>("block_shape", {2}, std::vector<int64_t>{2, 2});
test.AddOutput<float>("output", {1, 1, 2, 2}, std::vector<float>(4, 0.0f));
test.Run(OpTester::ExpectResult::kExpectFailure, "must be a positive multiple of prod(block_shape)",
{kDmlExecutionProvider});
}

} // namespace test
} // namespace onnxruntime
Loading