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 04331ee6ca7..ca312bcda88 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -1739,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++) { diff --git a/backends/webgpu/runtime/ops/compare/Compare.cpp b/backends/webgpu/runtime/ops/compare/Compare.cpp index f47f9491191..1870ee85213 100644 --- a/backends/webgpu/runtime/ops/compare/Compare.cpp +++ b/backends/webgpu/runtime/ops/compare/Compare.cpp @@ -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( @@ -142,8 +140,6 @@ void compare_impl( }; 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/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 3f9c8e29e3c..54d80016296 100644 --- a/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp +++ b/backends/webgpu/runtime/ops/to_copy/ToCopy.cpp @@ -70,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, @@ -115,9 +113,94 @@ void add_convert_op( wg_size, "to_copy(resize)"); }); +} - // Graph owns it so the resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); +// 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. @@ -223,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) { 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 2c78e9192ca..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 @@ -654,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, @@ -1030,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. @@ -1057,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 860f0661913..33b7940a9be 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -55,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() 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"))