It's not the prettiest code, but I used this to arrive at the conclusion that for large numbers of features, it's still slow:

import torch
import timeit
import gc
import numpy

batch_sizes = [8,16,32,64,128]
channels    = [2,32,256,1024]
features    = [(16,),
               (32,),
               (64,),
               (128,),
               (256,),
               (1024,),
               (10240,),
               (102400,),
               (32,32),
               (64,64),
               (128,128),
               (32,32,32),
               (64,64,64)]

def run_bn(input, running_mean, running_var, weight, bias, training, backward):
    out = torch.nn.functional.batch_norm(input, running_mean, running_var, weight, bias, training, 0.1, 1e-5)
    if backward:
        grads = torch.autograd.grad(out, [input, weight, bias], torch.ones_like(out))
    torch.cuda.synchronize()

for bs in batch_sizes:
    for c in channels:
        for f in features:
            shape = (bs, c)+f
            size = numpy.prod(shape)
            if size < 100_000_000:
                running_mean = torch.randn(c, device='cuda')
                running_var = torch.randn(c, device='cuda').exp()
                weight = torch.randn(c, device='cuda', requires_grad=True)
                bias = torch.randn(c, device='cuda', requires_grad=True)
                input = torch.randn(shape, device='cuda', requires_grad=True)
                gc.collect()
                for training in [True, False]:
                    run_bn(input, running_mean, running_var, weight, bias, training, training)
                    torch.cuda.synchronize()
                    res1 = timeit.timeit('run_bn(input, running_mean, running_var, weight, bias, training, training)', number=100, globals=globals())
                    with torch.backends.cudnn.flags(enabled=False):
                      res2 = timeit.timeit('run_bn(input, running_mean, running_var, weight, bias, training, training)', number=100, globals=globals())
                    print ('{} | {} | {} | {} | {:.2f} | {:.2f} | {:.1f}'.format(bs, c, f, training, res1, res2, res1/res2))

Seems that the last rebase was less harmless than it looked...

I know what the problem is. You need to add (aten, native_batch_norm) to aten/src/ATen/core/aten_intern ed_strings.h (and maybe others, but definitely this one).

CC @bwasti, this isn't going to be very good dev experience.

So it turns out that for fp16 even when using the variance in save_var, it doesn't work. Previously, this didn't affect the forward as we kept mean and var in accscalar_t variable within the kernel. Now that separated it into two kernels for performance reasons, we see that keeping it in fp16 isn't enough.
My conclusion is that we would would want to have save_mean/save_var as fp32 for fp16 batch norm. (This is part of what is done in CuDNN, but only for the stats that aren't handed to the user.)

I'm not entirely sure there isn't one, but I don't see the immediate connection to the failing TestScript.test_weak_script_function.

I think the correctness is there, though there are still performance improvements, in particular for the backward. What I'm not 100% certain about is whether removing thnn_batch_norm has compatibility implications.

Thanks a lot for doing this! I'm self-assigning so I remember to take a look. In the mean time, could you revert the submodule change?

No worries, thanks for looking!
The kernel logic itself didn't change much, I mostly switched from DeviceTensor to PackedTensorAccessors. However, I changed the splitting of the kernel.
The kernel bits that changed:

• I amended the old "inference" kernel to the batch_norm_transform_input_kernel to deal both with
  running_mean/running_var (train=False) and save_mean/save_invstd (train=True).
  The former needs to calculate invstd from running_var, while the latter takes it directly.
  So now this is the only thing calculating out = w * (inp - mean) * invstd + b.
• I removed the bit done by batch_norm_transform_input_kernel from the old training forward kernel
  and renamed it to batch_norm_collect_statistics_kernel. It now only computes save_mean / save_var and updates the running_mean / running_var as needed.
• This splitting of the training forward allows a favourable thread/block-parametrisation in the
  transform_input call, which used to be one main bottleneck in the forward.

save_mean/save_std have been changed to accscalar_t because we need the precision (previously, this was in a single kernel and also kept as accscalar_t and no-one cared about the poor backward). One might ask whether it would be desirable to (optionally) support accscalar_t weights for half (as cudnn does), but I didn't do this yet.

Finally, I also changed how the statistics gathering is parallelised by swapping dimensions 0 and 2 to have the larger as 2 in the calculation of the packed_accessor. This is because the parallelisation is done based on dimension 2 only.

Just a heads up: I've discussed kernels with the great folks working on apex sync batchnorm and implemented a backward kernel that is closer to what they use. If you wanted, I could drop in the updated backward kernel (which does similar changes as above, but in addition increases the accuracy of computation). Depending on your preference, I could add it in this PR, but by default, I'll keep it for a later PR (right now the backward kernel mostly the old THCUNN one).

ROCm doesn't seem to like my use of sqrt:

/opt/rocm/hcc/bin/ld.lld: error: relocation R_AMDGPU_REL32_LO cannot be used against symbol sqrtf; recompile with -fPIC
>>> defined in /tmp/tmp.PsfYz7ywhH/caffe2_hip_generated_Normalization.cu.kernel.bc-gfx900.isabin
>>> referenced by /tmp/tmp.PsfYz7ywhH/caffe2_hip_generated_Normalization.cu.kernel.bc-gfx900.isabin

Somehow I feel that I shouldn't change the compilation flags fundamentally here without consulting with someone. Any ideas?

Also, I assumed that CPU kernels didn't change much, right?

The CPU kernels are quite literally adapted from THNN.
I'll try to get a cleanish before/after benchmark (so far, I've mostly compared to cudnn).

Regarding further optimization:

• The welford-style forward statistics is the obvious thing I've not done. (I didn't change much kernel code except that immediately required).
• The backward would get a similar surgery (splitting into "reduction" and "pointwise") that seems, in extreme shapes (32x10x40_000) can have a rather large (10-100x) impact. Based on my discussion with the @mcarilli and @jjsjann123, my main focus for the backward kernel was getting the numerical precision a bit up for fp32 (and also we hope to share kernels between sync bn and this here).
  (I have an implementation of the kernel as an external module that looks correctish at https://gist.github.com/t-vi/82a46dc87eceae303a4f805147f82310 , but I didn't benchmark yet.)

With 32 bit indexing, anything from the "real-world" benchmark taking more than 1msec seems to be not slower in native:

At last!

Very nice speed-up! :)