From 0aa5965a729abd1a5243ff5bea490db5bb89af30 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 27 Aug 2026 12:44:17 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/native/runtime/graph/Argument.cpp | 18 ++ backends/native/runtime/graph/Argument.h | 251 +++++++++++++++++++++ backends/native/runtime/graph/Ids.h | 9 +- backends/native/runtime/graph/targets.bzl | 13 ++ 4 files changed, 287 insertions(+), 4 deletions(-) create mode 100644 backends/native/runtime/graph/Argument.cpp create mode 100644 backends/native/runtime/graph/Argument.h 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..8f544f6c495 --- /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 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; +}; + +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, for debug output only. +struct GraphArg { + std::string name; + GraphRef subgraph_ref = 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 e14ca279514..129497ffff4 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 NodeRef indexes the graph's node arena, a ValueRef -// its value arena. Plain int32_t aliases — they index, compare, and hash -// directly, at the cost of no NodeRef/ValueRef type distinction. kInvalid marks -// "no ref". +// Index-arena handles: a NodeRef indexes the graph's node arena, a ValueRef its +// value arena, a GraphRef 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 ref". using NodeRef = int32_t; using ValueRef = int32_t; +using GraphRef = int32_t; inline constexpr int32_t kInvalid = -1; constexpr bool valid(int32_t ref) { diff --git a/backends/native/runtime/graph/targets.bzl b/backends/native/runtime/graph/targets.bzl index 5da780d9d2f..37e085d90db 100644 --- a/backends/native/runtime/graph/targets.bzl +++ b/backends/native/runtime/graph/targets.bzl @@ -60,3 +60,16 @@ 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/..."], + )