diff --git a/backends/native/runtime/graph/Node.cpp b/backends/native/runtime/graph/Node.cpp new file mode 100644 index 00000000000..a5ce90e96a1 --- /dev/null +++ b/backends/native/runtime/graph/Node.cpp @@ -0,0 +1,69 @@ +// 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 + +namespace ptn { + +namespace { + +void push_id(std::vector& ids, ValueId id) { + if (valid(id)) { + ids.push_back(id); + } +} + +} // namespace + +std::vector Node::input_value_ids() const { + std::vector ids; + for (const NamedArgument& named : inputs) { + const Argument& arg = named.arg; + switch (arg.kind()) { + case ArgKind::Tensor: + push_id(ids, arg.as_tensor().id); + break; + case ArgKind::Int: + push_id(ids, arg.as_int().id); + break; + case ArgKind::Float: + push_id(ids, arg.as_float().id); + break; + case ArgKind::Bool: + push_id(ids, arg.as_bool().id); + break; + case ArgKind::IntList: + for (ValueId i : arg.as_int_list().ids) { + push_id(ids, i); + } + break; + case ArgKind::TensorList: + for (ValueId i : arg.as_tensor_list().ids) { + push_id(ids, i); + } + break; + case ArgKind::OptionalTensorList: + for (ValueId i : arg.as_optional_tensor_list().ids) { + push_id(ids, i); + } + break; + // Carry no ids. Listed rather than defaulted so a new ArgKind that does + // carry one is a compiler warning here, not a silently unwired operand. + case ArgKind::None: + case ArgKind::String: + case ArgKind::ScalarType: + case ArgKind::FloatList: + case ArgKind::BoolList: + case ArgKind::Graph: + break; + } + } + return ids; +} + +} // namespace ptn diff --git a/backends/native/runtime/graph/Node.h b/backends/native/runtime/graph/Node.h new file mode 100644 index 00000000000..63e7be5c0e8 --- /dev/null +++ b/backends/native/runtime/graph/Node.h @@ -0,0 +1,79 @@ +// 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 { + +// fx node kind (node.op). Pinned to the schema OpKind ids. +enum class OpKind : int8_t { + CallFunction = 0, + Placeholder = 1, + Output = 2, +}; + +// What a single node Output produces. Pinned to the schema OutputValueKind ids; +// named to match the schema (distinct from the graph-level OutputKind — user +// output vs buffer mutation). +enum class OutputValueKind : int8_t { + Tensor = 0, + TensorList = 1, + Int = 2, + Bool = 3, + Float = 4, +}; + +// One value produced by a node. Tensor / Int / Bool / Float use `value_id`; +// TensorList (e.g. split) uses `elem_ids`. The return-ABI grouping is +// preserved so engine translation can tell a single result from a tuple / list +// (topk emits two Tensor outputs; split emits one TensorList output). The +// storage-alias fact lives on the produced Value, not here. +struct Output { + OutputValueKind kind = OutputValueKind::Tensor; + ValueId value_id = kInvalid; + std::vector elem_ids; +}; + +// One fx node: an op invocation (CallFunction) or a graph-boundary marker +// (Placeholder / Output). For an Output node, `inputs` is the ordered return +// list (tensors and literals alike) and `outputs` is empty; for a Placeholder, +// `target` is empty and it produces a single output. `attrs` is a transient +// scratch map (the fx node.meta analog); it is not serialized. +struct Node { + std::string name; + OpKind op_kind = OpKind::CallFunction; + // fqn, e.g. "torch.ops.aten.addmm.default"; empty for placeholder / output. + std::string target; + std::vector inputs; + std::vector outputs; + std::unordered_map attrs; + + bool is_call() const { + return op_kind == OpKind::CallFunction; + } + bool is_placeholder() const { + return op_kind == OpKind::Placeholder; + } + bool is_output() const { + return op_kind == OpKind::Output; + } + + // Every ValueId this node consumes: tensor args, tensor-list / optional-list + // elements, and symbolic scalar ids (kInvalid entries skipped). Used to + // (re)build def-use wiring. + std::vector input_value_ids() const; +}; + +} // namespace ptn diff --git a/backends/native/runtime/graph/targets.bzl b/backends/native/runtime/graph/targets.bzl index 9d5f7d66baa..f4e6a1e7240 100644 --- a/backends/native/runtime/graph/targets.bzl +++ b/backends/native/runtime/graph/targets.bzl @@ -73,6 +73,19 @@ def define_common_targets(): visibility = ["//executorch/backends/native/..."], ) + runtime.cxx_library( + name = "node", + srcs = ["Node.cpp"], + exported_headers = [ + "Node.h", + ], + exported_deps = [ + ":argument", + ":ids", + ], + 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. @@ -84,6 +97,7 @@ def define_common_targets(): ], exported_deps = [ ":argument", + ":node", ":scalar", ":tensor_meta", ], diff --git a/backends/native/runtime/graph/utils/Print.cpp b/backends/native/runtime/graph/utils/Print.cpp index fde6e745634..2a93bb7432a 100644 --- a/backends/native/runtime/graph/utils/Print.cpp +++ b/backends/native/runtime/graph/utils/Print.cpp @@ -31,6 +31,13 @@ std::string id_list_str(const std::vector& ids) { return s + "]"; } +std::string output_str(const Output& out) { + if (out.kind == OutputValueKind::TensorList) { + return id_list_str(out.elem_ids); + } + return id_str(out.value_id); +} + } // namespace std::string to_string(const TensorMeta& meta) { @@ -122,4 +129,45 @@ std::string to_string(const Argument& arg) { return "?"; } +std::string to_string(const Node& node) { + std::string s = node.name.empty() ? "_" : node.name; + s += " = "; + switch (node.op_kind) { + case OpKind::CallFunction: + s += node.target; + break; + case OpKind::Placeholder: + s += ""; + break; + case OpKind::Output: + s += ""; + break; + } + s += "("; + for (size_t i = 0; i < node.inputs.size(); ++i) { + if (i) { + s += ", "; + } + const NamedArgument& named = node.inputs[i]; + if (!named.name.empty()) { + s += named.name + "="; + } + s += to_string(named.arg); + if (named.mutated) { + s += "!"; + } + } + s += ")"; + if (!node.outputs.empty()) { + s += " -> "; + for (size_t i = 0; i < node.outputs.size(); ++i) { + if (i) { + s += ", "; + } + s += output_str(node.outputs[i]); + } + } + return s; +} + } // namespace ptn diff --git a/backends/native/runtime/graph/utils/Print.h b/backends/native/runtime/graph/utils/Print.h index 6f17528947e..1272cac7fa4 100644 --- a/backends/native/runtime/graph/utils/Print.h +++ b/backends/native/runtime/graph/utils/Print.h @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -30,4 +31,7 @@ std::string to_string(const Scalar& scalar); // as "%"; a literal renders as its value. std::string to_string(const Argument& arg); +// Single line, e.g. "a = aten.add.Tensor(x, y, alpha=1) -> %3". +std::string to_string(const Node& node); + } // namespace ptn