Skip to content

Commit e55803a

Browse files
JCNTHJCNTH
committed
[ExecuTorch][Vulkan] Add et_vk.fused_ce kernels (fused cross-entropy loss + dlogits)
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: 404846443 @exported-using-ghexport Differential Revision: [D111761780](https://our.internmc.facebook.com/intern/diff/D111761780/)
1 parent 18e0e66 commit e55803a

5 files changed

Lines changed: 388 additions & 0 deletions

File tree

‎backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#version 450 core
10+
11+
${define_required_extensions(STORAGE, DTYPE)}
12+
13+
#define PRECISION ${PRECISION}
14+
#define T ${buffer_scalar_type(DTYPE)}
15+
16+
${define_active_storage_type(STORAGE)}
17+
18+
layout(std430) buffer;
19+
20+
${layout_declare_tensor(B, "w", "t_dlogits", DTYPE, "buffer")}
21+
${layout_declare_tensor(B, "w", "t_loss_partial", DTYPE, "buffer")}
22+
${layout_declare_tensor(B, "r", "t_logits", DTYPE, "buffer")}
23+
${layout_declare_tensor(B, "r", "t_labels", "int", "buffer")}
24+
25+
${layout_declare_ubo(B, "int", "vocab")}
26+
${layout_declare_ubo(B, "int", "n_rows")}
27+
${layout_declare_ubo(B, "float", "n_valid")}
28+
29+
layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
30+
31+
#define NWORKERS 64
32+
33+
shared float red_m[NWORKERS];
34+
shared float red_l[NWORKERS];
35+
36+
// Fused cross-entropy: one workgroup per row cooperatively reduces the vocab
37+
// dimension with a single-pass online softmax (running max + rescaled running
38+
// sum), then writes per-row loss and dlogits. Rows with label < 0 are masked.
39+
void main() {
40+
const uint row = gl_GlobalInvocationID.y;
41+
if (int(row) >= n_rows) {
42+
return;
43+
}
44+
45+
const uint tid = gl_LocalInvocationID.x;
46+
const uint V = uint(vocab);
47+
const uint base = row * V;
48+
const int lbl = t_labels[row];
49+
50+
if (lbl < 0) {
51+
for (uint j = tid; j < V; j += NWORKERS) {
52+
t_dlogits[base + j] = T(0);
53+
}
54+
if (tid == 0u) {
55+
t_loss_partial[row] = T(0);
56+
}
57+
return;
58+
}
59+
60+
// Single read pass: maintain a running max m and the sum l of exp(x - m),
61+
// rescaling l whenever a larger value updates m. Finite -3.4e38 init.
62+
float m = -3.4e38;
63+
float l = 0.0;
64+
for (uint j = tid; j < V; j += NWORKERS) {
65+
float x = float(t_logits[base + j]);
66+
if (x > m) {
67+
l = l * exp(m - x) + 1.0;
68+
m = x;
69+
} else {
70+
l = l + exp(x - m);
71+
}
72+
}
73+
red_m[tid] = m;
74+
red_l[tid] = l;
75+
memoryBarrierShared();
76+
barrier();
77+
78+
// Tree-combine the (m, l) pairs: m becomes the max, l is rescaled to it.
79+
for (uint s = NWORKERS / 2u; s > 0u; s >>= 1u) {
80+
if (tid < s) {
81+
float ma = red_m[tid];
82+
float la = red_l[tid];
83+
float mb = red_m[tid + s];
84+
float lb = red_l[tid + s];
85+
float mm = max(ma, mb);
86+
red_m[tid] = mm;
87+
red_l[tid] = la * exp(ma - mm) + lb * exp(mb - mm);
88+
}
89+
memoryBarrierShared();
90+
barrier();
91+
}
92+
93+
const float row_max = red_m[0];
94+
const float denom = red_l[0];
95+
const float inv = 1.0 / denom;
96+
const float scale = 1.0 / n_valid;
97+
98+
if (tid == 0u) {
99+
const float lse = row_max + log(denom);
100+
t_loss_partial[row] = T((lse - float(t_logits[base + uint(lbl)])) * scale);
101+
}
102+
103+
for (uint j = tid; j < V; j += NWORKERS) {
104+
float g = exp(float(t_logits[base + j]) - row_max) * inv * scale;
105+
if (j == uint(lbl)) {
106+
g = g - scale;
107+
}
108+
t_dlogits[base + j] = T(g);
109+
}
110+
}

‎backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
fused_ce:
8+
parameter_names_with_default_values:
9+
DTYPE: float
10+
STORAGE: buffer
11+
generate_variant_forall:
12+
DTYPE:
13+
- VALUE: float
14+
shader_variants:
15+
- NAME: fused_ce_buffer

‎backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#version 450 core
10+
11+
${define_required_extensions(STORAGE, DTYPE)}
12+
13+
#define PRECISION ${PRECISION}
14+
#define T ${buffer_scalar_type(DTYPE)}
15+
16+
${define_active_storage_type(STORAGE)}
17+
18+
layout(std430) buffer;
19+
20+
${layout_declare_tensor(B, "w", "t_loss", DTYPE, "buffer")}
21+
${layout_declare_tensor(B, "r", "t_loss_partial", DTYPE, "buffer")}
22+
23+
${layout_declare_ubo(B, "int", "n_rows")}
24+
25+
layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
26+
27+
#define NWORKERS 64
28+
29+
shared float red[NWORKERS];
30+
31+
// Self-contained [N] -> [1] tree-sum of the per-row losses in one workgroup, so
32+
// fused_ce carries no cross-op reduce dependency.
33+
void main() {
34+
const uint tid = gl_LocalInvocationID.x;
35+
36+
float s = 0.0;
37+
for (uint j = tid; j < uint(n_rows); j += NWORKERS) {
38+
s += float(t_loss_partial[j]);
39+
}
40+
red[tid] = s;
41+
memoryBarrierShared();
42+
barrier();
43+
44+
for (uint k = NWORKERS / 2u; k > 0u; k >>= 1u) {
45+
if (tid < k) {
46+
red[tid] += red[tid + k];
47+
}
48+
memoryBarrierShared();
49+
barrier();
50+
}
51+
52+
if (tid == 0u) {
53+
t_loss[0] = T(red[0]);
54+
}
55+
}

‎backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
fused_ce_sum:
8+
parameter_names_with_default_values:
9+
DTYPE: float
10+
STORAGE: buffer
11+
generate_variant_forall:
12+
DTYPE:
13+
- VALUE: float
14+
shader_variants:
15+
- NAME: fused_ce_sum_buffer

‎backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp‎

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
10+
11+
#include <executorch/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h>
12+
13+
#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
14+
#include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
15+
16+
#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
17+
#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
18+
19+
namespace vkcompute {
20+
21+
using namespace utils;
22+
23+
utils::uvec3 fused_ce_global_wg_size(
24+
ComputeGraph* graph,
25+
const vkapi::ShaderInfo& shader,
26+
const std::vector<ArgGroup>& args,
27+
const std::vector<ValueRef>& resize_args) {
28+
(void)shader;
29+
(void)resize_args;
30+
const ValueRef loss_partial = args.at(0).refs.at(1);
31+
return {
32+
1u, utils::safe_downcast<uint32_t>(graph->numel_of(loss_partial)), 1u};
33+
}
34+
35+
utils::uvec3 fused_ce_local_wg_size(
36+
ComputeGraph* graph,
37+
const vkapi::ShaderInfo& shader,
38+
const utils::uvec3& global_workgroup_size,
39+
const std::vector<ArgGroup>& args,
40+
const std::vector<ValueRef>& resize_args) {
41+
(void)graph;
42+
(void)shader;
43+
(void)global_workgroup_size;
44+
(void)args;
45+
(void)resize_args;
46+
return {64u, 1u, 1u};
47+
}
48+
49+
utils::uvec3 fused_ce_sum_global_wg_size(
50+
ComputeGraph* graph,
51+
const vkapi::ShaderInfo& shader,
52+
const std::vector<ArgGroup>& args,
53+
const std::vector<ValueRef>& resize_args) {
54+
(void)graph;
55+
(void)shader;
56+
(void)args;
57+
(void)resize_args;
58+
return {1u, 1u, 1u};
59+
}
60+
61+
utils::uvec3 fused_ce_sum_local_wg_size(
62+
ComputeGraph* graph,
63+
const vkapi::ShaderInfo& shader,
64+
const utils::uvec3& global_workgroup_size,
65+
const std::vector<ArgGroup>& args,
66+
const std::vector<ValueRef>& resize_args) {
67+
(void)graph;
68+
(void)shader;
69+
(void)global_workgroup_size;
70+
(void)args;
71+
(void)resize_args;
72+
return {64u, 1u, 1u};
73+
}
74+
75+
void resize_fused_ce_node(
76+
ComputeGraph* graph,
77+
const std::vector<ArgGroup>& args,
78+
const std::vector<ValueRef>& resize_args) {
79+
(void)resize_args;
80+
const ValueRef dlogits = args.at(0).refs.at(0);
81+
const ValueRef loss_partial = args.at(0).refs.at(1);
82+
const ValueRef logits = args.at(1).refs.at(0);
83+
84+
const std::vector<int64_t> logits_sizes = graph->sizes_of(logits);
85+
graph->virtual_resize(dlogits, logits_sizes);
86+
graph->virtual_resize(loss_partial, {logits_sizes.at(0)});
87+
}
88+
89+
void resize_fused_ce_sum_node(
90+
ComputeGraph* graph,
91+
const std::vector<ArgGroup>& args,
92+
const std::vector<ValueRef>& resize_args) {
93+
(void)resize_args;
94+
const ValueRef loss = args.at(0).refs.at(0);
95+
// Rank-0 scalar, matching fused_ce_meta's logits.new_empty([]).
96+
graph->virtual_resize(loss, {});
97+
}
98+
99+
void fused_ce(ComputeGraph& graph, const std::vector<ValueRef>& args) {
100+
int arg_idx = 0;
101+
const ValueRef logits = args[arg_idx++];
102+
const ValueRef labels = args[arg_idx++];
103+
const ValueRef n_valid_ref = args[arg_idx++];
104+
const ValueListPtr out_tuple = graph.get_value_list(args[arg_idx++]);
105+
const ValueRef loss = out_tuple->at(0);
106+
const ValueRef dlogits = out_tuple->at(1);
107+
108+
VK_CHECK_COND(
109+
graph.is_buffer_storage(logits) && graph.is_buffer_storage(dlogits),
110+
"fused_ce: logits and dlogits must use buffer storage");
111+
VK_CHECK_COND(
112+
graph.dim_of(logits) == 2, "fused_ce: logits must be 2D [n_rows, vocab]");
113+
VK_CHECK_COND(
114+
graph.sizes_of(dlogits) == graph.sizes_of(logits),
115+
"fused_ce: dlogits must match logits shape");
116+
VK_CHECK_COND(
117+
graph.dtype_of(labels) == vkapi::kInt, "fused_ce: labels must be int32");
118+
VK_CHECK_COND(graph.dim_of(labels) == 1, "fused_ce: labels must be 1D [N]");
119+
VK_CHECK_COND(
120+
graph.size_at<int64_t>(0, labels) == graph.size_at<int64_t>(0, logits),
121+
"fused_ce: labels length must equal number of rows");
122+
VK_CHECK_COND(graph.numel_of(loss) == 1, "fused_ce: loss must be a scalar");
123+
124+
const int32_t n_rows = graph.size_at<int32_t>(0, logits);
125+
const int32_t vocab = graph.size_at<int32_t>(1, logits);
126+
const float n_valid = graph.extract_scalar<float>(n_valid_ref);
127+
128+
// One workgroup per row; the Vulkan spec guarantees maxComputeWorkGroupCount
129+
// >= 65535 per dimension.
130+
VK_CHECK_COND(
131+
n_rows <= 65535, "fused_ce: n_rows exceeds max workgroup count");
132+
133+
// Per-row loss contributions, reduced to the scalar loss by the sum node.
134+
TmpTensor loss_partial(
135+
&graph, {n_rows}, graph.dtype_of(logits), utils::kBuffer);
136+
137+
std::string kernel_name = "fused_ce";
138+
kernel_name.reserve(kShaderNameReserve);
139+
add_storage_type_suffix(kernel_name, graph.storage_type_of(dlogits));
140+
add_dtype_suffix(kernel_name, graph.dtype_of(dlogits));
141+
142+
graph.execute_nodes().emplace_back(new DynamicDispatchNode(
143+
graph,
144+
VK_KERNEL_FROM_STR(kernel_name),
145+
fused_ce_global_wg_size,
146+
fused_ce_local_wg_size,
147+
// Inputs and Outputs
148+
{{{dlogits, loss_partial}, vkapi::kWrite},
149+
{{logits, labels}, vkapi::kRead}},
150+
// Shader params buffers
151+
{graph.create_params_buffer(vocab),
152+
graph.create_params_buffer(n_rows),
153+
graph.create_params_buffer(n_valid)},
154+
// Push Constants
155+
{},
156+
// Specialization Constants
157+
{},
158+
// Resize Args
159+
{},
160+
// Resizing Logic
161+
resize_fused_ce_node));
162+
163+
std::string sum_kernel_name = "fused_ce_sum";
164+
sum_kernel_name.reserve(kShaderNameReserve);
165+
add_storage_type_suffix(sum_kernel_name, graph.storage_type_of(loss));
166+
add_dtype_suffix(sum_kernel_name, graph.dtype_of(loss));
167+
168+
// Separate dispatch node so the loss_partial write (above) is ordered before
169+
// this read: the graph inserts a per-tensor pipeline barrier between them.
170+
graph.execute_nodes().emplace_back(new DynamicDispatchNode(
171+
graph,
172+
VK_KERNEL_FROM_STR(sum_kernel_name),
173+
fused_ce_sum_global_wg_size,
174+
fused_ce_sum_local_wg_size,
175+
// Inputs and Outputs
176+
{{loss, vkapi::kWrite}, {loss_partial, vkapi::kRead}},
177+
// Shader params buffers
178+
{graph.create_params_buffer(n_rows)},
179+
// Push Constants
180+
{},
181+
// Specialization Constants
182+
{},
183+
// Resize Args
184+
{},
185+
// Resizing Logic
186+
resize_fused_ce_sum_node));
187+
}
188+
189+
REGISTER_OPERATORS {
190+
VK_REGISTER_OP(et_vk.fused_ce.default, fused_ce);
191+
}
192+
193+
} // namespace vkcompute

0 commit comments

Comments
 (0)