Improvements to Normalization Functions

[avik-pal]

Current Status

• batchnorm
  • CUDA is heavily optimized and calls into cuDNN
  • batchnorm -- similar to what is done for groupnorm (added in feat: improved fallback BN implementation LuxLib.jl#106)

Native Julia Implementations

• groupnorm has custom kernels for the forward and backward pass. This is achieved by restructuring all the arrays into 4D and constructing either loops or KA kernels. This gives us around 2-3x performance boost over standard cases (verified against Flux as well).
  • This particular form however causes slowdown in the 2D case. Also it is better to rewrite the reduction directly as loops for the backward pass. Similar to perf: fusing activation functions and other misc perf improvements LuxLib.jl#126 -- fixed in perf: rework CPU groupnorm implementation LuxLib.jl#134
• instancenorm -- very similar to groupnorm.
• layernorm
  • We should change the default to what Pytorch does. This case is very simple to optimize using LoopVectorization on CPU and KernelAbstractions on GPU.
  • The general broadcasting case is very hard to optimize, best we could do is fuse into a single GPU kernel but this is not worth it much.

Integration into vendor libraries

• MIOpen for AMDGPU batchnorm. With the new batchnorm kernels, we should test whether this is even worth it. Though I would need someone who has access to ROCm capable GPUs to test this.
  • Seems like the kernels are quite performant, and we are close to cuDNN in performance even with the naive way the kernels are written.