Reset inputs after each execution. (pytorch#13229)

Summary:
Pull Request resolved: pytorch#13229

.

Reviewed By: JacobSzwejbka

Differential Revision: D79850580

Author: shoumikhin

runtime/executor/method.cpp

@@ -1451,6 +1451,10 @@ Error Method::reset_execution() {
       InvalidState,
       "Cannot reset until EndOfMethod has been reached.");
   step_state_ = StepState{0, 0};
+  const auto n_input = inputs_size();
+  for (size_t i = 0; i < n_input; ++i) {
+    input_set_[i] = false;
+  }
   return Error::Ok;
 }
 

runtime/executor/test/kernel_integration_test.cpp

@@ -294,6 +294,14 @@ TEST_F(KernelIntegrationTest, KernelHookIsCalled) {
   EXPECT_EQ(err, Error::Ok);
   EXPECT_EQ(control_->call_count, 1);
 
+  // Set up inputs again.
+  auto inputs_cleanup =
+      executorch::extension::prepare_input_tensors(*method_);
+  ASSERT_EQ(inputs_cleanup.error(), Error::Ok);
+  auto input_err =
+      method_->set_input(executorch::runtime::EValue(1.0), 2);
+  ASSERT_EQ(input_err, Error::Ok);
+
   // Calling it again bumps the count.
   err = method_->execute();
   EXPECT_EQ(err, Error::Ok);
@@ -336,6 +344,14 @@ TEST_F(KernelIntegrationTest, DefaultPlatformMemoryAllocator) {
   EXPECT_EQ(control_->call_count, 1);
   EXPECT_EQ(control_->total_allocated_size, 4);
 
+  // Set up inputs again.
+  auto inputs_cleanup =
+      executorch::extension::prepare_input_tensors(*method_);
+  ASSERT_EQ(inputs_cleanup.error(), Error::Ok);
+  auto input_err =
+      method_->set_input(executorch::runtime::EValue(1.0), 2);
+  ASSERT_EQ(input_err, Error::Ok);
+
   control_->temp_memory_size = 8;
   // This is not a simulation. This actually allocates memory, using the
   // default platform memory allocator.
@@ -371,6 +387,14 @@ TEST_F(KernelTempMemoryAllocatorIntegrationTest, UsingTempMemoryAllocator) {
   EXPECT_EQ(temp_allocator_->number_of_resets, 1);
   EXPECT_EQ(temp_allocator_->currently_allocated_size, 0);
 
+  // Set up inputs again.
+  auto inputs_cleanup =
+      executorch::extension::prepare_input_tensors(*method_);
+  ASSERT_EQ(inputs_cleanup.error(), Error::Ok);
+  auto input_err =
+      method_->set_input(executorch::runtime::EValue(1.0), 2);
+  ASSERT_EQ(input_err, Error::Ok);
+
   control_->temp_memory_size = 8;
   err = method_->execute();
   EXPECT_EQ(err, Error::Ok);

runtime/executor/test/method_test.cpp

@@ -114,6 +114,13 @@ TEST_F(MethodTest, MoveTest) {
   Error err = method->execute();
   ASSERT_EQ(err, Error::Ok);
 
+  // Set dummy inputs again.
+  auto input_cleanup2 = prepare_input_tensors(*method);
+  ASSERT_EQ(input_cleanup2.error(), Error::Ok);
+  input_err =
+      method->set_input(executorch::runtime::EValue(1.0), 2);
+  ASSERT_EQ(input_err, Error::Ok);
+
   // Move into a new Method.
   Method new_method(std::move(method.get()));