XNNPACK graph runtime: arena allocator and shape environment

backends/xnnpack/runtime/executor/arena.cpp

@@ -0,0 +1,18 @@
+#include <executorch/backends/xnnpack/runtime/executor/arena.h>
+
+namespace executorch::backends::xnnpack::executor {
+
+runtime::Error Arena::resize(size_t new_size) {
+  if (new_size <= size) {
+    return runtime::Error::Ok;
+  }
+  auto* new_buffer = new (std::nothrow) uint8_t[new_size];
+  if (new_buffer == nullptr) {
+    return runtime::Error::MemoryAllocationFailed;
+  }
+  buffer.reset(new_buffer);
+  size = new_size;
+  return runtime::Error::Ok;
+}
+
+} // namespace executorch::backends::xnnpack::executor

backends/xnnpack/runtime/executor/arena.h

@@ -0,0 +1,28 @@
+#pragma once
+
+#include <executorch/runtime/core/error.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+namespace executorch::backends::xnnpack::executor {
+
+/*
+ * Provides a block of growable, contiguous memory.
+ */
+struct Arena {
+  std::unique_ptr<uint8_t[]> buffer;
+  size_t size = 0;
+
+  inline void* data() {
+    return buffer.get();
+  }
+
+  // Grows the arena to at least `new_size` bytes. The arena is
+  // never shrunk. Re-allocation does not preserve existing contents.
+  // On allocation failure the arena is left unchanged.
+  runtime::Error resize(size_t new_size);
+};
+
+} // namespace executorch::backends::xnnpack::executor

backends/xnnpack/runtime/executor/shape_env.cpp

@@ -0,0 +1,70 @@
+#include <executorch/backends/xnnpack/runtime/executor/shape_env.h>
+
+namespace executorch::backends::xnnpack::executor {
+
+using executorch::runtime::Span;
+
+ShapeEnv::ShapeEnv(uint32_t num_symints) : bounds(num_symints) {}
+
+runtime::Error ShapeEnv::specialize(
+    Span<const graph::TensorSpec> specs,
+    Span<core::Tensor> values) {
+  for (auto& b : bounds) {
+    b.min = 1;
+    b.max = {};
+  }
+
+  if (specs.size() != values.size()) {
+    return runtime::Error::InvalidArgument;
+  }
+
+  for (size_t i = 0; i < specs.size(); i++) {
+    auto& spec = specs[i];
+    auto& tensor = values[i];
+
+    if (spec.sizes.size() != tensor.sizes.size()) {
+      return runtime::Error::InvalidArgument;
+    }
+
+    for (size_t d = 0; d < spec.sizes.size(); d++) {
+      auto& dim = spec.sizes[d];
+      auto concrete = static_cast<int64_t>(tensor.sizes[d]);
+
+      if (dim.is_constant()) {
+        if (dim.offset != concrete) {
+          return runtime::Error::InvalidArgument;
+        }
+        continue;
+      }
+
+      if (dim.coeffs.size() == 1 && dim.coeffs[0].coefficient == 1) {
+        auto sym = dim.coeffs[0].sym;
+        if (sym >= bounds.size()) {
+          // Spec references a symint outside this env's range.
+          return runtime::Error::Internal;
+        }
+        // Solve sym = concrete - offset. A dim is always >= 1, so a value
+        // smaller than the offset means the concrete shape can't satisfy
+        // the spec.
+        int64_t solved_signed = concrete - dim.offset;
+        if (solved_signed < 1) {
+          return runtime::Error::InvalidArgument;
+        }
+        auto solved = static_cast<uint64_t>(solved_signed);
+        auto& bound = bounds[sym];
+        // `min` accumulates the largest solved value and `max` the smallest;
+        // if any two occurrences disagree, min > max flags the contradiction.
+        bound.min = std::max(bound.min, solved);
+        bound.max = bound.max ? std::min(*bound.max, solved) : solved;
+        if (bound.min > *bound.max) {
+          return runtime::Error::InvalidArgument;
+        }
+        continue;
+      }
+    }
+  }
+
+  return runtime::Error::Ok;
+}
+
+} // namespace executorch::backends::xnnpack::executor

