Skip to content

Commit 566ded7

Browse files
JCNTHJCNTH
committed
[ExecuTorch][Vulkan] Add et_vk.adamw_step kernel (on-GPU AdamW optimizer step)
Pull Request resolved: #20937 Vulkan GLSL kernel + handler for the training custom op `et_vk.adamw_step` (elementwise in-place AdamW parameter update). Mirrors `impl/BinaryOp.cpp` elementwise structure + `unary_op.glsl` push-constant scalars; param/m/v are read-modify-write (`kReadWrite`). AOT registration already exists under the shared Vulkan partitioner. ghstack-source-id: 405026087 @exported-using-ghexport Differential Revision: [D111761777](https://our.internmc.facebook.com/intern/diff/D111761777/)
1 parent 4b8a685 commit 566ded7

3 files changed

Lines changed: 197 additions & 0 deletions

File tree

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
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 PRECISION ${PRECISION}
12+
13+
#define T ${buffer_scalar_type(DTYPE)}
14+
15+
${define_active_storage_type(STORAGE)}
16+
17+
${define_required_extensions(STORAGE, DTYPE)}
18+
19+
layout(std430) buffer;
20+
21+
${layout_declare_tensor(B, "rw", "t_param", DTYPE, STORAGE)}
22+
${layout_declare_tensor(B, "rw", "t_m", DTYPE, STORAGE)}
23+
${layout_declare_tensor(B, "rw", "t_v", DTYPE, STORAGE)}
24+
${layout_declare_tensor(B, "r", "t_grad", DTYPE, STORAGE)}
25+
26+
layout(push_constant) uniform restrict Block {
27+
int numel;
28+
float lr;
29+
float beta1;
30+
float beta2;
31+
float eps;
32+
float weight_decay;
33+
float bias_correction1;
34+
float bias_correction2;
35+
};
36+
37+
layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
38+
39+
void main() {
40+
const int i = int(gl_GlobalInvocationID.x);
41+
if (i >= numel) {
42+
return;
43+
}
44+
T g = t_grad[i];
45+
T p = t_param[i];
46+
p = p - lr * weight_decay * p;
47+
T m = beta1 * t_m[i] + (1.0 - beta1) * g;
48+
T v = beta2 * t_v[i] + (1.0 - beta2) * g * g;
49+
t_m[i] = m;
50+
t_v[i] = v;
51+
T mhat = m / bias_correction1;
52+
T vhat = v / bias_correction2;
53+
t_param[i] = p - lr * mhat / (sqrt(vhat) + eps);
54+
}

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
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+
adamw_step:
8+
parameter_names_with_default_values:
9+
DTYPE: float
10+
STORAGE: buffer
11+
generate_variant_forall:
12+
STORAGE:
13+
- VALUE: buffer
14+
DTYPE:
15+
- VALUE: float
16+
shader_variants:
17+
- NAME: adamw_step

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

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
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/DynamicDispatchNode.h>
10+
#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
11+
12+
#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
13+
14+
#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
15+
16+
#include <array>
17+
18+
namespace vkcompute {
19+
20+
void check_adamw_step_args(
21+
ComputeGraph& graph,
22+
const ValueRef param,
23+
const ValueRef m,
24+
const ValueRef v,
25+
const ValueRef grad) {
26+
VK_CHECK_COND(graph.dtype_of(param) == vkapi::kFloat);
27+
VK_CHECK_COND(graph.dtype_of(m) == vkapi::kFloat);
28+
VK_CHECK_COND(graph.dtype_of(v) == vkapi::kFloat);
29+
VK_CHECK_COND(graph.dtype_of(grad) == vkapi::kFloat);
30+
31+
const int32_t numel = graph.numel_of(param);
32+
VK_CHECK_COND(graph.numel_of(m) == numel);
33+
VK_CHECK_COND(graph.numel_of(v) == numel);
34+
VK_CHECK_COND(graph.numel_of(grad) == numel);
35+
}
36+
37+
void add_adamw_step_node(
38+
ComputeGraph& graph,
39+
const ValueRef param,
40+
const ValueRef m,
41+
const ValueRef v,
42+
const ValueRef grad,
43+
const ValueRef lr,
44+
const ValueRef beta1,
45+
const ValueRef beta2,
46+
const ValueRef eps,
47+
const ValueRef weight_decay,
48+
const ValueRef bias_correction1,
49+
const ValueRef bias_correction2) {
50+
check_adamw_step_args(graph, param, m, v, grad);
51+
52+
const float bc1 = graph.extract_scalar<float>(bias_correction1);
53+
const float bc2 = graph.extract_scalar<float>(bias_correction2);
54+
VK_CHECK_COND(bc1 != 0.0f);
55+
VK_CHECK_COND(bc2 != 0.0f);
56+
57+
// Split into <=16-byte push-constant entries (PushConstantData.h limit).
58+
const std::array<float, 4> pc0 = {
59+
graph.extract_scalar<float>(lr),
60+
graph.extract_scalar<float>(beta1),
61+
graph.extract_scalar<float>(beta2),
62+
graph.extract_scalar<float>(eps)};
63+
const std::array<float, 3> pc1 = {
64+
graph.extract_scalar<float>(weight_decay), bc1, bc2};
65+
66+
std::string kernel_name("adamw_step");
67+
add_storage_type_suffix(kernel_name, graph.storage_type_of(param));
68+
add_dtype_suffix(kernel_name, graph.dtype_of(param));
69+
70+
graph.execute_nodes().emplace_back(new DynamicDispatchNode(
71+
graph,
72+
VK_KERNEL_FROM_STR(kernel_name),
73+
default_pick_global_wg_size,
74+
default_pick_local_wg_size,
75+
// Inputs and Outputs
76+
{{param, vkapi::kReadWrite},
77+
{m, vkapi::kReadWrite},
78+
{v, vkapi::kReadWrite},
79+
{grad, vkapi::kRead}},
80+
// Shader params buffers
81+
{},
82+
// Push Constants
83+
{graph.numel_pc_of(param),
84+
PushConstantDataInfo(pc0.data(), sizeof(pc0)),
85+
PushConstantDataInfo(pc1.data(), sizeof(pc1))},
86+
// Specialization Constants
87+
{},
88+
// Resize Args
89+
{},
90+
// Resizing Logic
91+
nullptr));
92+
}
93+
94+
void adamw_step(ComputeGraph& graph, const std::vector<ValueRef>& args) {
95+
int arg_idx = 0;
96+
const ValueRef param = args[arg_idx++];
97+
const ValueRef m = args[arg_idx++];
98+
const ValueRef v = args[arg_idx++];
99+
const ValueRef grad = args[arg_idx++];
100+
const ValueRef lr = args[arg_idx++];
101+
const ValueRef beta1 = args[arg_idx++];
102+
const ValueRef beta2 = args[arg_idx++];
103+
const ValueRef eps = args[arg_idx++];
104+
const ValueRef weight_decay = args[arg_idx++];
105+
const ValueRef bias_correction1 = args[arg_idx++];
106+
const ValueRef bias_correction2 = args[arg_idx++];
107+
add_adamw_step_node(
108+
graph,
109+
param,
110+
m,
111+
v,
112+
grad,
113+
lr,
114+
beta1,
115+
beta2,
116+
eps,
117+
weight_decay,
118+
bias_correction1,
119+
bias_correction2);
120+
}
121+
122+
REGISTER_OPERATORS {
123+
VK_REGISTER_OP(et_vk.adamw_step.default, adamw_step);
124+
}
125+
126+
} // namespace vkcompute

0 commit comments

Comments
 (0)