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

backends/native/runtime/graph/Graph.cpp

@@ -0,0 +1,102 @@
+// 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 id, const char* what) {
+  if (!in_bounds(id, vec.size())) {
+    throw std::runtime_error(std::string(what) + ": invalid id");
+  }
+  return static_cast<size_t>(id);
+}
+
+// False for kInvalid — a legitimately absent operand. Throws when `id` is set
+// but out of range: skipping it would leave def-use half-wired with no signal.
+bool resolves(ValueId id, size_t size, const char* what) {
+  if (!valid(id)) {
+    return false;
+  }
+  if (!in_bounds(id, size)) {
+    throw std::runtime_error(
+        std::string("Graph::rebuild_def_use: ") + what + " id " +
+        std::to_string(id) + " does not address the value arena");
+  }
+  return true;
+}
+
+} // namespace
+
+Node& Graph::node(NodeId id) {
+  return nodes[checked_index(nodes, id, "Graph::node")];
+}
+const Node& Graph::node(NodeId id) const {
+  return nodes[checked_index(nodes, id, "Graph::node")];
+}
+
+Value& Graph::value(ValueId id) {
+  return values[checked_index(values, id, "Graph::value")];
+}
+const Value& Graph::value(ValueId id) const {
+  return values[checked_index(values, id, "Graph::value")];
+}
+
+Graph& Graph::subgraph(GraphId id) {
+  return subgraphs[checked_index(subgraphs, id, "Graph::subgraph")];
+}
+const Graph& Graph::subgraph(GraphId id) const {
+  return subgraphs[checked_index(subgraphs, id, "Graph::subgraph")];
+}
+
+void Graph::initialize_schedule() {
+  schedule.resize(nodes.size());
+  std::iota(schedule.begin(), schedule.end(), NodeId{0});
+}
+
+void Graph::rebuild_def_use() {
+  for (Value& v : values) {
+    v.producer_id = kInvalid;
+    v.consumer_ids.clear();
+  }
+  for (size_t i = 0; i < nodes.size(); ++i) {
+    const NodeId ni = static_cast<NodeId>(i);
+    const Node& n = nodes[i];
+    for (const Output& out : n.outputs) {
+      if (out.kind == OutputValueKind::TensorList) {
+        for (ValueId e : out.elem_ids) {
+          if (resolves(e, values.size(), "output element")) {
+            values[e].producer_id = ni;
+          }
+        }
+      } else if (resolves(out.value_id, values.size(), "output")) {
+        values[out.value_id].producer_id = ni;
+      }
+    }
+    for (ValueId in : n.input_value_ids()) {
+      if (!resolves(in, 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<NodeId>& consumers = values[in].consumer_ids;
+      if (consumers.empty() || consumers.back() != ni) {
+        consumers.push_back(ni);
+      }
+    }
+  }
+}
+
+} // namespace ptn

backends/native/runtime/graph/Graph.h

@@ -0,0 +1,78 @@
+// 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 <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 Id
+// handles point into, plus the ordered graph I/O and the subgraph storage for
+// higher-order-op branch bodies. Mirrors the schema Graph; stateful
+// method-level bindings (constants / output specs / mutable buffers) live on
+// Method, not here.
+//
+// Nodes and Values are separate id spaces: a node produces none (Output), one,
+// or several values (topk, split), so the two are not 1:1, though a
+// single-output node does share its SSA name with the value it produces. The
+// dataflow DAG lives in the Values, each carrying its producer node and its
+// consumers; `schedule` is not that graph, only the linear order a runtime
+// walks the nodes in.
+//
+// GraphId indexes `subgraphs` of the *enclosing* Graph, matching the schema
+// recursion and the per-Graph SSA namespace, so a subgraph is 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_ids`.
+//
+// Node storage and node *order* are decoupled. `nodes` is append-only — a new
+// node lands at the end, out of dataflow position — so NodeIds stay stable.
+// `schedule` carries the execution order instead: reorder or insert there,
+// which moves no storage and invalidates no id. Nodes are never erased, since
+// dropping one would shift every later NodeId.
+struct Graph {
+  std::vector<Node> nodes; // node arena, incl. placeholder / output nodes
+  std::vector<NodeId> schedule; // execution / topological order over `nodes`
+  std::vector<Value> values; // SSA-value arena; ValueId indexes this
+  std::vector<ValueId> input_ids; // graph input values, in order
+  std::vector<ValueId> output_ids; // graph output values, in order
+  std::vector<Graph> subgraphs; // HOP branch bodies; GraphId indexes this
+
+  // Bounds-checked id resolution; each throws std::runtime_error on an invalid
+  // (out-of-range or kInvalid) id.
+  Node& node(NodeId id);
+  const Node& node(NodeId id) const;
+  Value& value(ValueId id);
+  const Value& value(ValueId id) const;
+  Graph& subgraph(GraphId id);
+  const Graph& subgraph(GraphId id) const;
+
+  // Set `schedule` to the identity order, which is the execution order exactly
+  // when the nodes are already in dataflow position — as they are straight off
+  // the wire, before any mutation. A graph is not executable until this runs:
+  // `schedule` starts empty, and an engine's work list is seeded from it.
+  void initialize_schedule();
+
+  // Recompute every Value's producer / consumers from the nodes, clearing the
+  // existing wiring first: each node produces its output values and consumes
+  // its input_value_ids(). Independent of `schedule` — it walks `nodes`
+  // directly. Does NOT recurse into subgraphs, which have their own SSA
+  // namespaces; call it per graph.
+  //
+  // Throws std::runtime_error on an id that is set but does not address
+  // `values`, which can only mean a corrupt graph. kInvalid is left alone: an
+  // absent operand is legal.
+  void rebuild_def_use();
+};
+
+} // namespace ptn

