[ExecuTorch][Vulkan] Add et_vk.linear_dW kernel (dense fp32 weight-grad GEMM)

Differential Revision: D111761779

Pull Request resolved: #20939

Author: JCNTH

backends/vulkan/runtime/graph/ops/glsl/linear_dW.glsl

@@ -0,0 +1,78 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+#define PRECISION ${PRECISION}
+
+#define T ${texel_load_component_type(DTYPE, STORAGE)}
+
+#define TILE_N 4
+#define TILE_K 4
+
+${define_required_extensions(STORAGE, DTYPE)}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "w", "t_dw", DTYPE, STORAGE, is_scalar_array=True)}
+${layout_declare_tensor(B, "r", "t_dout", DTYPE, STORAGE, is_scalar_array=True)}
+${layout_declare_tensor(B, "r", "t_x", DTYPE, STORAGE, is_scalar_array=True)}
+
+${layout_declare_ubo(B, "ivec4", "dout_sizes")}
+${layout_declare_ubo(B, "ivec4", "x_sizes")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  // dW[N, K] = sum_m d_out[m, N] * x[m, K]; contraction over the flattened M.
+  const int N = dout_sizes.x;
+  const int M = dout_sizes.y * dout_sizes.z * dout_sizes.w;
+  const int K = x_sizes.x;
+
+  const int nnt = (N + TILE_N - 1) / TILE_N;
+  const int nkt = (K + TILE_K - 1) / TILE_K;
+  const int tiles = nnt * nkt;
+
+  const int tile_idx = int(gl_GlobalInvocationID.x);
+  if (tile_idx >= tiles) {
+    return;
+  }
+
+  const int n0 = (tile_idx / nkt) * TILE_N;
+  const int k0 = (tile_idx % nkt) * TILE_K;
+
+  float acc[TILE_N * TILE_K];
+  for (int i = 0; i < TILE_N * TILE_K; ++i) {
+    acc[i] = 0.0;
+  }
+
+  for (int m = 0; m < M; ++m) {
+    float dout_reg[TILE_N];
+    for (int nl = 0; nl < TILE_N; ++nl) {
+      const int n_eff = min(n0 + nl, N - 1);
+      dout_reg[nl] = float(t_dout[m * N + n_eff]);
+    }
+    for (int kl = 0; kl < TILE_K; ++kl) {
+      const int k_eff = min(k0 + kl, K - 1);
+      const float xv = float(t_x[m * K + k_eff]);
+      for (int nl = 0; nl < TILE_N; ++nl) {
+        acc[nl * TILE_K + kl] += dout_reg[nl] * xv;
+      }
+    }
+  }
+
+  for (int nl = 0; nl < TILE_N; ++nl) {
+    const int n = n0 + nl;
+    for (int kl = 0; kl < TILE_K; ++kl) {
+      const int k = k0 + kl;
+      if (n < N && k < K) {
+        t_dw[n * K + k] = T(acc[nl * TILE_K + kl]);
+      }
+    }
+  }
+}

backends/vulkan/runtime/graph/ops/glsl/linear_dW.yaml

@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+linear_dW:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+  generate_variant_forall:
+    STORAGE:
+      - VALUE: buffer
+    DTYPE:
+      - VALUE: float
+  shader_variants:
+    - NAME: linear_dW

backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp

@@ -0,0 +1,112 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+namespace vkcompute {
+
+void resize_linear_dW_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)resize_args;
+  const ValueRef dW = args.at(0).refs.at(0);
+  const ValueRef d_out = args.at(1).refs.at(0);
+  const ValueRef x = args.at(1).refs.at(1);
+  const int64_t N = graph->size_at<int64_t>(-1, d_out);
+  const int64_t K = graph->size_at<int64_t>(-1, x);
+  graph->virtual_resize(dW, {N, K});
+}
+
+utils::uvec3 linear_dW_global_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)shader;
+  (void)resize_args;
+  const ValueRef dW = args.at(0).refs.at(0);
+  const uint32_t N = graph->size_at<uint32_t>(-2, dW);
+  const uint32_t K = graph->size_at<uint32_t>(-1, dW);
+  const uint32_t tiles = utils::div_up_4(N) * utils::div_up_4(K);
+  return {tiles, 1u, 1u};
+}
+
+utils::uvec3 linear_dW_local_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const utils::uvec3& global_workgroup_size,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)graph;
+  (void)shader;
+  (void)global_workgroup_size;
+  (void)args;
+  (void)resize_args;
+  return {64u, 1u, 1u};
+}
+
+void linear_dW(ComputeGraph& graph, const std::vector<ValueRef>& args) {
+  int32_t i = 0;
+  const ValueRef d_out = args.at(i++);
+  const ValueRef x = args.at(i++);
+  const ValueRef dW = args.at(i);
+
+  VK_CHECK_COND(graph.dtype_of(d_out) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(x) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(dW) == vkapi::kFloat);
+  VK_CHECK_COND(graph.is_buffer_storage(d_out));
+  VK_CHECK_COND(graph.is_buffer_storage(x));
+  VK_CHECK_COND(graph.is_buffer_storage(dW));
+
+  const int64_t N = graph.size_at<int64_t>(-1, d_out);
+  const int64_t K = graph.size_at<int64_t>(-1, x);
+  VK_CHECK_COND(N > 0 && K > 0);
+  VK_CHECK_COND(graph.numel_of(d_out) % N == 0);
+  VK_CHECK_COND(graph.numel_of(x) % K == 0);
+  VK_CHECK_COND(graph.numel_of(d_out) / N == graph.numel_of(x) / K);
+
+  const uint32_t tiles = utils::div_up_4(static_cast<uint32_t>(N)) *
+      utils::div_up_4(static_cast<uint32_t>(K));
+  VK_CHECK_COND(
+      (tiles + 63u) / 64u <= 65535u,
+      "linear_dW: tile count exceeds max workgroup count");
+
+  std::string kernel_name = "linear_dW";
+  kernel_name.reserve(kShaderNameReserve);
+  add_storage_type_suffix(kernel_name, graph.storage_type_of(dW));
+  add_dtype_suffix(kernel_name, graph.dtype_of(dW));
+
+  graph.execute_nodes().emplace_back(new DynamicDispatchNode(
+      graph,
+      VK_KERNEL_FROM_STR(kernel_name),
+      linear_dW_global_wg_size,
+      linear_dW_local_wg_size,
+      // Inputs and Outputs
+      {{dW, vkapi::kWrite}, {{d_out, x}, vkapi::kRead}},
+      // Shader params buffers
+      {graph.sizes_ubo(d_out), graph.sizes_ubo(x)},
+      // Push Constants
+      {},
+      // Specialization Constants
+      {},
+      // Resize Args
+      {},
+      // Resizing Logic
+      resize_linear_dW_node));
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(et_vk.linear_dW.default, linear_dW);
+}
+
+} // namespace vkcompute