pytorch

21dc1d25 - [Vulkan] Add the 2D case to Layernorm operator (#110796)

[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