MethodMeta

runtime/executor/method_meta.cpp

@@ -0,0 +1,188 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/util/scalar_type_util.h>
+#include <executorch/runtime/core/result.h>
+#include <executorch/runtime/core/span.h>
+#include <executorch/runtime/core/tag.h>
+#include <executorch/runtime/executor/method_meta.h>
+#include <executorch/schema/program_generated.h>
+
+namespace torch {
+namespace executor {
+
+namespace {
+Result<Tag> get_tag(
+    flatbuffers::Vector<flatbuffers::Offset<executorch_flatbuffer::EValue>>::
+        return_type serialization_value,
+    size_t index) {
+  switch (serialization_value->val_type()) {
+    case executorch_flatbuffer::KernelTypes::Int: {
+      return Tag::Int;
+    } break;
+    case executorch_flatbuffer::KernelTypes::Double: {
+      return Tag::Double;
+    } break;
+    case executorch_flatbuffer::KernelTypes::Bool: {
+      return Tag::Bool;
+    } break;
+    case executorch_flatbuffer::KernelTypes::String: {
+      return Tag::String;
+    } break;
+    case executorch_flatbuffer::KernelTypes::Tensor: {
+      return Tag::Tensor;
+    } break;
+    default:
+      ET_LOG(
+          Error,
+          "Invalid tag: %zu input: %zu",
+          (size_t)serialization_value->val_type(),
+          index);
+      return Error::Internal;
+  }
+}
+
+size_t calculate_nbytes(
+    Span<const int32_t> sizes,
+    exec_aten::ScalarType scalar_type) {
+  ssize_t n = 1;
+  for (ssize_t i = 0; i < sizes.size(); i++) {
+    n *= sizes[i];
+  }
+  return n * sizeof_scalar_type(scalar_type);
+}
+
+} // namespace
+
+TensorInfo::TensorInfo(
+    Span<const int32_t> sizes,
+    Span<const uint8_t> dim_order,
+    exec_aten::ScalarType scalar_type)
+    : sizes_(sizes),
+      dim_order_(dim_order),
+      scalar_type_(scalar_type),
+      nbytes_(calculate_nbytes(sizes_, scalar_type_)) {}
+
+Span<const int32_t> TensorInfo::sizes() const {
+  return sizes_;
+}
+
+Span<const uint8_t> TensorInfo::dim_order() const {
+  return dim_order_;
+}
+
+exec_aten::ScalarType TensorInfo::scalar_type() const {
+  return scalar_type_;
+}
+
+size_t TensorInfo::nbytes() const {
+  return nbytes_;
+}
+
+MethodMeta::MethodMeta(const executorch_flatbuffer::ExecutionPlan* s_plan)
+    : s_plan_(s_plan) {}
+
+const char* MethodMeta::name() const {
+  return s_plan_->name()->c_str();
+}
+
+size_t MethodMeta::num_inputs() const {
+  return s_plan_->inputs()->size();
+}
+
+Result<Tag> MethodMeta::input_tag(size_t index) const {
+  auto num_inputs = this->num_inputs();
+  ET_CHECK_OR_RETURN_ERROR(
+      index >= 0 && index < num_inputs,
+      InvalidArgument,
+      "index %zu out of range. num_inputs: %zu",
+      index,
+      num_inputs);
+  auto input_index = s_plan_->inputs()->Get(index);
+  auto serialization_value = s_plan_->values()->Get(input_index);
+  return get_tag(serialization_value, index);
+}
+
+Result<TensorInfo> MethodMeta::input_tensor_meta(size_t index) const {
+  auto tag = this->input_tag(index);
+  if (!tag.ok()) {
+    return tag.error();
+  }
+  ET_CHECK_OR_RETURN_ERROR(
+      tag.get() == Tag::Tensor,
+      InvalidArgument,
+      "Tag: %zu input: %zu is not Tensor",
+      (size_t)tag.get(),
+      index);
+  auto input_index = s_plan_->inputs()->Get(index);
+  auto tensor_value = s_plan_->values()->Get(input_index)->val_as_Tensor();
+  return TensorInfo(
+      Span<const int32_t>(
+          tensor_value->sizes()->data(), tensor_value->sizes()->size()),
+      Span<const uint8_t>(
+          tensor_value->dim_order()->data(), tensor_value->dim_order()->size()),
+      static_cast<exec_aten::ScalarType>(tensor_value->scalar_type()));
+}
+
+size_t MethodMeta::num_outputs() const {
+  return s_plan_->outputs()->size();
+}
+
+Result<Tag> MethodMeta::output_tag(size_t index) const {
+  auto num_outputs = this->num_outputs();
+  ET_CHECK_OR_RETURN_ERROR(
+      index >= 0 && index < num_outputs,
+      InvalidArgument,
+      "index %zu out of range. num_outputs: %zu",
+      index,
+      num_outputs);
+  auto input_index = s_plan_->outputs()->Get(index);
+  auto serialization_value = s_plan_->values()->Get(input_index);
+  return get_tag(serialization_value, index);
+}
+
+Result<TensorInfo> MethodMeta::output_tensor_meta(size_t index) const {
+  auto tag = this->output_tag(index);
+  if (!tag.ok()) {
+    return tag.error();
+  }
+  ET_CHECK_OR_RETURN_ERROR(
+      tag.get() == Tag::Tensor,
+      InvalidArgument,
+      "Tag: %zu output: %zu is not Tensor",
+      (size_t)tag.get(),
+      index);
+  auto input_index = s_plan_->outputs()->Get(index);
+  auto tensor_value = s_plan_->values()->Get(input_index)->val_as_Tensor();
+  return TensorInfo(
+      Span<const int32_t>(
+          tensor_value->sizes()->data(), tensor_value->sizes()->size()),
+      Span<const uint8_t>(
+          tensor_value->dim_order()->data(), tensor_value->dim_order()->size()),
+      static_cast<exec_aten::ScalarType>(tensor_value->scalar_type()));
+}
+
+size_t MethodMeta::num_non_const_buffers() const {
+  return s_plan_->non_const_buffer_sizes()->size();
+}
+
+Result<int64_t> MethodMeta::non_const_buffer_size(size_t index) const {
+  auto num_buffers = this->num_non_const_buffers();
+  ET_CHECK_OR_RETURN_ERROR(
+      index >= 0 && index < num_buffers,
+      InvalidArgument,
+      "index %zu out of range. num_buffers: %zu",
+      index,
+      num_buffers);
+  return s_plan_->non_const_buffer_sizes()->Get(index);
+}
+
+} // namespace executor
+} // namespace torch

