diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 44525b54afc..130d81e600f 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -18,6 +18,8 @@ #include #include #include +#include +#include namespace executorch::backends::webgpu { @@ -50,6 +52,62 @@ constexpr int64_t kQ4gswShmemTileN = 32; constexpr uint32_t kQ4gswShmemMinDim = 4096u; constexpr uint32_t kQ4gswShmemNMinDim = 2048u; +// Workgroup count for a linear_q4gsw dispatch (bicol GEMV / shmem GEMM / tiled +// GEMM), with the range/limit guards shared by the build-time path and the +// resize hook. use_gemv/use_shmem_gemm are the build-time routing decision (the +// shader/pipeline is fixed at build); the resize hook re-runs this with live m. +uint32_t compute_q4gsw_workgroup_count( + WGPUDevice device, + bool use_gemv, + bool use_shmem_gemm, + uint32_t m, + uint32_t n, + uint32_t wg_size, + const char* op_name) { + if (use_gemv) { + // bicol: fixed 64 lanes, 2 output columns/workgroup, grid-strided over + // ceil(N/2) column-pairs (M == 1 on this decode path). + const uint64_t pairs = (static_cast(n) + 1u) / 2u; + if (pairs == 0u || pairs > UINT32_MAX) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": N/2 out of range"); + } + const uint32_t wgc = + utils::clamp_workgroup_count(device, static_cast(pairs)); + if (wgc == 0u) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": zero GEMV dispatch"); + } + return wgc; + } + if (use_shmem_gemm) { + // shmem GEMM: one workgroup per tile, no grid-stride -> throw over limit. + const int64_t total_wgs = utils::div_up(m, kQ4gswShmemTileM) * + utils::div_up(n, kQ4gswShmemTileN); + WGPULimits limits = {}; + const uint32_t max_wgs = + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupsPerDimension > 0 + ? limits.maxComputeWorkgroupsPerDimension + : 65535u; + if (total_wgs > static_cast(max_wgs)) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + + ": shmem GEMM tile count exceeds the 1D dispatch limit"); + } + return static_cast(total_wgs); + } + const int64_t total_tiles = utils::div_up(m, kQ4gswTileM) * + utils::div_up(n, kQ4gswTileN); + if (total_tiles > static_cast(UINT32_MAX)) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + + ": tile count exceeds the 1D dispatch limit"); + } + return utils::compute_1d_workgroup_count( + device, static_cast(total_tiles), wg_size, op_name); +} + // et_vk.linear_q4gsw args: [in, weight, scales, group_size, bias, out]. void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { const int in_id = args.at(0); @@ -137,44 +195,8 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { const char* shader_src = use_gemv ? kQ4gswLinearCoop4BicolWGSL : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL : kQ4gswLinearWGSL; - uint32_t workgroup_count; - if (use_gemv) { - // bicol: fixed 64 lanes, 2 output columns/workgroup, grid-strided over - // ceil(N/2) column-pairs (M == 1 on this decode path). - const uint64_t pairs = (static_cast(N) + 1u) / 2u; - if (pairs == 0u || pairs > UINT32_MAX) { - throw std::runtime_error("WebGPU linear_q4gsw: N/2 out of range"); - } - workgroup_count = - utils::clamp_workgroup_count(device, static_cast(pairs)); - if (workgroup_count == 0u) { - throw std::runtime_error("WebGPU linear_q4gsw: zero GEMV dispatch"); - } - } else if (use_shmem_gemm) { - // shmem GEMM: one workgroup per tile, no grid-stride -> throw over limit. - const int64_t total_wgs = utils::div_up(M, kQ4gswShmemTileM) * - utils::div_up(N, kQ4gswShmemTileN); - WGPULimits limits = {}; - const uint32_t max_wgs = - wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && - limits.maxComputeWorkgroupsPerDimension > 0 - ? limits.maxComputeWorkgroupsPerDimension - : 65535u; - if (total_wgs > static_cast(max_wgs)) { - throw std::runtime_error( - "WebGPU linear_q4gsw: shmem GEMM tile count exceeds the 1D dispatch limit"); - } - workgroup_count = static_cast(total_wgs); - } else { - const int64_t total_tiles = utils::div_up(M, kQ4gswTileM) * - utils::div_up(N, kQ4gswTileN); - if (total_tiles > static_cast(UINT32_MAX)) { - throw std::runtime_error( - "WebGPU linear_q4gsw: tile count exceeds the 1D dispatch limit"); - } - workgroup_count = utils::compute_1d_workgroup_count( - device, static_cast(total_tiles), wg_size, "linear_q4gsw"); - } + const uint32_t workgroup_count = compute_q4gsw_workgroup_count( + device, use_gemv, use_shmem_gemm, M, N, wg_size, "linear_q4gsw"); // Optional bias: real buffer if present, else a dummy for the fixed layout. uint32_t has_bias = 0; @@ -287,12 +309,78 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - graph.add_dispatch({pipeline, bind_group, workgroup_count, "linear_q4gsw"}); + const size_t dispatch_idx = graph.add_dispatch( + {pipeline, bind_group, workgroup_count, "linear_q4gsw"}); + + // Dynamic shapes: recompute dispatch + params.M for the live M. use_gemv and + // use_shmem_gemm are captured (routing is fixed at build); the helper re-runs + // the same path's workgroup-count formula with the live m. + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + M, + K, + N, + K_packed, + gs, + padded_N, + has_bias, + wg_size, + use_gemv, + use_shmem_gemm, + dispatch_idx, + uniform_buffer](WebGPUGraph& g) { + const auto& d = g.cur_dims(in_id); + if (d.empty()) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): empty input dims"); + } + const uint64_t numel = utils::numel_of(d); + if (numel % static_cast(K) != 0u) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): live input numel not a multiple " + "of K"); + } + const uint32_t m = + static_cast(numel / static_cast(K)); + if (m == 0u) { + throw std::runtime_error("WebGPU linear_q4gsw(resize): live M == 0"); + } + // Buffers/bind-groups were sized for the build-time max M; a larger + // live M would write out of bounds. + if (m > M) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): live M exceeds the build-time max"); + } + const uint32_t wgc = compute_q4gsw_workgroup_count( + g.device(), + use_gemv, + use_shmem_gemm, + m, + N, + wg_size, + "linear_q4gsw(resize)"); + Q4gswParams p = {}; + p.M = m; + p.N = N; + p.K = K; + p.K_packed = K_packed; + p.group_size = gs; + p.padded_N = padded_N; + p.has_bias = has_bias; + wgpuQueueWriteBuffer(g.queue(), uniform_buffer, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc; + std::vector od(d.begin(), d.end()); + od.back() = static_cast(N); + g.set_cur_dims(out_id, od); + }); wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); wgpuPipelineLayoutRelease(pipeline_layout); - wgpuBufferRelease(uniform_buffer); + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); } } // namespace