Skip to content

Commit a1c18f8

Browse files
JCNTHJCNTH
committed
Update
[ghstack-poisoned]
2 parents 46fcce5 + 3ca6ea6 commit a1c18f8

7 files changed

Lines changed: 205 additions & 132 deletions

File tree

‎backends/webgpu/runtime/WebGPUBackend.cpp‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -104,19 +104,18 @@ Error WebGPUBackend::execute(
104104
// Copy inputs from EValue tensors to GPU buffers
105105
std::vector<InputData> inputs;
106106
inputs.reserve(num_inputs);
107-
for (size_t i = 0; i < num_inputs; i++) {
108-
const auto& tensor = args[i]->toTensor();
109-
const bool host_is_int64 =
110-
tensor.scalar_type() == executorch::aten::ScalarType::Long;
111-
inputs.push_back({tensor.const_data_ptr(), tensor.nbytes(), host_is_int64});
112-
}
113107
// Fail loud as a runtime Error so a throw never crosses the backend boundary.
114108
try {
115-
// Dynamic shapes: shrink each input to its live sizes before upload
116-
// (mirrors Vulkan maybe_resize_input). No-op when unchanged, so a static
117-
// graph is byte-identical.
109+
// Build the input list and, for dynamic shapes, shrink each input to its
110+
// live sizes before upload (mirrors Vulkan maybe_resize_input). No-op when
111+
// unchanged, so a static graph is byte-identical.
118112
for (size_t i = 0; i < num_inputs; i++) {
119-
const auto sizes = args[i]->toTensor().sizes();
113+
const auto& tensor = args[i]->toTensor();
114+
const bool host_is_int64 =
115+
tensor.scalar_type() == executorch::aten::ScalarType::Long;
116+
inputs.push_back(
117+
{tensor.const_data_ptr(), tensor.nbytes(), host_is_int64});
118+
const auto sizes = tensor.sizes();
120119
std::vector<int64_t> new_dims(sizes.begin(), sizes.end());
121120
graph->resize_input(graph->input_ids()[i], new_dims);
122121
}

‎backends/webgpu/runtime/WebGPUGraph.cpp‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -142,16 +142,14 @@ void WebGPUGraph::update_symints_from_inputs(
142142
// Reads the [0,..,index,..,0] element; symint sources are scalar-ish.
143143
const int64_t offset = static_cast<int64_t>(index) * stride;
144144
const void* host = inputs[pos].data;
145-
// Stored elem_size (live nbytes/numel mis-derives for a dynamic source).
146-
const size_t elem_size = tensors_[src.input_tensor_id].elem_size;
145+
// Interpret the HOST buffer by its scalar type, not the tensor's serialized
146+
// elem_size: copy_inputs narrows an int64 host input to an int32 buffer, so
147+
// elem_size (buffer-derived) would misread int64 host data as int32.
147148
int32_t val;
148-
if (elem_size == sizeof(int64_t)) {
149+
if (inputs[pos].host_is_int64) {
149150
val = static_cast<int32_t>(static_cast<const int64_t*>(host)[offset]);
150-
} else if (elem_size == sizeof(int32_t)) {
151-
val = static_cast<const int32_t*>(host)[offset];
152151
} else {
153-
throw std::runtime_error(
154-
"select_as_symint: unsupported input element size");
152+
val = static_cast<const int32_t*>(host)[offset];
155153
}
156154
set_symint(src.symint_id, val);
157155
}

‎backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp‎

Lines changed: 56 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,50 @@ static_assert(
3737
sizeof(EmbeddingParams) == 32,
3838
"EmbeddingParams must be 32 bytes");
3939

40+
// Resize hook body: recompute counts/dispatch; out = indices dims +
41+
// [embed_dim].
42+
void resize_embedding_q4gsw(
43+
WebGPUGraph& g,
44+
int indices_id,
45+
int out_id,
46+
uint32_t embed_dim,
47+
uint32_t blocks_per_row,
48+
uint32_t gs_u,
49+
uint32_t groups_per_row,
50+
uint32_t bytes_per_row,
51+
uint32_t wg_size,
52+
size_t dispatch_idx,
53+
WGPUBuffer params_buf) {
54+
const auto& id = g.cur_dims(indices_id);
55+
const uint64_t ni = utils::numel_of(id);
56+
if (ni == 0) {
57+
throw std::runtime_error("WebGPU embedding_q4gsw: zero indices");
58+
}
59+
const uint64_t total_blocks = ni * blocks_per_row;
60+
if (total_blocks > UINT32_MAX) {
61+
throw std::runtime_error(
62+
"WebGPU embedding_q4gsw: total_blocks exceeds uint32");
63+
}
64+
std::vector<int64_t> od = id;
65+
od.push_back(static_cast<int64_t>(embed_dim));
66+
g.set_cur_dims(out_id, od);
67+
EmbeddingParams p = {};
68+
p.embed_dim = embed_dim;
69+
p.blocks_per_row = blocks_per_row;
70+
p.num_indices = static_cast<uint32_t>(ni);
71+
p.group_size = gs_u;
72+
p.groups_per_row = groups_per_row;
73+
p.bytes_per_row = bytes_per_row;
74+
p.total_blocks = static_cast<uint32_t>(total_blocks);
75+
wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
76+
g.dispatch_at(dispatch_idx).workgroup_count_x =
77+
utils::compute_1d_workgroup_count(
78+
g.device(),
79+
static_cast<uint32_t>(total_blocks),
80+
wg_size,
81+
"embedding_q4gsw(resize)");
82+
}
83+
4084
// arg order mirrors Vulkan EmbeddingQ4gsw.cpp.
4185
void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
4286
const int weight_id = args.at(0);
@@ -241,34 +285,18 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
241285
wg_size,
242286
dispatch_idx,
243287
params_buf](WebGPUGraph& g) {
244-
const auto& id = g.cur_dims(indices_id);
245-
const uint64_t ni = utils::numel_of(id);
246-
if (ni == 0) {
247-
throw std::runtime_error("WebGPU embedding_q4gsw: zero indices");
248-
}
249-
const uint64_t total_blocks = ni * blocks_per_row;
250-
if (total_blocks > UINT32_MAX) {
251-
throw std::runtime_error(
252-
"WebGPU embedding_q4gsw: total_blocks exceeds uint32");
253-
}
254-
std::vector<int64_t> od = id;
255-
od.push_back(static_cast<int64_t>(embed_dim));
256-
g.set_cur_dims(out_id, od);
257-
EmbeddingParams p = {};
258-
p.embed_dim = embed_dim;
259-
p.blocks_per_row = blocks_per_row;
260-
p.num_indices = static_cast<uint32_t>(ni);
261-
p.group_size = gs_u;
262-
p.groups_per_row = groups_per_row;
263-
p.bytes_per_row = bytes_per_row;
264-
p.total_blocks = static_cast<uint32_t>(total_blocks);
265-
wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
266-
g.dispatch_at(dispatch_idx).workgroup_count_x =
267-
utils::compute_1d_workgroup_count(
268-
g.device(),
269-
static_cast<uint32_t>(total_blocks),
270-
wg_size,
271-
"embedding_q4gsw(resize)");
288+
resize_embedding_q4gsw(
289+
g,
290+
indices_id,
291+
out_id,
292+
embed_dim,
293+
blocks_per_row,
294+
gs_u,
295+
groups_per_row,
296+
bytes_per_row,
297+
wg_size,
298+
dispatch_idx,
299+
params_buf);
272300
});
273301

274302
wgpuShaderModuleRelease(shader);

‎backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -277,7 +277,6 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
277277
{pipeline, bind_group, workgroup_count, "linear_q4gsw"});
278278

