update set_tensor_data to use MethodMeta and set_input (#524)

Summary:
Pull Request resolved: #524

Switch the util used in all of our examples to use "set_input" instead of get_input, set_data.

Theres still more work to be done on this util like make it follow our style guide, put it in a more specific file, use RAII but thatll be after MVP.

Differential Revision: D49747239

fbshipit-source-id: 2a7c2d518d381ade4f724e69dad744df14058e60

Author: JacobSzwejbka

examples/bundled_executor_runner/bundled_executor_runner.cpp

@@ -212,7 +212,8 @@ int main(int argc, char** argv) {
         status);
   } else {
     // Use ones-initialized inputs.
-    inputs = torch::executor::util::PrepareInputTensors(*method);
+    inputs = torch::executor::util::PrepareInputTensors(
+        *(program->method_meta(method_name)), *method);
   }
   ET_LOG(Info, "Inputs prepared.");
 

examples/executor_runner/executor_runner.cpp

@@ -157,7 +157,8 @@ int main(int argc, char** argv) {
 
   // Prepare the inputs.
   // Use ones-initialized inputs.
-  auto inputs = util::PrepareInputTensors(*method);
+  auto inputs =
+      util::PrepareInputTensors(*(program->method_meta(method_name)), *method);
   ET_LOG(Info, "Inputs prepared.");
 
   // Run the model.

runtime/executor/test/allocation_failure_stress_test.cpp

@@ -84,7 +84,8 @@ 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);
+        torch::executor::util::PrepareInputTensors(
+            *(program_->method_meta("forward")), *method);
     err = method->execute();
     torch::executor::util::FreeInputs(inputs);
     ASSERT_EQ(err, Error::Ok);
@@ -120,7 +121,8 @@ 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);
+        torch::executor::util::PrepareInputTensors(
+            *(program_->method_meta("forward")), *method);
     err = method->execute();
     torch::executor::util::FreeInputs(inputs);
     ASSERT_EQ(err, Error::Ok);

runtime/executor/test/backend_integration_test.cpp

