From 8b4516a28371a6ef42649358a183b7793808fb9a Mon Sep 17 00:00:00 2001 From: Min Guo Date: Wed, 28 Oct 2026 11:07:07 -0700 Subject: [PATCH] Skip xnn_reshape_runtime when the input shapes have not changed (#24219) Summary: `XNNExecutor::prepare_args` reshapes every external input and calls `xnn_reshape_runtime` on every execute, which re-runs each operator's reshape and re-plans memory even when no shape changed. Models built from many small XNNPACK partitions pay this per partition per call. In the Phoenix CARMM body synthesizer (38 XNNPACK delegates), a simpleperf profile of HandBodyReplay puts `prepare_args` at 7.8% of per-frame CPU, 5.7% of it inside `xnn_reshape_runtime`, versus 14.3% for the delegates' actual compute. prepare_args now collects each input's rank and dims and reshapes only when they differ from the last successful reshape (or the runtime was never reshaped). Input shapes are the whole reshape input: the runtime partitions work for the pthreadpool it was created with (`xnn_create_runtime_v4`), whose size cannot change; `UseNThreadsThreadPoolGuard` only limits threads per dispatch. This is safe because XNNPACK re-binds every data pointer in `xnn_setup_runtime_v2`, and when one runtime grows a shared workspace it rebases the value and op-workspace pointers of all runtimes sharing it (`runtime.c`). Two things captured at reshape are reused across skipped calls: per-thread scratch sized from that fixed pool, and the packed-weights address looked up from the weights cache, whose slots are only appended or cleared on release (`XNNWeightsCache.cpp`). A failed reshape clears the flag, so the next call reshapes again. Output resizing is unchanged. Edited in both the xplat and fbcode copies, which are kept byte-identical. Reviewed By: JakeStevens Differential Revision: D127875036 --- backends/xnnpack/runtime/XNNExecutor.cpp | 50 ++-- backends/xnnpack/runtime/XNNExecutor.h | 6 + .../xnnpack/test/runtime/test_xnnexecutor.cpp | 248 ++++++++++++++++++ 3 files changed, 288 insertions(+), 16 deletions(-) diff --git a/backends/xnnpack/runtime/XNNExecutor.cpp b/backends/xnnpack/runtime/XNNExecutor.cpp index 937054633c9..5001d400bb3 100644 --- a/backends/xnnpack/runtime/XNNExecutor.cpp +++ b/backends/xnnpack/runtime/XNNExecutor.cpp @@ -83,9 +83,10 @@ ET_NODISCARD Error XNNExecutor::initialize( * Prepares the args for XNNPACK Runtime. * * Creates an array of xnn_externals_values from the EValues passed in. - * Reshapes all the external input tensors, in case any input shapes have - * changed. The reshapes the entire runtime, propagating shape information - * through the runtime. + * If any input shape differs from the last reshape (or the runtime was never + * reshaped), reshapes all the external input tensors and then the entire + * runtime, propagating shape information through it. XNNPACK re-binds data + * pointers in setup, so unchanged shapes need no reshape. * * Note: the external ids given to the external tensors in the XNNPACK * runtime correspond to their index in the list of arg passed into @@ -114,6 +115,7 @@ ET_NODISCARD Error XNNExecutor::prepare_args(Span args) { // Create xnn_externals_value from evalue args xnn_status status; + input_dims_.clear(); for (uint32_t i = 0; i < externals_.size(); ++i) { if (i < input_ids_.size()) { externals_[i].id = input_ids_[i]; @@ -134,7 +136,7 @@ ET_NODISCARD Error XNNExecutor::prepare_args(Span args) { executorch::aten::DimOrderType dim_order[kTensorDimensionLimit]; - // Reshape runtime inputs + // Collect runtime input shapes if (i < input_ids_.size()) { size_t num_dims = tensor->dim(); Error err = @@ -143,7 +145,6 @@ ET_NODISCARD Error XNNExecutor::prepare_args(Span args) { err == Error::Ok, Internal, "Failed to retrieve dim order from tensor!"); - size_t dims[XNN_MAX_TENSOR_DIMS]; ET_CHECK_OR_RETURN_ERROR( num_dims <= XNN_MAX_TENSOR_DIMS, InvalidArgument, @@ -151,26 +152,43 @@ ET_NODISCARD Error XNNExecutor::prepare_args(Span args) { XNN_MAX_TENSOR_DIMS, num_dims); + input_dims_.push_back(num_dims); for (int j = 0; j < num_dims; ++j) { - dims[j] = tensor->size(static_cast(dim_order[j])); + input_dims_.push_back(tensor->size(static_cast(dim_order[j]))); } - status = - xnn_reshape_external_value(runtime_.get(), ext_id, num_dims, dims); + } + } + + if (!runtime_reshaped_ || input_dims_ != reshaped_input_dims_) { + runtime_reshaped_ = false; + size_t offset = 0; + for (uint32_t i = 0; i < input_ids_.size(); ++i) { + const size_t num_dims = input_dims_[offset]; + status = xnn_reshape_external_value( + runtime_.get(), + externals_[i].id, + num_dims, + input_dims_.data() + offset + 1); ET_CHECK_OR_RETURN_ERROR( status == xnn_status_success, Internal, "Internal Error: Reshape Input Tensor Failed with code: %s", xnn_status_to_string(status)); + offset += num_dims + 1; } - } - // Propagate Input Shape and Memory Plan for increased allocation - status = xnn_reshape_runtime(runtime_.get()); - ET_CHECK_OR_RETURN_ERROR( - status == xnn_status_success, - Internal, - "Internal Error: Propagating input shapes failed with code: %s", - xnn_status_to_string(status)); + // Propagate Input Shape and Memory Plan for increased allocation + status = xnn_reshape_runtime(runtime_.get()); + + ET_CHECK_OR_RETURN_ERROR( + status == xnn_status_success, + Internal, + "Internal Error: Propagating input shapes failed with code: %s", + xnn_status_to_string(status)); + + reshaped_input_dims_ = input_dims_; + runtime_reshaped_ = true; + } // Resize output tensors. Error err = resize_outputs(args); diff --git a/backends/xnnpack/runtime/XNNExecutor.h b/backends/xnnpack/runtime/XNNExecutor.h index 7d966b0f6a5..ad70a18e5e7 100644 --- a/backends/xnnpack/runtime/XNNExecutor.h +++ b/backends/xnnpack/runtime/XNNExecutor.h @@ -51,6 +51,12 @@ class XNNExecutor { std::shared_ptr weights_cache_; std::atomic in_use_{false}; std::atomic destroyed_{false}; + // Input shapes, each as its rank followed by its dims, as of the last + // xnn_reshape_runtime. They are the whole reshape input: the runtime's + // threadpool is fixed at xnn_create_runtime_v4 and cannot change size. + std::vector input_dims_; + std::vector reshaped_input_dims_; + bool runtime_reshaped_ = false; public: XNNExecutor(std::shared_ptr workspace) diff --git a/backends/xnnpack/test/runtime/test_xnnexecutor.cpp b/backends/xnnpack/test/runtime/test_xnnexecutor.cpp index 63b9c096b2b..cd7eadacdea 100644 --- a/backends/xnnpack/test/runtime/test_xnnexecutor.cpp +++ b/backends/xnnpack/test/runtime/test_xnnexecutor.cpp @@ -197,3 +197,251 @@ TEST(XNNExecutorTest, ResizeOutputsWithLongTensorConvertsInt32ToInt64) { EXPECT_EQ(result.const_data_ptr()[i * 4], 3); } } + +namespace { + +using executorch::runtime::TensorShapeDynamism; + +// Defines an fp32 external tensor with the given external id. +uint32_t define_external( + xnn_subgraph_t subgraph, + const std::vector& dims, + uint32_t external_id, + uint32_t flags) { + uint32_t id = XNN_INVALID_VALUE_ID; + EXPECT_EQ( + xnn_define_tensor_value( + subgraph, + xnn_datatype_fp32, + dims.size(), + dims.data(), + nullptr, + external_id, + flags, + &id), + xnn_status_success); + return id; +} + +// Runs prepare_args and forward on one executor call. +Error run(XNNExecutor& executor, std::vector args) { + Span span(args.data(), args.size()); + Error err = executor.prepare_args(span); + if (err != Error::Ok) { + return err; + } + executorch::ET_RUNTIME_NAMESPACE::BackendExecutionContext context; + return executor.forward(context); +} + +std::vector values(const EValue& ev) { + const Tensor& t = ev.toTensor(); + return std::vector( + t.const_data_ptr(), t.const_data_ptr() + t.numel()); +} + +} // namespace + +// Tests that a changed input shape is propagated through the runtime, and that +// repeating a shape with new input data still computes on that data. +TEST(XNNExecutorTest, ReshapesWhenInputShapeChanges) { + XNNExecutor executor({}); + et_pal_init(); + ASSERT_EQ(xnn_initialize(nullptr), xnn_status_success); + xnn_subgraph_t subgraph = nullptr; + ASSERT_EQ(xnn_create_subgraph(2, 0, &subgraph), xnn_status_success); + std::unique_ptr auto_subgraph( + subgraph, xnn_delete_subgraph); + const uint32_t input_id = + define_external(subgraph, {2, 3}, 0, XNN_VALUE_FLAG_EXTERNAL_INPUT); + const uint32_t output_id = + define_external(subgraph, {2, 3}, 1, XNN_VALUE_FLAG_EXTERNAL_OUTPUT); + ASSERT_EQ( + xnn_define_clamp(subgraph, 0.0f, 10.0f, input_id, output_id, 0), + xnn_status_success); + xnn_runtime_t rt = nullptr; + ASSERT_EQ(xnn_create_runtime(subgraph, &rt), xnn_status_success); + ASSERT_EQ(executor.initialize(rt, {0}, {1}, {}), Error::Ok); + + TensorFactory tf; + // Sized for the largest shape; prepare_args resizes it to each input's shape. + EValue out(tf.make( + {4, 3}, + std::vector(12, -1.0f), + {}, + TensorShapeDynamism::DYNAMIC_BOUND)); + + EValue in1(tf.make({2, 3}, {-1, 2, 3, 4, 5, 20})); + ASSERT_EQ(run(executor, {&in1, &out}), Error::Ok); + EXPECT_EQ(values(out), (std::vector{0, 2, 3, 4, 5, 10})); + + EValue in2(tf.make({4, 3}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, -12})); + ASSERT_EQ(run(executor, {&in2, &out}), Error::Ok); + EXPECT_EQ( + values(out), (std::vector{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 10, 0})); + + EValue in3(tf.make({2, 3}, {6, 5, 4, 3, 2, 1})); + ASSERT_EQ(run(executor, {&in3, &out}), Error::Ok); + EXPECT_EQ(values(out), (std::vector{6, 5, 4, 3, 2, 1})); + + EValue in4(tf.make({2, 3}, {7, 7, 7, -7, 70, 0})); + ASSERT_EQ(run(executor, {&in4, &out}), Error::Ok); + EXPECT_EQ(values(out), (std::vector{7, 7, 7, 0, 10, 0})); +} + +// Tests that a repeated input shape does not reshape the runtime: a shape set +// on the runtime behind the executor's back survives the next call. +TEST(XNNExecutorTest, SkipsReshapeForRepeatedInputShape) { + XNNExecutor executor({}); + et_pal_init(); + ASSERT_EQ(xnn_initialize(nullptr), xnn_status_success); + xnn_subgraph_t subgraph = nullptr; + ASSERT_EQ(xnn_create_subgraph(2, 0, &subgraph), xnn_status_success); + std::unique_ptr auto_subgraph( + subgraph, xnn_delete_subgraph); + const uint32_t input_id = + define_external(subgraph, {2, 3}, 0, XNN_VALUE_FLAG_EXTERNAL_INPUT); + const uint32_t output_id = + define_external(subgraph, {2, 3}, 1, XNN_VALUE_FLAG_EXTERNAL_OUTPUT); + ASSERT_EQ( + xnn_define_clamp(subgraph, 0.0f, 10.0f, input_id, output_id, 0), + xnn_status_success); + xnn_runtime_t rt = nullptr; + ASSERT_EQ(xnn_create_runtime(subgraph, &rt), xnn_status_success); + ASSERT_EQ(executor.initialize(rt, {0}, {1}, {}), Error::Ok); + + TensorFactory tf; + EValue out(tf.make( + {4, 3}, + std::vector(12, -1.0f), + {}, + TensorShapeDynamism::DYNAMIC_BOUND)); + // A {4, 3} buffer viewed as {2, 3}, so running the runtime at {4, 3} stays in + // bounds. + Tensor input = tf.make( + {4, 3}, + {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, + {}, + TensorShapeDynamism::DYNAMIC_BOUND); + const std::array view_sizes = {2, 3}; + ASSERT_EQ( + executorch::ET_RUNTIME_NAMESPACE::resize_tensor( + input, {view_sizes.data(), view_sizes.size()}), + Error::Ok); + EValue in(input); + ASSERT_EQ(run(executor, {&in, &out}), Error::Ok); + EXPECT_EQ(out.toTensor().size(0), 2); + + const std::vector other_dims = {4, 3}; + ASSERT_EQ( + xnn_reshape_external_value(rt, 0, other_dims.size(), other_dims.data()), + xnn_status_success); + ASSERT_EQ(xnn_reshape_runtime(rt), xnn_status_success); + + ASSERT_EQ(run(executor, {&in, &out}), Error::Ok); + EXPECT_EQ(out.toTensor().size(0), 4); +} + +// Tests that after a failed reshape the next call reshapes again, and that +// swapping the shapes of two inputs is seen as a shape change. +TEST(XNNExecutorTest, ReshapesAfterFailedReshapeAndSwappedInputShapes) { + XNNExecutor executor({}); + et_pal_init(); + ASSERT_EQ(xnn_initialize(nullptr), xnn_status_success); + xnn_subgraph_t subgraph = nullptr; + ASSERT_EQ(xnn_create_subgraph(3, 0, &subgraph), xnn_status_success); + std::unique_ptr auto_subgraph( + subgraph, xnn_delete_subgraph); + const uint32_t a_id = + define_external(subgraph, {2, 3}, 0, XNN_VALUE_FLAG_EXTERNAL_INPUT); + const uint32_t b_id = + define_external(subgraph, {2, 3}, 1, XNN_VALUE_FLAG_EXTERNAL_INPUT); + const uint32_t out_id = + define_external(subgraph, {2, 3}, 2, XNN_VALUE_FLAG_EXTERNAL_OUTPUT); + ASSERT_EQ( + xnn_define_binary( + subgraph, xnn_binary_add, nullptr, a_id, b_id, out_id, 0), + xnn_status_success); + xnn_runtime_t rt = nullptr; + ASSERT_EQ(xnn_create_runtime(subgraph, &rt), xnn_status_success); + ASSERT_EQ(executor.initialize(rt, {0, 1}, {2}, {}), Error::Ok); + + TensorFactory tf; + EValue out(tf.make( + {4, 3}, + std::vector(12, -1.0f), + {}, + TensorShapeDynamism::DYNAMIC_BOUND)); + + EValue a(tf.make({2, 3}, {1, 2, 3, 4, 5, 6})); + EValue b(tf.make({2, 3}, {10, 20, 30, 40, 50, 60})); + ASSERT_EQ(run(executor, {&a, &b, &out}), Error::Ok); + EXPECT_EQ(values(out), (std::vector{11, 22, 33, 44, 55, 66})); + + EValue b_bad(tf.make({4, 3}, std::vector(12, 1.0f))); + EXPECT_NE(run(executor, {&a, &b_bad, &out}), Error::Ok); + + EValue b2(tf.make({2, 3}, {1, 1, 1, 1, 1, 1})); + ASSERT_EQ(run(executor, {&a, &b2, &out}), Error::Ok); + EXPECT_EQ(values(out), (std::vector{2, 3, 4, 5, 6, 7})); + + EValue row(tf.make({1, 3}, {100, 200, 300})); + ASSERT_EQ(run(executor, {&row, &b2, &out}), Error::Ok); + EXPECT_EQ(values(out), (std::vector{101, 201, 301, 101, 201, 301})); + + ASSERT_EQ(run(executor, {&b2, &row, &out}), Error::Ok); + EXPECT_EQ(values(out), (std::vector{101, 201, 301, 101, 201, 301})); + EValue a2(tf.make({2, 3}, {0, 0, 0, 1, 1, 1})); + ASSERT_EQ(run(executor, {&a2, &row, &out}), Error::Ok); + EXPECT_EQ(values(out), (std::vector{100, 200, 300, 101, 201, 301})); +} + +// Tests a rank-0 input in the last position, whose shape is the final entry of +// the cached dims. +TEST(XNNExecutorTest, ReshapesWithRankZeroLastInput) { + XNNExecutor executor({}); + et_pal_init(); + ASSERT_EQ(xnn_initialize(nullptr), xnn_status_success); + xnn_subgraph_t subgraph = nullptr; + ASSERT_EQ(xnn_create_subgraph(3, 0, &subgraph), xnn_status_success); + std::unique_ptr auto_subgraph( + subgraph, xnn_delete_subgraph); + const uint32_t a_id = + define_external(subgraph, {2, 3}, 0, XNN_VALUE_FLAG_EXTERNAL_INPUT); + const uint32_t s_id = + define_external(subgraph, {}, 1, XNN_VALUE_FLAG_EXTERNAL_INPUT); + const uint32_t out_id = + define_external(subgraph, {2, 3}, 2, XNN_VALUE_FLAG_EXTERNAL_OUTPUT); + ASSERT_EQ( + xnn_define_binary( + subgraph, xnn_binary_add, nullptr, a_id, s_id, out_id, 0), + xnn_status_success); + xnn_runtime_t rt = nullptr; + ASSERT_EQ(xnn_create_runtime(subgraph, &rt), xnn_status_success); + ASSERT_EQ(executor.initialize(rt, {0, 1}, {2}, {}), Error::Ok); + + TensorFactory tf; + EValue out(tf.make( + {2, 3}, + std::vector(6, -1.0f), + {}, + TensorShapeDynamism::DYNAMIC_BOUND)); + EValue a(tf.make({2, 3}, {1, 2, 3, 4, 5, 6})); + // TensorFactory gives a rank-0 tensor a null dim order, which + // get_dim_order would pass to memcpy. + float scalar = 10; + executorch::aten::DimOrderType scalar_dim_order[1] = {0}; + executorch::aten::TensorImpl scalar_impl( + executorch::aten::ScalarType::Float, + 0, + nullptr, + &scalar, + scalar_dim_order); + EValue s{Tensor(&scalar_impl)}; + ASSERT_EQ(run(executor, {&a, &s, &out}), Error::Ok); + EXPECT_EQ(values(out), (std::vector{11, 12, 13, 14, 15, 16})); + + scalar = -1; + ASSERT_EQ(run(executor, {&a, &s, &out}), Error::Ok); + EXPECT_EQ(values(out), (std::vector{0, 1, 2, 3, 4, 5})); +}