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.

Reviewed By: dbort

Differential Revision: D49747239

fbshipit-source-id: 8124b7941e4a6a9232f1ce173e4aeab9b4d2572b

Author: JacobSzwejbka

runtime/executor/method.h

@@ -89,17 +89,18 @@ class Method final {
   }
 
   /**
-   * Sets a specific method input to the provided value.
+   * Sets the internal input value to be equivalent to the to the provided
+   * value.
    *
-   * NOTE: Based on the memory plan of the method, the inputs may not have
-   * buffer space pre-allocated for them, in this case the executor will alias
-   * the memory of the tensors provided as inputs here, so the user should take
-   * care that the life span of this memory outlasts the executor forward.
+   * @param[in] input_evalue The evalue to copy into the method input. If the
+   *     evalue is a tensor, the data is copied in most cases, so the tensor
+   *     passed in here does not always need to outlive this call. But there is
+   *     a case where the Method will keep a pointer to the tensor's data.
+   *     Based on the memory plan of the method, the inputs may not have
+   *     buffer space pre-allocated for them. In this case the executor will
+   *     alias the memory of the tensors provided as inputs here rather then
+   *     deepcopy the input into the memory planned arena.
    *
-   * @param[in] input_evalue The value to set the input to. The type of this
-   *     must match the type of the corresponding input. If this value is a
-   *     tensor, attempts to allow dynamic shape, but the dtype must always
-   *     agree.
    * @param[in] input_idx Zero-based index of the input to set. Must be less
    *     than the value returned by inputs_size().
    *
@@ -110,7 +111,7 @@ class Method final {
   /**
    * Sets the values of all method inputs.
    *
-   * See NOTE on set_input().
+   * See set_input() for a more detailed description of the behavior.
    *
    * @param[in] input_evalues The new values for all of the method inputs. The
    *     type of each element must match the type of corresponding input. If the

test/size_test.cpp

@@ -12,7 +12,6 @@
 #include <executorch/runtime/platform/log.h>
 #include <executorch/runtime/platform/profiler.h>
 #include <executorch/runtime/platform/runtime.h>
-#include <executorch/util/util.h>
 #include <stdio.h>
 
 using namespace torch::executor;
@@ -65,7 +64,13 @@ int main(int argc, char** argv) {
 
   // Prepare for inputs
   // It assumes the input is one tensor.
-  auto inputs = torch::executor::util::PrepareInputTensors(*method);
+  float data[] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0};
+  Tensor::SizesType sizes[] = {6};
+  Tensor::DimOrderType dim_order[] = {0};
+  TensorImpl impl(ScalarType::Float, 1, sizes, data, dim_order);
+  Tensor t(&impl);
+  Error set_input_error = method->set_input(t, 0);
+  ET_CHECK(set_input_error == Error::Ok);
 
   ET_LOG(Info, "Inputs prepared.");
 
@@ -90,7 +95,6 @@ int main(int argc, char** argv) {
       ET_LOG(Info, "%f", data_output[j]);
     }
   }
-  torch::executor::util::FreeInputs(inputs);
   prof_result_t prof_result;
   EXECUTORCH_DUMP_PROFILE_RESULTS(&prof_result);
   if (prof_result.num_bytes != 0) {

test/targets.bzl

@@ -7,7 +7,6 @@ SIZE_TEST_SOURCES = [
 SIZE_TEST_DEPS = [
     "//executorch/runtime/executor:program",
     "//executorch/extension/data_loader:file_data_loader",
-    "//executorch/util:util",
 ]
 
 def define_common_targets():

util/util.h

@@ -12,6 +12,7 @@
 
 #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 +62,77 @@ 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(Method& method) {
+  auto method_meta = method.method_meta();
   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 tensor_meta = method_meta.input_tensor_meta(i);
+    inputs[num_allocated++] = malloc(tensor_meta->nbytes());
+
+#ifdef USE_ATEN_LIB
+    std::vector<int64_t> at_tensor_sizes;
+    for (auto s : tensor_meta->sizes()) {
+      at_tensor_sizes.push_back(s);
     }
+    at::Tensor t = at::from_blob(
+        inputs[num_allocated - 1],
+        at_tensor_sizes,
+        at::TensorOptions(tensor_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.
+
+    TensorImpl::SizesType* sizes = static_cast<TensorImpl::SizesType*>(
+        malloc(sizeof(TensorImpl::SizesType) * tensor_meta->sizes().size()));
+    TensorImpl::DimOrderType* dim_order =
+        static_cast<TensorImpl::DimOrderType*>(malloc(
+            sizeof(TensorImpl::DimOrderType) *
+            tensor_meta->dim_order().size()));
+
+    for (size_t size_idx = 0; size_idx < tensor_meta->sizes().size();
+         size_idx++) {
+      sizes[size_idx] = tensor_meta->sizes()[size_idx];
     }
-    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 (size_t dim_idx = 0; dim_idx < tensor_meta->dim_order().size();
+         dim_idx++) {
+      dim_order[dim_idx] = tensor_meta->dim_order()[dim_idx];
     }
+
+    TensorImpl impl = TensorImpl(
+        tensor_meta->scalar_type(),
+        tensor_meta->sizes().size(),
+        sizes,
+        inputs[num_allocated - 1],
+        dim_order);
+    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);
+#ifndef USE_ATEN_LIB // Portable Tensor
+    free(sizes);
+    free(dim_order);
+#endif
+  }
   return {inputs, num_allocated};
 }