diff --git a/backends/vulkan/serialization/targets.bzl b/backends/vulkan/serialization/targets.bzl index caf18b255bc..f8a0712d3f3 100644 --- a/backends/vulkan/serialization/targets.bzl +++ b/backends/vulkan/serialization/targets.bzl @@ -26,6 +26,7 @@ def define_common_targets(is_fbcode = False): name = "vk_delegate_schema", srcs = [], visibility = [ + "//executorch/backends/webgpu/...", "//executorch/backends/vulkan/...", ], exported_headers = { diff --git a/backends/webgpu/BUCK b/backends/webgpu/BUCK new file mode 100644 index 00000000000..4e76682cd45 --- /dev/null +++ b/backends/webgpu/BUCK @@ -0,0 +1,15 @@ +# 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. + +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") + +oncall("executorch") + +load(":targets.bzl", "define_common_targets") + +non_fbcode_target(_kind = define_common_targets) + +fbcode_target(_kind = define_common_targets) diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 536348ca69f..adeea976fac 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -47,6 +47,17 @@ list(APPEND WEBGPU_SRCS ${WEBGPU_OP_SRCS}) add_library(webgpu_backend ${WEBGPU_SRCS}) +add_library(webgpu_model_loader runner/webgpu_model_loader.cpp) +target_include_directories( + webgpu_model_loader PUBLIC $ +) +target_link_libraries(webgpu_model_loader PUBLIC extension_module_static) +target_compile_options(webgpu_model_loader PRIVATE -fexceptions) +target_compile_definitions( + webgpu_model_loader PUBLIC C10_USING_CUSTOM_GENERATED_MACROS +) +set_property(TARGET webgpu_model_loader PROPERTY CXX_STANDARD 17) + # Verify committed *_wgsl.h match their *.wgsl (drift fails the build). resolve_python_executable() add_custom_target( @@ -165,8 +176,8 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) PRIVATE "${EXECUTORCH_ROOT}/third-party/json/single_include" ) - # Device-free util unit test: no backend/Dawn link (pure manifest/tolerance - # + dispatch-grid-math helpers), so it does NOT use the native-test helper. + # Device-free util unit test. WebGPUUtils.h needs Dawn declarations, but the + # test calls only pure helpers and does not link the Dawn implementation. add_executable( webgpu_op_test_util_test test/op_tests/test_driver_util.cpp test/op_tests/driver_util.cpp @@ -178,7 +189,13 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) "${EXECUTORCH_ROOT}/third-party/json/single_include" ) target_link_libraries( - webgpu_op_test_util_test PRIVATE GTest::gtest GTest::gtest_main + webgpu_op_test_util_test PRIVATE executorch_core GTest::gtest + GTest::gtest_main + ) + target_include_directories( + webgpu_op_test_util_test + PRIVATE + "$" ) target_compile_options(webgpu_op_test_util_test PRIVATE -fexceptions) set_property(TARGET webgpu_op_test_util_test PROPERTY CXX_STANDARD 17) @@ -235,5 +252,16 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) webgpu_compute_dispatch_test test/native/test_compute_dispatch.cpp ) target_link_libraries(webgpu_compute_dispatch_test PRIVATE GTest::gtest) + + add_executable(webgpu_model_loader_test test/native/test_model_loader.cpp) + target_include_directories( + webgpu_model_loader_test PRIVATE $ + ) + target_link_libraries( + webgpu_model_loader_test PRIVATE webgpu_model_loader GTest::gtest + GTest::gtest_main + ) + target_compile_options(webgpu_model_loader_test PRIVATE -fexceptions) + set_property(TARGET webgpu_model_loader_test PROPERTY CXX_STANDARD 17) endif() endif() diff --git a/backends/webgpu/runner/webgpu_model_loader.cpp b/backends/webgpu/runner/webgpu_model_loader.cpp new file mode 100644 index 00000000000..97148c14d61 --- /dev/null +++ b/backends/webgpu/runner/webgpu_model_loader.cpp @@ -0,0 +1,46 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +runtime::Result> load_webgpu_model( + WebGPUModelLoadSpec spec) { + if (spec.pte_path.empty() || spec.required_methods.empty()) { + return runtime::Error::InvalidArgument; + } + std::unordered_set methods; + for (const auto& method : spec.required_methods) { + if (method.empty() || !methods.insert(method).second) { + return runtime::Error::InvalidArgument; + } + } + std::unordered_set data_files; + for (const auto& path : spec.ptd_paths) { + if (path.empty() || path == spec.pte_path || + !data_files.insert(path).second) { + return runtime::Error::InvalidArgument; + } + } + + auto module = std::make_unique( + spec.pte_path, std::move(spec.ptd_paths), spec.load_mode); + for (const auto& method : spec.required_methods) { + const runtime::Error error = module->load_method(method); + if (error != runtime::Error::Ok) { + return error; + } + } + return module; +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runner/webgpu_model_loader.h b/backends/webgpu/runner/webgpu_model_loader.h new file mode 100644 index 00000000000..1aff780ed94 --- /dev/null +++ b/backends/webgpu/runner/webgpu_model_loader.h @@ -0,0 +1,30 @@ +/* + * 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 +#include +#include + +namespace executorch::backends::webgpu { + +struct WebGPUModelLoadSpec { + std::string pte_path; + std::vector ptd_paths; + std::vector required_methods; + extension::Module::LoadMode load_mode = extension::Module::LoadMode::File; +}; + +runtime::Result> load_webgpu_model( + WebGPUModelLoadSpec spec); + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index 35d4225bc7c..5e034137dab 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -44,6 +44,16 @@ using executorch::runtime::resize_tensor; using executorch::runtime::Result; using executorch::runtime::Span; +namespace { +thread_local WebGPUGraph* last_execution_graph = nullptr; +} // namespace + +std::string webgpu_backend_execution_attestation_json() { + return last_execution_graph == nullptr + ? "{\"schemaVersion\":1,\"unavailable\":true}" + : last_execution_graph->execution_attestation_json(); +} + Result parse_webgpu_graph_config( ArrayRef compile_specs) { WebGPUGraphConfig config; @@ -228,6 +238,7 @@ Error WebGPUBackend::execute( // the backend boundary. try { const WebGPUExecutionPlan plan = graph->make_execution_plan(graph_options); + last_execution_graph = graph; graph->execute(plan); // Copy outputs from GPU staging buffers to EValue tensor data pointers @@ -242,7 +253,9 @@ Error WebGPUBackend::execute( {tensor.mutable_data_ptr(), tensor.nbytes(), host_is_fp32}); } graph->copy_outputs(outputs, plan); + graph->complete_execution_attestation(); } catch (const std::exception& e) { + graph->fail_execution_attestation(e.what()); ET_LOG(Error, "WebGPU execute / output copy failed: %s", e.what()); return Error::Internal; } @@ -253,6 +266,9 @@ Error WebGPUBackend::execute( void WebGPUBackend::destroy(DelegateHandle* handle) const { if (handle != nullptr) { WebGPUGraph* graph = static_cast(handle); + if (last_execution_graph == graph) { + last_execution_graph = nullptr; + } graph->~WebGPUGraph(); } } diff --git a/backends/webgpu/runtime/WebGPUBackend.h b/backends/webgpu/runtime/WebGPUBackend.h index 59f7e33994c..f8320a631a8 100644 --- a/backends/webgpu/runtime/WebGPUBackend.h +++ b/backends/webgpu/runtime/WebGPUBackend.h @@ -19,6 +19,8 @@ executorch::runtime::Result parse_webgpu_graph_config( executorch::runtime::ArrayRef compile_specs); +std::string webgpu_backend_execution_attestation_json(); + class WebGPUBackend final : public ::executorch::runtime::BackendInterface { public: ~WebGPUBackend() override = default; diff --git a/backends/webgpu/runtime/WebGPUDevice.cpp b/backends/webgpu/runtime/WebGPUDevice.cpp index d4b148cda5f..a61cc767790 100644 --- a/backends/webgpu/runtime/WebGPUDevice.cpp +++ b/backends/webgpu/runtime/WebGPUDevice.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include #include @@ -184,16 +185,31 @@ WebGPUContext create_webgpu_context() { } namespace { -WebGPUContext* g_default_context = nullptr; +std::atomic g_default_context{nullptr}; } // namespace void set_default_webgpu_context(WebGPUContext* ctx) { - g_default_context = ctx; + g_default_context.store(ctx, std::memory_order_release); +} + +WebGPUContext* get_explicit_default_webgpu_context() { + return g_default_context.load(std::memory_order_acquire); +} + +bool compare_and_set_default_webgpu_context( + WebGPUContext* expected, + WebGPUContext* desired) { + return g_default_context.compare_exchange_strong( + expected, + desired, + std::memory_order_acq_rel, + std::memory_order_acquire); } WebGPUContext* get_default_webgpu_context() { - if (g_default_context) { - return g_default_context; + if (WebGPUContext* explicit_context = + get_explicit_default_webgpu_context()) { + return explicit_context; } #if !defined(__EMSCRIPTEN__) // Native-only lazy process-wide context, mirroring Vulkan api::context(). diff --git a/backends/webgpu/runtime/WebGPUDevice.h b/backends/webgpu/runtime/WebGPUDevice.h index 12f73c969a7..f9c4c57a56a 100644 --- a/backends/webgpu/runtime/WebGPUDevice.h +++ b/backends/webgpu/runtime/WebGPUDevice.h @@ -41,6 +41,13 @@ void destroy_webgpu_context(WebGPUContext& ctx); // Global context used by WebGPUGraph::build() when no device is pre-set. void set_default_webgpu_context(WebGPUContext* ctx); +// Returns only a caller-installed context, never the native lazy fallback. +WebGPUContext* get_explicit_default_webgpu_context(); +// Replaces the explicit context only when its current pointer equals expected. +// Registration is non-owning; the caller keeps the installed context alive. +bool compare_and_set_default_webgpu_context( + WebGPUContext* expected, + WebGPUContext* desired); WebGPUContext* get_default_webgpu_context(); } // namespace webgpu diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.cpp b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp index 19c0892030f..e022d977be8 100644 --- a/backends/webgpu/runtime/WebGPUExecutionOptions.cpp +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include #include @@ -18,8 +19,217 @@ namespace { thread_local WebGPUExecutionOptions execution_options; +void append_json_string(std::ostringstream& out, const std::string& value) { + out << '"'; + for (const unsigned char c : value) { + switch (c) { + case '"': + out << "\\\""; + break; + case '\\': + out << "\\\\"; + break; + case '\b': + out << "\\b"; + break; + case '\f': + out << "\\f"; + break; + case '\n': + out << "\\n"; + break; + case '\r': + out << "\\r"; + break; + case '\t': + out << "\\t"; + break; + default: + if (c < 0x20) { + constexpr char kHex[] = "0123456789abcdef"; + out << "\\u00" << kHex[c >> 4] << kHex[c & 0x0f]; + } else { + out << static_cast(c); + } + } + } + out << '"'; +} + +const char* json_bool(bool value) { + return value ? "true" : "false"; +} + } // namespace +WebGPUCommandInventory build_webgpu_command_inventory( + const std::vector& commands) { + WebGPUCommandInventory inventory; + std::vector run_ordinals(commands.size(), -1); + bool inside_compute_run = false; + for (size_t i = 0; i < commands.size(); i++) { + const auto& command = commands[i]; + if (command.kind != WebGPUCommandKind::OutputCopy) { + ++inventory.static_dispatch_records; + } + if (command.kind == WebGPUCommandKind::Compute) { + if (command.zero_grid) { + ++inventory.zero_grid_compute_count; + } + if (!command.enabled || command.zero_grid) { + continue; + } + if (!inside_compute_run) { + ++inventory.maximal_compute_runs; + inside_compute_run = true; + } + run_ordinals[i] = + static_cast(inventory.maximal_compute_runs - 1); + ++inventory.active_compute_count; + continue; + } + if (command.kind == WebGPUCommandKind::GraphCopy) { + if (command.enabled) { + ++inventory.graph_copy_count; + inside_compute_run = false; + } + continue; + } + if (command.enabled && !command.suppressed) { + ++inventory.output_copy_count; + } + inside_compute_run = false; + } + + std::vector preceding_runs(commands.size(), -1); + std::vector following_runs(commands.size(), -1); + int64_t last_run = -1; + for (size_t i = 0; i < commands.size(); i++) { + preceding_runs[i] = last_run; + if (run_ordinals[i] >= 0) { + last_run = run_ordinals[i]; + } + } + int64_t next_run = -1; + for (size_t i = commands.size(); i > 0; i--) { + following_runs[i - 1] = next_run; + if (run_ordinals[i - 1] >= 0) { + next_run = run_ordinals[i - 1]; + } + } + + std::ostringstream out; + out << "{\"commands\":["; + for (size_t i = 0; i < commands.size(); i++) { + if (i != 0) { + out << ','; + } + const auto& command = commands[i]; + if (command.kind == WebGPUCommandKind::Compute) { + out << "{\"enabled\":" << json_bool(command.enabled) + << ",\"expectedMaximalRunOrdinal\":" << run_ordinals[i] + << ",\"grid\":[" << command.workgroup_count_x << ',' + << command.workgroup_count_y << ",1],\"identity\":"; + append_json_string(out, command.identity); + out << ",\"kind\":\"compute\",\"staticDispatchIndex\":" + << command.static_dispatch_index << ",\"zeroGrid\":" + << json_bool(command.zero_grid) << '}'; + } else if (command.kind == WebGPUCommandKind::GraphCopy) { + out << "{\"byteCount\":" << command.byte_count + << ",\"destinationIdentity\":" << command.destination_identity + << ",\"destinationOffset\":" << command.destination_offset + << ",\"enabled\":" << json_bool(command.enabled) + << ",\"followingMaximalRunOrdinal\":" << following_runs[i] + << ",\"kind\":\"graph_copy\",\"precedingMaximalRunOrdinal\":" + << preceding_runs[i] << ",\"sourceIdentity\":" + << command.source_identity << ",\"sourceOffset\":" + << command.source_offset << ",\"staticDispatchIndex\":" + << command.static_dispatch_index << '}'; + } else { + out << "{\"byteCount\":" << command.byte_count + << ",\"destinationIdentity\":" << command.destination_identity + << ",\"destinationOffset\":" << command.destination_offset + << ",\"enabled\":" << json_bool(command.enabled) + << ",\"kind\":\"output_copy\",\"outputOrdinal\":" + << command.output_ordinal << ",\"sourceIdentity\":" + << command.source_identity << ",\"sourceOffset\":" + << command.source_offset << ",\"suppressed\":" + << json_bool(command.suppressed) << '}'; + } + } + out << "],\"schemaVersion\":1}"; + inventory.canonical_commands_json = out.str(); + return inventory; +} + +std::string serialize_webgpu_execution_attestation( + const WebGPUExecutionAttestation& attestation) { + const auto& inventory = attestation.inventory; + std::ostringstream out; + out << "{\"activeComputeCount\":" << inventory.active_compute_count + << ",\"applied\":" << json_bool(attestation.applied) + << ",\"canonicalCommands\":" + << (inventory.canonical_commands_json.empty() + ? "{\"commands\":[],\"schemaVersion\":1}" + : inventory.canonical_commands_json) + << ",\"completed\":" << json_bool(attestation.completed) + << ",\"encodedComputePasses\":" + << attestation.encoded_compute_passes << ",\"errorReason\":"; + append_json_string(out, attestation.error_reason); + out << ",\"executionOrdinal\":" << attestation.execution_ordinal + << ",\"graphCopyCount\":" << inventory.graph_copy_count + << ",\"maxComputeDispatchesPerPass\":" + << attestation.max_compute_dispatches_per_pass + << ",\"maximalComputeRuns\":" << inventory.maximal_compute_runs + << ",\"outputCopyCount\":" << inventory.output_copy_count + << ",\"queueSubmitCount\":" << attestation.queue_submit_count + << ",\"requested\":" << json_bool(attestation.requested) + << ",\"schemaVersion\":1,\"staticDispatchRecords\":" + << inventory.static_dispatch_records << ",\"zeroGridComputeCount\":" + << inventory.zero_grid_compute_count << '}'; + return out.str(); +} + +bool webgpu_pass_cap_reached( + size_t dispatches_in_current_pass, + size_t max_compute_dispatches_per_pass) { + return max_compute_dispatches_per_pass != 0 && + dispatches_in_current_pass >= max_compute_dispatches_per_pass; +} + +size_t count_webgpu_compute_passes( + const std::vector& commands, + bool single_compute_pass, + size_t max_compute_dispatches_per_pass) { + if (!single_compute_pass && max_compute_dispatches_per_pass != 0) { + throw std::invalid_argument( + "WebGPU: pass cap requires single_compute_pass"); + } + size_t pass_count = 0; + size_t dispatches_in_pass = 0; + for (const auto& command : commands) { + if (command.kind != WebGPUCommandKind::Compute) { + if (command.enabled) { + dispatches_in_pass = 0; + } + continue; + } + if (!command.enabled || command.zero_grid) { + continue; + } + if (!single_compute_pass || dispatches_in_pass == 0) { + ++pass_count; + } + ++dispatches_in_pass; + if (!single_compute_pass || webgpu_pass_cap_reached( + dispatches_in_pass, + max_compute_dispatches_per_pass)) { + dispatches_in_pass = 0; + } + } + return pass_count; +} + WebGPUExecutionOptions current_webgpu_execution_options() { return execution_options; } @@ -78,6 +288,9 @@ WebGPUExecutionPlan plan_webgpu_execution( WebGPUExecutionPlan plan; plan.copy_outputs = std::move(copy_outputs); + plan.single_compute_pass = options.single_compute_pass; + plan.max_compute_dispatches_per_pass = + options.max_compute_dispatches_per_pass; auto append_chunk = [&](size_t begin, size_t end) { std::vector indices; @@ -120,11 +333,15 @@ WebGPUExecutionPlan plan_webgpu_execution( WebGPUGraphExecutionOptions resolve_webgpu_graph_execution_options( const std::vector& delegate_outputs, WebGPUExecutionOptions options) { + WebGPUGraphExecutionOptions resolved; + resolved.single_compute_pass = options.single_compute_pass; + resolved.max_compute_dispatches_per_pass = + options.max_compute_dispatches_per_pass; if (options.discardable_output_data == nullptr) { - return {}; + return resolved; } if (!options.exact_method_certificate_verified) { - return {}; + return resolved; } size_t match = kNoOutputOrdinal; @@ -133,11 +350,12 @@ WebGPUGraphExecutionOptions resolve_webgpu_graph_execution_options( continue; } if (match != kNoOutputOrdinal) { - return {}; + return resolved; } match = i; } - return {match}; + resolved.suppress_output_ordinal = match; + return resolved; } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.h b/backends/webgpu/runtime/WebGPUExecutionOptions.h index 304f46839d8..ce14dfc204f 100644 --- a/backends/webgpu/runtime/WebGPUExecutionOptions.h +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.h @@ -9,6 +9,8 @@ #pragma once #include +#include +#include #include #include @@ -23,10 +25,14 @@ struct WebGPUExecutionOptions { // synchronous backend invocation in which these options are scoped. const void* discardable_output_data = nullptr; bool exact_method_certificate_verified = false; + bool single_compute_pass = false; + size_t max_compute_dispatches_per_pass = 0; }; struct WebGPUGraphExecutionOptions { size_t suppress_output_ordinal = kNoOutputOrdinal; + bool single_compute_pass = false; + size_t max_compute_dispatches_per_pass = 0; }; struct ExecuteConfig { @@ -44,8 +50,66 @@ struct SuppressibleOutput { struct WebGPUExecutionPlan { std::vector> dispatch_chunks; std::vector copy_outputs; + bool single_compute_pass = false; + size_t max_compute_dispatches_per_pass = 0; }; +enum class WebGPUCommandKind { Compute, GraphCopy, OutputCopy }; + +struct WebGPUCommandRecord { + WebGPUCommandKind kind = WebGPUCommandKind::Compute; + size_t static_dispatch_index = 0; + size_t output_ordinal = kNoOutputOrdinal; + std::string identity; + bool enabled = true; + bool zero_grid = false; + bool suppressed = false; + uint32_t workgroup_count_x = 1; + uint32_t workgroup_count_y = 1; + int64_t source_identity = -1; + int64_t destination_identity = -1; + size_t source_offset = 0; + size_t destination_offset = 0; + size_t byte_count = 0; +}; + +struct WebGPUCommandInventory { + size_t static_dispatch_records = 0; + size_t active_compute_count = 0; + size_t zero_grid_compute_count = 0; + size_t graph_copy_count = 0; + size_t output_copy_count = 0; + size_t maximal_compute_runs = 0; + std::string canonical_commands_json; +}; + +struct WebGPUExecutionAttestation { + uint64_t execution_ordinal = 0; + bool requested = false; + bool applied = false; + bool completed = false; + size_t encoded_compute_passes = 0; + size_t queue_submit_count = 0; + size_t max_compute_dispatches_per_pass = 0; + std::string error_reason; + WebGPUCommandInventory inventory; +}; + +bool webgpu_pass_cap_reached( + size_t dispatches_in_current_pass, + size_t max_compute_dispatches_per_pass); + +size_t count_webgpu_compute_passes( + const std::vector& commands, + bool single_compute_pass, + size_t max_compute_dispatches_per_pass); + +WebGPUCommandInventory build_webgpu_command_inventory( + const std::vector& commands); + +std::string serialize_webgpu_execution_attestation( + const WebGPUExecutionAttestation& attestation); + WebGPUExecutionPlan plan_webgpu_execution( size_t dispatch_count, size_t output_count, diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index d1e1d625ad7..b7332a1cd9c 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -35,6 +35,9 @@ namespace executorch::backends::webgpu { namespace { +constexpr int kHostBiasProjectionDim = -2; +constexpr int kIdentityPrepackConstantProjectionDim = -3; + const uint8_t* checked_inline_constant( const uint8_t* data, size_t data_size, @@ -520,29 +523,43 @@ void WebGPUGraph::update_symints_from_inputs( } // Live cur_dims: the source may be a dynamic-shape input. const auto& dims = tensors_[src.input_tensor_id].cur_dims; - int dim = normalize_dim( - src.dim, static_cast(dims.size()), "select_as_symint"); - int index = src.index; - if (index < 0) { - index += static_cast(dims[dim]); - } - if (index < 0 || index >= static_cast(dims[dim])) { - throw std::runtime_error("select_as_symint: index out of range"); + if (src.indices.size() != dims.size()) { + throw std::runtime_error("select_as_symint: source rank changed"); } - int64_t numel = 1; - for (int64_t d : dims) { - numel *= d; + size_t offset = 0; + for (size_t dim = 0; dim < dims.size(); dim++) { + if (dims[dim] <= 0) { + throw std::runtime_error("select_as_symint: empty input tensor"); + } + if (static_cast(dims[dim]) > + std::numeric_limits::max()) { + throw std::runtime_error("select_as_symint: dimension exceeds size_t"); + } + int64_t index = src.indices[dim]; + if (index < 0) { + index += dims[dim]; + } + if (index < 0 || index >= dims[dim]) { + throw std::runtime_error("select_as_symint: index out of range"); + } + const size_t dim_size = static_cast(dims[dim]); + const size_t normalized_index = static_cast(index); + if (offset > + (std::numeric_limits::max() - normalized_index) / dim_size) { + throw std::runtime_error("select_as_symint: offset overflow"); + } + offset = offset * dim_size + normalized_index; } - if (numel <= 0) { - throw std::runtime_error("select_as_symint: empty input tensor"); + const void* host = inputs[pos].data; + if (host == nullptr || inputs[pos].host_is_fp32) { + throw std::runtime_error("select_as_symint: source is not integral"); } - int64_t stride = 1; - for (size_t i = static_cast(dim) + 1; i < dims.size(); i++) { - stride *= dims[i]; + const size_t host_element_size = + inputs[pos].host_is_int64 ? sizeof(int64_t) : sizeof(int32_t); + if (inputs[pos].nbytes % host_element_size != 0 || + offset >= inputs[pos].nbytes / host_element_size) { + throw std::runtime_error("select_as_symint: host source is too small"); } - // Reads the [0,..,index,..,0] element; symint sources are scalar-ish. - const int64_t offset = static_cast(index) * stride; - const void* host = inputs[pos].data; // Interpret the HOST buffer by its scalar type, not the tensor's serialized // elem_size: copy_inputs narrows an int64 host input to an int32 buffer, so // elem_size (buffer-derived) would misread int64 host data as int32. @@ -558,6 +575,15 @@ void WebGPUGraph::update_symints_from_inputs( } else { val = static_cast(host)[offset]; } + if (src.bias != 0) { + const int64_t adjusted = static_cast(val) + src.bias; + if (adjusted < std::numeric_limits::min() || + adjusted > std::numeric_limits::max()) { + throw std::runtime_error( + "select_as_symint: bias result overflows int32"); + } + val = static_cast(adjusted); + } set_symint(src.symint_id, val); } // sym_size.int: SymInt = a tensor's live dim (cur_dims). Usually unused (ops @@ -570,7 +596,197 @@ void WebGPUGraph::update_symints_from_inputs( } } +void WebGPUGraph::add_input_select_projection( + int output_tensor_id, + int input_tensor_id, + int dim, + int index) { + const auto& input = tensors_.at(input_tensor_id); + const auto& output = tensors_.at(output_tensor_id); + if (!input.is_int || input.elem_size != sizeof(int32_t) || !output.is_int || + output.elem_size != sizeof(int32_t) || input.dims.empty() || + output.dims.size() + 1 != input.dims.size()) { + return; + } + const int normalized_dim = + normalize_dim(dim, static_cast(input.dims.size()), "select"); + int64_t normalized_index = index; + if (normalized_index < 0) { + normalized_index += input.dims[normalized_dim]; + } + if (normalized_index < 0 || normalized_index >= input.dims[normalized_dim]) { + return; + } + for (size_t in_dim = 0, out_dim = 0; in_dim < input.dims.size(); in_dim++) { + if (static_cast(in_dim) == normalized_dim) { + continue; + } + if (output.dims[out_dim++] != input.dims[in_dim]) { + return; + } + } + input_select_projections_[output_tensor_id] = { + input_tensor_id, normalized_dim, index}; +} + +void WebGPUGraph::add_input_bias_projection( + int output_tensor_id, + int input_tensor_id, + int64_t bias) { + if (bias < std::numeric_limits::min() || + bias > std::numeric_limits::max()) { + throw std::runtime_error("select_as_symint: host bias exceeds int"); + } + input_select_projections_[output_tensor_id] = { + input_tensor_id, kHostBiasProjectionDim, static_cast(bias)}; +} + +bool WebGPUGraph::try_read_constant_bytes( + int const_value_id, + std::vector& out) const { + const auto it = constant_sources_.find(const_value_id); + if (it == constant_sources_.end() || it->second.nbytes == 0u) { + return false; + } + const ConstantSource& source = it->second; + if (source.inline_offset != UINT64_MAX) { + if (constant_data_ == nullptr || + source.inline_offset > constant_data_size_ || + source.nbytes > + constant_data_size_ - static_cast(source.inline_offset)) { + return false; + } + const uint8_t* data = + constant_data_ + static_cast(source.inline_offset); + out.assign(data, data + source.nbytes); + return true; + } + if (source.named_key.empty() || named_data_map_ == nullptr) { + return false; + } + auto data = named_data_map_->get_data(source.named_key.c_str()); + if (!data.ok()) { + return false; + } + if (data->size() != source.nbytes) { + data->Free(); + return false; + } + const auto* bytes = static_cast(data->data()); + out.assign(bytes, bytes + source.nbytes); + data->Free(); + return true; +} + +bool WebGPUGraph::try_read_prepacked_int32_scalar(int value_id, int32_t& out) + const { + const auto projection = input_select_projections_.find(value_id); + if (projection == input_select_projections_.end() || + projection->second.dim != kIdentityPrepackConstantProjectionDim) { + return false; + } + const int source_id = projection->second.input_tensor_id; + if (source_id < 0 || value_id < 0 || source_id >= num_values() || + value_id >= num_values()) { + return false; + } + const auto& source = tensors_[source_id]; + const auto& prepacked = tensors_[value_id]; + if (!source.is_int || source.elem_size != sizeof(int32_t) || + source.nbytes != sizeof(int32_t) || !prepacked.is_int || + prepacked.elem_size != sizeof(int32_t) || + prepacked.nbytes != sizeof(int32_t) || source.dims != prepacked.dims) { + return false; + } + try { + if (utils::numel(source.dims) != 1u) { + return false; + } + } catch (const std::runtime_error&) { + return false; + } + std::vector bytes; + if (!try_read_constant_bytes(source_id, bytes) || + bytes.size() != sizeof(int32_t)) { + return false; + } + std::memcpy(&out, bytes.data(), sizeof(out)); + return true; +} + +void WebGPUGraph::add_symint_source( + int symint_id, + int source_tensor_id, + int dim, + int index) { + struct SelectStep { + int dim; + int index; + }; + std::vector reverse_steps; + int64_t bias = 0; + int root_id = source_tensor_id; + for (size_t hops = 0; hops <= input_select_projections_.size(); hops++) { + if (std::find(input_ids_.begin(), input_ids_.end(), root_id) != + input_ids_.end()) { + break; + } + const auto projection = input_select_projections_.find(root_id); + if (projection == input_select_projections_.end()) { + throw std::runtime_error( + "select_as_symint: source is not a graph input or supported " + "projection"); + } + const auto& current = projection->second; + if (current.dim == kHostBiasProjectionDim) { + if ((current.index > 0 && + bias > std::numeric_limits::max() - current.index) || + (current.index < 0 && + bias < std::numeric_limits::min() - current.index)) { + throw std::runtime_error( + "select_as_symint: bias chain overflows int64"); + } + bias += current.index; + root_id = current.input_tensor_id; + continue; + } + reverse_steps.push_back({current.dim, current.index}); + root_id = current.input_tensor_id; + } + if (std::find(input_ids_.begin(), input_ids_.end(), root_id) == + input_ids_.end()) { + throw std::runtime_error("select_as_symint: projection cycle"); + } + + const auto& root = tensors_.at(root_id); + if (!root.is_int || root.elem_size != sizeof(int32_t) || root.dims.empty()) { + throw std::runtime_error("select_as_symint: source is not int32"); + } + std::vector indices(root.dims.size(), 0); + std::vector remaining_dims = root.dims; + std::vector root_axes(root.dims.size()); + for (size_t axis = 0; axis < root_axes.size(); axis++) { + root_axes[axis] = axis; + } + for (auto step = reverse_steps.rbegin(); step != reverse_steps.rend(); + step++) { + const int selected_dim = normalize_dim( + step->dim, static_cast(remaining_dims.size()), "select_as_symint"); + indices[root_axes[selected_dim]] = step->index; + remaining_dims.erase(remaining_dims.begin() + selected_dim); + root_axes.erase(root_axes.begin() + selected_dim); + } + const int selected_dim = normalize_dim( + dim, static_cast(remaining_dims.size()), "select_as_symint"); + indices[root_axes[selected_dim]] = index; + symint_sources_.push_back({symint_id, root_id, std::move(indices), bias}); +} + void WebGPUGraph::set_symint(int id, int32_t val) { + if (in_post_resize_phase_) { + throw std::runtime_error( + "WebGPU resize: post hook cannot set a SymInt"); + } auto it = symints_.find(id); if (it == symints_.end()) { throw std::runtime_error("WebGPUGraph::set_symint: id is not a SymInt"); @@ -586,6 +802,10 @@ void WebGPUGraph::set_symint(int id, int32_t val) { void WebGPUGraph::set_cur_dims( int value_id, const std::vector& new_dims) { + if (in_post_resize_phase_) { + throw std::runtime_error( + "WebGPU resize: post hook cannot set tensor dimensions"); + } auto& t = tensors_[value_id]; if (new_dims.size() != t.dims.size()) { throw std::runtime_error("WebGPU resize: tensor rank changed"); @@ -625,23 +845,34 @@ void WebGPUGraph::propagate_resize() { if (dirty_symints_.empty() && dirty_tensors_.empty()) { return; } - // Hooks fire in registration (topological) order: operands update first. - for (auto& hook : resize_hooks_) { - if (dirty_symints_.count(hook.symint_id) != 0) { - hook.fn(*this); + + std::unordered_set retry_symints = dirty_symints_; + std::unordered_set retry_tensors = dirty_tensors_; + std::unordered_set changed_symints; + std::unordered_set changed_tensors; + + try { + // Hooks fire in registration (topological) order: operands update first. + for (auto& hook : resize_hooks_) { + if (dirty_symints_.count(hook.symint_id) != 0) { + hook.fn(*this); + } } - } - dirty_symints_.clear(); - // Tensor hooks: bounded fixpoint. A hook may dirty its output (cascading to a - // consumer); each pass handles the currently-dirty set. A forward DAG - // converges in <= depth passes (set_cur_dims re-dirties only on a change). - for (size_t pass = 0; - !dirty_tensors_.empty() && pass <= tensor_resize_hooks_.size(); - pass++) { - std::unordered_set processing; - processing.swap(dirty_tensors_); - pending_dynamic_dispatch_grids_.clear(); - try { + changed_symints.insert(dirty_symints_.begin(), dirty_symints_.end()); + retry_symints.insert(dirty_symints_.begin(), dirty_symints_.end()); + dirty_symints_.clear(); + + // Tensor hooks: bounded fixpoint. A hook may dirty its output (cascading + // to a consumer); each pass handles the currently-dirty set. + for (size_t pass = 0; + !dirty_tensors_.empty() && pass <= tensor_resize_hooks_.size(); + pass++) { + std::unordered_set processing; + processing.swap(dirty_tensors_); + changed_tensors.insert(processing.begin(), processing.end()); + retry_tensors.insert(processing.begin(), processing.end()); + pending_dynamic_dispatch_grids_.clear(); + for (auto& hook : tensor_resize_hooks_) { if (processing.count(hook.trigger_tensor_id) != 0) { hook.fn(*this); @@ -662,28 +893,50 @@ void WebGPUGraph::propagate_resize() { pending_dynamic_dispatch_grids_.push_back( {dynamic_grid.dispatch_index, grid}); } - } catch (...) { + for (const auto& pending : pending_dynamic_dispatch_grids_) { + auto& dispatch = dispatches_[pending.dispatch_index]; + dispatch.workgroup_count_x = pending.grid.x; + dispatch.workgroup_count_y = pending.grid.y; + } pending_dynamic_dispatch_grids_.clear(); - // Keep both the current triggers and any cascaded outputs dirty so the - // caller can fix the hook or picker and retry without rebuilding. - dirty_tensors_.insert(processing.begin(), processing.end()); - throw; } - for (const auto& pending : pending_dynamic_dispatch_grids_) { - auto& dispatch = dispatches_[pending.dispatch_index]; - dispatch.workgroup_count_x = pending.grid.x; - dispatch.workgroup_count_y = pending.grid.y; + if (!dirty_tensors_.empty()) { + throw std::runtime_error( + "WebGPU resize: tensor resize hooks did not converge"); + } + // Tensor hooks must not set_symint (dirty_symints_ already drained above). + if (!dirty_symints_.empty()) { + throw std::runtime_error( + "WebGPU resize: a tensor resize hook set a SymInt; not supported"); + } + + in_post_resize_phase_ = true; + for (auto& hook : post_resize_hooks_) { + bool matches = false; + for (int id : hook.tensor_trigger_ids) { + matches = matches || changed_tensors.count(id) != 0; + } + for (int id : hook.symint_trigger_ids) { + matches = matches || changed_symints.count(id) != 0; + } + if (!matches) { + continue; + } + hook.fn(*this); + if (!dirty_symints_.empty() || !dirty_tensors_.empty()) { + throw std::runtime_error( + "WebGPU resize: post hook started another propagation wave"); + } } + in_post_resize_phase_ = false; + } catch (...) { pending_dynamic_dispatch_grids_.clear(); - } - if (!dirty_tensors_.empty()) { - throw std::runtime_error( - "WebGPU resize: tensor resize hooks did not converge"); - } - // Tensor hooks must not set_symint (dirty_symints_ already drained above). - if (!dirty_symints_.empty()) { - throw std::runtime_error( - "WebGPU resize: a tensor resize hook set a SymInt; not supported"); + in_post_resize_phase_ = false; + retry_symints.insert(dirty_symints_.begin(), dirty_symints_.end()); + retry_tensors.insert(dirty_tensors_.begin(), dirty_tensors_.end()); + dirty_symints_ = std::move(retry_symints); + dirty_tensors_ = std::move(retry_tensors); + throw; } } @@ -747,6 +1000,10 @@ WebGPUGraph::~WebGPUGraph() { wgpuBindGroupLayoutRelease(bgl); } } + if (queue_) { + wgpuQueueRelease(queue_); + queue_ = nullptr; + } } void WebGPUGraph::build( @@ -811,6 +1068,15 @@ void WebGPUGraph::build( if (!a) { continue; } + if (is_prepack && a->size() >= 2) { + // This projection is valid only for et_vk.prepack.default: its handler + // verifies identical shape, dtype, and byte size, then materializes the + // source bytes unchanged. Layout-transforming prepacks must not use it. + input_select_projections_[static_cast(a->Get(1))] = { + static_cast(a->Get(0)), + kIdentityPrepackConstantProjectionDim, + 0}; + } if (oc->name()->str() == "sym_size.int" && a->size() >= 3 && values) { const auto* out = values->Get(a->Get(2)); if (out && out->value_type() == vkgraph::GraphTypes::SymInt) { @@ -1593,6 +1859,42 @@ void WebGPUGraph::record_active_route(const std::string& kernel_name) { } #endif // WGPU_BACKEND_ENABLE_PROFILING +namespace { + +class ScopedComputePass final { + public: + ~ScopedComputePass() { + close(); + } + + WGPUComputePassEncoder get() const { + return pass_; + } + + void reset(WGPUComputePassEncoder pass) { + close(); + pass_ = pass; + } + + void close() { + if (pass_ == nullptr) { + return; + } + wgpuComputePassEncoderEnd(pass_); + wgpuComputePassEncoderRelease(pass_); + pass_ = nullptr; + } + + private: + WGPUComputePassEncoder pass_ = nullptr; +}; + +} // namespace + +std::string WebGPUGraph::execution_attestation_json() const { + return serialize_webgpu_execution_attestation(execution_attestation_); +} + WebGPUExecutionPlan WebGPUGraph::make_execution_plan( const WebGPUGraphExecutionOptions& options) const { const size_t n = dispatches_.size(); @@ -1624,17 +1926,82 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { #endif // WGPU_BACKEND_ENABLE_PROFILING const size_t n = dispatches_.size(); const size_t chunk = execute_config_.chunk_size; + execution_attestation_ = {}; + execution_attestation_.execution_ordinal = ++execution_ordinal_; + execution_attestation_.requested = plan.single_compute_pass; + execution_attestation_.max_compute_dispatches_per_pass = + plan.max_compute_dispatches_per_pass; if (plan.copy_outputs.size() != output_copies_.size()) { + execution_attestation_.error_reason = + "execution plan output count mismatch"; throw std::runtime_error("WebGPU: execution plan output count mismatch"); } + std::vector selected_dispatches(n, false); for (const auto& dispatch_chunk : plan.dispatch_chunks) { for (size_t dispatch_index : dispatch_chunk) { if (dispatch_index >= n) { + execution_attestation_.error_reason = + "execution plan dispatch index out of range"; throw std::runtime_error( "WebGPU: execution plan dispatch index out of range"); } - } + selected_dispatches[dispatch_index] = true; + } + } + + std::vector command_records; + command_records.reserve(dispatches_.size() + output_copies_.size()); + for (size_t i = 0; i < dispatches_.size(); i++) { + const auto& dispatch = dispatches_[i]; + WebGPUCommandRecord record; + record.static_dispatch_index = i; + record.enabled = selected_dispatches[i]; + if (dispatch.kind == WebGPUDispatch::Kind::Compute) { + record.identity = dispatch.kernel_name; + record.workgroup_count_x = dispatch.workgroup_count_x; + record.workgroup_count_y = dispatch.workgroup_count_y; + record.zero_grid = + dispatch.workgroup_count_x == 0 || dispatch.workgroup_count_y == 0; + } else { + record.kind = WebGPUCommandKind::GraphCopy; + record.byte_count = dispatch.copy_nbytes; + } + command_records.push_back(std::move(record)); + } + for (size_t i = 0; i < output_copies_.size(); i++) { + WebGPUCommandRecord record; + record.kind = WebGPUCommandKind::OutputCopy; + record.output_ordinal = i; + record.source_identity = output_ids_[i]; + record.destination_identity = static_cast(i); + record.byte_count = output_copies_[i].nbytes; + record.enabled = plan.copy_outputs[i]; + record.suppressed = !record.enabled; + command_records.push_back(std::move(record)); + } + execution_attestation_.inventory = + build_webgpu_command_inventory(command_records); + + if (!plan.single_compute_pass && + plan.max_compute_dispatches_per_pass != 0) { + execution_attestation_.error_reason = + "pass cap requires single_compute_pass"; + throw std::runtime_error("WebGPU: " + execution_attestation_.error_reason); + } + const bool is_chunked = chunk != 0 && n > chunk; + if (plan.single_compute_pass && is_chunked) { + execution_attestation_.error_reason = + "single_compute_pass is incompatible with chunked execution"; + throw std::runtime_error("WebGPU: " + execution_attestation_.error_reason); } +#ifdef WGPU_BACKEND_ENABLE_PROFILING + if (plan.single_compute_pass && should_timestamp_query()) { + execution_attestation_.error_reason = + "single_compute_pass is incompatible with timestamp queries"; + throw std::runtime_error("WebGPU: " + execution_attestation_.error_reason); + } +#endif // WGPU_BACKEND_ENABLE_PROFILING + execution_attestation_.applied = plan.single_compute_pass; if (plan.dispatch_chunks.empty()) { return 0; @@ -1670,14 +2037,17 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device_, &enc_desc); - // One pass per dispatch: enforces storage RAW ordering across deps. #ifdef WGPU_BACKEND_ENABLE_PROFILING uint32_t query_index = 0; #endif + ScopedComputePass shared_pass; + size_t dispatches_in_shared_pass = 0; 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) { + shared_pass.close(); + dispatches_in_shared_pass = 0; wgpuCommandEncoderCopyBufferToBuffer( encoder, dispatch.copy_src, @@ -1699,15 +2069,31 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { pass_desc.timestampWrites = &tw; } #endif // WGPU_BACKEND_ENABLE_PROFILING - WGPUComputePassEncoder pass = - wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); + WGPUComputePassEncoder pass = shared_pass.get(); + if (pass == nullptr) { + pass = wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); + ++execution_attestation_.encoded_compute_passes; + if (plan.single_compute_pass) { + shared_pass.reset(pass); + } + } 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); + if (!plan.single_compute_pass) { + wgpuComputePassEncoderEnd(pass); + wgpuComputePassEncoderRelease(pass); + } else { + ++dispatches_in_shared_pass; + if (webgpu_pass_cap_reached( + dispatches_in_shared_pass, + plan.max_compute_dispatches_per_pass)) { + shared_pass.close(); + dispatches_in_shared_pass = 0; + } + } #ifdef WGPU_BACKEND_ENABLE_PROFILING if (qp) { qp->record( @@ -1720,6 +2106,7 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { #endif // WGPU_BACKEND_ENABLE_PROFILING } } + shared_pass.close(); for (size_t i = 0; i < output_copies_.size(); i++) { const size_t logical_nbytes = tensors_[output_ids_[i]].cur_nbytes; @@ -1741,6 +2128,7 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { WGPUCommandBufferDescriptor cmd_desc = {}; WGPUCommandBuffer cmd = wgpuCommandEncoderFinish(encoder, &cmd_desc); wgpuQueueSubmit(queue_, 1, &cmd); + execution_attestation_.queue_submit_count = 1; wgpuCommandBufferRelease(cmd); wgpuCommandEncoderRelease(encoder); @@ -1751,6 +2139,15 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { qp->print_results(); } #endif // WGPU_BACKEND_ENABLE_PROFILING + const size_t expected_passes = count_webgpu_compute_passes( + command_records, + plan.single_compute_pass, + plan.max_compute_dispatches_per_pass); + if (execution_attestation_.encoded_compute_passes != expected_passes) { + execution_attestation_.error_reason = + "encoded compute-pass count does not match command inventory"; + throw std::runtime_error("WebGPU: " + execution_attestation_.error_reason); + } return 1; } @@ -1785,6 +2182,7 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { WGPUComputePassDescriptor pass_desc = {}; WGPUComputePassEncoder pass = wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); + ++execution_attestation_.encoded_compute_passes; wgpuComputePassEncoderSetPipeline(pass, dispatches_[i].pipeline); wgpuComputePassEncoderSetBindGroup( pass, 0, dispatches_[i].bind_group, 0, nullptr); @@ -1813,10 +2211,18 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { WGPUCommandBufferDescriptor cmd_desc = {}; WGPUCommandBuffer cmd = wgpuCommandEncoderFinish(encoder, &cmd_desc); wgpuQueueSubmit(queue_, 1, &cmd); + ++execution_attestation_.queue_submit_count; wgpuCommandBufferRelease(cmd); wgpuCommandEncoderRelease(encoder); } + const size_t expected_passes = + count_webgpu_compute_passes(command_records, false, 0); + if (execution_attestation_.encoded_compute_passes != expected_passes) { + execution_attestation_.error_reason = + "encoded compute-pass count does not match command inventory"; + throw std::runtime_error("WebGPU: " + execution_attestation_.error_reason); + } return plan.dispatch_chunks.size(); } diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 23ce9df03ed..d639d0851e6 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -175,6 +175,20 @@ class WebGPUGraph { std::vector& outputs, const WebGPUExecutionPlan& plan); + std::string execution_attestation_json() const; + + void complete_execution_attestation() { + execution_attestation_.completed = true; + execution_attestation_.error_reason.clear(); + } + + void fail_execution_attestation(const std::string& reason) { + execution_attestation_.completed = false; + if (execution_attestation_.error_reason.empty()) { + execution_attestation_.error_reason = reason; + } + } + const std::vector& input_ids() const { return input_ids_; } @@ -230,13 +244,21 @@ class WebGPUGraph { struct SymIntSource { int symint_id; int input_tensor_id; - int dim; - int index; + std::vector indices; + int64_t bias = 0; }; void - add_symint_source(int symint_id, int input_tensor_id, int dim, int index) { - symint_sources_.push_back({symint_id, input_tensor_id, dim, index}); - } + add_symint_source(int symint_id, int source_tensor_id, int dim, int index); + void add_input_select_projection( + int output_tensor_id, + int input_tensor_id, + int dim, + int index); + void add_input_bias_projection( + int output_tensor_id, + int input_tensor_id, + int64_t bias); + bool try_read_prepacked_int32_scalar(int value_id, int32_t& out) const; const std::vector& symint_sources() const { return symint_sources_; } @@ -328,6 +350,55 @@ class WebGPUGraph { }); } + // Terminal notification after the SymInt DAG and tensor-shape fixpoint. + // Callbacks may refresh complete derived state, but cannot start another + // propagation wave through set_symint or set_cur_dims. + void add_post_resize_hook( + const std::vector& tensor_trigger_ids, + const std::vector& symint_trigger_ids, + std::function fn) { + if (tensor_trigger_ids.empty() && symint_trigger_ids.empty()) { + throw std::runtime_error( + "WebGPU resize: post hook requires at least one trigger"); + } + for (int id : tensor_trigger_ids) { + if (id < 0 || id >= num_values() || + get_value_type(id) != ValueType::Tensor) { + throw std::runtime_error( + "WebGPU resize: post-hook tensor trigger must be a Tensor"); + } + } + for (int id : symint_trigger_ids) { + if (id < 0 || id >= num_values() || + get_value_type(id) != ValueType::SymInt) { + throw std::runtime_error( + "WebGPU resize: post-hook SymInt trigger must be a SymInt"); + } + } + if (!fn) { + throw std::runtime_error("WebGPU resize: null post-resize hook"); + } + post_resize_hooks_.push_back( + {tensor_trigger_ids, symint_trigger_ids, std::move(fn)}); + } + + template + void add_post_resize_hook( + const std::vector& tensor_trigger_ids, + const std::vector& symint_trigger_ids, + void (*fn)(WebGPUGraph&, const Context&), + Context context) { + if (fn == nullptr) { + throw std::runtime_error("WebGPU resize: null post-resize hook"); + } + add_post_resize_hook( + tensor_trigger_ids, + symint_trigger_ids, + [fn, context = std::move(context)](WebGPUGraph& graph) { + fn(graph, context); + }); + } + // Run hooks for changed SymInts and tensors, then clear; call before execute. void propagate_resize(); @@ -565,6 +636,9 @@ class WebGPUGraph { } private: + bool try_read_constant_bytes(int const_value_id, std::vector& out) + const; + #ifdef WGPU_BACKEND_ENABLE_PROFILING void record_active_route(const std::string& kernel_name); #endif // WGPU_BACKEND_ENABLE_PROFILING @@ -596,6 +670,12 @@ class WebGPUGraph { }; std::unordered_map symints_; std::vector symint_sources_; + struct InputSelectProjection { + int input_tensor_id; + int dim; + int index; + }; + std::unordered_map input_select_projections_; std::vector symint_dim_sources_; std::unordered_set dynamic_tensor_ids_; @@ -616,6 +696,14 @@ class WebGPUGraph { std::vector tensor_resize_hooks_; std::unordered_set dirty_tensors_; + struct PostResizeHook { + std::vector tensor_trigger_ids; + std::vector symint_trigger_ids; + std::function fn; + }; + std::vector post_resize_hooks_; + bool in_post_resize_phase_ = false; + // Dynamic grids are stored separately so ordinary dispatches remain compact. // The graph owns each dispatch index and picker; ops only provide typed // context and a named grid function. @@ -683,6 +771,8 @@ class WebGPUGraph { std::unordered_map constant_sources_; ExecuteConfig execute_config_; + uint64_t execution_ordinal_ = 0; + WebGPUExecutionAttestation execution_attestation_; // Caches for reusing GPU objects across dispatches. std::unordered_map shader_cache_; diff --git a/backends/webgpu/runtime/WebGPUQueryPool.cpp b/backends/webgpu/runtime/WebGPUQueryPool.cpp index dc00414df04..efa38bbe6b0 100644 --- a/backends/webgpu/runtime/WebGPUQueryPool.cpp +++ b/backends/webgpu/runtime/WebGPUQueryPool.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include @@ -25,13 +26,16 @@ struct MapCallbackData { WGPUMapAsyncStatus status = WGPUMapAsyncStatus_Error; }; +using MapCallbackDataPtr = std::shared_ptr; + void map_callback( WGPUMapAsyncStatus status, WGPUStringView /*message*/, void* userdata1, void* /*userdata2*/) { - auto* data = static_cast(userdata1); - data->status = status; + std::unique_ptr data_owner( + static_cast(userdata1)); + (*data_owner)->status = status; } constexpr uint64_t kTimestampBytes = sizeof(uint64_t); @@ -88,6 +92,7 @@ void WebGPUQueryPool::reset(uint32_t num_dispatches) { } num_pairs_ = num_dispatches; durations_.clear(); + result_state_.invalidate(); } WGPUPassTimestampWrites WebGPUQueryPool::writes_for(uint32_t i) { @@ -154,33 +159,65 @@ void fill_shader_durations( } } +bool WebGPUQueryPool::finalize_extraction( + WGPUWaitStatus wait_status, + WGPUMapAsyncStatus map_status, + const uint64_t* ticks) { + result_state_.invalidate(); + if (!detail::query_result_extraction_succeeded( + num_pairs_, wait_status, map_status, ticks)) { + return false; + } + if (num_pairs_ != 0) { + fill_shader_durations(durations_, ticks, ns_per_tick_); + } + result_state_.complete(); + return true; +} + void WebGPUQueryPool::extract_results(WGPUInstance instance) { + result_state_.invalidate(); if (num_pairs_ == 0) { + (void)finalize_extraction( + WGPUWaitStatus_Success, WGPUMapAsyncStatus_Success, nullptr); return; } const uint32_t count = 2 * num_pairs_; const uint64_t bytes = static_cast(count) * kTimestampBytes; - MapCallbackData cb; + const auto cb_data = std::make_shared(); WGPUBufferMapCallbackInfo cb_info = {}; cb_info.mode = WGPUCallbackMode_WaitAnyOnly; cb_info.callback = map_callback; - cb_info.userdata1 = &cb; - webgpu_wait( - instance, - wgpuBufferMapAsync(readback_buf_, WGPUMapMode_Read, 0, bytes, cb_info)); + cb_info.userdata1 = new MapCallbackDataPtr(cb_data); + const WGPUFuture map_future = + wgpuBufferMapAsync(readback_buf_, WGPUMapMode_Read, 0, bytes, cb_info); + const WGPUWaitStatus wait_status = webgpu_wait(instance, map_future); - if (cb.status != WGPUMapAsyncStatus_Success) { + if (wait_status != WGPUWaitStatus_Success) { + wgpuBufferUnmap(readback_buf_); + (void)webgpu_wait(instance, map_future); + (void)finalize_extraction(wait_status, cb_data->status, nullptr); printf( - "WebGPUQueryPool: readback map failed (status %d)\n", (int)cb.status); + "WebGPUQueryPool: readback wait failed (status %d)\n", + (int)wait_status); + return; + } + if (cb_data->status != WGPUMapAsyncStatus_Success) { + (void)finalize_extraction(wait_status, cb_data->status, nullptr); + printf( + "WebGPUQueryPool: readback map failed (status %d)\n", + (int)cb_data->status); return; } const uint64_t* ticks = static_cast( wgpuBufferGetConstMappedRange(readback_buf_, 0, bytes)); - if (ticks != nullptr) { - fill_shader_durations(durations_, ticks, ns_per_tick_); - } + const bool valid = + finalize_extraction(wait_status, cb_data->status, ticks); wgpuBufferUnmap(readback_buf_); + if (!valid) { + printf("WebGPUQueryPool: readback mapped range is null\n"); + } } void WebGPUQueryPool::print_results(bool tsv) const { diff --git a/backends/webgpu/runtime/WebGPUQueryPool.h b/backends/webgpu/runtime/WebGPUQueryPool.h index 3f5089d5e9a..71b86f7c3ba 100644 --- a/backends/webgpu/runtime/WebGPUQueryPool.h +++ b/backends/webgpu/runtime/WebGPUQueryPool.h @@ -17,6 +17,41 @@ namespace executorch::backends::webgpu { +namespace detail { + +class WebGPUQueryResultState final { + public: + bool results_valid() const { + return results_valid_; + } + uint64_t result_generation() const { + return result_generation_; + } + void invalidate() { + results_valid_ = false; + } + void complete() { + results_valid_ = true; + result_generation_++; + } + + private: + bool results_valid_ = false; + uint64_t result_generation_ = 0; +}; + +inline bool query_result_extraction_succeeded( + uint32_t num_pairs, + WGPUWaitStatus wait_status, + WGPUMapAsyncStatus map_status, + const uint64_t* ticks) { + return num_pairs == 0 || + (wait_status == WGPUWaitStatus_Success && + map_status == WGPUMapAsyncStatus_Success && ticks != nullptr); +} + +} // namespace detail + #ifdef WGPU_BACKEND_ENABLE_PROFILING // Per-dispatch GPU timing; mirrors Vulkan QueryPool ShaderDuration. @@ -70,10 +105,21 @@ class WebGPUQueryPool { const std::vector& results() const { return durations_; } + bool results_valid() const { + return result_state_.results_valid(); + } + uint64_t result_generation() const { + return result_state_.result_generation(); + } void print_results(bool tsv = false) const; uint64_t get_mean_shader_ns(const std::string& kernel_name) const; private: + bool finalize_extraction( + WGPUWaitStatus wait_status, + WGPUMapAsyncStatus map_status, + const uint64_t* ticks); + WGPUQuerySet qset_ = nullptr; WGPUBuffer resolve_buf_ = nullptr; // QueryResolve | CopySrc WGPUBuffer readback_buf_ = nullptr; // MapRead | CopyDst @@ -81,6 +127,7 @@ class WebGPUQueryPool { uint32_t num_pairs_ = 0; double ns_per_tick_ = 1.0; // WebGPU timestamps are already nanoseconds std::vector durations_; + detail::WebGPUQueryResultState result_state_; }; // Per-op durations from begin/end tick pairs (consecutive-end delta). diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index 480944ea93d..ebc3dcc6cd7 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -13,6 +13,8 @@ #include #include #include +#include +#include #include #include #include @@ -153,7 +155,7 @@ namespace executorch::backends::webgpu { namespace { -constexpr std::array kShaderRegistry = {{ +constexpr std::array kShaderRegistry = {{ { "abs", kAbsWGSL, @@ -203,6 +205,20 @@ constexpr std::array kShaderRegistry = {{ kArgReduceWorkgroupSizeY, kArgReduceWorkgroupSizeZ, }, + { + "arg_reduce_final", + kArgReduceFinalWGSL, + kArgReduceFinalWorkgroupSizeX, + kArgReduceFinalWorkgroupSizeY, + kArgReduceFinalWorkgroupSizeZ, + }, + { + "arg_reduce_partial", + kArgReducePartialWGSL, + kArgReducePartialWorkgroupSizeX, + kArgReducePartialWorkgroupSizeY, + kArgReducePartialWorkgroupSizeZ, + }, { "avg_pool2d", kAvgPool2dWGSL, diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index 20ff14c34f1..b1373adfa7b 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -26,6 +26,37 @@ namespace executorch::backends::webgpu::utils { +struct RowChunking { + uint64_t rows_per_chunk; + uint32_t num_chunks; +}; + +inline RowChunking compute_row_chunking( + uint64_t max_binding_bytes, + uint64_t bytes_per_row, + uint64_t total_rows, + const char* label) { + const std::string name = label == nullptr ? "row chunking" : label; + if (max_binding_bytes == 0 || bytes_per_row == 0 || total_rows == 0) { + throw std::runtime_error( + "WebGPU " + name + ": row chunking arguments must be nonzero"); + } + if (bytes_per_row > max_binding_bytes) { + throw std::runtime_error( + "WebGPU " + name + ": one row exceeds the storage binding limit"); + } + + const uint64_t rows_per_chunk = + std::min(total_rows, max_binding_bytes / bytes_per_row); + const uint64_t num_chunks = total_rows / rows_per_chunk + + static_cast(total_rows % rows_per_chunk != 0); + if (num_chunks > std::numeric_limits::max()) { + throw std::runtime_error( + "WebGPU " + name + ": row chunk count exceeds uint32_t"); + } + return {rows_per_chunk, static_cast(num_chunks)}; +} + // Product of dims (live element count); used by dynamic-resize hooks. Delegates // to numel (single impl; keeps the negative-dim guard, no caller churn). inline uint64_t numel_of(const std::vector& dims) { @@ -332,8 +363,18 @@ struct BindingSpec { WGPUBufferBindingType type; WGPUBuffer buffer; uint64_t size; + uint64_t offset = 0; }; +inline WGPUBindGroupEntry make_bind_group_entry(const BindingSpec& binding) { + WGPUBindGroupEntry entry = {}; + entry.binding = binding.binding; + entry.buffer = binding.buffer; + entry.offset = binding.offset; + entry.size = binding.size; + return entry; +} + // Owns the shader module, bind-group layout, and pipeline layout, releasing // them on destruction. `pipeline` and `bind_group` are NOT released here — // every op hands them to WebGPUGraph::add_dispatch, which keeps them alive @@ -411,10 +452,7 @@ inline ComputePipelineBundle make_compute_pipeline( layout_entries[i].visibility = WGPUShaderStage_Compute; layout_entries[i].buffer.type = bindings[i].type; - bind_entries[i] = {}; - bind_entries[i].binding = bindings[i].binding; - bind_entries[i].buffer = bindings[i].buffer; - bind_entries[i].size = bindings[i].size; + bind_entries[i] = make_bind_group_entry(bindings[i]); } WGPUBindGroupLayoutDescriptor bgl_desc = {}; @@ -488,10 +526,7 @@ inline ComputePipelineBundle make_compute_pipeline( std::vector bind_entries(bindings.size()); for (size_t i = 0; i < bindings.size(); i++) { - bind_entries[i] = {}; - bind_entries[i].binding = bindings[i].binding; - bind_entries[i].buffer = bindings[i].buffer; - bind_entries[i].size = bindings[i].size; + bind_entries[i] = make_bind_group_entry(bindings[i]); } WGPUComputePipelineDescriptor pipeline_desc = {}; diff --git a/backends/webgpu/runtime/ops/argmax/Reduce.cpp b/backends/webgpu/runtime/ops/argmax/Reduce.cpp index 7b5b714f1d2..83953fc198b 100644 --- a/backends/webgpu/runtime/ops/argmax/Reduce.cpp +++ b/backends/webgpu/runtime/ops/argmax/Reduce.cpp @@ -9,11 +9,15 @@ #include #include #include +#include +#include +#include #include #include #include +#include #include #include @@ -35,6 +39,151 @@ static_assert( kArgReduceWorkgroupSizeX <= 256, "arg_reduce workgroup size exceeds the 256-wide shared partials"); +struct ArgReduceMultiWgParams { + uint32_t num_rows; + uint32_t reduce_size; + uint32_t is_argmin; + uint32_t num_parts; +}; +static_assert(sizeof(ArgReduceMultiWgParams) == 16); + +static_assert( + kArgReducePartialWorkgroupSizeX == kArgReduceFinalWorkgroupSizeX, + "arg_reduce stage workgroup widths must match"); + +void arg_reduce_multiwg_impl( + WebGPUGraph& graph, + int in_id, + int out_id, + uint32_t num_rows, + uint32_t reduce_size, + uint32_t is_argmin, + uint32_t parts, + bool keepdim) { + WGPUDevice device = graph.device(); + const WebGPUTensor& in_tensor = graph.get_tensor(in_id); + const WebGPUTensor& out_tensor = graph.get_tensor(out_id); + const uint32_t max_parts = parts; + const uint64_t max_partial_slots = + static_cast(num_rows) * max_parts; + const size_t partials_bytes = + static_cast(max_partial_slots) * 2u * sizeof(uint32_t); + WGPUBuffer partials = graph.acquire_scratch(partials_bytes); + WebGPUGraph::ScopedScratch partials_guard(&graph, partials); + + ArgReduceMultiWgParams params = {num_rows, reduce_size, is_argmin, parts}; + WGPUBuffer uniform_buffer = + graph.make_uniform_buffer(¶ms, sizeof(params)); + graph.own_uniform_buffer(uniform_buffer); + const utils::WgCount partial_grid = utils::compute_2d_workgroup_count( + device, num_rows * parts, 1u, "arg_reduce_partial"); + const utils::WgCount final_grid = utils::compute_2d_workgroup_count( + device, num_rows, 1u, "arg_reduce_final"); + + utils::ComputePipelineBundle partial_bundle = utils::make_compute_pipeline( + device, + kArgReducePartialWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, WGPUBufferBindingType_Storage, partials, partials_bytes}, + {2, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(params)}, + }); + const size_t partial_dispatch = graph.add_dispatch( + {partial_bundle.pipeline, + partial_bundle.bind_group, + partial_grid.x, + "arg_reduce_partial", + partial_grid.y}); + utils::ComputePipelineBundle final_bundle = utils::make_compute_pipeline( + device, + kArgReduceFinalWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, partials, partials_bytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(params)}, + }); + const size_t final_dispatch = graph.add_dispatch( + {final_bundle.pipeline, + final_bundle.bind_group, + final_grid.x, + "arg_reduce_final", + final_grid.y}); + + WGPUBuffer params_buffer = uniform_buffer; + graph.add_tensor_resize_hook( + in_id, + [in_id, + out_id, + keepdim, + is_argmin, + max_parts, + max_partial_slots, + partial_dispatch, + final_dispatch, + params_buffer](WebGPUGraph& resized_graph) { + const auto& dims = resized_graph.cur_dims(in_id); + if (dims.empty() || dims.back() <= 0 || + static_cast(dims.back()) > + std::numeric_limits::max()) { + throw std::runtime_error( + "arg_reduce(resize): reduce dim is out of range"); + } + const uint32_t live_reduce_size = static_cast(dims.back()); + const uint64_t total = utils::numel_of(dims); + const uint64_t live_rows = total / live_reduce_size; + if (live_rows == 0u || + live_rows > std::numeric_limits::max()) { + throw std::runtime_error( + "arg_reduce(resize): row count is out of range"); + } + const uint32_t rows = static_cast(live_rows); + uint32_t live_parts = select_arg_reduce_parts( + rows, + live_reduce_size, + utils::queried_max_workgroups(resized_graph.device())); + live_parts = std::max(1u, std::min(live_parts, max_parts)); + if (!arg_reduce_partial_slots_fit( + max_partial_slots, rows, live_parts)) { + throw std::runtime_error( + "arg_reduce(resize): partial scratch capacity exceeded"); + } + std::vector output_dims = dims; + if (keepdim) { + output_dims.back() = 1; + } else { + output_dims.pop_back(); + } + resized_graph.set_cur_dims(out_id, output_dims); + ArgReduceMultiWgParams live_params = { + rows, live_reduce_size, is_argmin, live_parts}; + wgpuQueueWriteBuffer( + resized_graph.queue(), + params_buffer, + 0, + &live_params, + sizeof(live_params)); + const utils::WgCount partial_count = utils::compute_2d_workgroup_count( + resized_graph.device(), + static_cast(live_rows * live_parts), + 1u, + "arg_reduce_partial"); + auto& partial = resized_graph.dispatch_at(partial_dispatch); + partial.workgroup_count_x = partial_count.x; + partial.workgroup_count_y = partial_count.y; + const utils::WgCount final_count = utils::compute_2d_workgroup_count( + resized_graph.device(), rows, 1u, "arg_reduce_final"); + auto& final = resized_graph.dispatch_at(final_dispatch); + final.workgroup_count_x = final_count.x; + final.workgroup_count_y = final_count.y; + }); +} + // Last-dim argmax/argmin -> int32 index; mirrors Vulkan arg_reduce_impl. void arg_reduce_impl( WebGPUGraph& graph, @@ -78,6 +227,15 @@ void arg_reduce_impl( throw std::runtime_error("arg_reduce: shape mismatch (rows * reduce_size)"); } + const bool keepdim = graph.get_bool(args.at(2)); + const uint32_t parts = select_arg_reduce_parts( + num_rows, reduce_size, utils::queried_max_workgroups(device)); + if (parts != 0u) { + arg_reduce_multiwg_impl( + graph, in_id, out_id, num_rows, reduce_size, is_argmin, parts, keepdim); + return; + } + uint32_t wg_size = utils::clamp_workgroup_size(device, kArgReduceWorkgroupSizeX); // One workgroup per row (cooperative reduction); grid = num_rows workgroups. @@ -125,7 +283,6 @@ void arg_reduce_impl( workgroup_count.y}); // Dynamic shapes: recompute reduce_size (last dim) + num_rows + dispatch. - const bool keepdim = graph.get_bool(args.at(2)); WGPUBuffer params_buf = uniform_buffer; graph.add_tensor_resize_hook( in_id, diff --git a/backends/webgpu/runtime/ops/argmax/arg_reduce_final.wgsl b/backends/webgpu/runtime/ops/argmax/arg_reduce_final.wgsl new file mode 100644 index 00000000000..4d6ffcc6899 --- /dev/null +++ b/backends/webgpu/runtime/ops/argmax/arg_reduce_final.wgsl @@ -0,0 +1,88 @@ +// 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. + +struct Partial { + val: f32, + idx: u32, +} +@group(0) @binding(0) var t_part: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + num_rows: u32, + reduce_size: u32, + is_argmin: u32, + num_parts: u32, +} +@group(0) @binding(2) var params: Params; + +// Stage 2: one workgroup per row merges that row's `num_parts` stage-1 +// partials and writes the winning index. Same lower-index-wins rule as +// arg_reduce.wgsl:57-61; every partial is a real (value, index) pair drawn +// from the row, so merging them in any order yields the row's lowest-index +// extremum -- the identical result to the single-workgroup kernel. +const WG: u32 = 64u; +var part_val: array; +var part_idx: array; + +@compute @workgroup_size(WG, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = wid.x + wid.y * num_workgroups.x; + if (row >= params.num_rows) { + return; + } + let pbase = row * params.num_parts; + + var best = t_part[pbase].val; + var best_idx = t_part[pbase].idx; + var t = lid.x; + while (t < params.num_parts) { + let v = t_part[pbase + t].val; + let idx = t_part[pbase + t].idx; + if (params.is_argmin != 0u) { + if (v < best || (v == best && idx < best_idx)) { best = v; best_idx = idx; } + } else { + if (v > best || (v == best && idx < best_idx)) { best = v; best_idx = idx; } + } + t = t + WG; + } + part_val[lid.x] = best; + part_idx[lid.x] = best_idx; + workgroupBarrier(); + + var stride: u32 = WG >> 1u; + loop { + if (stride == 0u) { + break; + } + if (lid.x < stride) { + let bv = part_val[lid.x]; + let bi = part_idx[lid.x]; + let v = part_val[lid.x + stride]; + let idx = part_idx[lid.x + stride]; + if (params.is_argmin != 0u) { + if (v < bv || (v == bv && idx < bi)) { + part_val[lid.x] = v; + part_idx[lid.x] = idx; + } + } else { + if (v > bv || (v == bv && idx < bi)) { + part_val[lid.x] = v; + part_idx[lid.x] = idx; + } + } + } + workgroupBarrier(); + stride = stride >> 1u; + } + + if (lid.x == 0u) { + t_out[row] = part_idx[0]; + } +} diff --git a/backends/webgpu/runtime/ops/argmax/arg_reduce_final_wgsl.h b/backends/webgpu/runtime/ops/argmax/arg_reduce_final_wgsl.h new file mode 100644 index 00000000000..bacc8a4bee8 --- /dev/null +++ b/backends/webgpu/runtime/ops/argmax/arg_reduce_final_wgsl.h @@ -0,0 +1,112 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from arg_reduce_final.wgsl - DO NOT EDIT. +// wgsl-sha256: efded2fdd0607954d05602c55f7ae6ec6e0dd732cb09e5468d16d7a788489284 +inline constexpr const char* kArgReduceFinalWGSL = R"( +// 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. + +struct Partial { + val: f32, + idx: u32, +} +@group(0) @binding(0) var t_part: array; +@group(0) @binding(1) var t_out: array; + +struct Params { + num_rows: u32, + reduce_size: u32, + is_argmin: u32, + num_parts: u32, +} +@group(0) @binding(2) var params: Params; + +// Stage 2: one workgroup per row merges that row's `num_parts` stage-1 +// partials and writes the winning index. Same lower-index-wins rule as +// arg_reduce.wgsl:57-61; every partial is a real (value, index) pair drawn +// from the row, so merging them in any order yields the row's lowest-index +// extremum -- the identical result to the single-workgroup kernel. +const WG: u32 = 64u; +var part_val: array; +var part_idx: array; + +@compute @workgroup_size(WG, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let row = wid.x + wid.y * num_workgroups.x; + if (row >= params.num_rows) { + return; + } + let pbase = row * params.num_parts; + + var best = t_part[pbase].val; + var best_idx = t_part[pbase].idx; + var t = lid.x; + while (t < params.num_parts) { + let v = t_part[pbase + t].val; + let idx = t_part[pbase + t].idx; + if (params.is_argmin != 0u) { + if (v < best || (v == best && idx < best_idx)) { best = v; best_idx = idx; } + } else { + if (v > best || (v == best && idx < best_idx)) { best = v; best_idx = idx; } + } + t = t + WG; + } + part_val[lid.x] = best; + part_idx[lid.x] = best_idx; + workgroupBarrier(); + + var stride: u32 = WG >> 1u; + loop { + if (stride == 0u) { + break; + } + if (lid.x < stride) { + let bv = part_val[lid.x]; + let bi = part_idx[lid.x]; + let v = part_val[lid.x + stride]; + let idx = part_idx[lid.x + stride]; + if (params.is_argmin != 0u) { + if (v < bv || (v == bv && idx < bi)) { + part_val[lid.x] = v; + part_idx[lid.x] = idx; + } + } else { + if (v > bv || (v == bv && idx < bi)) { + part_val[lid.x] = v; + part_idx[lid.x] = idx; + } + } + } + workgroupBarrier(); + stride = stride >> 1u; + } + + if (lid.x == 0u) { + t_out[row] = part_idx[0]; + } +} +)"; + +inline constexpr uint32_t kArgReduceFinalWorkgroupSizeX = 64; +inline constexpr uint32_t kArgReduceFinalWorkgroupSizeY = 1; +inline constexpr uint32_t kArgReduceFinalWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/argmax/arg_reduce_multiwg_route.h b/backends/webgpu/runtime/ops/argmax/arg_reduce_multiwg_route.h new file mode 100644 index 00000000000..6d8babc7d9a --- /dev/null +++ b/backends/webgpu/runtime/ops/argmax/arg_reduce_multiwg_route.h @@ -0,0 +1,50 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +constexpr uint32_t kArgReduceElemsPerWorkgroup = 1024u; +constexpr uint32_t kArgReduceMaxParts = 256u; +constexpr uint32_t kArgReduceMinReduceSize = 4096u; +constexpr uint32_t kArgReduceMaxPartials = 1u << 20; + +constexpr bool arg_reduce_partial_slots_fit( + uint64_t max_partial_slots, + uint32_t num_rows, + uint32_t num_parts) { + return max_partial_slots != 0u && num_rows != 0u && num_parts != 0u && + num_rows <= max_partial_slots / num_parts; +} + +constexpr uint32_t select_arg_reduce_parts( + uint32_t num_rows, + uint32_t reduce_size, + uint32_t max_workgroups_per_dim) { + if (num_rows == 0u || reduce_size < kArgReduceMinReduceSize) { + return 0u; + } + uint32_t parts = reduce_size / kArgReduceElemsPerWorkgroup + + static_cast(reduce_size % kArgReduceElemsPerWorkgroup != 0u); + if (parts > kArgReduceMaxParts) { + parts = kArgReduceMaxParts; + } + if (parts < 2u || num_rows > kArgReduceMaxPartials / parts || + max_workgroups_per_dim == 0u) { + return 0u; + } + const uint64_t total = static_cast(num_rows) * parts; + const uint64_t grid_capacity = + static_cast(max_workgroups_per_dim) * max_workgroups_per_dim; + return total <= grid_capacity ? parts : 0u; +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/argmax/arg_reduce_partial.wgsl b/backends/webgpu/runtime/ops/argmax/arg_reduce_partial.wgsl new file mode 100644 index 00000000000..d91aeae9475 --- /dev/null +++ b/backends/webgpu/runtime/ops/argmax/arg_reduce_partial.wgsl @@ -0,0 +1,106 @@ +// 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. + +@group(0) @binding(0) var t_in: array; + +struct Partial { + val: f32, + idx: u32, +} +@group(0) @binding(1) var t_part: array; + +struct Params { + num_rows: u32, + reduce_size: u32, + is_argmin: u32, + num_parts: u32, +} +@group(0) @binding(2) var params: Params; + +// Stage 1 of the two-stage arg-reduction. The single-workgroup kernel +// (arg_reduce.wgsl) puts an entire row on ONE 64-lane workgroup; here each row +// is split into `num_parts` contiguous chunks, one workgroup per chunk, so the +// row's elements are scanned by num_parts x WG lanes instead of WG. +// +// Semantics are held IDENTICAL to arg_reduce.wgsl: every lane is seeded with +// (t_in[base], 0) -- element 0 is always a valid candidate -- and scans with a +// STRICT compare, so the lowest index wins inside a lane. The in-workgroup +// merge is a halving tree using the same lower-index-wins rule; because no +// partial can ever hold a NaN unless t_in[base] itself is NaN (a strict +// compare never selects NaN), the merge operator is a total order and the tree +// is order-independent -> bit-exact index equality with the serial tail. +const WG: u32 = 64u; +var part_val: array; +var part_idx: array; + +@compute @workgroup_size(WG, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded grid of num_rows * num_parts workgroups. + let slot = wid.x + wid.y * num_workgroups.x; + let row = slot / params.num_parts; + if (row >= params.num_rows) { + return; + } + let part = slot - row * params.num_parts; + let base = row * params.reduce_size; + + let chunk = (params.reduce_size + params.num_parts - 1u) / params.num_parts; + let start = part * chunk; + var end = start + chunk; + if (end > params.reduce_size) { + end = params.reduce_size; + } + + var best = t_in[base]; + var best_idx: u32 = 0u; + var k = start + lid.x; + while (k < end) { + let v = t_in[base + k]; + if (params.is_argmin != 0u) { + if (v < best) { best = v; best_idx = k; } + } else { + if (v > best) { best = v; best_idx = k; } + } + k = k + WG; + } + part_val[lid.x] = best; + part_idx[lid.x] = best_idx; + workgroupBarrier(); + + var stride: u32 = WG >> 1u; + loop { + if (stride == 0u) { + break; + } + if (lid.x < stride) { + let bv = part_val[lid.x]; + let bi = part_idx[lid.x]; + let v = part_val[lid.x + stride]; + let idx = part_idx[lid.x + stride]; + if (params.is_argmin != 0u) { + if (v < bv || (v == bv && idx < bi)) { + part_val[lid.x] = v; + part_idx[lid.x] = idx; + } + } else { + if (v > bv || (v == bv && idx < bi)) { + part_val[lid.x] = v; + part_idx[lid.x] = idx; + } + } + } + workgroupBarrier(); + stride = stride >> 1u; + } + + if (lid.x == 0u) { + t_part[slot].val = part_val[0]; + t_part[slot].idx = part_idx[0]; + } +} diff --git a/backends/webgpu/runtime/ops/argmax/arg_reduce_partial_wgsl.h b/backends/webgpu/runtime/ops/argmax/arg_reduce_partial_wgsl.h new file mode 100644 index 00000000000..185f05899a1 --- /dev/null +++ b/backends/webgpu/runtime/ops/argmax/arg_reduce_partial_wgsl.h @@ -0,0 +1,130 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from arg_reduce_partial.wgsl - DO NOT EDIT. +// wgsl-sha256: 12dc45f7d391b2b531e0a028786bde612541e061934466b68dc0d1806899a9bb +inline constexpr const char* kArgReducePartialWGSL = R"( +// 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. + +@group(0) @binding(0) var t_in: array; + +struct Partial { + val: f32, + idx: u32, +} +@group(0) @binding(1) var t_part: array; + +struct Params { + num_rows: u32, + reduce_size: u32, + is_argmin: u32, + num_parts: u32, +} +@group(0) @binding(2) var params: Params; + +// Stage 1 of the two-stage arg-reduction. The single-workgroup kernel +// (arg_reduce.wgsl) puts an entire row on ONE 64-lane workgroup; here each row +// is split into `num_parts` contiguous chunks, one workgroup per chunk, so the +// row's elements are scanned by num_parts x WG lanes instead of WG. +// +// Semantics are held IDENTICAL to arg_reduce.wgsl: every lane is seeded with +// (t_in[base], 0) -- element 0 is always a valid candidate -- and scans with a +// STRICT compare, so the lowest index wins inside a lane. The in-workgroup +// merge is a halving tree using the same lower-index-wins rule; because no +// partial can ever hold a NaN unless t_in[base] itself is NaN (a strict +// compare never selects NaN), the merge operator is a total order and the tree +// is order-independent -> bit-exact index equality with the serial tail. +const WG: u32 = 64u; +var part_val: array; +var part_idx: array; + +@compute @workgroup_size(WG, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded grid of num_rows * num_parts workgroups. + let slot = wid.x + wid.y * num_workgroups.x; + let row = slot / params.num_parts; + if (row >= params.num_rows) { + return; + } + let part = slot - row * params.num_parts; + let base = row * params.reduce_size; + + let chunk = (params.reduce_size + params.num_parts - 1u) / params.num_parts; + let start = part * chunk; + var end = start + chunk; + if (end > params.reduce_size) { + end = params.reduce_size; + } + + var best = t_in[base]; + var best_idx: u32 = 0u; + var k = start + lid.x; + while (k < end) { + let v = t_in[base + k]; + if (params.is_argmin != 0u) { + if (v < best) { best = v; best_idx = k; } + } else { + if (v > best) { best = v; best_idx = k; } + } + k = k + WG; + } + part_val[lid.x] = best; + part_idx[lid.x] = best_idx; + workgroupBarrier(); + + var stride: u32 = WG >> 1u; + loop { + if (stride == 0u) { + break; + } + if (lid.x < stride) { + let bv = part_val[lid.x]; + let bi = part_idx[lid.x]; + let v = part_val[lid.x + stride]; + let idx = part_idx[lid.x + stride]; + if (params.is_argmin != 0u) { + if (v < bv || (v == bv && idx < bi)) { + part_val[lid.x] = v; + part_idx[lid.x] = idx; + } + } else { + if (v > bv || (v == bv && idx < bi)) { + part_val[lid.x] = v; + part_idx[lid.x] = idx; + } + } + } + workgroupBarrier(); + stride = stride >> 1u; + } + + if (lid.x == 0u) { + t_part[slot].val = part_val[0]; + t_part[slot].idx = part_idx[0]; + } +} +)"; + +inline constexpr uint32_t kArgReducePartialWorkgroupSizeX = 64; +inline constexpr uint32_t kArgReducePartialWorkgroupSizeY = 1; +inline constexpr uint32_t kArgReducePartialWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/cat/Cat.cpp b/backends/webgpu/runtime/ops/cat/Cat.cpp index 3c3c34d571c..2fd40022a87 100644 --- a/backends/webgpu/runtime/ops/cat/Cat.cpp +++ b/backends/webgpu/runtime/ops/cat/Cat.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -71,7 +72,7 @@ void cat_impl(WebGPUGraph& graph, const std::vector& args) { // Validate + cache input meta/wgc BEFORE any GPU alloc (no leak on throw). std::vector in_metas(ids.size()); - std::vector wg_counts(ids.size()); + std::vector wg_counts(ids.size()); int64_t concat_sum = 0; for (size_t k = 0; k < ids.size(); k++) { const int id = ids[k]; @@ -92,7 +93,7 @@ void cat_impl(WebGPUGraph& graph, const std::vector& args) { static_cast(in_metas[k].numel) * sizeof(float)) { throw std::runtime_error("cat: non-fp32 input (nbytes != numel * 4)"); } - wg_counts[k] = utils::compute_1d_workgroup_count( + wg_counts[k] = utils::compute_2d_workgroup_count( device, in_metas[k].numel, wg_size, "cat"); concat_sum += in_tensor.dims[dim]; } @@ -158,8 +159,8 @@ void cat_impl(WebGPUGraph& graph, const std::vector& args) { in_meta_bufs[k] = in_meta_buf; params_bufs[k] = params_buf; - dispatch_idxs[k] = - graph.add_dispatch({bundle.pipeline, bundle.bind_group, wg_counts[k]}); + dispatch_idxs[k] = graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, wg_counts[k].x, wg_counts[k].y); if (!shared_resources.has_value()) { shared_resources.emplace(std::move(bundle)); } @@ -211,9 +212,11 @@ void cat_impl(WebGPUGraph& graph, const std::vector& args) { params.off_k = off; wgpuQueueWriteBuffer( g.queue(), params_bufs[k], 0, ¶ms, sizeof(params)); - g.dispatch_at(dispatch_idxs[k]).workgroup_count_x = - utils::compute_1d_workgroup_count( - g.device(), in_meta.numel, wg_size, "cat(resize)"); + set_cat_dispatch_grid( + g, + dispatch_idxs[k], + utils::compute_2d_workgroup_count( + g.device(), in_meta.numel, wg_size, "cat(resize)")); off += static_cast(in_dims[cdim]); } }; diff --git a/backends/webgpu/runtime/ops/cat/CatDispatch.h b/backends/webgpu/runtime/ops/cat/CatDispatch.h new file mode 100644 index 00000000000..3a30f119032 --- /dev/null +++ b/backends/webgpu/runtime/ops/cat/CatDispatch.h @@ -0,0 +1,27 @@ +/* + * 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 { + +inline void set_cat_dispatch_grid( + WebGPUGraph& graph, + size_t dispatch_index, + const utils::WgCount& grid) { + WebGPUDispatch& dispatch = graph.dispatch_at(dispatch_index); + dispatch.workgroup_count_x = grid.x; + dispatch.workgroup_count_y = grid.y; +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/cat/cat.wgsl b/backends/webgpu/runtime/ops/cat/cat.wgsl index aa90f3979c2..02509c9591e 100644 --- a/backends/webgpu/runtime/ops/cat/cat.wgsl +++ b/backends/webgpu/runtime/ops/cat/cat.wgsl @@ -19,8 +19,10 @@ struct Params { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let in_bufi = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let in_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); if (in_bufi >= in_meta.numel) { return; } diff --git a/backends/webgpu/runtime/ops/cat/cat_wgsl.h b/backends/webgpu/runtime/ops/cat/cat_wgsl.h index 26e5f80f8a5..4153b4f6332 100644 --- a/backends/webgpu/runtime/ops/cat/cat_wgsl.h +++ b/backends/webgpu/runtime/ops/cat/cat_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from cat.wgsl - DO NOT EDIT. -// wgsl-sha256: 1a5f66607e2959c12d757989a3c1ae6fc6f2ede139bc9afb6e14dd3961a6fcb7 +// wgsl-sha256: 545ca763c55f5b15cfdbd5307a02349404e82dc0953268e853438def4be3e8fa inline constexpr const char* kCatWGSL = R"( @group(0) @binding(0) var input: array; @group(0) @binding(1) var output: array; @@ -36,8 +36,10 @@ struct Params { override wg_size: u32 = 64u; @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let in_bufi = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let in_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); if (in_bufi >= in_meta.numel) { return; } diff --git a/backends/webgpu/scripts/BUCK b/backends/webgpu/scripts/BUCK new file mode 100644 index 00000000000..4e76682cd45 --- /dev/null +++ b/backends/webgpu/scripts/BUCK @@ -0,0 +1,15 @@ +# 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. + +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") + +oncall("executorch") + +load(":targets.bzl", "define_common_targets") + +non_fbcode_target(_kind = define_common_targets) + +fbcode_target(_kind = define_common_targets) diff --git a/backends/webgpu/scripts/gen_wgsl_headers.py b/backends/webgpu/scripts/gen_wgsl_headers.py index 1e35b205888..9dd93531216 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -414,8 +414,30 @@ def _resolve_dim(tok: str, src: str) -> int: return int(m.group(1)) +def strip_wgsl_comments(src: str) -> str: + """Blank comments while preserving source offsets and line structure.""" + out = [] + i = 0 + while i < len(src): + if src.startswith("//", i): + end = src.find("\n", i) + end = len(src) if end < 0 else end + out.append(" " * (end - i)) + i = end + elif src.startswith("/*", i): + end = src.find("*/", i + 2) + end = len(src) if end < 0 else end + 2 + out.append("".join(c if c == "\n" else " " for c in src[i:end])) + i = end + else: + out.append(src[i]) + i += 1 + return "".join(out) + + def parse_workgroup_size(src: str) -> tuple[int, int, int]: """Resolve the (x, y, z) dims of @workgroup_size; y and z default to 1.""" + src = strip_wgsl_comments(src) m = re.search(r"@workgroup_size\s*\(([^)]*)\)", src) if not m: raise ValueError("no @workgroup_size found") diff --git a/backends/webgpu/scripts/targets.bzl b/backends/webgpu/scripts/targets.bzl new file mode 100644 index 00000000000..f448ec0ae35 --- /dev/null +++ b/backends/webgpu/scripts/targets.bzl @@ -0,0 +1,22 @@ +# 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. + +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + runtime.python_library( + name = "webgpu_artifact_manifest", + srcs = ["webgpu_artifact_manifest.py"], + base_module = "executorch.backends.webgpu.scripts", + visibility = ["PUBLIC"], + ) + + runtime.python_binary( + name = "webgpu_artifact_manifest_cli", + main_module = "executorch.backends.webgpu.scripts.webgpu_artifact_manifest", + deps = [":webgpu_artifact_manifest"], + visibility = ["PUBLIC"], + ) diff --git a/backends/webgpu/scripts/webgpu_artifact_manifest.py b/backends/webgpu/scripts/webgpu_artifact_manifest.py new file mode 100644 index 00000000000..5905da7d566 --- /dev/null +++ b/backends/webgpu/scripts/webgpu_artifact_manifest.py @@ -0,0 +1,176 @@ +# 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. + +"""Create and validate reproducible WebGPU artifact manifests.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path +from typing import Iterable, Mapping, Sequence + + +SCHEMA_VERSION = 1 +ALLOWED_SINGLE_ROLES = frozenset( + {"javascript", "wasm", "pte", "source", "object", "link_map"} +) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _contained_file(root: Path, path: Path) -> tuple[Path, str]: + resolved_root = root.resolve(strict=True) + candidate = path if path.is_absolute() else root / path + if candidate.is_symlink(): + raise ValueError(f"artifact symlink escapes are not allowed: {path}") + resolved = candidate.resolve(strict=True) + try: + relative = resolved.relative_to(resolved_root) + except ValueError as error: + raise ValueError(f"artifact escapes manifest root: {path}") from error + if not resolved.is_file(): + raise ValueError(f"artifact is not a regular file: {path}") + return resolved, relative.as_posix() + + +def create_manifest( + root: Path, + role_paths: Mapping[str, Path], + ptd_paths: Sequence[Path] = (), +) -> dict[str, object]: + unknown_roles = set(role_paths) - ALLOWED_SINGLE_ROLES + if unknown_roles: + raise ValueError(f"unsupported artifact roles: {sorted(unknown_roles)}") + + artifacts: list[dict[str, object]] = [] + for role in sorted(role_paths): + resolved, relative = _contained_file(root, role_paths[role]) + artifacts.append( + { + "bytes": resolved.stat().st_size, + "path": relative, + "role": role, + "sha256": _sha256(resolved), + } + ) + + ptd_order: list[str] = [] + for ptd in ptd_paths: + resolved, relative = _contained_file(root, ptd) + if relative in ptd_order: + raise ValueError(f"duplicate PTD path: {relative}") + ptd_order.append(relative) + artifacts.append( + { + "bytes": resolved.stat().st_size, + "path": relative, + "role": "ptd", + "sha256": _sha256(resolved), + } + ) + + return { + "artifacts": artifacts, + "ptd_order": ptd_order, + "schema_version": SCHEMA_VERSION, + } + + +def _validate_artifact_entry( + root: Path, + artifact: object, + seen_single_roles: set[str], +) -> tuple[str, str]: + if not isinstance(artifact, dict): + raise ValueError("artifact entry must be an object") + role = artifact.get("role") + path_value = artifact.get("path") + if not isinstance(role, str) or not isinstance(path_value, str): + raise ValueError("artifact role/path must be strings") + if Path(path_value).is_absolute(): + raise ValueError("manifest stores an absolute artifact path") + if role != "ptd": + if role not in ALLOWED_SINGLE_ROLES: + raise ValueError(f"unsupported artifact role: {role}") + if role in seen_single_roles: + raise ValueError(f"duplicate singleton artifact role: {role}") + seen_single_roles.add(role) + + resolved, normalized = _contained_file(root, Path(path_value)) + if normalized != path_value: + raise ValueError(f"non-canonical artifact path: {path_value}") + if artifact.get("bytes") != resolved.stat().st_size: + raise ValueError(f"artifact byte count mismatch: {path_value}") + if artifact.get("sha256") != _sha256(resolved): + raise ValueError(f"artifact SHA-256 mismatch: {path_value}") + return role, path_value + + +def validate_manifest(root: Path, manifest: Mapping[str, object]) -> None: + if manifest.get("schema_version") != SCHEMA_VERSION: + raise ValueError("unsupported manifest schema version") + artifacts = manifest.get("artifacts") + ptd_order = manifest.get("ptd_order") + if not isinstance(artifacts, list) or not isinstance(ptd_order, list): + raise ValueError("manifest artifacts/PTD order must be lists") + + seen_single_roles: set[str] = set() + observed_ptds: list[str] = [] + for artifact in artifacts: + role, path_value = _validate_artifact_entry(root, artifact, seen_single_roles) + if role == "ptd": + observed_ptds.append(path_value) + + if observed_ptds != ptd_order: + raise ValueError("manifest PTD order does not match artifact order") + + +def _parse_roles(values: Iterable[str]) -> dict[str, Path]: + result: dict[str, Path] = {} + for value in values: + role, separator, path = value.partition("=") + if not separator or not role or not path or role in result: + raise ValueError(f"invalid or duplicate ROLE=PATH: {value}") + result[role] = Path(path) + return result + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser() + commands = parser.add_subparsers(dest="command", required=True) + create = commands.add_parser("create") + create.add_argument("--root", type=Path, required=True) + create.add_argument("--output", type=Path, required=True) + create.add_argument("--role", action="append", default=[]) + create.add_argument("--ptd", action="append", type=Path, default=[]) + validate = commands.add_parser("validate") + validate.add_argument("--root", type=Path, required=True) + validate.add_argument("--manifest", type=Path, required=True) + args = parser.parse_args(argv) + + if args.command == "create": + manifest = create_manifest(args.root, _parse_roles(args.role), args.ptd) + args.output.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return 0 + + manifest = json.loads(args.manifest.read_text(encoding="utf-8")) + validate_manifest(args.root, manifest) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/backends/webgpu/targets.bzl b/backends/webgpu/targets.bzl new file mode 100644 index 00000000000..756e952ac82 --- /dev/null +++ b/backends/webgpu/targets.bzl @@ -0,0 +1,103 @@ +# 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. + +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +_DAWN_DEPS = [ + "fbsource//third-party/dawn:dawn_common", + "fbsource//third-party/dawn:dawn_native", + "fbsource//third-party/dawn:dawn_platform", + "fbsource//third-party/dawn:dawn_proc", +] + +def define_common_targets(): + runtime.cxx_library( + name = "webgpu_backend", + srcs = native.glob(["runtime/**/*.cpp"]), + exported_headers = native.glob(["runtime/**/*.h"]), + compiler_flags = [ + "-DWEBGPU_DAWN_INSTANCE_CAPABILITIES", + "-fexceptions", + ], + link_whole = True, + exported_deps = _DAWN_DEPS + [ + "//executorch/backends/vulkan/serialization:vk_delegate_schema", + "//executorch/runtime/backend:interface", + "//executorch/runtime/core:core", + "//executorch/runtime/core:named_data_map", + "//executorch/runtime/core/exec_aten/util:tensor_util", + ], + visibility = ["PUBLIC"], + ) + + runtime.cxx_library( + name = "webgpu_model_loader", + srcs = ["runner/webgpu_model_loader.cpp"], + exported_headers = ["runner/webgpu_model_loader.h"], + compiler_flags = [ + "-DC10_USING_CUSTOM_GENERATED_MACROS", + "-fexceptions", + ], + exported_preprocessor_flags = [ + "-DC10_USING_CUSTOM_GENERATED_MACROS", + ], + exported_deps = [ + "//executorch/extension/module:module", + ], + visibility = ["PUBLIC"], + ) + + runtime.cxx_test( + name = "webgpu_model_loader_test", + srcs = ["test/native/test_model_loader.cpp"], + deps = [":webgpu_model_loader"], + ) + + runtime.cxx_test( + name = "webgpu_utils_test", + srcs = ["test/native/test_webgpu_utils.cpp"], + deps = [":webgpu_backend"], + ) + + runtime.cxx_test( + name = "webgpu_device_header_test", + srcs = ["test/native/test_webgpu_device_header.cpp"], + deps = [":webgpu_backend"], + ) + + runtime.cxx_test( + name = "webgpu_default_context_test", + srcs = ["test/native/test_webgpu_default_context.cpp"], + deps = [":webgpu_backend"], + ) + + runtime.cxx_test( + name = "webgpu_query_pool_test", + srcs = ["test/native/test_webgpu_query_pool.cpp"], + deps = [":webgpu_backend"], + ) + + runtime.cxx_test( + name = "webgpu_query_pool_profiled_test", + srcs = [ + "runtime/WebGPUQueryPool.cpp", + "test/native/test_webgpu_query_pool.cpp", + ], + compiler_flags = [ + "-DWEBGPU_BACKEND_ENABLE_PROFILING", + "-DWEBGPU_DAWN_INSTANCE_CAPABILITIES", + "-fexceptions", + ], + deps = [":webgpu_backend"], + ) + + runtime.cxx_test( + name = "webgpu_execution_options_test", + srcs = [ + "runtime/WebGPUExecutionOptions.cpp", + "test/native/test_execution_options.cpp", + ], + ) diff --git a/backends/webgpu/test/BUCK b/backends/webgpu/test/BUCK index 49f54f27d3a..bcddaa0f66b 100644 --- a/backends/webgpu/test/BUCK +++ b/backends/webgpu/test/BUCK @@ -1,6 +1,6 @@ load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") -load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") oncall("executorch") @@ -19,6 +19,15 @@ fbcode_target( ], ) +non_fbcode_target( + _kind = runtime.python_test, + name = "test_webgpu_artifact_manifest", + srcs = ["test_webgpu_artifact_manifest.py"], + deps = [ + "//executorch/backends/webgpu/scripts:webgpu_artifact_manifest", + ], +) + fbcode_target( _kind = python_unittest, name = "test_index", @@ -42,3 +51,12 @@ fbcode_target( "//executorch/backends/vulkan:vulkan_preprocess", ], ) + +fbcode_target( + _kind = python_unittest, + name = "test_webgpu_artifact_manifest", + srcs = ["test_webgpu_artifact_manifest.py"], + deps = [ + "//executorch/backends/webgpu/scripts:webgpu_artifact_manifest", + ], +) diff --git a/backends/webgpu/test/native/test_dispatch_order.cpp b/backends/webgpu/test/native/test_dispatch_order.cpp index d8aa627eff3..2815dd1b510 100644 --- a/backends/webgpu/test/native/test_dispatch_order.cpp +++ b/backends/webgpu/test/native/test_dispatch_order.cpp @@ -7,10 +7,13 @@ */ #include +#include +#include #include #include #include +#include #include #include @@ -66,32 +69,96 @@ void run_case(const char* name, const std::vector& sizes) { ASSERT_EQ(input.size(), expected) << "input numel " << input.size() << " != expected " << expected << " for " << name; - auto x = make_tensor_ptr(sizes, std::vector(input)); - auto result = module.forward({EValue(x)}); - ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() - << ")"; - const auto& outputs = result.get(); - ASSERT_TRUE(!outputs.empty() && outputs[0].isTensor()) << "no tensor output"; - const auto& out_tensor = outputs[0].toTensor(); - ASSERT_EQ(static_cast(out_tensor.numel()), golden.size()) - << "output numel " << (size_t)out_tensor.numel() << " != golden " - << golden.size(); - const float* out_data = out_tensor.const_data_ptr(); - - float max_abs_err = 0.0f; - float max_rel_err = 0.0f; - for (size_t i = 0; i < golden.size(); i++) { - const float abs_err = std::abs(out_data[i] - golden[i]); - max_abs_err = std::max(max_abs_err, abs_err); - const float denom = std::max(std::abs(golden[i]), 1e-6f); - max_rel_err = std::max(max_rel_err, abs_err / denom); + struct Mode { + bool single_compute_pass; + size_t cap; + }; + const std::vector modes = { + {false, 0}, {true, 0}, {true, 2}, {false, 0}}; + std::string command_inventory; + uint64_t prior_ordinal = 0; + for (const Mode mode : modes) { + auto x = make_tensor_ptr(sizes, std::vector(input)); + WebGPUExecutionOptions options; + options.single_compute_pass = mode.single_compute_pass; + options.max_compute_dispatches_per_pass = mode.cap; + auto result = with_webgpu_execution_options( + options, [&]() { return module.forward({EValue(x)}); }); + ASSERT_TRUE(result.ok()) << "forward failed (error " << (int)result.error() + << ")"; + const auto& outputs = result.get(); + ASSERT_TRUE(!outputs.empty() && outputs[0].isTensor()) + << "no tensor output"; + const auto& out_tensor = outputs[0].toTensor(); + ASSERT_EQ(static_cast(out_tensor.numel()), golden.size()); + const float* out_data = out_tensor.const_data_ptr(); + + float max_abs_err = 0.0f; + float max_rel_err = 0.0f; + for (size_t i = 0; i < golden.size(); i++) { + const float abs_err = std::abs(out_data[i] - golden[i]); + max_abs_err = std::max(max_abs_err, abs_err); + const float denom = std::max(std::abs(golden[i]), 1e-6f); + max_rel_err = std::max(max_rel_err, abs_err / denom); + } + EXPECT_FALSE(max_abs_err > 1e-3f && max_rel_err > 1e-3f) + << "dispatch_order[" << name + << "] exceeds tolerance 1e-3 (max_abs_err=" << max_abs_err + << " max_rel_err=" << max_rel_err << ")"; + + const auto attestation = nlohmann::json::parse( + webgpu_backend_execution_attestation_json()); + EXPECT_EQ( + attestation.at("requested").get(), mode.single_compute_pass); + EXPECT_EQ( + attestation.at("applied").get(), mode.single_compute_pass); + EXPECT_TRUE(attestation.at("completed").get()); + EXPECT_EQ(attestation.at("queueSubmitCount").get(), 1u); + const size_t active = + attestation.at("activeComputeCount").get(); + const size_t runs = + attestation.at("maximalComputeRuns").get(); + ASSERT_GT(active, 0u); + ASSERT_GT(runs, 0u); + size_t expected_passes = 0; + size_t dispatches_in_pass = 0; + for (const auto& command : + attestation.at("canonicalCommands").at("commands")) { + if (command.at("kind") != "compute") { + if (command.at("enabled").get()) { + dispatches_in_pass = 0; + } + continue; + } + if (!command.at("enabled").get() || + command.at("zeroGrid").get()) { + continue; + } + if (!mode.single_compute_pass || dispatches_in_pass == 0) { + ++expected_passes; + } + ++dispatches_in_pass; + if (!mode.single_compute_pass || + (mode.cap != 0 && dispatches_in_pass >= mode.cap)) { + dispatches_in_pass = 0; + } + } + EXPECT_EQ( + attestation.at("encodedComputePasses").get(), expected_passes); + EXPECT_EQ( + attestation.at("maxComputeDispatchesPerPass").get(), mode.cap); + const uint64_t ordinal = + attestation.at("executionOrdinal").get(); + EXPECT_EQ(ordinal, prior_ordinal + 1); + prior_ordinal = ordinal; + const std::string current_inventory = + attestation.at("canonicalCommands").dump(); + if (command_inventory.empty()) { + command_inventory = current_inventory; + } else { + EXPECT_EQ(current_inventory, command_inventory); + } } - // Lenient gate: pass iff abs<=tol OR rel<=tol (near-zero goldens). - EXPECT_FALSE(max_abs_err > 1e-3f && max_rel_err > 1e-3f) - << "dispatch_order[" << name - << "] exceeds tolerance 1e-3 (max_abs_err=" << max_abs_err - << " max_rel_err=" << max_rel_err << ", " << golden.size() - << " elements)"; } } // namespace diff --git a/backends/webgpu/test/native/test_execution_options.cpp b/backends/webgpu/test/native/test_execution_options.cpp index 382bfbf7efa..1aef37cd7f6 100644 --- a/backends/webgpu/test/native/test_execution_options.cpp +++ b/backends/webgpu/test/native/test_execution_options.cpp @@ -17,8 +17,122 @@ namespace executorch::backends::webgpu { namespace { TEST(WebGPUExecutionOptionsTest, DefaultsToPreservingOutputs) { - EXPECT_EQ( - current_webgpu_execution_options().discardable_output_data, nullptr); + const auto options = current_webgpu_execution_options(); + EXPECT_EQ(options.discardable_output_data, nullptr); + EXPECT_FALSE(options.single_compute_pass); + EXPECT_EQ(options.max_compute_dispatches_per_pass, 0u); +} + +TEST(WebGPUExecutionOptionsTest, SinglePassOptionsSurviveResolution) { + WebGPUExecutionOptions options; + options.single_compute_pass = true; + options.max_compute_dispatches_per_pass = 7; + + const auto no_suppression = + resolve_webgpu_graph_execution_options({}, options); + EXPECT_TRUE(no_suppression.single_compute_pass); + EXPECT_EQ(no_suppression.max_compute_dispatches_per_pass, 7u); + + int output = 0; + options.discardable_output_data = &output; + options.exact_method_certificate_verified = false; + const auto uncertified = + resolve_webgpu_graph_execution_options({&output}, options); + EXPECT_TRUE(uncertified.single_compute_pass); + EXPECT_EQ(uncertified.max_compute_dispatches_per_pass, 7u); + EXPECT_EQ(uncertified.suppress_output_ordinal, kNoOutputOrdinal); +} + +TEST(WebGPUExecutionOptionsTest, SinglePassScopeRestoresEveryField) { + WebGPUExecutionOptions options; + options.single_compute_pass = true; + options.max_compute_dispatches_per_pass = 3; + { + ScopedWebGPUExecutionOptions scope(options); + const auto current = current_webgpu_execution_options(); + EXPECT_TRUE(current.single_compute_pass); + EXPECT_EQ(current.max_compute_dispatches_per_pass, 3u); + } + const auto restored = current_webgpu_execution_options(); + EXPECT_FALSE(restored.single_compute_pass); + EXPECT_EQ(restored.max_compute_dispatches_per_pass, 0u); +} + +TEST(WebGPUExecutionOptionsTest, PassCapZeroIsUnlimited) { + EXPECT_FALSE(webgpu_pass_cap_reached(0, 0)); + EXPECT_FALSE(webgpu_pass_cap_reached(1000, 0)); +} + +TEST(WebGPUExecutionOptionsTest, PassCapClosesAtExactDispatchCount) { + EXPECT_FALSE(webgpu_pass_cap_reached(2, 3)); + EXPECT_TRUE(webgpu_pass_cap_reached(3, 3)); + EXPECT_TRUE(webgpu_pass_cap_reached(4, 3)); +} + +TEST(WebGPUExecutionOptionsTest, CommandInventoryTracksComputeRunsAndCopies) { + std::vector commands(5); + commands[0].identity = "first"; + commands[1].identity = "disabled"; + commands[1].enabled = false; + commands[1].zero_grid = true; + commands[2].kind = WebGPUCommandKind::GraphCopy; + commands[3].identity = "second"; + commands[4].kind = WebGPUCommandKind::OutputCopy; + + const auto inventory = build_webgpu_command_inventory(commands); + EXPECT_EQ(inventory.static_dispatch_records, 4u); + EXPECT_EQ(inventory.active_compute_count, 2u); + EXPECT_EQ(inventory.zero_grid_compute_count, 1u); + EXPECT_EQ(inventory.graph_copy_count, 1u); + EXPECT_EQ(inventory.output_copy_count, 1u); + EXPECT_EQ(inventory.maximal_compute_runs, 2u); + EXPECT_NE(inventory.canonical_commands_json.find("first"), std::string::npos); + EXPECT_NE(inventory.canonical_commands_json.find("second"), std::string::npos); +} + +TEST(WebGPUExecutionOptionsTest, ComputePassCountHonorsCopiesAndCap) { + std::vector commands(7); + commands[0].identity = "a"; + commands[1].identity = "b"; + commands[2].identity = "disabled"; + commands[2].enabled = false; + commands[3].kind = WebGPUCommandKind::GraphCopy; + commands[4].identity = "c"; + commands[5].identity = "d"; + commands[6].identity = "e"; + + EXPECT_EQ(count_webgpu_compute_passes(commands, false, 0), 5u); + EXPECT_EQ(count_webgpu_compute_passes(commands, true, 0), 2u); + EXPECT_EQ(count_webgpu_compute_passes(commands, true, 2), 3u); + EXPECT_THROW( + count_webgpu_compute_passes(commands, false, 2), std::invalid_argument); +} + +TEST(WebGPUExecutionOptionsTest, DisabledCopyDoesNotSplitComputeRun) { + std::vector commands(3); + commands[0].identity = "first"; + commands[1].kind = WebGPUCommandKind::GraphCopy; + commands[1].enabled = false; + commands[2].identity = "second"; + EXPECT_EQ(count_webgpu_compute_passes(commands, true, 0), 1u); +} + +TEST(WebGPUExecutionOptionsTest, AttestationSerializesObservedPassCounts) { + WebGPUExecutionAttestation attestation; + attestation.execution_ordinal = 3; + attestation.requested = true; + attestation.applied = true; + attestation.completed = true; + attestation.encoded_compute_passes = 2; + attestation.queue_submit_count = 1; + attestation.max_compute_dispatches_per_pass = 4; + + const std::string json = serialize_webgpu_execution_attestation(attestation); + EXPECT_NE(json.find("\"executionOrdinal\":3"), std::string::npos); + EXPECT_NE(json.find("\"encodedComputePasses\":2"), std::string::npos); + EXPECT_NE( + json.find("\"maxComputeDispatchesPerPass\":4"), std::string::npos); + EXPECT_NE(json.find("\"completed\":true"), std::string::npos); } TEST(WebGPUExecutionOptionsTest, NestedScopesRestorePriorValue) { diff --git a/backends/webgpu/test/native/test_model_loader.cpp b/backends/webgpu/test/native/test_model_loader.cpp new file mode 100644 index 00000000000..84c58d822aa --- /dev/null +++ b/backends/webgpu/test/native/test_model_loader.cpp @@ -0,0 +1,72 @@ +/* + * 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 + +namespace executorch::backends::webgpu { +namespace { + +TEST(WebGPUModelLoaderTest, RejectsEmptyProgramPath) { + WebGPUModelLoadSpec spec; + spec.required_methods = {"forward"}; + const auto result = load_webgpu_model(std::move(spec)); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.error(), runtime::Error::InvalidArgument); +} + +TEST(WebGPUModelLoaderTest, RejectsEmptyRequiredMethods) { + WebGPUModelLoadSpec spec; + spec.pte_path = "not-opened.pte"; + const auto result = load_webgpu_model(std::move(spec)); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.error(), runtime::Error::InvalidArgument); +} + +TEST(WebGPUModelLoaderTest, RejectsDuplicateRequiredMethod) { + WebGPUModelLoadSpec spec; + spec.pte_path = "not-opened.pte"; + spec.required_methods = {"forward", "forward"}; + const auto result = load_webgpu_model(std::move(spec)); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.error(), runtime::Error::InvalidArgument); +} + +TEST(WebGPUModelLoaderTest, RejectsEmptyPtdPath) { + WebGPUModelLoadSpec spec; + spec.pte_path = "not-opened.pte"; + spec.ptd_paths = {""}; + spec.required_methods = {"forward"}; + const auto result = load_webgpu_model(std::move(spec)); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.error(), runtime::Error::InvalidArgument); +} + +TEST(WebGPUModelLoaderTest, RejectsDuplicatePtdPath) { + WebGPUModelLoadSpec spec; + spec.pte_path = "not-opened.pte"; + spec.ptd_paths = {"weights.ptd", "weights.ptd"}; + spec.required_methods = {"forward"}; + const auto result = load_webgpu_model(std::move(spec)); + EXPECT_FALSE(result.ok()); + EXPECT_EQ(result.error(), runtime::Error::InvalidArgument); +} + +TEST(WebGPUModelLoaderTest, InvalidProgramFailsWithoutPublishingModule) { + WebGPUModelLoadSpec spec; + spec.pte_path = "does-not-exist.pte"; + spec.ptd_paths = {"first.ptd", "second.ptd"}; + spec.required_methods = {"forward"}; + spec.load_mode = extension::Module::LoadMode::Mmap; + const auto result = load_webgpu_model(std::move(spec)); + EXPECT_FALSE(result.ok()); +} + +} // namespace +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_webgpu_default_context.cpp b/backends/webgpu/test/native/test_webgpu_default_context.cpp new file mode 100644 index 00000000000..7603cdcfce7 --- /dev/null +++ b/backends/webgpu/test/native/test_webgpu_default_context.cpp @@ -0,0 +1,93 @@ +/* + * 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 + +namespace executorch::backends::webgpu { +namespace { + +class ExplicitDefaultContextTest : public ::testing::Test { + protected: + void SetUp() override { + set_default_webgpu_context(nullptr); + } + + void TearDown() override { + set_default_webgpu_context(nullptr); + } +}; + +TEST_F(ExplicitDefaultContextTest, ClaimAndReleaseRequireExactOwner) { + WebGPUContext first; + WebGPUContext second; + + EXPECT_EQ(get_explicit_default_webgpu_context(), nullptr); + EXPECT_TRUE(compare_and_set_default_webgpu_context(nullptr, &first)); + EXPECT_EQ(get_explicit_default_webgpu_context(), &first); + EXPECT_EQ(get_default_webgpu_context(), &first); + + EXPECT_FALSE(compare_and_set_default_webgpu_context(nullptr, &second)); + EXPECT_FALSE(compare_and_set_default_webgpu_context(&second, nullptr)); + EXPECT_EQ(get_explicit_default_webgpu_context(), &first); + + EXPECT_TRUE(compare_and_set_default_webgpu_context(&first, nullptr)); + EXPECT_EQ(get_explicit_default_webgpu_context(), nullptr); +} + +TEST_F(ExplicitDefaultContextTest, ConcurrentClaimsHaveOneWinner) { + constexpr size_t kClaimants = 8; + std::array contexts; + std::array, kClaimants> claimed; + for (auto& value : claimed) { + value.store(false); + } + + std::vector threads; + threads.reserve(kClaimants); + for (size_t i = 0; i < kClaimants; i++) { + threads.emplace_back([&, i]() { + claimed[i].store( + compare_and_set_default_webgpu_context(nullptr, &contexts[i])); + }); + } + for (auto& thread : threads) { + thread.join(); + } + + size_t winner = kClaimants; + size_t winners = 0; + for (size_t i = 0; i < kClaimants; i++) { + if (claimed[i].load()) { + winner = i; + winners++; + } + } + ASSERT_EQ(winners, 1u); + ASSERT_LT(winner, kClaimants); + EXPECT_EQ(get_explicit_default_webgpu_context(), &contexts[winner]); + + for (size_t i = 0; i < kClaimants; i++) { + if (i != winner) { + EXPECT_FALSE( + compare_and_set_default_webgpu_context(&contexts[i], nullptr)); + } + } + EXPECT_TRUE( + compare_and_set_default_webgpu_context(&contexts[winner], nullptr)); + EXPECT_EQ(get_explicit_default_webgpu_context(), nullptr); +} + +} // namespace +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_webgpu_device_header.cpp b/backends/webgpu/test/native/test_webgpu_device_header.cpp new file mode 100644 index 00000000000..dbae66d7482 --- /dev/null +++ b/backends/webgpu/test/native/test_webgpu_device_header.cpp @@ -0,0 +1,24 @@ +/* + * 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 + +namespace executorch::backends::webgpu { +namespace { + +TEST(WebGPUDeviceHeader, CompilesWithoutDirectDawnDependency) { + WebGPUContext context; + EXPECT_EQ(context.instance, nullptr); + EXPECT_EQ(context.device, nullptr); + EXPECT_EQ(context.queue, nullptr); +} + +} // namespace +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_webgpu_query_pool.cpp b/backends/webgpu/test/native/test_webgpu_query_pool.cpp new file mode 100644 index 00000000000..c38dea9aa71 --- /dev/null +++ b/backends/webgpu/test/native/test_webgpu_query_pool.cpp @@ -0,0 +1,65 @@ +/* + * 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 + +namespace executorch::backends::webgpu { +namespace { + +TEST(WebGPUQueryResultState, StartsInvalidAtGenerationZero) { + detail::WebGPUQueryResultState state; + EXPECT_FALSE(state.results_valid()); + EXPECT_EQ(state.result_generation(), 0u); +} + +TEST(WebGPUQueryResultState, InvalidationPreservesMonotonicGeneration) { + detail::WebGPUQueryResultState state; + state.complete(); + EXPECT_TRUE(state.results_valid()); + EXPECT_EQ(state.result_generation(), 1u); + + state.invalidate(); + EXPECT_FALSE(state.results_valid()); + EXPECT_EQ(state.result_generation(), 1u); + + state.complete(); + EXPECT_TRUE(state.results_valid()); + EXPECT_EQ(state.result_generation(), 2u); +} + +TEST(WebGPUQueryResultState, ExtractionPredicateFailsClosed) { + const uint64_t ticks[] = {10, 20}; + + EXPECT_TRUE(detail::query_result_extraction_succeeded( + 0, WGPUWaitStatus_TimedOut, WGPUMapAsyncStatus_Error, nullptr)); + EXPECT_FALSE(detail::query_result_extraction_succeeded( + 1, + WGPUWaitStatus_TimedOut, + WGPUMapAsyncStatus_Success, + ticks)); + EXPECT_FALSE(detail::query_result_extraction_succeeded( + 1, + WGPUWaitStatus_Success, + WGPUMapAsyncStatus_Error, + ticks)); + EXPECT_FALSE(detail::query_result_extraction_succeeded( + 1, + WGPUWaitStatus_Success, + WGPUMapAsyncStatus_Success, + nullptr)); + EXPECT_TRUE(detail::query_result_extraction_succeeded( + 1, + WGPUWaitStatus_Success, + WGPUMapAsyncStatus_Success, + ticks)); +} + +} // namespace +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_webgpu_utils.cpp b/backends/webgpu/test/native/test_webgpu_utils.cpp index a839d224b16..f60dc7be2fd 100644 --- a/backends/webgpu/test/native/test_webgpu_utils.cpp +++ b/backends/webgpu/test/native/test_webgpu_utils.cpp @@ -6,11 +6,11 @@ * LICENSE file in the root directory of this source tree. */ -// Device-free unit tests for the dispatch-grid math (WebGPUDispatchMath.h has -// zero WebGPU/Dawn dependency, unlike WebGPUUtils.h which needs a WGPUDevice -// for its other helpers). +// Device-free unit tests for pure WebGPU utility math. The shared utility +// header also exposes device-taking helpers, but these tests do not call them. -#include +#include +#include #include @@ -70,3 +70,80 @@ TEST(WebGPUUtils, DispatchGridThrowsPastCapacity) { static_cast(max_dim) * max_dim + 1u, 1u, max_dim, "test"), std::runtime_error); } + +TEST(WebGPUUtils, RowChunkingKeepsFittingRowsTogether) { + const utils::RowChunking chunking = + utils::compute_row_chunking(1024u, 256u, 10u, "test"); + EXPECT_EQ(chunking.rows_per_chunk, 4u); + EXPECT_EQ(chunking.num_chunks, 3u); +} + +TEST(WebGPUUtils, BindingSpecCarriesExplicitBufferOffset) { + const utils::BindingSpec binding = { + 3u, WGPUBufferBindingType_Storage, nullptr, 64u, 256u}; + const WGPUBindGroupEntry entry = utils::make_bind_group_entry(binding); + EXPECT_EQ(entry.binding, 3u); + EXPECT_EQ(entry.buffer, nullptr); + EXPECT_EQ(entry.size, 64u); + EXPECT_EQ(entry.offset, 256u); +} + +TEST(WebGPUUtils, BindingSpecDefaultsBufferOffsetToZero) { + const utils::BindingSpec binding = { + 3u, WGPUBufferBindingType_Storage, nullptr, 64u}; + const WGPUBindGroupEntry entry = utils::make_bind_group_entry(binding); + EXPECT_EQ(entry.offset, 0u); +} + +TEST(WebGPUUtils, RowChunkingUsesOneChunkWhenAllRowsFit) { + const utils::RowChunking chunking = + utils::compute_row_chunking(1024u, 16u, 8u, "test"); + EXPECT_EQ(chunking.rows_per_chunk, 8u); + EXPECT_EQ(chunking.num_chunks, 1u); +} + +TEST(WebGPUUtils, RowChunkingRejectsInvalidArguments) { + EXPECT_THROW( + utils::compute_row_chunking(0u, 1u, 1u, "test"), std::runtime_error); + EXPECT_THROW( + utils::compute_row_chunking(1u, 0u, 1u, "test"), std::runtime_error); + EXPECT_THROW( + utils::compute_row_chunking(1u, 1u, 0u, "test"), std::runtime_error); + EXPECT_THROW( + utils::compute_row_chunking(1u, 2u, 1u, "test"), std::runtime_error); +} + +TEST(WebGPUUtils, RowChunkingRejectsChunkCountsAboveUint32) { + EXPECT_THROW( + utils::compute_row_chunking( + 1u, + 1u, + static_cast(std::numeric_limits::max()) + 1u, + "test"), + std::runtime_error); +} + +TEST(WebGPUUtils, ArgReduceRouteKeepsShortRowsOnGenericKernel) { + EXPECT_EQ(select_arg_reduce_parts(1u, 4095u, 65535u), 0u); + EXPECT_EQ(select_arg_reduce_parts(0u, 262144u, 65535u), 0u); +} + +TEST(WebGPUUtils, ArgReduceRouteSelectsLongVocabularyRows) { + EXPECT_EQ(select_arg_reduce_parts(1u, 4096u, 65535u), 4u); + EXPECT_EQ(select_arg_reduce_parts(1u, 262144u, 65535u), 256u); + EXPECT_EQ( + select_arg_reduce_parts(1u, std::numeric_limits::max(), 65535u), + 256u); +} + +TEST(WebGPUUtils, ArgReduceRouteFailsClosedWhenScratchOrGridWouldOverflow) { + EXPECT_EQ(select_arg_reduce_parts(4097u, 262144u, 65535u), 0u); + EXPECT_EQ(select_arg_reduce_parts(1u, 262144u, 1u), 0u); +} + +TEST(WebGPUUtils, ArgReduceResizeRejectsScratchGrowth) { + EXPECT_TRUE(arg_reduce_partial_slots_fit(256u, 1u, 256u)); + EXPECT_TRUE(arg_reduce_partial_slots_fit(256u, 256u, 1u)); + EXPECT_FALSE(arg_reduce_partial_slots_fit(256u, 257u, 1u)); + EXPECT_FALSE(arg_reduce_partial_slots_fit(256u, 129u, 2u)); +} diff --git a/backends/webgpu/test/test_webgpu_artifact_manifest.py b/backends/webgpu/test/test_webgpu_artifact_manifest.py new file mode 100644 index 00000000000..fed14389466 --- /dev/null +++ b/backends/webgpu/test/test_webgpu_artifact_manifest.py @@ -0,0 +1,81 @@ +# 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. + +import tempfile +import unittest +from pathlib import Path + +from executorch.backends.webgpu.scripts.webgpu_artifact_manifest import ( + create_manifest, + validate_manifest, +) + + +class WebGPUArtifactManifestTest(unittest.TestCase): + def test_round_trip_preserves_ordered_ptds(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + for name, data in { + "runner.js": b"js", + "runner.wasm": b"wasm", + "model.pte": b"pte", + "first.ptd": b"first", + "second.ptd": b"second", + }.items(): + (root / name).write_bytes(data) + manifest = create_manifest( + root, + { + "javascript": Path("runner.js"), + "wasm": Path("runner.wasm"), + "pte": Path("model.pte"), + }, + [Path("first.ptd"), Path("second.ptd")], + ) + validate_manifest(root, manifest) + self.assertEqual(manifest["ptd_order"], ["first.ptd", "second.ptd"]) + self.assertNotIn(str(root), str(manifest)) + + def test_tampered_bytes_fail_validation(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + artifact = root / "runner.wasm" + artifact.write_bytes(b"before") + manifest = create_manifest(root, {"wasm": artifact}) + artifact.write_bytes(b"after") + with self.assertRaisesRegex(ValueError, "mismatch"): + validate_manifest(root, manifest) + + def test_symlink_escape_is_rejected(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) / "root" + root.mkdir() + outside = Path(directory) / "outside.pte" + outside.write_bytes(b"outside") + (root / "escape.pte").symlink_to(outside) + with self.assertRaisesRegex(ValueError, "escapes"): + create_manifest(root, {"pte": Path("escape.pte")}) + + def test_in_root_symlink_is_rejected(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + target = root / "model.pte" + target.write_bytes(b"model") + (root / "alias.pte").symlink_to(target) + with self.assertRaisesRegex(ValueError, "symlink"): + create_manifest(root, {"pte": Path("alias.pte")}) + + def test_ptd_reordering_is_rejected(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "first.ptd").write_bytes(b"first") + (root / "second.ptd").write_bytes(b"second") + manifest = create_manifest( + root, {}, [Path("first.ptd"), Path("second.ptd")] + ) + manifest["ptd_order"] = ["second.ptd", "first.ptd"] + with self.assertRaisesRegex(ValueError, "order"): + validate_manifest(root, manifest) diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index 6448568a66e..5b647e1ae7c 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -2417,10 +2417,140 @@ void test_resize_hook(const std::string& blob_path) { << "after set(8)+propagate run_count=" << run_count << " last_seen=" << last_seen << " (want 2,8)"; + int tid = -1; + for (int id : graph.input_ids()) { + if (graph.get_value_type(id) == WebGPUGraph::ValueType::Tensor && + graph.tensor_has_dynamic_dims(id) && + graph.cur_dims(id).size() >= 2) { + tid = id; + break; + } + } + ASSERT_GE(tid, 0) << "no dynamic tensor input deserialized"; + const std::vector max_dims = graph.cur_dims(tid); + std::vector small_dims = max_dims; + ASSERT_GT(small_dims[small_dims.size() - 2], 1); + small_dims[small_dims.size() - 2] = 1; + + ASSERT_THROW( + graph.add_post_resize_hook({}, {}, [](WebGPUGraph&) {}), + std::runtime_error); + ASSERT_THROW( + graph.add_post_resize_hook({sid}, {}, [](WebGPUGraph&) {}), + std::runtime_error); + ASSERT_THROW( + graph.add_post_resize_hook({}, {tid}, [](WebGPUGraph&) {}), + std::runtime_error); + + int post_count = 0; + graph.add_post_resize_hook({tid}, {sid}, [&](WebGPUGraph& g) { + ++post_count; + EXPECT_EQ(g.read_symint(sid), last_seen); + }); + + // A tensor-only change reaches the terminal phase exactly once. + graph.resize_input(tid, small_dims); + graph.propagate_resize(); + ASSERT_EQ(post_count, 1); + ASSERT_EQ(graph.cur_dims(tid), small_dims); + + // Simultaneous tensor + SymInt changes still invoke the post hook once. + last_seen = 9; + graph.set_symint(sid, last_seen); + graph.resize_input(tid, max_dims); + graph.propagate_resize(); + ASSERT_EQ(post_count, 2); + ASSERT_EQ(graph.cur_dims(tid), max_dims); + + // A failed post hook restores every trigger. Retrying without a setter must + // execute it again, while a clean third call must remain inert. + bool fail_post_once = true; + int fail_post_count = 0; + graph.add_post_resize_hook({tid}, {sid}, [&](WebGPUGraph&) { + ++fail_post_count; + if (fail_post_once) { + fail_post_once = false; + throw std::runtime_error("injected post-hook failure"); + } + }); + last_seen = 10; + graph.set_symint(sid, last_seen); + ASSERT_THROW(graph.propagate_resize(), std::runtime_error); + ASSERT_EQ(fail_post_count, 1); + graph.propagate_resize(); + ASSERT_EQ(fail_post_count, 2); + graph.propagate_resize(); + ASSERT_EQ(fail_post_count, 2); + + // Terminal callbacks cannot start another propagation wave, including by + // writing the same value (which would otherwise leave no dirty-set trace). + bool attack_symint_once = true; + graph.add_post_resize_hook({}, {sid}, [&](WebGPUGraph& g) { + if (attack_symint_once) { + attack_symint_once = false; + g.set_symint(sid, g.read_symint(sid)); + } + }); + last_seen = 11; + graph.set_symint(sid, last_seen); + ASSERT_THROW(graph.propagate_resize(), std::runtime_error); + graph.propagate_resize(); + + bool attack_tensor_once = true; + graph.add_post_resize_hook({tid}, {}, [&](WebGPUGraph& g) { + if (attack_tensor_once) { + attack_tensor_once = false; + g.set_cur_dims(tid, g.cur_dims(tid)); + } + }); + graph.resize_input(tid, small_dims); + ASSERT_THROW(graph.propagate_resize(), std::runtime_error); + graph.propagate_resize(); + + // A tensor-hook failure also restores its direct trigger for a setter-free + // retry, and the terminal-phase guard is reset on every exception path. + bool fail_tensor_once = true; + int fail_tensor_count = 0; + graph.add_tensor_resize_hook(tid, [&](WebGPUGraph&) { + ++fail_tensor_count; + if (fail_tensor_once) { + fail_tensor_once = false; + throw std::runtime_error("injected tensor-hook failure"); + } + }); + graph.resize_input(tid, max_dims); + ASSERT_THROW(graph.propagate_resize(), std::runtime_error); + ASSERT_EQ(fail_tensor_count, 1); + graph.propagate_resize(); + ASSERT_EQ(fail_tensor_count, 2); + + // A non-converging hook restores both the original and cascading tensor + // triggers. Disarming it permits a retry without another resize_input call. + bool oscillate = true; + graph.add_tensor_resize_hook(tid, [&](WebGPUGraph& g) { + if (!oscillate) { + return; + } + std::vector dims = g.cur_dims(tid); + dims[dims.size() - 2] = dims[dims.size() - 2] == 1 ? 2 : 1; + g.set_cur_dims(tid, dims); + }); + graph.resize_input(tid, small_dims); + ASSERT_THROW(graph.propagate_resize(), std::runtime_error); + oscillate = false; + graph.propagate_resize(); + + // A normal setter still works after both post-phase failures. + last_seen = 12; + graph.set_symint(sid, last_seen); + graph.propagate_resize(); + ASSERT_EQ(graph.read_symint(sid), 12); + printf( - "PASS: resize-hook dirty-gating (SymInt %d: runs only on change, " - "once per change; saw 3 then 8)\n", - sid); + "PASS: resize hooks (SymInt %d, Tensor %d: dirty gating, terminal " + "coherence, retry, and setter guards)\n", + sid, + tid); } // q4gsw embedding_q4gsw on-GPU configs: small + llama1b (env-gated, diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 9990297b9f0..98812c0f144 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -145,6 +145,26 @@ def test_parse_workgroup_not_fooled_by_const(self) -> None: ) self.assertEqual(g.parse_workgroup_size(src), (64, 1, 1)) + def test_parse_workgroup_ignores_commented_attributes(self) -> None: + line = ( + "// prose mentioning @workgroup_size(1) inline\n" + "@compute @workgroup_size(64, 1, 1)\nfn main(){}" + ) + self.assertEqual(g.parse_workgroup_size(line), (64, 1, 1)) + + block = ( + "/* banner\n * @workgroup_size(1)\n */\n" + "@compute @workgroup_size(32, 2, 1)\nfn main(){}" + ) + self.assertEqual(g.parse_workgroup_size(block), (32, 2, 1)) + + def test_parse_workgroup_ignores_commented_constants(self) -> None: + src = ( + "// const WG: u32 = 1u;\nconst WG: u32 = 64u;\n" + "@compute @workgroup_size(WG, 1, 1)\nfn main(){}" + ) + self.assertEqual(g.parse_workgroup_size(src), (64, 1, 1)) + def test_render_header_shape(self) -> None: wgsl = "@compute @workgroup_size(64, 1, 1)\nfn main(){}\n" h = g.render_header(Path("runtime/ops/update_cache/update_cache.wgsl"), wgsl) @@ -220,14 +240,14 @@ def test_generated_output_manifest_digest(self) -> None: digest.update(b"\0") digest.update(output.read_bytes()) digest.update(b"\0") - self.assertEqual(len(outputs), 136) + self.assertEqual(len(outputs), 138) self.assertEqual( digest.hexdigest(), - "0512f8d258952e446ffaedcb653b6a3a720eccf8a6b5327d95fd454a912214a3", + "fee848cd069b4c09d3d2e9a7920331f46d5646b74bec30259542dde8f287e504", ) self.assertEqual( hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "28aaa7a8d3e916df43e407120e91d487d0d51cbc5ca93c56bd822d25d109890e", + "477721998b3cd8f3f0fdd485fa797c71035a20cbc10a8b4bf44893e37fa435b8", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: