diff --git a/backends/native/runtime/graph/Graph.cpp b/backends/native/runtime/graph/Graph.cpp new file mode 100644 index 00000000000..167a106db57 --- /dev/null +++ b/backends/native/runtime/graph/Graph.cpp @@ -0,0 +1,134 @@ +// 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 + +namespace ptn { + +namespace { + +template +size_t checked_index(const Vec& vec, int32_t ref, const char* what) { + if (!in_bounds(ref, vec.size())) { + throw std::runtime_error(std::string(what) + ": invalid ref"); + } + return static_cast(ref); +} + +// True when `ref` names a value in the arena. False for kInvalid, which is a +// legitimately absent operand. Throws when `ref` is set but out of range: that +// can only be a corrupt graph, and skipping it would leave def-use half-wired +// with no signal to the caller. +bool resolves(ValueRef ref, size_t size, const char* what) { + if (!valid(ref)) { + return false; + } + if (!in_bounds(ref, size)) { + throw std::runtime_error( + std::string("Graph::rebuild_def_use: ") + what + " ref " + + std::to_string(ref) + " does not address the value arena"); + } + return true; +} + +// Format a ref list as "[%0, %1]". +std::string join_refs(const std::vector& refs) { + std::string s = "["; + for (size_t i = 0; i < refs.size(); ++i) { + if (i) { + s += ", "; + } + s += "%" + std::to_string(refs[i]); + } + return s + "]"; +} + +} // namespace + +Node& Graph::node(NodeRef ref) { + return nodes[checked_index(nodes, ref, "Graph::node")]; +} +const Node& Graph::node(NodeRef ref) const { + return nodes[checked_index(nodes, ref, "Graph::node")]; +} + +Value& Graph::value(ValueRef ref) { + return values[checked_index(values, ref, "Graph::value")]; +} +const Value& Graph::value(ValueRef ref) const { + return values[checked_index(values, ref, "Graph::value")]; +} + +Graph& Graph::subgraph(GraphRef ref) { + return subgraphs[checked_index(subgraphs, ref, "Graph::subgraph")]; +} +const Graph& Graph::subgraph(GraphRef ref) const { + return subgraphs[checked_index(subgraphs, ref, "Graph::subgraph")]; +} + +void Graph::reset_schedule() { + schedule.resize(nodes.size()); + std::iota(schedule.begin(), schedule.end(), NodeRef{0}); +} + +void Graph::rebuild_def_use() { + for (Value& v : values) { + v.producer_ref = kInvalid; + v.consumer_refs.clear(); + } + for (size_t i = 0; i < nodes.size(); ++i) { + const NodeRef ni = static_cast(i); + const Node& n = nodes[i]; + for (const Output& out : n.outputs) { + if (out.kind == OutputValueKind::TensorList) { + for (ValueRef r : out.elem_refs) { + if (resolves(r, values.size(), "output element")) { + values[r].producer_ref = ni; + } + } + } else if (resolves(out.value_ref, values.size(), "output")) { + values[out.value_ref].producer_ref = ni; + } + } + for (ValueRef r : n.input_value_refs()) { + if (!resolves(r, values.size(), "input")) { + continue; + } + // Nodes are walked in arena order, so a repeated operand appends `ni` + // consecutively; checking the tail is enough to keep this a set. + std::vector& consumers = values[r].consumer_refs; + if (consumers.empty() || consumers.back() != ni) { + consumers.push_back(ni); + } + } + } +} + +std::string Graph::to_string() const { + std::string s = "inputs: " + join_refs(input_refs) + "\n"; + if (!schedule.empty()) { + for (NodeRef ref : schedule) { + s += " " + node(ref).to_string() + "\n"; + } + } else { + for (const Node& n : nodes) { + s += " " + n.to_string() + "\n"; + } + } + s += "outputs: " + join_refs(output_refs) + "\n"; + if (!subgraphs.empty()) { + s += "(" + std::to_string(subgraphs.size()) + " subgraphs)\n"; + } + return s; +} + +} // namespace ptn diff --git a/backends/native/runtime/graph/Graph.h b/backends/native/runtime/graph/Graph.h new file mode 100644 index 00000000000..a108d391f97 --- /dev/null +++ b/backends/native/runtime/graph/Graph.h @@ -0,0 +1,82 @@ +// 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 + +namespace ptn { + +// A pure function body: the index arena that owns the Nodes and Values the Ref +// handles point into, plus the ordered graph I/O and the subgraph arena for +// higher-order-op branch bodies. Mirrors the schema Graph; stateful +// method-level bindings (constants / output specs / mutable buffers) live on +// Method (deferred), not here. +// +// GraphRef indexes `subgraphs` of the *enclosing* Graph (per-graph arena), +// matching the schema recursion and the per-Graph SSA namespace. A subgraph is +// thus self-contained with its parent. +// +// Placeholder and Output nodes are real entries in `nodes` (schema OpKind), so +// def-use is uniform: a graph input value's producer is its placeholder node, +// not kInvalid; graph inputs are identified by membership in `input_refs`. +// +// Node storage and node *order* are decoupled. `nodes` is an append-only arena +// (a new node lands at the end, out of dataflow position) so NodeRefs stay +// stable — the index-arena invariant. `schedule` gives the topological / +// execution order over it: the order a runtime walks the nodes. At load the +// arena order equals the wire's topological order and `schedule` is the +// identity [0, 1, ..., n-1]; across mutation the arena order is no longer +// topological, so `schedule` is authoritative — reorder / insert there (moving +// NodeRefs, invalidating no ref) rather than moving storage. Deletion is still +// deferred; it will need a tombstone plus a compacting pass, since dropping a +// node outright would shift every later NodeRef. +struct Graph { + std::vector nodes; // node arena, incl. placeholder / output nodes + std::vector schedule; // execution / topological order over `nodes` + std::vector values; // SSA-value arena; ValueRef indexes this + std::vector input_refs; // graph input values, in order + std::vector output_refs; // graph output values, in order + std::vector subgraphs; // HOP branch/body bodies; GraphRef indexes this + + // Bounds-checked ref resolution; each throws std::runtime_error on an invalid + // (out-of-range or kInvalid) ref. + Node& node(NodeRef ref); + const Node& node(NodeRef ref) const; + Value& value(ValueRef ref); + const Value& value(ValueRef ref) const; + Graph& subgraph(GraphRef ref); + const Graph& subgraph(GraphRef ref) const; + + // Reset `schedule` to the identity order [0, 1, ..., nodes.size() - 1] (arena + // order). The deserializer calls this after appending the nodes in wire + // order. + void reset_schedule(); + + // Recompute every Value's producer / consumers from the nodes (clears the + // existing wiring first). Each node is the producer of its output values and + // a consumer of its input_value_refs(). Order-independent (walks the arena), + // so it does not depend on `schedule`. Does NOT recurse into subgraphs — each + // has its own SSA namespace, so call it per graph. + // + // consumer_refs is a set of consuming nodes, not a bag of uses: a node that + // reads a value twice (`add(x, x)`) is listed once. Throws + // std::runtime_error on a ref that is set but does not address the value + // arena, which can only mean a corrupt graph; kInvalid is left alone, since + // an absent operand is legal. + void rebuild_def_use(); + + // Multi-line debug dump: inputs, one line per node in `schedule` order (arena + // order if `schedule` is empty), outputs, subgraph count. + std::string to_string() const; +}; + +} // namespace ptn diff --git a/backends/native/runtime/graph/Value.h b/backends/native/runtime/graph/Value.h index d48ac44ec82..746ae583ae0 100644 --- a/backends/native/runtime/graph/Value.h +++ b/backends/native/runtime/graph/Value.h @@ -47,9 +47,11 @@ class Value { public: // SSA name, scoped to the enclosing Graph. std::string name; - // Defining node; invalid => graph input. + // Defining node (a placeholder node for a graph input); invalid => unwired. NodeRef producer_ref = kInvalid; - // Def-use, built by inverting node inputs. + // Def-use, built by inverting node inputs. The consuming nodes, each listed + // once: a node that reads this value twice (`add(x, x)`) appears once, so + // size() counts consumers rather than uses. std::vector consumer_refs; // Shares storage with this value (a view); fresh if invalid. ValueRef alias_ref = kInvalid; diff --git a/backends/native/runtime/graph/targets.bzl b/backends/native/runtime/graph/targets.bzl index 0b54abeb91f..be04d450874 100644 --- a/backends/native/runtime/graph/targets.bzl +++ b/backends/native/runtime/graph/targets.bzl @@ -98,3 +98,19 @@ def define_common_targets(): deps = [":format"], visibility = ["//executorch/backends/native/..."], ) + + # The index arena: a pure function body owning the Nodes and Values that Refs + # index into, plus ordered graph I/O and the per-graph subgraph arena. + runtime.cxx_library( + name = "graph", + srcs = ["Graph.cpp"], + exported_headers = [ + "Graph.h", + ], + exported_deps = [ + ":ids", + ":node", + ":value", + ], + visibility = ["//executorch/backends/native/..."], + )