From 2f263f328c16bbb6aba541ec117493f005465729 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 14:49:50 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../graph/ops/glsl/linear_q4gsw_dw.glsl | 78 ++++++++++++ .../graph/ops/glsl/linear_q4gsw_dw.yaml | 17 +++ .../graph/ops/impl/QuantizedLinearDw.cpp | 112 ++++++++++++++++++ 3 files changed, 207 insertions(+) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_dw.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_dw.yaml create mode 100644 backends/vulkan/runtime/graph/ops/impl/QuantizedLinearDw.cpp diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_dw.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_dw.glsl new file mode 100644 index 00000000000..1a8ed94031b --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_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]); + } + } + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_dw.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_dw.yaml new file mode 100644 index 00000000000..c6103b307fb --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_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_q4gsw_dw: + parameter_names_with_default_values: + DTYPE: float + STORAGE: buffer + generate_variant_forall: + STORAGE: + - VALUE: buffer + DTYPE: + - VALUE: float + shader_variants: + - NAME: linear_q4gsw_dw diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearDw.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearDw.cpp new file mode 100644 index 00000000000..a779a445cd6 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearDw.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 + +#include + +#include + +namespace vkcompute { + +void resize_linear_q4gsw_dw_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& 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(-1, d_out); + const int64_t K = graph->size_at(-1, x); + graph->virtual_resize(dW, {N, K}); +} + +utils::uvec3 linear_q4gsw_dw_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)resize_args; + const ValueRef dW = args.at(0).refs.at(0); + const uint32_t N = graph->size_at(-2, dW); + const uint32_t K = graph->size_at(-1, dW); + const uint32_t tiles = utils::div_up_4(N) * utils::div_up_4(K); + return {tiles, 1u, 1u}; +} + +utils::uvec3 linear_q4gsw_dw_local_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const utils::uvec3& global_workgroup_size, + const std::vector& args, + const std::vector& resize_args) { + (void)graph; + (void)shader; + (void)global_workgroup_size; + (void)args; + (void)resize_args; + return {64u, 1u, 1u}; +} + +void linear_q4gsw_dw(ComputeGraph& graph, const std::vector& 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(-1, d_out); + const int64_t K = graph.size_at(-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(N)) * + utils::div_up_4(static_cast(K)); + VK_CHECK_COND( + (tiles + 63u) / 64u <= 65535u, + "linear_q4gsw_dw: tile count exceeds max workgroup count"); + + std::string kernel_name = "linear_q4gsw_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_q4gsw_dw_global_wg_size, + linear_q4gsw_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_q4gsw_dw_node)); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.linear_q4gsw_dw.default, linear_q4gsw_dw); +} + +} // namespace vkcompute