From e39b785a11aba0265c25466533c230e5b8aef8e8 Mon Sep 17 00:00:00 2001
From: Gregory Comer <gjcomer@meta.com>
Date: Fri, 12 Jun 2026 14:30:27 -0700
Subject: [PATCH] Update

[ghstack-poisoned]
---
 backends/xnnpack/CMakeLists.txt               |   1 +
 backends/xnnpack/runtime/plan/memory_plan.cpp | 105 ++++++++++++++++++
 backends/xnnpack/runtime/plan/memory_plan.h   |  49 ++++++++
 3 files changed, 155 insertions(+)
 create mode 100644 backends/xnnpack/runtime/plan/memory_plan.cpp
 create mode 100644 backends/xnnpack/runtime/plan/memory_plan.h

diff --git a/backends/xnnpack/CMakeLists.txt b/backends/xnnpack/CMakeLists.txt
index c0cf4db0921..29e5279ba00 100644
--- a/backends/xnnpack/CMakeLists.txt
+++ b/backends/xnnpack/CMakeLists.txt
@@ -116,6 +116,7 @@ list(
   backends/xnnpack/runtime/plan/xnn_subgraph.cpp
   backends/xnnpack/runtime/plan/schedule.cpp
   backends/xnnpack/runtime/plan/execution_plan.cpp
+  backends/xnnpack/runtime/plan/memory_plan.cpp
 )
 
 list(TRANSFORM _xnnpack_backend__srcs PREPEND "${EXECUTORCH_ROOT}/")
diff --git a/backends/xnnpack/runtime/plan/memory_plan.cpp b/backends/xnnpack/runtime/plan/memory_plan.cpp
new file mode 100644
index 00000000000..c8ab137a868
--- /dev/null
+++ b/backends/xnnpack/runtime/plan/memory_plan.cpp
@@ -0,0 +1,105 @@
+#include <executorch/backends/xnnpack/runtime/plan/memory_plan.h>
+
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/result.h>
+#include <executorch/runtime/platform/log.h>
+
+#include <cstdlib>
+#include <variant>
+
+namespace executorch::backends::xnnpack::plan {
+
+using namespace graph;
+
+namespace {
+
+runtime::Result<uint64_t> resolve_dim_upper(
+    const DimSizeSpec& dim,
+    const executor::ShapeEnv& shape_env) {
+  int64_t result = dim.offset;
+  for (auto& term : dim.coeffs) {
+    auto& bound = shape_env.bounds[term.sym];
+    ET_CHECK_OR_RETURN_ERROR(
+        bound.max.has_value(),
+        InvalidState,
+        "Symint upper bound is unresolved; cannot size the arena");
+    result += term.coefficient * static_cast<int64_t>(*bound.max);
+  }
+  return static_cast<uint64_t>(result);
+}
+
+runtime::Result<std::vector<uint64_t>> resolve_sizes_upper(
+    const TensorSpec& spec,
+    const executor::ShapeEnv& shape_env) {
+  std::vector<uint64_t> sizes;
+  sizes.reserve(spec.sizes.size());
+  for (auto& dim : spec.sizes) {
+    ET_UNWRAP(resolved, resolve_dim_upper(dim, shape_env));
+    sizes.push_back(resolved);
+  }
+  return sizes;
+}
+
+} // namespace
+
+MemoryPlan create_memory_plan(Graph& graph, ExecutionPlan& execution_plan) {
+  uint32_t num_slots = 0;
+  for (auto& node : graph.nodes) {
+    uint32_t end = node.tag + node.output_count();
+    if (end > num_slots)
+      num_slots = end;
+  }
+
+  std::vector<AllocationInfo> allocations(num_slots);
+  std::vector<TensorSpec> specs(num_slots);
+
+  for (uint32_t n = 0; n < graph.nodes.size(); n++) {
+    auto& node = graph.nodes[n];
+    uint32_t base_slot = node.tag;
+
+    for (uint32_t o = 0; o < node.output_count(); o++) {
+      uint32_t slot = base_slot + o;
+      ValueHandle vh{n, o};
+      specs[slot] = graph.get_tensor_spec(vh);
+
+      if (std::holds_alternative<InputNode>(node.value) ||
+          std::holds_alternative<ConstantNode>(node.value)) {
+        allocations[slot] = ExternalAllocation{};
+      } else {
+        ArenaAllocation a;
+        a.offset = 0;
+        a.size = 0;
+        allocations[slot] = a;
+      }
+    }
+  }
+
+  MemoryPlan mp;
+  mp.arena_size = 0;
+  mp.value_allocations = std::move(allocations);
+  mp.value_specs = std::move(specs);
+  return mp;
+}
+
+runtime::Error MemoryPlan::replan(const executor::ShapeEnv& shape_env) {
+  size_t arena_offset = 0;
+
+  for (size_t i = 0; i < value_allocations.size(); i++) {
+    if (auto* arena = std::get_if<ArenaAllocation>(&value_allocations[i])) {
+      ET_UNWRAP(concrete_sizes, resolve_sizes_upper(value_specs[i], shape_env));
+      ET_UNWRAP(
+          size,
+          core::compute_storage_size(
+              {concrete_sizes.data(), concrete_sizes.size()},
+              value_specs[i].dtype));
+      arena->offset = arena_offset;
+      arena->size = size;
+      arena_offset += size;
+    }
+  }
+
+  arena_size = arena_offset;
+  return runtime::Error::Ok;
+}
+
+} // namespace executorch::backends::xnnpack::plan
diff --git a/backends/xnnpack/runtime/plan/memory_plan.h b/backends/xnnpack/runtime/plan/memory_plan.h
new file mode 100644
index 00000000000..ae5f41664d2
--- /dev/null
+++ b/backends/xnnpack/runtime/plan/memory_plan.h
@@ -0,0 +1,49 @@
+#pragma once
+
+#include <executorch/backends/xnnpack/runtime/core/tensor.h>
+#include <executorch/backends/xnnpack/runtime/executor/shape_env.h>
+#include <executorch/backends/xnnpack/runtime/graph/graph.h>
+#include <executorch/backends/xnnpack/runtime/graph/tensor_spec.h>
+#include <executorch/backends/xnnpack/runtime/plan/execution_plan.h>
+#include <executorch/runtime/core/error.h>
+
+#include <variant>
+#include <vector>
+
+namespace executorch::backends::xnnpack::plan {
+
+/* Describes an allocation planned in the primary memory arena. */
+struct ArenaAllocation {
+  size_t offset;
+  size_t size;
+};
+/*
+ * Describes a standalone dynamic allocation. This memory is not
+ * overlapped or otherwise memory planned.
+ */
+struct DynamicAllocation {};
+/*
+ * Describes externally-owned memory.
+ */
+struct ExternalAllocation {};
+
+using AllocationInfo =
+    std::variant<ArenaAllocation, DynamicAllocation, ExternalAllocation>;
+
+/*
+ * Describes the arena range and/or allocation strategy for each
+ * value in an execution plan.
+ */
+struct MemoryPlan {
+  size_t arena_size;
+  std::vector<AllocationInfo> value_allocations;
+  std::vector<graph::TensorSpec> value_specs;
+
+  runtime::Error replan(const executor::ShapeEnv& shape_env);
+};
+
+MemoryPlan create_memory_plan(
+    graph::Graph& graph,
+    ExecutionPlan& execution_plan);
+
+} // namespace executorch::backends::xnnpack::plan