[Pytoch][Vulkan] Create context for layernorm (#114701)

copyrightly · facebook-github-bot · commit f96df3b3fc72 · 2023-11-29T09:01:30.000-08:00
Summary: `Layernorm` has two arguments weight and bias which are stored as constant tensors on the CPU and they are transferred to GPU at every inference call. We create a context for this op to avoid the repeated passing. Specifically, we - created `create_layernorm_context` and `run_layernorm_context` in `Layernorm.h` and `Layernorm.cpp` - registered them in `Register.cpp` - rewrote the graph representation of the op in `vulkan_rewrite.cpp` Test Plan: ## Numerical test ``` [luwei@devbig984.prn1 /data/users/luwei/fbsource (b6ccc956c)]$ 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*" 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 Building: finished in 0.1 sec (100%) 339/339 jobs, 0/339 updated Total time: 0.2 sec BUILD SUCCEEDED Running main() from third-party/googletest/1.14.0/googletest/googletest/src/gtest_main.cc Note: Google Test filter = *layer_norm* [==========] Running 10 tests from 1 test suite. [----------] Global test environment set-up. [----------] 10 tests from VulkanAPITest [ RUN ] VulkanAPITest.packed_layer_norm_2d [ OK ] VulkanAPITest.packed_layer_norm_2d (342 ms) [ RUN ] VulkanAPITest.packed_layer_norm_3d [ OK ] VulkanAPITest.packed_layer_norm_3d (284 ms) [ RUN ] VulkanAPITest.packed_layer_norm_4d [ OK ] VulkanAPITest.packed_layer_norm_4d (5 ms) [ RUN ] VulkanAPITest.layer_norm_invalid_inputs [ OK ] VulkanAPITest.layer_norm_invalid_inputs (28 ms) [ RUN ] VulkanAPITest.layer_norm_2d [ OK ] VulkanAPITest.layer_norm_2d (1 ms) [ RUN ] VulkanAPITest.layer_norm_3d [ OK ] VulkanAPITest.layer_norm_3d (2 ms) [ RUN ] VulkanAPITest.layer_norm_4d [ OK ] VulkanAPITest.layer_norm_4d (4 ms) [ RUN ] VulkanAPITest.native_layer_norm_2d [ OK ] VulkanAPITest.native_layer_norm_2d (1 ms) [ RUN ] VulkanAPITest.native_layer_norm_3d [ OK ] VulkanAPITest.native_layer_norm_3d (2 ms) [ RUN ] VulkanAPITest.native_layer_norm_4d [ OK ] VulkanAPITest.native_layer_norm_4d (6 ms) [----------] 10 tests from VulkanAPITest (679 ms total) [----------] Global test environment tear-down [==========] 10 tests from 1 test suite ran. (679 ms total) [ PASSED ] 10 tests. ``` Full test result in P888496077, summary as below ``` [----------] 419 tests from VulkanAPITest (21652 ms total) [----------] Global test environment tear-down [==========] 419 tests from 1 test suite ran. (21652 ms total) [ PASSED ] 418 tests. [ SKIPPED ] 1 test, listed below: [ SKIPPED ] VulkanAPITest.querypool_flushed_shader_log ``` ## Graph representation comparison We created a model using `layer_norm` and traced it as below ``` class MyModel(torch.nn.Module): def __init__(self): super(MyModel, self).__init__() self.layer_norm = torch.nn.LayerNorm(normalized_shape=10) def forward(self, x): return self.layer_norm(x) # Create an instance of the model model = MyModel() # Create a dummy input tensor for tracing input_tensor = torch.randn(1, 10) # Use torch.jit.trace to trace the model and generate a graph traced_model = torch.jit.trace(model, input_tensor) ``` Then we converted the traced model to Vulkan backend using `optimize_for_mobile` ``` from torch.utils import mobile_optimizer vulkan_model = mobile_optimizer.optimize_for_mobile( traced_model, backend="vulkan", preserved_methods=to_preserve ) ``` Then we can print the graph of the `vulkan_model` as `print(vk_model.graph)` - Before this diff ``` %4 : bool = prim::Constant[value=1](), scope: __module.layer_norm # /mnt/xarfuse/uid-602118/33e18f68-seed-nspid4026531836_cgpid32066351-ns-4026531840/torch/nn/functional.py:2546:0 %5 : float = prim::Constant[value=1.0000000000000001e-05](), scope: __module.layer_norm # /mnt/xarfuse/uid-602118/33e18f68-seed-nspid4026531836_cgpid32066351-ns-4026531840/torch/nn/functional.py:2546:0 %14 : int[] = prim::Constant[value=[10]]() %33 : Tensor = aten::to(%x, %53, %30, %31, %31) %10 : Tensor = aten::layer_norm(%33, %14, %self.layer_norm.weight, %self.layer_norm.bias, %5, %4), scope: __module.layer_norm # /mnt/xarfuse/uid-602118/33e18f68-seed-nspid4026531836_cgpid32066351-ns-4026531840/torch/nn/functional.py:2546:0 ``` - after this diff ``` %14 : int[] = prim::Constant[value=[10]]() %47 : Tensor = aten::to(%x, %78, %44, %45, %45) %16 : Tensor = vulkan_prepack::run_layernorm_context(%47, %14, %17) ``` Reviewed By: SS-JIA Differential Revision: D51530478
diff --git a/aten/src/ATen/native/vulkan/ops/Layernorm.cpp b/aten/src/ATen/native/vulkan/ops/Layernorm.cpp
@@ -1,3 +1,9 @@
+#include <ATen/native/vulkan/ops/Layernorm.h>
+#include <ATen/native/vulkan/ops/Utils.h>
+
+#include <ATen/Context.h>
+#include <c10/util/irange.h>
+
 #include <ATen/native/vulkan/ops/Common.h>
 #include <torch/library.h>
 
@@ -11,26 +17,81 @@ namespace at {
 namespace native {
 namespace vulkan {
 namespace ops {
-namespace {
 
-using namespace api::utils;
+LayernormPackedContext::LayernormPackedContext(
+    const c10::optional<Tensor>& weight,
+    const c10::optional<Tensor>& bias,
+    double eps)
+    : unpacked_{c10::AnyType::get()} {
+  packed_.reserve(ListArgs::kNumArgs);
 
-Tensor layer_norm(
-    const at::Tensor& input_arg,
+  TORCH_CHECK(weight, "Weight must be provided!");
+  packed_.emplace_back(weight->vulkan());
+  TORCH_CHECK(bias, "Bias must be provided!");
+  packed_.emplace_back(bias->vulkan());
+  packed_.emplace_back(eps);
+
+  if (!at::globalContext().releaseWeightsWhenPrepacking()) {
+    unpacked_.reserve(ListArgs::kNumArgs);
+    unpacked_.emplace_back(weight);
+    unpacked_.emplace_back(bias);
+    unpacked_.emplace_back(eps);
+  }
+}
+
+LayernormPackedContext LayernormPackedContext::pack(
+    c10::impl::GenericList unpacked) {
+  return LayernormPackedContext(
+      get_optional_tensor(unpacked, ListArgs::kWeight),
+      get_optional_tensor(unpacked, ListArgs::kBias),
+      unpacked.get(ListArgs::kEps).toDouble());
+}
+
+c10::intrusive_ptr<LayernormPackedContext> create_layernorm_context(
+    c10::optional<Tensor>&& weight,
+    c10::optional<Tensor>&& bias,
+    double eps) {
+  return c10::make_intrusive<LayernormPackedContext>(
+      LayernormPackedContext(weight, bias, eps));
+}
+
+Tensor run_layernorm_context(
+    const Tensor& input_arg,
     IntArrayRef normalized_shape,
-    const c10::optional<Tensor>& weight_opt /* optional */,
-    const c10::optional<Tensor>& bias_opt /* optional */,
-    double eps,
-    bool /* cudnn_enable, deprecated */) {
+    const c10::intrusive_ptr<LayernormPackedContext>& layernorm_context) {
   const Tensor input = input_arg.is_vulkan() ? input_arg : input_arg.vulkan();
 
+  const c10::optional<Tensor>& weight_opt =
+      layernorm_context->get_val(LayernormPackedContext::ListArgs::kWeight)
+          .toTensor();
+  const c10::optional<Tensor>& bias_opt =
+      layernorm_context->get_val(LayernormPackedContext::ListArgs::kBias)
+          .toTensor();
+  const float eps = api::utils::safe_downcast<float>(
+      layernorm_context->get_val(LayernormPackedContext::ListArgs::kEps)
+          .toDouble());
+
   // We invoke native_layer_norm which returns a tuple of tensors: <layer_norm,
   // mean, 1/sqrt(var+eps)>, but we only need the first tensor (layer_norm).
   std::tuple<Tensor, Tensor, Tensor> native_layer_norm_output =
       at::native_layer_norm(input, normalized_shape, weight_opt, bias_opt, eps);
   return std::get<0>(native_layer_norm_output);
 }
 
+Tensor layer_norm(
+    const at::Tensor& input_arg,
+    IntArrayRef normalized_shape,
+    const c10::optional<Tensor>& weight_opt /* optional */,
+    const c10::optional<Tensor>& bias_opt /* optional */,
+    double eps,
+    bool /* cudnn_enable, deprecated */) {
+  return run_layernorm_context(
+      input_arg,
+      normalized_shape,
+      c10::make_intrusive<LayernormPackedContext>(
+          LayernormPackedContext(weight_opt, bias_opt, eps)));
+}
+
 #ifdef USE_VULKAN_API
 
 TORCH_LIBRARY_IMPL(aten, Vulkan, m) {
@@ -39,7 +100,6 @@ TORCH_LIBRARY_IMPL(aten, Vulkan, m) {
 
 #endif /* USE_VULKAN_API */
 
-} // namespace
 } // namespace ops
 } // namespace vulkan
 } // namespace native
diff --git a/aten/src/ATen/native/vulkan/ops/Layernorm.h b/aten/src/ATen/native/vulkan/ops/Layernorm.h
@@ -0,0 +1,60 @@
+#pragma once
+
+#ifdef USE_VULKAN_API
+
+#include <ATen/native/vulkan/ops/Common.h>
+#include <ATen/native/vulkan/ops/VulkanPackedContext.h>
+#include <torch/library.h>
+
+namespace at {
+namespace native {
+namespace vulkan {
+namespace ops {
+
+class LayernormPackedContext final : virtual public VulkanPackedContext,
+                                     public torch::jit::CustomClassHolder {
+ private:
+  c10::impl::GenericList unpacked_;
+
+ public:
+  LayernormPackedContext(
+      const c10::optional<Tensor>& weight,
+      const c10::optional<Tensor>& bias,
+      double eps);
+
+  /*
+   * Assigns a name to each index in the unpacked list.
+   */
+  struct ListArgs final {
+    static constexpr uint32_t kWeight = 0u;
+    static constexpr uint32_t kBias = 1u;
+    static constexpr uint32_t kEps = 2u;
+
+    static constexpr uint32_t kNumArgs = 3u;
+  };
+
+  static LayernormPackedContext pack(const c10::impl::GenericList);
+
+  const c10::impl::GenericList unpack() const override {
+    TORCH_CHECK(!unpacked_.empty(), "unpacked_ does not have any elements!");
+
+    return unpacked_;
+  }
+};
+
+c10::intrusive_ptr<LayernormPackedContext> create_layernorm_context(
+    c10::optional<Tensor>&& weight,
+    c10::optional<Tensor>&& bias,
+    double eps);
+
+Tensor run_layernorm_context(
+    const Tensor& input,
+    IntArrayRef normalized_shape,
+    const c10::intrusive_ptr<LayernormPackedContext>& context);
+
+} // namespace ops
+} // namespace vulkan
+} // namespace native
+} // namespace at
+
+#endif /* USE_VULKAN_API */
diff --git a/aten/src/ATen/native/vulkan/ops/Register.cpp b/aten/src/ATen/native/vulkan/ops/Register.cpp
@@ -5,6 +5,7 @@
 #include <ATen/native/vulkan/ops/Common.h>
 #include <ATen/native/vulkan/ops/Convolution.h>
 #include <ATen/native/vulkan/ops/Gru.h>
+#include <ATen/native/vulkan/ops/Layernorm.h>
 #include <ATen/native/vulkan/ops/Lstm.h>
 #include <ATen/native/vulkan/ops/Mm.h>
 #include <ATen/native/vulkan/ops/QuantizedFunctions.h>
@@ -55,6 +56,25 @@ int register_vulkan_linear_packed_context() {
   return 0;
 }
 
+int register_vulkan_layernorm_packed_context() {
+  static auto register_vulkan_layernorm_context =
+      torch::selective_class_<LayernormPackedContext>(
+          "vulkan", TORCH_SELECTIVE_CLASS("LayernormPackedContext"))
+          .def_pickle(
+              // __getstate__
+              [](const c10::intrusive_ptr<LayernormPackedContext>& context) {
+                // context is packed
+                return context->unpack();
+              },
+              // __setstate__
+              [](c10::impl::GenericList state) {
+                // state is unpacked
+                return c10::make_intrusive<LayernormPackedContext>(
+                    LayernormPackedContext::pack(state));
+              });
+  return 0;
+}
+
 namespace {
 
 TORCH_LIBRARY(vulkan, m) {
@@ -99,6 +119,7 @@ TORCH_LIBRARY(vulkan, m) {
           });
   register_vulkan_conv2d_packed_context();
   register_vulkan_linear_packed_context();
+  register_vulkan_layernorm_packed_context();
   // To maintain backwards compatibility.
   m.class_<Conv2dOpContext>("Conv2dOpContext")
       .def_pickle(
@@ -162,6 +183,12 @@ TORCH_LIBRARY(vulkan_prepack, m) {
   m.def(TORCH_SELECTIVE_SCHEMA(
       "vulkan_prepack::run_qlinear_context(Tensor X, float scale, int zero_point, "
       "__torch__.torch.classes.vulkan.LinearPackedContext vk_context) -> Tensor Y"));
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "vulkan_prepack::create_layernorm_context(Tensor? W, Tensor? B, float eps) "
+      "-> __torch__.torch.classes.vulkan.LayernormPackedContext"));
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "vulkan_prepack::run_layernorm_context(Tensor X, SymInt[] normalized_shape, "
+      "__torch__.torch.classes.vulkan.LayernormPackedContext BW_prepack) -> Tensor Y"));
   m.def(TORCH_SELECTIVE_SCHEMA(
       "vulkan_prepack::create_gru_context(Tensor[] params_cpu, "
       "bool has_biases, "
@@ -220,6 +247,9 @@ TORCH_LIBRARY_IMPL(vulkan_prepack, CPU, m) {
   m.impl(
       TORCH_SELECTIVE_NAME("vulkan_prepack::create_linear_context"),
       TORCH_FN(create_linear_context));
+  m.impl(
+      TORCH_SELECTIVE_NAME("vulkan_prepack::create_layernorm_context"),
+      TORCH_FN(create_layernorm_context));
   m.impl(
       TORCH_SELECTIVE_NAME("vulkan_prepack::create_gru_context"),
       TORCH_FN(create_gru_context));
@@ -253,6 +283,9 @@ TORCH_LIBRARY_IMPL(vulkan_prepack, Vulkan, m) {
   m.impl(
       TORCH_SELECTIVE_NAME("vulkan_prepack::run_linear_context"),
       TORCH_FN(run_linear_context));
+  m.impl(
+      TORCH_SELECTIVE_NAME("vulkan_prepack::run_layernorm_context"),
+      TORCH_FN(run_layernorm_context));
   m.impl(
       TORCH_SELECTIVE_NAME("vulkan_prepack::run_qlinear_context"),
       TORCH_FN(run_qlinear_context));
diff --git a/aten/src/ATen/native/vulkan/ops/Register.h b/aten/src/ATen/native/vulkan/ops/Register.h
@@ -7,6 +7,7 @@ namespace ops {
 
 int register_vulkan_conv2d_packed_context();
 int register_vulkan_linear_packed_context();
+int register_vulkan_layernorm_packed_context();
 
 } // namespace ops
 } // namespace vulkan
diff --git a/aten/src/ATen/test/vulkan_api_test.cpp b/aten/src/ATen/test/vulkan_api_test.cpp
@@ -2709,6 +2709,66 @@ TEST_F(VulkanAPITest, hardtanh_) {
   ASSERT_TRUE(check);
 }
 
+void test_packed_layer_norm(
+    const at::IntArrayRef input_shape,
+    const at::IntArrayRef normalized_shape,
+    const at::IntArrayRef weight_shape,
+    const at::IntArrayRef bias_shape,
+    const float eps) {
+  c10::InferenceMode mode;
+
+  const auto input_cpu =
+      at::rand(input_shape, at::device(at::kCPU).dtype(at::kFloat));
+  const auto input_vulkan = input_cpu.vulkan();
+
+  const auto weight_cpu =
+      at::rand(weight_shape, at::device(at::kCPU).dtype(at::kFloat));
+
+  const auto bias_cpu =
+      at::rand(bias_shape, at::device(at::kCPU).dtype(at::kFloat));
+
+  const auto output_cpu = at::layer_norm(
+      input_cpu, normalized_shape, weight_cpu, bias_cpu, eps, false);
+
+  auto prepack = callOpByName(
+      "vulkan_prepack::create_layernorm_context",
+      "",
+      weight_cpu, bias_cpu, eps);
+
+  auto vulkan_output = callOpByName(
+      "vulkan_prepack::run_layernorm_context",
+      "",
+      input_cpu.vulkan(), normalized_shape, prepack[0]);
+
+  auto output_vulkan = vulkan_output[0].toTensor();
+
+  const auto check = almostEqual(output_cpu, output_vulkan.cpu());
+  if (!check) {
+    showRtol(output_cpu, output_vulkan.cpu());
+  }
+
+  ASSERT_TRUE(check);
+}
+
+TEST_F(VulkanAPITest, packed_layer_norm_2d) {
+  test_packed_layer_norm({5, 7}, {7}, {7}, {7}, 1e-05);
+  test_packed_layer_norm({5, 7}, {5, 7}, {5, 7}, {5, 7}, 1e-05);
+}
+
+TEST_F(VulkanAPITest, packed_layer_norm_3d) {
+  test_packed_layer_norm({11, 5, 7}, {7}, {7}, {7}, 1e-05);
+  test_packed_layer_norm({11, 5, 7}, {5, 7}, {5, 7}, {5, 7}, 1e-05);
+  test_packed_layer_norm({11, 5, 7}, {11, 5, 7}, {11, 5, 7}, {11, 5, 7}, 1e-05);
+}
+
+TEST_F(VulkanAPITest, packed_layer_norm_4d) {
+  test_packed_layer_norm({3, 11, 5, 7}, {7}, {7}, {7}, 1e-05);
+  test_packed_layer_norm({3, 11, 5, 7}, {5, 7}, {5, 7}, {5, 7}, 1e-05);
+  test_packed_layer_norm({3, 11, 5, 7}, {11, 5, 7}, {11, 5, 7}, {11, 5, 7}, 1e-05);
+  test_packed_layer_norm(
+      {3, 11, 5, 7}, {3, 11, 5, 7}, {3, 11, 5, 7}, {3, 11, 5, 7}, 1e-05);
+}
+
 TEST_F(VulkanAPITest, layer_norm_invalid_inputs) {
   c10::InferenceMode mode;
 
diff --git a/torch/csrc/jit/passes/vulkan_rewrite.cpp b/torch/csrc/jit/passes/vulkan_rewrite.cpp
@@ -57,6 +57,24 @@ void insertPrePackedLinearOp(std::shared_ptr<Graph>& graph) {
   linear_rewriter.runOnGraph(graph);
 }
 
+void insertPrePackedLayernormOp(std::shared_ptr<Graph>& graph) {
+  std::string layernorm_pattern = R"(
+    graph(%input, %normalized_shape, %weight, %bias, %eps, %cudnn_enable):
+        %r = aten::layer_norm(%input, %normalized_shape, %weight, %bias, %eps, %cudnn_enable)
+        return (%r))";
+  std::string prepacked_ops_pattern = R"(
+    graph(%input, %normalized_shape, %weight, %bias, %eps, %cudnn_enable):
+        %op_context : __torch__.torch.classes.vulkan.LayernormPackedContext = vulkan_prepack::create_layernorm_context(
+            %weight, %bias, %eps)
+        %res = vulkan_prepack::run_layernorm_context(%input, %normalized_shape, %op_context)
+        return (%res))";
+
+  SubgraphRewriter layernorm_rewriter;
+  layernorm_rewriter.RegisterRewritePattern(
+      layernorm_pattern, prepacked_ops_pattern);
+  layernorm_rewriter.runOnGraph(graph);
+}
+
 void insertPrePackedConv2dOp(std::shared_ptr<Graph>& graph) {
   graph_rewrite_helper::replaceConvolutionWithAtenConv(graph);
 
@@ -369,6 +387,7 @@ void fuseReluWithPackedOps(std::shared_ptr<Graph>& graph) {
 
 void vulkanInsertPrePackedOps(std::shared_ptr<Graph>& graph) {
   insertPrePackedLinearOp(graph);
+  insertPrePackedLayernormOp(graph);
   insertPrePackedConv2dOp(graph);
   rewriteQuantizedOps(graph);
   insertPrePackedGruOp(graph);
@@ -409,6 +428,8 @@ void vulkanFoldPrePackingOps(script::Module& m) {
              "vulkan_quantized_prepack::convert_linear_context")) ||
         (n->kind() ==
          Symbol::fromQualString("vulkan_prepack::create_linear_context")) ||
+        (n->kind() ==
+         Symbol::fromQualString("vulkan_prepack::create_layernorm_context")) ||
         (n->kind() ==
          Symbol::fromQualString("vulkan_prepack::create_gru_context")) ||
         (n->kind() ==