backends/native/runtime/graph/Value.h

@@ -45,9 +45,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.
   NodeId producer_id = 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<NodeId> consumer_ids;
   // Shares storage with this value (a view); fresh if invalid.
   ValueId alias_id = kInvalid;

backends/native/runtime/graph/targets.bzl

@@ -86,6 +86,20 @@ def define_common_targets():
         visibility = ["//executorch/backends/native/..."],
     )
 
+    runtime.cxx_library(
+        name = "graph",
+        srcs = ["Graph.cpp"],
+        exported_headers = [
+            "Graph.h",
+        ],
+        exported_deps = [
+            ":ids",
+            ":node",
+            ":value",
+        ],
+        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.
@@ -97,6 +111,7 @@ def define_common_targets():
         ],
         exported_deps = [
             ":argument",
+            ":graph",
             ":node",
             ":scalar",
             ":tensor_meta",

backends/native/runtime/graph/utils/Print.cpp

@@ -31,6 +31,19 @@ std::string id_list_str(const std::vector<ValueId>& ids) {
   return s + "]";
 }
 
+// Like id_list_str, but every entry is expected to resolve: a graph's I/O list
+// has no absent slots, so kInvalid there is corruption worth showing as "%-1".
+std::string join_ids(const std::vector<ValueId>& ids) {
+  std::string s = "[";
+  for (size_t i = 0; i < ids.size(); ++i) {
+    if (i) {
+      s += ", ";
+    }
+    s += "%" + std::to_string(ids[i]);
+  }
+  return s + "]";
+}
+
 std::string output_str(const Output& out) {
   if (out.kind == OutputValueKind::TensorList) {
     return id_list_str(out.elem_ids);
@@ -177,4 +190,22 @@ std::string to_string(const Node& node) {
   return s;
 }
 
+std::string to_string(const Graph& graph) {
+  std::string s = "inputs: " + join_ids(graph.input_ids) + "\n";
+  if (!graph.schedule.empty()) {
+    for (NodeId id : graph.schedule) {
+      s += "  " + to_string(graph.node(id)) + "\n";
+    }
+  } else {
+    for (const Node& n : graph.nodes) {
+      s += "  " + to_string(n) + "\n";
+    }
+  }
+  s += "outputs: " + join_ids(graph.output_ids) + "\n";
+  if (!graph.subgraphs.empty()) {
+    s += "(" + std::to_string(graph.subgraphs.size()) + " subgraphs)\n";
+  }
+  return s;
+}
+
 } // namespace ptn

backends/native/runtime/graph/utils/Print.h

@@ -9,6 +9,7 @@
 #include <string>
 
 #include <executorch/backends/native/runtime/graph/Argument.h>
+#include <executorch/backends/native/runtime/graph/Graph.h>
 #include <executorch/backends/native/runtime/graph/Node.h>
 #include <executorch/backends/native/runtime/graph/Scalar.h>
 #include <executorch/backends/native/runtime/graph/TensorMeta.h>
@@ -34,4 +35,8 @@ 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);
 
+// Multi-line: inputs, one line per node in `schedule` order (declaration order
+// if `schedule` is empty), outputs, subgraph count.
+std::string to_string(const Graph& graph);
+
 } // namespace ptn