diff --git a/backends/native/runtime/graph/Argument.cpp b/backends/native/runtime/graph/Argument.cpp new file mode 100644 index 00000000000..1eca016a758 --- /dev/null +++ b/backends/native/runtime/graph/Argument.cpp @@ -0,0 +1,18 @@ +// 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 + +namespace ptn { + +void Argument::throw_bad_kind(const char* what) { + throw std::runtime_error(std::string("Argument::") + what); +} + +} // namespace ptn diff --git a/backends/native/runtime/graph/Argument.h b/backends/native/runtime/graph/Argument.h new file mode 100644 index 00000000000..cdff7a22cab --- /dev/null +++ b/backends/native/runtime/graph/Argument.h @@ -0,0 +1,251 @@ +// 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 +#include +#include +#include + +#include +#include + +namespace ptn { + +// One payload struct per kind of value an fx arg/kwarg can hold; the same set +// of kinds as the schema ArgumentValue union, but in-graph references are +// resolved to ValueIds (the deserializer turns SSA names into arena indices). + +struct NoneArg {}; + +struct TensorArg { + ValueId id = kInvalid; +}; + +// A scalar int operand: a literal `value`, or (when `id` is valid) a reference +// to an in-graph int value — a sym_size / arith node — with `value` ignored. +struct IntArg { + int64_t value = 0; + ValueId id = kInvalid; +}; + +struct FloatArg { + double value = 0.0; + ValueId id = kInvalid; +}; + +struct BoolArg { + bool value = false; + ValueId id = kInvalid; +}; + +struct StringArg { + std::string value; +}; + +struct ScalarTypeArg { + ScalarType value = ScalarType::Float; +}; + +// A list of ints. Element i is a literal (values[i]), or a symbolic reference +// when `ids` is non-empty and ids[i] is valid (values[i] ignored). Empty `ids` +// means all literal; otherwise ids.size() == values.size(). E.g. a dynamic view +// size [s0, -1] is values={0, -1}, ids={, kInvalid}. +struct IntListArg { + std::vector values; + std::vector ids; +}; + +struct FloatListArg { + std::vector values; +}; + +struct BoolListArg { + std::vector values; +}; + +struct TensorListArg { + std::vector ids; +}; + +// A list of optional tensor references (Tensor?[]); a kInvalid entry is None. +struct OptionalTensorListArg { + std::vector ids; +}; + +// A subgraph passed to a higher-order op (torch.cond / while_loop / map). +// `name` is the original submodule attr label, for debug output only. +struct GraphArg { + std::string name; + GraphId subgraph_id = kInvalid; +}; + +// Same order as the schema ArgumentValue union, with every id one lower: +// flatbuffers reserves 0 for an absent union, so its tags start at 1. Nothing +// casts between the two; the deserializer switches on the schema tag by name. +enum class ArgKind : int8_t { + Tensor = 0, + None = 1, + Int = 2, + Float = 3, + Bool = 4, + String = 5, + ScalarType = 6, + IntList = 7, + FloatList = 8, + BoolList = 9, + TensorList = 10, + OptionalTensorList = 11, + Graph = 12, +}; + +// A single fx argument value: a std::variant over the payload structs, with the +// alternatives listed in ArgKind order so kind() is the variant's index. The +// static_asserts below hold the two in step. +// +// Construct from a payload directly: `Argument a = IntArg{5};`. +class Argument { + private: + using Storage = std::variant< + TensorArg, + NoneArg, + IntArg, + FloatArg, + BoolArg, + StringArg, + ScalarTypeArg, + IntListArg, + FloatListArg, + BoolListArg, + TensorListArg, + OptionalTensorListArg, + GraphArg>; + + // kind() casts the variant index straight to ArgKind, so inserting an + // alternative or an enumerator without the other would silently mis-tag every + // switch on kind(). These make that a compile error instead. + template + static constexpr bool alt_is = std:: + is_same_v(K), Storage>, T>; + + static_assert(alt_is); + static_assert(alt_is); + static_assert(alt_is); + static_assert(alt_is); + static_assert(alt_is); + static_assert(alt_is); + static_assert(alt_is); + static_assert(alt_is); + static_assert(alt_is); + static_assert(alt_is); + static_assert(alt_is); + static_assert(alt_is); + static_assert(alt_is); + static_assert( + std::variant_size_v == static_cast(ArgKind::Graph) + 1); + + Storage value_; + + // The live payload, or a std::runtime_error reading "Argument::". + template + const T& payload(const char* what) const { + const T* p = std::get_if(&value_); + if (p == nullptr) { + throw_bad_kind(what); + } + return *p; + } + + [[noreturn]] static void throw_bad_kind(const char* what); + + public: + // Pinned to NoneArg rather than defaulted, so the default stays None however + // the alternatives are ordered. + Argument() : value_(NoneArg{}) {} + + // Implicit by design: a payload struct is the natural spelling of an + // Argument. Spelled out per alternative rather than as one constrained + // template, because several payloads are aggregates that would make a + // template's overload resolution depend on which of them happens to accept + // the argument. The move is omitted on the trivially copyable payloads, + // where it buys nothing. + // cppcheck-suppress-begin noExplicitConstructor + /* implicit */ Argument(TensorArg a) : value_(a) {} + /* implicit */ Argument(NoneArg a) : value_(a) {} + /* implicit */ Argument(IntArg a) : value_(a) {} + /* implicit */ Argument(FloatArg a) : value_(a) {} + /* implicit */ Argument(BoolArg a) : value_(a) {} + /* implicit */ Argument(StringArg a) : value_(std::move(a)) {} + /* implicit */ Argument(ScalarTypeArg a) : value_(a) {} + /* implicit */ Argument(IntListArg a) : value_(std::move(a)) {} + /* implicit */ Argument(FloatListArg a) : value_(std::move(a)) {} + /* implicit */ Argument(BoolListArg a) : value_(std::move(a)) {} + /* implicit */ Argument(TensorListArg a) : value_(std::move(a)) {} + /* implicit */ Argument(OptionalTensorListArg a) : value_(std::move(a)) {} + /* implicit */ Argument(GraphArg a) : value_(std::move(a)) {} + // cppcheck-suppress-end noExplicitConstructor + + ArgKind kind() const { + return static_cast(value_.index()); + } + + // Typed payload accessors: throw std::runtime_error unless the kind matches. + const TensorArg& as_tensor() const { + return payload("as_tensor: argument is not a Tensor"); + } + const IntArg& as_int() const { + return payload("as_int: argument is not an Int"); + } + const FloatArg& as_float() const { + return payload("as_float: argument is not a Float"); + } + const BoolArg& as_bool() const { + return payload("as_bool: argument is not a Bool"); + } + const StringArg& as_string() const { + return payload("as_string: argument is not a String"); + } + const ScalarTypeArg& as_scalar_type() const { + return payload( + "as_scalar_type: argument is not a ScalarType"); + } + const IntListArg& as_int_list() const { + return payload("as_int_list: argument is not an IntList"); + } + const FloatListArg& as_float_list() const { + return payload("as_float_list: argument is not a FloatList"); + } + const BoolListArg& as_bool_list() const { + return payload("as_bool_list: argument is not a BoolList"); + } + const TensorListArg& as_tensor_list() const { + return payload( + "as_tensor_list: argument is not a TensorList"); + } + const OptionalTensorListArg& as_optional_tensor_list() const { + return payload( + "as_optional_tensor_list: argument is not an OptionalTensorList"); + } + const GraphArg& as_graph() const { + return payload("as_graph: argument is not a Graph"); + } +}; + +// A positional or keyword argument. `name` is the operator-schema parameter +// name (NOT a value reference; empty for positional-only). `mutated` is true +// when the op writes this input in place (op-schema Tensor(a!), e.g. counter / +// kv_cache). +struct NamedArgument { + std::string name; + Argument arg; + bool mutated = false; +}; + +} // namespace ptn diff --git a/backends/native/runtime/graph/Ids.h b/backends/native/runtime/graph/Ids.h index baa33415006..5206c2f7d06 100644 --- a/backends/native/runtime/graph/Ids.h +++ b/backends/native/runtime/graph/Ids.h @@ -12,12 +12,13 @@ namespace ptn { -// Index-arena handles: a NodeId indexes the graph's node arena, a ValueId -// its value arena. Plain int32_t aliases — they index, compare, and hash -// directly, at the cost of no NodeId/ValueId type distinction. kInvalid marks -// "no id". +// Index-arena handles: a NodeId indexes the graph's node arena, a ValueId its +// value arena, a GraphId a subgraph arena (HOP branch bodies). Plain int32_t +// aliases — they index, compare, and hash directly, at the cost of no type +// distinction between them. kInvalid marks "no id". using NodeId = int32_t; using ValueId = int32_t; +using GraphId = int32_t; inline constexpr int32_t kInvalid = -1; constexpr bool valid(int32_t id) { diff --git a/backends/native/runtime/graph/targets.bzl b/backends/native/runtime/graph/targets.bzl index c18b4f522b1..9d5f7d66baa 100644 --- a/backends/native/runtime/graph/targets.bzl +++ b/backends/native/runtime/graph/targets.bzl @@ -60,6 +60,19 @@ def define_common_targets(): visibility = ["//executorch/backends/native/..."], ) + runtime.cxx_library( + name = "argument", + srcs = ["Argument.cpp"], + exported_headers = [ + "Argument.h", + ], + exported_deps = [ + ":ids", + ":scalar_type", + ], + visibility = ["//executorch/backends/native/..."], + ) + # utils/ has no BUCK of its own, so the IR printer's target lives here. Kept # separate from the IR libraries so only a consumer that dumps the IR links # the formatting code. @@ -70,6 +83,7 @@ def define_common_targets(): "utils/Print.h", ], exported_deps = [ + ":argument", ":scalar", ":tensor_meta", ], diff --git a/backends/native/runtime/graph/utils/Print.cpp b/backends/native/runtime/graph/utils/Print.cpp index 0e3078d913f..9e0f59a25d9 100644 --- a/backends/native/runtime/graph/utils/Print.cpp +++ b/backends/native/runtime/graph/utils/Print.cpp @@ -8,11 +8,31 @@ #include #include +#include #include namespace ptn { +namespace { + +std::string id_str(ValueId id) { + return valid(id) ? "%" + std::to_string(id) : "None"; +} + +std::string id_list_str(const std::vector& ids) { + std::string s = "["; + for (size_t i = 0; i < ids.size(); ++i) { + if (i) { + s += ", "; + } + s += id_str(ids[i]); + } + return s + "]"; +} + +} // namespace + std::string to_string(const TensorMeta& meta) { std::string s = scalar_type_name(meta.dtype); s += "["; @@ -43,4 +63,70 @@ std::string to_string(const Scalar& scalar) { return format_double(scalar.to_double()); } +std::string to_string(const Argument& arg) { + switch (arg.kind()) { + case ArgKind::None: + return "None"; + case ArgKind::Tensor: + return id_str(arg.as_tensor().id); + case ArgKind::Int: { + const IntArg& a = arg.as_int(); + return valid(a.id) ? id_str(a.id) : std::to_string(a.value); + } + case ArgKind::Float: { + const FloatArg& a = arg.as_float(); + return valid(a.id) ? id_str(a.id) : format_double(a.value); + } + case ArgKind::Bool: { + const BoolArg& a = arg.as_bool(); + return valid(a.id) ? id_str(a.id) : (a.value ? "true" : "false"); + } + case ArgKind::String: + return "\"" + arg.as_string().value + "\""; + case ArgKind::ScalarType: + return scalar_type_name(arg.as_scalar_type().value); + case ArgKind::IntList: { + const IntListArg& a = arg.as_int_list(); + std::string s = "["; + for (size_t i = 0; i < a.values.size(); ++i) { + if (i) { + s += ", "; + } + const bool sym = i < a.ids.size() && valid(a.ids[i]); + s += sym ? id_str(a.ids[i]) : std::to_string(a.values[i]); + } + return s + "]"; + } + case ArgKind::FloatList: { + const FloatListArg& a = arg.as_float_list(); + std::string s = "["; + for (size_t i = 0; i < a.values.size(); ++i) { + if (i) { + s += ", "; + } + s += format_double(a.values[i]); + } + return s + "]"; + } + case ArgKind::BoolList: { + const BoolListArg& a = arg.as_bool_list(); + std::string s = "["; + for (size_t i = 0; i < a.values.size(); ++i) { + if (i) { + s += ", "; + } + s += a.values[i] ? "true" : "false"; + } + return s + "]"; + } + case ArgKind::TensorList: + return id_list_str(arg.as_tensor_list().ids); + case ArgKind::OptionalTensorList: + return id_list_str(arg.as_optional_tensor_list().ids); + case ArgKind::Graph: + return "graph(" + arg.as_graph().name + ")"; + } + return "?"; +} + } // namespace ptn diff --git a/backends/native/runtime/graph/utils/Print.h b/backends/native/runtime/graph/utils/Print.h index 4fb3fd566c9..6f17528947e 100644 --- a/backends/native/runtime/graph/utils/Print.h +++ b/backends/native/runtime/graph/utils/Print.h @@ -8,6 +8,7 @@ #include +#include #include #include @@ -25,4 +26,8 @@ std::string to_string(const TensorMeta& meta); // The live alternative only: "true", "-3", "1.5e-08". std::string to_string(const Scalar& scalar); +// Compact one-line form. A symbolic operand (one carrying a valid id) renders +// as "%"; a literal renders as its value. +std::string to_string(const Argument& arg); + } // namespace ptn