diff --git a/src/core/shape_inference/include/space_to_batch_shape_inference.hpp b/src/core/shape_inference/include/space_to_batch_shape_inference.hpp index ffeea95253f1a6..af563e2801ccc6 100644 --- a/src/core/shape_inference/include/space_to_batch_shape_inference.hpp +++ b/src/core/shape_inference/include/space_to_batch_shape_inference.hpp @@ -70,7 +70,8 @@ std::vector shape_infer(const SpaceToBatch* op, for (auto idx = spatial_dim_offset; idx < data_rank_size; ++idx) { NODE_VALIDATION_CHECK(op, (*blocks)[idx] > 0, "block_shape values must be greater than 0"); - const auto padded_dim = data_shape[idx] + static_cast((*pads_begin)[idx] + (*pads_end)[idx]); + const auto padded_dim = + data_shape[idx] + static_cast((*pads_begin)[idx]) + static_cast((*pads_end)[idx]); const auto divisor = static_cast((*blocks)[idx]); if (static_cast(padded_dim.get_max_length()) == dim::inf_bound) { diff --git a/src/core/src/op/space_to_batch.cpp b/src/core/src/op/space_to_batch.cpp index 948f5ed824af89..98d533faabf299 100644 --- a/src/core/src/op/space_to_batch.cpp +++ b/src/core/src/op/space_to_batch.cpp @@ -6,10 +6,12 @@ #include #include +#include #include #include #include "itt.hpp" +#include "openvino/core/memory_util.hpp" #include "openvino/core/shape.hpp" #include "openvino/op/util/attr_types.hpp" #include "openvino/op/util/precision_sensitive_attribute.hpp" @@ -100,10 +102,29 @@ bool evaluate(TensorVector& outputs, const TensorVector& inputs) { Shape padded_shape(data_shape.size()); for (size_t i = 0; i < data_shape.size(); ++i) { - padded_shape[i] = data_shape[i] + pads_begin_vec[i] + pads_end_vec[i]; + OPENVINO_ASSERT(pads_begin_vec[i] >= 0, + "SpaceToBatch: pads_begin[", + i, + "] must be non-negative, got ", + pads_begin_vec[i]); + OPENVINO_ASSERT(pads_end_vec[i] >= 0, + "SpaceToBatch: pads_end[", + i, + "] must be non-negative, got ", + pads_end_vec[i]); + const auto pb = static_cast(pads_begin_vec[i]); + const auto pe = static_cast(pads_end_vec[i]); + OPENVINO_ASSERT(data_shape[i] <= std::numeric_limits::max() - pb && + data_shape[i] + pb <= std::numeric_limits::max() - pe, + "SpaceToBatch: padded dimension ", + i, + " overflows"); + padded_shape[i] = data_shape[i] + pb + pe; } - std::vector padded_data(shape_size(padded_shape) * elem_size); + const auto padded_byte_size = ov::util::get_memory_size_safe(data.get_element_type(), padded_shape); + OPENVINO_ASSERT(padded_byte_size.has_value(), "SpaceToBatch: padded shape size overflows"); + std::vector padded_data(*padded_byte_size); reference::pad(static_cast(data.data()), pad_value, padded_data.data(), @@ -122,8 +143,8 @@ bool evaluate(TensorVector& outputs, const TensorVector& inputs) { std::iota(plain_axes_order.begin(), plain_axes_order.end(), 0); std::vector flat_data(padded_data.begin(), padded_data.end()); - std::vector dispersed_data(shape_size(data_shape) * elem_size); - std::vector post_transpose_data(shape_size(data_shape) * elem_size); + std::vector dispersed_data(*padded_byte_size); + std::vector post_transpose_data(*padded_byte_size); for (int64_t block_idx = block_values_size - 1; block_idx >= 0; --block_idx) { int64_t sq_shape_idx = block_values_size - 1; diff --git a/src/plugins/template/tests/functional/op_reference/space_to_batch.cpp b/src/plugins/template/tests/functional/op_reference/space_to_batch.cpp index 5960d6dc9d72fe..69f326234c1c5c 100644 --- a/src/plugins/template/tests/functional/op_reference/space_to_batch.cpp +++ b/src/plugins/template/tests/functional/op_reference/space_to_batch.cpp @@ -6,8 +6,10 @@ #include +#include + #include "base_reference_test.hpp" -#include "openvino/op/constant.hpp" +#include "openvino/core/except.hpp" #include "openvino/op/parameter.hpp" using namespace reference_tests; @@ -158,4 +160,37 @@ INSTANTIATE_TEST_SUITE_P(smoke_SpaceToBatch_With_Hardcoded_Refs, ReferenceSpaceToBatchLayerTest, testing::ValuesIn(generateCombinedParams()), ReferenceSpaceToBatchLayerTest::getTestCaseName); + +class ReferenceSpaceToBatchLayerNegativeTest : public ReferenceSpaceToBatchLayerTest {}; + +TEST_P(ReferenceSpaceToBatchLayerNegativeTest, InvalidPaddingThrows) { + EXPECT_THROW(Exec(), ov::Exception); +} + +std::vector generateNegativeParams() { + constexpr auto I64_MAX = std::numeric_limits::max(); + const reference_tests::Tensor dummyExpected({1}, element::f32, std::vector{0}); + + return { + // INT64_MAX padding causes size_t overflow in padded_shape computation + SpaceToBatchParams(reference_tests::Tensor({1, 4, 4}, element::f32, std::vector(16, 1.0f)), + reference_tests::Tensor({3}, element::i64, std::vector{1, 1, 1}), + reference_tests::Tensor({3}, element::i64, std::vector{0, I64_MAX, I64_MAX}), + reference_tests::Tensor({3}, element::i64, std::vector{0, I64_MAX, I64_MAX}), + dummyExpected, + "padding_overflow"), + // Negative padding values + SpaceToBatchParams(reference_tests::Tensor({1, 1, 2, 2}, element::f32, std::vector{1, 1, 1, 1}), + reference_tests::Tensor({4}, element::i64, std::vector{1, 1, 1, 1}), + reference_tests::Tensor({4}, element::i64, std::vector{0, 0, -1, 0}), + reference_tests::Tensor({4}, element::i64, std::vector{0, 0, 0, 0}), + dummyExpected, + "negative_padding"), + }; +} + +INSTANTIATE_TEST_SUITE_P(smoke_SpaceToBatch_Negative, + ReferenceSpaceToBatchLayerNegativeTest, + testing::ValuesIn(generateNegativeParams()), + ReferenceSpaceToBatchLayerNegativeTest::getTestCaseName); } // namespace