From 37a4c97f38a1049f56586b8cd67a6b2a92566f4f Mon Sep 17 00:00:00 2001 From: Gregory Comer Date: Fri, 12 Jun 2026 14:30:10 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/xnnpack/CMakeLists.txt | 2 + backends/xnnpack/runtime/executor/arena.cpp | 18 ++ backends/xnnpack/runtime/executor/arena.h | 28 +++ .../xnnpack/runtime/executor/shape_env.cpp | 70 ++++++ backends/xnnpack/runtime/executor/shape_env.h | 47 ++++ backends/xnnpack/test/CMakeLists.txt | 5 +- backends/xnnpack/test/runtime/test_arena.cpp | 38 +++ .../xnnpack/test/runtime/test_shape_env.cpp | 217 ++++++++++++++++++ 8 files changed, 423 insertions(+), 2 deletions(-) create mode 100644 backends/xnnpack/runtime/executor/arena.cpp create mode 100644 backends/xnnpack/runtime/executor/arena.h create mode 100644 backends/xnnpack/runtime/executor/shape_env.cpp create mode 100644 backends/xnnpack/runtime/executor/shape_env.h create mode 100644 backends/xnnpack/test/runtime/test_arena.cpp create mode 100644 backends/xnnpack/test/runtime/test_shape_env.cpp diff --git a/backends/xnnpack/CMakeLists.txt b/backends/xnnpack/CMakeLists.txt index f02d48e1141..32bafa4ec59 100644 --- a/backends/xnnpack/CMakeLists.txt +++ b/backends/xnnpack/CMakeLists.txt @@ -109,6 +109,8 @@ list( backends/xnnpack/runtime/graph/graph.cpp backends/xnnpack/runtime/graph/graph_builder.cpp backends/xnnpack/runtime/operators/operator.cpp + backends/xnnpack/runtime/executor/arena.cpp + backends/xnnpack/runtime/executor/shape_env.cpp ) list(TRANSFORM _xnnpack_backend__srcs PREPEND "${EXECUTORCH_ROOT}/") diff --git a/backends/xnnpack/runtime/executor/arena.cpp b/backends/xnnpack/runtime/executor/arena.cpp new file mode 100644 index 00000000000..29cd9b11f66 --- /dev/null +++ b/backends/xnnpack/runtime/executor/arena.cpp @@ -0,0 +1,18 @@ +#include + +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 diff --git a/backends/xnnpack/runtime/executor/arena.h b/backends/xnnpack/runtime/executor/arena.h new file mode 100644 index 00000000000..2b6b5faee12 --- /dev/null +++ b/backends/xnnpack/runtime/executor/arena.h @@ -0,0 +1,28 @@ +#pragma once + +#include + +#include +#include +#include + +namespace executorch::backends::xnnpack::executor { + +/* + * Provides a block of growable, contiguous memory. + */ +struct Arena { + std::unique_ptr 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 diff --git a/backends/xnnpack/runtime/executor/shape_env.cpp b/backends/xnnpack/runtime/executor/shape_env.cpp new file mode 100644 index 00000000000..35ac70f646f --- /dev/null +++ b/backends/xnnpack/runtime/executor/shape_env.cpp @@ -0,0 +1,70 @@ +#include + +namespace executorch::backends::xnnpack::executor { + +using executorch::runtime::Span; + +ShapeEnv::ShapeEnv(uint32_t num_symints) : bounds(num_symints) {} + +runtime::Error ShapeEnv::specialize( + Span specs, + Span 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(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(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 diff --git a/backends/xnnpack/runtime/executor/shape_env.h b/backends/xnnpack/runtime/executor/shape_env.h new file mode 100644 index 00000000000..6ef2e5bb93c --- /dev/null +++ b/backends/xnnpack/runtime/executor/shape_env.h @@ -0,0 +1,47 @@ +#pragma once + +#include +#include +#include +#include + +#include +#include + +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 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 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 specs, + runtime::Span values); +}; + +} // namespace executorch::backends::xnnpack::executor diff --git a/backends/xnnpack/test/CMakeLists.txt b/backends/xnnpack/test/CMakeLists.txt index d8d57e81f9d..a9b432d2de4 100644 --- a/backends/xnnpack/test/CMakeLists.txt +++ b/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( diff --git a/backends/xnnpack/test/runtime/test_arena.cpp b/backends/xnnpack/test/runtime/test_arena.cpp new file mode 100644 index 00000000000..294e482a31f --- /dev/null +++ b/backends/xnnpack/test/runtime/test_arena.cpp @@ -0,0 +1,38 @@ +#include + +#include + +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); +} diff --git a/backends/xnnpack/test/runtime/test_shape_env.cpp b/backends/xnnpack/test/runtime/test_shape_env.cpp new file mode 100644 index 00000000000..80eb1a81e4c --- /dev/null +++ b/backends/xnnpack/test/runtime/test_shape_env.cpp @@ -0,0 +1,217 @@ +#include + +#include +#include + +#include + +using namespace executorch::backends::xnnpack::core; +using namespace executorch::backends::xnnpack::executor; +using namespace executorch::backends::xnnpack::graph; +using executorch::runtime::Error; + +static TensorSpec make_spec(std::vector sizes) { + return TensorSpec{.dtype = DType::Float32, .sizes = std::move(sizes)}; +} + +static Tensor make_tensor(std::vector sizes) { + return Tensor{.dtype = DType::Float32, .sizes = std::move(sizes)}; +} + +template +static std::vector make_tensors(Args&&... args) { + std::vector v; + v.reserve(sizeof...(args)); + (v.push_back(std::forward(args)), ...); + return v; +} + +TEST(TestShapeEnv, single_sym_solves) { + // spec: [s0], tensor: [42] => s0 = 42 + ShapeEnv env(1); + std::vector specs = {make_spec({DimSizeSpec::sym(0)})}; + auto values = make_tensors(make_tensor({42})); + + EXPECT_TRUE( + env.specialize( + {specs.data(), specs.size()}, {values.data(), values.size()}) == + Error::Ok); + EXPECT_EQ(env.bounds[0].min, 42u); + ASSERT_TRUE(env.bounds[0].max.has_value()); + EXPECT_EQ(*env.bounds[0].max, 42u); +} + +TEST(TestShapeEnv, constant_dim_valid) { + // spec: [const(10)], tensor: [10] => ok + ShapeEnv env(0); + std::vector specs = {make_spec({DimSizeSpec::constant(10)})}; + auto values = make_tensors(make_tensor({10})); + + EXPECT_TRUE( + env.specialize( + {specs.data(), specs.size()}, {values.data(), values.size()}) == + Error::Ok); +} + +TEST(TestShapeEnv, constant_dim_mismatch) { + // spec: [const(10)], tensor: [5] => fail + ShapeEnv env(0); + std::vector specs = {make_spec({DimSizeSpec::constant(10)})}; + auto values = make_tensors(make_tensor({5})); + + EXPECT_FALSE( + env.specialize( + {specs.data(), specs.size()}, {values.data(), values.size()}) == + Error::Ok); +} + +TEST(TestShapeEnv, sym_with_offset) { + // spec: [1*s0 + 3], tensor: [13] => s0 = 10 + ShapeEnv env(1); + DimSizeSpec dim = {.coeffs = {{.sym = 0, .coefficient = 1}}, .offset = 3}; + std::vector specs = {make_spec({dim})}; + auto values = make_tensors(make_tensor({13})); + + EXPECT_TRUE( + env.specialize( + {specs.data(), specs.size()}, {values.data(), values.size()}) == + Error::Ok); + EXPECT_EQ(env.bounds[0].min, 10u); + ASSERT_TRUE(env.bounds[0].max.has_value()); + EXPECT_EQ(*env.bounds[0].max, 10u); +} + +TEST(TestShapeEnv, same_sym_consistent) { + // Two dims both use s0, both give s0 = 5 => ok, bounds tightened to [5, 5] + ShapeEnv env(1); + std::vector specs = { + make_spec({DimSizeSpec::sym(0)}), + make_spec({DimSizeSpec::sym(0)}), + }; + auto values = make_tensors(make_tensor({5}), make_tensor({5})); + + EXPECT_TRUE( + env.specialize( + {specs.data(), specs.size()}, {values.data(), values.size()}) == + Error::Ok); + EXPECT_EQ(env.bounds[0].min, 5u); + ASSERT_TRUE(env.bounds[0].max.has_value()); + EXPECT_EQ(*env.bounds[0].max, 5u); +} + +TEST(TestShapeEnv, same_sym_contradictory) { + // Two dims both use s0, one gives s0=5, other gives s0=7 => fail + ShapeEnv env(1); + std::vector specs = { + make_spec({DimSizeSpec::sym(0)}), + make_spec({DimSizeSpec::sym(0)}), + }; + auto values = make_tensors(make_tensor({5}), make_tensor({7})); + + EXPECT_FALSE( + env.specialize( + {specs.data(), specs.size()}, {values.data(), values.size()}) == + Error::Ok); +} + +TEST(TestShapeEnv, multiple_syms) { + // spec: [s0, s1], tensor: [3, 7] => s0=3, s1=7 + ShapeEnv env(2); + std::vector specs = { + make_spec({DimSizeSpec::sym(0), DimSizeSpec::sym(1)}), + }; + auto values = make_tensors(make_tensor({3, 7})); + + EXPECT_TRUE( + env.specialize( + {specs.data(), specs.size()}, {values.data(), values.size()}) == + Error::Ok); + EXPECT_EQ(env.bounds[0].min, 3u); + EXPECT_EQ(*env.bounds[0].max, 3u); + EXPECT_EQ(env.bounds[1].min, 7u); + EXPECT_EQ(*env.bounds[1].max, 7u); +} + +TEST(TestShapeEnv, spec_value_count_mismatch) { + ShapeEnv env(1); + std::vector specs = { + make_spec({DimSizeSpec::sym(0)}), + make_spec({DimSizeSpec::sym(0)}), + }; + auto values = make_tensors(make_tensor({5})); + + EXPECT_FALSE( + env.specialize( + {specs.data(), specs.size()}, {values.data(), values.size()}) == + Error::Ok); +} + +TEST(TestShapeEnv, dim_count_mismatch) { + // spec has 2 dims, tensor has 1 dim + ShapeEnv env(1); + std::vector specs = { + make_spec({DimSizeSpec::sym(0), DimSizeSpec::constant(10)}), + }; + auto values = make_tensors(make_tensor({5})); + + EXPECT_FALSE( + env.specialize( + {specs.data(), specs.size()}, {values.data(), values.size()}) == + Error::Ok); +} + +TEST(TestShapeEnv, multi_term_skipped) { + // Multi-term dim: 2*s0 + 3*s1 + 0 is left unsolved, bounds stay at defaults + ShapeEnv env(2); + DimSizeSpec dim = { + .coeffs = {{.sym = 0, .coefficient = 2}, {.sym = 1, .coefficient = 3}}, + .offset = 0}; + std::vector specs = {make_spec({dim})}; + auto values = make_tensors(make_tensor({100})); + + EXPECT_TRUE( + env.specialize( + {specs.data(), specs.size()}, {values.data(), values.size()}) == + Error::Ok); + // Bounds should remain at defaults (min=1, max=nullopt) + EXPECT_EQ(env.bounds[0].min, 1u); + EXPECT_FALSE(env.bounds[0].max.has_value()); + EXPECT_EQ(env.bounds[1].min, 1u); + EXPECT_FALSE(env.bounds[1].max.has_value()); +} + +TEST(TestShapeEnv, non_unit_coefficient_skipped) { + // Single term but coefficient != 1: 2*s0, left unsolved + ShapeEnv env(1); + DimSizeSpec dim = {.coeffs = {{.sym = 0, .coefficient = 2}}, .offset = 0}; + std::vector specs = {make_spec({dim})}; + auto values = make_tensors(make_tensor({10})); + + EXPECT_TRUE( + env.specialize( + {specs.data(), specs.size()}, {values.data(), values.size()}) == + Error::Ok); + EXPECT_EQ(env.bounds[0].min, 1u); + EXPECT_FALSE(env.bounds[0].max.has_value()); +} + +TEST(TestShapeEnv, repeated_specialize_resets) { + // Calling specialize twice should reset bounds from first call + ShapeEnv env(1); + std::vector specs = {make_spec({DimSizeSpec::sym(0)})}; + auto values1 = make_tensors(make_tensor({42})); + auto values2 = make_tensors(make_tensor({7})); + + EXPECT_TRUE( + env.specialize( + {specs.data(), specs.size()}, {values1.data(), values1.size()}) == + Error::Ok); + EXPECT_EQ(env.bounds[0].min, 42u); + + EXPECT_TRUE( + env.specialize( + {specs.data(), specs.size()}, {values2.data(), values2.size()}) == + Error::Ok); + EXPECT_EQ(env.bounds[0].min, 7u); + EXPECT_EQ(*env.bounds[0].max, 7u); +}