SIGN IN SIGN UP

native_layer_norm (for width dim) (#3001)

Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/3001

We implement `native_layer_norm` which has 3 outputs
- normalization of the input tensor according to the given `normalized_shape`
- mean
- 1/sqrt(var + eps)

```
func: native_layer_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight, Tensor? bias, float eps) -> (Tensor, Tensor, Tensor)
```

According to SS-JIA's suggestion, a model specific implementation is more performant and preferred to a generic one. So we implemented the op in the following optimized way
- our current use case has `normalized_shape` of len 1, namely we do the normalization through computing the mean and var at the last width dim
- we do the computation in just one shader `native_layer_norm.glsl` without invoking the shaders to compute mean and var respectively
- we use [Welford's online algorithm](https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm) to compute mean and variance in one pass

Reviewed By: SS-JIA, jorgep31415

Differential Revision: D56005629

fbshipit-source-id: 096c2e2f04b95f1f5c9205c4827091169771978c
W
Wei Lu committed
74576e83ddc4995baa4b544178f68df8a26db4a6
Parent: 075fe40
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> on 4/15/2024, 9:10:57 PM