Infinite loop in XNNExecutor::resize_outputs (pytorch#17181)

lucylq · web-flow · commit 8e99f3b01447 · 2026-02-03T18:04:02.000-08:00
### Summary
size_t --&gt; ssize_t so it's signed and won't loop indefinitely.

### Test plan
```
cd build
cmake .. -DEXECUTORCH_BUILD_TESTS=ON -DEXECUTORCH_BUILD_XNNPACK=ON
cmake --build . --target backends_xnnpack_test
ctest -R  backends_xnnpack_test -V
```
diff --git a/backends/xnnpack/runtime/XNNExecutor.cpp b/backends/xnnpack/runtime/XNNExecutor.cpp
@@ -246,7 +246,7 @@ ET_NODISCARD Error XNNExecutor::resize_outputs(Span<EValue*> args) const {
     if (out_tensor->scalar_type() == ScalarType::Long) {
       int64_t* data_64 = out_tensor->mutable_data_ptr<int64_t>();
       const int32_t* data_32 = out_tensor->const_data_ptr<int32_t>();
-      for (size_t j = out_tensor->numel() - 1; j >= 0; --j) {
+      for (ssize_t j = out_tensor->numel() - 1; j >= 0; --j) {
         data_64[j] = data_32[j];
       }
     }
diff --git a/backends/xnnpack/test/runtime/test_xnnexecutor.cpp b/backends/xnnpack/test/runtime/test_xnnexecutor.cpp
@@ -11,6 +11,7 @@
 #include <gtest/gtest.h>
 #include <xnnpack.h>
 
+using executorch::aten::Tensor;
 using executorch::backends::xnnpack::delegate::XNNExecutor;
 using executorch::runtime::Error;
 using executorch::runtime::EValue;
@@ -95,3 +96,96 @@ TEST(XNNExecutorTest, ArgumentWithTooManyDimensions) {
   // Check for invalid number of dimensions should fail without stack overflow.
   EXPECT_EQ(executor.prepare_args(stack_args), Error::InvalidArgument);
 }
+
+// Tests that resize_outputs correctly converts int32 indices to int64.
+TEST(XNNExecutorTest, ResizeOutputsWithLongTensorConvertsInt32ToInt64) {
+  XNNExecutor executor({});
+  xnn_runtime_t rt = nullptr;
+  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);
+
+  std::vector<size_t> in_dims = {1, 4, 4, 1}, out_dims = {1, 2, 2, 1};
+  uint32_t input_id = XNN_INVALID_VALUE_ID;
+  uint32_t value_id = XNN_INVALID_VALUE_ID;
+  uint32_t index_id = XNN_INVALID_VALUE_ID;
+
+  ASSERT_EQ(
+      xnn_status_success,
+      xnn_define_tensor_value(
+          subgraph,
+          xnn_datatype_fp32,
+          in_dims.size(),
+          in_dims.data(),
+          nullptr,
+          0,
+          XNN_VALUE_FLAG_EXTERNAL_INPUT,
+          &input_id));
+  ASSERT_EQ(
+      xnn_status_success,
+      xnn_define_tensor_value(
+          subgraph,
+          xnn_datatype_fp32,
+          out_dims.size(),
+          out_dims.data(),
+          nullptr,
+          1,
+          XNN_VALUE_FLAG_EXTERNAL_OUTPUT,
+          &value_id));
+  ASSERT_EQ(
+      xnn_status_success,
+      xnn_define_tensor_value(
+          subgraph,
+          xnn_datatype_int32,
+          out_dims.size(),
+          out_dims.data(),
+          nullptr,
+          2,
+          XNN_VALUE_FLAG_EXTERNAL_OUTPUT,
+          &index_id));
+  ASSERT_EQ(
+      xnn_status_success,
+      xnn_define_argmax_pooling_2d(
+          subgraph, 0, 0, 0, 0, 2, 2, input_id, value_id, index_id, 0));
+
+  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_float;
+  TensorFactory<executorch::aten::ScalarType::Long> tf_long;
+
+  auto input = tf_float.make(
+      {1, 4, 4, 1}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+  auto out_value = tf_float.make({1, 2, 2, 1}, {0, 0, 0, 0});
+  auto out_index = tf_long.make({1, 2, 2, 1}, {0, 0, 0, 0});
+
+  EValue ev_in(input), ev_val(out_value), ev_idx(out_index);
+  std::array<EValue*, 3> args = {&ev_in, &ev_val, &ev_idx};
+  Span<EValue*> span(args.data(), 3);
+
+  ASSERT_EQ(executor.prepare_args(span), Error::Ok);
+  executorch::ET_RUNTIME_NAMESPACE::BackendExecutionContext context;
+  ASSERT_EQ(executor.forward(context), Error::Ok);
+  ASSERT_EQ(executor.resize_outputs(span), Error::Ok);
+
+  Tensor& result = args[2]->toTensor();
+  ASSERT_EQ(result.scalar_type(), executorch::aten::ScalarType::Long);
+
+  /*
+  Input 4x4:          Output values:   Output indices:
+  1  2  | 3  4        6  | 8           3 | 3
+  5  6  | 7  8        14 | 16          3 | 3
+  ------|-----
+  9  10 |11 12
+  13 14 |15 16
+
+  Each 2x2 quadrant → max value + index of max (3 = bottom-right).
+  */
+  for (ssize_t i = 0; i < result.numel(); ++i) {
+    EXPECT_EQ(result.const_data_ptr<int64_t>()[i], 3);
+  }
+}

Original file line number	Diff line number	Diff line change
`@@ -246,7 +246,7 @@ ET_NODISCARD Error XNNExecutor::resize_outputs(Span<EValue*> args) const {`
`246`	`246`	`if (out_tensor->scalar_type() == ScalarType::Long) {`
`247`	`247`	`int64_t* data_64 = out_tensor->mutable_data_ptr<int64_t>();`
`248`	`248`	`const int32_t* data_32 = out_tensor->const_data_ptr<int32_t>();`
`249`		`- for (size_t j = out_tensor->numel() - 1; j >= 0; --j) {`
	`249`	`+ for (ssize_t j = out_tensor->numel() - 1; j >= 0; --j) {`
`250`	`250`	`data_64[j] = data_32[j];`
`251`	`251`	`}`
`252`	`252`	`}`