Skip to content

Commit 73005fd

Browse files
JCNTHJCNTH
committed
Update
[ghstack-poisoned]
2 parents 127852e + db9be5a commit 73005fd

33 files changed

Lines changed: 119 additions & 136 deletions

‎backends/vulkan/custom_ops_lib.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -335,7 +335,13 @@ def linear_q4gsw_backward(ctx, grad_out):
335335
d_x = torch.ops.et_vk.linear_q4gsw_backward(
336336
grad_out, weights, weight_scales, ctx.group_size
337337
)
338-
return d_x, None, None, None, None # grads for (x, weights, scales, group_size, bias)
338+
return (
339+
d_x,
340+
None,
341+
None,
342+
None,
343+
None,
344+
) # grads for (x, weights, scales, group_size, bias)
339345

340346

341347
torch.library.register_autograd(
@@ -1197,7 +1203,6 @@ def fused_ce_backward(ctx, grad_loss, grad_dlogits):
11971203
fused_ce_op = getattr(getattr(torch.ops, namespace), name)
11981204

11991205

1200-
12011206
###########################
12021207
## adamw_step (training) ##
12031208
###########################
@@ -1272,7 +1277,6 @@ def linear_q4gsw_dw_meta(
12721277
linear_q4gsw_dw_op = getattr(getattr(torch.ops, namespace), name)
12731278

12741279

1275-
12761280
##################
12771281
## q4gsw_requant ##
12781282
##################

‎backends/vulkan/op_registry.py‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1769,8 +1769,6 @@ def register_fused_ce():
17691769
)
17701770

17711771

1772-
1773-
17741772
@update_features(
17751773
[
17761774
exir_ops.edge.aten.ne.Scalar,
@@ -1789,8 +1787,6 @@ def register_compare_scalar_ops():
17891787
)
17901788

17911789

1792-
1793-
17941790
@update_features(exir_ops.edge.aten.logical_not.default)
17951791
def register_logical_not():
17961792
return OpFeatures(
@@ -1801,7 +1797,6 @@ def register_logical_not():
18011797
)
18021798

18031799

1804-
18051800
@update_features("et_vk::adamw_step")
18061801
def register_adamw_step():
18071802
return OpFeatures(
@@ -1819,8 +1814,6 @@ def register_linear_q4gsw_dw():
18191814
)
18201815

18211816

1822-
1823-
18241817
@update_features(exir_ops.edge.et_vk.q4gsw_requant.default)
18251818
def register_q4gsw_requant():
18261819
return OpFeatures(

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

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,7 @@ utils::uvec3 linear_q4gsw_backward_global_wg_size(
4141
(void)resize_args;
4242
const ValueRef d_x = args.at(0).refs.at(0);
4343
const uint32_t K = graph->size_at<uint32_t>(-1, d_x);
44-
const uint32_t M =
45-
utils::safe_downcast<uint32_t>(graph->numel_of(d_x) / K);
44+
const uint32_t M = utils::safe_downcast<uint32_t>(graph->numel_of(d_x) / K);
4645
const uint32_t tiles = utils::div_up_4(M) * utils::div_up_4(K);
4746
return {tiles, 1u, 1u};
4847
}
@@ -92,8 +91,7 @@ void linear_q4gsw_backward(
9291
const ValueRef packed_scales =
9392
prepack_q4_scales(graph, weight_scales_data, in_dtype);
9493

95-
const uint32_t M =
96-
utils::safe_downcast<uint32_t>(graph.numel_of(d_out) / N);
94+
const uint32_t M = utils::safe_downcast<uint32_t>(graph.numel_of(d_out) / N);
9795
const uint32_t tiles =
9896
utils::div_up_4(M) * utils::div_up_4(static_cast<uint32_t>(K));
9997
VK_CHECK_COND(

‎backends/vulkan/test/op_tests/CMakeLists.txt‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -117,16 +117,14 @@ if(TARGET vulkan_backend AND LIB_TORCH)
117117
test_utils
118118
)
119119
vulkan_op_test(
120-
adamw_step_test ${CMAKE_CURRENT_SOURCE_DIR}/adamw_step_test.cpp
121-
test_utils
120+
adamw_step_test ${CMAKE_CURRENT_SOURCE_DIR}/adamw_step_test.cpp test_utils
122121
)
123122
vulkan_op_test(
124123
linear_q4gsw_dw_test ${CMAKE_CURRENT_SOURCE_DIR}/linear_q4gsw_dw_test.cpp
125124
test_utils
126125
)
127126
vulkan_op_test(
128-
fused_ce_test ${CMAKE_CURRENT_SOURCE_DIR}/fused_ce_test.cpp
129-
test_utils
127+
fused_ce_test ${CMAKE_CURRENT_SOURCE_DIR}/fused_ce_test.cpp test_utils
130128
)
131129
vulkan_op_test(
132130
quantized_linear_backward_test

‎backends/vulkan/test/op_tests/linear_q4gsw_dw_test.cpp‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,7 @@ at::Tensor linear_q4gsw_dw_reference_impl(
2727
const at::Tensor& x) {
2828
const int64_t N = d_out.size(-1);
2929
const int64_t K = x.size(-1);
30-
return d_out.reshape({-1, N})
31-
.t()
32-
.matmul(x.reshape({-1, K}))
33-
.contiguous();
30+
return d_out.reshape({-1, N}).t().matmul(x.reshape({-1, K})).contiguous();
3431
}
3532

3633
//
@@ -40,8 +37,7 @@ at::Tensor linear_q4gsw_dw_reference_impl(
4037
void test_vulkan_linear_q4gsw_dw_impl(
4138
const std::vector<int64_t>& d_out_sizes,
4239
const std::vector<int64_t>& x_sizes,
43-
const vkcompute::utils::StorageType storage =
44-
vkcompute::utils::kBuffer) {
40+
const vkcompute::utils::StorageType storage = vkcompute::utils::kBuffer) {
4541
at::Tensor d_out =
4642
at::rand(d_out_sizes, at::device(at::kCPU).dtype(at::kFloat));
4743
at::Tensor x = at::rand(x_sizes, at::device(at::kCPU).dtype(at::kFloat));

‎backends/webgpu/runtime/ops/compare/Compare.cpp‎

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,8 @@ void compare_impl(
6969
const size_t out_bind_size = (out_tensor.nbytes + 3) & ~size_t(3);
7070
const uint32_t n_words = (numel + 3u) / 4u;
7171

72-
uint32_t wg_size = utils::clamp_workgroup_size(device, kCompareWorkgroupSizeX);
72+
uint32_t wg_size =
73+
utils::clamp_workgroup_size(device, kCompareWorkgroupSizeX);
7374
uint32_t workgroup_count =
7475
utils::compute_1d_workgroup_count(device, n_words, wg_size, op_name);
7576

@@ -139,20 +140,21 @@ void compare_impl(
139140
graph.add_dispatch({pipeline, bind_group, workgroup_count});
140141

141142
WGPUBuffer p_buf = params_buf;
142-
auto cmp_resize = [self_id, out_id, mode, scalar, wg_size, dispatch_idx,
143-
p_buf, op_name](WebGPUGraph& g) {
144-
const auto& d = g.cur_dims(self_id);
145-
uint32_t n = 1u;
146-
for (auto x : d) {
147-
n *= static_cast<uint32_t>(x);
148-
}
149-
g.set_cur_dims(out_id, d);
150-
CompareParams p = {n, mode, scalar, 0u};
151-
wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p));
152-
const uint32_t nw = (n + 3u) / 4u;
153-
g.dispatch_at(dispatch_idx).workgroup_count_x =
154-
utils::compute_1d_workgroup_count(g.device(), nw, wg_size, op_name);
155-
};
143+
auto cmp_resize =
144+
[self_id, out_id, mode, scalar, wg_size, dispatch_idx, p_buf, op_name](
145+
WebGPUGraph& g) {
146+
const auto& d = g.cur_dims(self_id);
147+
uint32_t n = 1u;
148+
for (auto x : d) {
149+
n *= static_cast<uint32_t>(x);
150+
}
151+
g.set_cur_dims(out_id, d);
152+
CompareParams p = {n, mode, scalar, 0u};
153+
wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p));
154+
const uint32_t nw = (n + 3u) / 4u;
155+
g.dispatch_at(dispatch_idx).workgroup_count_x =
156+
utils::compute_1d_workgroup_count(g.device(), nw, wg_size, op_name);
157+
};
156158
graph.add_tensor_resize_hook(self_id, cmp_resize);
157159

158160
wgpuShaderModuleRelease(shader);

‎backends/webgpu/runtime/ops/embedding/Embedding.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,8 @@ void embedding_impl(WebGPUGraph& graph, const std::vector<int>& args) {
5454
if (out.nbytes != out_numel * sizeof(float) ||
5555
weight.nbytes % sizeof(float) != 0 ||
5656
indices.nbytes != index_numel * sizeof(int32_t)) {
57-
throw std::runtime_error("embedding: fp32 weight/out + i32 indices required");
57+
throw std::runtime_error(
58+
"embedding: fp32 weight/out + i32 indices required");
5859
}
5960
if (out_numel != index_numel * dim) {
6061
throw std::runtime_error("embedding: out numel != num_indices * dim");

‎backends/webgpu/runtime/ops/fill/Fill.cpp‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -144,20 +144,26 @@ void add_fill(
144144

145145
void full_impl(WebGPUGraph& graph, const std::vector<int>& args) {
146146
add_fill(
147-
graph, args.at(args.size() - 1), read_scalar(graph, args.at(1), "full"),
147+
graph,
148+
args.at(args.size() - 1),
149+
read_scalar(graph, args.at(1), "full"),
148150
"full");
149151
}
150152

151153
void full_like_impl(WebGPUGraph& graph, const std::vector<int>& args) {
152154
add_fill(
153-
graph, args.at(args.size() - 1),
154-
read_scalar(graph, args.at(1), "full_like"), "full_like");
155+
graph,
156+
args.at(args.size() - 1),
157+
read_scalar(graph, args.at(1), "full_like"),
158+
"full_like");
155159
}
156160

157161
void scalar_tensor_impl(WebGPUGraph& graph, const std::vector<int>& args) {
158162
add_fill(
159-
graph, args.at(args.size() - 1),
160-
read_scalar(graph, args.at(0), "scalar_tensor"), "scalar_tensor");
163+
graph,
164+
args.at(args.size() - 1),
165+
read_scalar(graph, args.at(0), "scalar_tensor"),
166+
"scalar_tensor");
161167
}
162168

163169
} // namespace

‎backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,9 @@ struct LogSoftmaxParams {
2929
uint32_t inner;
3030
uint32_t pad_;
3131
};
32-
static_assert(sizeof(LogSoftmaxParams) == 16, "LogSoftmaxParams must be 16 bytes");
32+
static_assert(
33+
sizeof(LogSoftmaxParams) == 16,
34+
"LogSoftmaxParams must be 16 bytes");
3335

3436
// Decompose dims into [outer, R, inner] for a reduction along `dim`.
3537
void decompose(
@@ -85,7 +87,8 @@ void log_softmax_impl(WebGPUGraph& graph, const std::vector<int>& args) {
8587
for (int64_t d : in.dims) {
8688
numel *= static_cast<uint64_t>(d);
8789
}
88-
if (in.nbytes != numel * sizeof(float) || out.nbytes != numel * sizeof(float)) {
90+
if (in.nbytes != numel * sizeof(float) ||
91+
out.nbytes != numel * sizeof(float)) {
8992
throw std::runtime_error(
9093
"WebGPU log_softmax: fp32-only (byte-size mismatch)");
9194
}
@@ -205,7 +208,9 @@ void log_softmax_impl(WebGPUGraph& graph, const std::vector<int>& args) {
205208
wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
206209
g.dispatch_at(dispatch_idx).workgroup_count_x =
207210
utils::compute_1d_workgroup_count(
208-
g.device(), static_cast<uint32_t>(live_lines), wg_size,
211+
g.device(),
212+
static_cast<uint32_t>(live_lines),
213+
wg_size,
209214
"log_softmax(resize)");
210215
g.set_cur_dims(out_id, std::vector<int64_t>(d.begin(), d.end()));
211216
});

‎backends/webgpu/runtime/ops/logical_not/LogicalNot.cpp‎

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -120,20 +120,21 @@ void logical_not_impl(WebGPUGraph& graph, const std::vector<int>& args) {
120120
graph.add_dispatch({pipeline, bind_group, workgroup_count});
121121

122122
WGPUBuffer p_buf = params_buf;
123-
auto resize = [self_id, out_id, wg_size, dispatch_idx, p_buf](
124-
WebGPUGraph& g) {
125-
const auto& d = g.cur_dims(self_id);
126-
uint32_t n = 1u;
127-
for (auto x : d) {
128-
n *= static_cast<uint32_t>(x);
129-
}
130-
g.set_cur_dims(out_id, d);
131-
LogicalNotParams p = {n, 0u, 0u, 0u};
132-
wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p));
133-
const uint32_t nw = (n + 3u) / 4u;
134-
g.dispatch_at(dispatch_idx).workgroup_count_x =
135-
utils::compute_1d_workgroup_count(g.device(), nw, wg_size, "logical_not");
136-
};
123+
auto resize =
124+
[self_id, out_id, wg_size, dispatch_idx, p_buf](WebGPUGraph& g) {
125+
const auto& d = g.cur_dims(self_id);
126+
uint32_t n = 1u;
127+
for (auto x : d) {
128+
n *= static_cast<uint32_t>(x);
129+
}
130+
g.set_cur_dims(out_id, d);
131+
LogicalNotParams p = {n, 0u, 0u, 0u};
132+
wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p));
133+
const uint32_t nw = (n + 3u) / 4u;
134+
g.dispatch_at(dispatch_idx).workgroup_count_x =
135+
utils::compute_1d_workgroup_count(
136+
g.device(), nw, wg_size, "logical_not");
137+
};
137138
graph.add_tensor_resize_hook(self_id, resize);
138139

139140
wgpuShaderModuleRelease(shader);

0 commit comments

Comments
 (0)