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.

❌ 2 Cancelled Jobs, 1 Unrelated Failure, 1 Unclassified Failure

As of commit 36674e6 with merge base 6683757 ():

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• periodic / gather-models (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• pull / test-binary-size-linux / linux-job (gh)
  API rate limit exceeded
• pull / test-build-wasm-linux / linux-job (gh)

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-llama-runner-linux (fp32, xnnpack+quantize_kv, linux.2xlarge, executorch-ubuntu-22.04-clang12) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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