[ExecuTorch][Vulkan] Add et_vk.linear_q4gsw_backward kernel (4-bit input-grad)

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

@@ -0,0 +1,104 @@
+/*
+ * 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_M 4
+#define TILE_K 4
+
+${define_required_extensions(STORAGE, DTYPE)}
+${define_required_extensions("buffer", DTYPE)}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "w", "t_dx", 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_q4_weights", "int", "buffer", is_scalar_array=False, vec_size=4)}
+${layout_declare_tensor(B, "r", "t_scales", DTYPE, "buffer", is_scalar_array=True)}
+
+${layout_declare_ubo(B, "ivec4", "dout_sizes")}
+${layout_declare_ubo(B, "ivec4", "dx_sizes")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+${layout_declare_spec_const(C, "int", "group_size", "32")}
+
+// d_x[M, K] = d_out[M, N] @ dequant(W)[N, K], contracting over N.
+// dequant(W[n, k]) = (code - 8) * scale, with code read from the same W_4X8
+// block-packed weight the forward reads (mirrors q4gsw_linear_gemm__w_4x8.glsl).
+void main() {
+  const int N = dout_sizes.x;
+  const int M = dout_sizes.y * dout_sizes.z * dout_sizes.w;
+  const int K = dx_sizes.x;
+
+  const int nmt = (M + TILE_M - 1) / TILE_M;
+  const int nkt = (K + TILE_K - 1) / TILE_K;
+  const int tiles = nmt * nkt;
+
+  const int tile_idx = int(gl_GlobalInvocationID.x);
+  if (tile_idx >= tiles) {
+    return;
+  }
+
+  const int m0 = (tile_idx / nkt) * TILE_M;
+  const int k0 = (tile_idx % nkt) * TILE_K;
+
+  // K and N are multiples of 4 (prepack guarantees), so k0 is 4-aligned: the
+  // tile's 4 K lanes are byte b = kl of one k4 group and share one scale group.
+  const int k4 = k0 >> 2;
+  const int N4 = (N + 3) >> 2;
+  const int N4_padded = (N4 + 1) & ~1;
+  const int N8 = N4_padded >> 1;
+  const int group = k0 / group_size;
+
+  float acc[TILE_M * TILE_K];
+  for (int i = 0; i < TILE_M * TILE_K; ++i) {
+    acc[i] = 0.0;
+  }
+
+  for (int n = 0; n < N; ++n) {
+    float dout_reg[TILE_M];
+    for (int ml = 0; ml < TILE_M; ++ml) {
+      const int m_eff = min(m0 + ml, M - 1);
+      dout_reg[ml] = float(t_dout[m_eff * N + n]);
+    }
+
+    // W_4X8 address for column n: ivec4 at (k4, n8); component by (n4 parity,
+    // n-in-tile half); low/high nibble by n parity (even-N low, odd-N high).
+    const int n4 = n >> 2;
+    const int ni = n & 3;
+    const int n8 = n4 >> 1;
+    const int comp = (n4 & 1) * 2 + (ni >> 1);
+    const int nib_hi = (ni & 1) * 4;
+    const ivec4 w_block = t_q4_weights[k4 * N8 + n8];
+    const int w_int = w_block[comp];
+    const float scale = float(t_scales[group * N + n]);
+
+    for (int kl = 0; kl < TILE_K; ++kl) {
+      const int code = int((uint(w_int) >> (8 * kl + nib_hi)) & 0xFu);
+      const float dq = float(code - 8) * scale;
+      for (int ml = 0; ml < TILE_M; ++ml) {
+        acc[ml * TILE_K + kl] += dout_reg[ml] * dq;
+      }
+    }
+  }
+
+  for (int ml = 0; ml < TILE_M; ++ml) {
+    const int m = m0 + ml;
+    for (int kl = 0; kl < TILE_K; ++kl) {
+      const int k = k0 + kl;
+      if (m < M && k < K) {
+        t_dx[m * K + k] = T(acc[ml * TILE_K + kl]);
+      }
+    }
+  }
+}

backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.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.
+
+q4gsw_backward:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+  generate_variant_forall:
+    STORAGE:
+      - VALUE: buffer
+    DTYPE:
+      - VALUE: float
+  shader_variants:
+    - NAME: q4gsw_backward

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

@@ -0,0 +1,130 @@
+/*
+ * 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/DynamicDispatchNode.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+namespace vkcompute {
+
+// Resize d_x to d_out.shape[:-1] + (K,), mirroring linear_q4gsw_backward_meta.
+// extra_args = { weight_data, d_out }.
+void resize_linear_q4gsw_backward_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& extra_args) {
+  const ValueRef d_x = args.at(0).refs.at(0);
+  const ValueRef weight_data = extra_args.at(0);
+  const ValueRef d_out = extra_args.at(1);
+  const int64_t K = graph->sizes_of(weight_data).at(1) * 2;
+  std::vector<int64_t> new_sizes = graph->sizes_of(d_out);
+  new_sizes.back() = K;
+  graph->virtual_resize(d_x, new_sizes);
+}
+
+utils::uvec3 linear_q4gsw_backward_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 d_x = args.at(0).refs.at(0);
+  const uint32_t K = graph->size_at<uint32_t>(-1, d_x);
+  const uint32_t M = utils::safe_downcast<uint32_t>(graph->numel_of(d_x) / K);
+  const uint32_t tiles = utils::div_up_4(M) * utils::div_up_4(K);
+  return {tiles, 1u, 1u};
+}
+
+utils::uvec3 linear_q4gsw_backward_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_q4gsw_backward(
+    ComputeGraph& graph,
+    const std::vector<ValueRef>& args) {
+  int32_t i = 0;
+  const ValueRef d_out = args.at(i++);
+  const ValueRef weight_data = args.at(i++);
+  const ValueRef weight_scales_data = args.at(i++);
+  const ValueRef group_size_ref = args.at(i++);
+  const ValueRef d_x = args.at(i++);
+
+  VK_CHECK_COND(graph.dtype_of(d_out) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(d_x) == vkapi::kFloat);
+  VK_CHECK_COND(graph.is_buffer_storage(d_out));
+  VK_CHECK_COND(graph.is_buffer_storage(d_x));
+
+  const vkapi::ScalarType in_dtype = graph.dtype_of(d_out);
+  const int64_t group_size_val = graph.extract_scalar<int64_t>(group_size_ref);
+  VK_CHECK_COND(group_size_val > 0 && group_size_val % 4 == 0);
+
+  const std::vector<int64_t> weight_sizes = graph.sizes_of(weight_data);
+  const int64_t N = weight_sizes.at(0);
+  const int64_t K = weight_sizes.at(1) * 2;
+  VK_CHECK_COND(N > 0 && K > 0);
+  VK_CHECK_COND(N % 4 == 0 && K % 4 == 0);
+  VK_CHECK_COND(K % group_size_val == 0);
+  VK_CHECK_COND(graph.size_at<int64_t>(-1, d_out) == N);
+
+  const ValueRef packed_weight = prepack_q4_w_4x8_nc_buffer(graph, weight_data);
+  const ValueRef packed_scales =
+      prepack_q4_scales(graph, weight_scales_data, in_dtype);
+
+  const uint32_t M = utils::safe_downcast<uint32_t>(graph.numel_of(d_out) / N);
+  const uint32_t tiles =
+      utils::div_up_4(M) * utils::div_up_4(static_cast<uint32_t>(K));
+  VK_CHECK_COND(
+      (tiles + 63u) / 64u <= 65535u,
+      "linear_q4gsw_backward: tile count exceeds max workgroup count");
+
+  std::string kernel_name = "q4gsw_backward";
+  kernel_name.reserve(kShaderNameReserve);
+  add_storage_type_suffix(kernel_name, graph.storage_type_of(d_x));
+  add_dtype_suffix(kernel_name, graph.dtype_of(d_x));
+
+  graph.execute_nodes().emplace_back(new DynamicDispatchNode(
+      graph,
+      VK_KERNEL_FROM_STR(kernel_name),
+      linear_q4gsw_backward_global_wg_size,
+      linear_q4gsw_backward_local_wg_size,
+      // Inputs and Outputs
+      {{d_x, vkapi::kWrite},
+       {{d_out, packed_weight, packed_scales}, vkapi::kRead}},
+      // Shader params buffers
+      {graph.sizes_ubo(d_out), graph.sizes_ubo(d_x)},
+      // Push Constants
+      {},
+      // Specialization Constants
+      {static_cast<uint32_t>(group_size_val)},
+      // Resize Args
+      {weight_data, d_out},
+      // Resizing Logic
+      resize_linear_q4gsw_backward_node));
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(et_vk.linear_q4gsw_backward.default, linear_q4gsw_backward);
+}
+
+} // namespace vkcompute