runtime/executor/method_meta.h

@@ -0,0 +1,190 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/result.h>
+#include <executorch/runtime/core/span.h>
+#include <executorch/runtime/core/tag.h>
+
+// Forward declare flatbuffer types. This is a public header and must not
+// include the generated flatbuffer header.
+namespace executorch_flatbuffer {
+struct ExecutionPlan;
+} // namespace executorch_flatbuffer
+
+namespace torch {
+namespace executor {
+
+/**
+ * Metadata about a specific tensor of an Executorch Program.
+ *
+ * The program used to create the MethodMeta object that created this
+ * TensorInfo must outlive this TensorInfo.
+ */
+class TensorInfo final {
+ public:
+  TensorInfo() = delete;
+  TensorInfo(const TensorInfo&) = default;
+  TensorInfo(TensorInfo&&) = default;
+  TensorInfo& operator=(const TensorInfo&) = default;
+  TensorInfo& operator=(TensorInfo&& other) = default;
+  ~TensorInfo() = default;
+
+  /**
+   * Returns the sizes of the tensor.
+   */
+  Span<const int32_t> sizes() const;
+
+  /**
+   * Returns the dim order of the tensor.
+   */
+  Span<const uint8_t> dim_order() const;
+
+  /**
+   * Returns the scalar type of the input/output.
+   */
+  exec_aten::ScalarType scalar_type() const;
+
+  /**
+   * Returns the size of the tensor in bytes.
+   */
+  size_t nbytes() const;
+
+ private:
+  // Let MethodMeta create TensorInfo.
+  friend class MethodMeta;
+
+  TensorInfo(
+      Span<const int32_t> sizes,
+      Span<const uint8_t> dim_order,
+      exec_aten::ScalarType scalar_type);
+
+  /**
+   * The sizes of the tensor.
+   *
+   * NOTE: References data from the Program, so the Program must outlive the
+   * TensorInfo.
+   */
+  Span<const int32_t> sizes_;
+
+  /**
+   * The dim order of the tensor.
+   *
+   * NOTE: References data from the Program, so the Program must outlive the
+   * TensorInfo.
+   */
+  Span<const uint8_t> dim_order_;
+
+  /// The scalar type of the tensor.
+  exec_aten::ScalarType scalar_type_;
+
+  /// The size in bytes of the tensor.
+  size_t nbytes_;
+};
+
+/**
+ * Describes a a method in an Executorch program.
+ *
+ * The program used to create a MethodMeta object must outlive the MethodMeta.
+ * It is separate from Method so that this information can be accessed without
+ * paying the initialization cost of loading the full Method.
+ */
+class MethodMeta final {
+ public:
+  MethodMeta() = delete;
+  MethodMeta(const MethodMeta&) = default;
+  MethodMeta(MethodMeta&&) = default;
+  MethodMeta& operator=(const MethodMeta&) = default;
+  MethodMeta& operator=(MethodMeta&& other) = default;
+  ~MethodMeta() = default;
+
+  /**
+   * Get the name of this method.
+   *
+   * @returns The method name.
+   */
+  const char* name() const;
+
+  /**
+   * Get the number of inputs to this method.
+   *
+   * @returns The number of inputs.
+   */
+  size_t num_inputs() const;
+
+  /**
+   * Get the tag of the specified input.
+   *
+   * @param[in] index The index of the input to look up.
+   * @returns The tag of input, can only be [Tensor, Int, Bool, Double, String].
+   */
+  Result<Tag> input_tag(size_t index) const;
+
+  /**
+   * Get metadata about the specified input.
+   *
+   * @param[in] index The index of the input to look up.
+   * @returns The metadata on success, or an error on failure. Only valid for
+   * tag::Tensor
+   */
+  Result<TensorInfo> input_tensor_meta(size_t index) const;
+
+  /**
+   * Get the number of outputs to this method.
+   *
+   * @returns The number of outputs.
+   */
+  size_t num_outputs() const;
+
+  /**
+   * Get the tag of the specified output.
+   *
+   * @param[in] index The index of the output to look up.
+   * @returns The tag of output, can only be [Tensor, Int, Bool, Double,
+   * String].
+   */
+  Result<Tag> output_tag(size_t index) const;
+
+  /**
+   * Get metadata about the specified output.
+   *
+   * @param[in] index The index of the output to look up.
+   * @returns The metadata on success, or an error on failure. Only valid for
+   * tag::Tensor
+   */
+  Result<TensorInfo> output_tensor_meta(size_t index) const;
+
+  /**
+   * Get the number of non-constant buffers this method requires.
+   *
+   * @returns The number of non-constant buffers.
+   */
+  size_t num_non_const_buffers() const;
+
+  /**
+   * Get the size in bytes of the specified non-constant buffer.
+   *
+   * @param[in] index The index of the buffer to look up.
+   * @returns The size in bytes on success, or an error on failure.
+   */
+  Result<int64_t> non_const_buffer_size(size_t index) const;
+
+ private:
+  // Let Program create MethodMeta.
+  friend class Program;
+
+  explicit MethodMeta(const executorch_flatbuffer::ExecutionPlan* s_plan);
+
+  /// Source of truth for method information
+  const executorch_flatbuffer::ExecutionPlan* s_plan_;
+};
+
+} // namespace executor
+} // namespace torch

runtime/executor/program.cpp

@@ -206,6 +206,18 @@ Result<Method> Program::load_method(
   return Error::InvalidArgument;
 }
 
+Result<MethodMeta> Program::method_meta(const char* method_name) const {
+  EXECUTORCH_SCOPE_PROF("Program::method_meta");
+  auto execution_plans = internal_program_->execution_plan();
+  for (size_t i = 0; i < execution_plans->size(); i++) {
+    auto serialization_plan = execution_plans->GetMutableObject(i);
+    if (std::strcmp(serialization_plan->name()->c_str(), method_name) == 0) {
+      return MethodMeta(serialization_plan);
+    }
+  }
+  return Error::InvalidArgument;
+}
+
 const void* Program::get_constant_buffer_data(size_t buffer_idx) const {
   ET_CHECK(is_valid());
   auto internal_program =