[Vulkan] Add the 2D case to Layernorm operator (#110796)
Summary:
We add a 2D implementation to the op [LayerNorm](https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html?fbclid=IwAR00Xi7gt-qo4_LDFo18aaKTxnC4s1vlqk5EREqL0KE0Iz_97-WTlvi0muY)
The current implementation of layer_norm D37407311 supports
- input of 3D and normalized_shape also of 3D, or
- input of 4D with batch dim equal to 1 and normalized_shape of 3D
Since a 2D tensor of [H, W] can be represented as [1, H, W] in shader, we make a straightforward generalization to the case where both input and normalized_shape are of 2D.
Test Plan:
## Before
```
[luwei@devbig984.prn1 ~/fbsource (e09fe4ae4|remote/fbsource/stable...)]$ LD_LIBRARY_PATH=third-party/swiftshader/lib/linux-x64/ buck run fbcode/mode/dev-nosan //xplat/caffe2:pt_vulkan_api_test_bin -- --gtest_filter="*layer_norm_2d*"
Recommended: For faster builds try buck2: replace 'buck' with 'buck2'
NOTE: buck-out/ has changed: look for files in fbsource/buck-out/v2/
'buck2 build --show-output //xplat/caffe2:pt_vulkan_api_test_bin' will print the new output paths.
If you are building in fbsource//xplat and have questions, post in 'Cross Platform Dev Discussions': https://fb.workplace.com/groups/xplat.qa
Targets matching .buckconfig buck2.supported_projects:
{'//xplat/caffe2:pt_vulkan_api_test_bin': '//xplat'}
To suppress this warning: touch ~/.config/.dont_hint_buck2
clang-12: warning: argument unused during compilation: '-pthread' [-Wunused-command-line-argument]
Downloaded 2/4 artifacts, 125.45 Kbytes, 33.3% cache miss (for updated rules)
Building: finished in 4.9 sec (100%) 2637/2637 jobs, 3/2637 updated
Total time: 4.9 sec
BUILD SUCCEEDED
Running main() from third-party/googletest/1.11.0/googletest/googletest/src/gtest_main.cc
Note: Google Test filter = *layer_norm_2d*
[==========] Running 3 tests from 1 test suite.
[----------] Global test environment set-up.
[----------] 3 tests from VulkanAPITest
[ RUN ] VulkanAPITest.layer_norm_2d_small
unknown file: Failure
C++ exception with description "Vulkan layernorm expects 3-dim or 4-dim input!
Exception raised from layer_norm at xplat/caffe2/aten/src/ATen/native/vulkan/ops/Layernorm.cpp:66 (most recent call first):
(no backtrace available)" thrown in the test body.
[ FAILED ] VulkanAPITest.layer_norm_2d_small (56 ms)
[ RUN ] VulkanAPITest.layer_norm_2d_medium
unknown file: Failure
C++ exception with description "Vulkan layernorm expects 3-dim or 4-dim input!
Exception raised from layer_norm at xplat/caffe2/aten/src/ATen/native/vulkan/ops/Layernorm.cpp:66 (most recent call first):
(no backtrace available)" thrown in the test body.
[ FAILED ] VulkanAPITest.layer_norm_2d_medium (0 ms)
[ RUN ] VulkanAPITest.layer_norm_2d_large
unknown file: Failure
C++ exception with description "Vulkan layernorm expects 3-dim or 4-dim input!
Exception raised from layer_norm at xplat/caffe2/aten/src/ATen/native/vulkan/ops/Layernorm.cpp:66 (most recent call first):
(no backtrace available)" thrown in the test body.
[ FAILED ] VulkanAPITest.layer_norm_2d_large (27 ms)
[----------] 3 tests from VulkanAPITest (84 ms total)
[----------] Global test environment tear-down
[==========] 3 tests from 1 test suite ran. (84 ms total)
[ PASSED ] 0 tests.
[ FAILED ] 3 tests, listed below:
[ FAILED ] VulkanAPITest.layer_norm_2d_small
[ FAILED ] VulkanAPITest.layer_norm_2d_medium
[ FAILED ] VulkanAPITest.layer_norm_2d_large
3 FAILED TESTS
```
## After
```
[luwei@devbig984.prn1 ~/fbsource (e09fe4ae4|remote/fbsource/stable...)]$ LD_LIBRARY_PATH=third-party/swiftshader/lib/linux-x64/ buck run fbcode/mode/dev-nosan //xplat/caffe2:pt_vulkan_api_test_bin -- --gtest_filter="*layer_norm_2d*"
Recommended: For faster builds try buck2: replace 'buck' with 'buck2'
NOTE: buck-out/ has changed: look for files in fbsource/buck-out/v2/
'buck2 build --show-output //xplat/caffe2:pt_vulkan_api_test_bin' will print the new output paths.
If you are building in fbsource//xplat and have questions, post in 'Cross Platform Dev Discussions': https://fb.workplace.com/groups/xplat.qa
Targets matching .buckconfig buck2.supported_projects:
{'//xplat/caffe2:pt_vulkan_api_test_bin': '//xplat'}
To suppress this warning: touch ~/.config/.dont_hint_buck2
clang-12: warning: argument unused during compilation: '-pthread' [-Wunused-command-line-argument]
Downloaded 1/3 artifacts, 1.40 Mbytes, 50.0% cache miss (for updated rules)
Building: finished in 5.0 sec (100%) 2637/2637 jobs, 2/2637 updated
Total time: 5.0 sec
BUILD SUCCEEDED
Running main() from third-party/googletest/1.11.0/googletest/googletest/src/gtest_main.cc
Note: Google Test filter = *layer_norm_2d*
[==========] Running 3 tests from 1 test suite.
[----------] Global test environment set-up.
[----------] 3 tests from VulkanAPITest
[ RUN ] VulkanAPITest.layer_norm_2d_small
[ OK ] VulkanAPITest.layer_norm_2d_small (282 ms)
[ RUN ] VulkanAPITest.layer_norm_2d_medium
[ OK ] VulkanAPITest.layer_norm_2d_medium (0 ms)
[ RUN ] VulkanAPITest.layer_norm_2d_large
[ OK ] VulkanAPITest.layer_norm_2d_large (214 ms)
[----------] 3 tests from VulkanAPITest (497 ms total)
[----------] Global test environment tear-down
[==========] 3 tests from 1 test suite ran. (497 ms total)
[ PASSED ] 3 tests.
```
full test result: P848167714
Differential Revision: D50048054
Pull Request resolved: https://github.com/pytorch/pytorch/pull/110796
Approved by: https://github.com/yipjustin
