Remove all uses of PrepareInputTensors

Differential Revision: D61750837

Pull Request resolved: #4910

Author: dbort

runtime/executor/test/allocation_failure_stress_test.cpp

@@ -11,27 +11,28 @@
 #include <memory>
 
 #include <executorch/extension/data_loader/file_data_loader.h>
+#include <executorch/extension/runner_util/inputs.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/executor/method.h>
 #include <executorch/runtime/executor/program.h>
 #include <executorch/runtime/executor/test/managed_memory_manager.h>
 #include <executorch/runtime/platform/runtime.h>
-#include <executorch/util/util.h>
 
 #include <gtest/gtest.h>
 
 using namespace ::testing;
 using exec_aten::ArrayRef;
 using exec_aten::Scalar;
 using exec_aten::Tensor;
+using executorch::extension::FileDataLoader;
+using executorch::extension::prepare_input_tensors;
 using executorch::runtime::Error;
 using executorch::runtime::MemoryAllocator;
 using executorch::runtime::MemoryManager;
 using executorch::runtime::Method;
 using executorch::runtime::Program;
 using executorch::runtime::Result;
 using executorch::runtime::testing::ManagedMemoryManager;
-using torch::executor::util::FileDataLoader;
 
 constexpr size_t kDefaultNonConstMemBytes = 32 * 1024U;
 constexpr size_t kDefaultRuntimeMemBytes = 32 * 1024U;
@@ -85,10 +86,9 @@ TEST_F(AllocationFailureStressTest, End2EndIncreaseRuntimeMemUntilSuccess) {
 
     // Execution does not use the runtime allocator, so it should always succeed
     // once load was successful.
-    exec_aten::ArrayRef<void*> inputs =
-        torch::executor::util::PrepareInputTensors(*method);
+    auto input_cleanup = prepare_input_tensors(*method);
+    ASSERT_EQ(input_cleanup.error(), Error::Ok);
     err = method->execute();
-    torch::executor::util::FreeInputs(inputs);
     ASSERT_EQ(err, Error::Ok);
   }
   EXPECT_GT(num_load_failures, 0) << "Expected at least some failures";
@@ -121,10 +121,9 @@ TEST_F(AllocationFailureStressTest, End2EndNonConstantMemUntilSuccess) {
 
     // Execution does not use the runtime allocator, so it should always succeed
     // once load was successful.
-    exec_aten::ArrayRef<void*> inputs =
-        torch::executor::util::PrepareInputTensors(*method);
+    auto input_cleanup = prepare_input_tensors(*method);
+    ASSERT_EQ(input_cleanup.error(), Error::Ok);
     err = method->execute();
-    torch::executor::util::FreeInputs(inputs);
     ASSERT_EQ(err, Error::Ok);
   }
   EXPECT_GT(num_load_failures, 0) << "Expected at least some failures";

runtime/executor/test/backend_integration_test.cpp

@@ -14,6 +14,7 @@
 
 #include <executorch/extension/data_loader/buffer_data_loader.h>
 #include <executorch/extension/data_loader/file_data_loader.h>
+#include <executorch/extension/runner_util/inputs.h>
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/result.h>
@@ -23,7 +24,6 @@
 #include <executorch/runtime/platform/runtime.h>
 #include <executorch/test/utils/DeathTest.h>
 #include <executorch/test/utils/alignment.h>
-#include <executorch/util/util.h>
 
 #include <gtest/gtest.h>
 
