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

copyrightly · facebook-github-bot · commit c45e45f44e15 · 2023-10-11T15:02:34.000-07:00
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 Reviewed By: yipjustin Differential Revision: D50048054
diff --git a/aten/src/ATen/native/vulkan/ops/Layernorm.cpp b/aten/src/ATen/native/vulkan/ops/Layernorm.cpp
@@ -62,20 +62,22 @@ Tensor layer_norm(
   _check_layer_norm_inputs(input_arg, normalized_shape, weight_opt, bias_opt);
 
   TORCH_CHECK(
-      input_arg.dim() == 3 || input_arg.dim() == 4,
-      "Vulkan layernorm expects 3-dim or 4-dim input!");
+      input_arg.dim() >= 2 && input_arg.dim() <= 4,
+      "Vulkan layernorm expects input of 2d, 3d or 4d!");
 
   const Tensor input = input_arg.is_vulkan() ? input_arg : input_arg.vulkan();
   const vTensor& v_input = convert(input);
   const IntArrayRef v_input_sizes = v_input.sizes();
 
   TORCH_CHECK(
-      input_arg.dim() == 3 || v_input_sizes[Layout::Activation4D::batch] == 1,
+      input_arg.dim() == 2 || input_arg.dim() == 3 ||
+          v_input_sizes[Layout::Activation4D::batch] == 1,
       "Vulkan layernorm expects batch dim == 1 when the input is 4-dimensional!");
 
   TORCH_CHECK(
-      normalized_shape.size() == 3,
-      "Vulkan layernorm expects normalized_shape to have length 3, i.e. [C, H, W]");
+      (normalized_shape.size() == 2 && input_arg.dim() == 2) ||
+          (normalized_shape.size() == 3 && input_arg.dim() >= 3),
+      "Vulkan layernorm expects normalized_shape and input dimensions match for 2d and 3d input; normalized_shape of 3d is expected for 4d input!");
 
   TORCH_CHECK(
       weight_opt->defined() && bias_opt->defined(),
@@ -107,7 +109,9 @@ Tensor layer_norm(
   } block{
       v_input.extents(),
       safe_downcast<int32_t>(volume),
-      safe_downcast<int32_t>((v_input_sizes[input_arg.dim() - 3] - 1) % 4),
+      safe_downcast<int32_t>(
+          input_arg.dim() >= 3 ? (v_input_sizes[input_arg.dim() - 3] - 1) % 4
+                               : 0),
       safe_downcast<float>(eps)};
 
   api::UniformParamsBuffer params(context, block);
diff --git a/aten/src/ATen/test/vulkan_api_test.cpp b/aten/src/ATen/test/vulkan_api_test.cpp
@@ -2391,17 +2391,75 @@ TEST_F(VulkanAPITest, layer_norm_invalid_inputs) {
       1e-05,
       false);
   }, ::c10::Error);
+}
 
-  // Act: input has too few dimensions
-  EXPECT_THROW({
-    at::layer_norm(
-      at::rand({3, 5}, at::device(at::kCPU).dtype(at::kFloat)).vulkan(),
-      {3, 5},
-      at::rand({3, 5}, at::device(at::kCPU).dtype(at::kFloat)).vulkan(),
-      at::rand({3, 5}, at::device(at::kCPU).dtype(at::kFloat)).vulkan(),
-      1e-05,
-      false);
-  }, ::c10::Error);
+TEST_F(VulkanAPITest, layer_norm_2d_small) {
+  c10::InferenceMode mode;
+
+  const auto input_cpu = at::rand({1, 1}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto input_vulkan = input_cpu.vulkan();
+
+  const auto weight_cpu = at::rand({1, 1}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto weight_vulkan = weight_cpu.vulkan();
+
+  const auto bias_cpu = at::rand({1, 1}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto bias_vulkan = bias_cpu.vulkan();
+
+  const auto output_cpu = at::layer_norm(input_cpu, {1, 1}, weight_cpu, bias_cpu, 1e-05, false);
+  const auto output_vulkan = at::layer_norm(input_vulkan, {1, 1}, weight_vulkan, bias_vulkan, 1e-05, false);
+
+  const auto check = almostEqual(output_cpu, output_vulkan.cpu());
+  if (!check) {
+    showRtol(output_cpu, output_vulkan.cpu());
+  }
+
+  ASSERT_TRUE(check);
+}
+
+TEST_F(VulkanAPITest, layer_norm_2d_medium) {
+  c10::InferenceMode mode;
+
+  const auto input_cpu = at::rand({5, 7}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto input_vulkan = input_cpu.vulkan();
+
+  const auto weight_cpu = at::rand({5, 7}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto weight_vulkan = weight_cpu.vulkan();
+
+  const auto bias_cpu = at::rand({5, 7}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto bias_vulkan = bias_cpu.vulkan();
+
+  const auto output_cpu = at::layer_norm(input_cpu, {5, 7}, weight_cpu, bias_cpu, 1e-05, false);
+  const auto output_vulkan = at::layer_norm(input_vulkan, {5, 7}, weight_vulkan, bias_vulkan, 1e-05, false);
+
+  const auto check = almostEqual(output_cpu, output_vulkan.cpu());
+  if (!check) {
+    showRtol(output_cpu, output_vulkan.cpu());
+  }
+
+  ASSERT_TRUE(check);
+}
+
+TEST_F(VulkanAPITest, layer_norm_2d_large) {
+  c10::InferenceMode mode;
+
+  const auto input_cpu = at::rand({139, 109}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto input_vulkan = input_cpu.vulkan();
+
+  const auto weight_cpu = at::rand({139, 109}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto weight_vulkan = weight_cpu.vulkan();
+
+  const auto bias_cpu = at::rand({139, 109}, at::device(at::kCPU).dtype(at::kFloat));
+  const auto bias_vulkan = bias_cpu.vulkan();
+
+  const auto output_cpu = at::layer_norm(input_cpu, {139, 109}, weight_cpu, bias_cpu, 1e-05, false);
+  const auto output_vulkan = at::layer_norm(input_vulkan, {139, 109}, weight_vulkan, bias_vulkan, 1e-05, false);
+
+  const auto check = almostEqual(output_cpu, output_vulkan.cpu());
+  if (!check) {
+    showRtol(output_cpu, output_vulkan.cpu());
+  }
+
+  ASSERT_TRUE(check);
 }
 
 TEST_F(VulkanAPITest, layer_norm_3d_small) {
