Zero grad_input in max_pool2d_with_indices_backward

[anurag2796]

Fixes #21686.

max_pool_backward_impl scatter-adds into grad_input, writing only the argmax positions:

grad_input_ptr[maxindex] += grad_output_ptr[index];

but the kernel only resizes grad_input and never clears it. Since the memory planner recycles arena
buffers across ops and across iterations, every non-argmax element accumulates onto stale data.

ATen does the same accumulation and zeroes first —
DilatedMaxPool2d.cpp
calls gradInput.zero_() before dispatching, and
MaxPoolKernel.cpp
holds the character-for-character same += loop. The loop was ported here; the zero_() was not.

Effect

Any trainable graph containing Conv2d -> MaxPool2d gets corrupted weight gradients, because the
polluted grad_input feeds convolution_backward. Pool-only graphs are unaffected — there the maxpool
grad_input is the gradient w.r.t. the network input and is discarded before reaching a parameter,
which is why the bug looks so selective.

The failure mode depends on model size, and the large-model case is the concerning one:

model	before	after
Conv2d->MaxPool2d->Linear (67,642 params)	2.309544 -> NaN (291/300 NaN steps)	2.309544 -> 0.003952 (0 NaN)
ResNet-18 GroupNorm (11,177,538 params)	0.693080 -> 0.814618, 0 NaN, never improves	0.693080 -> 0.018742

At 11M parameters there is no NaN at all — the loss just wanders and ends above where it started.

Validation

• No regression: five other trainable models (strided-conv CNN, conv-only, pool-only, MLP, a second
  strided variant) produce byte-identical loss curves before and after, under identical flags and a
  fixed batch, with only this file changed.
• Gradient correctness: against an independently-built LiteRT/TF graph on the same task and init, the
  patched curve is bit-identical for the first 5 steps (max abs deviation 0.000e+00) and tracks
  within <1% from step 100 to 300.
• On device: cross-compiled arm64-v8a (NDK 27.1.12297006), run on a Snapdragon 845 / Android 10
  handset pinned to the big cluster. Same before/after, and the fixed curve agrees with the macOS arm64
  host to 3.07e-04 max deviation, within 1 ULP on the final loss.
• Cost: +0.06% median step latency on the handset (239.94 ms -> 240.09 ms); peak RSS 10.38 -> 10.48 MB.

All-zero is the correct bit pattern for every dtype ET_SWITCH_FLOATHBF16_TYPES covers
(float / half / bfloat16).

Reproduced on v1.3.1 and verified still present on main and in v1.4.0.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 12 Awaiting Approval

As of commit efc8221 with merge base 48741ac ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: anurag2796 / name: anurag2796 (efc8221)

[meta-cla]

Hi @anurag2796!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[Copilot]

Pull request overview

This PR fixes incorrect gradient computation in the portable CPU implementation of max_pool2d_with_indices_backward by ensuring grad_input is zero-initialized before the kernel’s scatter-add loop runs. This aligns ExecuTorch behavior with ATen’s implementation and prevents stale arena-buffer contents from corrupting downstream gradients (notably convolution_backward) in trainable graphs.

Changes:

• Add <cstring> include for byte-level initialization utilities.
• Zero grad_input via memset(..., 0, grad_input.nbytes()) after resize_tensor and before max_pool_backward_impl performs += scatter-add.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!