diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index 9b8b9e34008..487f0da012b 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -1339,13 +1339,44 @@ Method::set_input(const EValue& input_evalue, size_t input_idx) { input_idx, numel, executorch::runtime::elementSize(t_src.scalar_type())); + auto tensor_meta = this->method_meta().input_tensor_meta(input_idx); + ET_CHECK_OK_OR_RETURN_ERROR( + tensor_meta.error(), + "Error getting tensor metadata at input %" ET_PRIsize_t, + input_idx); + + // For DYNAMIC_BOUND the serialized sizes are the per-dimension upper + // bound, but numel_bound_ is initialized to their product, so the resize + // below only rejects shapes whose total numel exceeds it. A shape that + // trades size between dimensions stays under that product and would + // otherwise be accepted. Rank mismatches are left to resize_tensor. + const auto* s_tensor = serialization_plan_->values() + ->Get(get_input_index(input_idx)) + ->val_as_Tensor(); + if (s_tensor->shape_dynamism() == + executorch_flatbuffer::TensorShapeDynamism::DYNAMIC_BOUND) { + const Span bound_sizes = tensor_meta->sizes(); + if (static_cast(t_src.dim()) == bound_sizes.size()) { + for (size_t i = 0; i < bound_sizes.size(); ++i) { + ET_CHECK_OR_RETURN_ERROR( + t_src.size(i) <= bound_sizes[i], + InvalidArgument, + "Input %" ET_PRIsize_t " dimension %" ET_PRIsize_t + " is %" ET_PRIssize_t " but its upper bound is %" PRId32 ".", + input_idx, + i, + static_cast(t_src.size(i)), + bound_sizes[i]); + } + } + } + // Reset the shape for the Method's input as the size of forwarded input // tensor for shape dynamism. Also is a safety check if need memcpy. ET_CHECK_OK_OR_RETURN_ERROR( resize_tensor(t_dst, t_src.sizes()), "Error resizing tensor at input %" ET_PRIsize_t, input_idx); - auto tensor_meta = this->method_meta().input_tensor_meta(input_idx); if (tensor_meta->is_memory_planned()) { ET_CHECK_OK_OR_RETURN_ERROR( internal::copy_tensor_data(t_dst, t_src), diff --git a/runtime/executor/test/method_test.cpp b/runtime/executor/test/method_test.cpp index 0310c5d648b..bcbe3e64316 100644 --- a/runtime/executor/test/method_test.cpp +++ b/runtime/executor/test/method_test.cpp @@ -325,7 +325,8 @@ TEST_F(MethodTest, SetInputRejectsOverflowingSizes) { ASSERT_EQ(method.error(), Error::Ok); // Create a 2D tensor with enormous sizes. On 32-bit platforms the numel - // multiplication overflows; on 64-bit the resize bounds check rejects it. + // multiplication overflows; on 64-bit the per-dimension bound check in + // set_input rejects it. int32_t sizes[2] = {2000000000, 2000000000}; uint8_t dim_order[2] = {0, 1}; int32_t strides[2] = {1, 1}; @@ -342,6 +343,163 @@ TEST_F(MethodTest, SetInputRejectsOverflowingSizes) { EXPECT_NE(input_err, Error::Ok); } +// The "cat" model (ModuleDynamicCatUnallocatedIO) is DYNAMIC_BOUND with a +// serialized upper-bound shape of (3, 4), so numel_bound_ is 12. +TEST_F(MethodTest, SetInputAcceptsShapesWithinPerDimBound) { + float data[12] = {}; + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["cat"]->load_method("forward", &mmm.get()); + ASSERT_EQ(method.error(), Error::Ok); + + // (3, 4) is the bound itself, and must still be accepted. + const int32_t accepted[][2] = {{3, 4}, {2, 4}, {1, 1}}; + for (const auto& shape : accepted) { + SCOPED_TRACE( + testing::Message() << "shape (" << shape[0] << ", " << shape[1] << ")"); + int32_t sizes[2] = {shape[0], shape[1]}; + uint8_t dim_order[2] = {0, 1}; + int32_t strides[2] = {shape[1], 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, + 2, + sizes, + data, + dim_order, + strides); + EXPECT_EQ( + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0), + Error::Ok); + } +} + +TEST_F(MethodTest, SetInputRejectsPerDimOverflowWithinNumelBound) { + float data[12] = {}; + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["cat"]->load_method("forward", &mmm.get()); + ASSERT_EQ(method.error(), Error::Ok); + + // Both shapes have numel 12, so they satisfy the numel-only bound while + // exceeding the upper bound of 4 on dim 1. + const int32_t rejected[][2] = {{2, 6}, {1, 12}}; + for (const auto& shape : rejected) { + SCOPED_TRACE( + testing::Message() << "shape (" << shape[0] << ", " << shape[1] << ")"); + int32_t sizes[2] = {shape[0], shape[1]}; + uint8_t dim_order[2] = {0, 1}; + int32_t strides[2] = {shape[1], 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, + 2, + sizes, + data, + dim_order, + strides); + EXPECT_EQ( + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0), + Error::InvalidArgument); + } +} + +TEST_F(MethodTest, SetInputRejectsPerDimOverflowBeyondNumelBound) { + float data[16] = {}; + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["cat"]->load_method("forward", &mmm.get()); + ASSERT_EQ(method.error(), Error::Ok); + + // (4, 4) exceeds both dim 0 and numel_bound_. It was already rejected, but + // as NotSupported from resize_tensor; the per-dim check now reports it + // first, as InvalidArgument. + int32_t sizes[2] = {4, 4}; + uint8_t dim_order[2] = {0, 1}; + int32_t strides[2] = {4, 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, 2, sizes, data, dim_order, strides); + + EXPECT_EQ( + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0), + Error::InvalidArgument); +} + +TEST_F(MethodTest, SetInputRejectsRankMismatchOnBoundedInput) { + float data[12] = {}; + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["cat"]->load_method("forward", &mmm.get()); + ASSERT_EQ(method.error(), Error::Ok); + + // The per-dim check skips these, so resize_tensor still reports the rank + // change. Covered here because the check indexes the bound shape by the + // input's rank. + { + int32_t sizes[1] = {12}; + uint8_t dim_order[1] = {0}; + int32_t strides[1] = {1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, + 1, + sizes, + data, + dim_order, + strides); + EXPECT_EQ( + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0), + Error::NotSupported); + } + { + int32_t sizes[3] = {1, 3, 4}; + uint8_t dim_order[3] = {0, 1, 2}; + int32_t strides[3] = {12, 4, 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, + 3, + sizes, + data, + dim_order, + strides); + EXPECT_EQ( + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0), + Error::NotSupported); + } +} + +TEST_F(MethodTest, SetInputStaticShapeBehaviorUnchanged) { + float data[4] = {}; + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["add"]->load_method("forward", &mmm.get()); + ASSERT_EQ(method.error(), Error::Ok); + + // "add" has STATIC (2, 2) inputs, which the per-dim check leaves alone. + { + int32_t sizes[2] = {2, 2}; + uint8_t dim_order[2] = {0, 1}; + int32_t strides[2] = {2, 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, + 2, + sizes, + data, + dim_order, + strides); + EXPECT_EQ( + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0), + Error::Ok); + } + { + int32_t sizes[2] = {4, 1}; + uint8_t dim_order[2] = {0, 1}; + int32_t strides[2] = {1, 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, + 2, + sizes, + data, + dim_order, + strides); + EXPECT_EQ( + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0), + Error::NotSupported); + } +} + TEST_F(MethodTest, ConstantSegmentTest) { // Execute model with constants stored in segment. ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);