[executorch][native] Add Graph index arena to the in-memory IR

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 <executorch/backends/native/runtime/graph/Graph.h>
+
+#include <cstddef>
+#include <numeric>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace ptn {
+
+namespace {
+
+template <typename Vec>
+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<size_t>(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<ValueRef>& 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<NodeRef>(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<NodeRef>& 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

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 <string>
+#include <vector>
+
+#include <executorch/backends/native/runtime/graph/Ids.h>
+#include <executorch/backends/native/runtime/graph/Node.h>
+#include <executorch/backends/native/runtime/graph/Value.h>
+
+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<Node> nodes; // node arena, incl. placeholder / output nodes
+  std::vector<NodeRef> schedule; // execution / topological order over `nodes`
+  std::vector<Value> values; // SSA-value arena; ValueRef indexes this
+  std::vector<ValueRef> input_refs; // graph input values, in order
+  std::vector<ValueRef> output_refs; // graph output values, in order
+  std::vector<Graph> 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

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<NodeRef> consumer_refs;
   // Shares storage with this value (a view); fresh if invalid.
   ValueRef alias_ref = kInvalid;

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/..."],
+    )