[decomp] Use var_mean in native_batch_norm decomposition

[peterbell10]

Stack from ghstack (oldest at bottom):

• std/var: support floating point correction value #94073
• WelfordOps: Remove combine_t and use acc_scalar_t instead #94522
• -> [decomp] Use var_mean in native_batch_norm decomposition #94140
• [inductor] Avoid re-computing mean in lowering for aten.var_mean #94139
• [inductor] Count bytes can't read from buffers that are never written #94142
• [decomp] Decompose std/std_mean into aten.var/var_mean #94072

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• ROCm jobs fail to access AMD apt repo

✅ No Failures

As of commit c9e2a65:
💚 Looks good so far! There are no failures yet. 💚

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

[ngimel]

Did you check that inductor perf didn't change?

[ngimel]

why does tolerance change here?

[peterbell10]

aten.var_mean uses a different algorithm from aten.mean which ends up being slightly more precise. 5e-7 is still incredibly good for half-precision though. The default rtol for torch.testing is 1e-5.

[peterbell10]

Did you check that inductor perf didn't change?

This improves perf by removing the duplicate mean calculation. However it's very slight since the second mean was being fused with the sum of square deviations in the variance. In the following example, I see a 0.6% speedup from 366 us to 364 us

import torch
import torch._dynamo
from torch._inductor import config
config.debug = True

a = torch.nn.BatchNorm3d(10).train().cuda()
b = torch.rand(10, 10, 16, 64, 64, device="cuda")

@torch._dynamo.optimize()
def fn(x):
    return a(x)

_ = fn(b)
%timeit fn(b)