279279
// Dynamic shapes: recompute dispatch + params.M for the live M.
280-
WGPUBuffer params_buf = uniform_buffer;
281280
graph.add_tensor_resize_hook(
282281
in_id,
283282
[in_id,
@@ -292,26 +291,28 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
292291
wg_size,
293292
use_gemv,
294293
dispatch_idx,
295-
params_buf](WebGPUGraph& g) {
294+
uniform_buffer](WebGPUGraph& g) {
296295
const auto& d = g.cur_dims(in_id);
297296
if (d.empty()) {
298-
throw std::runtime_error("WebGPU linear_q4gsw: empty input dims");
297+
throw std::runtime_error(
298+
"WebGPU linear_q4gsw(resize): empty input dims");
299299
}
300300
const uint64_t numel = utils::numel_of(d);
301301
if (numel % static_cast<uint64_t>(K) != 0u) {
302302
throw std::runtime_error(
303-
"WebGPU linear_q4gsw: live input numel not a multiple of K");
303+
"WebGPU linear_q4gsw(resize): live input numel not a multiple "
304+
"of K");
304305
}
305306
const uint32_t m =
306307
static_cast<uint32_t>(numel / static_cast<uint64_t>(K));
307308
if (m == 0u) {
308-
throw std::runtime_error("WebGPU linear_q4gsw: live M == 0");
309+
throw std::runtime_error("WebGPU linear_q4gsw(resize): live M == 0");
309310
}
310311
// Buffers/bind-groups were sized for the build-time max M; a larger
311312
// live M would write out of bounds.
312313
if (m > M) {
313314
throw std::runtime_error(
314-
"WebGPU linear_q4gsw: live M exceeds the build-time max");
315+
"WebGPU linear_q4gsw(resize): live M exceeds the build-time max");
315316
}
316317
const uint32_t wgc = compute_q4gsw_workgroup_count(
317318
g.device(), use_gemv, m, N, wg_size, "linear_q4gsw(resize)");
@@ -323,7 +324,7 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
323324
p.group_size = gs;
324325
p.padded_N = padded_N;
325326
p.has_bias = has_bias;
326-
wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
327+
wgpuQueueWriteBuffer(g.queue(), uniform_buffer, 0, &p, sizeof(p));
327328
g.dispatch_at(dispatch_idx).workgroup_count_x = wgc;
328329
std::vector<int64_t> od(d.begin(), d.end());
329330
od.back() = static_cast<int64_t>(N);