backends/xnnpack/runtime/executor/shape_env.h

@@ -0,0 +1,47 @@
+#pragma once
+
+#include <executorch/backends/xnnpack/runtime/core/tensor.h>
+#include <executorch/backends/xnnpack/runtime/graph/tensor_spec.h>
+#include <executorch/runtime/core/error.h>
+#include <executorch/runtime/core/span.h>
+
+#include <cstdint>
+#include <optional>
+
+namespace executorch::backends::xnnpack::executor {
+
+/*
+ * Specifies the lower and optional upper bounds for a shape value.
+ * When max is empty, the value can be arbitrarily large.
+ */
+struct ShapeBound {
+  uint64_t min = 1;
+  std::optional<uint64_t> max = {};
+};
+
+/*
+ * Tracks symint values and provides logic to specialize symints
+ * based on concrete inputs. This class implements a restricted
+ * subset of the PyTorch ShapeEnv logic.
+ */
+struct ShapeEnv {
+  /*
+   * Symint bounds, solved from concrete inputs by specialize(). For
+   * example, if an input tensor is [1, s0] and given concrete
+   * shape [1, 10], then s0 is known to be 10.
+   */
+  std::vector<ShapeBound> bounds;
+
+  ShapeEnv() = default;
+  ShapeEnv(uint32_t num_symints);
+
+  /*
+   * Specialize the bounds for a given set of concrete tensors, solving for the
+   * symints that appear in the specs. Each call resets the bounds.
+   */
+  runtime::Error specialize(
+      runtime::Span<const graph::TensorSpec> specs,
+      runtime::Span<core::Tensor> values);
+};
+
+} // namespace executorch::backends::xnnpack::executor

backends/xnnpack/test/CMakeLists.txt

@@ -42,8 +42,9 @@ target_include_directories(
 )
 
 # Graph runtime unit tests.
-set(_graph_runtime_test_srcs runtime/test_quant_params.cpp
-                             runtime/test_graph_builder.cpp
+set(_graph_runtime_test_srcs
+    runtime/test_quant_params.cpp runtime/test_graph_builder.cpp
+    runtime/test_shape_env.cpp runtime/test_arena.cpp
 )
 
 et_cxx_test(

backends/xnnpack/test/runtime/test_arena.cpp

@@ -0,0 +1,38 @@
+#include <gtest/gtest.h>
+
+#include <executorch/backends/xnnpack/runtime/executor/arena.h>
+
+using namespace executorch::backends::xnnpack::executor;
+using executorch::runtime::Error;
+
+TEST(TestArena, initial_empty) {
+  Arena arena;
+  EXPECT_EQ(arena.size, 0u);
+  EXPECT_EQ(arena.data(), nullptr);
+}
+
+TEST(TestArena, grow_allocates) {
+  Arena arena;
+  EXPECT_EQ(arena.resize(128), Error::Ok);
+  EXPECT_EQ(arena.size, 128u);
+  EXPECT_NE(arena.data(), nullptr);
+}
+
+TEST(TestArena, shrink_is_noop) {
+  Arena arena;
+  ASSERT_EQ(arena.resize(128), Error::Ok);
+  void* data_before = arena.data();
+
+  // A smaller (or equal) request neither reallocates nor shrinks.
+  EXPECT_EQ(arena.resize(64), Error::Ok);
+  EXPECT_EQ(arena.size, 128u);
+  EXPECT_EQ(arena.data(), data_before);
+}
+
+TEST(TestArena, grow_again) {
+  Arena arena;
+  ASSERT_EQ(arena.resize(64), Error::Ok);
+  EXPECT_EQ(arena.resize(256), Error::Ok);
+  EXPECT_EQ(arena.size, 256u);
+  EXPECT_NE(arena.data(), nullptr);
+}