Parallelize LayerNormalization operator

[robertknight]

This is an operation which appears in most transformer models, although some more recent ones use relatives such as RMSNorm. It can be parallelized by splitting the input over a non-normalized axis and applying normalization to each chunk separately.

From a quick experiment on a 4-core system I can get ~2x speedup quite easily.