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..28a74942c2c --- /dev/null +++ b/backends/native/runtime/graph/Argument.h @@ -0,0 +1,229 @@ +// 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 + +namespace ptn { + +// One payload struct per kind of value an fx arg/kwarg can hold; mirrors the +// schema ArgumentValue union, but in-graph references are resolved to ValueRefs +// (the deserializer turns SSA names into arena indices). + +struct NoneArg {}; + +struct TensorArg { + ValueRef ref = kInvalid; +}; + +// A scalar int operand: a literal `value`, or (when `ref` is valid) a reference +// to an in-graph int value — a sym_size / arith node — with `value` ignored. +struct IntArg { + int64_t value = 0; + ValueRef ref = kInvalid; +}; + +struct FloatArg { + double value = 0.0; + ValueRef ref = kInvalid; +}; + +struct BoolArg { + bool value = false; + ValueRef ref = 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 `refs` is non-empty and refs[i] is valid (values[i] ignored). Empty +// `refs` means all literal; otherwise refs.size() == values.size(). E.g. a +// dynamic view size [s0, -1] is values={0, -1}, refs={, kInvalid}. +struct IntListArg { + std::vector values; + std::vector refs; +}; + +struct FloatListArg { + std::vector values; +}; + +struct BoolListArg { + std::vector values; +}; + +// A list of tensor references (e.g. cat's input list). +struct TensorListArg { + std::vector refs; +}; + +// A list of optional tensor references (Tensor?[]); a kInvalid entry is None. +struct OptionalTensorListArg { + std::vector refs; +}; + +// A subgraph passed to a higher-order op (torch.cond / while_loop / map). +// `name` is the original submodule attr label (debug only); `subgraph_ref` +// indexes the subgraph arena (populated once Graph/Model land). +struct GraphArg { + std::string name; + GraphRef subgraph_ref = kInvalid; +}; + +enum class ArgKind : int8_t { + None = 0, + Tensor = 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 and the +// two can never drift apart. +// +// Construct from a payload directly (`Argument a = IntArg{5};`); read the +// payload via the as_*() accessors after checking kind() (each throws +// std::runtime_error on a kind mismatch). +class Argument { + private: + using Storage = std::variant< + NoneArg, + TensorArg, + IntArg, + FloatArg, + BoolArg, + StringArg, + ScalarTypeArg, + IntListArg, + FloatListArg, + BoolListArg, + TensorListArg, + OptionalTensorListArg, + GraphArg>; + + 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: + Argument() = default; // None + + // 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 reads every one of these as an unintended converting constructor; + // the suppression is scoped to this block so a genuinely accidental implicit + // constructor added elsewhere in the class is still reported. + // cppcheck-suppress-begin noExplicitConstructor + /* implicit */ Argument(TensorArg 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. Each throws std::runtime_error unless kind() + // matches; guard with kind(). + 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 65295791111..90a0f98ce8a 100644 --- a/backends/native/runtime/graph/Ids.h +++ b/backends/native/runtime/graph/Ids.h @@ -13,12 +13,14 @@ namespace ptn { // Index-arena handles. A NodeRef indexes Graph.nodes; a ValueRef indexes -// Graph.values. Plain int32_t aliases — they index, compare, and hash -// directly, at the cost of no NodeRef/ValueRef type distinction. kInvalid -// marks "no ref" (e.g. a graph input has no producer node; a fresh value has -// no alias). +// Graph.values; a GraphRef indexes a subgraph arena (HOP branch bodies; the +// arena itself lands with Graph/Model). Plain int32_t aliases — they index, +// compare, and hash directly, at the cost of no type distinction between them. +// kInvalid marks "no ref" (e.g. a graph input has no producer node; a fresh +// value has no alias). using NodeRef = int32_t; using ValueRef = int32_t; +using GraphRef = int32_t; constexpr int32_t kInvalid = -1; inline bool valid(int32_t ref) { diff --git a/backends/native/runtime/graph/targets.bzl b/backends/native/runtime/graph/targets.bzl index 606c9dac305..15196df7d51 100644 --- a/backends/native/runtime/graph/targets.bzl +++ b/backends/native/runtime/graph/targets.bzl @@ -67,3 +67,18 @@ def define_common_targets(): ], visibility = ["//executorch/backends/native/..."], ) + + # One fx argument value: a std::variant over the schema ArgumentValue + # payload kinds (in-graph refs resolved to ValueRefs), plus NamedArgument. + runtime.cxx_library( + name = "argument", + srcs = ["Argument.cpp"], + exported_headers = [ + "Argument.h", + ], + exported_deps = [ + ":ids", + ":scalar_type", + ], + visibility = ["//executorch/backends/native/..."], + )