From 33df6d73362bed496caca55db9d7f9e90d99e421 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Thu, 10 Aug 2023 10:06:47 -0700 Subject: [PATCH] MethodMeta (#42) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/42 There is a growing class of usecases that want to be able to inspect meta data about methods without paying the full init cost. These classes provide a safe and cheap way to view this information Reviewed By: dbort Differential Revision: D48039273 fbshipit-source-id: 78fe8c65711448eb2c3ec3edd29a6110a427ce14 --- runtime/executor/method_meta.cpp | 188 ++++++++++++++++++++ runtime/executor/method_meta.h | 190 +++++++++++++++++++++ runtime/executor/program.cpp | 12 ++ runtime/executor/program.h | 8 + runtime/executor/targets.bzl | 2 + runtime/executor/test/method_meta_test.cpp | 129 ++++++++++++++ runtime/executor/test/targets.bzl | 13 ++ 7 files changed, 542 insertions(+) create mode 100644 runtime/executor/method_meta.cpp create mode 100644 runtime/executor/method_meta.h create mode 100644 runtime/executor/test/method_meta_test.cpp diff --git a/runtime/executor/method_meta.cpp b/runtime/executor/method_meta.cpp new file mode 100644 index 00000000000..26e97dcaad4 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include +#include + +namespace torch { +namespace executor { + +namespace { +Result get_tag( + flatbuffers::Vector>:: + 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 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 sizes, + Span dim_order, + exec_aten::ScalarType scalar_type) + : sizes_(sizes), + dim_order_(dim_order), + scalar_type_(scalar_type), + nbytes_(calculate_nbytes(sizes_, scalar_type_)) {} + +Span TensorInfo::sizes() const { + return sizes_; +} + +Span 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 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 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( + tensor_value->sizes()->data(), tensor_value->sizes()->size()), + Span( + tensor_value->dim_order()->data(), tensor_value->dim_order()->size()), + static_cast(tensor_value->scalar_type())); +} + +size_t MethodMeta::num_outputs() const { + return s_plan_->outputs()->size(); +} + +Result 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 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( + tensor_value->sizes()->data(), tensor_value->sizes()->size()), + Span( + tensor_value->dim_order()->data(), tensor_value->dim_order()->size()), + static_cast(tensor_value->scalar_type())); +} + +size_t MethodMeta::num_non_const_buffers() const { + return s_plan_->non_const_buffer_sizes()->size(); +} + +Result 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 diff --git a/runtime/executor/method_meta.h b/runtime/executor/method_meta.h new file mode 100644 index 00000000000..e772cd5ef58 --- /dev/null +++ b/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 +#include +#include +#include + +// 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 sizes() const; + + /** + * Returns the dim order of the tensor. + */ + Span 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 sizes, + Span 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 sizes_; + + /** + * The dim order of the tensor. + * + * NOTE: References data from the Program, so the Program must outlive the + * TensorInfo. + */ + Span 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 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 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 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 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 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 diff --git a/runtime/executor/program.cpp b/runtime/executor/program.cpp index 5bd085698be..6fa096126bc 100644 --- a/runtime/executor/program.cpp +++ b/runtime/executor/program.cpp @@ -206,6 +206,18 @@ Result Program::load_method( return Error::InvalidArgument; } +Result 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 = diff --git a/runtime/executor/program.h b/runtime/executor/program.h index 68ae243353c..06bf396fe70 100644 --- a/runtime/executor/program.h +++ b/runtime/executor/program.h @@ -17,6 +17,7 @@ #include #include #include +#include #include // Forward declare flatbuffer types. This is a public header and must not @@ -147,6 +148,13 @@ class Program final { const char* method_name, MemoryManager* memory_manager) const; + /** + * Gathers metadata for the named method. + * + * @param[in] method_name The name of the method to get metadata for. + */ + Result method_meta(const char* method_name) const; + /** * Get the size of constant buffer * @return The size of whole constant buffer diff --git a/runtime/executor/targets.bzl b/runtime/executor/targets.bzl index a9083c22d7d..5f3711b4972 100644 --- a/runtime/executor/targets.bzl +++ b/runtime/executor/targets.bzl @@ -49,6 +49,7 @@ def define_common_targets(): name = "program" + aten_suffix, srcs = [ "method.cpp", + "method_meta.cpp", "program.cpp", "tensor_parser{}.cpp".format(aten_suffix), ], @@ -57,6 +58,7 @@ def define_common_targets(): ], exported_headers = [ "method.h", + "method_meta.h", "program.h", ], deps = [ diff --git a/runtime/executor/test/method_meta_test.cpp b/runtime/executor/test/method_meta_test.cpp new file mode 100644 index 00000000000..640acbb708a --- /dev/null +++ b/runtime/executor/test/method_meta_test.cpp @@ -0,0 +1,129 @@ +/* + * 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 +#include + +#include +#include +#include +#include +#include +#include + +using namespace ::testing; +using torch::executor::Error; +using torch::executor::MethodMeta; +using torch::executor::Program; +using torch::executor::Result; +using torch::executor::TensorInfo; +using torch::executor::util::FileDataLoader; + +class MethodMetaTest : public ::testing::Test { + protected: + void SetUp() override { + // Create a loader for the serialized ModuleAdd program. + const char* path = std::getenv("ET_MODULE_ADD_PATH"); + Result loader = FileDataLoader::From(path); + ASSERT_EQ(loader.error(), Error::Ok); + loader_ = std::make_unique(std::move(loader.get())); + + // Use it to load the program. + Result program = Program::Load( + loader_.get(), Program::Verification::InternalConsistency); + ASSERT_EQ(program.error(), Error::Ok); + program_ = std::make_unique(std::move(program.get())); + } + + private: + // Must outlive program_, but tests shouldn't need to touch it. + std::unique_ptr loader_; + + protected: + std::unique_ptr program_; +}; + +namespace { + +// Check TensorInfo against hard coded values from AddModule. +void check_tensor(const TensorInfo& tensor_info) { + auto sizes = tensor_info.sizes(); + auto dim_order = tensor_info.dim_order(); + EXPECT_EQ(sizes.size(), 2); + EXPECT_EQ(sizes[0], 2); + EXPECT_EQ(sizes[1], 2); + EXPECT_EQ(tensor_info.scalar_type(), exec_aten::ScalarType::Float); + EXPECT_EQ(dim_order.size(), 2); + EXPECT_EQ(dim_order[0], 0); + EXPECT_EQ(dim_order[1], 1); + EXPECT_EQ(tensor_info.nbytes(), 16); +} +} // namespace + +TEST_F(MethodMetaTest, MethodMetaApi) { + Result method_meta = program_->method_meta("forward"); + ASSERT_EQ(method_meta.error(), Error::Ok); + + // Appropriate amount of inputs + EXPECT_EQ(method_meta->num_inputs(), 2); + + // Appropriate amount of outputs + EXPECT_EQ(method_meta->num_outputs(), 1); + + // Appropriate amount of non_const_buffers + EXPECT_EQ(method_meta->num_non_const_buffers(), 2); + + // Appropriate content of non_const_buffers + EXPECT_EQ(method_meta->non_const_buffer_size(1).get(), 48); + + // Invalid index Errors + EXPECT_EQ( + method_meta->non_const_buffer_size(2).error(), Error::InvalidArgument); + + EXPECT_EQ( + program_->method_meta("not_a_method").error(), Error::InvalidArgument); +} + +TEST_F(MethodMetaTest, TensorInfoApi) { + Result method_meta = program_->method_meta("forward"); + ASSERT_EQ(method_meta.error(), Error::Ok); + + // Input 1 + Result in_1 = method_meta->input_tensor_meta(0); + ASSERT_TRUE(in_1.ok()); + check_tensor(in_1.get()); + + // Input 2 + Result in_2 = method_meta->input_tensor_meta(1); + ASSERT_TRUE(in_2.ok()); + check_tensor(in_2.get()); + + // Output 1 + Result out_1 = method_meta->output_tensor_meta(0); + ASSERT_TRUE(out_1.ok()); + check_tensor(out_1.get()); + + // Copyable + Result info = method_meta->input_tensor_meta(0); + TensorInfo info_copy_ctor(info.get()); + TensorInfo info_copy_assign(out_1.get()); + info_copy_assign = info.get(); + check_tensor(info_copy_ctor); + check_tensor(info_copy_assign); + + // Move-able + TensorInfo info_move_ctor(std::move(info.get())); + check_tensor(info_move_ctor); + + // Errors + EXPECT_EQ(method_meta->input_tensor_meta(3).error(), Error::InvalidArgument); + EXPECT_EQ(method_meta->input_tensor_meta(-1).error(), Error::InvalidArgument); + EXPECT_EQ(method_meta->output_tensor_meta(3).error(), Error::InvalidArgument); + EXPECT_EQ( + method_meta->output_tensor_meta(-1).error(), Error::InvalidArgument); +} diff --git a/runtime/executor/test/targets.bzl b/runtime/executor/test/targets.bzl index d49952410c5..ae1640ea71e 100644 --- a/runtime/executor/test/targets.bzl +++ b/runtime/executor/test/targets.bzl @@ -127,6 +127,19 @@ def define_common_targets(is_fbcode = False): env = modules_env, ) + runtime.cxx_test( + name = "method_meta_test", + srcs = [ + "method_meta_test.cpp", + ], + deps = [ + "//executorch/runtime/executor:program", + "//executorch/util:util", + "//executorch/extension/data_loader:file_data_loader", + ], + env = modules_env, + ) + runtime.cxx_test( name = "program_test", srcs = [