🐛 Describe the bug

torch.compile(backend='inductor') silently produces wrong results when adaptive_avg_pool2d is fused with flatten + sum (or mean). The bug triggers when batch > 1 and the flattened reduction dimension (C × pool_h × pool_w) exceeds 1024 with channels not a multiple of 16.

This is a silent correctness bug — no error is raised, the output is simply wrong. The error is large (up to 1e5 in max-autotune mode) and persists in float64, ruling out precision noise.

Minimal repro

import torch

x = torch.randn(3, 33, 8, 8, dtype=torch.float64, device='cuda')

def f(x):
    return torch.nn.functional.adaptive_avg_pool2d(x, 7).flatten(1).sum(dim=-1)

# Eager — correct
print(f"eager = {f(x).tolist()}")

# Inductor — wrong
torch._dynamo.reset()
compiled = torch.compile(f, backend='inductor')(x)
print(f"inductor = {compiled.tolist()}")
print(f"max_diff = {(f(x) - compiled).abs().max().item():.2e}")
# max_diff ~ 5e+00 to 8e+01 depending on seed

Expected behavior

Inductor output should match eager output. adaptive_avg_pool2d(x, 7).flatten(1).sum(dim=-1) is a standard global-pool-then-reduce pattern used in CNN classifiers.

Trigger conditions

Root cause analysis

The bug is not in adaptive_avg_pool2d itself — pool-only output is correct. It occurs when inductor fuses the pool with flatten + sum into a single reduction kernel. The tiling/indexing for the fused kernel appears to miscalculate offsets for batch elements beyond the first (batch[0] is always correct, batch[1+] are wrong).

The alignment to multiples of 16 and the 1024-element threshold suggest a Triton reduction tile size boundary issue in the fused kernel's index computation.

Additional observations

• All compile modes affected: default, reduce-overhead, max-autotune (worst: diff up to 2.2e5)
• float64 diffs range from 1e0 to 1e2 — not a precision issue
• Replacing sum with mean also triggers the bug
• pool + sum(spatial_dims) without flatten is correct — the flatten reshape changes the reduction structure
• Non-square spatial sizes (e.g., 8×10) also trigger
• Different pool output sizes (5, 7) all trigger under the same conditions

Versions

• PyTorch: 2.13.0.dev20260420+cu126 (nightly)
• CUDA: 12.6
• GPU: Tesla T4
• OS: Ubuntu 20.04, Python 3.11

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo