[ExecuTorch][Vulkan] Add et_vk.fused_ce kernels (fused cross-entropy loss + dlogits)

JCNTH · JCNTH · commit 1a2c4da709dc · 2026-07-21T09:48:09.000-07:00
Pull Request resolved: #20941 Vulkan GLSL kernels + handler for the training custom op `et_vk.fused_ce` (fused CE: per-row online-softmax loss + dlogits, then an N->1 loss sum). Per-row shared-memory reduction structured like `glsl/reduce_per_row_buffer.glsl` (`NWORKERS=64`, one workgroup/row, barrier tree-combine); two dispatch nodes share `loss_partial`, ordered by the runtime per-`vTensor` barrier (same pattern as `impl/SDPA.cpp` QK->softmax->AV). AOT registration already exists under the shared Vulkan partitioner. ghstack-source-id: 405059103 @exported-using-ghexport Differential Revision: [D111761780](https://our.internmc.facebook.com/intern/diff/D111761780/)
diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl 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
@@ -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
@@ -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
@@ -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
@@ -0,0 +1,197 @@
+/*
+ * 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/impl/utils/TensorUtils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+namespace vkcompute {
+
+using namespace utils;
+
+utils::uvec3 fused_ce_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 loss_partial = args.at(0).refs.at(1);
+  return {
+      1u, utils::safe_downcast<uint32_t>(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<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};
+}
+
+utils::uvec3 fused_ce_sum_global_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& 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<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 resize_fused_ce_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& 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<int64_t> 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<ArgGroup>& args,
+    const std::vector<ValueRef>& 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<ValueRef>& 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 ValueRef out_tuple_ref = args[arg_idx];
+
+  // Read loss/dlogits via short-lived get_value_list() temporaries: no get_*()
+  // pointer may be alive when loss_partial is added below, else adding Values
+  // can reallocate values_ and check_no_active_value_ptrs() throws.
+  const ValueRef loss = graph.get_value_list(out_tuple_ref)->at(0);
+  const ValueRef dlogits = graph.get_value_list(out_tuple_ref)->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<int64_t>(0, labels) == graph.size_at<int64_t>(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<int32_t>(0, logits);
+  const int32_t vocab = graph.size_at<int32_t>(1, logits);
+  const float n_valid = graph.extract_scalar<float>(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