@@ -454,10 +454,9 @@ TEST_P(BackendIntegrationTest, EndToEndTestWithProcessedAsHandle) {
     EXPECT_FALSE(spy_loader.WasFreed(init_processed->data()));
     auto method(std::move(method_res.get()));
     // Execute the model.
-    exec_aten::ArrayRef<void*> inputs =
-        torch::executor::util::PrepareInputTensors(method);
+    auto input_cleanup = executorch::extension::prepare_input_tensors(method);
+    ASSERT_EQ(input_cleanup.error(), Error::Ok);
     auto err = method.execute();
-    torch::executor::util::FreeInputs(inputs);
     EXPECT_EQ(err, Error::Ok);
 
     // Check that the processed buffer was passed to execute() as the handle.

runtime/executor/test/kernel_integration_test.cpp

@@ -13,6 +13,7 @@
 #include <memory>
 
 #include <executorch/extension/data_loader/file_data_loader.h>
+#include <executorch/extension/runner_util/inputs.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/result.h>
 #include <executorch/runtime/executor/method.h>
@@ -22,7 +23,6 @@
 #include <executorch/runtime/kernel/operator_registry.h>
 #include <executorch/runtime/platform/compiler.h>
 #include <executorch/runtime/platform/runtime.h>
-#include <executorch/util/util.h>
 
 #include <gtest/gtest.h>
 
@@ -158,12 +158,15 @@ class KernelIntegrationTest : public ::testing::Test {
     method_ = std::make_unique<Method>(std::move(method.get()));
 
     // Set up its inputs.
-    inputs_ = torch::executor::util::PrepareInputTensors(*method_);
+    auto inputs_cleanup =
+        executorch::extension::prepare_input_tensors(*method_);
+    ASSERT_EQ(inputs_cleanup.error(), Error::Ok);
+    inputs_cleanup_ = std::make_unique<executorch::extension::BufferCleanup>(
+        std::move(*inputs_cleanup));
   }
 
   void TearDown() override {
-    torch::executor::util::FreeInputs(inputs_);
-    inputs_ = {};
+    inputs_cleanup_.reset();
   }
 
  private:
@@ -173,7 +176,7 @@ class KernelIntegrationTest : public ::testing::Test {
   // Must outlive method_
   std::unique_ptr<Program> program_;
   std::unique_ptr<ManagedMemoryManager> mmm_;
-  ArrayRef<void*> inputs_;
+  std::unique_ptr<executorch::extension::BufferCleanup> inputs_cleanup_;
 
  protected:
   // An executable method that will call the kernel associated with control_.

runtime/executor/test/method_test.cpp

@@ -10,17 +10,18 @@
 #include <filesystem>
 
 #include <executorch/extension/data_loader/file_data_loader.h>
+#include <executorch/extension/runner_util/inputs.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/executor/method.h>
 #include <executorch/runtime/executor/program.h>
 #include <executorch/runtime/executor/test/managed_memory_manager.h>
 #include <executorch/runtime/platform/runtime.h>
 #include <executorch/test/utils/DeathTest.h>
-#include <executorch/util/util.h>
 #include <gtest/gtest.h>
 
 using namespace ::testing;
 using exec_aten::ArrayRef;
+using executorch::extension::prepare_input_tensors;
 using executorch::runtime::Error;
 using executorch::runtime::EValue;
 using executorch::runtime::Method;
@@ -80,8 +81,8 @@ TEST_F(MethodTest, MoveTest) {
   ASSERT_EQ(method.error(), Error::Ok);
 
   // Can execute the method.
-  exec_aten::ArrayRef<void*> inputs =
-      torch::executor::util::PrepareInputTensors(*method);
+  auto input_cleanup = prepare_input_tensors(*method);
+  ASSERT_EQ(input_cleanup.error(), Error::Ok);
   Error err = method->execute();
   ASSERT_EQ(err, Error::Ok);
 
@@ -95,8 +96,6 @@ TEST_F(MethodTest, MoveTest) {
   // Can execute the new method.
   err = new_method.execute();
   ASSERT_EQ(err, Error::Ok);
-
-  torch::executor::util::FreeInputs(inputs);
 }
 
 TEST_F(MethodTest, GetInputTests) {
@@ -173,8 +172,8 @@ TEST_F(MethodTest, SetPrimInputTest) {
   ASSERT_EQ(method.error(), Error::Ok);
 
   // Can execute the method.
-  exec_aten::ArrayRef<void*> inputs =
-      torch::executor::util::PrepareInputTensors(*method);
+  auto input_cleanup = prepare_input_tensors(*method);
+  ASSERT_EQ(input_cleanup.error(), Error::Ok);
 
   // The args to the method are x, y, alpha. x and y are tensors handled above
   // alpha is a prim.
@@ -189,8 +188,6 @@ TEST_F(MethodTest, SetPrimInputTest) {
 
   Error err = method->execute();
   EXPECT_EQ(err, Error::Ok);
-
-  torch::executor::util::FreeInputs(inputs);
 }
 
 TEST_F(MethodTest, MethodMetaTest) {
@@ -297,28 +294,28 @@ TEST_F(MethodTest, ConstantBufferTest) {
   ASSERT_EQ(err, Error::Ok);
 }
 
-// TODO(T161163608): Test is disabled due to a resize bug in tensor_index_out of
-// the portable op lib
-
-// TEST_F(MethodTest, OptionalTensorListDeserialization) {
-//   ManagedMemoryManager mmm(kDefaultNonConstMemBytes,
-//   kDefaultRuntimeMemBytes); Result<Method> method =
-//   index_program_->load_method("forward", &mmm.get());
-//   ASSERT_EQ(method.error(), Error::Ok);
+/*
+ * TODO(T161163608): Test is disabled due to a resize bug in tensor_index_out of
+ * the portable op lib
 
-//   // Can execute the method.
-//   exec_aten::ArrayRef<void*> inputs =
-//       executorch::runtime::util::PrepareInputTensors(*method);
-//   Error err = method->execute();
-//   ASSERT_EQ(err, Error::Ok);
+TEST_F(MethodTest, OptionalTensorListDeserialization) {
+  ManagedMemoryManager mmm(kDefaultNonConstMemBytes,
+  kDefaultRuntimeMemBytes); Result<Method> method =
+  index_program_->load_method("forward", &mmm.get());
+  ASSERT_EQ(method.error(), Error::Ok);
 
-//   EXPECT_EQ(method->inputs_size(), 1);
+  // Can execute the method.
+  auto input_cleanup = prepare_input_tensors(*method);
+  ASSERT_EQ(input_cleanup.error(), Error::Ok);
+  Error err = method->execute();
+  ASSERT_EQ(err, Error::Ok);
 
-//   auto outputs = method->get_output(0);
-//   EXPECT_EQ(outputs.toTensor().dim(), 3);
-//   EXPECT_EQ(outputs.toTensor().size(0), 5);
-//   EXPECT_EQ(outputs.toTensor().size(1), 2);
-//   EXPECT_EQ(outputs.toTensor().size(2), 10);
+  EXPECT_EQ(method->inputs_size(), 1);
 
-//   executorch::runtime::util::FreeInputs(inputs);
-// }
+  auto outputs = method->get_output(0);
+  EXPECT_EQ(outputs.toTensor().dim(), 3);
+  EXPECT_EQ(outputs.toTensor().size(0), 5);
+  EXPECT_EQ(outputs.toTensor().size(1), 2);
+  EXPECT_EQ(outputs.toTensor().size(2), 10);
+}
+*/

runtime/executor/test/targets.bzl

@@ -120,7 +120,7 @@ def define_common_targets(is_fbcode = False):
                 "//executorch/runtime/executor:program",
                 "//executorch/kernels/portable:generated_lib",
                 "//executorch/extension/data_loader:file_data_loader",
-                "//executorch/util:util",
+                "//executorch/extension/runner_util:inputs",
             ],
             env = modules_env,
         )
@@ -133,8 +133,8 @@ def define_common_targets(is_fbcode = False):
             deps = [
                 ":managed_memory_manager",
                 "//executorch/runtime/executor:program",
-                "//executorch/util:util",
                 "//executorch/extension/data_loader:file_data_loader",
+                "//executorch/extension/runner_util:inputs",
                 "//executorch/kernels/portable:generated_lib",
             ],
             env = modules_env,
@@ -189,12 +189,12 @@ def define_common_targets(is_fbcode = False):
             deps = [
                 ":managed_memory_manager",
                 "//executorch/extension/data_loader:file_data_loader",
+                "//executorch/extension/runner_util:inputs",
                 "//executorch/runtime/core:core",
                 "//executorch/runtime/executor:program",
                 "//executorch/runtime/kernel:kernel_runtime_context",
                 "//executorch/runtime/kernel:operator_registry",
                 "//executorch/runtime/platform:platform",
-                "//executorch/util:util",
             ],
             env = modules_env,
         )
@@ -210,7 +210,7 @@ def define_common_targets(is_fbcode = False):
                 "//executorch/runtime/executor:program",
                 "//executorch/extension/data_loader:buffer_data_loader",
                 "//executorch/extension/data_loader:file_data_loader",
-                "//executorch/util:util",
+                "//executorch/extension/runner_util:inputs",
             ],
             env = {
                 # The tests use these vars to find the program files to load.