diff --git a/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc b/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc index 3f5d07cef4b7c..f8b78e80e7cac 100644 --- a/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc +++ b/onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc @@ -1,11 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -// TODO: x86 release builds produce warning C4723: potential divide by 0 -#ifdef _MSC_VER -#pragma warning(disable : 4723) -#endif - #include "core/providers/cpu/tensor/space_depth_ops.h" #include "core/common/eigen_common_wrapper.h" #include diff --git a/onnxruntime/core/providers/cpu/tensor/space_depth_ops.h b/onnxruntime/core/providers/cpu/tensor/space_depth_ops.h index 9fe8958097e50..3db5f2565226b 100644 --- a/onnxruntime/core/providers/cpu/tensor/space_depth_ops.h +++ b/onnxruntime/core/providers/cpu/tensor/space_depth_ops.h @@ -3,6 +3,7 @@ #pragma once +#include #include #if !defined(SHARED_PROVIDER) && !defined(BUILD_CUDA_EP_AS_PLUGIN) @@ -52,6 +53,10 @@ inline Status InputValidationsAndOutputDimsCalc(int64_t blocksize, input_shape.NumDimensions()); } + if (blocksize <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "SpaceDepth ops require blocksize to be greater than 0."); + } + batch = input_shape[0]; if constexpr (IsNHWC) { input_depth = input_shape[3]; @@ -72,16 +77,26 @@ inline Status InputValidationsAndOutputDimsCalc(int64_t blocksize, return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "SpaceToDepth requires input width to be a multiple of block_size"); } + const auto int64_max = std::numeric_limits::max(); + if (input_depth > int64_max / blocksize || input_depth * blocksize > int64_max / blocksize) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "SpaceToDepth output depth exceeds int64_t limits"); + } + output_depth = input_depth * blocksize * blocksize; output_height = input_height / blocksize; output_width = input_width / blocksize; } else { // DepthToSpace op - if ((input_depth % (blocksize * blocksize) != 0)) { + if (input_depth % blocksize != 0 || (input_depth / blocksize) % blocksize != 0) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "DepthToSpace requires input depth to be a multiple of (block_size * block_size)"); } + const auto int64_max = std::numeric_limits::max(); + if (input_height > int64_max / blocksize || input_width > int64_max / blocksize) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "DepthToSpace output dimensions exceed int64_t limits"); + } + output_depth = input_depth / blocksize / blocksize; output_height = input_height * blocksize; output_width = input_width * blocksize; diff --git a/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc b/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc index 05c4374df1df8..80927a95afb86 100644 --- a/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/space_depth_ops_test.cc @@ -6,6 +6,7 @@ #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "core/framework/allocator.h" #include "core/providers/cpu/tensor/space_depth_ops.h" #include "core/mlas/inc/mlas.h" @@ -19,6 +20,46 @@ class TensorOpTest : public ::testing::Test { using TensorOpTestTypes = ::testing::Types; TYPED_TEST_SUITE(TensorOpTest, TensorOpTestTypes); +TEST(TensorOpTest, RejectsNonPositiveBlocksizeBeforeShapeArithmetic) { + auto allocator = std::make_shared(); + Tensor input(DataTypeImpl::GetType(), TensorShape({1, 1, 2, 2}), allocator); + + for (const auto& [blocksize, is_space_to_depth] : + std::vector>{{0, true}, {-1, false}}) { + int64_t batch = 0; + int64_t input_depth = 0; + int64_t input_height = 0; + int64_t input_width = 0; + int64_t output_depth = 0; + int64_t output_height = 0; + int64_t output_width = 0; + const auto status = space_depth_internal::InputValidationsAndOutputDimsCalc( + blocksize, input, batch, input_depth, input_height, input_width, + output_depth, output_height, output_width, is_space_to_depth); + EXPECT_EQ(status.Code(), common::StatusCode::INVALID_ARGUMENT); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("blocksize to be greater than 0")); + } +} + +TEST(TensorOpTest, RejectsBlocksizeWhoseSquareExceedsInt64) { + auto allocator = std::make_shared(); + Tensor input(DataTypeImpl::GetType(), TensorShape({1, 1, 2, 2}), allocator); + int64_t batch = 0; + int64_t input_depth = 0; + int64_t input_height = 0; + int64_t input_width = 0; + int64_t output_depth = 0; + int64_t output_height = 0; + int64_t output_width = 0; + + const auto status = space_depth_internal::InputValidationsAndOutputDimsCalc( + int64_t{1} << 32, input, batch, input_depth, input_height, input_width, + output_depth, output_height, output_width, false); + + EXPECT_EQ(status.Code(), common::StatusCode::INVALID_ARGUMENT); + EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("multiple of (block_size * block_size)")); +} + TEST(TensorOpTest, SpaceToDepthTest_1) { OpTester test("SpaceToDepth"); constexpr int64_t blocksize = 2;