From 4601d15365d77739ce3441683ba54980bd0ce2c6 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 14 Jul 2026 14:50:00 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../runtime/graph/ops/glsl/fused_ce.glsl | 110 ++++++++++ .../runtime/graph/ops/glsl/fused_ce.yaml | 15 ++ .../runtime/graph/ops/glsl/fused_ce_sum.glsl | 55 +++++ .../runtime/graph/ops/glsl/fused_ce_sum.yaml | 15 ++ .../vulkan/runtime/graph/ops/impl/FusedCe.cpp | 193 ++++++++++++++++++ 5 files changed, 388 insertions(+) create mode 100644 backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml create mode 100644 backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl create mode 100644 backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml create mode 100644 backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl new file mode 100644 index 00000000000..044301d5330 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl @@ -0,0 +1,110 @@ +/* + * 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_required_extensions(STORAGE, DTYPE)} + +#define PRECISION ${PRECISION} +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type(STORAGE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_dlogits", DTYPE, "buffer")} +${layout_declare_tensor(B, "w", "t_loss_partial", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_logits", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_labels", "int", "buffer")} + +${layout_declare_ubo(B, "int", "vocab")} +${layout_declare_ubo(B, "int", "n_rows")} +${layout_declare_ubo(B, "float", "n_valid")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#define NWORKERS 64 + +shared float red_m[NWORKERS]; +shared float red_l[NWORKERS]; + +// Fused cross-entropy: one workgroup per row cooperatively reduces the vocab +// dimension with a single-pass online softmax (running max + rescaled running +// sum), then writes per-row loss and dlogits. Rows with label < 0 are masked. +void main() { + const uint row = gl_GlobalInvocationID.y; + if (int(row) >= n_rows) { + return; + } + + const uint tid = gl_LocalInvocationID.x; + const uint V = uint(vocab); + const uint base = row * V; + const int lbl = t_labels[row]; + + if (lbl < 0) { + for (uint j = tid; j < V; j += NWORKERS) { + t_dlogits[base + j] = T(0); + } + if (tid == 0u) { + t_loss_partial[row] = T(0); + } + return; + } + + // Single read pass: maintain a running max m and the sum l of exp(x - m), + // rescaling l whenever a larger value updates m. Finite -3.4e38 init. + float m = -3.4e38; + float l = 0.0; + for (uint j = tid; j < V; j += NWORKERS) { + float x = float(t_logits[base + j]); + if (x > m) { + l = l * exp(m - x) + 1.0; + m = x; + } else { + l = l + exp(x - m); + } + } + red_m[tid] = m; + red_l[tid] = l; + memoryBarrierShared(); + barrier(); + + // Tree-combine the (m, l) pairs: m becomes the max, l is rescaled to it. + for (uint s = NWORKERS / 2u; s > 0u; s >>= 1u) { + if (tid < s) { + float ma = red_m[tid]; + float la = red_l[tid]; + float mb = red_m[tid + s]; + float lb = red_l[tid + s]; + float mm = max(ma, mb); + red_m[tid] = mm; + red_l[tid] = la * exp(ma - mm) + lb * exp(mb - mm); + } + memoryBarrierShared(); + barrier(); + } + + const float row_max = red_m[0]; + const float denom = red_l[0]; + const float inv = 1.0 / denom; + const float scale = 1.0 / n_valid; + + if (tid == 0u) { + const float lse = row_max + log(denom); + t_loss_partial[row] = T((lse - float(t_logits[base + uint(lbl)])) * scale); + } + + for (uint j = tid; j < V; j += NWORKERS) { + float g = exp(float(t_logits[base + j]) - row_max) * inv * scale; + if (j == uint(lbl)) { + g = g - scale; + } + t_dlogits[base + j] = T(g); + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml new file mode 100644 index 00000000000..8cf9cdac04e --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml @@ -0,0 +1,15 @@ +# 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. + +fused_ce: + parameter_names_with_default_values: + DTYPE: float + STORAGE: buffer + generate_variant_forall: + DTYPE: + - VALUE: float + shader_variants: + - NAME: fused_ce_buffer diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl new file mode 100644 index 00000000000..c61820e074c --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl @@ -0,0 +1,55 @@ +/* + * 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_required_extensions(STORAGE, DTYPE)} + +#define PRECISION ${PRECISION} +#define T ${buffer_scalar_type(DTYPE)} + +${define_active_storage_type(STORAGE)} + +layout(std430) buffer; + +${layout_declare_tensor(B, "w", "t_loss", DTYPE, "buffer")} +${layout_declare_tensor(B, "r", "t_loss_partial", DTYPE, "buffer")} + +${layout_declare_ubo(B, "int", "n_rows")} + +layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; + +#define NWORKERS 64 + +shared float red[NWORKERS]; + +// Self-contained [N] -> [1] tree-sum of the per-row losses in one workgroup, so +// fused_ce carries no cross-op reduce dependency. +void main() { + const uint tid = gl_LocalInvocationID.x; + + float s = 0.0; + for (uint j = tid; j < uint(n_rows); j += NWORKERS) { + s += float(t_loss_partial[j]); + } + red[tid] = s; + memoryBarrierShared(); + barrier(); + + for (uint k = NWORKERS / 2u; k > 0u; k >>= 1u) { + if (tid < k) { + red[tid] += red[tid + k]; + } + memoryBarrierShared(); + barrier(); + } + + if (tid == 0u) { + t_loss[0] = T(red[0]); + } +} diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml new file mode 100644 index 00000000000..cc1bd16cebb --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml @@ -0,0 +1,15 @@ +# 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. + +fused_ce_sum: + parameter_names_with_default_values: + DTYPE: float + STORAGE: buffer + generate_variant_forall: + DTYPE: + - VALUE: float + shader_variants: + - NAME: fused_ce_sum_buffer diff --git a/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp b/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp new file mode 100644 index 00000000000..75bdf7c54c5 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp @@ -0,0 +1,193 @@ +/* + * 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 +#include + +#include +#include + +namespace vkcompute { + +using namespace utils; + +utils::uvec3 fused_ce_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 loss_partial = args.at(0).refs.at(1); + return { + 1u, utils::safe_downcast(graph->numel_of(loss_partial)), 1u}; +} + +utils::uvec3 fused_ce_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}; +} + +utils::uvec3 fused_ce_sum_global_wg_size( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)graph; + (void)shader; + (void)args; + (void)resize_args; + return {1u, 1u, 1u}; +} + +utils::uvec3 fused_ce_sum_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 resize_fused_ce_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + (void)resize_args; + const ValueRef dlogits = args.at(0).refs.at(0); + const ValueRef loss_partial = args.at(0).refs.at(1); + const ValueRef logits = args.at(1).refs.at(0); + + const std::vector logits_sizes = graph->sizes_of(logits); + graph->virtual_resize(dlogits, logits_sizes); + graph->virtual_resize(loss_partial, {logits_sizes.at(0)}); +} + +void resize_fused_ce_sum_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + (void)resize_args; + const ValueRef loss = args.at(0).refs.at(0); + // Rank-0 scalar, matching fused_ce_meta's logits.new_empty([]). + graph->virtual_resize(loss, {}); +} + +void fused_ce(ComputeGraph& graph, const std::vector& args) { + int arg_idx = 0; + const ValueRef logits = args[arg_idx++]; + const ValueRef labels = args[arg_idx++]; + const ValueRef n_valid_ref = args[arg_idx++]; + const ValueListPtr out_tuple = graph.get_value_list(args[arg_idx++]); + const ValueRef loss = out_tuple->at(0); + const ValueRef dlogits = out_tuple->at(1); + + VK_CHECK_COND( + graph.is_buffer_storage(logits) && graph.is_buffer_storage(dlogits), + "fused_ce: logits and dlogits must use buffer storage"); + VK_CHECK_COND( + graph.dim_of(logits) == 2, "fused_ce: logits must be 2D [n_rows, vocab]"); + VK_CHECK_COND( + graph.sizes_of(dlogits) == graph.sizes_of(logits), + "fused_ce: dlogits must match logits shape"); + VK_CHECK_COND( + graph.dtype_of(labels) == vkapi::kInt, "fused_ce: labels must be int32"); + VK_CHECK_COND(graph.dim_of(labels) == 1, "fused_ce: labels must be 1D [N]"); + VK_CHECK_COND( + graph.size_at(0, labels) == graph.size_at(0, logits), + "fused_ce: labels length must equal number of rows"); + VK_CHECK_COND(graph.numel_of(loss) == 1, "fused_ce: loss must be a scalar"); + + const int32_t n_rows = graph.size_at(0, logits); + const int32_t vocab = graph.size_at(1, logits); + const float n_valid = graph.extract_scalar(n_valid_ref); + + // One workgroup per row; the Vulkan spec guarantees maxComputeWorkGroupCount + // >= 65535 per dimension. + VK_CHECK_COND( + n_rows <= 65535, "fused_ce: n_rows exceeds max workgroup count"); + + // Per-row loss contributions, reduced to the scalar loss by the sum node. + TmpTensor loss_partial( + &graph, {n_rows}, graph.dtype_of(logits), utils::kBuffer); + + std::string kernel_name = "fused_ce"; + kernel_name.reserve(kShaderNameReserve); + add_storage_type_suffix(kernel_name, graph.storage_type_of(dlogits)); + add_dtype_suffix(kernel_name, graph.dtype_of(dlogits)); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + fused_ce_global_wg_size, + fused_ce_local_wg_size, + // Inputs and Outputs + {{{dlogits, loss_partial}, vkapi::kWrite}, + {{logits, labels}, vkapi::kRead}}, + // Shader params buffers + {graph.create_params_buffer(vocab), + graph.create_params_buffer(n_rows), + graph.create_params_buffer(n_valid)}, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + {}, + // Resizing Logic + resize_fused_ce_node)); + + std::string sum_kernel_name = "fused_ce_sum"; + sum_kernel_name.reserve(kShaderNameReserve); + add_storage_type_suffix(sum_kernel_name, graph.storage_type_of(loss)); + add_dtype_suffix(sum_kernel_name, graph.dtype_of(loss)); + + // Separate dispatch node so the loss_partial write (above) is ordered before + // this read: the graph inserts a per-tensor pipeline barrier between them. + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(sum_kernel_name), + fused_ce_sum_global_wg_size, + fused_ce_sum_local_wg_size, + // Inputs and Outputs + {{{loss}, vkapi::kWrite}, {{loss_partial}, vkapi::kRead}}, + // Shader params buffers + {graph.create_params_buffer(n_rows)}, + // Push Constants + {}, + // Specialization Constants + {}, + // Resize Args + {}, + // Resizing Logic + resize_fused_ce_sum_node)); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(et_vk.fused_ce.default, fused_ce); +} + +} // namespace vkcompute