update set_tensor_data to use MethodMeta and set_input

runtime/executor/method.cpp

@@ -1100,6 +1100,17 @@ Error Method::execute() {
   return experimental_reset_execution();
 }
 
+MethodMeta Method::method_meta() const {
+  auto name = serialization_plan_->name()->c_str();
+  auto method_meta = program_->method_meta(name);
+  ET_CHECK_MSG(
+      method_meta.ok(),
+      "Internal error: method_meta(%s) returned 0x%" PRIx32,
+      name,
+      static_cast<uint32_t>(method_meta.error()));
+  return method_meta.get();
+}
+
 size_t Method::values_size() const {
   return n_value_;
 }

runtime/executor/method.h

@@ -12,6 +12,7 @@
 #include <executorch/runtime/core/event_tracer.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/executor/memory_manager.h>
+#include <executorch/runtime/executor/method_meta.h>
 #include <executorch/runtime/platform/compiler.h>
 
 // Forward declare flatbuffer types. This is a public header and must not
@@ -88,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().
    *
@@ -109,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
@@ -196,17 +198,43 @@ class Method final {
    */
   __ET_NODISCARD Error experimental_reset_execution();
 
-  size_t values_size() const;
-  const EValue& get_value(size_t i) const;
-  EValue& mutable_value(size_t i);
+  /**
+   * Returns the MethodMeta that corresponds to the calling Method.
+   */
+  MethodMeta method_meta() const;
+
+  /**
+   * Returns the number of inputs the Method expects.
+   */
   size_t inputs_size() const;
-  size_t get_input_index(size_t i) const;
-  const EValue& get_input(size_t i) const;
-  EValue& mutable_input(size_t i);
+
+  /**
+   * Returns the number of outputs the Method returns.
+   */
   size_t outputs_size() const;
-  size_t get_output_index(size_t i) const;
+
+  /**
+   * Retrieves the output at the specified index.
+   */
   const EValue& get_output(size_t i) const;
-  EValue& mutable_output(size_t i);
+
+  __ET_DEPRECATED size_t values_size() const;
+  __ET_DEPRECATED const EValue& get_value(size_t i) const;
+  __ET_DEPRECATED EValue& mutable_value(size_t i);
+  /// DEPRECATED: Use MethodMeta instead to access metadata, and set_input to
+  /// update Method inputs.
+  __ET_DEPRECATED size_t get_input_index(size_t i) const;
+  /// DEPRECATED: Use MethodMeta instead to access metadata, and set_input to
+  /// update Method inputs.
+  __ET_DEPRECATED const EValue& get_input(size_t i) const;
+  /// DEPRECATED: Use MethodMeta instead to access metadata, and set_input to
+  /// update Method inputs.
+  __ET_DEPRECATED EValue& mutable_input(size_t i);
+  __ET_DEPRECATED size_t get_output_index(size_t i) const;
+  /// DEPRECATED: Use MethodMeta instead to access metadata, and get_output to
+  /// retrieve Method outputs.
+  __ET_DEPRECATED EValue& mutable_output(size_t i);
+
   ~Method();
 
  private:

runtime/executor/test/method_test.cpp

@@ -116,6 +116,17 @@ TEST_F(MethodTest, SetPrimInputTest) {
   torch::executor::util::FreeInputs(inputs);
 }
 
+TEST_F(MethodTest, MethodMetaTest) {
+  ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
+  Result<Method> method = programs_["add"]->load_method("forward", &mmm.get());
+  ASSERT_EQ(method.error(), Error::Ok);
+
+  auto method_meta = method->method_meta();
+
+  EXPECT_EQ(method_meta.num_inputs(), method->inputs_size());
+  EXPECT_EQ(method_meta.num_outputs(), method->outputs_size());
+}
+
 TEST_F(MethodTest, AliasedIOTest) {
   // TODO(T163238401)
   ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);

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};
 }