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})); +}