‎backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp‎

Lines changed: 35 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,39 @@ struct RmsNormParams {
3232
};
3333
static_assert(sizeof(RmsNormParams) == 16, "RmsNormParams must be 16 bytes");
3434

35+
// Resize hook body: recompute num_rows + rewrite the UBO for the live input.
36+
void resize_rms_norm(
37+
WebGPUGraph& g,
38+
int in_id,
39+
int out_id,
40+
uint32_t row_width,
41+
float epsilon,
42+
size_t dispatch_idx,
43+
WGPUBuffer params_buf) {
44+
const auto& d = g.cur_dims(in_id);
45+
const uint64_t numel = utils::numel_of(d);
46+
if (numel % static_cast<uint64_t>(row_width) != 0) {
47+
throw std::runtime_error(
48+
"WebGPU rms_norm: numel not a multiple of row_width");
49+
}
50+
const uint32_t rows =
51+
static_cast<uint32_t>(numel / static_cast<uint64_t>(row_width));
52+
if (rows == 0) {
53+
throw std::runtime_error("WebGPU rms_norm: zero rows");
54+
}
55+
if (rows > 65535u) {
56+
throw std::runtime_error(
57+
"WebGPU rms_norm: num_rows exceeds the 1D dispatch limit (65535)");
58+
}
59+
RmsNormParams p = {};
60+
p.num_rows = rows;
61+
p.row_width = row_width;
62+
p.epsilon = epsilon;
63+
wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
64+
g.dispatch_at(dispatch_idx).workgroup_count_x = rows;
65+
g.set_cur_dims(out_id, d);
66+
}
67+
3568
void rms_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
3669
// et_vk.rms_norm.default args: [in, weight, eps, out]
3770
const int in_id = args.at(0);
@@ -197,28 +230,8 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
197230
in_id,
198231
[in_id, out_id, row_width, epsilon, dispatch_idx, params_buf](
199232
WebGPUGraph& g) {
200-
const auto& d = g.cur_dims(in_id);
201-
const uint64_t numel = utils::numel_of(d);
202-
if (numel % static_cast<uint64_t>(row_width) != 0) {
203-
throw std::runtime_error(
204-
"WebGPU rms_norm: numel not a multiple of row_width");
205-
}
206-
const uint32_t rows =
207-
static_cast<uint32_t>(numel / static_cast<uint64_t>(row_width));
208-
if (rows == 0) {
209-
throw std::runtime_error("WebGPU rms_norm: zero rows");
210-
}
211-
if (rows > 65535u) {
212-
throw std::runtime_error(
213-
"WebGPU rms_norm: num_rows exceeds the 1D dispatch limit (65535)");
214-
}
215-
RmsNormParams p = {};
216-
p.num_rows = rows;
217-
p.row_width = row_width;
218-
p.epsilon = epsilon;
219-
wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
220-
g.dispatch_at(dispatch_idx).workgroup_count_x = rows;
221-
g.set_cur_dims(out_id, d);
233+
resize_rms_norm(
234+
g, in_id, out_id, row_width, epsilon, dispatch_idx, params_buf);
222235
});
223236

224237
// Release intermediate objects (pipeline and bind_group are kept by dispatch)

0 commit comments

Comments
 (0)