diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h index 561bebc6b87..60638b499bb 100644 --- a/backends/webgpu/runtime/WebGPUDispatchMath.h +++ b/backends/webgpu/runtime/WebGPUDispatchMath.h @@ -24,7 +24,7 @@ namespace executorch::backends::webgpu::utils { // Ceiling division for non-negative integers (mirrors Vulkan's utils::div_up). template inline T div_up(T a, T b) { - return (a + b - 1) / b; + return a / b + (a % b != 0); } // Product of a tensor's dims; the same accumulation was duplicated per-op. diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 58d7178ddbf..ca312bcda88 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -212,6 +212,31 @@ bool vk_datatype_is_int(vkgraph::VkDataType dtype) { } } +size_t storage_buffer_size(size_t nbytes) { + const size_t at_least_four = std::max(nbytes, size_t(4)); + if (at_least_four > std::numeric_limits::max() - 3u) { + throw std::runtime_error("WebGPU: storage buffer size overflows alignment"); + } + return (at_least_four + 3u) & ~size_t(3); +} + +void write_storage_buffer( + WGPUQueue queue, + WGPUBuffer buffer, + const void* data, + size_t nbytes) { + if (nbytes == 0u) { + return; + } + if (nbytes % 4u == 0u) { + wgpuQueueWriteBuffer(queue, buffer, 0, data, nbytes); + return; + } + std::vector padded(storage_buffer_size(nbytes), 0u); + std::memcpy(padded.data(), data, nbytes); + wgpuQueueWriteBuffer(queue, buffer, 0, padded.data(), padded.size()); +} + // Normalize a possibly-negative dim against rank; throws (fail-loud) if OOR. int normalize_dim(int dim, int rank, const char* op) { if (dim < 0) { @@ -276,7 +301,7 @@ WebGPUGraph::WebGPUGraph() = default; WGPUBuffer WebGPUGraph::create_scratch_buffer(size_t nbytes) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = nbytes > 0 ? nbytes : 4; + buf_desc.size = storage_buffer_size(nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -286,7 +311,7 @@ WGPUBuffer WebGPUGraph::create_scratch_buffer(size_t nbytes) { } WGPUBuffer WebGPUGraph::acquire_scratch(size_t nbytes) { - nbytes = nbytes > 0 ? nbytes : 4; + nbytes = storage_buffer_size(nbytes); // Best-fit reuse: smallest free slot with size in [nbytes, 2*nbytes] -- the // 2x cap stops a large Cmax-sized buffer from backing a tiny request. Never // reuse an in_use slot (co-live safety). @@ -890,6 +915,7 @@ void WebGPUGraph::build( throw std::runtime_error("WebGPU: tensor byte size overflows"); } tensor.is_int = vk_datatype_is_int(vk_tensor->datatype()); + tensor.is_bool = vk_tensor->datatype() == vkgraph::VkDataType::BOOL; tensor.is_int8 = vk_tensor->datatype() == vkgraph::VkDataType::INT8; tensor.nbytes = numel * tensor.elem_size; // Live dims start == max (serialized upper bound); resize_input shrinks @@ -911,7 +937,7 @@ void WebGPUGraph::build( tensor.cur_nbytes = tensor.nbytes; tensor_mem_obj_ids_[i] = -1; WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(tensor.nbytes, size_t(4)); + buf_desc.size = storage_buffer_size(tensor.nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -936,10 +962,9 @@ void WebGPUGraph::build( std::memcpy(&value, src + e * sizeof(float), sizeof(float)); converted[e] = executorch::runtime::etensor::Half(value); } - wgpuQueueWriteBuffer( + write_storage_buffer( queue_, tensor.buffer, - 0, converted.data(), converted.size() * sizeof(converted[0])); }; @@ -1013,7 +1038,7 @@ void WebGPUGraph::build( prepack_src_ids.count(i) != 0 && direct_use_ids.count(i) == 0; if (!defer) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(tensor.nbytes, size_t(4)); + buf_desc.size = storage_buffer_size(tensor.nbytes); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -1110,7 +1135,7 @@ void WebGPUGraph::build( shared_buffers_.resize(shared_buffer_sizes_.size(), nullptr); for (size_t id = 0; id < shared_buffer_sizes_.size(); id++) { WGPUBufferDescriptor buf_desc = {}; - buf_desc.size = std::max(shared_buffer_sizes_[id], size_t(4)); + buf_desc.size = storage_buffer_size(shared_buffer_sizes_[id]); buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst | WGPUBufferUsage_CopySrc; buf_desc.mappedAtCreation = false; @@ -1138,7 +1163,7 @@ void WebGPUGraph::build( // Create staging buffer for output readback WGPUBufferDescriptor staging_desc = {}; - staging_desc.size = std::max(tensors_[oid].nbytes, size_t(4)); + staging_desc.size = storage_buffer_size(tensors_[oid].nbytes); staging_desc.usage = WGPUBufferUsage_MapRead | WGPUBufferUsage_CopyDst; staging_desc.mappedAtCreation = false; output_staging_buffers_.push_back( @@ -1313,7 +1338,7 @@ void WebGPUGraph::materialize_constant(int const_value_id, WGPUBuffer dst) { cs.nbytes, "WebGPU: inline constant exceeds constant data"); if (cs.nbytes != 0) { - wgpuQueueWriteBuffer(queue_, dst, 0, data, cs.nbytes); + write_storage_buffer(queue_, dst, data, cs.nbytes); } } else if (cs.nbytes == 0) { return; @@ -1327,7 +1352,7 @@ void WebGPUGraph::materialize_constant(int const_value_id, WGPUBuffer dst) { throw std::runtime_error( "WebGPU: named constant '" + cs.named_key + "' undersized"); } - wgpuQueueWriteBuffer(queue_, dst, 0, buf->data(), cs.nbytes); + write_storage_buffer(queue_, dst, buf->data(), cs.nbytes); buf->Free(); } else { throw std::runtime_error("WebGPU: constant has no source"); @@ -1419,7 +1444,7 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { // Fast path: host and GPU element types match byte-for-byte. if (in.nbytes == live_nbytes) { - wgpuQueueWriteBuffer(queue_, tensor.buffer, 0, in.data, live_nbytes); + write_storage_buffer(queue_, tensor.buffer, in.data, live_nbytes); continue; } @@ -1439,8 +1464,7 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { #endif narrowed[e] = static_cast(src[e]); } - wgpuQueueWriteBuffer( - queue_, tensor.buffer, 0, narrowed.data(), live_nbytes); + write_storage_buffer(queue_, tensor.buffer, narrowed.data(), live_nbytes); continue; } @@ -1454,8 +1478,7 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { for (size_t e = 0; e < numel; e++) { narrowed[e] = executorch::runtime::etensor::Half(src[e]); } - wgpuQueueWriteBuffer( - queue_, tensor.buffer, 0, narrowed.data(), live_nbytes); + write_storage_buffer(queue_, tensor.buffer, narrowed.data(), live_nbytes); continue; } @@ -1684,10 +1707,11 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { } for (size_t i = 0; i < output_copies_.size(); i++) { - const size_t copy_nbytes = tensors_[output_ids_[i]].cur_nbytes; - if (!plan.copy_outputs[i] || copy_nbytes == 0) { + const size_t logical_nbytes = tensors_[output_ids_[i]].cur_nbytes; + if (!plan.copy_outputs[i] || logical_nbytes == 0) { continue; } + const size_t copy_nbytes = storage_buffer_size(logical_nbytes); const auto& copy = output_copies_[i]; wgpuCommandEncoderCopyBufferToBuffer( encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy_nbytes); @@ -1715,12 +1739,13 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { return 1; } - // GPU timestamp queries assume one submit; chunked execute is multi-submit. +#ifdef WGPU_BACKEND_ENABLE_PROFILING if (should_timestamp_query()) { throw std::runtime_error( "WebGPU: WEBGPU_TIMESTAMP_QUERY is incompatible with chunked execute " "(multi-submit); disable chunking to use GPU timestamp queries"); } +#endif // WGPU_BACKEND_ENABLE_PROFILING for (size_t chunk_index = 0; chunk_index < plan.dispatch_chunks.size(); chunk_index++) { @@ -1759,10 +1784,11 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { if (chunk_index + 1 == plan.dispatch_chunks.size()) { for (size_t i = 0; i < output_copies_.size(); i++) { - const size_t copy_nbytes = tensors_[output_ids_[i]].cur_nbytes; - if (!plan.copy_outputs[i] || copy_nbytes == 0) { + const size_t logical_nbytes = tensors_[output_ids_[i]].cur_nbytes; + if (!plan.copy_outputs[i] || logical_nbytes == 0) { continue; } + const size_t copy_nbytes = storage_buffer_size(logical_nbytes); const auto& copy = output_copies_[i]; wgpuCommandEncoderCopyBufferToBuffer( encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy_nbytes); @@ -1813,13 +1839,13 @@ void WebGPUGraph::copy_outputs( continue; } const auto& tensor = tensors_[output_ids_[i]]; - const size_t map_nbytes = tensor.cur_nbytes; - if (map_nbytes == 0) { + const size_t logical_nbytes = tensor.cur_nbytes; + if (logical_nbytes == 0) { continue; } const size_t dst_nbytes = outputs[i].nbytes; const bool is_double_width = - dst_nbytes % 2 == 0 && dst_nbytes / 2 == map_nbytes; + dst_nbytes % 2 == 0 && dst_nbytes / 2 == logical_nbytes; const bool widen_fp16 = is_double_width && !tensor.is_int && tensor.elem_size == 2; const bool widen_int32 = @@ -1832,7 +1858,7 @@ void WebGPUGraph::copy_outputs( if (outputs[i].host_is_fp32 && buffer_is_fp16 && !widen_fp16) { throw std::runtime_error("WebGPU: fp16 output buffer size mismatch"); } - if (dst_nbytes != map_nbytes && !widen_fp16 && !widen_int32) { + if (dst_nbytes != logical_nbytes && !widen_fp16 && !widen_int32) { throw std::runtime_error("WebGPU: output buffer size mismatch"); } } @@ -1842,13 +1868,14 @@ void WebGPUGraph::copy_outputs( continue; } const auto& tensor = tensors_[output_ids_[i]]; - const size_t map_nbytes = tensor.cur_nbytes; - if (map_nbytes == 0) { + const size_t logical_nbytes = tensor.cur_nbytes; + if (logical_nbytes == 0) { continue; } + const size_t map_nbytes = storage_buffer_size(logical_nbytes); const size_t dst_nbytes = outputs[i].nbytes; const bool is_double_width = - dst_nbytes % 2 == 0 && dst_nbytes / 2 == map_nbytes; + dst_nbytes % 2 == 0 && dst_nbytes / 2 == logical_nbytes; const bool widen_fp16 = is_double_width && !tensor.is_int && tensor.elem_size == 2; const bool widen_int32 = @@ -1887,7 +1914,7 @@ void WebGPUGraph::copy_outputs( const auto* src = static_cast(mapped); auto* dst = static_cast(outputs[i].data); - const size_t n = map_nbytes / sizeof(*src); + const size_t n = logical_nbytes / sizeof(*src); for (size_t k = 0; k < n; k++) { dst[k] = static_cast(src[k]); } @@ -1895,12 +1922,12 @@ void WebGPUGraph::copy_outputs( // int64 host output backed by an int32 GPU buffer: widen (sign-extend). const int32_t* src = static_cast(mapped); int64_t* dst = static_cast(outputs[i].data); - const size_t n = map_nbytes / sizeof(int32_t); + const size_t n = logical_nbytes / sizeof(int32_t); for (size_t k = 0; k < n; k++) { dst[k] = static_cast(src[k]); } } else { - std::memcpy(outputs[i].data, mapped, map_nbytes); + std::memcpy(outputs[i].data, mapped, logical_nbytes); } wgpuBufferUnmap(output_staging_buffers_[i]); } diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index a634ebb4172..23ce9df03ed 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -39,6 +39,8 @@ struct WebGPUTensor { // Serialized (GPU-side) element type, used to narrow wider host inputs. size_t elem_size = 0; bool is_int = false; + // Exact BOOL tag for byte-packed WGSL storage. + bool is_bool = false; // Exactly int8 (not uint8/bool), so int8-only ops can guard their dtype. bool is_int8 = false; }; diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index 5374390722b..480944ea93d 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -38,6 +38,7 @@ #include #include #include +#include #include #include #include @@ -124,6 +125,7 @@ #include #include #include +#include #include #include #include @@ -151,7 +153,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -362,6 +364,13 @@ constexpr std::array kShaderRegistry = {{ kConstantPadNdWorkgroupSizeY, kConstantPadNdWorkgroupSizeZ, }, + { + "conv1d", + kConv1dWGSL, + kConv1dWorkgroupSizeX, + kConv1dWorkgroupSizeY, + kConv1dWorkgroupSizeZ, + }, { "conv1d_dw", kConv1dDwWGSL, @@ -1034,6 +1043,13 @@ constexpr std::array kShaderRegistry = {{ kTanhWorkgroupSizeY, kTanhWorkgroupSizeZ, }, + { + "to_copy_bool_to_float", + kToCopyBoolToFloatWGSL, + kToCopyBoolToFloatWorkgroupSizeX, + kToCopyBoolToFloatWorkgroupSizeY, + kToCopyBoolToFloatWorkgroupSizeZ, + }, { "to_copy_float_to_int", kToCopyFloatToIntWGSL, diff --git a/backends/webgpu/runtime/ops/compare/Compare.cpp b/backends/webgpu/runtime/ops/compare/Compare.cpp index 99200cb2425..1870ee85213 100644 --- a/backends/webgpu/runtime/ops/compare/Compare.cpp +++ b/backends/webgpu/runtime/ops/compare/Compare.cpp @@ -58,13 +58,13 @@ void compare_impl( in2_tensor.elem_size != 4) { throw std::runtime_error("compare: fp32 inputs only"); } - if (!out_tensor.is_int || out_tensor.elem_size != 1) { + if (!out_tensor.is_bool || out_tensor.elem_size != 1) { throw std::runtime_error("compare: output must be a 1-byte bool tensor"); } const uint64_t numel = out_tensor.nbytes; - // out bool packed 4/word (array); numel%4==0 gates the readback map. - if (numel == 0u || numel % 4u != 0u || numel > UINT32_MAX) { - throw std::runtime_error("compare: numel must be a nonzero mult of 4"); + // out bool is byte-packed into ceil(numel / 4) u32 storage words. + if (numel == 0u || numel > UINT32_MAX) { + throw std::runtime_error("compare: numel must be nonzero and fit u32"); } const uint64_t in_numel = in1_tensor.nbytes / sizeof(float); if (in1_tensor.nbytes != in2_tensor.nbytes || in_numel != numel) { @@ -75,7 +75,7 @@ void compare_impl( params.num_elements = static_cast(numel); params.op = op; - const uint32_t words = static_cast(numel / 4u); + const uint32_t words = static_cast((numel + 3u) / 4u); uint32_t wg_size = utils::clamp_workgroup_size(device, kCompareWorkgroupSizeX); utils::WgCount workgroup_count = @@ -85,9 +85,7 @@ void compare_impl( wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(CompareParams)); - graph.add_uniform_buffer_bytes(sizeof(CompareParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); // out (rw storage) + in1/in2 (ro storage) + params (uniform). utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( @@ -97,7 +95,7 @@ void compare_impl( {0, WGPUBufferBindingType_Storage, out_tensor.buffer, - out_tensor.nbytes}, + static_cast(words) * sizeof(uint32_t)}, {1, WGPUBufferBindingType_ReadOnlyStorage, in1_tensor.buffer, @@ -127,9 +125,8 @@ void compare_impl( WebGPUGraph& g) { const auto& d = g.cur_dims(in1_id); const uint64_t n = utils::numel_of(d); - if (n == 0u || n % 4u != 0u || n > UINT32_MAX || - utils::numel_of(g.cur_dims(in2_id)) != n) { - throw std::runtime_error("compare(resize): numel must be a mult of 4"); + if (n == 0u || n > UINT32_MAX || utils::numel_of(g.cur_dims(in2_id)) != n) { + throw std::runtime_error("compare(resize): invalid numel"); } g.set_cur_dims(out_id, d); CompareParams p = {}; @@ -137,14 +134,12 @@ void compare_impl( p.op = op; wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); const utils::WgCount wgc = utils::compute_2d_workgroup_count( - g.device(), static_cast(n / 4u), wg_size, "compare"); + g.device(), static_cast((n + 3u) / 4u), wg_size, "compare"); g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }; graph.add_tensor_resize_hook(in1_id, resize); graph.add_tensor_resize_hook(in2_id, resize); - - graph.own_uniform_buffer(uniform_buffer); } void eq_op(WebGPUGraph& graph, const std::vector& args) { diff --git a/backends/webgpu/runtime/ops/compare/compare.wgsl b/backends/webgpu/runtime/ops/compare/compare.wgsl index a16568e5533..a650a6298b5 100644 --- a/backends/webgpu/runtime/ops/compare/compare.wgsl +++ b/backends/webgpu/runtime/ops/compare/compare.wgsl @@ -16,27 +16,29 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // One thread per output word = 4 bool bytes; num_elements%4==0 (host). + // One thread per output word = up to 4 bool bytes. let widx = gid.x + gid.y * (num_workgroups.x * wg_size); - let words = (params.num_elements + 3u) / 4u; + let words = (params.num_elements - 1u) / 4u + 1u; if (widx >= words) { return; } var packed: u32 = 0u; for (var j: u32 = 0u; j < 4u; j = j + 1u) { let i = widx * 4u + j; - let a = input1[i]; - let b = input2[i]; - var r: bool; - switch params.op { - case 0u: { r = a == b; } // eq - case 1u: { r = a < b; } // lt - case 2u: { r = a <= b; } // le - case 3u: { r = a > b; } // gt - default: { r = a >= b; } // ge - } - if (r) { - packed = packed | (1u << (j * 8u)); + if (i < params.num_elements) { + let a = input1[i]; + let b = input2[i]; + var r: bool; + switch params.op { + case 0u: { r = a == b; } // eq + case 1u: { r = a < b; } // lt + case 2u: { r = a <= b; } // le + case 3u: { r = a > b; } // gt + default: { r = a >= b; } // ge + } + if (r) { + packed = packed | (1u << (j * 8u)); + } } } t_out[widx] = packed; diff --git a/backends/webgpu/runtime/ops/compare/compare_wgsl.h b/backends/webgpu/runtime/ops/compare/compare_wgsl.h index 672c99b62d8..c1c4ac23e4f 100644 --- a/backends/webgpu/runtime/ops/compare/compare_wgsl.h +++ b/backends/webgpu/runtime/ops/compare/compare_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from compare.wgsl - DO NOT EDIT. -// wgsl-sha256: 241e7e6762b1eded07d28a3767936c970f509f6591e7cbf599d0b1eb61efb181 +// wgsl-sha256: 8f330b5a1e29a1fb8135e64600dadb1dc64c98f8f5371357f8de73a1808b76d6 inline constexpr const char* kCompareWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var input1: array; @@ -33,27 +33,29 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // One thread per output word = 4 bool bytes; num_elements%4==0 (host). + // One thread per output word = up to 4 bool bytes. let widx = gid.x + gid.y * (num_workgroups.x * wg_size); - let words = (params.num_elements + 3u) / 4u; + let words = (params.num_elements - 1u) / 4u + 1u; if (widx >= words) { return; } var packed: u32 = 0u; for (var j: u32 = 0u; j < 4u; j = j + 1u) { let i = widx * 4u + j; - let a = input1[i]; - let b = input2[i]; - var r: bool; - switch params.op { - case 0u: { r = a == b; } // eq - case 1u: { r = a < b; } // lt - case 2u: { r = a <= b; } // le - case 3u: { r = a > b; } // gt - default: { r = a >= b; } // ge - } - if (r) { - packed = packed | (1u << (j * 8u)); + if (i < params.num_elements) { + let a = input1[i]; + let b = input2[i]; + var r: bool; + switch params.op { + case 0u: { r = a == b; } // eq + case 1u: { r = a < b; } // lt + case 2u: { r = a <= b; } // le + case 3u: { r = a > b; } // gt + default: { r = a >= b; } // ge + } + if (r) { + packed = packed | (1u << (j * 8u)); + } } } t_out[widx] = packed; diff --git a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp index ae86ccffba8..64d8c2d2380 100644 --- a/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp +++ b/backends/webgpu/runtime/ops/conv1d_dw/Conv1dDW.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include @@ -49,8 +50,24 @@ uint32_t conv1d_out_len( int64_t stride, int64_t padding, int64_t dilation) { - return static_cast( - (in_len + 2 * padding - dilation * (k - 1) - 1) / stride + 1); + if (in_len <= 0 || k <= 0 || stride <= 0 || padding < 0 || dilation <= 0) { + throw std::runtime_error("conv1d: invalid geometry parameter"); + } + constexpr int64_t kMaxShaderIndex = std::numeric_limits::max(); + if (in_len > kMaxShaderIndex || k > kMaxShaderIndex || + stride > kMaxShaderIndex || padding > kMaxShaderIndex || + dilation > kMaxShaderIndex) { + throw std::runtime_error("conv1d: geometry parameter exceeds i32"); + } + const int64_t numerator = in_len + 2 * padding - dilation * (k - 1) - 1; + if (numerator < 0) { + throw std::runtime_error("conv1d: kernel exceeds padded input"); + } + const int64_t out_len = numerator / stride + 1; + if (static_cast(out_len) > UINT32_MAX) { + throw std::runtime_error("conv1d: output length exceeds u32"); + } + return static_cast(out_len); } int64_t first_int(const std::vector& v) { @@ -71,6 +88,24 @@ static_assert( sizeof(Conv1dPwParams) == 32, "Conv1dPwParams must match the WGSL Params struct (32 bytes)"); +struct Conv1dParams { + uint32_t in_channels; + uint32_t out_channels; + uint32_t in_len; + uint32_t out_len; + uint32_t kernel_size; + uint32_t stride; + uint32_t padding; + uint32_t dilation; + uint32_t numel; + uint32_t has_bias; +}; +static_assert( + sizeof(Conv1dParams) == 40, + "Conv1dParams must match the WGSL Params struct (40 bytes)"); +constexpr uint64_t kMaxConv1dDispatchElements = + static_cast(std::numeric_limits::max()); + // Pointwise conv1d (K=1, groups=1): a per-position matmul over channels. void add_conv1d_pw_node( WebGPUGraph& graph, @@ -203,6 +238,154 @@ void add_conv1d_pw_node( graph.own_uniform_buffer(params_buf); } +// General groups=1 conv1d. Voxtral uses K=3 with stride 1 then 2. +void add_conv1d_node( + WebGPUGraph& graph, + int in_id, + int weight_id, + int bias_id, + int out_id, + uint32_t stride, + uint32_t padding, + uint32_t dilation) { + WGPUDevice device = graph.device(); + const auto& in = graph.get_tensor(in_id); + const auto& weight = graph.get_tensor(weight_id); + const auto& out = graph.get_tensor(out_id); + const bool has_bias = + graph.get_value_type(bias_id) == WebGPUGraph::ValueType::Tensor; + if (!utils::is_fp32_tensor(in) || !utils::is_fp32_tensor(weight) || + !utils::is_fp32_tensor(out)) { + throw std::runtime_error("conv1d: input, weight, and output must be fp32"); + } + + const uint32_t expected_out_len = conv1d_out_len( + in.dims.at(2), weight.dims.at(2), stride, padding, dilation); + const uint32_t batch = static_cast(in.dims.at(0)); + const uint32_t in_channels = static_cast(in.dims.at(1)); + const uint32_t in_len = static_cast(in.dims.at(2)); + const uint32_t out_channels = static_cast(out.dims.at(1)); + const uint32_t out_len = static_cast(out.dims.at(2)); + const uint32_t kernel_size = static_cast(weight.dims.at(2)); + if (out.dims.at(0) != in.dims.at(0) || out_len != expected_out_len || + weight.dims.at(0) != out.dims.at(1) || + weight.dims.at(1) != in.dims.at(1)) { + throw std::runtime_error("conv1d: shape mismatch"); + } + + const uint64_t in_numel = utils::check_fp32(in, "conv1d", "input"); + const uint64_t out_numel = utils::check_fp32(out, "conv1d", "output"); + const uint64_t weight_numel = utils::check_fp32(weight, "conv1d", "weight"); + if (in_numel != static_cast(batch) * in_channels * in_len || + out_numel != static_cast(batch) * out_channels * out_len || + weight_numel != + static_cast(out_channels) * in_channels * kernel_size || + in_numel > UINT32_MAX || weight_numel > UINT32_MAX || + out_numel > kMaxConv1dDispatchElements) { + throw std::runtime_error("conv1d: fp32 byte-size or u32 mismatch"); + } + if (has_bias) { + const auto& bias = graph.get_tensor(bias_id); + if (!utils::is_fp32_tensor(bias) || bias.dims.size() != 1 || + bias.dims.at(0) != out.dims.at(1) || + utils::check_fp32(bias, "conv1d", "bias") != out_channels) { + throw std::runtime_error("conv1d: bias shape mismatch"); + } + } + + Conv1dParams params = {}; + params.in_channels = in_channels; + params.out_channels = out_channels; + params.in_len = in_len; + params.out_len = out_len; + params.kernel_size = kernel_size; + params.stride = stride; + params.padding = padding; + params.dilation = dilation; + params.numel = static_cast(out_numel); + params.has_bias = has_bias ? 1u : 0u; + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kConv1dWorkgroupSizeX); + const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( + device, params.numel, wg_size, "conv1d"); + WGPUConstantEntry wg_size_constant = utils::make_wg_size_constant(wg_size); + WGPUBuffer params_buf = graph.create_params_buffer(params); + WGPUBuffer bias_buf = + has_bias ? graph.get_tensor(bias_id).buffer : weight.buffer; + const uint64_t bias_size = + has_bias ? graph.get_tensor(bias_id).nbytes : weight.nbytes; + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kConv1dWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, bias_buf, bias_size}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(Conv1dParams)}, + }, + &wg_size_constant, + 1); + const size_t dispatch_idx = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "conv1d", + workgroup_count.y}); + + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + has_bias, + wg_size, + dispatch_idx, + params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + if (dims.size() != 3 || dims[0] <= 0 || dims[1] <= 0 || dims[2] <= 0 || + dims[1] != static_cast(in_channels)) { + throw std::runtime_error("conv1d(resize): input shape changed"); + } + Conv1dParams p = {}; + p.in_channels = in_channels; + p.out_channels = out_channels; + p.in_len = static_cast(dims[2]); + p.out_len = + conv1d_out_len(dims[2], kernel_size, stride, padding, dilation); + p.kernel_size = kernel_size; + p.stride = stride; + p.padding = padding; + p.dilation = dilation; + const uint64_t input_numel = utils::numel(dims); + const uint64_t numel = utils::numel( + {dims[0], out_channels, static_cast(p.out_len)}); + if (input_numel > UINT32_MAX || numel > kMaxConv1dDispatchElements) { + throw std::runtime_error( + "conv1d(resize): tensor numel exceeds shader index range"); + } + p.numel = static_cast(numel); + p.has_bias = has_bias ? 1u : 0u; + const utils::WgCount wgc = utils::compute_2d_workgroup_count( + g.device(), p.numel, wg_size, "conv1d(resize)"); + g.set_cur_dims( + out_id, {dims[0], out_channels, static_cast(p.out_len)}); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; + }); +} + // depthwise-conv1d (groups==C); mirrors Vulkan conv1d_dw (Convolution.cpp:755). void convolution_impl(WebGPUGraph& graph, const std::vector& args) { // args mirror Vulkan conv1d_dw; bias (arg 2) may be Null; out=args.back(). @@ -242,38 +425,44 @@ void convolution_impl(WebGPUGraph& graph, const std::vector& args) { const bool transposed = graph.get_bool(transposed_id); const int64_t groups = graph.get_int(groups_id); - // Pointwise (K=1, groups=1): a matmul over channels; stride-1 / no-pad only. - if (!transposed && groups == 1 && weight_tensor.dims.at(2) == 1 && - first_int(graph.get_int_list(stride_id)) == 1 && - first_int(graph.get_int_list(padding_id)) == 0) { - add_conv1d_pw_node(graph, in_id, weight_id, bias_id, out_id); - return; - } - - // Otherwise only the depthwise config (groups==C, weight [C,1,K]). - if (transposed || groups != static_cast(channels) || - weight_tensor.dims.at(0) != static_cast(channels) || - weight_tensor.dims.at(1) != 1) { - throw std::runtime_error( - "convolution: only depthwise or pointwise conv1d supported"); - } - const int64_t stride_i = first_int(graph.get_int_list(stride_id)); const int64_t padding_i = first_int(graph.get_int_list(padding_id)); const int64_t dilation_i = first_int(graph.get_int_list(dilation_id)); - if (stride_i < 1) { - throw std::runtime_error("convolution: stride must be >= 1"); + if (stride_i < 1 || stride_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: stride must fit positive i32"); } - if (padding_i < 0) { - throw std::runtime_error("convolution: padding must be >= 0"); + if (padding_i < 0 || padding_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: padding must fit nonnegative i32"); } - if (dilation_i < 1) { - throw std::runtime_error("convolution: dilation must be >= 1"); + if (dilation_i < 1 || dilation_i > std::numeric_limits::max()) { + throw std::runtime_error("convolution: dilation must fit positive i32"); } const uint32_t stride = static_cast(stride_i); const uint32_t padding = static_cast(padding_i); const uint32_t dilation = static_cast(dilation_i); + // Pointwise (K=1, groups=1): a matmul over channels; stride-1 / no-pad only. + if (!transposed && groups == 1 && weight_tensor.dims.at(2) == 1 && + stride_i == 1 && padding_i == 0) { + add_conv1d_pw_node(graph, in_id, weight_id, bias_id, out_id); + return; + } + + const bool is_depthwise = !transposed && + groups == static_cast(channels) && + weight_tensor.dims.at(0) == static_cast(channels) && + weight_tensor.dims.at(1) == 1; + if (!is_depthwise && !transposed && groups == 1) { + add_conv1d_node( + graph, in_id, weight_id, bias_id, out_id, stride, padding, dilation); + return; + } + + if (!is_depthwise) { + throw std::runtime_error( + "convolution: only depthwise, pointwise, or groups=1 conv1d supported"); + } + uint64_t out_numel = 1; for (int64_t d : out_tensor.dims) { out_numel *= static_cast(d); diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl b/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl new file mode 100644 index 00000000000..c51594d2005 --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d.wgsl @@ -0,0 +1,53 @@ +override wg_size: u32 = 64u; + +struct Params { + in_channels: u32, + out_channels: u32, + in_len: u32, + out_len: u32, + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, + numel: u32, + has_bias: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + + let out_t = idx % params.out_len; + let out_c = (idx / params.out_len) % params.out_channels; + let batch = idx / (params.out_channels * params.out_len); + var sum = 0.0; + + for (var in_c = 0u; in_c < params.in_channels; in_c = in_c + 1u) { + for (var k = 0u; k < params.kernel_size; k = k + 1u) { + let in_t = i32(out_t * params.stride + k * params.dilation) - + i32(params.padding); + if (in_t >= 0 && in_t < i32(params.in_len)) { + let input_idx = + (batch * params.in_channels + in_c) * params.in_len + u32(in_t); + let weight_idx = + (out_c * params.in_channels + in_c) * params.kernel_size + k; + sum = fma(input[input_idx], weight[weight_idx], sum); + } + } + } + if (params.has_bias != 0u) { + sum = sum + bias[out_c]; + } + output[idx] = sum; +} diff --git a/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h new file mode 100644 index 00000000000..e29af80459b --- /dev/null +++ b/backends/webgpu/runtime/ops/conv1d_dw/conv1d_wgsl.h @@ -0,0 +1,77 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from conv1d.wgsl - DO NOT EDIT. +// wgsl-sha256: 7bc955b7f43473aab96222e7a1228973b65e38e2e6a19e2b9793e0ed7f3768d9 +inline constexpr const char* kConv1dWGSL = R"( +override wg_size: u32 = 64u; + +struct Params { + in_channels: u32, + out_channels: u32, + in_len: u32, + out_len: u32, + kernel_size: u32, + stride: u32, + padding: u32, + dilation: u32, + numel: u32, + has_bias: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + + let out_t = idx % params.out_len; + let out_c = (idx / params.out_len) % params.out_channels; + let batch = idx / (params.out_channels * params.out_len); + var sum = 0.0; + + for (var in_c = 0u; in_c < params.in_channels; in_c = in_c + 1u) { + for (var k = 0u; k < params.kernel_size; k = k + 1u) { + let in_t = i32(out_t * params.stride + k * params.dilation) - + i32(params.padding); + if (in_t >= 0 && in_t < i32(params.in_len)) { + let input_idx = + (batch * params.in_channels + in_c) * params.in_len + u32(in_t); + let weight_idx = + (out_c * params.in_channels + in_c) * params.kernel_size + k; + sum = fma(input[input_idx], weight[weight_idx], sum); + } + } + } + if (params.has_bias != 0u) { + sum = sum + bias[out_c]; + } + output[idx] = sum; +} +)"; + +inline constexpr uint32_t kConv1dWorkgroupSizeX = 64; +inline constexpr uint32_t kConv1dWorkgroupSizeY = 1; +inline constexpr uint32_t kConv1dWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp index 5bfc0fa3bf1..9f6a4e68ba6 100644 --- a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp +++ b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp @@ -14,6 +14,7 @@ #include +#include #include namespace executorch::backends::webgpu { @@ -34,6 +35,22 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { const auto& in_tensor = graph.get_tensor(in_id); const auto& out_tensor = graph.get_tensor(out_id); + if (graph.get_value_type(args.at(1)) != WebGPUGraph::ValueType::IntList) { + throw std::runtime_error( + "WebGPU expand_copy: dynamic target sizes are unsupported"); + } + for (int64_t target_size : graph.get_int_list(args.at(1))) { + if (target_size == -1) { + throw std::runtime_error( + "WebGPU expand_copy: inferred target sizes are unsupported"); + } + } + if (graph.tensor_has_dynamic_dims(in_id) || + graph.tensor_has_dynamic_dims(out_id)) { + throw std::runtime_error( + "WebGPU expand_copy: dynamic shapes are unsupported"); + } + TensorMeta out_meta; TensorMeta in_meta; fill_tensor_meta(out_tensor, &out_meta); @@ -44,21 +61,24 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error( "expand_copy: non-fp32 operand (nbytes != numel*4)"); } + if (out_meta.numel > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU expand_copy: element count exceeds the flattened 2D dispatch " + "limit"); + } uint32_t wg_size = utils::clamp_workgroup_size(device, kExpandCopyWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, out_meta.numel, wg_size, "expand_copy"); WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUBuffer out_meta_buf = - utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); - WGPUBuffer in_meta_buf = - utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); - graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta)); + WGPUBuffer out_meta_buf = graph.create_params_buffer(out_meta); + WGPUBuffer in_meta_buf = graph.create_params_buffer(in_meta); utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, @@ -78,10 +98,8 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { &wg_size_constant, 1); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); - - wgpuBufferRelease(out_meta_buf); - wgpuBufferRelease(in_meta_buf); + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y); } } // namespace diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl index 053311a69f4..fab4df15a90 100644 --- a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl @@ -13,8 +13,10 @@ struct TensorMeta { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h index 83c5881604f..f1449f61793 100644 --- a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h +++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from expand_copy.wgsl - DO NOT EDIT. -// wgsl-sha256: 99953670bea89e42bc9c689ab80addfd9a331442c8c8f1a5b0c39dbe11c19370 +// wgsl-sha256: b3c032ab961ffde245fc44289b67df3b5e4ca93eedb9ada2f20a3eaa6f10e9c6 inline constexpr const char* kExpandCopyWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -30,8 +30,10 @@ struct TensorMeta { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let idx = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); if (idx >= out_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/gelu/Gelu.cpp b/backends/webgpu/runtime/ops/gelu/Gelu.cpp index 943012515ff..023ed7b5d88 100644 --- a/backends/webgpu/runtime/ops/gelu/Gelu.cpp +++ b/backends/webgpu/runtime/ops/gelu/Gelu.cpp @@ -13,6 +13,7 @@ #include +#include #include #include #include @@ -42,13 +43,18 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { const auto& out_tensor = graph.get_tensor(out_id); utils::check_elementwise_fp32_io(in_tensor, out_tensor, "gelu"); - uint32_t num_elements = - static_cast(out_tensor.nbytes / sizeof(float)); + const uint64_t num_elements64 = out_tensor.nbytes / sizeof(float); + if (num_elements64 > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU gelu: element count exceeds the flattened 2D dispatch limit"); + } + const uint32_t num_elements = static_cast(num_elements64); // Each thread handles up to 4 elements (vec4 body + scalar-tail idiom). uint32_t num_vec4_threads = utils::div_up(num_elements, 4u); uint32_t wg_size = utils::clamp_workgroup_size(device, kGeluWorkgroupSizeX); - uint32_t workgroup_count = utils::compute_1d_workgroup_count( + utils::WgCount workgroup_count = utils::compute_2d_workgroup_count( device, num_vec4_threads, wg_size, "gelu"); WGPUConstantEntry wg_constant = utils::make_wg_size_constant(wg_size); @@ -56,9 +62,7 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { GeluParams params = {}; params.num_elements = num_elements; - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(GeluParams)); - graph.add_uniform_buffer_bytes(sizeof(GeluParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); // input (read storage) + output (storage) + params. The exact/approximate // choice is baked into the compiled pipeline via the entry point (mirrors @@ -85,10 +89,38 @@ void gelu_impl(WebGPUGraph& graph, const std::vector& args) { 1, exact ? "main_erf" : "main_tanh"); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); - - // Drop our ref; the bind group keeps the uniform buffer alive until release. - wgpuBufferRelease(uniform_buffer); + const size_t dispatch_idx = graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y); + + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + const uint64_t num_elements64 = utils::numel_of(dims); + if (num_elements64 > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + "WebGPU gelu(resize): element count exceeds the flattened 2D " + "dispatch limit"); + } + const uint32_t num_elements = static_cast(num_elements64); + g.set_cur_dims(out_id, dims); + + GeluParams params = {}; + params.num_elements = num_elements; + wgpuQueueWriteBuffer( + g.queue(), params_buf, 0, ¶ms, sizeof(GeluParams)); + + const uint32_t num_vec4_threads = utils::div_up(num_elements, 4u); + const utils::WgCount resized_workgroup_count = + utils::compute_2d_workgroup_count( + g.device(), num_vec4_threads, wg_size, "gelu(resize)"); + g.dispatch_at(dispatch_idx).workgroup_count_x = + resized_workgroup_count.x; + g.dispatch_at(dispatch_idx).workgroup_count_y = + resized_workgroup_count.y; + }); } } // namespace diff --git a/backends/webgpu/runtime/ops/gelu/gelu.wgsl b/backends/webgpu/runtime/ops/gelu/gelu.wgsl index 4f7eb68bc96..9583ef81551 100644 --- a/backends/webgpu/runtime/ops/gelu/gelu.wgsl +++ b/backends/webgpu/runtime/ops/gelu/gelu.wgsl @@ -33,8 +33,11 @@ fn gelu_erf4(x: vec4) -> vec4 { // before use), computes GELU as one vec4 op, then scatters back only the // in-bounds lanes. @compute @workgroup_size(wg_size, 1, 1) -fn main_tanh(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_tanh( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } @@ -49,8 +52,11 @@ fn main_tanh(@builtin(global_invocation_id) gid: vec3) { } @compute @workgroup_size(wg_size, 1, 1) -fn main_erf(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_erf( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h index 6da12e229af..f8af0f8d2c3 100644 --- a/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h +++ b/backends/webgpu/runtime/ops/gelu/gelu_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from gelu.wgsl - DO NOT EDIT. -// wgsl-sha256: 18f4a82d3bad1ef8703397b871c708804140c4cb382451661f7a77367ac2425f +// wgsl-sha256: 96570753688590fa009ee5503f754cf3eb572dcb3dcae6818220fe06fe3139ee inline constexpr const char* kGeluWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -50,8 +50,11 @@ fn gelu_erf4(x: vec4) -> vec4 { // before use), computes GELU as one vec4 op, then scatters back only the // in-bounds lanes. @compute @workgroup_size(wg_size, 1, 1) -fn main_tanh(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_tanh( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } @@ -66,8 +69,11 @@ fn main_tanh(@builtin(global_invocation_id) gid: vec3) { } @compute @workgroup_size(wg_size, 1, 1) -fn main_erf(@builtin(global_invocation_id) gid: vec3) { - let base = gid.x * 4u; +fn main_erf( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let thread_idx = gid.x + gid.y * (num_workgroups.x * wg_size); + let base = thread_idx * 4u; if (base >= params.num_elements) { return; } diff --git a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp index 05f7ff6fe24..54d80016296 100644 --- a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp +++ b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp @@ -10,12 +10,14 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -68,9 +70,7 @@ void add_convert_op( ConvertParams params = {}; params.num_elements = num_elements; - WGPUBuffer uniform_buffer = - utils::make_uniform(device, ¶ms, sizeof(ConvertParams)); - graph.add_uniform_buffer_bytes(sizeof(ConvertParams)); + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, @@ -113,8 +113,185 @@ void add_convert_op( wg_size, "to_copy(resize)"); }); +} + +// Decode byte-packed bool storage into numeric fp32 values. +void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("to_copy_bool_to_float: null buffer binding"); + } + if (!in_tensor.is_bool || in_tensor.elem_size != 1 || out_tensor.is_int || + out_tensor.elem_size != sizeof(float) || + out_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes / sizeof(float) != in_tensor.nbytes) { + throw std::runtime_error("to_copy_bool_to_float: dtype/numel mismatch"); + } + if (in_tensor.nbytes == 0u || + in_tensor.nbytes > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: numel must be nonzero and fit u32"); + } + + const uint32_t num_elements = static_cast(in_tensor.nbytes); + const uint64_t input_bind_size_u64 = + (static_cast(in_tensor.nbytes) + 3u) & ~uint64_t(3); + if (input_bind_size_u64 > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: input binding size overflows"); + } + const size_t input_bind_size = static_cast(input_bind_size_u64); + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kToCopyBoolToFloatWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, num_elements, wg_size, "to_copy_bool_to_float"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + ConvertParams params = {}; + params.num_elements = num_elements; + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kToCopyBoolToFloatWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + input_bind_size}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ConvertParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + const uint64_t numel = utils::numel_of(dims); + if (numel == 0u || numel > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float(resize): invalid numel"); + } + g.set_cur_dims(out_id, dims); + ConvertParams p = {}; + p.num_elements = static_cast(numel); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), + static_cast(numel), + wg_size, + "to_copy_bool_to_float(resize)"); + }); +} + +// Decode byte-packed bool storage into numeric fp32 values. +void add_bool_to_float_op(WebGPUGraph& graph, int in_id, int out_id) { + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.buffer == nullptr || out_tensor.buffer == nullptr) { + throw std::runtime_error("to_copy_bool_to_float: null buffer binding"); + } + if (!in_tensor.is_bool || in_tensor.elem_size != 1 || out_tensor.is_int || + out_tensor.elem_size != sizeof(float) || + out_tensor.nbytes % sizeof(float) != 0 || + out_tensor.nbytes / sizeof(float) != in_tensor.nbytes) { + throw std::runtime_error("to_copy_bool_to_float: dtype/numel mismatch"); + } + if (in_tensor.nbytes == 0u || + in_tensor.nbytes > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: numel must be nonzero and fit u32"); + } + + const uint32_t num_elements = static_cast(in_tensor.nbytes); + const uint64_t input_bind_size_u64 = + (static_cast(in_tensor.nbytes) + 3u) & ~uint64_t(3); + if (input_bind_size_u64 > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float: input binding size overflows"); + } + const size_t input_bind_size = static_cast(input_bind_size_u64); + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kToCopyBoolToFloatWorkgroupSizeX); + const uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, num_elements, wg_size, "to_copy_bool_to_float"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + ConvertParams params = {}; + params.num_elements = num_elements; + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(ConvertParams)); + graph.add_uniform_buffer_bytes(sizeof(ConvertParams)); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kToCopyBoolToFloatWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + input_bind_size}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ConvertParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + + WGPUBuffer params_buf = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, out_id, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& dims = g.cur_dims(in_id); + const uint64_t numel = utils::numel_of(dims); + if (numel == 0u || numel > std::numeric_limits::max()) { + throw std::runtime_error( + "to_copy_bool_to_float(resize): invalid numel"); + } + g.set_cur_dims(out_id, dims); + ConvertParams p = {}; + p.num_elements = static_cast(numel); + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + utils::compute_1d_workgroup_count( + g.device(), + static_cast(numel), + wg_size, + "to_copy_bool_to_float(resize)"); + }); - // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } @@ -129,6 +306,12 @@ void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id) { const auto& in_tensor = graph.get_tensor(in_id); const auto& out_tensor = graph.get_tensor(out_id); + if (in_tensor.is_bool != out_tensor.is_bool && in_tensor.is_int && + out_tensor.is_int) { + throw std::runtime_error( + "WebGPU to_copy: bool and integer conversions are unsupported"); + } + // Same is_int+width = flat byte copy; unique dtype key in the 32-bit domain. if (in_tensor.is_int == out_tensor.is_int && in_tensor.elem_size == out_tensor.elem_size) { @@ -137,7 +320,10 @@ void add_to_copy_node(WebGPUGraph& graph, int in_id, int out_id) { } // int<->float = numeric convert (mirrors Vulkan add_view_copy_convert_node). - if (in_tensor.is_int && !out_tensor.is_int) { + if (in_tensor.is_bool && !out_tensor.is_int && + out_tensor.elem_size == sizeof(float)) { + add_bool_to_float_op(graph, in_id, out_id); + } else if (in_tensor.is_int && !out_tensor.is_int) { add_convert_op( graph, in_id, diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl new file mode 100644 index 00000000000..239730de65d --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float.wgsl @@ -0,0 +1,24 @@ +override wg_size: u32 = 256u; + +struct Params { + num_elements: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let word = input[idx / 4u]; + let byte_shift = (idx % 4u) * 8u; + let value = (word >> byte_shift) & 0xffu; + output[idx] = select(0.0, 1.0, value != 0u); +} diff --git a/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h new file mode 100644 index 00000000000..ef7e40976f3 --- /dev/null +++ b/backends/webgpu/runtime/ops/to_copy/to_copy_bool_to_float_wgsl.h @@ -0,0 +1,48 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from to_copy_bool_to_float.wgsl - DO NOT EDIT. +// wgsl-sha256: 29fd43b2f638489e9b8d72b2cc9140d07174c750cbdc291e63793319a2fa5961 +inline constexpr const char* kToCopyBoolToFloatWGSL = R"( +override wg_size: u32 = 256u; + +struct Params { + num_elements: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let word = input[idx / 4u]; + let byte_shift = (idx % 4u) * 8u; + let value = (word >> byte_shift) & 0xffu; + output[idx] = select(0.0, 1.0, value != 0u); +} +)"; + +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeX = 256; +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeY = 1; +inline constexpr uint32_t kToCopyBoolToFloatWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index 6ed7229604b..949ad3c0de6 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -153,42 +153,90 @@ void expect_dual_q4_topology(WebGPUGraph& graph) { 1); } -TEST(WebGPUShaderRegistry, FindsKnownShaderAndRejectsUnknownName) { - const WebGPUShaderInfo& sigmoid = get_webgpu_shader_info("sigmoid"); - EXPECT_EQ(sigmoid.name, "sigmoid"); - EXPECT_NE(sigmoid.source, nullptr); - EXPECT_GT(sigmoid.workgroup_size_x, 0u); - EXPECT_THROW( - get_webgpu_shader_info("not_a_registered_shader"), std::runtime_error); -} +struct Conv1dRouteCase { + const char* name; + std::vector input_dims; + std::vector weight_dims; + std::vector output_dims; + int64_t stride; + int64_t padding; + int64_t dilation; + int64_t groups; + const char* expected_kernel; +}; + +void build_conv1d_route_graph( + WebGPUGraph& graph, + const Conv1dRouteCase& test_case) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](const std::vector& dims, int mem_obj_id) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + /*constant_id=*/-1, + mem_obj_id) + .Union())); + return id; + }; + auto add_int = [&](int64_t value) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, value).Union())); + return id; + }; + auto add_int_list = [&](int64_t value) { + const int id = static_cast(values.size()); + const std::vector items = {value}; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::IntList, + vk::CreateIntListDirect(fbb, &items).Union())); + return id; + }; + + const int input = add_tensor(test_case.input_dims, 0); + const int weight = add_tensor(test_case.weight_dims, 1); + const int bias = static_cast(values.size()); + values.push_back(vk::CreateVkValue(fbb)); + const int stride = add_int_list(test_case.stride); + const int padding = add_int_list(test_case.padding); + const int dilation = add_int_list(test_case.dilation); + const int transposed = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Bool, vk::CreateBool(fbb, false).Union())); + const int output_padding = add_int_list(0); + const int groups = add_int(test_case.groups); + const int output = add_tensor(test_case.output_dims, 2); + const std::vector args = { + input, + weight, + bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + output}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten.convolution.default", &args)); + const std::vector input_ids = { + static_cast(input), static_cast(weight)}; + const std::vector output_ids = {static_cast(output)}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); -TEST(WebGPUQ4RouteSignal, PreservesStaticAndRecordsBothDynamicSignals) { - WebGPUGraph static_graph; - static_graph.set_device(g_device); - build_q4_route_graph(static_graph, Q4RouteSignal::Static); - const auto static_dispatches = q4_dispatches(static_graph); - ASSERT_EQ(static_dispatches.size(), 1); - EXPECT_EQ(static_graph.num_dispatches(), 1); - EXPECT_NE(static_dispatches[0]->pipeline, nullptr); - EXPECT_NE(static_dispatches[0]->bind_group, nullptr); - EXPECT_FALSE(static_graph.has_dynamic_shapes()); - EXPECT_FALSE(static_graph.config().record_q4gsw_decode_route); - - WebGPUGraph legacy_graph; - legacy_graph.set_device(g_device); - build_q4_route_graph(legacy_graph, Q4RouteSignal::LegacyGraphMarker); - EXPECT_TRUE(legacy_graph.has_dynamic_shapes()); - EXPECT_FALSE(legacy_graph.config().record_q4gsw_decode_route); - EXPECT_EQ(legacy_graph.num_dispatches(), 3); - expect_dual_q4_topology(legacy_graph); - - WebGPUGraph explicit_graph; - explicit_graph.set_device(g_device); - build_q4_route_graph(explicit_graph, Q4RouteSignal::ExplicitOption); - EXPECT_FALSE(explicit_graph.has_dynamic_shapes()); - EXPECT_TRUE(explicit_graph.config().record_q4gsw_decode_route); - EXPECT_EQ(explicit_graph.num_dispatches(), 2); - expect_dual_q4_topology(explicit_graph); + graph.set_device(g_device); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); } TEST(WebGPUComputeDispatch, PipelineKeyCanonicalizesConstants) { @@ -899,6 +947,58 @@ TEST(WebGPURopeValidation, RejectsMalformedGraphsBeforeDispatchAllocation) { } } +TEST(WebGPUToCopyValidation, RejectsBoolAndByteIntegerConversions) { + ASSERT_TRUE(webgpu_operator_registry().has_op("aten._to_copy.default")); + namespace vk = vkgraph; + struct TestCase { + const char* name; + vk::VkDataType input_dtype; + vk::VkDataType output_dtype; + }; + const TestCase cases[] = { + {"bool_to_int8", vk::VkDataType::BOOL, vk::VkDataType::INT8}, + {"bool_to_uint8", vk::VkDataType::BOOL, vk::VkDataType::UINT8}, + {"int8_to_bool", vk::VkDataType::INT8, vk::VkDataType::BOOL}, + {"uint8_to_bool", vk::VkDataType::UINT8, vk::VkDataType::BOOL}, + }; + for (const TestCase& test_case : cases) { + SCOPED_TRACE(test_case.name); + ::flatbuffers::FlatBufferBuilder fbb; + const std::vector dims = {4}; + std::vector<::flatbuffers::Offset> values; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect(fbb, test_case.input_dtype, &dims, -1, 0) + .Union())); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect(fbb, test_case.output_dtype, &dims, -1, 1) + .Union())); + const std::vector args = {0, 1}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten._to_copy.default", &args)); + const std::vector input_ids = {0}; + const std::vector output_ids = {1}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); + FAIL() << test_case.name << " unexpectedly built"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "WebGPU to_copy: bool and integer conversions are unsupported"); + } + EXPECT_EQ(graph.memory_stats().num_dispatches, 0); + } +} + TEST(WebGPUExecution, FullySuppressedPlanPerformsNoQueueSubmission) { WebGPUGraph graph; const WebGPUExecutionPlan plan; diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 7b3c1d6b0a7..cf16cd7e87c 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -18,6 +18,8 @@ // F dyn_rms_chain (rms(rms(x))) at 3 S -> golden (resize CASCADE, DD-4) // G rms+residual H rms*x I dyn_linear J sdpa_dyn K emb_dyn L rope_dyn // M dyn_sigmoid N dyn_select (select_copy(0,-1), dynamic S) +// O ONE dyn_conv1d graph reused across live input lengths +// P ONE dyn_gelu graph reused above -> at -> above the old 1D dispatch cap // .pte + goldens from test/ops/dynamic_shape/test_dynamic_shape_export.py. // // Artifacts dir: $WEBGPU_DYNAMIC_SHAPE_DIR, else argv[1], else @@ -33,6 +35,7 @@ #include #include +#include #include #include #include @@ -127,6 +130,74 @@ void check_s(Module& m, const std::string& prefix, int s) { << " golden.size=" << golden.size() << ")"; } +constexpr int kConv1dInChannels = 3; +constexpr int kConv1dOutChannels = 4; +constexpr int kConv1dKernel = 3; +constexpr int kConv1dStride = 2; +constexpr int kConv1dPadding = 1; +constexpr int kConv1dDilation = 2; + +void check_conv1d(Module& module, int length) { + const std::string prefix = g_dir + "/dyn_conv1d.S" + std::to_string(length); + auto input = read_bin(prefix + ".input.bin"); + auto golden = read_bin(prefix + ".golden.bin"); + ASSERT_EQ(input.size(), static_cast(kConv1dInChannels * length)); + ASSERT_FALSE(golden.empty()); + auto tensor = + make_tensor_ptr({1, kConv1dInChannels, length}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << "conv1d length=" << length << " forward failed"; + const auto& output = result.get()[0].toTensor(); + const int output_length = (length + 2 * kConv1dPadding - + kConv1dDilation * (kConv1dKernel - 1) - 1) / + kConv1dStride + + 1; + ASSERT_EQ(output.dim(), 3); + ASSERT_EQ(output.size(0), 1); + ASSERT_EQ(output.size(1), kConv1dOutChannels); + ASSERT_EQ(output.size(2), output_length); + const size_t numel = static_cast(kConv1dOutChannels * output_length); + ASSERT_EQ(static_cast(output.numel()), numel); + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + const float error = max_err(got, golden); + EXPECT_LT(error, 1e-3f) << "conv1d length=" << length << " max_err=" << error; +} + +constexpr int kGeluOld1dDispatchCap = 4 * 64 * 65535; +constexpr int kGelu2dDispatchBoundary = kGeluOld1dDispatchCap + 1; +constexpr int kGeluPatternSize = 257; + +void check_gelu_2d(Module& module, int elements) { + std::array golden = {}; + std::vector input(static_cast(elements)); + for (int i = 0; i < kGeluPatternSize; i++) { + const float value = -4.0f + 8.0f * i / (kGeluPatternSize - 1); + golden[i] = 0.5f * value * (1.0f + std::erf(value * 0.7071067811865476f)); + } + for (int i = 0; i < elements; i++) { + input[i] = -4.0f + 8.0f * (i % kGeluPatternSize) / (kGeluPatternSize - 1); + } + auto tensor = make_tensor_ptr({elements}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << "gelu elements=" << elements << " forward failed"; + const auto& output = result.get()[0].toTensor(); + ASSERT_EQ(output.dim(), 1); + ASSERT_EQ(output.size(0), elements); + ASSERT_EQ(output.numel(), elements); + const float* data = output.const_data_ptr(); + float error = 0.0f; + for (int i = 0; i < elements; i++) { + error = std::fmax(error, std::fabs(data[i] - golden[i % kGeluPatternSize])); + } + EXPECT_LT(error, 1e-4f) << "gelu elements=" << elements + << " max_err=" << error; +} + // Dynamic quantized linear: input [M, kLinK] -> output [M, n]. kLinN is the // register-tiled/bicol config; kLinNShmem (N>=2048) routes to the shmem GEMM. constexpr int kLinK = 64; @@ -893,6 +964,43 @@ TEST(DynamicShape, RmsNormReusedGraph) { } } +TEST(DynamicShape, Conv1dReusedGraph) { + Module module(g_dir + "/dyn_conv1d.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_conv1d.pte"; + for (int length : {16, 9, 5, 16}) { + check_conv1d(module, length); + } +} + +TEST(DynamicShape, GeluCrosses2dDispatchBoundary) { + if (std::getenv("WEBGPU_TEST_HEAVY") == nullptr) { + GTEST_SKIP() << "WEBGPU_TEST_HEAVY not set"; + } + Module module(g_dir + "/dyn_gelu_2d.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_gelu_2d.pte"; + for (int elements : + {kGelu2dDispatchBoundary, + kGeluOld1dDispatchCap, + kGelu2dDispatchBoundary}) { + check_gelu_2d(module, elements); + } +} + +TEST(DynamicShape, ExpandCopyRejectsDynamicShapesAtLoad) { + const std::string path = g_dir + "/dyn_expand_copy.pte"; + ASSERT_TRUE(std::ifstream(path).good()) << "missing dyn_expand_copy.pte"; + Module module(path); + EXPECT_NE(module.load_forward(), Error::Ok); +} + +TEST(DynamicShape, ExpandCopyRejectsInferredDynamicShapesAtLoad) { + const std::string path = g_dir + "/dyn_expand_copy_inferred.pte"; + ASSERT_TRUE(std::ifstream(path).good()) + << "missing dyn_expand_copy_inferred.pte"; + Module module(path); + EXPECT_NE(module.load_forward(), Error::Ok); +} + // C2: grow-only reuse — one loaded rms graph run smallest -> largest, so the // FIRST resize grows the dispatch (every other reuse test starts at MAXS and // only shrinks; this catches a hook with a shrink-only short-circuit). diff --git a/backends/webgpu/test/native/test_webgpu_utils.cpp b/backends/webgpu/test/native/test_webgpu_utils.cpp index edc0f315294..a839d224b16 100644 --- a/backends/webgpu/test/native/test_webgpu_utils.cpp +++ b/backends/webgpu/test/native/test_webgpu_utils.cpp @@ -14,8 +14,16 @@ #include +#include + using namespace executorch::backends::webgpu; +TEST(WebGPUUtils, DivUpDoesNotOverflowAtUint32Max) { + constexpr uint32_t kMax = std::numeric_limits::max(); + EXPECT_EQ(utils::div_up(kMax, 4u), 1073741824u); + EXPECT_EQ(utils::div_up(kMax, kMax), 1u); +} + TEST(WebGPUUtils, DispatchGridStaysOneDimUnderCeiling) { utils::DispatchGrid g = utils::compute_dispatch_grid_from_limits(1000u, 256u, 65535u, "test"); diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index bbd6e13dc9d..6a019ab8e9f 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -53,8 +53,16 @@ CompareModule, ) from executorch.backends.webgpu.test.ops.test_conv1d_dw import Conv1dDWModule -from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dPwModule +from executorch.backends.webgpu.test.ops.test_conv1d_pw import ( + Conv1dModule, + Conv1dPwModule, + GENERAL_CONFIGS as _CONV1D_CONFIGS, +) from executorch.backends.webgpu.test.ops.test_conv_with_clamp import ConvWithClampModule +from executorch.backends.webgpu.test.ops.test_expand_copy import ( + CONFIGS as _EXPAND_COPY_CONFIGS, + ExpandCopyModule, +) from executorch.backends.webgpu.test.ops.test_flip import FlipModule from executorch.backends.webgpu.test.ops.test_floor_divide import FloorDivideModule from executorch.backends.webgpu.test.ops.test_grid_priors import GridPriorsModule @@ -153,8 +161,13 @@ ) from executorch.backends.webgpu.test.ops.test_to_copy import ( + bool_tail_input, + compare_to_copy_input_a, + compare_to_copy_input_b, + CompareToCopyBoolToFloatModule, to_copy_float_input, to_copy_int_input, + ToCopyBoolToFloatModule, ToCopyFloatToIntToFloatModule, ToCopyIntToFloatModule, ) @@ -191,6 +204,39 @@ def _add_factory(variant: str = "regular") -> torch.nn.Module: }[variant]() +@register_op_test("to_copy_bool_to_float") +def _to_copy_bool_to_float_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=CompareToCopyBoolToFloatModule, + cases=[ + Case( + inputs=( + InputSpec((n,), gen=compare_to_copy_input_a), + InputSpec((n,), gen=compare_to_copy_input_b), + ), + name=f"length_{n}", + ) + for n in (1, 4, 5, 67) + ], + golden_dtype="float32", + ) + + +@register_op_test("to_copy_bool_input_to_float") +def _to_copy_bool_input_to_float_suite() -> WebGPUTestSuite: + return WebGPUTestSuite( + module_factory=ToCopyBoolToFloatModule, + cases=[ + Case( + inputs=(InputSpec((n,), gen=bool_tail_input),), + name=f"length_{n}", + ) + for n in (1, 4, 5, 67) + ], + golden_dtype="float32", + ) + + @register_op_test("add") def _add_suite() -> WebGPUTestSuite: # Same-shape numeric coverage only: broadcast adds stay export-smoke in @@ -295,10 +341,7 @@ def _minimum_suite() -> WebGPUTestSuite: def _compare_suite(op: str) -> WebGPUTestSuite: - # Elementwise fp32 comparison -> bool (byte-exact golden). The two inputs use - # DIFFERENT discrete-range seeds so a!=b (real lt/gt mix) while colliding - # often (eq/le/ge ties); all shapes have numel % 4 == 0 (bool output packs 4 - # bytes/word). Same-shape only (flat kernel; broadcast=smoke). + # Distinct inputs and tail shapes cover byte-exact packed BOOL output. def case(name, shape): return Case( name=name, @@ -310,7 +353,14 @@ def case(name, shape): return WebGPUTestSuite( module_factory=lambda: CompareModule(op), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + cases=[ + case("tail_1", (1,)), + case("tail_5", (5,)), + case("tail_67", (67,)), + case("2d", (4, 8)), + case("3d", (2, 3, 8)), + case("sq", (16, 16)), + ], golden_dtype="bool", ) @@ -612,6 +662,54 @@ def case(name, C, L, kernel, stride, padding, dilation, bias): case("k3s2p1", 4, 8, 3, 2, 1, 1, True), case("dil2", 3, 10, 3, 1, 2, 2, True), case("k5_nobias", 5, 7, 5, 1, 0, 1, False), + case("single_channel_route", 1, 8, 3, 1, 1, 1, True), + ], + atol=1e-3, + rtol=1e-3, + ) + + +@register_op_test("conv1d") +def _conv1d_suite() -> WebGPUTestSuite: + # General NCL conv1d; neighboring suites cover the retained fast paths. + def case(name, cfg): + n, ic, oc, length, kernel, stride, padding, dilation, bias = cfg + return Case( + name=name, + construct={ + "in_channels": ic, + "out_channels": oc, + "kernel_size": kernel, + "stride": stride, + "padding": padding, + "dilation": dilation, + "bias": bias, + }, + inputs=((n, ic, length),), + ) + + dynamic_cfg = _CONV1D_CONFIGS["voxtral_stride1"] + n, ic, oc, length, kernel, stride, padding, dilation, bias = dynamic_cfg + dynamic_length = torch.export.Dim("manifest_conv1d_length", min=7, max=length) + return WebGPUTestSuite( + module_factory=Conv1dModule, + cases=[ + *[case(name, cfg) for name, cfg in _CONV1D_CONFIGS.items()], + Case( + name="dynamic_length_10_to_7", + construct={ + "in_channels": ic, + "out_channels": oc, + "kernel_size": kernel, + "stride": stride, + "padding": padding, + "dilation": dilation, + "bias": bias, + }, + export_inputs=((n, ic, length),), + inputs=((n, ic, 7),), + dynamic_shapes=({2: dynamic_length},), + ), ], atol=1e-3, rtol=1e-3, @@ -988,12 +1086,36 @@ def _cat_suite() -> WebGPUTestSuite: N as _GELU_N, ) +_GELU_2D_DISPATCH_BOUNDARY = 4 * 64 * 65535 + 1 +_EXPAND_COPY_2D_DISPATCH_BOUNDARY = 64 * 65535 + 1 + def _gelu_full_range(_shape) -> torch.Tensor: # Reuse the deterministic linspace(-6, 6) spanning negatives/zero/positives. return _gelu_det_input() +@register_op_test("expand_copy") +def _expand_copy_suite() -> WebGPUTestSuite: + cases = [ + Case(name=name, construct={"shape": out_shape}, inputs=(in_shape,)) + for name, (in_shape, out_shape) in _EXPAND_COPY_CONFIGS.items() + ] + cases.append( + Case( + name="dispatch_2d_boundary", + construct={"shape": (_EXPAND_COPY_2D_DISPATCH_BOUNDARY,)}, + inputs=((1,),), + heavy=True, + ) + ) + return WebGPUTestSuite( + module_factory=ExpandCopyModule, + cases=cases, + golden_dtype="float32", + ) + + @register_op_test("gelu") def _gelu_suite() -> WebGPUTestSuite: # erf ("none") is the Florence-2/BART + PyTorch default; tanh is the approx. @@ -1015,6 +1137,12 @@ def _gelu_suite() -> WebGPUTestSuite: construct={"approximate": "none"}, inputs=(InputSpec(shape=(_GELU_N,), gen=_gelu_full_range),), ), + Case( + name="erf_dispatch_2d_boundary", + construct={"approximate": "none"}, + inputs=(InputSpec(shape=(_GELU_2D_DISPATCH_BOUNDARY,), gen="ramp"),), + heavy=True, + ), ], atol=1e-4, rtol=1e-3, diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index 72f819f94ce..33b7940a9be 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -39,6 +39,8 @@ def _materialize(spec) -> torch.Tensor: shape, gen = spec, "randn" if callable(gen): _t = gen(shape) + if _t.dtype == torch.bool: + return _t return ( _t.to(torch.int32) if not _t.is_floating_point() else _t.to(torch.float32) ) @@ -53,13 +55,17 @@ def _materialize(spec) -> torch.Tensor: def export_case(suite: WebGPUTestSuite, case) -> tuple[torch.nn.Module, tuple, object]: - """Build the module + forward inputs and export to an ExecuTorch program.""" + """Build the module and export it, returning the live runtime inputs.""" module = suite.module_factory(**case.construct) # Seed so an unseeded-randn input is reproducible across generations (the golden uses # the SAME tensor, so this only affects which bytes a case sees, never pass/fail). torch.manual_seed(0) inputs = tuple(_materialize(s) for s in case.inputs) - ep = torch.export.export(module, inputs) + export_inputs = inputs + if case.export_inputs is not None: + torch.manual_seed(0) + export_inputs = tuple(_materialize(s) for s in case.export_inputs) + ep = torch.export.export(module, export_inputs, dynamic_shapes=case.dynamic_shapes) prog = to_edge_transform_and_lower( ep, partitioner=[VulkanPartitioner()] ).to_executorch() @@ -169,7 +175,10 @@ def generate_case(op: str, suite: WebGPUTestSuite, case, out_dir: str) -> list[d input_entries: list[dict] = [] for i, t in enumerate(inputs): rel = f"{case_id}.in{i}.bin" - if t.dtype == torch.int32: + if t.dtype == torch.bool: + _write_int8(t.to(torch.int8), os.path.join(out_dir, rel)) + in_dtype = "bool" + elif t.dtype == torch.int32: t.detach().cpu().numpy().astype(" +#include #include #include #include @@ -61,7 +62,15 @@ class OpCase : public ::testing::Test { const size_t n = numel(in.shape); std::vector sizes( in.shape.begin(), in.shape.end()); - if (in.dtype == "int32") { + if (in.dtype == "bool") { + auto data = load_int8_bin(in.path, n); + ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; + std::vector raw(data.begin(), data.end()); + tensors.push_back(make_tensor_ptr( + std::move(sizes), + std::move(raw), + executorch::aten::ScalarType::Bool)); + } else if (in.dtype == "int32") { auto data = load_int32_bin(in.path, n); ASSERT_FALSE(data.empty()) << "missing/short input: " << in.path; tensors.push_back(make_tensor_ptr(std::move(sizes), std::move(data))); @@ -96,10 +105,11 @@ class OpCase : public ::testing::Test { auto golden = load_int8_bin(e_.golden.path, gn); ASSERT_FALSE(golden.empty()) << "missing/short golden: " << e_.golden.path; - const bool* out_p = out_tensor.const_data_ptr(); + ASSERT_EQ(out_tensor.scalar_type(), executorch::aten::ScalarType::Bool); + const uint8_t* out_p = out_tensor.const_data_ptr(); int mism = -1; for (size_t i = 0; i < gn; i++) { - if (static_cast(out_p[i]) != golden[i]) { + if (out_p[i] != static_cast(golden[i])) { mism = static_cast(i); break; } diff --git a/backends/webgpu/test/op_tests/test_generator.py b/backends/webgpu/test/op_tests/test_generator.py index 65f765812be..ec4125a2818 100644 --- a/backends/webgpu/test/op_tests/test_generator.py +++ b/backends/webgpu/test/op_tests/test_generator.py @@ -61,6 +61,27 @@ def test_generate_case_writes_artifacts(tmp_path): assert entry["golden"]["output_index"] == 0 +def test_export_case_separates_upper_bound_from_runtime_inputs(monkeypatch): + suite = op_test_registry["conv1d"] + case = next(c for c in suite.cases if c.name == "dynamic_length_10_to_7") + export_shapes = [] + exported_dynamic_shapes = [] + real_export = torch.export.export + + def capture_export(module, inputs, **kwargs): + export_shapes.append(tuple(inputs[0].shape)) + exported_dynamic_shapes.append(kwargs.get("dynamic_shapes")) + return real_export(module, inputs, **kwargs) + + monkeypatch.setattr(torch.export, "export", capture_export) + _module, runtime_inputs, prog = g.export_case(suite, case) + + assert export_shapes == [(1, 4, 10)] + assert exported_dynamic_shapes == [case.dynamic_shapes] + assert tuple(runtime_inputs[0].shape) == (1, 4, 7) + assert g._has_vulkan_delegate(prog) + + def test_generate_manifest(tmp_path): g.generate(str(tmp_path), ops=["add"]) manifest = tmp_path / "manifest.json" diff --git a/backends/webgpu/test/op_tests/test_suite.py b/backends/webgpu/test/op_tests/test_suite.py index f2714125c84..17542cd2e55 100644 --- a/backends/webgpu/test/op_tests/test_suite.py +++ b/backends/webgpu/test/op_tests/test_suite.py @@ -58,6 +58,9 @@ class Case: required: bool = True heavy: bool = False golden_fn: Callable | None = None + # Optional upper-bound export inputs; `inputs` stay live manifest tensors. + export_inputs: tuple[Input, ...] | None = None + dynamic_shapes: object | None = None def __post_init__(self) -> None: # Mirror kQ4gswConfigs: every heavy config is required=False (export-gated, never FAILs on absence). diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index b302b3c120d..d9f666622ee 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -19,6 +19,8 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dModule +from executorch.backends.webgpu.test.ops.test_gelu import GeluModule from executorch.exir import to_edge_transform_and_lower from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes @@ -181,6 +183,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.select(0, -1) +class DynamicExpandCopyModule(torch.nn.Module): + """Dynamic expand_copy is rejected until its TensorMeta can be resized.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.expand((4, x.shape[1])).clone() + + +class DynamicExpandCopyInferredModule(torch.nn.Module): + """Dynamic expand_copy whose -1 target hides symbolic provenance.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.expand((4, -1)).clone() + + def _ramp(shape) -> torch.Tensor: n = 1 for d in shape: @@ -252,9 +268,75 @@ def _write_goldens(model, prefix: str, out_dir: str, s_values) -> None: print(f" golden {prefix} S={s}") +def export_dynamic_conv1d_cases(out_dir: str) -> None: + """Write one dynamic Conv1d program and live-length runtime fixtures.""" + os.makedirs(out_dir, exist_ok=True) + max_length = 16 + lengths = (max_length, 9, 5) + model = Conv1dModule( + in_channels=3, + out_channels=4, + kernel_size=3, + stride=2, + padding=1, + dilation=2, + bias=True, + ).eval() + length_dim = torch.export.Dim("conv1d_length", min=5, max=max_length) + _export( + model, + (_ramp((1, 3, max_length)),), + {"x": {2: length_dim}}, + os.path.join(out_dir, "dyn_conv1d.pte"), + ) + for length in lengths: + x = _ramp((1, 3, length)) + with torch.no_grad(): + golden = model(x) + prefix = os.path.join(out_dir, f"dyn_conv1d.S{length}") + x.detach().numpy().astype(" None: + """Write a dynamic GELU fixture crossing the old 1D dispatch cap.""" + os.makedirs(out_dir, exist_ok=True) + max_elements = 4 * 64 * 65535 + 1 + model = GeluModule("none").eval() + elements_dim = torch.export.Dim("gelu_elements", min=1024, max=max_elements) + _export( + model, + (torch.empty((max_elements,), dtype=torch.float32),), + {"x": {0: elements_dim}}, + os.path.join(out_dir, "dyn_gelu_2d.pte"), + ) + + +def export_dynamic_expand_copy_rejection_case(out_dir: str) -> None: + """Write a dynamic expand_copy graph that the runtime must reject at load.""" + model = DynamicExpandCopyModule().eval() + elements_dim = torch.export.Dim("expand_elements", min=1, max=8) + _export( + model, + (_ramp((1, 8)),), + {"x": {1: elements_dim}}, + os.path.join(out_dir, "dyn_expand_copy.pte"), + ) + _export( + DynamicExpandCopyInferredModule().eval(), + (_ramp((1, 8)),), + {"x": {1: elements_dim}}, + os.path.join(out_dir, "dyn_expand_copy_inferred.pte"), + ) + + def export_dynamic_shape_cases(out_dir: str) -> None: """Write the dynamic + static .pte's and per-S goldens for the native test.""" os.makedirs(out_dir, exist_ok=True) + export_dynamic_conv1d_cases(out_dir) + export_dynamic_expand_copy_rejection_case(out_dir) + if os.environ.get("WEBGPU_TEST_HEAVY"): + export_dynamic_gelu_boundary_cases(out_dir) s_dim = torch.export.Dim("s", min=1, max=MAXS) # 1) Single dynamic rms_norm, graph built at S=MAXS (upper bound). @@ -1536,6 +1618,13 @@ def test_export_dynamic_rms(self) -> None: self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.pte"))) self.assertTrue(os.path.exists(os.path.join(d, "dyn_rms.S1.golden.bin"))) expected = [ + "dyn_conv1d.pte", + "dyn_conv1d.S16.input.bin", + "dyn_conv1d.S16.golden.bin", + "dyn_conv1d.S9.input.bin", + "dyn_conv1d.S9.golden.bin", + "dyn_conv1d.S5.input.bin", + "dyn_conv1d.S5.golden.bin", "dyn_linear_bk64.pte", "dyn_linear_bk64.S512.input.bin", "dyn_linear_bk64.S512.golden.bin", diff --git a/backends/webgpu/test/ops/test_conv1d_pw.py b/backends/webgpu/test/ops/test_conv1d_pw.py index a0de8988f81..07036fe4a42 100644 --- a/backends/webgpu/test/ops/test_conv1d_pw.py +++ b/backends/webgpu/test/ops/test_conv1d_pw.py @@ -27,6 +27,16 @@ "batch2": (2, 3, 4, 5, True), } +# name -> N, C_in, C_out, L, K, stride, padding, dilation, bias +GENERAL_CONFIGS = { + "voxtral_stride1": (1, 4, 6, 10, 3, 1, 0, 1, True), + "voxtral_stride2": (1, 6, 5, 10, 3, 2, 0, 1, True), + "no_bias": (1, 3, 2, 9, 3, 1, 0, 1, False), + "padded": (1, 3, 4, 9, 3, 1, 1, 1, True), + "dilated": (1, 2, 3, 11, 3, 1, 2, 2, True), + "batch2": (2, 3, 4, 8, 3, 2, 1, 1, True), +} + class Conv1dPwModule(torch.nn.Module): def __init__(self, in_channels, out_channels, bias) -> None: @@ -42,6 +52,37 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x) +class Conv1dModule(torch.nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride, + padding, + dilation, + bias, + ) -> None: + super().__init__() + g = torch.Generator().manual_seed(0) + self.conv = torch.nn.Conv1d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=bias, + ) + with torch.no_grad(): + self.conv.weight.normal_(generator=g) + if bias: + self.conv.bias.normal_(generator=g) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + def _det_input(shape): g = torch.Generator().manual_seed(1) return torch.randn(*shape, generator=g, dtype=torch.float32) @@ -54,6 +95,14 @@ def _lower(cfg): return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) +def _lower_general(cfg, dynamic_shapes=None): + n, ic, oc, length, kernel, stride, padding, dilation, bias = cfg + module = Conv1dModule(ic, oc, kernel, stride, padding, dilation, bias).eval() + inputs = (_det_input((n, ic, length)),) + ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes) + return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + + def _delegated(et) -> bool: return any( d.id == "VulkanBackend" @@ -63,9 +112,17 @@ def _delegated(et) -> bool: def _op_delegated(edge, op_substr: str) -> bool: - # op must be absorbed into the delegate, not left as a top-level CPU-fallback node. + # Require the op in a delegate, not merely absent from the host graph. + from executorch.exir.lowered_backend_module import get_lowered_submodules + gm = edge.exported_program().graph_module - return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes) + if any(op_substr in str(getattr(n, "target", "")) for n in gm.graph.nodes): + return False + return any( + op_substr in str(getattr(dn, "target", "")) + for _, lowered, _ in get_lowered_submodules(gm) + for dn in lowered.original_module.graph_module.graph.nodes + ) class Conv1dPwTest(unittest.TestCase): @@ -82,3 +139,24 @@ def test_export_delegates(self) -> None: _op_delegated(edge, "convolution"), f"conv1d not delegated (fell back to CPU) for {name}", ) + + +class Conv1dTest(unittest.TestCase): + def test_export_delegates(self) -> None: + for name, cfg in GENERAL_CONFIGS.items(): + with self.subTest(name=name): + edge = _lower_general(cfg) + self.assertTrue( + _delegated(edge.to_executorch()), + f"Expected a VulkanBackend delegate (conv1d {name})", + ) + self.assertTrue( + _op_delegated(edge, "convolution"), + f"conv1d not delegated (fell back to CPU) for {name}", + ) + + def test_dynamic_length_export_delegates(self) -> None: + length = torch.export.Dim("conv1d_length", min=5, max=16) + edge = _lower_general(GENERAL_CONFIGS["padded"], dynamic_shapes=({2: length},)) + self.assertTrue(_delegated(edge.to_executorch())) + self.assertTrue(_op_delegated(edge, "convolution")) diff --git a/backends/webgpu/test/ops/test_to_copy.py b/backends/webgpu/test/ops/test_to_copy.py index 1fa2375f248..54b400ea9ef 100644 --- a/backends/webgpu/test/ops/test_to_copy.py +++ b/backends/webgpu/test/ops/test_to_copy.py @@ -46,6 +46,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.to(torch.float32, copy=True) +class ToCopyBoolToFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float32) + + +class ToCopyInt8ToFloatModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.to(torch.float32) + + +class CompareToCopyBoolToFloatModule(torch.nn.Module): + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return (a > b).to(torch.float32) + + def to_copy_int_input(shape: tuple[int, ...]) -> torch.Tensor: n = math.prod(shape) return (torch.arange(n, dtype=torch.int32) - n // 2).reshape(shape) @@ -61,14 +76,32 @@ def to_copy_float_input(shape: tuple[int, ...]) -> torch.Tensor: return pattern.repeat(repeats)[:n].reshape(shape) -def _lower(model: torch.nn.Module, x: torch.Tensor): - ep = torch.export.export(model.eval(), (x,)) +def bool_tail_input(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + pattern = torch.tensor([True, False, True, True, False, False, True]) + repeats = (n + pattern.numel() - 1) // pattern.numel() + return pattern.repeat(repeats)[:n].reshape(shape) + + +def compare_to_copy_input_a(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + pattern = torch.tensor([1.0, -1.0, 2.0, -2.0, 3.0, -3.0, 4.0]) + repeats = (n + pattern.numel() - 1) // pattern.numel() + return pattern.repeat(repeats)[:n].reshape(shape) + + +def compare_to_copy_input_b(shape: tuple[int, ...]) -> torch.Tensor: + return torch.zeros(shape, dtype=torch.float32) + + +def _lower(model: torch.nn.Module, *inputs: torch.Tensor): + ep = torch.export.export(model.eval(), inputs) edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) return ep, edge -def _export(model: torch.nn.Module, x: torch.Tensor): - _, edge = _lower(model, x) +def _export(model: torch.nn.Module, *inputs: torch.Tensor): + _, edge = _lower(model, *inputs) return edge.to_executorch() @@ -142,3 +175,22 @@ def test_float_passthrough_delegates(self) -> None: self.assertTrue( _delegated(et), "Expected a VulkanBackend delegate (to_copy float->float)" ) + + def test_bool_to_float_delegates(self) -> None: + x = bool_tail_input((5,)) + ep, edge = _lower(ToCopyBoolToFloatModule(), x) + self.assertEqual(_prepartition_cast_dtypes(ep), [torch.float32]) + self.assertEqual(_delegated_cast_dtypes(edge), [torch.float32]) + self.assertTrue(_delegated(edge.to_executorch())) + + def test_compare_bool_to_float_delegates(self) -> None: + a = compare_to_copy_input_a((5,)) + b = compare_to_copy_input_b((5,)) + ep, edge = _lower(CompareToCopyBoolToFloatModule(), a, b) + self.assertEqual(_prepartition_cast_dtypes(ep), [torch.float32]) + self.assertEqual(_delegated_cast_dtypes(edge), [torch.float32]) + self.assertTrue(_delegated(edge.to_executorch())) + + def test_int8_to_float_does_not_delegate(self) -> None: + x = torch.tensor([-2, 0, 3], dtype=torch.int8) + self.assertFalse(_delegated(_export(ToCopyInt8ToFloatModule(), x))) diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 574c3869864..9990297b9f0 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -77,6 +77,16 @@ def _function_source(text: str, name: str) -> str: class WgslCodegenTest(unittest.TestCase): + def test_compare_word_count_does_not_overflow_u32(self) -> None: + source = (g.BACKEND_ROOT / "runtime/ops/compare/compare.wgsl").read_text() + expression = "(params.num_elements - 1u) / 4u + 1u" + self.assertIn(expression, source) + for num_elements in (1, 4, 5, (1 << 32) - 3, (1 << 32) - 2, (1 << 32) - 1): + self.assertEqual( + (num_elements - 1) // 4 + 1, + (num_elements + 3) // 4, + ) + def test_registry_entries_match_concrete_headers(self) -> None: entries = g.registry_entries() names = [entry.name for entry in entries] @@ -210,14 +220,14 @@ def test_generated_output_manifest_digest(self) -> None: digest.update(b"\0") digest.update(output.read_bytes()) digest.update(b"\0") - self.assertEqual(len(outputs), 134) + self.assertEqual(len(outputs), 136) self.assertEqual( digest.hexdigest(), - "e502196846f0f8100f468e5d9f8f9c006b67e08df54e1e2e667daa2fc50d8844", + "0512f8d258952e446ffaedcb653b6a3a720eccf8a6b5327d95fd454a912214a3", ) self.assertEqual( hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "492535b396833ad6ebfed29b093057e38f40b1182b5c7d6b3eb2f3577cab024e", + "28aaa7a8d3e916df43e407120e91d487d0d51cbc5ca93c56bd822d25d109890e", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: