Skip xnn_reshape_runtime when the input shapes have not changed (#24219)

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<EValue*> 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<EValue*> 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,34 +145,50 @@ ET_NODISCARD Error XNNExecutor::prepare_args(Span<EValue*> 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,
           "XNNPACK backend accepts tensors with at most %d dims, but got %zu",
           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<int>(dim_order[j]));
+        input_dims_.push_back(tensor->size(static_cast<int>(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);

backends/xnnpack/runtime/XNNExecutor.h

@@ -51,6 +51,12 @@ class XNNExecutor {
   std::shared_ptr<XNNWeightsCache> weights_cache_;
   std::atomic<bool> in_use_{false};
   std::atomic<bool> 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<size_t> input_dims_;
+  std::vector<size_t> reshaped_input_dims_;
+  bool runtime_reshaped_ = false;
 
  public:
   XNNExecutor(std::shared_ptr<XNNWorkspace> workspace)

backends/xnnpack/test/runtime/test_xnnexecutor.cpp

@@ -197,3 +197,251 @@ TEST(XNNExecutorTest, ResizeOutputsWithLongTensorConvertsInt32ToInt64) {
     EXPECT_EQ(result.const_data_ptr<int64_t>()[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<size_t>& 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<EValue*> args) {
+  Span<EValue*> 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<float> values(const EValue& ev) {
+  const Tensor& t = ev.toTensor();
+  return std::vector<float>(
+      t.const_data_ptr<float>(), t.const_data_ptr<float>() + 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<xnn_subgraph, decltype(&xnn_delete_subgraph)> 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<executorch::aten::ScalarType::Float> tf;
+  // Sized for the largest shape; prepare_args resizes it to each input's shape.
+  EValue out(tf.make(
+      {4, 3},
+      std::vector<float>(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<float>{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<float>{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<float>{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<float>{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<xnn_subgraph, decltype(&xnn_delete_subgraph)> 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<executorch::aten::ScalarType::Float> tf;
+  EValue out(tf.make(
+      {4, 3},
+      std::vector<float>(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<executorch::aten::SizesType, 2> 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<size_t> 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<xnn_subgraph, decltype(&xnn_delete_subgraph)> 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<executorch::aten::ScalarType::Float> tf;
+  EValue out(tf.make(
+      {4, 3},
+      std::vector<float>(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<float>{11, 22, 33, 44, 55, 66}));
+
+  EValue b_bad(tf.make({4, 3}, std::vector<float>(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<float>{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<float>{101, 201, 301, 101, 201, 301}));
+
+  ASSERT_EQ(run(executor, {&b2, &row, &out}), Error::Ok);
+  EXPECT_EQ(values(out), (std::vector<float>{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<float>{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<xnn_subgraph, decltype(&xnn_delete_subgraph)> 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<executorch::aten::ScalarType::Float> tf;
+  EValue out(tf.make(
+      {2, 3},
+      std::vector<float>(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<float>{11, 12, 13, 14, 15, 16}));
+
+  scalar = -1;
+  ASSERT_EQ(run(executor, {&a, &s, &out}), Error::Ok);
+  EXPECT_EQ(values(out), (std::vector<float>{0, 1, 2, 3, 4, 5}));
+}