Set inputs directly on Method without caching. (pytorch#13215)

Summary: .

Differential Revision: D79850621

Author: shoumikhin

extension/module/module.cpp

@@ -210,7 +210,6 @@ runtime::Error Module::load_method(
         method_holder.memory_manager.get(),
         event_tracer ? event_tracer : this->event_tracer(),
         data_map_.get()));
-    method_holder.inputs.resize(method_holder.method->inputs_size());
     methods_.emplace(method_name, std::move(method_holder));
   }
   return runtime::Error::Ok;
@@ -233,28 +232,10 @@ runtime::Result<std::vector<runtime::EValue>> Module::execute(
     const std::vector<runtime::EValue>& input_values) {
   ET_CHECK_OK_OR_RETURN_ERROR(load_method(method_name));
   auto& method = methods_.at(method_name).method;
-  auto& inputs = methods_.at(method_name).inputs;
-
-  ET_CHECK_OR_RETURN_ERROR(
-      input_values.size() <= inputs.size(),
-      InvalidArgument,
-      "input size: %zu does not match method input size: %zu",
-      input_values.size(),
-      inputs.size());
-  for (size_t i = 0; i < input_values.size(); ++i) {
-    if (!input_values[i].isNone()) {
-      inputs[i] = input_values[i];
-    }
+  for (auto index = 0; index < input_values.size(); ++index) {
+    ET_CHECK_OK_OR_RETURN_ERROR(method->set_input(input_values[index], index));
   }
-  for (size_t i = 0; i < inputs.size(); ++i) {
-    ET_CHECK_OR_RETURN_ERROR(
-        !inputs[i].isNone(), InvalidArgument, "input %zu is none", i);
-  }
-  ET_CHECK_OK_OR_RETURN_ERROR(
-      method->set_inputs(executorch::aten::ArrayRef<runtime::EValue>(
-          inputs.data(), inputs.size())));
   ET_CHECK_OK_OR_RETURN_ERROR(method->execute());
-
   const auto outputs_size = method->outputs_size();
   std::vector<runtime::EValue> outputs(outputs_size);
   ET_CHECK_OK_OR_RETURN_ERROR(
@@ -268,23 +249,17 @@ runtime::Error Module::set_input(
     const runtime::EValue& input_value,
     size_t input_index) {
   ET_CHECK_OK_OR_RETURN_ERROR(load_method(method_name));
-  methods_.at(method_name).inputs.at(input_index) = input_value;
-  return runtime::Error::Ok;
+  auto& method = methods_.at(method_name).method;
+  return method->set_input(input_value, input_index);
 }
 
 runtime::Error Module::set_inputs(
     const std::string& method_name,
     const std::vector<runtime::EValue>& input_values) {
   ET_CHECK_OK_OR_RETURN_ERROR(load_method(method_name));
-  auto& inputs = methods_.at(method_name).inputs;
-  ET_CHECK_OR_RETURN_ERROR(
-      inputs.size() == input_values.size(),
-      InvalidArgument,
-      "input size: %zu does not match method input size: %zu",
-      input_values.size(),
-      inputs.size());
-  inputs = input_values;
-  return runtime::Error::Ok;
+  auto& method = methods_.at(method_name).method;
+  return method->set_inputs(executorch::aten::ArrayRef<runtime::EValue>(
+      input_values.data(), input_values.size()));
 }
 
 runtime::Error Module::set_output(

extension/module/module.h

@@ -522,7 +522,6 @@ class Module {
     std::unique_ptr<runtime::HierarchicalAllocator> planned_memory;
     std::unique_ptr<runtime::MemoryManager> memory_manager;
     std::unique_ptr<Method> method;
-    std::vector<runtime::EValue> inputs;
   };
 
   std::string file_path_;

extension/module/test/module_test.cpp

@@ -267,7 +267,7 @@ TEST_F(ModuleTest, TestForward) {
   EXPECT_TENSOR_CLOSE(result->at(0).toTensor(), *expected.get());
 
   auto tensor2 = make_tensor_ptr({2, 2}, {2.f, 3.f, 4.f, 5.f});
-  const auto result2 = module->forward({tensor2, tensor2});
+  const auto result2 = module->forward({tensor2, tensor2, 1.0});
   EXPECT_EQ(result2.error(), Error::Ok);
 
   const auto expected2 = make_tensor_ptr({2, 2}, {4.f, 6.f, 8.f, 10.f});