Skip to content

Commit b900c6f

Browse files
JCNTHJCNTH
committed
Update (base update)
[ghstack-poisoned]
1 parent e9fa1d7 commit b900c6f

18 files changed

Lines changed: 76 additions & 79 deletions

File tree

‎backends/vulkan/op_registry.py‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1746,8 +1746,6 @@ def register_rms_norm():
17461746
)
17471747

17481748

1749-
1750-
17511749
@update_features(
17521750
[
17531751
exir_ops.edge.aten.ne.Scalar,
@@ -1766,8 +1764,6 @@ def register_compare_scalar_ops():
17661764
)
17671765

17681766

1769-
1770-
17711767
@update_features(exir_ops.edge.aten.logical_not.default)
17721768
def register_logical_not():
17731769
return OpFeatures(

‎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/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);

‎backends/webgpu/runtime/ops/reduce/Reduce.cpp‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,9 @@ void reduce_impl(
227227
wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
228228
g.dispatch_at(dispatch_idx).workgroup_count_x =
229229
utils::compute_1d_workgroup_count(
230-
g.device(), static_cast<uint32_t>(live_outputs), wg_size,
230+
g.device(),
231+
static_cast<uint32_t>(live_outputs),
232+
wg_size,
231233
"reduce(resize)");
232234
// Propagate reduced output dims for downstream resize hooks.
233235
int64_t nd = static_cast<int64_t>(d.size());

‎backends/webgpu/runtime/ops/softmax/Softmax.cpp‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -173,8 +173,7 @@ void softmax_impl(WebGPUGraph& graph, const std::vector<int>& args) {
173173

174174
// Dynamic shapes: recompute the decomposition for the reduced dim + dispatch.
175175
WGPUBuffer params_buf = uniform_buffer;
176-
const uint64_t build_numel =
177-
static_cast<uint64_t>(outer) * r * inner;
176+
const uint64_t build_numel = static_cast<uint64_t>(outer) * r * inner;
178177
graph.add_tensor_resize_hook(
179178
in_id,
180179
[in_id, out_id, dim, build_numel, wg_size, dispatch_idx, params_buf](
@@ -198,7 +197,9 @@ void softmax_impl(WebGPUGraph& graph, const std::vector<int>& args) {
198197
wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
199198
g.dispatch_at(dispatch_idx).workgroup_count_x =
200199
utils::compute_1d_workgroup_count(
201-
g.device(), static_cast<uint32_t>(live_lines), wg_size,
200+
g.device(),
201+
static_cast<uint32_t>(live_lines),
202+
wg_size,
202203
"softmax(resize)");
203204
g.set_cur_dims(out_id, d);
204205
});

‎backends/webgpu/runtime/ops/where/Where.cpp‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -60,10 +60,12 @@ void where_impl(WebGPUGraph& graph, const std::vector<int>& args) {
6060
fill_tensor_meta_broadcast(b_tensor, out_ndim, &b_meta);
6161

6262
// a/b/out are fp32; cond is 1-byte bool (read byte-packed as array<u32>).
63-
if (out_tensor.nbytes != static_cast<size_t>(out_meta.numel) * sizeof(float) ||
63+
if (out_tensor.nbytes !=
64+
static_cast<size_t>(out_meta.numel) * sizeof(float) ||
6465
a_tensor.nbytes != static_cast<size_t>(a_meta.numel) * sizeof(float) ||
6566
b_tensor.nbytes != static_cast<size_t>(b_meta.numel) * sizeof(float)) {
66-
throw std::runtime_error("where: non-fp32 self/other (nbytes != numel * 4)");
67+
throw std::runtime_error(
68+
"where: non-fp32 self/other (nbytes != numel * 4)");
6769
}
6870
if (cond_tensor.nbytes != static_cast<size_t>(cond_meta.numel)) {
6971
throw std::runtime_error("where: condition is not a 1-byte (bool) tensor");
@@ -169,8 +171,16 @@ void where_impl(WebGPUGraph& graph, const std::vector<int>& args) {
169171
// Dynamic shapes: rebuild the 4 broadcast TensorMeta UBOs + dispatch count.
170172
WGPUBuffer o_buf = out_meta_buf, c_buf = cond_meta_buf, a_buf = a_meta_buf,
171173
bb_buf = b_meta_buf;
172-
auto where_resize = [cond_id, a_id, b_id, out_id, wg_size, dispatch_idx, o_buf,
173-
c_buf, a_buf, bb_buf](WebGPUGraph& g) {
174+
auto where_resize = [cond_id,
175+
a_id,
176+
b_id,
177+
out_id,
178+
wg_size,
179+
dispatch_idx,
180+
o_buf,
181+
c_buf,
182+
a_buf,
183+
bb_buf](WebGPUGraph& g) {
174184
const auto& c = g.cur_dims(cond_id);
175185
const auto& a = g.cur_dims(a_id);
176186
const auto& b = g.cur_dims(b_id);
@@ -184,7 +194,8 @@ void where_impl(WebGPUGraph& graph, const std::vector<int>& args) {
184194
int64_t m = std::max({cv, av, bv});
185195
if ((cv != m && cv != 1) || (av != m && av != 1) ||
186196
(bv != m && bv != 1)) {
187-
throw std::runtime_error("where(resize): operands not broadcast-compatible");
197+
throw std::runtime_error(
198+
"where(resize): operands not broadcast-compatible");
188199
}
189200
out_d[i] = m;
190201
}

‎backends/webgpu/test/ops/test_bmm.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,7 @@
1919

2020
import torch
2121

22-
from executorch.backends.vulkan.partitioner.vulkan_partitioner import (
23-
VulkanPartitioner,
24-
)
22+
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
2523
from executorch.exir import to_edge_transform_and_lower
2624

2725
# name -> (shape_a [B, M, K], shape_b [B, K, N]). Output is [B, M, N].

‎backends/webgpu/test/ops/test_compare.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,7 @@
2121

2222
import torch
2323

24-
from executorch.backends.vulkan.partitioner.vulkan_partitioner import (
25-
VulkanPartitioner,
26-
)
24+
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
2725
from executorch.exir import to_edge_transform_and_lower
2826

2927
SCALAR = 0.0

0 commit comments

Comments
 (0)