Use LoopVectorization for faster Loops

[avik-pal]

Recently I refactored the code base to use loops on CPU instead of broadcasting. This makes the code quite a bit faster but more importantly allows us to easily swap in LoopVectorization

See commit history of LuxDL/LuxLib.jl#97 for more details.

Improvements over NNlib Functions

I don't think NNlib will accept LoopVectorization as a dependency. So we implement them here itself

• batched_mul --> A batched_matmul that checks if the Array is on CPU and if it can be loop vectorized then we loop vectorize else we forward to NNlib.batched_mul.
• conv --> Bypass the CPU conv routines in fused_conv with ones written using Loop Vectoization.
• pooling operations

Implementations where LoopVectorization will help

• Activation
• AlphaDropout
• Dropout
• Group Normalization
• CPU implementations for Normalization Improvements to Normalization Functions #1012
• Bias Activation

Reductions

• We have the infrastructure setup in impl/fast_ops.jl. For LoopedArrayOp we need to simply use VectorizedStatistics.jl (and maybe VectorizedReductions.jl).

Automatic Differentiation

We don't need to worry about ChainRules. It anyways has rrule defined as of now. But Enzyme really not happy with Loop Vectorization. Use custom rules for the following:

• Activation
• Dropout
• Alpha Dropout
• Group Norm
• CPU implementations for Normalization Improvements to Normalization Functions #1012
• Bias Activation