pytorch - 💡(How to fix) Fix `torch.compile` produces different output for model with `aten.select` parent + batch pointwise fusion pattern [2 comments, 2 participants]

import os
os.environ["TRITON_BACKENDS_IN_TREE"] = "1"

import torch
import torch.nn as nn

class SelectParentPointwiseModel(nn.Module):
    def __init__(self, channels=64, height=32, width=32):
        super().__init__()
        self.channels = channels
        self.conv = nn.Conv2d(3, channels, kernel_size=3, padding=1)
        self.linear1 = nn.Linear(channels, channels)
        self.linear2 = nn.Linear(channels, channels)
        self.bn = nn.BatchNorm2d(channels)

    def forward(self, x):
        batch_size = x.shape[0]
        base = self.conv(x)
        base = self.bn(base)

        base_flat = base.flatten(2)
        y1 = self.linear1(base_flat.transpose(1, 2))
        y2 = self.linear2(base_flat.transpose(1, 2))
        y1 = y1.transpose(1, 2).view_as(base)
        y2 = y2.transpose(1, 2).view_as(base)

        # Multiple aten.select from same parent → expand → pointwise
        slice1 = base.select(1, 0).unsqueeze(1).expand_as(base)
        slice2 = base.select(1, 1).unsqueeze(1).expand_as(base)
        slice3 = base.select(1, 2).unsqueeze(1).expand_as(base)
        slice4 = base.select(1, 3).unsqueeze(1).expand_as(base)

        out1 = slice1 + y1
        out2 = slice2 * y2
        out3 = slice3 + y1
        out4 = slice4 * y2
        out5 = torch.add(slice1, y1, alpha=0.5)
        out6 = torch.mul(slice2, y2)

        combined = out1 + out2 + out3 + out4 + out5 + out6
        result = combined * base + combined / (base + 1e-8)
        return result


device = "cuda"
torch.manual_seed(42)
model = SelectParentPointwiseModel().to(device).eval()
x = torch.randn(4, 3, 32, 32, device=device)

# Eager
with torch.no_grad():
    eager_out = model(x)

# Compiled
torch._dynamo.reset()
compiled = torch.compile(model, backend="inductor")
with torch.no_grad():
    comp_out = compiled(x)

diff = (eager_out.float() - comp_out.float()).abs().max().item()
print(f"Max diff: {diff}")
print(f"Match: {torch.allclose(eager_out, comp_out, atol=1e-5, rtol=1e-4)}")
# Expected: max_diff > 1.0 — significant mismatch

---

PyTorch version: 2.12.0.dev20260327+cu126
CUDA used to build PyTorch: 12.6
OS: Ubuntu 22.04.5 LTS (x86_64) — WSL2
Python version: 3.10.12
GPU: NVIDIA GeForce RTX 3060 Laptop GPU