| 1272 | } |
| 1273 | |
| 1274 | static std::optional<webgpu_encoded_op> ggml_webgpu_set_rows(webgpu_context & ctx, |
| 1275 | ggml_tensor * src, |
| 1276 | ggml_tensor * idx, |
| 1277 | ggml_tensor * dst) { |
| 1278 | // For set rows specifically, we need to check if src and idx are empty |
| 1279 | // tensors. |
| 1280 | if (ggml_is_empty(src) || ggml_is_empty(idx)) { |
| 1281 | return std::nullopt; |
| 1282 | } |
| 1283 | |
| 1284 | ggml_webgpu_shader_lib_context shader_lib_ctx = {}; |
| 1285 | shader_lib_ctx.src0 = src; |
| 1286 | shader_lib_ctx.src1 = idx; |
| 1287 | shader_lib_ctx.dst = dst; |
| 1288 | shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup; |
| 1289 | |
| 1290 | webgpu_pipeline pipeline = ctx->shader_lib->get_set_rows_pipeline(shader_lib_ctx); |
| 1291 | |
| 1292 | auto * decisions = static_cast<ggml_webgpu_set_rows_shader_decisions *>(pipeline.context.get()); |
| 1293 | |
| 1294 | std::vector<uint32_t> params = { |
| 1295 | (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src) / ggml_type_size(src->type)), |
| 1296 | (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, idx) / ggml_type_size(idx->type)), |
| 1297 | (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)), |
| 1298 | // Convert byte-strides to element-strides |
| 1299 | (uint32_t) (src->nb[1] / ggml_type_size(src->type)), (uint32_t) (src->nb[2] / ggml_type_size(src->type)), |
| 1300 | (uint32_t) (src->nb[3] / ggml_type_size(src->type)), (uint32_t) (idx->nb[0] / ggml_type_size(idx->type)), |
| 1301 | (uint32_t) (idx->nb[1] / ggml_type_size(idx->type)), (uint32_t) (idx->nb[2] / ggml_type_size(idx->type)), |
| 1302 | (uint32_t) (dst->nb[1] / ggml_type_size(dst->type)), (uint32_t) (dst->nb[2] / ggml_type_size(dst->type)), |
| 1303 | (uint32_t) (dst->nb[3] / ggml_type_size(dst->type)), |
| 1304 | // Shape of src |
| 1305 | (uint32_t) src->ne[0], (uint32_t) src->ne[1], (uint32_t) src->ne[2], (uint32_t) src->ne[3], |
| 1306 | // Shape of idx |
| 1307 | (uint32_t) (idx->ne[1]), (uint32_t) (idx->ne[2]) |
| 1308 | }; |
| 1309 | |
| 1310 | std::vector<wgpu::BindGroupEntry> entries = { |
| 1311 | ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src), |
| 1312 | ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, idx), |
| 1313 | ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst), |
| 1314 | }; |
| 1315 | |
| 1316 | if (decisions->i64_idx) { |
| 1317 | entries.push_back(ggml_webgpu_make_bind_group_entry(3, ctx->set_rows_dev_error_buf, 0, |
| 1318 | ctx->set_rows_dev_error_buf.GetSize())); |
| 1319 | } |
| 1320 | |
| 1321 | uint32_t threads; |
| 1322 | if (decisions->vec4) { |
| 1323 | threads = (src->ne[1] * src->ne[2] * src->ne[3]) * (src->ne[0] / 4); |
| 1324 | } else { |
| 1325 | threads = src->ne[0] * src->ne[1] * src->ne[2] * src->ne[3]; |
| 1326 | } |
| 1327 | uint32_t wg_x = CEIL_DIV(threads, decisions->wg_size); |
| 1328 | return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, 1); |
| 1329 | } |
| 1330 | |
| 1331 | // Workgroup size is a common constant |
no test coverage detected