Fix out-of-bounds traversal in reduce_util for non-contiguous tensors

kernels/portable/cpu/util/reduce_util.h

@@ -107,11 +107,10 @@ void apply_on_flat_ix_with_dim_mask_and_base(
         // Reset dim_index[curr_dim] to 0. We need to update curr_index
         // accordingly. Reseting dim_index[curr_dim] from in.size(curr_dim)
         // to 0 means we need to subtract in.size(curr_dim) * strides[curr_dim]
-        // from curr_index. However in.size(curr_dim) * strides[curr_dim] is
-        // equal to strides[curr_dim - 1]. Notice that curr_dim > 0 at this
-        // point in the execution
+        // from curr_index. This is only equal to strides[curr_dim - 1] when
+        // the tensor is contiguous, so it cannot be simplified to that.
         dim_index[curr_dim] = 0;
-        curr_index -= strides[curr_dim - 1];
+        curr_index -= in.size(curr_dim) * strides[curr_dim];
 
         // Decrease current dim
         curr_dim--;

kernels/quantized/test/op_dequantize_test.cpp

@@ -317,3 +317,54 @@ TEST(OpDequantizeOutTest, DequantizePerChannel) {
   test_per_channel_dtype<ScalarType::Byte>();
   test_per_channel_dtype<ScalarType::Char>();
 }
+
+// Per-channel dequantize on a channels-last input. Each element's expected
+// value depends only on its channel, never on where it sits in memory, so this
+// only passes if the kernel honors the tensor's dim order. Before the
+// reduce_util.h carry-over fix this wrote past the end of the output. See
+// issue #16429.
+TEST(OpDequantizeOutTest, DequantizePerChannelChannelsLast) {
+  et_pal_init();
+  TensorFactory<ScalarType::Char> tf;
+  TensorFactory<ScalarType::Double> tf_double;
+  TensorFactory<ScalarType::Long> tf_long;
+  TensorFactory<ScalarType::Float> tfo;
+
+  std::vector<int8_t> logical = {
+      -20, -59,  -22, -40, 127, -108, -57, 117, 24,  -103, 48,  -110,
+      80,  15,   -10, -75, -77, -46,  -12, -66, 35,  -87,  -50, -80,
+      12,  -127, 107, 91,  115, -54,  -6,  -6,  -41, 46,   42,  -83};
+  const double s0 = 0.0016989057185128331;
+  const double s1 = 0.001776964869350195;
+
+  Tensor input = tf.channels_last_like(tf.make({2, 2, 3, 3}, logical));
+  Tensor scale = tf_double.make({2}, {s0, s1});
+  Tensor zero_point = tf_long.make({2}, {0, 0});
+
+  std::vector<float> expected_logical(36);
+  for (int n = 0; n < 2; ++n) {
+    for (int c = 0; c < 2; ++c) {
+      for (int hw = 0; hw < 9; ++hw) {
+        size_t i = n * 18 + c * 9 + hw;
+        expected_logical[i] = static_cast<float>(logical[i]) *
+            static_cast<float>(c == 0 ? s0 : s1);
+      }
+    }
+  }
+  Tensor expected =
+      tfo.channels_last_like(tfo.make({2, 2, 3, 3}, expected_logical));
+
+  Tensor out = tfo.zeros_channels_last({2, 2, 3, 3});
+  dequantize_per_channel_out(
+      input,
+      scale,
+      zero_point,
+      /*axis=*/1,
+      /*quant_min=*/-128,
+      /*quant_max=*/127,
+      ScalarType::Char,
+      optional<ScalarType>(),
+      out);
+
+  EXPECT_TENSOR_CLOSE(out, expected);
+}

kernels/test/CMakeLists.txt

@@ -411,6 +411,7 @@ if(TARGET quantized_kernels)
   set(_quantized_kernels_test_sources
       "${EXECUTORCH_ROOT}/kernels/quantized/test/op_add_test.cpp"
       "${EXECUTORCH_ROOT}/kernels/quantized/test/op_choose_qparams_test.cpp"
+      "${EXECUTORCH_ROOT}/kernels/quantized/test/op_dequantize_test.cpp"
       "${EXECUTORCH_ROOT}/kernels/quantized/test/op_embedding2b_test.cpp"
       "${EXECUTORCH_ROOT}/kernels/quantized/test/op_embedding4b_test.cpp"
       "${EXECUTORCH_ROOT}/kernels/quantized/test/op_embedding_test.cpp"

kernels/test/op_any_test.cpp

@@ -208,3 +208,29 @@ TEST_F(OpAnyOutTest, TestAnyDimsOutEmptyDimList) {
   op_any_dims_out(self, opt_dim_list, keepdim, out);
   EXPECT_TENSOR_CLOSE(out, out_expected);
 }
+
+// Reducing over more than one dim walks the shared dim-mask traversal in
+// reduce_util.h. Its carry-over used to assume
+// size(d) * stride(d) == stride(d - 1), which only holds for a contiguous
+// tensor, so on a channels-last input it skipped a block of elements and read
+// past the end of the buffer. See issue #16429.
+//
+// The single true value below sits at channels-last physical index 19, inside
+// the range the broken traversal never visited, so channel 1 was reported as
+// false while channel 0 picked up out-of-bounds memory and was reported true.
+TEST_F(OpAnyOutTest, ChannelsLastMultiDimReduction) {
+  TensorFactory<ScalarType::Float> tf;
+  TensorFactory<ScalarType::Bool> tf_bool;
+
+  std::vector<float> data(36, 0.0f);
+  data[27] = 1.0f; // logical (n=1, c=1, h=0, w=0)
+  Tensor in = tf.channels_last_like(tf.make({2, 2, 3, 3}, data));
+
+  Tensor out = tf_bool.zeros_channels_last({1, 2, 1, 1});
+  int64_t dims[3] = {0, 2, 3};
+  op_any_dims_out(in, ArrayRef<int64_t>{dims, 3}, /*keepdim=*/true, out);
+
+  Tensor expected =
+      tf_bool.channels_last_like(tf_bool.make({1, 2, 1, 1}, {false, true}));
+  EXPECT_TENSOR_EQ(out, expected);
+}