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

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);
+  }
+}

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

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]);
+  }
+}

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

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