From 8aafb78f1039c624e3399038e5d7089d40b05205 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 21:10:49 -0700 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 28 ++- backends/webgpu/runtime/WebGPUBackend.cpp | 15 +- .../webgpu/runtime/WebGPUExecutionOptions.cpp | 128 +++++++++++ .../webgpu/runtime/WebGPUExecutionOptions.h | 84 +++++++ backends/webgpu/runtime/WebGPUGraph.cpp | 79 +++++-- backends/webgpu/runtime/WebGPUGraph.h | 14 +- .../test/native/test_execution_options.cpp | 127 +++++++++++ .../test/native/test_output_suppression.cpp | 174 ++++++++++++++ .../webgpu/test/ops/test_quantized_linear.py | 215 +++++++++++++++++- 9 files changed, 826 insertions(+), 38 deletions(-) create mode 100644 backends/webgpu/runtime/WebGPUExecutionOptions.cpp create mode 100644 backends/webgpu/runtime/WebGPUExecutionOptions.h create mode 100644 backends/webgpu/test/native/test_execution_options.cpp create mode 100644 backends/webgpu/test/native/test_output_suppression.cpp diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 78d9c4f30dd..4c9f1265c4e 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -26,9 +26,13 @@ if(NOT TARGET vulkan_schema) endif() set(WEBGPU_SRCS - runtime/WebGPUBackend.cpp runtime/WebGPUGraph.cpp - runtime/WebGPUDelegateHeader.cpp runtime/WebGPUDevice.cpp - runtime/WebGPUQueryPool.cpp runtime/ops/OperatorRegistry.cpp + runtime/WebGPUBackend.cpp + runtime/WebGPUGraph.cpp + runtime/WebGPUDelegateHeader.cpp + runtime/WebGPUDevice.cpp + runtime/WebGPUExecutionOptions.cpp + runtime/WebGPUQueryPool.cpp + runtime/ops/OperatorRegistry.cpp ) # Op handlers: glob so adding an op needs no CMakeLists edit. CONFIGURE_DEPENDS @@ -199,5 +203,23 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) target_link_libraries( webgpu_dispatch_2d_test PRIVATE GTest::gtest GTest::gtest_main ) + + add_executable( + webgpu_execution_options_test test/native/test_execution_options.cpp + runtime/WebGPUExecutionOptions.cpp + ) + target_include_directories( + webgpu_execution_options_test + PRIVATE $ + ) + target_link_libraries( + webgpu_execution_options_test PRIVATE GTest::gtest GTest::gtest_main + ) + target_compile_options(webgpu_execution_options_test PRIVATE -fexceptions) + set_property(TARGET webgpu_execution_options_test PROPERTY CXX_STANDARD 17) + add_webgpu_native_test( + webgpu_output_suppression_test test/native/test_output_suppression.cpp + ) + target_link_libraries(webgpu_output_suppression_test PRIVATE GTest::gtest) endif() endif() diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index e1497bc4ed8..51cca64fd50 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -121,9 +122,11 @@ Error WebGPUBackend::execute( DelegateHandle* handle, Span args) const { WebGPUGraph* graph = static_cast(handle); + const WebGPUExecutionOptions options = current_webgpu_execution_options(); const size_t num_inputs = graph->input_ids().size(); const size_t num_outputs = graph->output_ids().size(); + WebGPUGraphExecutionOptions graph_options; // Copy inputs from EValue tensors to GPU buffers std::vector inputs; @@ -159,13 +162,21 @@ Error WebGPUBackend::execute( return Error::Internal; } } + std::vector delegate_outputs; + delegate_outputs.reserve(num_outputs); + for (size_t i = 0; i < num_outputs; i++) { + delegate_outputs.push_back( + args[num_inputs + i]->toTensor().mutable_data_ptr()); + } + graph_options = + resolve_webgpu_graph_execution_options(delegate_outputs, options); } catch (const std::exception& e) { ET_LOG(Error, "WebGPU input/output resize / copy failed: %s", e.what()); return Error::Internal; } // Execute the compute graph - graph->execute(); + graph->execute(graph_options); // Copy outputs from GPU staging buffers to EValue tensor data pointers std::vector> outputs; @@ -175,7 +186,7 @@ Error WebGPUBackend::execute( auto& tensor = args[arg_idx]->toTensor(); outputs.emplace_back(tensor.mutable_data_ptr(), tensor.nbytes()); } - graph->copy_outputs(outputs); + graph->copy_outputs(outputs, graph_options); return Error::Ok; } diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.cpp b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp new file mode 100644 index 00000000000..7362df98cfb --- /dev/null +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp @@ -0,0 +1,128 @@ +/* + * 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. + */ + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { +namespace { + +thread_local WebGPUExecutionOptions execution_options; + +} // namespace + +WebGPUExecutionOptions current_webgpu_execution_options() { + return execution_options; +} + +ScopedWebGPUExecutionOptions::ScopedWebGPUExecutionOptions( + WebGPUExecutionOptions options) + : previous_(execution_options) { + execution_options = options; +} + +ScopedWebGPUExecutionOptions::~ScopedWebGPUExecutionOptions() { + execution_options = previous_; +} + +WebGPUExecutionPlan plan_webgpu_execution( + size_t dispatch_count, + size_t output_count, + ExecuteConfig config, + const std::vector& suppressible_outputs, + WebGPUGraphExecutionOptions options) { + std::vector suppressed_dispatches(dispatch_count, false); + std::vector copy_outputs(output_count, true); + std::vector suppressed_outputs(output_count, false); + + for (const auto& output : suppressible_outputs) { + if (output.output_ordinal >= output_count || + output.dispatch_begin >= output.dispatch_end || + output.dispatch_end > dispatch_count) { + throw std::runtime_error("WebGPU: invalid suppressible output range"); + } + if (suppressed_outputs[output.output_ordinal]) { + throw std::runtime_error("WebGPU: duplicate suppressible output"); + } + suppressed_outputs[output.output_ordinal] = true; + if (output.output_ordinal != options.suppress_output_ordinal) { + continue; + } + copy_outputs[output.output_ordinal] = false; + for (size_t i = output.dispatch_begin; i < output.dispatch_end; i++) { + if (suppressed_dispatches[i]) { + throw std::runtime_error( + "WebGPU: overlapping suppressible dispatch ranges"); + } + suppressed_dispatches[i] = true; + } + } + + WebGPUExecutionPlan plan; + plan.copy_outputs = std::move(copy_outputs); + + auto append_chunk = [&](size_t begin, size_t end) { + std::vector indices; + indices.reserve(end - begin); + for (size_t i = begin; i < end; i++) { + if (!suppressed_dispatches[i]) { + indices.push_back(i); + } + } + if (!indices.empty()) { + plan.dispatch_chunks.push_back(std::move(indices)); + } + }; + + if (config.chunk_size == 0 || dispatch_count <= config.chunk_size) { + append_chunk(0, dispatch_count); + } else { + size_t start = 0; + size_t current_chunk = config.initial_chunk_size > 0 + ? config.initial_chunk_size + : config.chunk_size; + while (start < dispatch_count) { + const size_t end = std::min(start + current_chunk, dispatch_count); + append_chunk(start, end); + start = end; + current_chunk = config.chunk_size; + } + } + if (plan.dispatch_chunks.empty()) { + plan.dispatch_chunks.emplace_back(); + } + return plan; +} + +WebGPUGraphExecutionOptions resolve_webgpu_graph_execution_options( + const std::vector& delegate_outputs, + WebGPUExecutionOptions options) { + if (options.discardable_output_data == nullptr) { + return {}; + } + if (!options.exact_method_certificate_verified) { + return {}; + } + + size_t match = kNoOutputOrdinal; + for (size_t i = 0; i < delegate_outputs.size(); i++) { + if (delegate_outputs[i] != options.discardable_output_data) { + continue; + } + if (match != kNoOutputOrdinal) { + return {}; + } + match = i; + } + return {match}; +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.h b/backends/webgpu/runtime/WebGPUExecutionOptions.h new file mode 100644 index 00000000000..691cd328e90 --- /dev/null +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.h @@ -0,0 +1,84 @@ +/* + * 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 +#include +#include + +namespace executorch::backends::webgpu { + +constexpr size_t kNoOutputOrdinal = static_cast(-1); + +struct WebGPUExecutionOptions { + // The certificate must bind the exact PTE and method and prove one delegate, + // no portable nodes, and a unique leaf method output at this data pointer. + const void* discardable_output_data = nullptr; + bool exact_method_certificate_verified = false; +}; + +struct WebGPUGraphExecutionOptions { + size_t suppress_output_ordinal = kNoOutputOrdinal; +}; + +struct ExecuteConfig { + size_t chunk_size = 0; + size_t initial_chunk_size = 0; +}; + +struct SuppressibleOutput { + int output_id = -1; + size_t output_ordinal = 0; + size_t dispatch_begin = 0; + size_t dispatch_end = 0; +}; + +struct WebGPUExecutionPlan { + std::vector> dispatch_chunks; + std::vector copy_outputs; +}; + +WebGPUExecutionPlan plan_webgpu_execution( + size_t dispatch_count, + size_t output_count, + ExecuteConfig config, + const std::vector& suppressible_outputs, + WebGPUGraphExecutionOptions options); + +WebGPUGraphExecutionOptions resolve_webgpu_graph_execution_options( + const std::vector& delegate_outputs, + WebGPUExecutionOptions options); + +WebGPUExecutionOptions current_webgpu_execution_options(); + +class ScopedWebGPUExecutionOptions final { + public: + explicit ScopedWebGPUExecutionOptions(WebGPUExecutionOptions options); + ~ScopedWebGPUExecutionOptions(); + + ScopedWebGPUExecutionOptions(const ScopedWebGPUExecutionOptions&) = delete; + ScopedWebGPUExecutionOptions& operator=(const ScopedWebGPUExecutionOptions&) = + delete; + ScopedWebGPUExecutionOptions(ScopedWebGPUExecutionOptions&&) = delete; + ScopedWebGPUExecutionOptions& operator=(ScopedWebGPUExecutionOptions&&) = + delete; + + private: + WebGPUExecutionOptions previous_; +}; + +template +decltype(auto) with_webgpu_execution_options( + WebGPUExecutionOptions options, + Fn&& fn) { + ScopedWebGPUExecutionOptions scope(options); + return std::forward(fn)(); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 2bb091d391e..5a567a36500 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -33,6 +33,8 @@ namespace { // Op name the AOT exporter emits for a prepacked constant (must match the // serialized schema); compared in the prepack pre-scan below. constexpr const char* kPrepackOpName = "et_vk.prepack.default"; +constexpr const char* kQ4gswLinearOpName = "et_vk.linear_q4gsw.default"; +constexpr size_t kQ4gswOutputArg = 5; size_t vk_datatype_size(vkgraph::VkDataType dtype) { switch (dtype) { @@ -1109,7 +1111,7 @@ void WebGPUGraph::build( create_scratch_buffer(tensors_[g.out_v].nbytes); } } - + const size_t dispatch_begin = dispatches_.size(); webgpu_operator_registry().get_op_fn(op_name)(*this, args); { @@ -1145,6 +1147,24 @@ void WebGPUGraph::build( release_scratch(tensors_[swiglu_groups[orl->second][2]].buffer); } } + + const size_t dispatch_end = dispatches_.size(); + + if (i + 1 == chain->size() && op_name == kQ4gswLinearOpName && + args.size() > kQ4gswOutputArg && dispatch_end > dispatch_begin) { + const int output_id = args[kQ4gswOutputArg]; + const auto output_it = + std::find(output_ids_.begin(), output_ids_.end(), output_id); + if (output_it != output_ids_.end() && + std::count(output_ids_.begin(), output_ids_.end(), output_id) == + 1) { + suppressible_outputs_.push_back( + {output_id, + static_cast(output_it - output_ids_.begin()), + dispatch_begin, + dispatch_end}); + } + } } } @@ -1710,9 +1730,15 @@ bool should_timestamp_query() { } } // namespace -void WebGPUGraph::execute() { +void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { const size_t n = dispatches_.size(); const size_t chunk = execute_config_.chunk_size; + const WebGPUExecutionPlan plan = plan_webgpu_execution( + n, + output_copies_.size(), + execute_config_, + suppressible_outputs_, + options); if (chunk == 0 || n <= chunk) { #ifdef WGPU_BACKEND_ENABLE_PROFILING @@ -1737,7 +1763,7 @@ void WebGPUGraph::execute() { wgpuDeviceCreateCommandEncoder(device_, &enc_desc); // One pass per dispatch: enforces storage RAW ordering across deps. - for (size_t i = 0; i < n; i++) { + for (size_t i : plan.dispatch_chunks.front()) { const auto& dispatch = dispatches_[i]; if (dispatch.kind == WebGPUDispatch::Kind::Copy) { wgpuCommandEncoderCopyBufferToBuffer( @@ -1778,7 +1804,11 @@ void WebGPUGraph::execute() { #endif // WGPU_BACKEND_ENABLE_PROFILING } - for (const auto& copy : output_copies_) { + for (size_t i = 0; i < output_copies_.size(); i++) { + if (!plan.copy_outputs[i]) { + continue; + } + const auto& copy = output_copies_[i]; wgpuCommandEncoderCopyBufferToBuffer( encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy.nbytes); } @@ -1812,21 +1842,13 @@ void WebGPUGraph::execute() { "(multi-submit); disable chunking to use GPU timestamp queries"); } - const size_t first_chunk = execute_config_.initial_chunk_size > 0 - ? execute_config_.initial_chunk_size - : chunk; - - size_t start = 0; - size_t current_chunk = first_chunk; - - while (start < n) { - size_t end = std::min(start + current_chunk, n); - + for (size_t chunk_index = 0; chunk_index < plan.dispatch_chunks.size(); + chunk_index++) { WGPUCommandEncoderDescriptor enc_desc = {}; WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device_, &enc_desc); - for (size_t i = start; i < end; i++) { + for (size_t i : plan.dispatch_chunks[chunk_index]) { if (dispatches_[i].kind == WebGPUDispatch::Kind::Copy) { wgpuCommandEncoderCopyBufferToBuffer( encoder, @@ -1852,8 +1874,12 @@ void WebGPUGraph::execute() { wgpuComputePassEncoderRelease(pass); } - if (end == n) { - for (const auto& copy : output_copies_) { + if (chunk_index + 1 == plan.dispatch_chunks.size()) { + for (size_t i = 0; i < output_copies_.size(); i++) { + if (!plan.copy_outputs[i]) { + continue; + } + const auto& copy = output_copies_[i]; wgpuCommandEncoderCopyBufferToBuffer( encoder, copy.src_buffer, 0, copy.staging_buffer, 0, copy.nbytes); } @@ -1865,9 +1891,6 @@ void WebGPUGraph::execute() { wgpuCommandBufferRelease(cmd); wgpuCommandEncoderRelease(encoder); - - start = end; - current_chunk = chunk; } } @@ -1888,14 +1911,22 @@ void buffer_map_callback( } // namespace -void WebGPUGraph::copy_outputs(std::vector>& outputs) { +void WebGPUGraph::copy_outputs( + std::vector>& outputs, + const WebGPUGraphExecutionOptions& options) { const size_t count = std::min(outputs.size(), output_staging_buffers_.size()); + const WebGPUExecutionPlan plan = plan_webgpu_execution( + dispatches_.size(), + output_copies_.size(), + execute_config_, + suppressible_outputs_, + options); std::vector cb_data(count); std::vector map_futures(count, WGPUFuture{}); for (size_t i = 0; i < count; i++) { - if (outputs[i].second == 0) { + if (!plan.copy_outputs[i] || outputs[i].second == 0) { cb_data[i].status = WGPUMapAsyncStatus_Success; continue; } @@ -1912,14 +1943,14 @@ void WebGPUGraph::copy_outputs(std::vector>& outputs) { } for (size_t i = 0; i < count; i++) { - if (outputs[i].second != 0 && + if (plan.copy_outputs[i] && outputs[i].second != 0 && webgpu_wait(instance_, map_futures[i]) != WGPUWaitStatus_Success) { throw std::runtime_error("WebGPU: WaitAny failed for output map"); } } for (size_t i = 0; i < count; i++) { - if (outputs[i].second == 0) { + if (!plan.copy_outputs[i] || outputs[i].second == 0) { continue; } if (cb_data[i].status == WGPUMapAsyncStatus_Success) { diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 6f99c6162a4..c520cb06f84 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -17,6 +17,7 @@ #include #include +#include #include namespace executorch::backends::webgpu { @@ -72,11 +73,6 @@ struct ConstantSource { size_t nbytes = 0; }; -struct ExecuteConfig { - size_t chunk_size = 0; - size_t initial_chunk_size = 0; -}; - struct WebGPUMemoryStats { size_t tensor_buffer_bytes = 0; size_t shared_buffer_bytes = 0; @@ -112,11 +108,13 @@ class WebGPUGraph { void copy_inputs(const std::vector& inputs); // Execute all recorded dispatches. - void execute(); + void execute(const WebGPUGraphExecutionOptions& options); // Copy output tensor data from GPU buffers back to host pointers. // Uses mapAsync + ASYNCIFY in Wasm. - void copy_outputs(std::vector>& outputs); + void copy_outputs( + std::vector>& outputs, + const WebGPUGraphExecutionOptions& options); const std::vector& input_ids() const { return input_ids_; @@ -433,6 +431,8 @@ class WebGPUGraph { // Pre-computed output copy descriptors for execute(). std::vector output_copies_; + std::vector suppressible_outputs_; + std::vector dispatches_; // Prepack-routed constant sources (offset/named-key + size); the prepack node diff --git a/backends/webgpu/test/native/test_execution_options.cpp b/backends/webgpu/test/native/test_execution_options.cpp new file mode 100644 index 00000000000..2f96636438d --- /dev/null +++ b/backends/webgpu/test/native/test_execution_options.cpp @@ -0,0 +1,127 @@ +/* + * 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. + */ + +#include + +#include + +#include +#include + +namespace executorch::backends::webgpu { +namespace { + +TEST(WebGPUExecutionOptionsTest, DefaultsToPreservingOutputs) { + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, nullptr); +} + +TEST(WebGPUExecutionOptionsTest, NestedScopesRestorePriorValue) { + int outer_output = 0; + int inner_output = 0; + { + ScopedWebGPUExecutionOptions outer({&outer_output, true}); + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, + &outer_output); + { + ScopedWebGPUExecutionOptions inner({&inner_output, true}); + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, + &inner_output); + } + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, + &outer_output); + } + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, nullptr); +} + +TEST(WebGPUExecutionOptionsTest, ExceptionRestoresPriorValue) { + int output = 0; + EXPECT_THROW( + with_webgpu_execution_options( + {&output, true}, + []() -> void { throw std::runtime_error("expected"); }), + std::runtime_error); + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, nullptr); +} + +TEST(WebGPUExecutionOptionsTest, ErrorReturnRestoresPriorValue) { + int output = 0; + const bool result = + with_webgpu_execution_options({&output, true}, []() { return false; }); + EXPECT_FALSE(result); + EXPECT_EQ( + current_webgpu_execution_options().discardable_output_data, nullptr); +} + +TEST(WebGPUExecutionOptionsTest, ResolvesOnlyOneExactDelegateOutput) { + int method_output = 0; + int delegate_intermediate = 0; + const std::vector delegate_outputs = {&delegate_intermediate}; + + EXPECT_EQ( + resolve_webgpu_graph_execution_options( + delegate_outputs, WebGPUExecutionOptions{&method_output, true}) + .suppress_output_ordinal, + kNoOutputOrdinal); + EXPECT_EQ( + resolve_webgpu_graph_execution_options( + {&delegate_intermediate, &method_output}, + WebGPUExecutionOptions{&method_output, true}) + .suppress_output_ordinal, + 1); + EXPECT_EQ( + resolve_webgpu_graph_execution_options( + {&method_output, &method_output}, + WebGPUExecutionOptions{&method_output, true}) + .suppress_output_ordinal, + kNoOutputOrdinal); + EXPECT_EQ( + resolve_webgpu_graph_execution_options( + {&method_output}, WebGPUExecutionOptions{&method_output, false}) + .suppress_output_ordinal, + kNoOutputOrdinal); +} + +TEST(WebGPUExecutionPlanTest, DefaultPlanPreservesDispatchesAndOutputs) { + const std::vector suppressible = {{9, 1, 4, 6}}; + const WebGPUExecutionPlan plan = plan_webgpu_execution( + 6, 2, ExecuteConfig{}, suppressible, WebGPUGraphExecutionOptions{}); + + EXPECT_EQ( + plan.dispatch_chunks, + (std::vector>{{0, 1, 2, 3, 4, 5}})); + EXPECT_EQ(plan.copy_outputs, (std::vector{true, true})); +} + +TEST(WebGPUExecutionPlanTest, SuppressionIsPerOutputAndSupportsChunking) { + const std::vector suppressible = {{9, 1, 4, 6}}; + const ExecuteConfig config = {2, 1}; + const WebGPUExecutionPlan plan = plan_webgpu_execution( + 6, 2, config, suppressible, WebGPUGraphExecutionOptions{1}); + + EXPECT_EQ( + plan.dispatch_chunks, + (std::vector>{{0}, {1, 2}, {3}})); + EXPECT_EQ(plan.copy_outputs, (std::vector{true, false})); +} + +TEST(WebGPUExecutionPlanTest, RejectsInvalidSuppressibleRange) { + const std::vector suppressible = {{9, 0, 3, 7}}; + EXPECT_THROW( + plan_webgpu_execution( + 6, 1, ExecuteConfig{}, suppressible, WebGPUGraphExecutionOptions{0}), + std::runtime_error); +} + +} // namespace +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_output_suppression.cpp b/backends/webgpu/test/native/test_output_suppression.cpp new file mode 100644 index 00000000000..5cfae6a14f2 --- /dev/null +++ b/backends/webgpu/test/native/test_output_suppression.cpp @@ -0,0 +1,174 @@ +/* + * 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. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace executorch::backends::webgpu; +using namespace executorch::extension; +using namespace executorch::runtime; + +namespace { + +constexpr int kWidth = 64; +constexpr float kPoison = -12345.0f; + +std::string g_dir; // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) + +std::vector read_bin(const std::string& path) { + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file) { + return {}; + } + const std::streamsize bytes = file.tellg(); + if (bytes < 0 || bytes % sizeof(float) != 0) { + return {}; + } + file.seekg(0); + std::vector values(static_cast(bytes) / sizeof(float)); + file.read(reinterpret_cast(values.data()), bytes); + return values; +} + +void expect_matches( + const std::vector& actual, + const std::vector& expected) { + ASSERT_EQ(actual.size(), expected.size()); + for (size_t i = 0; i < actual.size(); i++) { + const float abs_err = std::fabs(actual[i] - expected[i]); + const float denom = std::max(std::fabs(expected[i]), 1e-12f); + const float rel_err = abs_err / denom; + EXPECT_TRUE(abs_err <= 1e-3f || rel_err <= 1e-3f) + << "index=" << i << " actual=" << actual[i] + << " expected=" << expected[i] << " abs=" << abs_err + << " rel=" << rel_err; + } +} + +std::vector input_values() { + auto input = read_bin(g_dir + "/input.bin"); + EXPECT_EQ(input.size(), static_cast(kWidth)); + return input; +} + +Result> forward_with_suppression( + Module& module, + const TensorPtr& input, + const void* discardable_output_data) { + return with_webgpu_execution_options( + WebGPUExecutionOptions{discardable_output_data, true}, + [&]() { return module.forward({EValue(input)}); }); +} + +TEST(OutputSuppression, DirectFinalQ4SkipsOnlyDiscardedInvocation) { + Module module(g_dir + "/direct_final_q4.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok); + + auto input = make_tensor_ptr({1, kWidth}, input_values()); + std::vector suppressed(kWidth, kPoison); + auto suppressed_tensor = + make_tensor_ptr({1, kWidth}, static_cast(suppressed.data())); + ASSERT_EQ(module.set_output(EValue(suppressed_tensor)), Error::Ok); + const auto result = + forward_with_suppression(module, input, suppressed.data()); + ASSERT_TRUE(result.ok()); + EXPECT_TRUE(std::all_of(suppressed.begin(), suppressed.end(), [](float v) { + return v == kPoison; + })); + + std::vector terminal(kWidth, kPoison); + auto terminal_tensor = + make_tensor_ptr({1, kWidth}, static_cast(terminal.data())); + ASSERT_EQ(module.set_output(EValue(terminal_tensor)), Error::Ok); + const auto terminal_result = module.forward({EValue(input)}); + ASSERT_TRUE(terminal_result.ok()); + expect_matches( + terminal, read_bin(g_dir + "/direct_final_q4.output0.golden.bin")); +} + +TEST(OutputSuppression, Q4FeedingLaterAddIsNotSuppressible) { + Module module(g_dir + "/q4_then_add.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok); + + auto input = make_tensor_ptr({1, kWidth}, input_values()); + std::vector output0(kWidth, kPoison); + std::vector output1(kWidth, kPoison); + auto tensor0 = + make_tensor_ptr({1, kWidth}, static_cast(output0.data())); + auto tensor1 = + make_tensor_ptr({1, kWidth}, static_cast(output1.data())); + ASSERT_EQ(module.set_outputs({EValue(tensor0), EValue(tensor1)}), Error::Ok); + const auto result = forward_with_suppression(module, input, output0.data()); + ASSERT_TRUE(result.ok()); + expect_matches(output0, read_bin(g_dir + "/q4_then_add.output0.golden.bin")); + expect_matches(output1, read_bin(g_dir + "/q4_then_add.output1.golden.bin")); +} + +TEST(OutputSuppression, UnrelatedOutputStillCopies) { + Module module(g_dir + "/unrelated_then_final_q4.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok); + + auto input = make_tensor_ptr({1, kWidth}, input_values()); + std::vector unrelated(kWidth, kPoison); + std::vector suppressed(kWidth, kPoison); + auto unrelated_tensor = + make_tensor_ptr({1, kWidth}, static_cast(unrelated.data())); + auto suppressed_tensor = + make_tensor_ptr({1, kWidth}, static_cast(suppressed.data())); + ASSERT_EQ( + module.set_outputs({EValue(unrelated_tensor), EValue(suppressed_tensor)}), + Error::Ok); + const auto result = + forward_with_suppression(module, input, suppressed.data()); + ASSERT_TRUE(result.ok()); + expect_matches( + unrelated, + read_bin(g_dir + "/unrelated_then_final_q4.output0.golden.bin")); + EXPECT_TRUE(std::all_of(suppressed.begin(), suppressed.end(), [](float v) { + return v == kPoison; + })); +} + +} // namespace + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + g_dir = "/tmp/output_suppression"; + if (argc > 1) { + g_dir = argv[1]; + } + if (const char* env = std::getenv("WEBGPU_OUTPUT_SUPPRESSION_DIR")) { + g_dir = env; + } + + WebGPUContext context; + try { + context = create_webgpu_context(); + } catch (const std::exception& error) { + std::printf("SKIP: no WebGPU device (%s)\n", error.what()); + return 0; + } + set_default_webgpu_context(&context); + const int result = RUN_ALL_TESTS(); + set_default_webgpu_context(nullptr); + destroy_webgpu_context(context); + return result; +} diff --git a/backends/webgpu/test/ops/test_quantized_linear.py b/backends/webgpu/test/ops/test_quantized_linear.py index 54040a54bd7..451f478d871 100644 --- a/backends/webgpu/test/ops/test_quantized_linear.py +++ b/backends/webgpu/test/ops/test_quantized_linear.py @@ -15,6 +15,9 @@ CONFIGS table and reconstructs the identical deterministic ramp input bit-for-bit. """ +import hashlib +import hmac +import json import os import unittest from dataclasses import dataclass @@ -22,8 +25,10 @@ import numpy as np import torch -from executorch.backends.vulkan import VulkanPartitioner -from executorch.exir import to_edge_transform_and_lower +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.exir import ExecutorchBackendConfig, to_edge_transform_and_lower +from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes +from executorch.exir.passes import MemoryPlanningPass from torchao.quantization.granularity import PerGroup from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_ @@ -127,6 +132,149 @@ def _export(m: torch.nn.Module, x: torch.Tensor): ).to_executorch() +OUTPUT_SUPPRESSION_K = 64 +OUTPUT_SUPPRESSION_N = 64 +OUTPUT_SUPPRESSION_GROUP = 32 + + +class _DirectFinalQ4(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.q4 = _make_quantized_model( + OUTPUT_SUPPRESSION_K, + OUTPUT_SUPPRESSION_N, + OUTPUT_SUPPRESSION_GROUP, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.q4(x) + + +class _Q4ThenAdd(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.q4 = _make_quantized_model( + OUTPUT_SUPPRESSION_K, + OUTPUT_SUPPRESSION_N, + OUTPUT_SUPPRESSION_GROUP, + ) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + q4 = self.q4(x) + return q4, q4 + q4 + + +class _UnrelatedThenFinalQ4(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.q4 = _make_quantized_model( + OUTPUT_SUPPRESSION_K, + OUTPUT_SUPPRESSION_N, + OUTPUT_SUPPRESSION_GROUP, + ) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + unrelated = x + x + return unrelated, self.q4(unrelated) + + +class _DuplicateFinalQ4(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.q4 = _make_quantized_model( + OUTPUT_SUPPRESSION_K, + OUTPUT_SUPPRESSION_N, + OUTPUT_SUPPRESSION_GROUP, + ) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + q4 = self.q4(x) + return q4, q4 + + +class _Q4ThenPortableSum(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.q4 = _make_quantized_model( + OUTPUT_SUPPRESSION_K, + OUTPUT_SUPPRESSION_N, + OUTPUT_SUPPRESSION_GROUP, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.q4(x).sum() + + +def _output_suppression_cases() -> list[tuple[str, torch.nn.Module]]: + return [ + ("direct_final_q4", _DirectFinalQ4()), + ("q4_then_add", _Q4ThenAdd()), + ("unrelated_then_final_q4", _UnrelatedThenFinalQ4()), + ] + + +def _verify_output_suppression_certificate( + certificate: dict[str, object], pte: bytes, method: str = "forward" +) -> None: + if certificate.get("schema") != 1: + raise ValueError("invalid output-suppression certificate schema") + if certificate.get("method") != method: + raise ValueError("output-suppression certificate method mismatch") + if certificate.get("delegate_count") != 1: + raise ValueError("output suppression requires exactly one delegate") + if certificate.get("portable_node_count") != 0: + raise ValueError("output suppression rejects portable nodes") + outputs = certificate.get("method_output_nodes") + if not isinstance(outputs, list) or not outputs: + raise ValueError("output-suppression certificate has no method outputs") + if len(set(outputs)) != len(outputs): + raise ValueError("aliased method outputs cannot be suppressed") + actual_hash = hashlib.sha256(pte).hexdigest() + expected_hash = certificate.get("pte_sha256") + if not isinstance(expected_hash, str) or not hmac.compare_digest( + actual_hash, expected_hash + ): + raise ValueError("output-suppression certificate PTE hash mismatch") + + +def _lower_and_certify(model: torch.nn.Module, x: torch.Tensor): + edge = to_edge_transform_and_lower( + torch.export.export(model, (x,)), partitioner=[VulkanPartitioner()] + ) + graph = edge.exported_program().graph_module.graph + delegates = get_delegates(graph) + non_lowered = get_non_lowered_nodes(graph) + output_node = next(node for node in graph.nodes if node.op == "output") + method_outputs = output_node.args[0] + if not isinstance(method_outputs, (tuple, list)): + method_outputs = (method_outputs,) + if len(delegates) != 1: + raise ValueError(f"expected one delegate, got {len(delegates)}") + if non_lowered: + raise ValueError(f"non-lowered nodes: {non_lowered}") + if len({id(output) for output in method_outputs}) != len(method_outputs): + raise ValueError("aliased method outputs cannot be suppressed") + et = edge.to_executorch( + ExecutorchBackendConfig( + memory_planning_pass=MemoryPlanningPass(alloc_graph_output=False) + ) + ) + certificate = { + "schema": 1, + "method": "forward", + "pte_sha256": hashlib.sha256(et.buffer).hexdigest(), + "delegate_count": len(delegates), + "portable_node_count": len(non_lowered), + "method_output_nodes": [output.name for output in method_outputs], + } + _verify_output_suppression_certificate(certificate, et.buffer) + return et, certificate + + +def _lower_fully_delegated(model: torch.nn.Module, x: torch.Tensor): + return _lower_and_certify(model, x)[0] + + class TestQuantizedLinear(unittest.TestCase): def test_export_delegates(self) -> None: # Each (non-heavy) config must fuse to a VulkanBackend delegate (q4gsw); @@ -144,6 +292,37 @@ def test_export_delegates(self) -> None: ) self.assertTrue(found, f"no VulkanBackend delegate in {cfg.name}") + def test_output_suppression_models_are_fully_delegated(self) -> None: + x = _ramp_input(1, OUTPUT_SUPPRESSION_K) + for name, model in _output_suppression_cases(): + with self.subTest(case=name): + et = _lower_fully_delegated(model, x) + delegate_ids = [ + delegate.id + for plan in et.executorch_program.execution_plan + for delegate in plan.delegates + ] + self.assertEqual(delegate_ids, ["VulkanBackend"]) + + def test_output_suppression_rejects_aliased_method_outputs(self) -> None: + x = _ramp_input(1, OUTPUT_SUPPRESSION_K) + with self.assertRaisesRegex(ValueError, "aliased method outputs"): + _lower_fully_delegated(_DuplicateFinalQ4(), x) + + def test_output_suppression_rejects_portable_continuation(self) -> None: + x = _ramp_input(1, OUTPUT_SUPPRESSION_K) + with self.assertRaisesRegex(ValueError, "non-lowered nodes"): + _lower_fully_delegated(_Q4ThenPortableSum(), x) + + def test_output_suppression_certificate_binds_exact_pte(self) -> None: + x = _ramp_input(1, OUTPUT_SUPPRESSION_K) + et, certificate = _lower_and_certify(_DirectFinalQ4(), x) + _verify_output_suppression_certificate(certificate, et.buffer) + with self.assertRaises(ValueError): + _verify_output_suppression_certificate( + certificate, et.buffer + b"different" + ) + def test_golden_matches_eager(self) -> None: # Dual oracle (mirrors SDPA test_golden_matches_eager_op): the fp64 dequant- # matmul truth and torchao's own fp32 quantized forward are independent refs @@ -189,5 +368,37 @@ def export_all_quantized_linear_models( export_quantized_linear_model(cfg, pte, golden) +def export_output_suppression_models(out_dir: str) -> None: + os.makedirs(out_dir, exist_ok=True) + x = _ramp_input(1, OUTPUT_SUPPRESSION_K) + x.numpy().astype(" Date: Wed, 22 Jul 2026 19:11:00 -0700 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 3 +- backends/webgpu/runtime/WebGPUBackend.cpp | 13 +- .../webgpu/runtime/WebGPUExecutionOptions.cpp | 12 +- .../webgpu/runtime/WebGPUExecutionOptions.h | 2 + backends/webgpu/runtime/WebGPUGraph.cpp | 101 +++++---- backends/webgpu/runtime/WebGPUGraph.h | 10 +- .../test/native/test_compute_dispatch.cpp | 206 ++++++++++++++++++ .../test/native/test_execution_options.cpp | 19 +- 8 files changed, 302 insertions(+), 64 deletions(-) create mode 100644 backends/webgpu/test/native/test_compute_dispatch.cpp diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 4c9f1265c4e..a0d206aa9f0 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -27,10 +27,10 @@ endif() set(WEBGPU_SRCS runtime/WebGPUBackend.cpp + runtime/WebGPUExecutionOptions.cpp runtime/WebGPUGraph.cpp runtime/WebGPUDelegateHeader.cpp runtime/WebGPUDevice.cpp - runtime/WebGPUExecutionOptions.cpp runtime/WebGPUQueryPool.cpp runtime/ops/OperatorRegistry.cpp ) @@ -203,7 +203,6 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) target_link_libraries( webgpu_dispatch_2d_test PRIVATE GTest::gtest GTest::gtest_main ) - add_executable( webgpu_execution_options_test test/native/test_execution_options.cpp runtime/WebGPUExecutionOptions.cpp diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index a57d3566724..5cb96595f19 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -171,10 +171,8 @@ Error WebGPUBackend::execute( graph_options = resolve_webgpu_graph_execution_options(delegate_outputs, options); - // Execute the compute graph. Kept inside this try so a plan-validation - // throw from plan_webgpu_execution (or a GPU error) surfaces as - // Error::Internal instead of escaping across the delegate C API boundary. - graph->execute(graph_options); + const WebGPUExecutionPlan plan = graph->make_execution_plan(graph_options); + graph->execute(plan); // Copy outputs from GPU staging buffers to EValue tensor data pointers std::vector> outputs; @@ -184,9 +182,12 @@ Error WebGPUBackend::execute( auto& tensor = args[arg_idx]->toTensor(); outputs.emplace_back(tensor.mutable_data_ptr(), tensor.nbytes()); } - graph->copy_outputs(outputs, graph_options); + graph->copy_outputs(outputs, plan); } catch (const std::exception& e) { - ET_LOG(Error, "WebGPU execute / output copy failed: %s", e.what()); + ET_LOG( + Error, + "WebGPU input preparation / execute / output copy failed: %s", + e.what()); return Error::Internal; } diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.cpp b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp index 2aeede5d32b..67cfb97bbc7 100644 --- a/backends/webgpu/runtime/WebGPUExecutionOptions.cpp +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp @@ -42,7 +42,7 @@ WebGPUExecutionPlan plan_webgpu_execution( WebGPUGraphExecutionOptions options) { std::vector suppressed_dispatches(dispatch_count, false); std::vector copy_outputs(output_count, true); - std::vector suppressed_outputs(output_count, false); + std::vector seen_output_ordinals(output_count, false); for (const auto& output : suppressible_outputs) { if (output.output_ordinal >= output_count || @@ -52,12 +52,12 @@ WebGPUExecutionPlan plan_webgpu_execution( "WebGPU: invalid suppressible output range (output_id " + std::to_string(output.output_id) + ")"); } - if (suppressed_outputs[output.output_ordinal]) { + if (seen_output_ordinals[output.output_ordinal]) { throw std::runtime_error( "WebGPU: duplicate suppressible output (output_id " + std::to_string(output.output_id) + ")"); } - suppressed_outputs[output.output_ordinal] = true; + seen_output_ordinals[output.output_ordinal] = true; if (output.output_ordinal != options.suppress_output_ordinal) { continue; } @@ -101,7 +101,11 @@ WebGPUExecutionPlan plan_webgpu_execution( current_chunk = config.chunk_size; } } - if (plan.dispatch_chunks.empty()) { + if (plan.dispatch_chunks.empty() && + std::any_of( + plan.copy_outputs.begin(), plan.copy_outputs.end(), [](bool copy) { + return copy; + })) { plan.dispatch_chunks.emplace_back(); } return plan; diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.h b/backends/webgpu/runtime/WebGPUExecutionOptions.h index 691cd328e90..d59affd541b 100644 --- a/backends/webgpu/runtime/WebGPUExecutionOptions.h +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.h @@ -19,6 +19,8 @@ constexpr size_t kNoOutputOrdinal = static_cast(-1); struct WebGPUExecutionOptions { // The certificate must bind the exact PTE and method and prove one delegate, // no portable nodes, and a unique leaf method output at this data pointer. + // The caller must keep this pointer valid and unchanged for the complete + // synchronous backend invocation in which these options are scoped. const void* discardable_output_data = nullptr; bool exact_method_certificate_verified = false; }; diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 5a567a36500..1381449aa90 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -1730,15 +1730,23 @@ bool should_timestamp_query() { } } // namespace -void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { - const size_t n = dispatches_.size(); - const size_t chunk = execute_config_.chunk_size; - const WebGPUExecutionPlan plan = plan_webgpu_execution( - n, +WebGPUExecutionPlan WebGPUGraph::make_execution_plan( + const WebGPUGraphExecutionOptions& options) const { + return plan_webgpu_execution( + dispatches_.size(), output_copies_.size(), execute_config_, suppressible_outputs_, options); +} + +size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { + const size_t n = dispatches_.size(); + const size_t chunk = execute_config_.chunk_size; + + if (plan.dispatch_chunks.empty()) { + return 0; + } if (chunk == 0 || n <= chunk) { #ifdef WGPU_BACKEND_ENABLE_PROFILING @@ -1763,45 +1771,47 @@ void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { wgpuDeviceCreateCommandEncoder(device_, &enc_desc); // One pass per dispatch: enforces storage RAW ordering across deps. - for (size_t i : plan.dispatch_chunks.front()) { - const auto& dispatch = dispatches_[i]; - if (dispatch.kind == WebGPUDispatch::Kind::Copy) { - wgpuCommandEncoderCopyBufferToBuffer( - encoder, - dispatch.copy_src, - 0, - dispatch.copy_dst, - 0, - dispatch.copy_nbytes); - continue; - } - WGPUComputePassDescriptor pass_desc = {}; + for (const auto& dispatch_chunk : plan.dispatch_chunks) { + for (size_t i : dispatch_chunk) { + const auto& dispatch = dispatches_[i]; + if (dispatch.kind == WebGPUDispatch::Kind::Copy) { + wgpuCommandEncoderCopyBufferToBuffer( + encoder, + dispatch.copy_src, + 0, + dispatch.copy_dst, + 0, + dispatch.copy_nbytes); + continue; + } + WGPUComputePassDescriptor pass_desc = {}; #ifdef WGPU_BACKEND_ENABLE_PROFILING - // tw must outlive BeginComputePass (the descriptor points at it). - WGPUPassTimestampWrites tw = {}; - if (qp) { - tw = qp->writes_for(static_cast(i)); - pass_desc.timestampWrites = &tw; - } + // tw must outlive BeginComputePass (the descriptor points at it). + WGPUPassTimestampWrites tw = {}; + if (qp) { + tw = qp->writes_for(static_cast(i)); + pass_desc.timestampWrites = &tw; + } #endif // WGPU_BACKEND_ENABLE_PROFILING - WGPUComputePassEncoder pass = - wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); - wgpuComputePassEncoderSetPipeline(pass, dispatch.pipeline); - wgpuComputePassEncoderSetBindGroup( - pass, 0, dispatch.bind_group, 0, nullptr); - wgpuComputePassEncoderDispatchWorkgroups( - pass, dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1); - wgpuComputePassEncoderEnd(pass); - wgpuComputePassEncoderRelease(pass); + WGPUComputePassEncoder pass = + wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); + wgpuComputePassEncoderSetPipeline(pass, dispatch.pipeline); + wgpuComputePassEncoderSetBindGroup( + pass, 0, dispatch.bind_group, 0, nullptr); + wgpuComputePassEncoderDispatchWorkgroups( + pass, dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1); + wgpuComputePassEncoderEnd(pass); + wgpuComputePassEncoderRelease(pass); #ifdef WGPU_BACKEND_ENABLE_PROFILING - if (qp) { - qp->record( - static_cast(i), - dispatch.kernel_name, - {dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1}, - {1, 1, 1}); - } + if (qp) { + qp->record( + static_cast(i), + dispatch.kernel_name, + {dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1}, + {1, 1, 1}); + } #endif // WGPU_BACKEND_ENABLE_PROFILING + } } for (size_t i = 0; i < output_copies_.size(); i++) { @@ -1832,7 +1842,7 @@ void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { qp->print_results(); } #endif // WGPU_BACKEND_ENABLE_PROFILING - return; + return 1; } // GPU timestamp queries assume one submit; chunked execute is multi-submit. @@ -1892,6 +1902,7 @@ void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { wgpuCommandBufferRelease(cmd); wgpuCommandEncoderRelease(encoder); } + return plan.dispatch_chunks.size(); } namespace { @@ -1913,14 +1924,8 @@ void buffer_map_callback( void WebGPUGraph::copy_outputs( std::vector>& outputs, - const WebGPUGraphExecutionOptions& options) { + const WebGPUExecutionPlan& plan) { const size_t count = std::min(outputs.size(), output_staging_buffers_.size()); - const WebGPUExecutionPlan plan = plan_webgpu_execution( - dispatches_.size(), - output_copies_.size(), - execute_config_, - suppressible_outputs_, - options); std::vector cb_data(count); std::vector map_futures(count, WGPUFuture{}); diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index c520cb06f84..7f194a71400 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -107,14 +107,18 @@ class WebGPUGraph { // Copy input tensor data from host pointers into GPU buffers. void copy_inputs(const std::vector& inputs); - // Execute all recorded dispatches. - void execute(const WebGPUGraphExecutionOptions& options); + WebGPUExecutionPlan make_execution_plan( + const WebGPUGraphExecutionOptions& options) const; + + // Execute the dispatches selected by a plan created for this graph. Returns + // the number of GPU queue submissions performed. + size_t execute(const WebGPUExecutionPlan& plan); // Copy output tensor data from GPU buffers back to host pointers. // Uses mapAsync + ASYNCIFY in Wasm. void copy_outputs( std::vector>& outputs, - const WebGPUGraphExecutionOptions& options); + const WebGPUExecutionPlan& plan); const std::vector& input_ids() const { return input_ids_; diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp new file mode 100644 index 00000000000..22bade723bf --- /dev/null +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -0,0 +1,206 @@ +/* + * 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. + */ + +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { +namespace { + +// NOLINTBEGIN(cppcoreguidelines-avoid-non-const-global-variables) +WGPUDevice g_device = nullptr; +// NOLINTEND(cppcoreguidelines-avoid-non-const-global-variables) + +struct SigmoidParams { + uint32_t num_elements; + uint32_t padding[3]; +}; + +WGPUBuffer create_storage_buffer(size_t nbytes) { + WGPUBufferDescriptor descriptor = {}; + descriptor.size = nbytes; + descriptor.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst; + WGPUBuffer buffer = wgpuDeviceCreateBuffer(g_device, &descriptor); + if (buffer == nullptr) { + throw std::runtime_error("failed to create compute-dispatch test buffer"); + } + return buffer; +} + +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); +} + +TEST(WebGPUComputeDispatch, PipelineKeyCanonicalizesConstants) { + WebGPUComputeDispatchDescriptor first; + first.shader_name = "sigmoid"; + first.entry_point = "main"; + first.constants = {{"beta", 2.0}, {"alpha", 1.0}}; + + WebGPUComputeDispatchDescriptor reordered = first; + reordered.constants = {{"alpha", 1.0}, {"beta", 2.0}}; + + EXPECT_EQ( + make_compute_pipeline_key(first), make_compute_pipeline_key(reordered)); +} + +TEST(WebGPUComputeDispatch, PipelineKeyTracksCompileIdentityOnly) { + WebGPUComputeDispatchDescriptor base; + base.shader_name = "sigmoid"; + base.entry_point = "main"; + base.constants = {{"wg_size", 64.0}}; + base.grid = {1u, 1u}; + base.bindings = {{reinterpret_cast(1), 0u, 16u}}; + + WebGPUComputeDispatchDescriptor runtime_change = base; + runtime_change.grid = {17u, 3u}; + runtime_change.bindings = {{reinterpret_cast(2), 128u, 4096u}}; + EXPECT_EQ( + make_compute_pipeline_key(base), + make_compute_pipeline_key(runtime_change)); + + WebGPUComputeDispatchDescriptor shader_change = base; + shader_change.shader_name = "binary_add"; + EXPECT_NE( + make_compute_pipeline_key(base), + make_compute_pipeline_key(shader_change)); + + WebGPUComputeDispatchDescriptor entry_change = base; + entry_change.entry_point = "alternate"; + EXPECT_NE( + make_compute_pipeline_key(base), make_compute_pipeline_key(entry_change)); + + WebGPUComputeDispatchDescriptor constant_change = base; + constant_change.constants = {{"wg_size", 128.0}}; + EXPECT_NE( + make_compute_pipeline_key(base), + make_compute_pipeline_key(constant_change)); +} + +TEST(WebGPUComputeDispatch, PipelineKeyRejectsInvalidConstants) { + WebGPUComputeDispatchDescriptor duplicate; + duplicate.shader_name = "sigmoid"; + duplicate.constants = {{"wg_size", 64.0}, {"wg_size", 128.0}}; + EXPECT_THROW(make_compute_pipeline_key(duplicate), std::runtime_error); + + WebGPUComputeDispatchDescriptor non_finite; + non_finite.shader_name = "sigmoid"; + non_finite.constants = {{"wg_size", std::numeric_limits::infinity()}}; + EXPECT_THROW(make_compute_pipeline_key(non_finite), std::runtime_error); +} + +TEST(WebGPUComputeDispatch, DescriptorRejectsInvalidBindings) { + WebGPUComputeDispatchDescriptor null_buffer; + null_buffer.shader_name = "sigmoid"; + null_buffer.bindings = {{nullptr, 0u, 16u}}; + EXPECT_THROW( + validate_compute_dispatch_descriptor(null_buffer), std::runtime_error); + + WebGPUComputeDispatchDescriptor zero_size; + zero_size.shader_name = "sigmoid"; + zero_size.bindings = {{reinterpret_cast(1), 0u, 0u}}; + EXPECT_THROW( + validate_compute_dispatch_descriptor(zero_size), std::runtime_error); + + WebGPUComputeDispatchDescriptor overflow; + overflow.shader_name = "sigmoid"; + overflow.bindings = { + {reinterpret_cast(1), + std::numeric_limits::max(), + 2u}}; + EXPECT_THROW( + validate_compute_dispatch_descriptor(overflow), std::runtime_error); +} + +TEST(WebGPUComputeDispatch, ReusesPipelineAndReleasesDawnObjects) { + constexpr size_t kNumElements = 64; + constexpr size_t kBufferBytes = kNumElements * sizeof(float); + + for (int iteration = 0; iteration < 32; iteration++) { + WGPUBuffer input = create_storage_buffer(kBufferBytes); + WGPUBuffer output = create_storage_buffer(kBufferBytes); + { + WebGPUGraph graph; + graph.set_device(g_device); + WGPUBuffer params = + graph.create_params_buffer(SigmoidParams{kNumElements, {0, 0, 0}}); + + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = "sigmoid"; + descriptor.kernel_name = "sigmoid_test"; + descriptor.bindings = { + {input, 0u, kBufferBytes}, + {output, 0u, kBufferBytes}, + {params, 0u, sizeof(SigmoidParams)}}; + descriptor.constants = {{"wg_size", 64.0}}; + descriptor.grid = {1u, 1u}; + + graph.add_compute_dispatch(descriptor); + graph.add_compute_dispatch(descriptor); + + const WebGPUMemoryStats stats = graph.memory_stats(); + EXPECT_EQ(stats.num_dispatches, 2); + EXPECT_EQ(stats.num_cached_shaders, 1); + EXPECT_EQ(stats.num_cached_pipelines, 1); + } + wgpuBufferRelease(output); + wgpuBufferRelease(input); + } +} + +TEST(WebGPUComputeDispatch, RejectsBindingRangeBeyondDawnBuffer) { + WGPUBuffer buffer = create_storage_buffer(16u); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = "sigmoid"; + descriptor.bindings = {{buffer, 8u, 12u}}; + + EXPECT_THROW( + validate_compute_dispatch_descriptor(descriptor), std::runtime_error); + wgpuBufferRelease(buffer); +} + +TEST(WebGPUExecution, FullySuppressedPlanPerformsNoQueueSubmission) { + WebGPUGraph graph; + const WebGPUExecutionPlan plan; + + EXPECT_EQ(graph.execute(plan), 0u); +} + +} // namespace +} // namespace executorch::backends::webgpu + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + + executorch::backends::webgpu::WebGPUContext context; + try { + context = executorch::backends::webgpu::create_webgpu_context(); + } catch (const std::exception& error) { + std::printf("SKIP: %s\n", error.what()); + return 0; + } + executorch::backends::webgpu::set_default_webgpu_context(&context); + executorch::backends::webgpu::g_device = context.device; + + const int result = RUN_ALL_TESTS(); + executorch::backends::webgpu::set_default_webgpu_context(nullptr); + executorch::backends::webgpu::destroy_webgpu_context(context); + return result; +} diff --git a/backends/webgpu/test/native/test_execution_options.cpp b/backends/webgpu/test/native/test_execution_options.cpp index 2f96636438d..3b58c48534b 100644 --- a/backends/webgpu/test/native/test_execution_options.cpp +++ b/backends/webgpu/test/native/test_execution_options.cpp @@ -54,7 +54,7 @@ TEST(WebGPUExecutionOptionsTest, ExceptionRestoresPriorValue) { current_webgpu_execution_options().discardable_output_data, nullptr); } -TEST(WebGPUExecutionOptionsTest, ErrorReturnRestoresPriorValue) { +TEST(WebGPUExecutionOptionsTest, BooleanReturnRestoresPriorValue) { int output = 0; const bool result = with_webgpu_execution_options({&output, true}, []() { return false; }); @@ -123,5 +123,22 @@ TEST(WebGPUExecutionPlanTest, RejectsInvalidSuppressibleRange) { std::runtime_error); } +TEST(WebGPUExecutionPlanTest, AllSuppressedHasNoSyntheticDispatchChunk) { + const std::vector suppressible = {{9, 0, 0, 2}}; + const WebGPUExecutionPlan plan = plan_webgpu_execution( + 2, 1, ExecuteConfig{}, suppressible, WebGPUGraphExecutionOptions{0}); + + EXPECT_TRUE(plan.dispatch_chunks.empty()); + EXPECT_EQ(plan.copy_outputs, (std::vector{false})); +} + +TEST(WebGPUExecutionPlanTest, CopyOnlyPlanRetainsOneSubmissionChunk) { + const WebGPUExecutionPlan plan = plan_webgpu_execution( + 0, 1, ExecuteConfig{}, {}, WebGPUGraphExecutionOptions{}); + + EXPECT_EQ(plan.dispatch_chunks, (std::vector>{{}})); + EXPECT_EQ(plan.copy_outputs, (std::vector{true})); +} + } // namespace } // namespace executorch::backends::webgpu