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

[SuryanshSS1011]

Summary

Fixes #16429.

dequantize_per_channel on a channels-last input writes past the end of its output. The reported symptom is wrong values, but the underlying problem is a memory error and it is not in the quantized kernel.

apply_on_flat_ix_with_dim_mask_and_base in kernels/portable/cpu/util/reduce_util.h walks the elements that map to one output index. When the innermost dim is exhausted it rewinds curr_index and carries into the next dim, using:

curr_index -= strides[curr_dim - 1];

The comment above that line states the assumption: it treats strides[curr_dim - 1] as equal to in.size(curr_dim) * strides[curr_dim]. That identity only holds for a contiguous tensor. For a (2, 2, 3, 3) channels-last tensor, strides are [18, 1, 6, 2], and the two disagree at curr_dim = 1 (2 vs 18) and curr_dim = 2 (18 vs 1). The rewind is then wrong and the walk leaves the buffer.

For dequantize_per_channel on that tensor the walk covers indices 0 to 17, skips 18 to 34 entirely, and writes 17 times past the end of a 36 element output. Running it under the existing kernel tests produces wrong values followed by double free or corruption (!prev).

This changes the line to subtract the value the comment describes. For a contiguous tensor the two expressions evaluate to the same number, so that path is unaffected, and it is in the carry-over rather than the inner loop, so it is not on the per-element path.

Most reduction ops do not reach this. op_amax, op_amin, op_mean, op_sum and op_var all guard with tensor_is_default_dim_order(in) and reject non-contiguous input before the traversal runs. The ops that accept it are op_any, op_var_mean and dequantize_per_channel, and I have confirmed the fault in two of them: dequantize_per_channel writes out of bounds, and any.dims_out reads out of bounds and returns a wrong answer.

I took this over adding a tensor_is_default_dim_order guard to dequantize_per_channel to match the other five ops. A guard would stop the corruption but leaves the traversal wrong for op_any and op_var_mean, and #16429 asks for non-contiguous dim order to be supported in portable ops rather than rejected. Happy to switch if you would rather have the guard.

Relaxing the five existing guards is not part of this change.

Test plan

OpDequantizeOutTest.DequantizePerChannelChannelsLast dequantizes a (2, 2, 3, 3) channels-last int8 tensor per channel on axis=1. Each expected value is derived from the element's channel alone, so it passes only if the
kernel honors the tensor's dim order.

OpAnyOutTest.ChannelsLastMultiDimReduction reduces a channels-last tensor over dims {0, 2, 3} with keepdim=true. The single true value sits at channels-last physical index 19, inside the range the old traversal skipped, so before this change channel 1 was reported false while channel 0 picked up out-of-bounds memory and was reported true.

Both fail before this change and pass after it.

This also adds op_dequantize_test.cpp to _quantized_kernels_test_sources in kernels/test/CMakeLists.txt. It is already in kernels/quantized/test/targets.bzl, but the CMake target lists every other op in that directory and omits this one, so those tests do not run in a CMake build and the new one would not either.

Full runs after the change: quantized_kernels_test 73 passed, portable_kernels_test 1624 ran with 1520 passed and 0 failures.

cc @larryliu0820 @manuelcandales @JakeStevens

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21517

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 79c3303 with merge base 6683757 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[SuryanshSS1011]

@pytorchbot label "release notes: ops & kernels"

[nil-is-all]

Thanks for the PR, @SuryanshSS1011. Running CI currently.
cc @MartinPavella @metascroy @Gasoonjia @JacobSzwejbka @JakeStevens

[Gasoonjia]

please highlight it is contiguous memory format; there're two contiguous concepts one is memory format the other is layout.

[Gasoonjia]

Thanks for your great work! Please add an extra test here: https://github.com/pytorch/executorch/blob/main/kernels/portable/cpu/util/test/reduce_test.cpp

[SuryanshSS1011]

Thanks for the review, @Gasoonjia! Rephrased the marked comment to unambiguously highlight contiguous memory format; and added an extra test to reduce_test.cpp: ReduceUtilTest.ApplyOverDimListChannelsLast, covering the traversal at the util level rather than only through ops.

The added test counts visited indices instead of writing through the tensor like the tests above it, so it fails on the out-of-bounds index rather than on the heap corruption that follows.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Suppressed comments (3)

kernels/test/op_any_test.cpp:227

• This test is trying to reason about a channels-last physical index (per the comment above), but the value is injected using a hard-coded flat index into a pre-conversion buffer (data[27]) and then transformed with channels_last_like. That makes it harder to verify the intended physical placement (and easy to “fix” incorrectly later). Consider making the mapping explicit in the test (e.g., compute the flat index from (n,c,h,w) for the pre-conversion layout, or add a brief comment that channels_last_like moves this logical element to physical index 19).

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

kernels/quantized/test/op_dequantize_test.cpp:353

• The expected-value construction relies on hard-coded shape-derived constants (18, 9) and a flattened hw loop. This makes the test brittle if the shape changes (and obscures which logical indices are being exercised). Prefer deriving these from the tensor sizes (e.g., loop h/w explicitly or compute i using H, W, C variables from the shape) so the intent stays correct and readable.

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

kernels/portable/cpu/util/reduce_util.h:113

• Consider making the arithmetic type explicit here to avoid signed/unsigned surprises and to keep the index math consistent with curr_index/strides types (e.g., ensure the multiplication happens in the same width as curr_index). This helps prevent compiler warnings and reduces the risk of overflow/underflow bugs in large tensors.

        curr_index -= in.size(curr_dim) * strides[curr_dim];