Enforce per-dim upper bounds for DYNAMIC_BOUND inputs in set_input

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<const int32_t> bound_sizes = tensor_meta->sizes();
+      if (static_cast<size_t>(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<ssize_t>(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),

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> 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> 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> 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> 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> 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);