[ExecuTorch][Vulkan] Add et_vk.q4gsw_requant kernel (STE re-quant to W_4X8)

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

@@ -0,0 +1,97 @@
+/*
+ * 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_required_extensions(STORAGE, DTYPE)}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "w", "t_packed", "int", "buffer", is_scalar_array=False, vec_size=4)}
+${layout_declare_tensor(B, "r", "t_latent", DTYPE, STORAGE, is_scalar_array=True)}
+${layout_declare_tensor(B, "r", "t_scales", DTYPE, "buffer", is_scalar_array=True)}
+
+${layout_declare_ubo(B, "ivec4", "latent_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")}
+
+// STE re-quant of fp32 latent [N, K] into the W_4X8 block-packed 4-bit codes
+// the forward reads (mirrors pack_q4_linear_weight__w_4x8.glsl). Each thread
+// writes one 4K x 8N block (an ivec4) at (k4, n8).
+uint quant_nibble(const int n, const int k, const int N, const int K) {
+  const float s = float(t_scales[(k / group_size) * N + n]);
+  // roundEven + clamp-in-float match torch.round (half-to-even) then clamp.
+  float qf = 0.0;
+  if (s != 0.0) {
+    qf = clamp(roundEven(float(t_latent[n * K + k]) / s), -8.0, 7.0);
+  }
+  return uint(int(qf) + 8) & 0xFu;
+}
+
+// Pack a 4K x 4N tile into the (packed_x, packed_y) int pair. Byte b holds one
+// (even-N low nibble, odd-N high nibble) pair at K = k4*4 + b; rows N0,N1 go to
+// packed_x, rows N2,N3 to packed_y.
+void pack_tile(
+    out uint packed_x,
+    out uint packed_y,
+    const int k4,
+    const int n4,
+    const int N,
+    const int K) {
+  packed_x = 0u;
+  packed_y = 0u;
+  for (int ni = 0; ni < 4; ++ni) {
+    const int n = n4 * 4 + ni;
+    for (int b = 0; b < 4; ++b) {
+      const uint code = quant_nibble(n, k4 * 4 + b, N, K);
+      const int shift = 8 * b + (ni & 1) * 4;
+      if (ni < 2) {
+        packed_x |= code << shift;
+      } else {
+        packed_y |= code << shift;
+      }
+    }
+  }
+}
+
+void main() {
+  const int k4 = int(gl_GlobalInvocationID.x);
+  const int n8 = int(gl_GlobalInvocationID.y);
+
+  const int K = latent_sizes.x;
+  const int N = latent_sizes.y;
+  const int K4 = K >> 2;
+  const int N4 = (N + 3) >> 2;
+  const int N8 = (N4 + 1) >> 1;
+
+  if (k4 >= K4 || n8 >= N8) {
+    return;
+  }
+
+  const int n4_a = 2 * n8;
+  const int n4_b = n4_a + 1;
+
+  uint packed_x_a = 0u;
+  uint packed_y_a = 0u;
+  // OOB upper tile (odd N4 boundary) is the bias-zero pattern, matching the
+  // forward's prepack padding so the whole block is always readable.
+  uint packed_x_b = 0x88888888u;
+  uint packed_y_b = 0x88888888u;
+
+  pack_tile(packed_x_a, packed_y_a, k4, n4_a, N, K);
+  if (n4_b < N4) {
+    pack_tile(packed_x_b, packed_y_b, k4, n4_b, N, K);
+  }
+
+  t_packed[k4 * N8 + n8] = ivec4(
+      int(packed_x_a), int(packed_y_a), int(packed_x_b), int(packed_y_b));
+}

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

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

@@ -0,0 +1,131 @@
+/*
+ * 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/Staging.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+namespace vkcompute {
+
+// Resize the packed output to the W_4X8 int buffer size K4 * N4_padded * 2,
+// matching prepack_q4_w_4x8_nc_buffer's layout. extra_args = { latent }.
+void resize_q4gsw_requant_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& extra_args) {
+  const ValueRef packed = args.at(0).refs.at(0);
+  const ValueRef latent = extra_args.at(0);
+  const std::vector<int64_t> latent_sizes = graph->sizes_of(latent);
+  const int64_t N = latent_sizes.at(0);
+  const int64_t K = latent_sizes.at(1);
+  const int64_t K4 = K / 4;
+  const int64_t N4 = N / 4;
+  const int64_t N4_padded = (N4 + 1) & ~int64_t{1};
+  graph->virtual_resize(packed, {K4 * N4_padded * 2});
+}
+
+utils::uvec3 q4gsw_requant_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 latent = args.at(1).refs.at(0);
+  const std::vector<int64_t> latent_sizes = graph->sizes_of(latent);
+  const uint32_t N = utils::safe_downcast<uint32_t>(latent_sizes.at(0));
+  const uint32_t K = utils::safe_downcast<uint32_t>(latent_sizes.at(1));
+  const uint32_t K4 = K / 4u;
+  const uint32_t N4 = (N + 3u) / 4u;
+  const uint32_t N8 = (N4 + 1u) / 2u;
+  return {K4, N8, 1u};
+}
+
+utils::uvec3 q4gsw_requant_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 {8u, 8u, 1u};
+}
+
+void q4gsw_requant(ComputeGraph& graph, const std::vector<ValueRef>& args) {
+  int32_t i = 0;
+  const ValueRef latent = args.at(i++);
+  const ValueRef scales = args.at(i++);
+  const ValueRef group_size_ref = args.at(i++);
+  const ValueRef packed = args.at(i++);
+
+  VK_CHECK_COND(graph.dtype_of(latent) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(scales) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(packed) == vkapi::kInt);
+  VK_CHECK_COND(graph.is_buffer_storage(latent));
+  VK_CHECK_COND(graph.is_buffer_storage(packed));
+
+  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> latent_sizes = graph.sizes_of(latent);
+  VK_CHECK_COND(latent_sizes.size() == 2);
+  const int64_t N = latent_sizes.at(0);
+  const int64_t K = latent_sizes.at(1);
+  VK_CHECK_COND(N > 0 && K > 0);
+  VK_CHECK_COND(N % 4 == 0 && K % 4 == 0);
+  VK_CHECK_COND(K % group_size_val == 0);
+
+  const uint32_t K4 = utils::safe_downcast<uint32_t>(K / 4);
+  const uint32_t N4 = utils::safe_downcast<uint32_t>(N / 4);
+  const uint32_t N8 = (N4 + 1u) / 2u;
+  VK_CHECK_COND(
+      K4 <= 65535u && N8 <= 65535u,
+      "q4gsw_requant: dispatch grid exceeds max workgroup count");
+
+  // Scales are a frozen constant; materialize them to a GPU buffer once.
+  const ValueRef packed_scales =
+      prepack_standard(graph, scales, utils::kBuffer, utils::kWidthPacked);
+
+  std::string kernel_name = "q4gsw_requant";
+  kernel_name.reserve(kShaderNameReserve);
+  add_storage_type_suffix(kernel_name, graph.storage_type_of(latent));
+  add_dtype_suffix(kernel_name, graph.dtype_of(latent));
+
+  graph.execute_nodes().emplace_back(new DynamicDispatchNode(
+      graph,
+      VK_KERNEL_FROM_STR(kernel_name),
+      q4gsw_requant_global_wg_size,
+      q4gsw_requant_local_wg_size,
+      // Inputs and Outputs
+      {{packed, vkapi::kWrite}, {{latent, packed_scales}, vkapi::kRead}},
+      // Shader params buffers
+      {graph.sizes_ubo(latent)},
+      // Push Constants
+      {},
+      // Specialization Constants
+      {static_cast<uint32_t>(group_size_val)},
+      // Resize Args
+      {latent},
+      // Resizing Logic
+      resize_q4gsw_requant_node));
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(et_vk.q4gsw_requant.default, q4gsw_requant);
+}
+
+} // namespace vkcompute