Require setting all inputs before execution. (pytorch#13207)

Summary:

Method never cared to check all the inputs were properly set before execution and basically used some default allocated memory for unset inputs. Here we introduce a safety check to guarantee users don't forget to set all inputs explicitly.

Reviewed By: JacobSzwejbka

Differential Revision: D79849134

Author: shoumikhin

extension/runner_util/test/inputs_test.cpp

@@ -75,6 +75,9 @@ class InputsTest : public ::testing::Test {
 TEST_F(InputsTest, Smoke) {
   Result<BufferCleanup> input_buffers = prepare_input_tensors(*method_);
   ASSERT_EQ(input_buffers.error(), Error::Ok);
+  auto input_err =
+      method_->set_input(executorch::runtime::EValue(1.0), 2);
+  ASSERT_EQ(input_err, Error::Ok);
 
   // We can't look at the input tensors, but we can check that the outputs make
   // sense after executing the method.

runtime/executor/method.cpp

@@ -407,11 +407,22 @@ Error Method::parse_values(const NamedDataMap* external_data_map) {
   auto flatbuffer_values = serialization_plan_->values();
   ET_CHECK_OR_RETURN_ERROR(
       flatbuffer_values != nullptr, InvalidProgram, "Missing values");
-  size_t n_value = flatbuffer_values->size();
+  const size_t n_value = flatbuffer_values->size();
   values_ = memory_manager_->method_allocator()->allocateList<EValue>(n_value);
   if (values_ == nullptr) {
     return Error::MemoryAllocationFailed;
   }
+  const size_t n_input = inputs_size();
+  if (n_input > 0) {
+    input_set_ =
+        memory_manager_->method_allocator()->allocateList<bool>(n_input);
+    if (input_set_ == nullptr) {
+      return Error::MemoryAllocationFailed;
+    }
+    for (size_t i = 0; i < n_input; ++i) {
+      input_set_[i] = false;
+    }
+  }
 
   // Count the number of tensors marked as EXTERNAL for this method. The actual
   // number of external constants may be smaller, eg. if multiple tensors point
@@ -1076,26 +1087,22 @@ Method::set_input(const EValue& input_evalue, size_t input_idx) {
         executorch::runtime::toString(t_src.scalar_type()));
     // 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.
-    Error err = resize_tensor(t_dst, t_src.sizes());
-    ET_CHECK_OR_RETURN_ERROR(
-        err == Error::Ok,
-        InvalidArgument,
-        "Error setting input %" ET_PRIsize_t ": 0x%" PRIx32,
-        input_idx,
-        static_cast<uint32_t>(err));
-    Error error;
+    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()) {
-      error = internal::copy_tensor_data(t_dst, t_src);
+      ET_CHECK_OK_OR_RETURN_ERROR(
+        internal::copy_tensor_data(t_dst, t_src),
+        "Error setting copying tensor data at input %" ET_PRIsize_t,
+        input_idx);
     } else {
-      error = internal::share_tensor_data(t_dst, t_src);
+      ET_CHECK_OK_OR_RETURN_ERROR(
+        internal::share_tensor_data(t_dst, t_src),
+        "Error setting sharing tensor data at input %" ET_PRIsize_t,
+        input_idx);
     }
-    ET_CHECK_OR_RETURN_ERROR(
-        error == Error::Ok,
-        InvalidArgument,
-        "Error setting data_ptr %" ET_PRIsize_t ": 0x%" PRIx32,
-        input_idx,
-        static_cast<uint32_t>(error));
     // Prims have to be the same as what was traced
   } else if (e.isInt()) {
     ET_CHECK_OR_RETURN_ERROR(
@@ -1163,35 +1170,16 @@ Method::set_input(const EValue& input_evalue, size_t input_idx) {
 
     return Error::InvalidArgument;
   }
+  input_set_[input_idx] = true;
+
   return Error::Ok;
 }
 
 ET_NODISCARD Error
 Method::set_inputs(const executorch::aten::ArrayRef<EValue>& input_evalues) {
-  ET_CHECK_OR_RETURN_ERROR(
-      initialized(),
-      InvalidState,
-      "Inputs can not be set until method has been initialized.");
-
-  ET_CHECK_OR_RETURN_ERROR(
-      step_state_.instr_idx == 0 && step_state_.chain_idx == 0,
-      InvalidState,
-      "Inputs can not be set mid execution.");
-
-  size_t input_size = inputs_size();
-  ET_CHECK_OR_RETURN_ERROR(
-      input_size == input_evalues.size(),
-      InvalidArgument,
-      "The length of given input array (%" ET_PRIsize_t
-      ") must be same as the number of inputs in method (%" ET_PRIsize_t ").",
-      input_evalues.size(),
-      input_size);
-
-  for (size_t i = 0; i < input_size; i++) {
-    Error status = set_input(input_evalues[i], i);
-    if (status != Error::Ok) {
-      return status;
-    }
+  const size_t n_input = inputs_size();
+  for (size_t i = 0; i < n_input; ++i) {
+    ET_CHECK_OK_OR_RETURN_ERROR(set_input(input_evalues[i], i));
   }
   return Error::Ok;
 }
@@ -1284,20 +1272,18 @@ ET_NODISCARD Error Method::get_inputs(EValue* input_evalues, size_t length) {
       initialized(),
       InvalidState,
       "Inputs can not be retrieved until method has been initialized.");
-
+  const size_t n_input = inputs_size();
   ET_CHECK_OR_RETURN_ERROR(
-      length >= inputs_size(),
+      length >= n_input,
       InvalidArgument,
       "The given array is not large enough to hold all inputs.");
 
-  for (size_t i = 0; i < inputs_size(); i++) {
+  for (size_t i = 0; i < n_input; ++i) {
     input_evalues[i] = values_[get_input_index(i)];
   }
-
-  for (size_t i = inputs_size(); i < length; i++) {
+  for (size_t i = n_input; i < length; ++i) {
     input_evalues[i] = EValue();
   }
-
   return Error::Ok;
 }
 
@@ -1545,6 +1531,14 @@ Error Method::execute() {
       initialized(),
       NotSupported,
       "Cannot execute until method has been initialized.");
+  const size_t n_input = inputs_size();
+  for (size_t i = 0; i < n_input; ++i) {
+    ET_CHECK_OR_RETURN_ERROR(
+      input_set_[i],
+      InvalidArgument,
+      "Input %" ET_PRIsize_t " has not been set.",
+      i);
+  }
   ET_LOG(Debug, "Executing method: %s.", method_meta().name());
 
   // Chains are executed sequentially today, but future async designs may

runtime/executor/method.h

@@ -73,6 +73,7 @@ class Method final {
         event_tracer_(rhs.event_tracer_),
         n_value_(rhs.n_value_),
         values_(rhs.values_),
+        input_set_(rhs.input_set_),
         n_delegate_(rhs.n_delegate_),
         delegates_(rhs.delegates_),
         n_chains_(rhs.n_chains_),
@@ -85,6 +86,7 @@ class Method final {
     // anything twice.
     rhs.n_value_ = 0;
     rhs.values_ = nullptr;
+    rhs.input_set_ = nullptr;
     rhs.n_delegate_ = 0;
     rhs.delegates_ = nullptr;
 
@@ -314,6 +316,7 @@ class Method final {
         event_tracer_(event_tracer),
         n_value_(0),
         values_(nullptr),
+        input_set_(nullptr),
         n_delegate_(0),
         delegates_(nullptr),
         n_chains_(0),
@@ -362,6 +365,7 @@ class Method final {
 
   size_t n_value_;
   EValue* values_;
+  bool* input_set_;
 
   size_t n_delegate_;
   BackendDelegate* delegates_;

runtime/executor/test/allocation_failure_stress_test.cpp

@@ -88,6 +88,8 @@ TEST_F(AllocationFailureStressTest, End2EndIncreaseRuntimeMemUntilSuccess) {
     // once load was successful.
     auto input_cleanup = prepare_input_tensors(*method);
     ASSERT_EQ(input_cleanup.error(), Error::Ok);
+    auto input_err = method->set_input(executorch::runtime::EValue(1.0), 2);
+    ASSERT_EQ(input_err, Error::Ok);
     err = method->execute();
     ASSERT_EQ(err, Error::Ok);
   }
@@ -123,6 +125,8 @@ TEST_F(AllocationFailureStressTest, End2EndNonConstantMemUntilSuccess) {
     // once load was successful.
     auto input_cleanup = prepare_input_tensors(*method);
     ASSERT_EQ(input_cleanup.error(), Error::Ok);
+    auto input_err = method->set_input(executorch::runtime::EValue(1.0), 2);
+    ASSERT_EQ(input_err, Error::Ok);
     err = method->execute();
     ASSERT_EQ(err, Error::Ok);
   }

runtime/executor/test/backend_data_separation_test.cpp

@@ -95,6 +95,21 @@ TEST_F(BackendDataSeparationTest, TestSeparation) {
       /*named_data_map=*/linear_data_map_.get());
   ASSERT_EQ(method.error(), Error::Ok);
 
+  // Set a dummy input.
+  int32_t sizes[1] = {3};
+  uint8_t dim_order[1] = {0};
+  int32_t strides[1] = {1};
+  executorch::aten::TensorImpl impl(
+      executorch::aten::ScalarType::Float,
+      1,
+      sizes,
+      nullptr,
+      dim_order,
+      strides);
+  auto input_err = method->set_input(
+      executorch::runtime::EValue(executorch::aten::Tensor(&impl)), 0);
+  ASSERT_EQ(input_err, Error::Ok);
+
   // Can execute the method.
   Error err = method->execute();
   ASSERT_EQ(err, Error::Ok);

runtime/executor/test/backend_integration_test.cpp

@@ -603,6 +603,25 @@ TEST_P(BackendIntegrationTest, GetMethodNameDuringExecuteSuccess) {
   ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
   Result<Method> method = program->load_method("forward", &mmm.get());
   EXPECT_TRUE(method.ok());
+
+  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,
+      nullptr,
+      dim_order,
+      strides);
+  auto input_err = method->set_input(
+      executorch::runtime::EValue(executorch::aten::Tensor(&impl)), 0);
+  input_err = method->set_input(
+      executorch::runtime::EValue(executorch::aten::Tensor(&impl)), 1);
+  input_err = method->set_input(
+      executorch::runtime::EValue(executorch::aten::Tensor(&impl)), 2);
+  ASSERT_EQ(input_err, Error::Ok);
+
   Error err = method->execute();
   ASSERT_EQ(err, Error::Ok);
 }

runtime/executor/test/kernel_integration_test.cpp

@@ -248,6 +248,8 @@ class KernelIntegrationTest : public ::testing::Test {
     ASSERT_EQ(inputs_cleanup.error(), Error::Ok);
     inputs_cleanup_ = std::make_unique<executorch::extension::BufferCleanup>(
         std::move(*inputs_cleanup));
+    auto input_err = method_->set_input(executorch::runtime::EValue(1.0), 2);
+    ASSERT_EQ(input_err, Error::Ok);
   }
 
   void TearDown() override {

runtime/executor/test/method_test.cpp

@@ -104,9 +104,13 @@ TEST_F(MethodTest, MoveTest) {
   Result<Method> method = programs_["add"]->load_method("forward", &mmm.get());
   ASSERT_EQ(method.error(), Error::Ok);
 
-  // Can execute the method.
+  // Set dummy inputs.
   auto input_cleanup = prepare_input_tensors(*method);
   ASSERT_EQ(input_cleanup.error(), Error::Ok);
+  auto input_err = method->set_input(executorch::runtime::EValue(1.0), 2);
+  ASSERT_EQ(input_err, Error::Ok);
+
+  // Can execute the method.
   Error err = method->execute();
   ASSERT_EQ(err, Error::Ok);
 
@@ -312,6 +316,21 @@ TEST_F(MethodTest, ConstantSegmentTest) {
       programs_["add_mul"]->load_method("forward", &mmm.get());
   ASSERT_EQ(method.error(), Error::Ok);
 
+  // Set a dummy input.
+  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,
+      nullptr,
+      dim_order,
+      strides);
+  auto input_err = method->set_input(
+      executorch::runtime::EValue(executorch::aten::Tensor(&impl)), 0);
+  ASSERT_EQ(input_err, Error::Ok);
+
   // Can execute the method.
   Error err = method->execute();
   ASSERT_EQ(err, Error::Ok);
@@ -324,6 +343,21 @@ TEST_F(MethodTest, ConstantBufferTest) {
       programs_["linear_constant_buffer"]->load_method("forward", &mmm.get());
   ASSERT_EQ(method.error(), Error::Ok);
 
+  // Set a dummy input.
+  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,
+      nullptr,
+      dim_order,
+      strides);
+  auto input_err = method->set_input(
+      executorch::runtime::EValue(executorch::aten::Tensor(&impl)), 0);
+  ASSERT_EQ(input_err, Error::Ok);
+
   // Can execute the method.
   Error err = method->execute();
   ASSERT_EQ(err, Error::Ok);
@@ -335,6 +369,21 @@ TEST_F(MethodTest, ProgramDataSeparationTest) {
       "forward", &mmm.get(), nullptr, data_maps_["add_mul_data"].get());
   ASSERT_EQ(method.error(), Error::Ok);
 
+  // Set a dummy input.
+  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,
+      nullptr,
+      dim_order,
+      strides);
+  auto input_err = method->set_input(
+      executorch::runtime::EValue(executorch::aten::Tensor(&impl)), 0);
+  ASSERT_EQ(input_err, Error::Ok);
+
   // Can execute the method.
   Error err = method->execute();
   ASSERT_EQ(err, Error::Ok);
@@ -357,6 +406,21 @@ TEST_F(MethodTest, MethodGetAttributeTest) {
   // expect data to be set
   EXPECT_EQ(res->const_data_ptr(), &data);
 
+  // Set a dummy input.
+  int32_t sizes[1] = {1};
+  uint8_t dim_order[1] = {0};
+  int32_t strides[1] = {1};
+  executorch::aten::TensorImpl impl(
+      executorch::aten::ScalarType::Float,
+      1,
+      sizes,
+      nullptr,
+      dim_order,
+      strides);
+  auto input_err = method->set_input(
+      executorch::runtime::EValue(executorch::aten::Tensor(&impl)), 0);
+  ASSERT_EQ(input_err, Error::Ok);
+
   // Can execute the method.
   Error err = method->execute();
   ASSERT_EQ(err, Error::Ok);