@@ -415,7 +415,8 @@ TEST_P(BackendIntegrationTest, EndToEndTestWithProcessedAsHandle) {
     auto method(std::move(method_res.get()));
     // Execute the model.
     exec_aten::ArrayRef<void*> inputs =
-        torch::executor::util::PrepareInputTensors(method);
+        torch::executor::util::PrepareInputTensors(
+            program->method_meta("forward").get(), method);
     auto err = method.execute();
     torch::executor::util::FreeInputs(inputs);
     EXPECT_EQ(err, Error::Ok);

runtime/executor/test/kernel_integration_test.cpp

@@ -156,7 +156,8 @@ class KernelIntegrationTest : public ::testing::Test {
     method_ = std::make_unique<Method>(std::move(method.get()));
 
     // Set up its inputs.
-    inputs_ = torch::executor::util::PrepareInputTensors(*method_);
+    inputs_ = torch::executor::util::PrepareInputTensors(
+        *(program->method_meta("forward")), *method_);
   }
 
   void TearDown() override {

runtime/executor/test/method_test.cpp

@@ -72,7 +72,8 @@ TEST_F(MethodTest, MoveTest) {
 
   // Can execute the method.
   exec_aten::ArrayRef<void*> inputs =
-      torch::executor::util::PrepareInputTensors(*method);
+      torch::executor::util::PrepareInputTensors(
+          *(programs_["add"]->method_meta("forward")), *method);
   Error err = method->execute();
   ASSERT_EQ(err, Error::Ok);
 
@@ -97,7 +98,8 @@ TEST_F(MethodTest, SetPrimInputTest) {
 
   // Can execute the method.
   exec_aten::ArrayRef<void*> inputs =
-      torch::executor::util::PrepareInputTensors(*method);
+      torch::executor::util::PrepareInputTensors(
+          *(programs_["add"]->method_meta("forward")), *method);
 
   // The args to the method are x, y, alpha. x and y are tensors handled above
   // alpha is a prim.
@@ -191,12 +193,13 @@ TEST_F(MethodTest, AliasedIOTest) {
 // TEST_F(MethodTest, OptionalTensorListDeserialization) {
 //   ManagedMemoryManager mmm(kDefaultNonConstMemBytes,
 //   kDefaultRuntimeMemBytes); Result<Method> method =
-//   index_program_->load_method("forward", &mmm.get());
+//   programs_["index"]->load_method("forward", &mmm.get());
 //   ASSERT_EQ(method.error(), Error::Ok);
 
 //   // Can execute the method.
 //   exec_aten::ArrayRef<void*> inputs =
-//       torch::executor::util::PrepareInputTensors(*method);
+//       torch::executor::util::PrepareInputTensors(*(programs_["index"]->method_meta("forward")),
+//       *method);
 //   Error err = method->execute();
 //   ASSERT_EQ(err, Error::Ok);
 

sdk/runners/executor_runner.cpp

@@ -363,7 +363,8 @@ int main(int argc, char** argv) {
           status);
     } else {
       // Use ones-initialized inputs.
-      inputs = torch::executor::util::PrepareInputTensors(*method);
+      inputs = torch::executor::util::PrepareInputTensors(
+          *(program->method_meta(method_name)), *method);
     }
     ET_LOG(Info, "Inputs prepared.");
 

test/multi_runner.cpp

@@ -89,7 +89,8 @@ class PreparedModel final {
         memory_manager_(non_const_mem_bytes, runtime_mem_bytes),
         method_(load_method_or_die(program_, &memory_manager_.get())),
         has_run_(false) {
-    inputs_ = torch::executor::util::PrepareInputTensors(method_);
+    inputs_ = torch::executor::util::PrepareInputTensors(
+        program_.method_meta(name.c_str()).get(), method_);
   }
 
   void run() {

test/size_test.cpp

@@ -65,7 +65,8 @@ int main(int argc, char** argv) {
 
   // Prepare for inputs
   // It assumes the input is one tensor.
-  auto inputs = torch::executor::util::PrepareInputTensors(*method);
+  auto inputs = torch::executor::util::PrepareInputTensors(
+      *(program->method_meta(method_name)), *method);
 
   ET_LOG(Info, "Inputs prepared.");
 

util/util.h

@@ -9,9 +9,11 @@
 #pragma once
 
 #include <algorithm>
+#include <functional>
 
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/executor/method.h>
+#include <executorch/runtime/executor/method_meta.h>
 #include <executorch/runtime/platform/log.h>
 #ifdef USE_ATEN_LIB
 #include <ATen/ATen.h> // @manual=//caffe2/aten:ATen-core
@@ -61,51 +63,66 @@ inline void FillOnes(Tensor tensor) {
  * @returns An array of pointers that must be passed to `FreeInputs()` after
  *     the Method is no longer needed.
  */
-inline exec_aten::ArrayRef<void*> PrepareInputTensors(const Method& method) {
+inline exec_aten::ArrayRef<void*> PrepareInputTensors(
+    const MethodMeta& method_meta,
+    Method& method) {
   size_t input_size = method.inputs_size();
   size_t num_allocated = 0;
   void** inputs = (void**)malloc(input_size * sizeof(void*));
-#ifdef USE_ATEN_LIB
-  auto deleteByNone = [](void* p) {};
+
   for (size_t i = 0; i < input_size; i++) {
-    if (!method.get_input(i).isTensor()) {
+    if (*method_meta.input_tag(i) != Tag::Tensor) {
       ET_LOG(Info, "input %zu is not a tensor, skipping", i);
       continue;
     }
-    const auto& t = method.get_input(i).toTensor();
-    at::StorageImpl* storage =
-        t.unsafeGetTensorImpl()->unsafe_storage().unsafeGetStorageImpl();
-    if (storage->data_ptr().get() == nullptr) {
-      ET_LOG(Info, "input not initialized.");
-      inputs[num_allocated++] = malloc(t.nbytes());
-      storage->set_data_ptr(at::DataPtr(
-          inputs[num_allocated - 1],
-          inputs[num_allocated - 1],
-          deleteByNone,
-          DeviceType::CPU));
-      storage->set_nbytes(t.nbytes());
-    } else {
-      ET_LOG(Info, "input already initialized, refilling.");
+
+    // Tensor Input. Grab meta data and allocate buffer
+    auto input_meta = method_meta.input_tensor_meta(i);
+    inputs[num_allocated++] = malloc(input_meta->nbytes());
+
+#ifdef USE_ATEN_LIB
+    std::vector<int64_t> at_tensor_sizes;
+    for (auto s : input_meta->sizes()) {
+      at_tensor_sizes.push_back(s);
     }
+    at::Tensor t = at::from_blob(
+        inputs[num_allocated - 1],
+        at_tensor_sizes,
+        at::TensorOptions(input_meta->scalar_type()));
     t.fill_(1.0f);
-  }
-#else
-  for (size_t i = 0; i < input_size; i++) {
-    if (!method.get_input(i).isTensor()) {
-      ET_LOG(Info, "input %zu is not a tensor, skipping", i);
-      continue;
+
+#else // Portable Tensor
+    // The only memory that needs to persist after set_input is called is the
+    // data ptr of the input tensor, and that is only if the Method did not
+    // memory plan buffer space for the inputs and instead is expecting the user
+    // to provide them. Meta data like sizes and dim order are used to ensure
+    // the input aligns with the values expected by the plan, but references to
+    // them are not held onto.
+    std::vector<TensorImpl::SizesType> sizes;
+    std::vector<TensorImpl::DimOrderType> dim_order;
+    for (auto s : input_meta->sizes()) {
+      sizes.push_back(s);
     }
-    const auto& t = method.get_input(i).toTensor();
-    if (t.const_data_ptr() == nullptr) {
-      ET_LOG(Info, "input not initialized.");
-      inputs[num_allocated++] = malloc(t.nbytes());
-      t.set_data(inputs[num_allocated - 1]);
-    } else {
-      ET_LOG(Info, "input already initialized, refilling.");
+    for (auto d : input_meta->dim_order()) {
+      dim_order.push_back(d);
     }
+
+    TensorImpl impl = TensorImpl(
+        input_meta->scalar_type(),
+        input_meta->sizes().size(),
+        sizes.data(),
+        inputs[num_allocated - 1],
+        dim_order.data());
+    Tensor t(&impl);
     FillOnes(t);
-  }
 #endif
+    auto error = method.set_input(t, i);
+    ET_CHECK_MSG(
+        error == Error::Ok,
+        "Error: 0x%" PRIx32 " setting input %zu.",
+        error,
+        i);
+  }
   return {inputs, num_allocated};
 }