diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index a25c5567d3d..4d3fbccc538 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -32,6 +32,7 @@ set(WEBGPU_SRCS runtime/WebGPUDelegateHeader.cpp runtime/WebGPUDevice.cpp runtime/WebGPUQueryPool.cpp + runtime/WebGPUShaderRegistry.cpp runtime/ops/OperatorRegistry.cpp ) @@ -222,5 +223,10 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) webgpu_output_suppression_test test/native/test_output_suppression.cpp ) target_link_libraries(webgpu_output_suppression_test PRIVATE GTest::gtest) + + add_webgpu_native_test( + webgpu_compute_dispatch_test test/native/test_compute_dispatch.cpp + ) + target_link_libraries(webgpu_compute_dispatch_test PRIVATE GTest::gtest) endif() endif() diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index 7573a2f36cc..bedde5abdb7 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -18,6 +18,7 @@ #include #include +#include #include #include @@ -43,6 +44,27 @@ using executorch::runtime::resize_tensor; using executorch::runtime::Result; using executorch::runtime::Span; +Result parse_webgpu_graph_config( + ArrayRef compile_specs) { + WebGPUGraphConfig config; + for (const CompileSpec& spec : compile_specs) { + if (spec.key == nullptr || + std::strcmp(spec.key, "webgpu_record_q4gsw_decode_route") != 0) { + continue; + } + if (spec.value.nbytes != sizeof(uint8_t) || spec.value.buffer == nullptr) { + ET_LOG( + Error, + "WebGPU compile option webgpu_record_q4gsw_decode_route must be " + "exactly one byte"); + return Error::DelegateInvalidCompatibility; + } + config.record_q4gsw_decode_route = + *static_cast(spec.value.buffer) != 0; + } + return config; +} + bool WebGPUBackend::is_available() const { return true; } @@ -51,6 +73,13 @@ Result WebGPUBackend::init( BackendInitContext& context, FreeableBuffer* processed, ArrayRef compile_specs) const { + Result parsed_config = + parse_webgpu_graph_config(compile_specs); + if (!parsed_config.ok()) { + return parsed_config.error(); + } + WebGPUGraphConfig config = parsed_config.get(); + // Allocate graph on the runtime allocator WebGPUGraph* graph = context.get_runtime_allocator()->allocateInstance(); @@ -83,29 +112,23 @@ Result WebGPUBackend::init( // Load-time backend option (BackendOption / LoadBackendOptionsMap), keyed by // the registered backend name; default false. Mirrors the CoreML/XNNPACK // runtime-spec pattern -- no compile flag and no .pte re-export needed. - bool enable_f16_kv_cache = false; { Result spec = context.get_runtime_spec("enable_f16_kv_cache"); if (spec.ok()) { - enable_f16_kv_cache = spec.get(); + config.f16_kv_cache = spec.get(); } } - bool enable_f16_accumulate_gemm = false; { Result spec = context.get_runtime_spec("enable_f16_accumulate_gemm"); if (spec.ok()) { - enable_f16_accumulate_gemm = spec.get(); + config.f16_accumulate_gemm = spec.get(); } } try { graph->build( - flatbuffer_data, - constant_data, - context.get_named_data_map(), - enable_f16_kv_cache, - enable_f16_accumulate_gemm); + flatbuffer_data, constant_data, context.get_named_data_map(), config); } catch (const std::exception& e) { ET_LOG(Error, "WebGPU graph build failed: %s", e.what()); graph->~WebGPUGraph(); diff --git a/backends/webgpu/runtime/WebGPUBackend.h b/backends/webgpu/runtime/WebGPUBackend.h index 9c20a3d53be..59f7e33994c 100644 --- a/backends/webgpu/runtime/WebGPUBackend.h +++ b/backends/webgpu/runtime/WebGPUBackend.h @@ -8,12 +8,17 @@ #pragma once +#include #include namespace executorch { namespace backends { namespace webgpu { +executorch::runtime::Result parse_webgpu_graph_config( + executorch::runtime::ArrayRef + compile_specs); + class WebGPUBackend final : public ::executorch::runtime::BackendInterface { public: ~WebGPUBackend() override = default; diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h index 56544c6d496..7f0e7f1927a 100644 --- a/backends/webgpu/runtime/WebGPUDispatchMath.h +++ b/backends/webgpu/runtime/WebGPUDispatchMath.h @@ -13,6 +13,7 @@ // requires for its device-facing functions). #include +#include #include #include #include @@ -74,6 +75,110 @@ struct DispatchGrid { uint32_t stride_x; }; +struct WgCount { + uint32_t x; + uint32_t y; +}; + +struct DispatchRange { + size_t begin; + size_t end; +}; + +constexpr bool should_record_q4gsw_dual_route( + uint32_t max_m, + bool bicol_eligible, + bool has_dynamic_shapes, + bool record_q4gsw_decode_route) { + return max_m > 1u && bicol_eligible && + (has_dynamic_shapes || record_q4gsw_decode_route); +} + +constexpr bool should_record_sdpa_dual_route( + bool fd_eligible, + bool has_dynamic_sequence) { + return fd_eligible && has_dynamic_sequence; +} + +class DispatchRouteRegistry { + public: + template + size_t register_group( + size_t dispatch_count, + const std::vector& ranges, + IsCompute&& is_compute) { + if (dispatch_count < owners_.size() || ranges.size() < 2) { + throw std::runtime_error("invalid WebGPU dispatch route group"); + } + + std::vector claimed(dispatch_count, false); + for (const auto& range : ranges) { + if (range.begin >= range.end || range.end > dispatch_count) { + throw std::runtime_error("invalid WebGPU dispatch route range"); + } + for (size_t i = range.begin; i < range.end; i++) { + if (!is_compute(i)) { + throw std::runtime_error( + "WebGPU dispatch route contains a copy command"); + } + if (claimed[i] || (i < owners_.size() && owners_[i] != kNoOwner)) { + throw std::runtime_error("overlapping WebGPU dispatch route ranges"); + } + claimed[i] = true; + } + } + + const size_t group = groups_.size(); + owners_.resize(dispatch_count, kNoOwner); + for (size_t i = 0; i < claimed.size(); i++) { + if (claimed[i]) { + owners_[i] = group; + } + } + groups_.push_back(ranges); + return group; + } + + template + void select( + size_t group, + size_t active_route, + const std::vector& active_grids, + SetGrid&& set_grid) const { + if (group >= groups_.size()) { + throw std::runtime_error("invalid WebGPU dispatch route group"); + } + const auto& ranges = groups_[group]; + if (active_route >= ranges.size()) { + throw std::runtime_error("invalid active WebGPU dispatch route"); + } + const auto& active = ranges[active_route]; + if (active_grids.size() != active.end - active.begin) { + throw std::runtime_error("WebGPU dispatch route grid count mismatch"); + } + for (const auto& grid : active_grids) { + if (grid.x == 0 || grid.y == 0) { + throw std::runtime_error( + "active WebGPU dispatch route has a zero grid"); + } + } + + for (const auto& range : ranges) { + for (size_t i = range.begin; i < range.end; i++) { + set_grid(i, {0, 0}); + } + } + for (size_t i = 0; i < active_grids.size(); i++) { + set_grid(active.begin + i, active_grids[i]); + } + } + + private: + static constexpr size_t kNoOwner = static_cast(-1); + std::vector> groups_; + std::vector owners_; +}; + // Given the workgroup count needed (1D) and the device's per-dimension // dispatch-count ceiling, compute a near-square 2D grid rather than // {max_dim, div_up(total, max_dim)} — maxing one dim pads the other with diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.cpp b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp index 67cfb97bbc7..19c0892030f 100644 --- a/backends/webgpu/runtime/WebGPUExecutionOptions.cpp +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp @@ -39,7 +39,12 @@ WebGPUExecutionPlan plan_webgpu_execution( size_t output_count, ExecuteConfig config, const std::vector& suppressible_outputs, - WebGPUGraphExecutionOptions options) { + WebGPUGraphExecutionOptions options, + const std::vector& enabled_dispatches) { + if (!enabled_dispatches.empty() && + enabled_dispatches.size() != dispatch_count) { + throw std::runtime_error("WebGPU: enabled dispatch count mismatch"); + } std::vector suppressed_dispatches(dispatch_count, false); std::vector copy_outputs(output_count, true); std::vector seen_output_ordinals(output_count, false); @@ -78,7 +83,8 @@ WebGPUExecutionPlan plan_webgpu_execution( std::vector indices; indices.reserve(end - begin); for (size_t i = begin; i < end; i++) { - if (!suppressed_dispatches[i]) { + if (!suppressed_dispatches[i] && + (enabled_dispatches.empty() || enabled_dispatches[i])) { indices.push_back(i); } } diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.h b/backends/webgpu/runtime/WebGPUExecutionOptions.h index d59affd541b..304f46839d8 100644 --- a/backends/webgpu/runtime/WebGPUExecutionOptions.h +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.h @@ -51,7 +51,8 @@ WebGPUExecutionPlan plan_webgpu_execution( size_t output_count, ExecuteConfig config, const std::vector& suppressible_outputs, - WebGPUGraphExecutionOptions options); + WebGPUGraphExecutionOptions options, + const std::vector& enabled_dispatches = {}); WebGPUGraphExecutionOptions resolve_webgpu_graph_execution_options( const std::vector& delegate_outputs, diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 1968a9a8dee..bc5869c8ec1 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -7,6 +7,7 @@ */ #include +#include #include #include #include @@ -16,9 +17,11 @@ #include #include +#include #include #include +#include #include #include #include @@ -31,6 +34,132 @@ namespace executorch::backends::webgpu { namespace { +class ScopedBindGroupLayout final { + public: + explicit ScopedBindGroupLayout(WGPUBindGroupLayout handle) + : handle_(handle) {} + ~ScopedBindGroupLayout() { + if (handle_ != nullptr) { + wgpuBindGroupLayoutRelease(handle_); + } + } + ScopedBindGroupLayout(const ScopedBindGroupLayout&) = delete; + ScopedBindGroupLayout& operator=(const ScopedBindGroupLayout&) = delete; + + WGPUBindGroupLayout get() const { + return handle_; + } + + private: + WGPUBindGroupLayout handle_; +}; + +class ScopedBindGroup final { + public: + explicit ScopedBindGroup(WGPUBindGroup handle) : handle_(handle) {} + ~ScopedBindGroup() { + if (handle_ != nullptr) { + wgpuBindGroupRelease(handle_); + } + } + ScopedBindGroup(const ScopedBindGroup&) = delete; + ScopedBindGroup& operator=(const ScopedBindGroup&) = delete; + + WGPUBindGroup get() const { + return handle_; + } + + WGPUBindGroup release() { + WGPUBindGroup handle = handle_; + handle_ = nullptr; + return handle; + } + + private: + WGPUBindGroup handle_; +}; + +class ScopedComputePipeline final { + public: + explicit ScopedComputePipeline(WGPUComputePipeline handle) + : handle_(handle) {} + ~ScopedComputePipeline() { + if (handle_ != nullptr) { + wgpuComputePipelineRelease(handle_); + } + } + ScopedComputePipeline(const ScopedComputePipeline&) = delete; + ScopedComputePipeline& operator=(const ScopedComputePipeline&) = delete; + + WGPUComputePipeline release() { + WGPUComputePipeline handle = handle_; + handle_ = nullptr; + return handle; + } + + private: + WGPUComputePipeline handle_; +}; + +class ScopedComputePipelineRef final { + public: + explicit ScopedComputePipelineRef(WGPUComputePipeline handle) + : handle_(handle) { + wgpuComputePipelineAddRef(handle_); + } + ~ScopedComputePipelineRef() { + if (handle_ != nullptr) { + wgpuComputePipelineRelease(handle_); + } + } + ScopedComputePipelineRef(const ScopedComputePipelineRef&) = delete; + ScopedComputePipelineRef& operator=(const ScopedComputePipelineRef&) = delete; + + WGPUComputePipeline release() { + WGPUComputePipeline handle = handle_; + handle_ = nullptr; + return handle; + } + + private: + WGPUComputePipeline handle_; +}; + +void append_key_component(std::string& key, const std::string& value) { + const uint64_t size = value.size(); + key.append(reinterpret_cast(&size), sizeof(size)); + key.append(value); +} + +std::vector canonical_constants( + const WebGPUComputeDispatchDescriptor& descriptor) { + std::vector constants = descriptor.constants; + std::sort( + constants.begin(), + constants.end(), + [](const auto& left, const auto& right) { + return left.name < right.name; + }); + for (size_t i = 0; i < constants.size(); ++i) { + if (constants[i].name.empty()) { + throw std::runtime_error( + "WebGPU compute dispatch: empty specialization constant name"); + } + if (!std::isfinite(constants[i].value)) { + throw std::runtime_error( + "WebGPU compute dispatch: non-finite specialization constant"); + } + if (i > 0 && constants[i - 1].name == constants[i].name) { + throw std::runtime_error( + "WebGPU compute dispatch: duplicate specialization constant"); + } + if (constants[i].value == 0.0) { + constants[i].value = 0.0; + } + } + return constants; +} + // Op name the AOT exporter emits for a prepacked constant (must match the // serialized schema); compared in the prepack pre-scan below. constexpr const char* kPrepackOpName = "et_vk.prepack.default"; @@ -97,6 +226,52 @@ static_assert(sizeof(QkvFusedParams) == 32, "QkvFusedParams must be 32 bytes"); } // namespace +std::string make_compute_pipeline_key( + const WebGPUComputeDispatchDescriptor& descriptor) { + if (descriptor.shader_name.empty()) { + throw std::runtime_error("WebGPU compute dispatch: empty shader name"); + } + if (descriptor.entry_point.empty()) { + throw std::runtime_error("WebGPU compute dispatch: empty entry point"); + } + + std::string key; + append_key_component(key, descriptor.shader_name); + append_key_component(key, descriptor.entry_point); + for (const auto& constant : canonical_constants(descriptor)) { + append_key_component(key, constant.name); + uint64_t value_bits = 0; + static_assert(sizeof(value_bits) == sizeof(constant.value)); + std::memcpy(&value_bits, &constant.value, sizeof(value_bits)); + key.append(reinterpret_cast(&value_bits), sizeof(value_bits)); + } + return key; +} + +void validate_compute_dispatch_descriptor( + const WebGPUComputeDispatchDescriptor& descriptor) { + (void)make_compute_pipeline_key(descriptor); + if (descriptor.bindings.empty()) { + throw std::runtime_error("WebGPU compute dispatch: no buffer bindings"); + } + for (const auto& binding : descriptor.bindings) { + if (binding.buffer == nullptr) { + throw std::runtime_error("WebGPU compute dispatch: null buffer binding"); + } + if (binding.size == 0) { + throw std::runtime_error("WebGPU compute dispatch: zero-size binding"); + } + if (binding.offset > UINT64_MAX - binding.size) { + throw std::runtime_error( + "WebGPU compute dispatch: binding range overflow"); + } + if (binding.offset + binding.size > wgpuBufferGetSize(binding.buffer)) { + throw std::runtime_error( + "WebGPU compute dispatch: binding range exceeds buffer"); + } + } +} + WebGPUGraph::WebGPUGraph() = default; WGPUBuffer WebGPUGraph::create_scratch_buffer(size_t nbytes) { @@ -154,18 +329,156 @@ void WebGPUGraph::release_scratch(WGPUBuffer buffer) { } WGPUBuffer WebGPUGraph::make_uniform_buffer(const void* data, size_t size) { + if (data == nullptr || size == 0u) { + throw std::runtime_error("WebGPU: invalid uniform buffer data"); + } WGPUBufferDescriptor desc = {}; desc.size = size; desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; desc.mappedAtCreation = true; WGPUBuffer buffer = wgpuDeviceCreateBuffer(device_, &desc); + if (buffer == nullptr) { + throw std::runtime_error("WebGPU: failed to create uniform buffer"); + } void* mapped = wgpuBufferGetMappedRange(buffer, 0, size); + if (mapped == nullptr) { + wgpuBufferRelease(buffer); + throw std::runtime_error("WebGPU: failed to map uniform buffer"); + } std::memcpy(mapped, data, size); wgpuBufferUnmap(buffer); uniform_buffer_bytes_ += size; return buffer; } +size_t WebGPUGraph::add_compute_dispatch( + const WebGPUComputeDispatchDescriptor& descriptor) { + validate_compute_dispatch_descriptor(descriptor); + const WebGPUShaderInfo& shader_info = + get_webgpu_shader_info(descriptor.shader_name); + WGPUShaderModule shader = + get_or_create_shader(descriptor.shader_name, shader_info.source); + + const std::string pipeline_key = make_compute_pipeline_key(descriptor); + WGPUComputePipeline pipeline = nullptr; + auto pipeline_it = pipeline_cache_.find(pipeline_key); + if (pipeline_it != pipeline_cache_.end()) { + pipeline = pipeline_it->second; + } else { + const auto constants = canonical_constants(descriptor); + std::vector entries(constants.size()); + for (size_t i = 0; i < constants.size(); ++i) { + entries[i].key = {constants[i].name.data(), constants[i].name.size()}; + entries[i].value = constants[i].value; + } + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = nullptr; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = { + descriptor.entry_point.data(), descriptor.entry_point.size()}; + pipeline_desc.compute.constantCount = entries.size(); + pipeline_desc.compute.constants = entries.data(); + ScopedComputePipeline created_pipeline( + wgpuDeviceCreateComputePipeline(device_, &pipeline_desc)); + pipeline = created_pipeline.release(); + if (pipeline == nullptr) { + throw std::runtime_error("WebGPU: failed to create compute pipeline"); + } + ScopedComputePipeline pipeline_owner(pipeline); + pipeline_cache_.emplace(pipeline_key, pipeline); + pipeline_owner.release(); + } + + ScopedBindGroupLayout layout( + wgpuComputePipelineGetBindGroupLayout(pipeline, 0)); + if (layout.get() == nullptr) { + throw std::runtime_error("WebGPU: failed to get bind-group layout"); + } + + std::vector entries(descriptor.bindings.size()); + for (size_t i = 0; i < descriptor.bindings.size(); ++i) { + entries[i].binding = i; + entries[i].buffer = descriptor.bindings[i].buffer; + entries[i].offset = descriptor.bindings[i].offset; + entries[i].size = descriptor.bindings[i].size; + } + WGPUBindGroupDescriptor bind_group_desc = {}; + bind_group_desc.layout = layout.get(); + bind_group_desc.entryCount = entries.size(); + bind_group_desc.entries = entries.data(); + ScopedBindGroup bind_group( + wgpuDeviceCreateBindGroup(device_, &bind_group_desc)); + if (bind_group.get() == nullptr) { + throw std::runtime_error("WebGPU: failed to create bind group"); + } + + ScopedComputePipelineRef dispatch_pipeline(pipeline); + const size_t dispatch_index = add_dispatch( + {pipeline, + bind_group.get(), + descriptor.grid.x, + descriptor.kernel_name.empty() ? descriptor.shader_name + : descriptor.kernel_name, + descriptor.grid.y}); + bind_group.release(); + dispatch_pipeline.release(); + return dispatch_index; +} + +size_t WebGPUGraph::add_dynamic_compute_dispatch_impl( + const WebGPUComputeDispatchDescriptor& descriptor, + int trigger_tensor_id, + std::function pick_grid) { + if (trigger_tensor_id < 0 || trigger_tensor_id >= num_values() || + get_value_type(trigger_tensor_id) != ValueType::Tensor) { + throw std::runtime_error( + "WebGPU dynamic dispatch: trigger must be a Tensor"); + } + if (!pick_grid) { + throw std::runtime_error("WebGPU dynamic dispatch: null grid picker"); + } + + const WebGPUDispatchGrid initial_grid = pick_grid(*this); + if (initial_grid.x == 0 || initial_grid.y == 0) { + throw std::runtime_error("WebGPU dynamic dispatch: zero grid"); + } + + WebGPUComputeDispatchDescriptor initial_descriptor = descriptor; + initial_descriptor.grid = initial_grid; + + // Reserve both vectors before creating GPU objects, then stage the sidecar. + // If GPU-object creation fails, removing the sidecar restores the graph; no + // operation that can fail remains after add_compute_dispatch succeeds. + const size_t new_size = dynamic_dispatch_grids_.size() + 1; + dynamic_dispatch_grids_.reserve(new_size); + pending_dynamic_dispatch_grids_.reserve(new_size); + + const size_t expected_index = dispatches_.size(); + dynamic_dispatch_grids_.push_back( + {expected_index, trigger_tensor_id, std::move(pick_grid)}); + try { + add_compute_dispatch(initial_descriptor); + } catch (...) { + dynamic_dispatch_grids_.pop_back(); + throw; + } + return expected_index; +} + +void WebGPUGraph::validate_dynamic_dispatch_route_ranges( + const std::vector& ranges) const { + for (const auto& dynamic_grid : dynamic_dispatch_grids_) { + for (const auto& range : ranges) { + if (range.begin <= dynamic_grid.dispatch_index && + dynamic_grid.dispatch_index < range.end) { + throw std::runtime_error( + "WebGPU dispatch cannot have both dynamic-grid and route ownership"); + } + } + } +} + void WebGPUGraph::update_symints_from_inputs( const std::vector& inputs) { for (const auto& src : symint_sources_) { @@ -210,7 +523,13 @@ void WebGPUGraph::update_symints_from_inputs( // elem_size (buffer-derived) would misread int64 host data as int32. int32_t val; if (inputs[pos].host_is_int64) { - val = static_cast(static_cast(host)[offset]); + const int64_t raw = static_cast(host)[offset]; + if (raw < std::numeric_limits::min() || + raw > std::numeric_limits::max()) { + throw std::runtime_error( + "select_as_symint: selected value is outside int32 range"); + } + val = static_cast(raw); } else { val = static_cast(host)[offset]; } @@ -296,11 +615,41 @@ void WebGPUGraph::propagate_resize() { pass++) { std::unordered_set processing; processing.swap(dirty_tensors_); - for (auto& hook : tensor_resize_hooks_) { - if (processing.count(hook.trigger_tensor_id) != 0) { - hook.fn(*this); + pending_dynamic_dispatch_grids_.clear(); + try { + for (auto& hook : tensor_resize_hooks_) { + if (processing.count(hook.trigger_tensor_id) != 0) { + hook.fn(*this); + } } + + // A hook or picker may fail, so compute and validate every affected grid + // before changing any dispatch. The graph-owned staging vector has + // capacity for every registered dynamic grid and is reused on execute. + for (const auto& dynamic_grid : dynamic_dispatch_grids_) { + if (processing.count(dynamic_grid.trigger_tensor_id) == 0) { + continue; + } + const WebGPUDispatchGrid grid = dynamic_grid.pick_grid(*this); + if (grid.x == 0 || grid.y == 0) { + throw std::runtime_error("WebGPU dynamic dispatch: zero grid"); + } + pending_dynamic_dispatch_grids_.push_back( + {dynamic_grid.dispatch_index, grid}); + } + } catch (...) { + 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; } + pending_dynamic_dispatch_grids_.clear(); } if (!dirty_tensors_.empty()) { throw std::runtime_error( @@ -379,8 +728,7 @@ void WebGPUGraph::build( const void* flatbuffer_data, const uint8_t* constant_data, const executorch::runtime::NamedDataMap* named_data_map, - bool f16_kv_cache, - bool f16_accumulate_gemm) { + WebGPUGraphConfig config) { if (!device_) { auto* ctx = get_default_webgpu_context(); if (ctx) { @@ -403,12 +751,10 @@ void WebGPUGraph::build( // f16 KV cache (runtime opt-in): store K/V caches as f16 iff the opt-in is // set AND the device negotiated shader-f16 (fail-closed). + config_ = config; const WebGPUContext* kv_ctx = get_default_webgpu_context(); - kv_f16_ = f16_kv_cache && (kv_ctx != nullptr && kv_ctx->shader_f16_supported); - - // f16-accumulate q4gsw steel prefill GEMM (runtime opt-in). QuantizedLinear - // additionally gates the kernel on the negotiated shader-f16 feature. - f16_accumulate_gemm_ = f16_accumulate_gemm; + kv_f16_ = config_.f16_kv_cache && + (kv_ctx != nullptr && kv_ctx->shader_f16_supported); // Phase 1: Create all values const auto* values = graph->values(); @@ -438,6 +784,12 @@ void WebGPUGraph::build( if (!a) { continue; } + 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) { + dynamic_tensor_ids_.insert(static_cast(a->Get(0))); + } + } // f16 KV: tag sdpa K/V cache values (args[3],[4]) for half-size alloc. // Inert unless kv_f16_ (runtime opt-in) is set. if (kv_f16_ && a->size() > 4 && @@ -1236,8 +1588,16 @@ WGPUShaderModule WebGPUGraph::get_or_create_shader( WGPUShaderModuleDescriptor shader_desc = {}; shader_desc.nextInChain = &wgsl_desc.chain; WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device_, &shader_desc); + if (shader == nullptr) { + throw std::runtime_error("WebGPU: failed to create shader module"); + } - shader_cache_[key] = shader; + try { + shader_cache_.emplace(key, shader); + } catch (...) { + wgpuShaderModuleRelease(shader); + throw; + } return shader; } @@ -1626,59 +1986,66 @@ void WebGPUGraph::add_qkv_fused_hook(const QkvFusionGroup& g) { const size_t fused_idx = g.fused_dispatch, sep0 = g.sep_dispatch[0], sep1 = g.sep_dispatch[1], sep2 = g.sep_dispatch[2]; WGPUBuffer params_buf = g.fused_params; - add_tensor_resize_hook( - input_id, - [input_id, - out_q_id, - out_k_id, - out_v_id, - K, - Kp, - gs, - Nq, - Nk, - Nv, - Nf, - fused_idx, - sep0, - sep1, - sep2, - params_buf](WebGPUGraph& gr) { - const auto& d = gr.cur_dims(input_id); - uint64_t numel = 1; - for (int64_t v : d) { - numel *= static_cast(v); - } - const uint32_t m = static_cast(numel / K); - std::vector oq = d; - oq.back() = static_cast(Nq); - std::vector ok = d; - ok.back() = static_cast(Nk); - std::vector ov = d; - ov.back() = static_cast(Nv); - gr.set_cur_dims(out_q_id, oq); - gr.set_cur_dims(out_k_id, ok); - gr.set_cur_dims(out_v_id, ov); - QkvFusedParams p = {}; - p.M = m; - p.N = Nf; - p.K = K; - p.K_packed = Kp; - p.group_size = gs; - p.padded_N = Nf; - p.has_bias = 0; - wgpuQueueWriteBuffer(gr.queue(), params_buf, 0, &p, sizeof(p)); - if (m > 1u) { - const uint32_t nbN2 = (Nf + 63u) / 64u; - const uint32_t nbM2 = (m + 63u) / 64u; - gr.dispatch_at(fused_idx).workgroup_count_x = nbN2 * nbM2; - gr.dispatch_at(sep0).workgroup_count_x = 0u; - gr.dispatch_at(sep1).workgroup_count_x = 0u; - gr.dispatch_at(sep2).workgroup_count_x = 0u; - } else { - gr.dispatch_at(fused_idx).workgroup_count_x = 0u; - } - }); + auto update_route = [input_id, + out_q_id, + out_k_id, + out_v_id, + K, + Kp, + gs, + Nq, + Nk, + Nv, + Nf, + fused_idx, + sep0, + sep1, + sep2, + params_buf](WebGPUGraph& gr) { + const auto& d = gr.cur_dims(input_id); + uint64_t numel = 1; + for (int64_t v : d) { + numel *= static_cast(v); + } + const uint32_t m = static_cast(numel / K); + std::vector oq = d; + oq.back() = static_cast(Nq); + std::vector ok = d; + ok.back() = static_cast(Nk); + std::vector ov = d; + ov.back() = static_cast(Nv); + gr.set_cur_dims(out_q_id, oq); + gr.set_cur_dims(out_k_id, ok); + gr.set_cur_dims(out_v_id, ov); + QkvFusedParams p = {}; + p.M = m; + p.N = Nf; + p.K = K; + p.K_packed = Kp; + p.group_size = gs; + p.padded_N = Nf; + p.has_bias = 0; + wgpuQueueWriteBuffer(gr.queue(), params_buf, 0, &p, sizeof(p)); + if (m > 1u) { + const uint32_t nbN2 = (Nf + 63u) / 64u; + const uint32_t nbM2 = (m + 63u) / 64u; + gr.dispatch_at(fused_idx).workgroup_count_x = nbN2 * nbM2; + gr.dispatch_at(fused_idx).workgroup_count_y = 1u; + gr.dispatch_at(sep0).workgroup_count_x = 0u; + gr.dispatch_at(sep0).workgroup_count_y = 0u; + gr.dispatch_at(sep1).workgroup_count_x = 0u; + gr.dispatch_at(sep1).workgroup_count_y = 0u; + gr.dispatch_at(sep2).workgroup_count_x = 0u; + gr.dispatch_at(sep2).workgroup_count_y = 0u; + } else { + gr.dispatch_at(fused_idx).workgroup_count_x = 0u; + gr.dispatch_at(fused_idx).workgroup_count_y = 0u; + } + }; + // Apply the max-shape route immediately. Resize hooks do not run before the + // first execution when cur_dims already equal the serialized max shape. + update_route(*this); + add_tensor_resize_hook(input_id, std::move(update_route)); } void WebGPUGraph::copy_inputs(const std::vector& inputs) { @@ -1743,12 +2110,26 @@ bool should_timestamp_query() { WebGPUExecutionPlan WebGPUGraph::make_execution_plan( const WebGPUGraphExecutionOptions& options) const { + const size_t n = dispatches_.size(); + std::vector enabled_dispatches(n, true); + for (size_t i = 0; i < n; i++) { + if (dispatches_[i].kind != WebGPUDispatch::Kind::Compute) { + continue; + } + const bool zero_x = dispatches_[i].workgroup_count_x == 0; + const bool zero_y = dispatches_[i].workgroup_count_y == 0; + if (zero_x != zero_y) { + throw std::runtime_error("WebGPU: dispatch has a half-zero grid"); + } + enabled_dispatches[i] = !zero_x; + } return plan_webgpu_execution( - dispatches_.size(), + n, output_copies_.size(), execute_config_, suppressible_outputs_, - options); + options, + enabled_dispatches); } size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { @@ -1772,17 +2153,25 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { if (chunk == 0 || n <= chunk) { #ifdef WGPU_BACKEND_ENABLE_PROFILING + size_t active_compute_count = 0; + for (size_t i : plan.dispatch_chunks.front()) { + if (dispatches_[i].kind == WebGPUDispatch::Kind::Compute) { + active_compute_count++; + } + } // Bench: timestamp-query pool, null unless env-gated + feature present. WebGPUQueryPool* qp = nullptr; - if (should_timestamp_query() && n > 0) { + if (should_timestamp_query() && active_compute_count > 0) { if (auto* ctx = get_default_webgpu_context()) { if (ctx->timestamp_supported) { - if (!ctx->querypool || ctx->querypool->capacity() < n) { + if (!ctx->querypool || + ctx->querypool->capacity() < active_compute_count) { ctx->querypool = std::make_unique(); - ctx->querypool->initialize(device_, static_cast(n)); + ctx->querypool->initialize( + device_, static_cast(active_compute_count)); } qp = ctx->querypool.get(); - qp->reset(static_cast(n)); + qp->reset(static_cast(active_compute_count)); } } } @@ -1793,6 +2182,9 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { 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 for (const auto& dispatch_chunk : plan.dispatch_chunks) { for (size_t i : dispatch_chunk) { const auto& dispatch = dispatches_[i]; @@ -1811,7 +2203,7 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { // tw must outlive BeginComputePass (the descriptor points at it). WGPUPassTimestampWrites tw = {}; if (qp) { - tw = qp->writes_for(static_cast(i)); + tw = qp->writes_for(query_index); pass_desc.timestampWrites = &tw; } #endif // WGPU_BACKEND_ENABLE_PROFILING @@ -1827,10 +2219,11 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { #ifdef WGPU_BACKEND_ENABLE_PROFILING if (qp) { qp->record( - static_cast(i), + query_index, dispatch.kernel_name, {dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1}, {1, 1, 1}); + query_index++; } #endif // WGPU_BACKEND_ENABLE_PROFILING } diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index b67f426c3f2..45f21ccdc6a 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -12,11 +12,15 @@ #include #include +#include #include +#include #include #include +#include #include +#include #include #include @@ -60,6 +64,37 @@ struct WebGPUDispatch { size_t copy_nbytes = 0; }; +struct WebGPUBufferBinding { + WGPUBuffer buffer = nullptr; + uint64_t offset = 0; + uint64_t size = 0; +}; + +struct WebGPUSpecializationConstant { + std::string name; + double value = 0.0; +}; + +struct WebGPUDispatchGrid { + uint32_t x = 1; + uint32_t y = 1; +}; + +struct WebGPUComputeDispatchDescriptor { + std::string shader_name; + std::string entry_point = "main"; + std::string kernel_name; + std::vector bindings; + std::vector constants; + WebGPUDispatchGrid grid; +}; + +std::string make_compute_pipeline_key( + const WebGPUComputeDispatchDescriptor& descriptor); + +void validate_compute_dispatch_descriptor( + const WebGPUComputeDispatchDescriptor& descriptor); + struct OutputCopy { WGPUBuffer src_buffer = nullptr; WGPUBuffer staging_buffer = nullptr; @@ -92,6 +127,12 @@ struct WebGPUMemoryStats { } }; +struct WebGPUGraphConfig { + bool f16_kv_cache = false; + bool f16_accumulate_gemm = false; + bool record_q4gsw_decode_route = false; +}; + class WebGPUGraph { public: WebGPUGraph(); @@ -103,8 +144,7 @@ class WebGPUGraph { const void* flatbuffer_data, const uint8_t* constant_data, const executorch::runtime::NamedDataMap* named_data_map = nullptr, - bool f16_kv_cache = false, - bool f16_accumulate_gemm = false); + WebGPUGraphConfig config = {}); // Copy input tensor data from host pointers into GPU buffers. void copy_inputs(const std::vector& inputs); @@ -199,14 +239,43 @@ class WebGPUGraph { symint_dim_sources_.push_back({symint_id, tensor_id, dim}); } + bool tensor_has_dynamic_dims(int tensor_id) const { + return dynamic_tensor_ids_.count(tensor_id) != 0; + } + + bool has_dynamic_shapes() const { + return !dynamic_tensor_ids_.empty(); + } + // Execute-time select_as_symint read; mirrors Vulkan select_as_symint_impl. void update_symints_from_inputs(const std::vector& inputs); // Per-SymInt resize hook; mirrors Vulkan DynamicDispatchNode::trigger_resize. void add_resize_hook(int symint_id, std::function fn) { + if (symint_id < 0 || symint_id >= num_values() || + get_value_type(symint_id) != ValueType::SymInt) { + throw std::runtime_error("WebGPU resize: trigger must be a SymInt"); + } + if (!fn) { + throw std::runtime_error("WebGPU resize: null SymInt resize hook"); + } resize_hooks_.push_back({symint_id, std::move(fn)}); } + template + void add_resize_hook( + int symint_id, + void (*fn)(WebGPUGraph&, const Context&), + Context context) { + if (fn == nullptr) { + throw std::runtime_error("WebGPU resize: null SymInt resize hook"); + } + add_resize_hook( + symint_id, [fn, context = std::move(context)](WebGPUGraph& graph) { + fn(graph, context); + }); + } + // Set a graph input's live dims (<= max) + dirty it; static path stays inert. void resize_input(int value_id, const std::vector& new_dims); // Set a tensor's live dims (an op resize hook calls this for its output to @@ -221,9 +290,31 @@ class WebGPUGraph { void add_tensor_resize_hook( int trigger_tensor_id, std::function fn) { + if (trigger_tensor_id < 0 || trigger_tensor_id >= num_values() || + get_value_type(trigger_tensor_id) != ValueType::Tensor) { + throw std::runtime_error("WebGPU resize: trigger must be a Tensor"); + } + if (!fn) { + throw std::runtime_error("WebGPU resize: null tensor resize hook"); + } tensor_resize_hooks_.push_back({trigger_tensor_id, std::move(fn)}); } + template + void add_tensor_resize_hook( + int trigger_tensor_id, + void (*fn)(WebGPUGraph&, const Context&), + Context context) { + if (fn == nullptr) { + throw std::runtime_error("WebGPU resize: null tensor resize hook"); + } + add_tensor_resize_hook( + trigger_tensor_id, + [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(); @@ -235,6 +326,26 @@ class WebGPUGraph { return dispatches_.size(); } + size_t register_dispatch_route_group( + const std::vector& ranges) { + validate_dynamic_dispatch_route_ranges(ranges); + return dispatch_routes_.register_group( + dispatches_.size(), ranges, [&](size_t i) { + return dispatches_[i].kind == WebGPUDispatch::Kind::Compute; + }); + } + + void select_dispatch_route( + size_t group, + size_t active_route, + const std::vector& active_grids) { + dispatch_routes_.select( + group, active_route, active_grids, [&](size_t i, utils::WgCount grid) { + dispatches_[i].workgroup_count_x = grid.x; + dispatches_[i].workgroup_count_y = grid.y; + }); + } + WGPUDevice device() const { return device_; } @@ -288,6 +399,24 @@ class WebGPUGraph { owned_uniform_buffers_.push_back(buffer); } + template + WGPUBuffer create_params_buffer(const Block& data) { + static_assert( + std::is_trivially_copyable::value, + "WebGPU parameter blocks must be trivially copyable"); + static_assert( + sizeof(Block) % 4u == 0u, + "WebGPU parameter blocks must have a 4-byte-aligned size"); + WGPUBuffer buffer = make_uniform_buffer(&data, sizeof(Block)); + try { + own_uniform_buffer(buffer); + } catch (...) { + wgpuBufferRelease(buffer); + throw; + } + return buffer; + } + // Graph-owned scratch storage buffer for fused-op intermediates (e.g. SDPA). WGPUBuffer create_scratch_buffer(size_t nbytes); @@ -324,6 +453,26 @@ class WebGPUGraph { // in the memory stats. Shared helper for ops needing a uniform Params buffer. WGPUBuffer make_uniform_buffer(const void* data, size_t size); + size_t add_compute_dispatch( + const WebGPUComputeDispatchDescriptor& descriptor); + + template + size_t add_dynamic_compute_dispatch( + const WebGPUComputeDispatchDescriptor& descriptor, + int trigger_tensor_id, + WebGPUDispatchGrid (*pick_grid)(const WebGPUGraph&, const Context&), + Context context) { + if (pick_grid == nullptr) { + throw std::runtime_error("WebGPU dynamic dispatch: null grid picker"); + } + return add_dynamic_compute_dispatch_impl( + descriptor, + trigger_tensor_id, + [pick_grid, context = std::move(context)](const WebGPUGraph& graph) { + return pick_grid(graph, context); + }); + } + WGPUShaderModule get_or_create_shader( const std::string& key, const char* wgsl_source); @@ -376,13 +525,17 @@ class WebGPUGraph { // True when the q4gsw steel prefill GEMM uses the lossy f16-accumulate kernel // (runtime opt-in; perplexity-gated, not bit-exact). bool f16_accumulate_gemm() const { - return f16_accumulate_gemm_; + return config_.f16_accumulate_gemm; + } + + const WebGPUGraphConfig& config() const { + return config_; } private: bool kv_f16_ = false; std::unordered_set kv_cache_ids_; - bool f16_accumulate_gemm_ = false; + WebGPUGraphConfig config_; private: WGPUInstance instance_ = nullptr; @@ -408,6 +561,7 @@ class WebGPUGraph { std::unordered_map symints_; std::vector symint_sources_; std::vector symint_dim_sources_; + std::unordered_set dynamic_tensor_ids_; // Resize hooks + the set of SymInts changed since the last propagate_resize. struct ResizeHook { @@ -426,6 +580,29 @@ class WebGPUGraph { std::vector tensor_resize_hooks_; std::unordered_set dirty_tensors_; + // 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. + struct DynamicDispatchGrid { + size_t dispatch_index; + int trigger_tensor_id; + std::function pick_grid; + }; + std::vector dynamic_dispatch_grids_; + + struct PendingDynamicDispatchGrid { + size_t dispatch_index; + WebGPUDispatchGrid grid; + }; + std::vector pending_dynamic_dispatch_grids_; + + size_t add_dynamic_compute_dispatch_impl( + const WebGPUComputeDispatchDescriptor& descriptor, + int trigger_tensor_id, + std::function pick_grid); + void validate_dynamic_dispatch_route_ranges( + const std::vector& ranges) const; + std::vector input_ids_; std::vector output_ids_; @@ -459,6 +636,7 @@ class WebGPUGraph { std::vector suppressible_outputs_; std::vector dispatches_; + utils::DispatchRouteRegistry dispatch_routes_; // Prepack-routed constant sources (offset/named-key + size); the prepack node // materializes these once. constant_data_/named_data_map_ point at the .pte diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp new file mode 100644 index 00000000000..36722aa2c39 --- /dev/null +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -0,0 +1,1068 @@ +/* + * 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. + */ + +// @generated by scripts/gen_wgsl_headers.py - DO NOT EDIT. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { +namespace { + +constexpr std::array kShaderRegistry = {{ + { + "abs", + kAbsWGSL, + kAbsWorkgroupSizeX, + kAbsWorkgroupSizeY, + kAbsWorkgroupSizeZ, + }, + { + "adamw_step", + kAdamwStepWGSL, + kAdamwStepWorkgroupSizeX, + kAdamwStepWorkgroupSizeY, + kAdamwStepWorkgroupSizeZ, + }, + { + "addmm_tiled", + kAddmmTiledWGSL, + kAddmmTiledWorkgroupSizeX, + kAddmmTiledWorkgroupSizeY, + kAddmmTiledWorkgroupSizeZ, + }, + { + "amax", + kAmaxWGSL, + kAmaxWorkgroupSizeX, + kAmaxWorkgroupSizeY, + kAmaxWorkgroupSizeZ, + }, + { + "amin", + kAminWGSL, + kAminWorkgroupSizeX, + kAminWorkgroupSizeY, + kAminWorkgroupSizeZ, + }, + { + "apply_rotary_emb_interleaved", + kApplyRotaryEmbInterleavedWGSL, + kApplyRotaryEmbInterleavedWorkgroupSizeX, + kApplyRotaryEmbInterleavedWorkgroupSizeY, + kApplyRotaryEmbInterleavedWorkgroupSizeZ, + }, + { + "arg_reduce", + kArgReduceWGSL, + kArgReduceWorkgroupSizeX, + kArgReduceWorkgroupSizeY, + kArgReduceWorkgroupSizeZ, + }, + { + "avg_pool2d", + kAvgPool2dWGSL, + kAvgPool2dWorkgroupSizeX, + kAvgPool2dWorkgroupSizeY, + kAvgPool2dWorkgroupSizeZ, + }, + { + "batch_norm", + kBatchNormWGSL, + kBatchNormWorkgroupSizeX, + kBatchNormWorkgroupSizeY, + kBatchNormWorkgroupSizeZ, + }, + { + "binary_add", + kBinaryAddWGSL, + kBinaryAddWorkgroupSizeX, + kBinaryAddWorkgroupSizeY, + kBinaryAddWorkgroupSizeZ, + }, + { + "binary_div", + kBinaryDivWGSL, + kBinaryDivWorkgroupSizeX, + kBinaryDivWorkgroupSizeY, + kBinaryDivWorkgroupSizeZ, + }, + { + "binary_floor_divide", + kBinaryFloorDivideWGSL, + kBinaryFloorDivideWorkgroupSizeX, + kBinaryFloorDivideWorkgroupSizeY, + kBinaryFloorDivideWorkgroupSizeZ, + }, + { + "binary_minimum", + kBinaryMinimumWGSL, + kBinaryMinimumWorkgroupSizeX, + kBinaryMinimumWorkgroupSizeY, + kBinaryMinimumWorkgroupSizeZ, + }, + { + "binary_mul", + kBinaryMulWGSL, + kBinaryMulWorkgroupSizeX, + kBinaryMulWorkgroupSizeY, + kBinaryMulWorkgroupSizeZ, + }, + { + "binary_pow", + kBinaryPowWGSL, + kBinaryPowWorkgroupSizeX, + kBinaryPowWorkgroupSizeY, + kBinaryPowWorkgroupSizeZ, + }, + { + "binary_sub", + kBinarySubWGSL, + kBinarySubWorkgroupSizeX, + kBinarySubWorkgroupSizeY, + kBinarySubWorkgroupSizeZ, + }, + { + "bitwise_not", + kBitwiseNotWGSL, + kBitwiseNotWorkgroupSizeX, + kBitwiseNotWorkgroupSizeY, + kBitwiseNotWorkgroupSizeZ, + }, + { + "bmm_tiled", + kBmmTiledWGSL, + kBmmTiledWorkgroupSizeX, + kBmmTiledWorkgroupSizeY, + kBmmTiledWorkgroupSizeZ, + }, + { + "bmm_vec4", + kBmmVec4WGSL, + kBmmVec4WorkgroupSizeX, + kBmmVec4WorkgroupSizeY, + kBmmVec4WorkgroupSizeZ, + }, + { + "cat", + kCatWGSL, + kCatWorkgroupSizeX, + kCatWorkgroupSizeY, + kCatWorkgroupSizeZ, + }, + { + "choose_qparams_affine", + kChooseQparamsAffineWGSL, + kChooseQparamsAffineWorkgroupSizeX, + kChooseQparamsAffineWorkgroupSizeY, + kChooseQparamsAffineWorkgroupSizeZ, + }, + { + "clamp", + kClampWGSL, + kClampWorkgroupSizeX, + kClampWorkgroupSizeY, + kClampWorkgroupSizeZ, + }, + { + "compare", + kCompareWGSL, + kCompareWorkgroupSizeX, + kCompareWorkgroupSizeY, + kCompareWorkgroupSizeZ, + }, + { + "compare_eq", + kCompareEqWGSL, + kCompareEqWorkgroupSizeX, + kCompareEqWorkgroupSizeY, + kCompareEqWorkgroupSizeZ, + }, + { + "compare_ge", + kCompareGeWGSL, + kCompareGeWorkgroupSizeX, + kCompareGeWorkgroupSizeY, + kCompareGeWorkgroupSizeZ, + }, + { + "compare_gt", + kCompareGtWGSL, + kCompareGtWorkgroupSizeX, + kCompareGtWorkgroupSizeY, + kCompareGtWorkgroupSizeZ, + }, + { + "compare_le", + kCompareLeWGSL, + kCompareLeWorkgroupSizeX, + kCompareLeWorkgroupSizeY, + kCompareLeWorkgroupSizeZ, + }, + { + "compare_lt", + kCompareLtWGSL, + kCompareLtWorkgroupSizeX, + kCompareLtWorkgroupSizeY, + kCompareLtWorkgroupSizeZ, + }, + { + "compare_ne", + kCompareNeWGSL, + kCompareNeWorkgroupSizeX, + kCompareNeWorkgroupSizeY, + kCompareNeWorkgroupSizeZ, + }, + { + "constant_pad_nd", + kConstantPadNdWGSL, + kConstantPadNdWorkgroupSizeX, + kConstantPadNdWorkgroupSizeY, + kConstantPadNdWorkgroupSizeZ, + }, + { + "conv1d_dw", + kConv1dDwWGSL, + kConv1dDwWorkgroupSizeX, + kConv1dDwWorkgroupSizeY, + kConv1dDwWorkgroupSizeZ, + }, + { + "conv1d_pw", + kConv1dPwWGSL, + kConv1dPwWorkgroupSizeX, + kConv1dPwWorkgroupSizeY, + kConv1dPwWorkgroupSizeZ, + }, + { + "conv2d", + kConv2dWGSL, + kConv2dWorkgroupSizeX, + kConv2dWorkgroupSizeY, + kConv2dWorkgroupSizeZ, + }, + { + "conv2d_gemm", + kConv2dGemmWGSL, + kConv2dGemmWorkgroupSizeX, + kConv2dGemmWorkgroupSizeY, + kConv2dGemmWorkgroupSizeZ, + }, + { + "conv2d_vec4", + kConv2dVec4WGSL, + kConv2dVec4WorkgroupSizeX, + kConv2dVec4WorkgroupSizeY, + kConv2dVec4WorkgroupSizeZ, + }, + { + "conv_transpose2d", + kConvTranspose2dWGSL, + kConvTranspose2dWorkgroupSizeX, + kConvTranspose2dWorkgroupSizeY, + kConvTranspose2dWorkgroupSizeZ, + }, + { + "conv_with_clamp", + kConvWithClampWGSL, + kConvWithClampWorkgroupSizeX, + kConvWithClampWorkgroupSizeY, + kConvWithClampWorkgroupSizeZ, + }, + { + "cos", + kCosWGSL, + kCosWorkgroupSizeX, + kCosWorkgroupSizeY, + kCosWorkgroupSizeZ, + }, + { + "dequantize_per_tensor", + kDequantizePerTensorWGSL, + kDequantizePerTensorWorkgroupSizeX, + kDequantizePerTensorWorkgroupSizeY, + kDequantizePerTensorWorkgroupSizeZ, + }, + { + "embedding", + kEmbeddingWGSL, + kEmbeddingWorkgroupSizeX, + kEmbeddingWorkgroupSizeY, + kEmbeddingWorkgroupSizeZ, + }, + { + "embedding_q4gsw", + kEmbeddingQ4gswWGSL, + kEmbeddingQ4gswWorkgroupSizeX, + kEmbeddingQ4gswWorkgroupSizeY, + kEmbeddingQ4gswWorkgroupSizeZ, + }, + { + "et_vk_sdpa_av", + kEtVkSdpaAvWGSL, + kEtVkSdpaAvWorkgroupSizeX, + kEtVkSdpaAvWorkgroupSizeY, + kEtVkSdpaAvWorkgroupSizeZ, + }, + { + "et_vk_sdpa_qk", + kEtVkSdpaQkWGSL, + kEtVkSdpaQkWorkgroupSizeX, + kEtVkSdpaQkWorkgroupSizeY, + kEtVkSdpaQkWorkgroupSizeZ, + }, + { + "et_vk_sdpa_qk_entry", + kEtVkSdpaQkEntryWGSL, + kEtVkSdpaQkEntryWorkgroupSizeX, + kEtVkSdpaQkEntryWorkgroupSizeY, + kEtVkSdpaQkEntryWorkgroupSizeZ, + }, + { + "exp", + kExpWGSL, + kExpWorkgroupSizeX, + kExpWorkgroupSizeY, + kExpWorkgroupSizeZ, + }, + { + "expand_copy", + kExpandCopyWGSL, + kExpandCopyWorkgroupSizeX, + kExpandCopyWorkgroupSizeY, + kExpandCopyWorkgroupSizeZ, + }, + { + "fill", + kFillWGSL, + kFillWorkgroupSizeX, + kFillWorkgroupSizeY, + kFillWorkgroupSizeZ, + }, + { + "flip", + kFlipWGSL, + kFlipWorkgroupSizeX, + kFlipWorkgroupSizeY, + kFlipWorkgroupSizeZ, + }, + { + "fused_ce", + kFusedCeWGSL, + kFusedCeWorkgroupSizeX, + kFusedCeWorkgroupSizeY, + kFusedCeWorkgroupSizeZ, + }, + { + "gather", + kGatherWGSL, + kGatherWorkgroupSizeX, + kGatherWorkgroupSizeY, + kGatherWorkgroupSizeZ, + }, + { + "gelu", + kGeluWGSL, + kGeluWorkgroupSizeX, + kGeluWorkgroupSizeY, + kGeluWorkgroupSizeZ, + }, + { + "grid_priors", + kGridPriorsWGSL, + kGridPriorsWorkgroupSizeX, + kGridPriorsWorkgroupSizeY, + kGridPriorsWorkgroupSizeZ, + }, + { + "grid_sampler_2d", + kGridSampler2dWGSL, + kGridSampler2dWorkgroupSizeX, + kGridSampler2dWorkgroupSizeY, + kGridSampler2dWorkgroupSizeZ, + }, + { + "group_norm", + kGroupNormWGSL, + kGroupNormWorkgroupSizeX, + kGroupNormWorkgroupSizeY, + kGroupNormWorkgroupSizeZ, + }, + { + "group_norm_reduce", + kGroupNormReduceWGSL, + kGroupNormReduceWorkgroupSizeX, + kGroupNormReduceWorkgroupSizeY, + kGroupNormReduceWorkgroupSizeZ, + }, + { + "hardswish", + kHardswishWGSL, + kHardswishWorkgroupSizeX, + kHardswishWorkgroupSizeY, + kHardswishWorkgroupSizeZ, + }, + { + "index", + kIndexWGSL, + kIndexWorkgroupSizeX, + kIndexWorkgroupSizeY, + kIndexWorkgroupSizeZ, + }, + { + "index_select", + kIndexSelectWGSL, + kIndexSelectWorkgroupSizeX, + kIndexSelectWorkgroupSizeY, + kIndexSelectWorkgroupSizeZ, + }, + { + "leaky_relu", + kLeakyReluWGSL, + kLeakyReluWorkgroupSizeX, + kLeakyReluWorkgroupSizeY, + kLeakyReluWorkgroupSizeZ, + }, + { + "linear_dW", + kLinearDwWGSL, + kLinearDwWorkgroupSizeX, + kLinearDwWorkgroupSizeY, + kLinearDwWorkgroupSizeZ, + }, + { + "linear_dq8ca_q4gsw", + kLinearDq8caQ4gswWGSL, + kLinearDq8caQ4gswWorkgroupSizeX, + kLinearDq8caQ4gswWorkgroupSizeY, + kLinearDq8caQ4gswWorkgroupSizeZ, + }, + { + "linear_q8ta_q8csw", + kLinearQ8taQ8cswWGSL, + kLinearQ8taQ8cswWorkgroupSizeX, + kLinearQ8taQ8cswWorkgroupSizeY, + kLinearQ8taQ8cswWorkgroupSizeZ, + }, + { + "linear_tiled", + kLinearTiledWGSL, + kLinearTiledWorkgroupSizeX, + kLinearTiledWorkgroupSizeY, + kLinearTiledWorkgroupSizeZ, + }, + { + "linear_vec4", + kLinearVec4WGSL, + kLinearVec4WorkgroupSizeX, + kLinearVec4WorkgroupSizeY, + kLinearVec4WorkgroupSizeZ, + }, + { + "log_softmax", + kLogSoftmaxWGSL, + kLogSoftmaxWorkgroupSizeX, + kLogSoftmaxWorkgroupSizeY, + kLogSoftmaxWorkgroupSizeZ, + }, + { + "logical_and", + kLogicalAndWGSL, + kLogicalAndWorkgroupSizeX, + kLogicalAndWorkgroupSizeY, + kLogicalAndWorkgroupSizeZ, + }, + { + "logical_not", + kLogicalNotWGSL, + kLogicalNotWorkgroupSizeX, + kLogicalNotWorkgroupSizeY, + kLogicalNotWorkgroupSizeZ, + }, + { + "logical_or", + kLogicalOrWGSL, + kLogicalOrWorkgroupSizeX, + kLogicalOrWorkgroupSizeY, + kLogicalOrWorkgroupSizeZ, + }, + { + "max_pool2d", + kMaxPool2dWGSL, + kMaxPool2dWorkgroupSizeX, + kMaxPool2dWorkgroupSizeY, + kMaxPool2dWorkgroupSizeZ, + }, + { + "mm", + kMmWGSL, + kMmWorkgroupSizeX, + kMmWorkgroupSizeY, + kMmWorkgroupSizeZ, + }, + { + "mm_tiled", + kMmTiledWGSL, + kMmTiledWorkgroupSizeX, + kMmTiledWorkgroupSizeY, + kMmTiledWorkgroupSizeZ, + }, + { + "mm_vec4", + kMmVec4WGSL, + kMmVec4WorkgroupSizeX, + kMmVec4WorkgroupSizeY, + kMmVec4WorkgroupSizeZ, + }, + { + "native_layer_norm", + kNativeLayerNormWGSL, + kNativeLayerNormWorkgroupSizeX, + kNativeLayerNormWorkgroupSizeY, + kNativeLayerNormWorkgroupSizeZ, + }, + { + "neg", + kNegWGSL, + kNegWorkgroupSizeX, + kNegWorkgroupSizeY, + kNegWorkgroupSizeZ, + }, + { + "permute", + kPermuteWGSL, + kPermuteWorkgroupSizeX, + kPermuteWorkgroupSizeY, + kPermuteWorkgroupSizeZ, + }, + { + "pixel_shuffle", + kPixelShuffleWGSL, + kPixelShuffleWorkgroupSizeX, + kPixelShuffleWorkgroupSizeY, + kPixelShuffleWorkgroupSizeZ, + }, + { + "pow_scalar", + kPowScalarWGSL, + kPowScalarWorkgroupSizeX, + kPowScalarWorkgroupSizeY, + kPowScalarWorkgroupSizeZ, + }, + { + "q4gsw_backward", + kQ4gswBackwardWGSL, + kQ4gswBackwardWorkgroupSizeX, + kQ4gswBackwardWorkgroupSizeY, + kQ4gswBackwardWorkgroupSizeZ, + }, + { + "q4gsw_linear", + kQ4gswLinearWGSL, + kQ4gswLinearWorkgroupSizeX, + kQ4gswLinearWorkgroupSizeY, + kQ4gswLinearWorkgroupSizeZ, + }, + { + "q4gsw_linear_coop4_bicol", + kQ4gswLinearCoop4BicolWGSL, + kQ4gswLinearCoop4BicolWorkgroupSizeX, + kQ4gswLinearCoop4BicolWorkgroupSizeY, + kQ4gswLinearCoop4BicolWorkgroupSizeZ, + }, + { + "q4gsw_linear_gemm_qkv_fused", + kQ4gswLinearGemmQkvFusedWGSL, + kQ4gswLinearGemmQkvFusedWorkgroupSizeX, + kQ4gswLinearGemmQkvFusedWorkgroupSizeY, + kQ4gswLinearGemmQkvFusedWorkgroupSizeZ, + }, + { + "q4gsw_linear_gemm_shmem", + kQ4gswLinearGemmShmemWGSL, + kQ4gswLinearGemmShmemWorkgroupSizeX, + kQ4gswLinearGemmShmemWorkgroupSizeY, + kQ4gswLinearGemmShmemWorkgroupSizeZ, + }, + { + "q4gsw_linear_gemm_steel", + kQ4gswLinearGemmSteelWGSL, + kQ4gswLinearGemmSteelWorkgroupSizeX, + kQ4gswLinearGemmSteelWorkgroupSizeY, + kQ4gswLinearGemmSteelWorkgroupSizeZ, + }, + { + "q4gsw_linear_gemm_steel_half", + kQ4gswLinearGemmSteelHalfWGSL, + kQ4gswLinearGemmSteelHalfWorkgroupSizeX, + kQ4gswLinearGemmSteelHalfWorkgroupSizeY, + kQ4gswLinearGemmSteelHalfWorkgroupSizeZ, + }, + { + "q4gsw_linear_gemm_steel_half_pwdq", + kQ4gswLinearGemmSteelHalfPwdqWGSL, + kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeX, + kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeY, + kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeZ, + }, + { + "q4gsw_linear_gemm_steel_half_pwdq_f16acc", + kQ4gswLinearGemmSteelHalfPwdqF16accWGSL, + kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeX, + kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeY, + kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeZ, + }, + { + "q4gsw_requant", + kQ4gswRequantWGSL, + kQ4gswRequantWorkgroupSizeX, + kQ4gswRequantWorkgroupSizeY, + kQ4gswRequantWorkgroupSizeZ, + }, + { + "q8ta_add", + kQ8taAddWGSL, + kQ8taAddWorkgroupSizeX, + kQ8taAddWorkgroupSizeY, + kQ8taAddWorkgroupSizeZ, + }, + { + "q8ta_conv2d", + kQ8taConv2dWGSL, + kQ8taConv2dWorkgroupSizeX, + kQ8taConv2dWorkgroupSizeY, + kQ8taConv2dWorkgroupSizeZ, + }, + { + "q8ta_conv2d_dw", + kQ8taConv2dDwWGSL, + kQ8taConv2dDwWorkgroupSizeX, + kQ8taConv2dDwWorkgroupSizeY, + kQ8taConv2dDwWorkgroupSizeZ, + }, + { + "q8ta_conv2d_pw", + kQ8taConv2dPwWGSL, + kQ8taConv2dPwWorkgroupSizeX, + kQ8taConv2dPwWorkgroupSizeY, + kQ8taConv2dPwWorkgroupSizeZ, + }, + { + "q8ta_conv2d_transposed", + kQ8taConv2dTransposedWGSL, + kQ8taConv2dTransposedWorkgroupSizeX, + kQ8taConv2dTransposedWorkgroupSizeY, + kQ8taConv2dTransposedWorkgroupSizeZ, + }, + { + "q8ta_linear", + kQ8taLinearWGSL, + kQ8taLinearWorkgroupSizeX, + kQ8taLinearWorkgroupSizeY, + kQ8taLinearWorkgroupSizeZ, + }, + { + "q8ta_pixel_shuffle", + kQ8taPixelShuffleWGSL, + kQ8taPixelShuffleWorkgroupSizeX, + kQ8taPixelShuffleWorkgroupSizeY, + kQ8taPixelShuffleWorkgroupSizeZ, + }, + { + "q8ta_relu", + kQ8taReluWGSL, + kQ8taReluWorkgroupSizeX, + kQ8taReluWorkgroupSizeY, + kQ8taReluWorkgroupSizeZ, + }, + { + "qcs4w_linear", + kQcs4wLinearWGSL, + kQcs4wLinearWorkgroupSizeX, + kQcs4wLinearWorkgroupSizeY, + kQcs4wLinearWorkgroupSizeZ, + }, + { + "quantize_per_tensor", + kQuantizePerTensorWGSL, + kQuantizePerTensorWorkgroupSizeX, + kQuantizePerTensorWorkgroupSizeY, + kQuantizePerTensorWorkgroupSizeZ, + }, + { + "reduce", + kReduceWGSL, + kReduceWorkgroupSizeX, + kReduceWorkgroupSizeY, + kReduceWorkgroupSizeZ, + }, + { + "relu", + kReluWGSL, + kReluWorkgroupSizeX, + kReluWorkgroupSizeY, + kReluWorkgroupSizeZ, + }, + { + "repeat", + kRepeatWGSL, + kRepeatWorkgroupSizeX, + kRepeatWorkgroupSizeY, + kRepeatWorkgroupSizeZ, + }, + { + "rms_norm", + kRmsNormWGSL, + kRmsNormWorkgroupSizeX, + kRmsNormWorkgroupSizeY, + kRmsNormWorkgroupSizeZ, + }, + { + "rms_norm_vec4", + kRmsNormVec4WGSL, + kRmsNormVec4WorkgroupSizeX, + kRmsNormVec4WorkgroupSizeY, + kRmsNormVec4WorkgroupSizeZ, + }, + { + "rotary_embedding", + kRotaryEmbeddingWGSL, + kRotaryEmbeddingWorkgroupSizeX, + kRotaryEmbeddingWorkgroupSizeY, + kRotaryEmbeddingWorkgroupSizeZ, + }, + { + "rotary_embedding_hf", + kRotaryEmbeddingHfWGSL, + kRotaryEmbeddingHfWorkgroupSizeX, + kRotaryEmbeddingHfWorkgroupSizeY, + kRotaryEmbeddingHfWorkgroupSizeZ, + }, + { + "round", + kRoundWGSL, + kRoundWorkgroupSizeX, + kRoundWorkgroupSizeY, + kRoundWorkgroupSizeZ, + }, + { + "rsqrt", + kRsqrtWGSL, + kRsqrtWorkgroupSizeX, + kRsqrtWorkgroupSizeY, + kRsqrtWorkgroupSizeZ, + }, + { + "sdpa_compute_attn_weights", + kSdpaComputeAttnWeightsWGSL, + kSdpaComputeAttnWeightsWorkgroupSizeX, + kSdpaComputeAttnWeightsWorkgroupSizeY, + kSdpaComputeAttnWeightsWorkgroupSizeZ, + }, + { + "sdpa_compute_attn_weights_half", + kSdpaComputeAttnWeightsHalfWGSL, + kSdpaComputeAttnWeightsHalfWorkgroupSizeX, + kSdpaComputeAttnWeightsHalfWorkgroupSizeY, + kSdpaComputeAttnWeightsHalfWorkgroupSizeZ, + }, + { + "sdpa_compute_out", + kSdpaComputeOutWGSL, + kSdpaComputeOutWorkgroupSizeX, + kSdpaComputeOutWorkgroupSizeY, + kSdpaComputeOutWorkgroupSizeZ, + }, + { + "sdpa_compute_out_half", + kSdpaComputeOutHalfWGSL, + kSdpaComputeOutHalfWorkgroupSizeX, + kSdpaComputeOutHalfWorkgroupSizeY, + kSdpaComputeOutHalfWorkgroupSizeZ, + }, + { + "sdpa_fd_reduce", + kSdpaFdReduceWGSL, + kSdpaFdReduceWorkgroupSizeX, + kSdpaFdReduceWorkgroupSizeY, + kSdpaFdReduceWorkgroupSizeZ, + }, + { + "sdpa_fd_split", + kSdpaFdSplitWGSL, + kSdpaFdSplitWorkgroupSizeX, + kSdpaFdSplitWorkgroupSizeY, + kSdpaFdSplitWorkgroupSizeZ, + }, + { + "sdpa_fd_split_half", + kSdpaFdSplitHalfWGSL, + kSdpaFdSplitHalfWorkgroupSizeX, + kSdpaFdSplitHalfWorkgroupSizeY, + kSdpaFdSplitHalfWorkgroupSizeZ, + }, + { + "sdpa_softmax", + kSdpaSoftmaxWGSL, + kSdpaSoftmaxWorkgroupSizeX, + kSdpaSoftmaxWorkgroupSizeY, + kSdpaSoftmaxWorkgroupSizeZ, + }, + { + "select", + kSelectWGSL, + kSelectWorkgroupSizeX, + kSelectWorkgroupSizeY, + kSelectWorkgroupSizeZ, + }, + { + "sigmoid", + kSigmoidWGSL, + kSigmoidWorkgroupSizeX, + kSigmoidWorkgroupSizeY, + kSigmoidWorkgroupSizeZ, + }, + { + "silu_mul_fused", + kSiluMulFusedWGSL, + kSiluMulFusedWorkgroupSizeX, + kSiluMulFusedWorkgroupSizeY, + kSiluMulFusedWorkgroupSizeZ, + }, + { + "sin", + kSinWGSL, + kSinWorkgroupSizeX, + kSinWorkgroupSizeY, + kSinWorkgroupSizeZ, + }, + { + "slice", + kSliceWGSL, + kSliceWorkgroupSizeX, + kSliceWorkgroupSizeY, + kSliceWorkgroupSizeZ, + }, + { + "softmax", + kSoftmaxWGSL, + kSoftmaxWorkgroupSizeX, + kSoftmaxWorkgroupSizeY, + kSoftmaxWorkgroupSizeZ, + }, + { + "sqrt", + kSqrtWGSL, + kSqrtWorkgroupSizeX, + kSqrtWorkgroupSizeY, + kSqrtWorkgroupSizeZ, + }, + { + "tanh", + kTanhWGSL, + kTanhWorkgroupSizeX, + kTanhWorkgroupSizeY, + kTanhWorkgroupSizeZ, + }, + { + "to_copy_float_to_int", + kToCopyFloatToIntWGSL, + kToCopyFloatToIntWorkgroupSizeX, + kToCopyFloatToIntWorkgroupSizeY, + kToCopyFloatToIntWorkgroupSizeZ, + }, + { + "to_copy_int_to_float", + kToCopyIntToFloatWGSL, + kToCopyIntToFloatWorkgroupSizeX, + kToCopyIntToFloatWorkgroupSizeY, + kToCopyIntToFloatWorkgroupSizeZ, + }, + { + "update_cache", + kUpdateCacheWGSL, + kUpdateCacheWorkgroupSizeX, + kUpdateCacheWorkgroupSizeY, + kUpdateCacheWorkgroupSizeZ, + }, + { + "update_cache_half", + kUpdateCacheHalfWGSL, + kUpdateCacheHalfWorkgroupSizeX, + kUpdateCacheHalfWorkgroupSizeY, + kUpdateCacheHalfWorkgroupSizeZ, + }, + { + "upsample_bilinear2d", + kUpsampleBilinear2dWGSL, + kUpsampleBilinear2dWorkgroupSizeX, + kUpsampleBilinear2dWorkgroupSizeY, + kUpsampleBilinear2dWorkgroupSizeZ, + }, + { + "upsample_nearest2d", + kUpsampleNearest2dWGSL, + kUpsampleNearest2dWorkgroupSizeX, + kUpsampleNearest2dWorkgroupSizeY, + kUpsampleNearest2dWorkgroupSizeZ, + }, + { + "where", + kWhereWGSL, + kWhereWorkgroupSizeX, + kWhereWorkgroupSizeY, + kWhereWorkgroupSizeZ, + }, +}}; + +} // namespace + +const WebGPUShaderInfo& get_webgpu_shader_info(std::string_view name) { + for (const auto& shader : kShaderRegistry) { + if (shader.name == name) { + return shader; + } + } + throw std::runtime_error( + "WebGPU shader registry: unknown shader '" + std::string(name) + "'"); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.h b/backends/webgpu/runtime/WebGPUShaderRegistry.h new file mode 100644 index 00000000000..64b86e36110 --- /dev/null +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.h @@ -0,0 +1,26 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +struct WebGPUShaderInfo { + std::string_view name; + const char* source; + uint32_t workgroup_size_x; + uint32_t workgroup_size_y; + uint32_t workgroup_size_z; +}; + +const WebGPUShaderInfo& get_webgpu_shader_info(std::string_view name); + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index 06c7f312dcd..4f801d817e0 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -50,11 +51,6 @@ inline uint32_t clamp_workgroup_size_pow2(WGPUDevice device, uint32_t desired) { return p; } -struct WgCount { - uint32_t x; - uint32_t y; -}; - // Device's max workgroups per dispatch dimension; the WebGPU spec-default floor // (65535) if the query fails — never under-reports a real device's capacity. inline uint32_t queried_max_workgroups(WGPUDevice device) { @@ -514,6 +510,53 @@ inline ComputePipelineBundle make_compute_pipeline( return bundle; } +// Builds a pipeline for a different shader that uses the exact layout and +// bind group of an earlier bundle. Multi-route ops use this when alternate +// shaders have an identical binding contract: only the shader module and +// pipeline are new; layout and bind-group construction stay single-copy. +inline ComputePipelineBundle make_compute_pipeline( + WGPUDevice device, + const char* wgsl_source, + const ComputePipelineBundle& shared_resources, + const WGPUConstantEntry* constants = nullptr, + size_t constant_count = 0, + const char* entry_point = "main") { + if (shared_resources.bind_group_layout == nullptr || + shared_resources.pipeline_layout == nullptr || + shared_resources.bind_group == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: shared resources are not available"); + } + + ComputePipelineBundle bundle; + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {wgsl_source, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + bundle.shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + if (bundle.shader == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: shader module creation failed"); + } + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = shared_resources.pipeline_layout; + pipeline_desc.compute.module = bundle.shader; + pipeline_desc.compute.entryPoint = {entry_point, WGPU_STRLEN}; + pipeline_desc.compute.constantCount = constant_count; + pipeline_desc.compute.constants = constants; + bundle.pipeline = wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + if (bundle.pipeline == nullptr) { + throw std::runtime_error( + "make_compute_pipeline: compute pipeline creation failed"); + } + + wgpuBindGroupAddRef(shared_resources.bind_group); + bundle.bind_group = shared_resources.bind_group; + return bundle; +} + // The {wg_size, stride_x} override-constant pair every 2D-spill dispatch // builds from its DispatchGrid; was hand-rolled identically at 7 call sites. inline std::array make_grid_constants( diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 5280430df45..16d78a6f923 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -21,7 +21,6 @@ #include #include -#include #include #include #include @@ -162,6 +161,75 @@ uint32_t compute_q4gsw_workgroup_count( device, static_cast(total_tiles), wg_size, op_name); } +struct Q4gswExecutionState { + Q4gswParams params; + std::vector output_dims; + size_t active_route; + utils::WgCount active_grid; +}; + +constexpr size_t kQ4gswBicolRoute = 0; +constexpr size_t kQ4gswPrefillRoute = 1; + +Q4gswExecutionState make_q4gsw_execution_state( + WGPUDevice device, + const std::vector& input_dims, + uint32_t max_m, + uint32_t K, + uint32_t N, + uint32_t K_packed, + uint32_t gs, + uint32_t padded_N, + uint32_t has_bias, + uint32_t wg_size, + bool use_single_gemv, + bool use_dual_route, + bool prefill_use_steel, + bool prefill_use_shmem_gemm) { + if (input_dims.empty()) { + throw std::runtime_error("WebGPU linear_q4gsw(resize): empty input dims"); + } + const uint64_t numel = utils::numel_of(input_dims); + if (numel % static_cast(K) != 0u) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): live input numel not a multiple of K"); + } + const uint64_t live_m = numel / static_cast(K); + if (live_m == 0u) { + throw std::runtime_error("WebGPU linear_q4gsw(resize): live M == 0"); + } + if (live_m > max_m) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): live M exceeds the build-time max"); + } + const uint32_t m = static_cast(live_m); + const bool use_gemv = use_single_gemv || (use_dual_route && m == 1u); + const uint32_t workgroup_count = compute_q4gsw_workgroup_count( + device, + use_gemv, + !use_gemv && prefill_use_steel, + !use_gemv && prefill_use_shmem_gemm, + m, + N, + wg_size, + "linear_q4gsw(resize)"); + + Q4gswExecutionState state = {}; + state.params.M = m; + state.params.N = N; + state.params.K = K; + state.params.K_packed = K_packed; + state.params.group_size = gs; + state.params.padded_N = padded_N; + state.params.has_bias = has_bias; + state.output_dims = input_dims; + state.output_dims.back() = static_cast(N); + state.active_route = + use_dual_route ? (use_gemv ? kQ4gswBicolRoute : kQ4gswPrefillRoute) : 0u; + state.active_grid = {workgroup_count, 1u}; + return state; +} + // et_vk.linear_q4gsw args: [in, weight, scales, group_size, bias, out]. void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { const int in_id = args.at(0); @@ -243,9 +311,15 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { // M==1 -> bicol GEMV; M>1 -> steel GEMM (preferred) else shmem else tiled. const uint32_t wg_size = utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX); - const bool use_gemv = (M == 1u && K % 8u == 0u && gs % 8u == 0u); + const bool bicol_eligible = K % 8u == 0u && gs % 8u == 0u; + const bool use_gemv = M == 1u && bicol_eligible; + const bool use_dual_route = utils::should_record_q4gsw_dual_route( + M, + bicol_eligible, + graph.has_dynamic_shapes(), + graph.config().record_q4gsw_decode_route); // GEMV (bicol) is a pow2 tree reduction; compute its size only when used. - const uint32_t gemv_wg_size = use_gemv + const uint32_t gemv_wg_size = (use_gemv || use_dual_route) ? utils::clamp_workgroup_size_pow2( device, kQ4gswLinearCoop4BicolWorkgroupSizeX) : 0u; @@ -258,6 +332,9 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { // large K/N thresholds; otherwise the register-tiled path handles it. const bool use_shmem_gemm = !use_gemv && !use_steel && (K >= kQ4gswShmemMinDim || N >= kQ4gswShmemNMinDim); + const char* prefill_shader_src = use_steel ? kQ4gswLinearGemmSteelWGSL + : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL + : kQ4gswLinearWGSL; const char* shader_src = use_gemv ? kQ4gswLinearCoop4BicolWGSL : use_steel ? kQ4gswLinearGemmSteelWGSL : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL @@ -271,9 +348,10 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { // each u32 weight word once + hoists the per-column scale (half re-reads // them ~8x/~16x). Needs group_size % BK == 0 so the hoisted scale is // constant across the BK tile; else the per-nibble `half` kernel. - shader_src = (gs % kQ4gswSteelBK == 0u) + prefill_shader_src = (gs % kQ4gswSteelBK == 0u) ? kQ4gswLinearGemmSteelHalfPwdqWGSL : kQ4gswLinearGemmSteelHalfWGSL; + shader_src = prefill_shader_src; } } // f16-accumulate: pwdq staging with an f16 register accumulator. @@ -284,18 +362,10 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { if (use_steel && graph.f16_accumulate_gemm() && (gs % kQ4gswSteelBK == 0u)) { const WebGPUContext* ctx = get_default_webgpu_context(); if (ctx != nullptr && ctx->shader_f16_supported) { - shader_src = kQ4gswLinearGemmSteelHalfPwdqF16accWGSL; + prefill_shader_src = kQ4gswLinearGemmSteelHalfPwdqF16accWGSL; + shader_src = prefill_shader_src; } } - const uint32_t workgroup_count = compute_q4gsw_workgroup_count( - device, - use_gemv, - use_steel, - use_shmem_gemm, - M, - N, - wg_size, - "linear_q4gsw"); // Optional bias: real buffer if present, else a dummy for the fixed layout. uint32_t has_bias = 0; @@ -315,62 +385,96 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { bias_buffer = graph.create_scratch_buffer(4); } - Q4gswParams params = {}; - params.M = M; - params.N = N; - params.K = K; - params.K_packed = K_packed; - params.group_size = gs; - params.padded_N = padded_N; - params.has_bias = has_bias; - - WGPUBufferDescriptor uniform_desc = {}; - uniform_desc.size = sizeof(Q4gswParams); - uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; - uniform_desc.mappedAtCreation = true; - WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc); - void* mapped = - wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(Q4gswParams)); - std::memcpy(mapped, ¶ms, sizeof(Q4gswParams)); - wgpuBufferUnmap(uniform_buffer); - graph.add_uniform_buffer_bytes(sizeof(Q4gswParams)); - - // GEMV/tiled wire an override wg_size; steel (256) + shmem (64) are fixed. - const bool fixed_wg = use_steel || use_shmem_gemm; - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = - static_cast(use_gemv ? gemv_wg_size : wg_size); - - utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + const Q4gswExecutionState initial_state = make_q4gsw_execution_state( device, - shader_src, - { - {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, - {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, - {2, - WGPUBufferBindingType_ReadOnlyStorage, - weight.buffer, - weight.nbytes}, - {3, - WGPUBufferBindingType_ReadOnlyStorage, - scales.buffer, - scales.nbytes}, - {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buffer, bias_size}, - {5, - WGPUBufferBindingType_Uniform, - uniform_buffer, - sizeof(Q4gswParams)}, - }, - fixed_wg ? nullptr : &wg_size_constant, - fixed_wg ? 0u : 1u); - - const size_t dispatch_idx = graph.add_dispatch( - {bundle.pipeline, bundle.bind_group, workgroup_count, "linear_q4gsw"}); - - // Dynamic shapes: recompute dispatch + params.M for the live M. use_gemv and - // use_shmem_gemm are captured (routing is fixed at build); the helper re-runs - // the same path's workgroup-count formula with the live m. + in.dims, + M, + K, + N, + K_packed, + gs, + padded_N, + has_bias, + wg_size, + use_gemv, + use_dual_route, + use_steel, + use_shmem_gemm); + + WGPUBuffer params_buffer = graph.create_params_buffer(initial_state.params); + + const std::vector bindings = { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {2, WGPUBufferBindingType_ReadOnlyStorage, weight.buffer, weight.nbytes}, + {3, WGPUBufferBindingType_ReadOnlyStorage, scales.buffer, scales.nbytes}, + {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buffer, bias_size}, + {5, WGPUBufferBindingType_Uniform, params_buffer, sizeof(Q4gswParams)}, + }; + auto make_bundle = + [&](const char* source, bool fixed_wg, uint32_t override_wg_size) { + const WGPUConstantEntry wg_size_constant = + utils::make_wg_size_constant(override_wg_size); + return utils::make_compute_pipeline( + device, + source, + bindings, + fixed_wg ? nullptr : &wg_size_constant, + fixed_wg ? 0u : 1u); + }; + auto make_shared_bundle = + [&](const char* source, + const utils::ComputePipelineBundle& shared_resources, + bool fixed_wg, + uint32_t override_wg_size) { + const WGPUConstantEntry wg_size_constant = + utils::make_wg_size_constant(override_wg_size); + return utils::make_compute_pipeline( + device, + source, + shared_resources, + fixed_wg ? nullptr : &wg_size_constant, + fixed_wg ? 0u : 1u); + }; + + const bool fixed_prefill_wg = use_steel || use_shmem_gemm; + const char* prefill_label = use_steel ? "linear_q4gsw_steel" + : use_shmem_gemm ? "linear_q4gsw_shmem" + : "linear_q4gsw_tiled"; + size_t dispatch_idx = 0; + size_t route_group = 0; + if (use_dual_route) { + utils::ComputePipelineBundle bicol_bundle = + make_bundle(kQ4gswLinearCoop4BicolWGSL, false, gemv_wg_size); + const size_t bicol_idx = graph.add_dispatch( + {bicol_bundle.pipeline, + bicol_bundle.bind_group, + initial_state.active_grid.x, + "linear_q4gsw_coop4_bicol"}); + utils::ComputePipelineBundle prefill_bundle = make_shared_bundle( + prefill_shader_src, bicol_bundle, fixed_prefill_wg, wg_size); + const size_t prefill_idx = graph.add_dispatch( + {prefill_bundle.pipeline, + prefill_bundle.bind_group, + initial_state.active_grid.x, + prefill_label}); + route_group = graph.register_dispatch_route_group( + {{bicol_idx, bicol_idx + 1}, {prefill_idx, prefill_idx + 1}}); + graph.select_dispatch_route( + route_group, initial_state.active_route, {initial_state.active_grid}); + } else { + const bool fixed_wg = use_gemv ? false : fixed_prefill_wg; + utils::ComputePipelineBundle bundle = + make_bundle(shader_src, fixed_wg, use_gemv ? gemv_wg_size : wg_size); + dispatch_idx = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + initial_state.active_grid.x, + use_gemv ? "linear_q4gsw_coop4_bicol" : prefill_label}); + } + + // Dynamic shapes: recompute one shared Params block and select exactly one + // writer. The prefill pipeline remains the route chosen from max M. graph.add_tensor_resize_hook( in_id, [in_id, @@ -384,58 +488,40 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { has_bias, wg_size, use_gemv, + use_dual_route, use_steel, use_shmem_gemm, dispatch_idx, - uniform_buffer](WebGPUGraph& g) { + route_group, + params_buffer](WebGPUGraph& g) { const auto& d = g.cur_dims(in_id); - if (d.empty()) { - throw std::runtime_error( - "WebGPU linear_q4gsw(resize): empty input dims"); - } - const uint64_t numel = utils::numel_of(d); - if (numel % static_cast(K) != 0u) { - throw std::runtime_error( - "WebGPU linear_q4gsw(resize): live input numel not a multiple " - "of K"); - } - const uint32_t m = - static_cast(numel / static_cast(K)); - if (m == 0u) { - throw std::runtime_error("WebGPU linear_q4gsw(resize): live M == 0"); - } - // Buffers/bind-groups were sized for the build-time max M; a larger - // live M would write out of bounds. - if (m > M) { - throw std::runtime_error( - "WebGPU linear_q4gsw(resize): live M exceeds the build-time max"); - } - const uint32_t wgc = compute_q4gsw_workgroup_count( + const Q4gswExecutionState state = make_q4gsw_execution_state( g.device(), - use_gemv, - use_steel, - use_shmem_gemm, - m, + d, + M, + K, N, + K_packed, + gs, + padded_N, + has_bias, wg_size, - "linear_q4gsw(resize)"); - Q4gswParams p = {}; - p.M = m; - p.N = N; - p.K = K; - p.K_packed = K_packed; - p.group_size = gs; - p.padded_N = padded_N; - p.has_bias = has_bias; - wgpuQueueWriteBuffer(g.queue(), uniform_buffer, 0, &p, sizeof(p)); - g.dispatch_at(dispatch_idx).workgroup_count_x = wgc; - std::vector od(d.begin(), d.end()); - od.back() = static_cast(N); - g.set_cur_dims(out_id, od); + use_gemv, + use_dual_route, + use_steel, + use_shmem_gemm); + wgpuQueueWriteBuffer( + g.queue(), params_buffer, 0, &state.params, sizeof(state.params)); + if (use_dual_route) { + g.select_dispatch_route( + route_group, state.active_route, {state.active_grid}); + } else { + auto& dispatch = g.dispatch_at(dispatch_idx); + dispatch.workgroup_count_x = state.active_grid.x; + dispatch.workgroup_count_y = state.active_grid.y; + } + g.set_cur_dims(out_id, state.output_dims); }); - - // Graph owns it so the resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); } } // namespace diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 96ec782c72e..3c06d60747f 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -75,6 +75,22 @@ struct ComputeOutParams { }; static_assert(sizeof(ComputeOutParams) == 32, "ComputeOutParams must be 32B"); +struct SdpaLiveState { + int64_t s; + int64_t pos; + int64_t context_len; + UpdateCacheParams update_cache; + AttnWeightsParams attn_weights; + SoftmaxParams softmax; + ComputeOutParams compute_out; + utils::WgCount update_cache_grid; + utils::WgCount qk_grid; + utils::WgCount softmax_grid; + utils::WgCount av_grid; + bool use_fd; + SdpaFdDecodeState fd; +}; + // Param-struct builder helpers — used in both initial build and resize hook. static UpdateCacheParams make_update_cache_params( uint64_t kv_numel, @@ -383,20 +399,6 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error("WebGPU sdpa: only is_causal=true is supported"); } - // KV cache written in place; only attn_weights/softmax need scratch. - const uint64_t aw_floats = static_cast(Hq) * - static_cast(S) * static_cast(context_len); - // Dynamic input_pos: size+bind scratch for Cmax (no realloc; covers any ctx). - const uint64_t aw_cap_floats = static_cast(Hq) * - static_cast(S) * - static_cast(dynamic_pos ? Cmax : context_len); - const uint64_t aw_bytes = aw_cap_floats * sizeof(float); - - // Dynamic input_pos: the resize hook rewrites these per step. - WGPUBuffer uc_k_buf = nullptr, uc_v_buf = nullptr, qk_buf = nullptr, - softmax_buf = nullptr, av_buf = nullptr; - size_t qk_idx = 0, uc_k_idx = 0, uc_v_idx = 0, softmax_idx = 0, av_idx = 0; - const WGPUDevice device = graph.device(); const uint32_t uc_wg = utils::clamp_workgroup_size(device, kUpdateCacheWorkgroupSizeX); @@ -404,258 +406,332 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { device, kSdpaComputeAttnWeightsWorkgroupSizeX); const uint32_t av_wg = utils::clamp_workgroup_size(device, kSdpaComputeOutWorkgroupSizeX); + const uint32_t sm_wg = + utils::clamp_workgroup_size_pow2(device, kSdpaSoftmaxWorkgroupSizeX); + const bool fd_eligible = D <= kSdpaFdMaxHeadDim; + const int64_t pos_const = input_pos; + + auto compute_live_state = [q_id, + k_id, + v_id, + out_id, + qn, + kn, + S, + dynamic_pos, + input_pos_id, + pos_const, + Hq, + Hkv, + D, + Cmax, + g, + scale, + uc_wg, + qk_wg, + av_wg, + fd_eligible](WebGPUGraph& gr) { + SdpaLiveState state = {}; + const auto& q_live_dims = gr.cur_dims(q_id); + state.s = q_live_dims[qn - 3]; + state.pos = dynamic_pos ? static_cast(gr.read_symint(input_pos_id)) + : pos_const; + if (state.s <= 0 || state.pos < 0 || state.s > S) { + throw std::runtime_error("WebGPU sdpa: invalid live S or input_pos"); + } + if (gr.cur_dims(k_id)[kn - 3] != state.s || + gr.cur_dims(v_id)[gr.cur_dims(v_id).size() - 3] != state.s) { + throw std::runtime_error("WebGPU sdpa: live q/k/v seq_len mismatch"); + } + const auto& out_max_dims = gr.get_tensor(out_id).dims; + if (out_max_dims.size() != q_live_dims.size()) { + throw std::runtime_error("WebGPU sdpa: output rank must match q"); + } + for (size_t i = 0; i < q_live_dims.size(); i++) { + if (q_live_dims[i] <= 0 || q_live_dims[i] > out_max_dims[i]) { + throw std::runtime_error( + "WebGPU sdpa: live output shape exceeds allocation"); + } + } + state.context_len = state.s + state.pos; + if (state.context_len <= 0 || state.context_len > Cmax || + state.s > UINT32_MAX || state.pos > UINT32_MAX || + state.context_len > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa: live dimensions exceed cache or uint32 capacity"); + } + + const uint64_t kv_numel = static_cast(state.s) * + static_cast(Hkv) * static_cast(D); + const uint64_t kv_offset = static_cast(state.pos) * + static_cast(Hkv) * static_cast(D); + const uint64_t cache_numel = static_cast(Cmax) * + static_cast(Hkv) * static_cast(D); + const uint64_t aw_floats = static_cast(Hq) * + static_cast(state.s) * + static_cast(state.context_len); + const uint64_t qk_tiles = static_cast(Hq) * + static_cast(utils::div_up(state.s, kSdpaTileM)) * + static_cast(utils::div_up(state.context_len, kSdpaTileN)); + const uint64_t softmax_rows = + static_cast(Hq) * static_cast(state.s); + const uint64_t av_tiles = static_cast(Hq) * + static_cast(utils::div_up(state.s, kSdpaTileM)) * + static_cast(utils::div_up(D, kSdpaTileN)); + if (kv_numel > UINT32_MAX || kv_offset > UINT32_MAX || + cache_numel > UINT32_MAX || aw_floats > UINT32_MAX || + qk_tiles > UINT32_MAX || softmax_rows > UINT32_MAX || + av_tiles > UINT32_MAX) { + throw std::runtime_error("WebGPU sdpa: live workload exceeds uint32"); + } + + state.update_cache = make_update_cache_params( + kv_numel, static_cast(kv_offset), cache_numel); + state.attn_weights = make_attn_weights_params( + state.s, Hq, Hkv, D, state.context_len, state.pos, g, scale); + state.softmax = make_softmax_params(Hq, state.s, state.context_len); + state.compute_out = + make_compute_out_params(state.s, Hq, Hkv, D, state.context_len, g); + state.update_cache_grid = { + utils::compute_1d_workgroup_count( + gr.device(), static_cast(kv_numel), uc_wg, "uc(resize)"), + 1u}; + state.qk_grid = utils::compute_2d_workgroup_count( + gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); + state.softmax_grid = utils::compute_2d_workgroup_count( + gr.device(), static_cast(softmax_rows), 1, "softmax(resize)"); + state.av_grid = utils::compute_2d_workgroup_count( + gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); + state.use_fd = fd_eligible && state.s == 1; + // make_sdpa_fd_decode_state requires D % 4 == 0; the op-level guard above + // ("head_dim (D) must be a multiple of 4") rejects any other D before this + // lambda runs, so eager construction here can never throw on it. + if (fd_eligible) { + state.fd = make_sdpa_fd_decode_state( + gr.device(), Hq, Hkv, D, state.context_len, g, scale); + } + return state; + }; + + const SdpaLiveState initial_state = compute_live_state(graph); + const uint64_t aw_cap_floats = static_cast(Hq) * + static_cast(S) * + static_cast(dynamic_pos ? Cmax : context_len); + const uint64_t aw_bytes = aw_cap_floats * sizeof(float); - // Dispatches 1-2: write new K/V into the caches (reuses update_cache). - const uint64_t kv_numel = static_cast(S) * - static_cast(Hkv) * static_cast(D); - const uint32_t kv_dst_offset = static_cast( - static_cast(input_pos) * static_cast(Hkv) * - static_cast(D)); - uc_k_buf = record_update_cache_dispatch( + WGPUBuffer uc_k_buf = record_update_cache_dispatch( graph, device, k_cache, k, - kv_numel, - kv_dst_offset, - numel(k_cache), + initial_state.update_cache.numel, + initial_state.update_cache.dst_offset, + initial_state.update_cache.cache_numel, uc_wg, true, "update_cache(K)"); - uc_v_buf = record_update_cache_dispatch( + WGPUBuffer uc_v_buf = record_update_cache_dispatch( graph, device, v_cache, v, - kv_numel, - kv_dst_offset, - numel(v_cache), + initial_state.update_cache.numel, + initial_state.update_cache.dst_offset, + initial_state.update_cache.cache_numel, uc_wg, true, "update_cache(V)"); - uc_k_idx = graph.num_dispatches() - 2; - uc_v_idx = graph.num_dispatches() - 1; - - // FlashDecoding decode (S==1, static pos). Shapes FD can't handle (head dim - // > kSdpaFdMaxHeadDim) fall through to the materialized path below. - if (S == 1 && !dynamic_pos && D <= kSdpaFdMaxHeadDim) { - sdpa_fd_decode_dispatch( - graph, q, k_cache, v_cache, out, Hq, Hkv, D, context_len, g, scale); - return; - } - - // QK/softmax scratch — allocated only on the non-FD path (Hq*S*Cmax prefill). - WGPUBuffer attn_weights = graph.acquire_scratch(aw_bytes); - WebGPUGraph::ScopedScratch attn_weights_guard(&graph, attn_weights); - WGPUBuffer attn_weights_softmax = graph.acquire_scratch(aw_bytes); - WebGPUGraph::ScopedScratch attn_weights_softmax_guard( - &graph, attn_weights_softmax); - - // --- Dispatch 3: QK -> attn_weights. One thread per TM x TN tile. - { - if (aw_floats > UINT32_MAX) { - throw std::runtime_error( - "WebGPU sdpa: Hq*S*context_len exceeds uint32 max"); - } - const int64_t qk_tiles = Hq * utils::div_up(S, kSdpaTileM) * - utils::div_up(context_len, kSdpaTileN); - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - device, static_cast(qk_tiles), qk_wg, "QK"); - AttnWeightsParams p = make_attn_weights_params( - S, Hq, Hkv, D, context_len, input_pos, g, scale); - WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); - BufferBinding bindings[3] = { + const size_t uc_k_idx = graph.num_dispatches() - 2; + const size_t uc_v_idx = graph.num_dispatches() - 1; + const bool dynamic_sequence = graph.tensor_has_dynamic_dims(q_id) || + graph.tensor_has_dynamic_dims(k_id) || + graph.tensor_has_dynamic_dims(v_id); + const bool dual_route = + utils::should_record_sdpa_dual_route(fd_eligible, dynamic_sequence); + const bool record_materialized = dual_route || !initial_state.use_fd; + const bool record_fd = dual_route || initial_state.use_fd; + + WGPUBuffer qk_buf = nullptr; + WGPUBuffer softmax_buf = nullptr; + WGPUBuffer av_buf = nullptr; + size_t qk_idx = 0; + size_t softmax_idx = 0; + size_t av_idx = 0; + utils::DispatchRange materialized_range = {}; + if (record_materialized) { + WGPUBuffer attn_weights = graph.acquire_scratch(aw_bytes); + WebGPUGraph::ScopedScratch attn_weights_guard(&graph, attn_weights); + WGPUBuffer attn_weights_softmax = graph.acquire_scratch(aw_bytes); + WebGPUGraph::ScopedScratch attn_weights_softmax_guard( + &graph, attn_weights_softmax); + + materialized_range.begin = graph.num_dispatches(); + qk_buf = graph.make_uniform_buffer( + &initial_state.attn_weights, sizeof(AttnWeightsParams)); + BufferBinding qk_bindings[3] = { {attn_weights, aw_bytes}, {q.buffer, q.nbytes}, {k_cache.buffer, k_cache.nbytes}}; - const char* qk_src = kSdpaComputeAttnWeightsWGSL; - if (graph.kv_f16()) { - qk_src = kSdpaComputeAttnWeightsHalfWGSL; - } + const char* qk_src = graph.kv_f16() ? kSdpaComputeAttnWeightsHalfWGSL + : kSdpaComputeAttnWeightsWGSL; build_dispatch( graph, qk_src, - bindings, + qk_bindings, 3, - ubuf, - sizeof(p), - wgc.x, - wgc.y, + qk_buf, + sizeof(AttnWeightsParams), + initial_state.qk_grid.x, + initial_state.qk_grid.y, qk_wg, true, "sdpa_compute_attn_weights"); - qk_buf = ubuf; qk_idx = graph.num_dispatches() - 1; - } - // Dispatch 4: softmax, one workgroup per (h,s) row of width context_len. - { - // One workgroup per (h,s) row; wg_size 1 keeps the device dispatch check. - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - device, static_cast(Hq * S), 1, "softmax"); - const uint32_t sm_wg = - utils::clamp_workgroup_size_pow2(device, kSdpaSoftmaxWorkgroupSizeX); - SoftmaxParams p = make_softmax_params(Hq, S, context_len); - WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); - BufferBinding bindings[2] = { + softmax_buf = graph.make_uniform_buffer( + &initial_state.softmax, sizeof(SoftmaxParams)); + BufferBinding softmax_bindings[2] = { {attn_weights_softmax, aw_bytes}, {attn_weights, aw_bytes}}; build_dispatch( graph, kSdpaSoftmaxWGSL, - bindings, + softmax_bindings, 2, - ubuf, - sizeof(p), - wgc.x, - wgc.y, + softmax_buf, + sizeof(SoftmaxParams), + initial_state.softmax_grid.x, + initial_state.softmax_grid.y, sm_wg, true, "sdpa_softmax"); - softmax_buf = ubuf; softmax_idx = graph.num_dispatches() - 1; - } - // --- Dispatch 5: AV -> out. One thread per TM x TN tile. - { - const int64_t av_tiles = - Hq * utils::div_up(S, kSdpaTileM) * utils::div_up(D, kSdpaTileN); - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - device, static_cast(av_tiles), av_wg, "AV"); - ComputeOutParams p = make_compute_out_params(S, Hq, Hkv, D, context_len, g); - WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); - BufferBinding bindings[3] = { + av_buf = graph.make_uniform_buffer( + &initial_state.compute_out, sizeof(ComputeOutParams)); + BufferBinding av_bindings[3] = { {out.buffer, out.nbytes}, {attn_weights_softmax, aw_bytes}, {v_cache.buffer, v_cache.nbytes}}; - const char* av_src = kSdpaComputeOutWGSL; - if (graph.kv_f16()) { - av_src = kSdpaComputeOutHalfWGSL; - } + const char* av_src = + graph.kv_f16() ? kSdpaComputeOutHalfWGSL : kSdpaComputeOutWGSL; build_dispatch( graph, av_src, - bindings, + av_bindings, 3, - ubuf, - sizeof(p), - wgc.x, - wgc.y, + av_buf, + sizeof(ComputeOutParams), + initial_state.av_grid.x, + initial_state.av_grid.y, av_wg, true, "sdpa_compute_out"); - av_buf = ubuf; av_idx = graph.num_dispatches() - 1; + materialized_range.end = graph.num_dispatches(); } - // Per-step recompute: live S (q resize) or input_pos (SymInt); inert if - // static. - const int64_t pos_const = input_pos; - auto sdpa_resize = [q_id, - qn, - S, - out_id, - dynamic_pos, - input_pos_id, - pos_const, - Hq, - Hkv, - D, - Cmax, - g, - scale, - qk_idx, - uc_k_idx, - uc_v_idx, - softmax_idx, - av_idx, - uc_wg, - qk_wg, - av_wg, - uc_k_buf, - uc_v_buf, - qk_buf, - softmax_buf, - av_buf](WebGPUGraph& gr) { - const int64_t s = gr.cur_dims(q_id)[qn - 3]; - const int64_t pos = dynamic_pos - ? static_cast(gr.read_symint(input_pos_id)) - : pos_const; - if (s <= 0 || pos < 0) { - throw std::runtime_error("WebGPU sdpa: invalid live S or input_pos"); - } - // Scratch (attn_weights/softmax) is sized at build for S=max; a larger live - // S would overrun it. Make that invariant load-bearing. - if (s > S) { - throw std::runtime_error( - "WebGPU sdpa: live S exceeds the build-time max (scratch capacity)"); - } - const int64_t ctx = s + pos; - if (ctx <= 0 || ctx > Cmax) { - throw std::runtime_error( - "WebGPU sdpa: context_len exceeds cache capacity"); + SdpaFdDecodeResources fd_resources = {}; + size_t route_group = 0; + if (record_fd) { + fd_resources = record_sdpa_fd_decode_dispatches( + graph, q, k_cache, v_cache, out, initial_state.fd); + } + if (dual_route) { + route_group = graph.register_dispatch_route_group( + {materialized_range, fd_resources.dispatch_range}); + } + + auto refresh_state = [compute_live_state, + q_id, + out_id, + dual_route, + record_materialized, + record_fd, + fixed_use_fd = initial_state.use_fd, + route_group, + uc_k_idx, + uc_v_idx, + qk_idx, + softmax_idx, + av_idx, + uc_k_buf, + uc_v_buf, + qk_buf, + softmax_buf, + av_buf, + fd_resources](WebGPUGraph& gr) { + const SdpaLiveState state = compute_live_state(gr); + + wgpuQueueWriteBuffer( + gr.queue(), + uc_k_buf, + 0, + &state.update_cache, + sizeof(state.update_cache)); + wgpuQueueWriteBuffer( + gr.queue(), + uc_v_buf, + 0, + &state.update_cache, + sizeof(state.update_cache)); + if (record_materialized) { + wgpuQueueWriteBuffer( + gr.queue(), + qk_buf, + 0, + &state.attn_weights, + sizeof(state.attn_weights)); + wgpuQueueWriteBuffer( + gr.queue(), softmax_buf, 0, &state.softmax, sizeof(state.softmax)); + wgpuQueueWriteBuffer( + gr.queue(), av_buf, 0, &state.compute_out, sizeof(state.compute_out)); } - const uint32_t kv_off = static_cast( - static_cast(pos) * static_cast(Hkv) * - static_cast(D)); - const uint64_t aw_floats = static_cast(Hq) * - static_cast(s) * static_cast(ctx); - if (aw_floats > UINT32_MAX) { - throw std::runtime_error("WebGPU sdpa: Hq*S*context_len exceeds uint32"); + if (record_fd) { + write_sdpa_fd_decode_uniforms(gr.queue(), fd_resources, state.fd); } - const uint64_t kv_numel = static_cast(s) * - static_cast(Hkv) * static_cast(D); - if (kv_numel > UINT32_MAX) { - throw std::runtime_error("WebGPU sdpa: S*Hkv*D exceeds uint32"); - } - const uint64_t k_cache_numel = static_cast(Cmax) * - static_cast(Hkv) * static_cast(D); - // update_cache K/V: dispatch (kv_numel) + dst offset scale with live S/pos. - UpdateCacheParams uc = - make_update_cache_params(kv_numel, kv_off, k_cache_numel); - wgpuQueueWriteBuffer(gr.queue(), uc_k_buf, 0, &uc, sizeof(uc)); - wgpuQueueWriteBuffer(gr.queue(), uc_v_buf, 0, &uc, sizeof(uc)); - const uint32_t uc_wgc = utils::compute_1d_workgroup_count( - gr.device(), static_cast(kv_numel), uc_wg, "uc(resize)"); - gr.dispatch_at(uc_k_idx).workgroup_count_x = uc_wgc; - gr.dispatch_at(uc_v_idx).workgroup_count_x = uc_wgc; - - // QK: one thread per TM x TN tile; grid = Hq*ceil(S/TM)*ceil(ctx/TN). - AttnWeightsParams qp = - make_attn_weights_params(s, Hq, Hkv, D, ctx, pos, g, scale); - wgpuQueueWriteBuffer(gr.queue(), qk_buf, 0, &qp, sizeof(qp)); - const int64_t qk_tiles = - Hq * utils::div_up(s, kSdpaTileM) * utils::div_up(ctx, kSdpaTileN); - const utils::WgCount qk_wgc = utils::compute_2d_workgroup_count( - gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); - gr.dispatch_at(qk_idx).workgroup_count_x = qk_wgc.x; - gr.dispatch_at(qk_idx).workgroup_count_y = qk_wgc.y; - - // softmax: one workgroup per (h,s) row. - SoftmaxParams sp = make_softmax_params(Hq, s, ctx); - wgpuQueueWriteBuffer(gr.queue(), softmax_buf, 0, &sp, sizeof(sp)); - const utils::WgCount sm_wgc = utils::compute_2d_workgroup_count( - gr.device(), static_cast(Hq * s), 1, "softmax(resize)"); - gr.dispatch_at(softmax_idx).workgroup_count_x = sm_wgc.x; - gr.dispatch_at(softmax_idx).workgroup_count_y = sm_wgc.y; - - // AV: one thread per TM x TN tile; grid = Hq*ceil(S/TM)*ceil(D/TN). - ComputeOutParams op = make_compute_out_params(s, Hq, Hkv, D, ctx, g); - wgpuQueueWriteBuffer(gr.queue(), av_buf, 0, &op, sizeof(op)); - const int64_t av_tiles = - Hq * utils::div_up(s, kSdpaTileM) * utils::div_up(D, kSdpaTileN); - const utils::WgCount av_wgc = utils::compute_2d_workgroup_count( - gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); - gr.dispatch_at(av_idx).workgroup_count_x = av_wgc.x; - gr.dispatch_at(av_idx).workgroup_count_y = av_wgc.y; - - // Output attn has the same shape as q: [.., S, Hq, D]. + gr.dispatch_at(uc_k_idx).workgroup_count_x = state.update_cache_grid.x; + gr.dispatch_at(uc_k_idx).workgroup_count_y = state.update_cache_grid.y; + gr.dispatch_at(uc_v_idx).workgroup_count_x = state.update_cache_grid.x; + gr.dispatch_at(uc_v_idx).workgroup_count_y = state.update_cache_grid.y; + if (dual_route) { + const size_t active_route = state.use_fd ? 1 : 0; + const std::vector active_grids = state.use_fd + ? std::vector< + utils::WgCount>{state.fd.split_grid, state.fd.reduce_grid} + : std::vector{ + state.qk_grid, state.softmax_grid, state.av_grid}; + gr.select_dispatch_route(route_group, active_route, active_grids); + } else if (state.use_fd) { + if (!fixed_use_fd) { + throw std::runtime_error("WebGPU sdpa: static route changed"); + } + gr.dispatch_at(fd_resources.dispatch_range.begin).workgroup_count_x = + state.fd.split_grid.x; + gr.dispatch_at(fd_resources.dispatch_range.begin).workgroup_count_y = + state.fd.split_grid.y; + gr.dispatch_at(fd_resources.dispatch_range.begin + 1).workgroup_count_x = + state.fd.reduce_grid.x; + gr.dispatch_at(fd_resources.dispatch_range.begin + 1).workgroup_count_y = + state.fd.reduce_grid.y; + } else { + if (fixed_use_fd) { + throw std::runtime_error("WebGPU sdpa: static route changed"); + } + gr.dispatch_at(qk_idx).workgroup_count_x = state.qk_grid.x; + gr.dispatch_at(qk_idx).workgroup_count_y = state.qk_grid.y; + gr.dispatch_at(softmax_idx).workgroup_count_x = state.softmax_grid.x; + gr.dispatch_at(softmax_idx).workgroup_count_y = state.softmax_grid.y; + gr.dispatch_at(av_idx).workgroup_count_x = state.av_grid.x; + gr.dispatch_at(av_idx).workgroup_count_y = state.av_grid.y; + } gr.set_cur_dims(out_id, gr.cur_dims(q_id)); }; - // q and input_pos share one idempotent recompute; a double-fire is harmless. - graph.add_tensor_resize_hook(q_id, sdpa_resize); + + refresh_state(graph); + graph.add_tensor_resize_hook(q_id, refresh_state); if (dynamic_pos) { - graph.add_resize_hook(input_pos_id, sdpa_resize); + graph.add_resize_hook(input_pos_id, refresh_state); } } diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp index 82ead71b597..b39c2ae9014 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp @@ -17,6 +17,7 @@ #include +#include #include #include #include @@ -70,6 +71,7 @@ void build_dispatch( WGPUBuffer uniform_buffer, uint64_t uniform_size, uint32_t workgroup_count_x, + bool retain_uniform, const char* kernel_name) { WGPUDevice device = graph.device(); @@ -101,26 +103,55 @@ void build_dispatch( graph.add_dispatch( {bundle.pipeline, bundle.bind_group, workgroup_count_x, kernel_name}); - wgpuBufferRelease(uniform_buffer); + if (retain_uniform) { + graph.own_uniform_buffer(uniform_buffer); + } else { + wgpuBufferRelease(uniform_buffer); + } +} + +FdSplitParams make_split_params(const SdpaFdDecodeState& state) { + FdSplitParams params = {}; + params.Hkv = state.Hkv; + params.D = state.D; + params.context_len = state.context_len; + params.g = state.g; + params.num_splits = state.num_splits; + params.split_len = state.split_len; + params.scale = state.scale; + return params; +} + +FdReduceParams make_reduce_params(const SdpaFdDecodeState& state) { + FdReduceParams params = {}; + params.D = state.D; + params.num_splits = state.num_splits; + return params; } } // namespace -void sdpa_fd_decode_dispatch( - WebGPUGraph& graph, - const WebGPUTensor& q, - const WebGPUTensor& k_cache, - const WebGPUTensor& v_cache, - const WebGPUTensor& out, +SdpaFdDecodeState make_sdpa_fd_decode_state( + WGPUDevice device, int64_t Hq, int64_t Hkv, int64_t D, int64_t context_len, int64_t g, float scale) { - // Defensive contract guard: the Sdpa.cpp gate only routes D <= this here, but - // keep the check (lane-owns-D reach) so a future caller can't silently - // overrun. + if (Hq <= 0 || Hkv <= 0 || D <= 0 || context_len <= 0 || g <= 0) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: dimensions must be positive"); + } + if (Hq != Hkv * g) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: inconsistent GQA dimensions"); + } + if (Hq > UINT32_MAX || Hkv > UINT32_MAX || D > UINT32_MAX || + context_len > UINT32_MAX || g > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: parameter exceeds uint32 max"); + } if (D > kSdpaFdMaxHeadDim) { throw std::runtime_error( "WebGPU sdpa FlashDecoding: head dim must be <= " + @@ -130,26 +161,59 @@ void sdpa_fd_decode_dispatch( throw std::runtime_error( "WebGPU sdpa FlashDecoding: head dim must be a multiple of 4"); } - // context_len 0 -> split_len 0 -> empty KV loop -> silent zero output; the - // Sdpa.cpp gate guarantees ctx >= 1, but fail loud if called directly. - if (context_len <= 0) { - throw std::runtime_error( - "WebGPU sdpa FlashDecoding: context_len must be positive"); - } - // Split factor: one split per kSdpaFdSplitTile KV rows, capped. uint32_t num_splits = static_cast( (context_len + kSdpaFdSplitTile - 1) / kSdpaFdSplitTile); - if (num_splits > kSdpaFdMaxSplits) { - num_splits = kSdpaFdMaxSplits; - } + num_splits = std::min(num_splits, kSdpaFdMaxSplits); const uint32_t split_len = static_cast((context_len + num_splits - 1) / num_splits); + const uint64_t split_threads = static_cast(Hq) * + static_cast(num_splits) * + static_cast(kSdpaFdSplitWorkgroupSizeX); + const uint64_t reduce_threads = + static_cast(Hq) * kSdpaFdReduceWorkgroupSizeX; + if (split_threads > UINT32_MAX || reduce_threads > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: thread count exceeds uint32 max"); + } + + const uint32_t split_wgc = utils::compute_1d_workgroup_count( + device, + static_cast(split_threads), + kSdpaFdSplitWorkgroupSizeX, + "fd_split"); + const uint32_t reduce_wgc = utils::compute_1d_workgroup_count( + device, + static_cast(reduce_threads), + kSdpaFdReduceWorkgroupSizeX, + "fd_reduce"); + return { + static_cast(Hq), + static_cast(Hkv), + static_cast(D), + static_cast(context_len), + static_cast(g), + num_splits, + split_len, + scale, + {split_wgc, 1u}, + {reduce_wgc, 1u}}; +} + +SdpaFdDecodeResources record_sdpa_fd_decode_dispatches( + WebGPUGraph& graph, + const WebGPUTensor& q, + const WebGPUTensor& k_cache, + const WebGPUTensor& v_cache, + const WebGPUTensor& out, + const SdpaFdDecodeState& state) { + const size_t dispatch_begin = graph.num_dispatches(); + // Scratch: per-(head,split) partials at kSdpaFdMaxSplits stride. - const uint64_t po_floats = static_cast(Hq) * - static_cast(kSdpaFdMaxSplits) * static_cast(D); - const uint64_t pml_floats = static_cast(Hq) * + const uint64_t po_floats = static_cast(state.Hq) * + static_cast(kSdpaFdMaxSplits) * static_cast(state.D); + const uint64_t pml_floats = static_cast(state.Hq) * static_cast(kSdpaFdMaxSplits) * 2ull; WGPUBuffer part_o = graph.acquire_scratch(po_floats * sizeof(float)); WebGPUGraph::ScopedScratch part_o_guard(&graph, part_o); @@ -157,14 +221,7 @@ void sdpa_fd_decode_dispatch( WebGPUGraph::ScopedScratch part_ml_guard(&graph, part_ml); // Pass 1: split (Hq*num_splits WGs) -> writes part_o, part_ml. - FdSplitParams sp = {}; - sp.Hkv = static_cast(Hkv); - sp.D = static_cast(D); - sp.context_len = static_cast(context_len); - sp.g = static_cast(g); - sp.num_splits = num_splits; - sp.split_len = split_len; - sp.scale = scale; + FdSplitParams sp = make_split_params(state); WGPUBuffer ub_split = graph.make_uniform_buffer(&sp, sizeof(sp)); BufferBinding split_bindings[5] = { {part_o, po_floats * sizeof(float)}, @@ -172,20 +229,6 @@ void sdpa_fd_decode_dispatch( {q.buffer, q.nbytes}, {k_cache.buffer, k_cache.nbytes}, {v_cache.buffer, v_cache.nbytes}}; - // Compute the thread product in 64-bit + guard before the u32 cast, mirroring - // the Sdpa.cpp aw_floats > UINT32_MAX guards. - const uint64_t split_threads = static_cast(Hq) * - static_cast(num_splits) * - static_cast(kSdpaFdSplitWorkgroupSizeX); - if (split_threads > UINT32_MAX) { - throw std::runtime_error( - "WebGPU sdpa FlashDecoding: split thread count exceeds uint32 max"); - } - const uint32_t wgc_split = utils::compute_1d_workgroup_count( - graph.device(), - static_cast(split_threads), - kSdpaFdSplitWorkgroupSizeX, - "fd_split"); const char* split_shader = kSdpaFdSplitWGSL; if (graph.kv_f16()) { split_shader = kSdpaFdSplitHalfWGSL; @@ -198,23 +241,17 @@ void sdpa_fd_decode_dispatch( 2, ub_split, sizeof(sp), - wgc_split, + state.split_grid.x, + true, "fd_split"); // Pass 2: reduce (Hq WGs) -> reads part_o, part_ml; writes out. - FdReduceParams rp = {}; - rp.D = static_cast(D); - rp.num_splits = num_splits; + FdReduceParams rp = make_reduce_params(state); WGPUBuffer ub_reduce = graph.make_uniform_buffer(&rp, sizeof(rp)); BufferBinding reduce_bindings[3] = { {out.buffer, out.nbytes}, {part_o, po_floats * sizeof(float)}, {part_ml, pml_floats * sizeof(float)}}; - const uint32_t wgc_reduce = utils::compute_1d_workgroup_count( - graph.device(), - static_cast(Hq) * kSdpaFdReduceWorkgroupSizeX, - kSdpaFdReduceWorkgroupSizeX, - "fd_reduce"); build_dispatch( graph, kSdpaFdReduceWGSL, @@ -223,8 +260,27 @@ void sdpa_fd_decode_dispatch( 1, ub_reduce, sizeof(rp), - wgc_reduce, + state.reduce_grid.x, + true, "fd_reduce"); + + return {ub_split, ub_reduce, {dispatch_begin, graph.num_dispatches()}}; +} + +void write_sdpa_fd_decode_uniforms( + WGPUQueue queue, + const SdpaFdDecodeResources& resources, + const SdpaFdDecodeState& state) { + FdSplitParams split_params = make_split_params(state); + FdReduceParams reduce_params = make_reduce_params(state); + wgpuQueueWriteBuffer( + queue, resources.split_uniform, 0, &split_params, sizeof(split_params)); + wgpuQueueWriteBuffer( + queue, + resources.reduce_uniform, + 0, + &reduce_params, + sizeof(reduce_params)); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h index 1b161cd8ec0..d86b4c624b8 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include @@ -19,15 +20,27 @@ namespace executorch::backends::webgpu { // SDPA path (the FD selection predicate in Sdpa.cpp checks this). constexpr int64_t kSdpaFdMaxHeadDim = 128; -// Split-KV FlashDecoding decode dispatch (S==1): a split pass over -// Hq*num_splits workgroups + a reduce pass over Hq workgroups. Called from the -// Sdpa.cpp WEBGPU_SDPA_FD branch. -void sdpa_fd_decode_dispatch( - WebGPUGraph& graph, - const WebGPUTensor& q, - const WebGPUTensor& k_cache, - const WebGPUTensor& v_cache, - const WebGPUTensor& out, +struct SdpaFdDecodeState { + uint32_t Hq; + uint32_t Hkv; + uint32_t D; + uint32_t context_len; + uint32_t g; + uint32_t num_splits; + uint32_t split_len; + float scale; + utils::WgCount split_grid; + utils::WgCount reduce_grid; +}; + +struct SdpaFdDecodeResources { + WGPUBuffer split_uniform; + WGPUBuffer reduce_uniform; + utils::DispatchRange dispatch_range; +}; + +SdpaFdDecodeState make_sdpa_fd_decode_state( + WGPUDevice device, int64_t Hq, int64_t Hkv, int64_t D, @@ -35,4 +48,19 @@ void sdpa_fd_decode_dispatch( int64_t g, float scale); +// Records split + reduce with retained UBOs. Route selection is owned by the +// caller so this helper never mutates recorded dispatch counts. +SdpaFdDecodeResources record_sdpa_fd_decode_dispatches( + WebGPUGraph& graph, + const WebGPUTensor& q, + const WebGPUTensor& k_cache, + const WebGPUTensor& v_cache, + const WebGPUTensor& out, + const SdpaFdDecodeState& state); + +void write_sdpa_fd_decode_uniforms( + WGPUQueue queue, + const SdpaFdDecodeResources& resources, + const SdpaFdDecodeState& state); + } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/scripts/gen_wgsl_headers.py b/backends/webgpu/scripts/gen_wgsl_headers.py index c66aff2039a..cf044232dde 100644 --- a/backends/webgpu/scripts/gen_wgsl_headers.py +++ b/backends/webgpu/scripts/gen_wgsl_headers.py @@ -30,7 +30,7 @@ import sys from itertools import product from pathlib import Path -from typing import Any, Dict, List, Optional, Set +from typing import Any, Dict, List, NamedTuple, Optional, Set import yaml from yaml.constructor import ConstructorError @@ -124,6 +124,12 @@ def escape(line: str) -> str: def preprocess( input_text: str, variables: Dict[str, Any], input_path: str = "codegen" ) -> str: + # Normalize line endings first. Templates checked out with CRLF (common on + # Windows) otherwise break the trailing-backslash handling below: in + # re.MULTILINE, $ matches immediately before \n, so a CR would sit between a + # trailing \ and the line end and defeat the r"\\$" match, leaving a lone + # backslash that escape() turns into an unterminated Python string literal. + input_text = input_text.replace("\r\n", "\n").replace("\r", "\n") # Workaround to handle source files using \ to extend mecros to a new line input_text = re.sub(r"\\$", r"\\\\", input_text, flags=re.MULTILINE) @@ -496,6 +502,93 @@ def discover(): return sorted((BACKEND_ROOT / "runtime/ops").glob("**/*.wgsl")) +class RegistryEntry(NamedTuple): + name: str + include: str + symbol: str + + +def registry_path() -> Path: + return BACKEND_ROOT / "runtime/WebGPUShaderRegistry.cpp" + + +def registry_entries() -> List[RegistryEntry]: + """Return one registry entry for every concrete generated shader.""" + entries = [] + for wgsl in discover(): + for header, _ in headers_for_shader(wgsl): + suffix = "_wgsl.h" + if not header.name.endswith(suffix): + raise ValueError(f"unexpected generated header name: {header.name}") + name = header.name[: -len(suffix)] + entries.append( + RegistryEntry( + name=name, + include=header.relative_to(BACKEND_ROOT).as_posix(), + symbol=symbol_base(name), + ) + ) + return sorted(entries) + + +def render_registry(entries: List[RegistryEntry]) -> str: + """Render the generated name-to-WGSL registry implementation.""" + ordered = sorted(entries) + names = [entry.name for entry in ordered] + if len(names) != len(set(names)): + raise ValueError("duplicate shader registry name") + + includes = "\n".join( + sorted( + f"#include " + for entry in ordered + ) + ) + values = "\n".join( + " {\n" + f' "{entry.name}",\n' + f" k{entry.symbol}WGSL,\n" + f" k{entry.symbol}WorkgroupSizeX,\n" + f" k{entry.symbol}WorkgroupSizeY,\n" + f" k{entry.symbol}WorkgroupSizeZ,\n" + " }," + for entry in ordered + ) + return f"""{_BSD_HEADER} + +// @generated by scripts/gen_wgsl_headers.py - DO NOT EDIT. + +#include + +{includes} + +#include +#include +#include + +namespace executorch::backends::webgpu {{ +namespace {{ + +constexpr std::array kShaderRegistry = {{{{ +{values} +}}}}; + +}} // namespace + +const WebGPUShaderInfo& get_webgpu_shader_info(std::string_view name) {{ + for (const auto& shader : kShaderRegistry) {{ + if (shader.name == name) {{ + return shader; + }} + }} + throw std::runtime_error( + "WebGPU shader registry: unknown shader '" + std::string(name) + "'"); +}} + +}} // namespace executorch::backends::webgpu +""" + + def headers_for_shader(wgsl): """Yield (header_path, rendered_text) pairs for one shader source. @@ -538,6 +631,16 @@ def _report_drift(missing, stale) -> None: print(f" {h.relative_to(BACKEND_ROOT)}") +def _sync_generated_output(output, want, check, missing, stale) -> None: + """Write one generated file, or record its --check drift.""" + if output.exists() and output.read_text() == want: + return + if check: + (missing if not output.exists() else stale).append(output) + else: + output.write_text(want) + + def main(argv=None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( @@ -562,12 +665,15 @@ def main(argv=None) -> int: continue for header, want in rendered: # Full-content compare (not just the sha) catches generator-logic drift too. - if header.exists() and header.read_text() == want: - continue - if args.check: - (missing if not header.exists() else stale).append(header) - else: - header.write_text(want) + _sync_generated_output(header, want, args.check, missing, stale) + + if not errors: + try: + registry = render_registry(registry_entries()) + output = registry_path() + _sync_generated_output(output, registry, args.check, missing, stale) + except ValueError as e: + errors.append(f"shader registry: {e}") if errors: print("Cannot generate header (malformed shader):") diff --git a/backends/webgpu/scripts/test_webgpu_native_ci.sh b/backends/webgpu/scripts/test_webgpu_native_ci.sh index 810d8165303..75af3ff84e6 100644 --- a/backends/webgpu/scripts/test_webgpu_native_ci.sh +++ b/backends/webgpu/scripts/test_webgpu_native_ci.sh @@ -17,10 +17,8 @@ # source .ci/scripts/setup-webgpu-linux-deps.sh # bash backends/webgpu/scripts/test_webgpu_native_ci.sh # -# Builds whatever native test targets are present in the landed tree (NOT a fixed -# list): webgpu_native_test (base) + webgpu_dispatch_order_test, -# webgpu_scratch_buffer_test (D107576199) + webgpu_update_cache_test -# (D107547307). SDPA executables join once they land. +# Builds and runs the fixed native target matrix defined by this tree. A missing +# target or fixture is a CI failure, not an optional skip. set -e @@ -37,18 +35,38 @@ fi cd "${EXECUTORCH_ROOT}" -# ── Exports for the model-driven executables (best-effort) ─────────────────── -# native_test (quantized_linear/SDPA/update_cache) + dispatch_order read .pte/ -# golden inputs via env/dir and self-skip if absent; scratch is standalone. -# native_test itself is gated below on the executorch wheel being importable. +# ── Exports for the model-driven executables ───────────────────────────────── +if ! "${PYTHON_EXECUTABLE}" -c "import executorch" 2>/dev/null; then + echo "ERROR: executorch wheel unavailable; required fixture exports cannot run" >&2 + exit 1 +fi + +require_file() { + if [[ ! -f "$1" ]]; then + echo "ERROR: required WebGPU fixture missing: $1" >&2 + exit 1 + fi +} + +run_with_required_device() { + local output + if ! output="$("$@" 2>&1)"; then + printf '%s\n' "${output}" + return 1 + fi + printf '%s\n' "${output}" + if ! grep -q '^WebGPU device acquired (native)$' <<<"${output}"; then + echo "ERROR: WebGPU native test did not acquire a device" >&2 + return 1 + fi +} + DISPATCH_ORDER_DIR="/tmp/dispatch_order" -DISPATCH_ORDER_OK=1 UPDATE_CACHE_DIR="/tmp/update_cache" -UPDATE_CACHE_OK=1 INDEX_DIR="/tmp/index" -INDEX_OK=1 DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape" -DYNAMIC_SHAPE_OK=1 +SYMINT_BLOB="/tmp/sdpa_dyn_small.pte" +OUTPUT_SUPPRESSION_DIR="/tmp/output_suppression" EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte" EMBEDDING_INDICES="/tmp/webgpu_embedding_q4gsw_indices.bin" EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin" @@ -69,33 +87,34 @@ PREPACK_TIED_MODEL="/tmp/webgpu_prepack_tied_const.pte" PREPACK_TIED_GOLDEN="/tmp/webgpu_prepack_tied_const_golden.bin" $PYTHON_EXECUTABLE -c " -from executorch.backends.webgpu.test.ops.test_quantized_linear import export_all_quantized_linear_models +from executorch.backends.webgpu.test.ops.test_quantized_linear import export_all_quantized_linear_models, export_output_suppression_models export_all_quantized_linear_models('/tmp') -" || echo "WARN: q4gsw export failed; required configs will FAIL in webgpu_native_test" +export_output_suppression_models('${OUTPUT_SUPPRESSION_DIR}') +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_embedding_q4gsw import export_embedding_q4gsw_model export_embedding_q4gsw_model('${EMBEDDING_MODEL}', '${EMBEDDING_GOLDEN}', '${EMBEDDING_INDICES}') export_embedding_q4gsw_model('${EMBEDDING_LLAMA1B_MODEL}', '${EMBEDDING_LLAMA1B_GOLDEN}', '${EMBEDDING_LLAMA1B_INDICES}', 'llama1b') -" || echo "WARN: embedding_q4gsw export failed; embedding configs will FAIL in webgpu_native_test" +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_rope import export_rope_model export_rope_model('${ROPE_MODEL}', '${ROPE_XQ_GOLDEN}', '${ROPE_XK_GOLDEN}') export_rope_model('${ROPE_DECODE_MODEL}', '${ROPE_DECODE_XQ_GOLDEN}', '${ROPE_DECODE_XK_GOLDEN}', 'decode') -" || echo "WARN: rope export failed; apply_rotary_emb configs will FAIL in webgpu_native_test" +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_prepack import export_prepack_model, export_prepack_two_const_model, export_prepack_tied_const_model export_prepack_model('${PREPACK_MODEL}', '${PREPACK_GOLDEN}') export_prepack_two_const_model('${PREPACK2_MODEL}', '${PREPACK2_GOLDEN}') export_prepack_tied_const_model('${PREPACK_TIED_MODEL}', '${PREPACK_TIED_GOLDEN}') -" || echo "WARN: prepack export failed; prepack configs will FAIL in webgpu_native_test" +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_dispatch_order import export_dispatch_order_cases export_dispatch_order_cases('${DISPATCH_ORDER_DIR}') -" || { echo "WARN: dispatch_order export failed; skipping dispatch_order native test"; DISPATCH_ORDER_OK=0; } +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_update_cache import ( @@ -106,20 +125,18 @@ from executorch.backends.webgpu.test.ops.test_update_cache import ( export_update_cache_cases('${UPDATE_CACHE_DIR}') export_update_cache_replay('${UPDATE_CACHE_DIR}') export_update_cache_negative('${UPDATE_CACHE_DIR}') -" || { echo "WARN: update_cache export failed; skipping update_cache native test"; UPDATE_CACHE_OK=0; } +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.index.test_index import export_all_index_models export_all_index_models('${INDEX_DIR}') -" || { echo "WARN: index export failed; skipping index native test"; INDEX_OK=0; } +" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.dynamic_shape.test_dynamic_shape_export import export_dynamic_shape_cases export_dynamic_shape_cases('${DYNAMIC_SHAPE_DIR}') -" || { echo "WARN: dynamic_shape export failed; skipping dynamic_shape native test"; DYNAMIC_SHAPE_OK=0; } +" -# Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in -# webgpu_native_test below (precise per-config error), so don't exit/mask here. $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_sdpa import ( export_all_sdpa_models, @@ -131,12 +148,16 @@ export_all_sdpa_models('/tmp') export_replay_sequences('/tmp') export_dynamic_decode('/tmp') export_incache_decode('/tmp') -" || echo "WARN: sdpa export failed; required 4k/8k configs will FAIL in webgpu_native_test" +" + +require_file "${SYMINT_BLOB}" +require_file "${OUTPUT_SUPPRESSION_DIR}/input.bin" # ── Configure (Dawn-only: no -DWEBGPU_IMPL; Dawn is the sole backend) ───────── echo "=== Configure WebGPU native tests on Dawn ===" rm -rf "${BUILD_DIR}" cmake \ + -DPYTHON_EXECUTABLE="${PYTHON_EXECUTABLE}" \ -DEXECUTORCH_BUILD_WEBGPU=ON \ -DEXECUTORCH_BUILD_WEBGPU_TEST=ON \ -DEXECUTORCH_BUILD_TESTS=ON \ @@ -150,102 +171,63 @@ cmake \ -B "${BUILD_DIR}" \ "${EXECUTORCH_ROOT}" -# ── Build + run every native test target that exists in this tree ──────────── -TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test) +# ── Build + run every fixed native test target in this tree ────────────────── +REQUIRED_TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test webgpu_compute_dispatch_test webgpu_execution_options_test webgpu_output_suppression_test webgpu_op_test_util_test) BIN_DIR="${BUILD_DIR}/backends/webgpu" -# Which targets are defined depends on which diffs are landed (native_test + -# rms_norm here; dispatch_order + scratch from D107576199). Query the configured -# target list ONCE so a not-yet-landed target is skipped WITHOUT masking a real -# compile failure of a target that IS defined (CI uses the Make generator). DEFINED_TARGETS="$(cmake --build "${BUILD_DIR}" --target help 2>/dev/null || true)" -# Fail loud if the probe found nothing (e.g. a non-Make generator or a cmake -# regression): otherwise every target would skip and the job would go green -# having tested nothing. webgpu_native_test is always defined at/after this diff. -if ! printf '%s\n' "${DEFINED_TARGETS}" | grep -qw webgpu_native_test; then - echo "ERROR: cmake target probe returned no webgpu_native_test; aborting" >&2 - exit 1 -fi - -for t in "${TARGETS[@]}"; do - if printf '%s\n' "${DEFINED_TARGETS}" | grep -qw "${t}"; then - # Defined target: build with stderr visible; set -e fails the job on a real - # build error (never silently skipped). - cmake --build "${BUILD_DIR}" --target "${t}" -j"${NPROC}" - echo "built ${t}" - else - echo "(target ${t} not defined in this tree — skipping)" +for t in "${REQUIRED_TARGETS[@]}"; do + if ! printf '%s\n' "${DEFINED_TARGETS}" | grep -qw "${t}"; then + echo "ERROR: required CMake target is not defined: ${t}" >&2 + exit 1 fi + cmake --build "${BUILD_DIR}" --target "${t}" -j"${NPROC}" + echo "built ${t}" done echo "=== Run native tests on Dawn + SwiftShader ===" -# webgpu_native_test hosts the quantized_linear / SDPA / update_cache / symint -# sweeps. Gate on the executorch wheel being importable (the proxy for "the -# exports above ran"): CI has the wheel so they ran; a bare local run without it -# skips here rather than hard-failing the required-config guards. -if [[ -x "${BIN_DIR}/webgpu_native_test" ]] && - "${PYTHON_EXECUTABLE}" -c "import executorch" 2>/dev/null; then - env WEBGPU_TEST_SDPA_DIR=/tmp/ \ - WEBGPU_TEST_QUANTIZED_LINEAR_DIR=/tmp/ \ - WEBGPU_TEST_EMBEDDING_Q4GSW_MODEL="${EMBEDDING_MODEL}" \ - WEBGPU_TEST_EMBEDDING_Q4GSW_INDICES="${EMBEDDING_INDICES}" \ - WEBGPU_TEST_EMBEDDING_Q4GSW_GOLDEN="${EMBEDDING_GOLDEN}" \ - WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_MODEL="${EMBEDDING_LLAMA1B_MODEL}" \ - WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_INDICES="${EMBEDDING_LLAMA1B_INDICES}" \ - WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_GOLDEN="${EMBEDDING_LLAMA1B_GOLDEN}" \ - WEBGPU_TEST_ROPE_MODEL="${ROPE_MODEL}" \ - WEBGPU_TEST_ROPE_XQ_GOLDEN="${ROPE_XQ_GOLDEN}" \ - WEBGPU_TEST_ROPE_XK_GOLDEN="${ROPE_XK_GOLDEN}" \ - WEBGPU_TEST_ROPE_DECODE_MODEL="${ROPE_DECODE_MODEL}" \ - WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN="${ROPE_DECODE_XQ_GOLDEN}" \ - WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN="${ROPE_DECODE_XK_GOLDEN}" \ - WEBGPU_TEST_PREPACK_MODEL="${PREPACK_MODEL}" \ - WEBGPU_TEST_PREPACK_GOLDEN="${PREPACK_GOLDEN}" \ - WEBGPU_TEST_PREPACK2_MODEL="${PREPACK2_MODEL}" \ - WEBGPU_TEST_PREPACK2_GOLDEN="${PREPACK2_GOLDEN}" \ - WEBGPU_TEST_PREPACK_TIED_MODEL="${PREPACK_TIED_MODEL}" \ - WEBGPU_TEST_PREPACK_TIED_GOLDEN="${PREPACK_TIED_GOLDEN}" \ - "${BIN_DIR}/webgpu_native_test" -else - echo "(skipping webgpu_native_test: executorch wheel absent — exports did not run)" -fi -if [[ "${UPDATE_CACHE_OK}" == "1" && -x "${BIN_DIR}/webgpu_update_cache_test" ]]; then - "${BIN_DIR}/webgpu_update_cache_test" "${UPDATE_CACHE_DIR}" -fi -if [[ "${DISPATCH_ORDER_OK}" == "1" && -x "${BIN_DIR}/webgpu_dispatch_order_test" ]]; then - "${BIN_DIR}/webgpu_dispatch_order_test" "${DISPATCH_ORDER_DIR}" -fi -if [[ "${INDEX_OK}" == "1" && -x "${BIN_DIR}/webgpu_index_test" ]]; then - "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}" -fi -if [[ "${DYNAMIC_SHAPE_OK}" == "1" && -x "${BIN_DIR}/webgpu_dynamic_shape_test" ]]; then - "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}" -fi -[[ -x "${BIN_DIR}/webgpu_scratch_buffer_test" ]] && "${BIN_DIR}/webgpu_scratch_buffer_test" -# Device-free: pure 2D workgroup-count fold unit test (no .pte, no GPU). -[[ -x "${BIN_DIR}/webgpu_dispatch_2d_test" ]] && "${BIN_DIR}/webgpu_dispatch_2d_test" +run_with_required_device env WEBGPU_TEST_SDPA_DIR=/tmp/ \ + WEBGPU_TEST_QUANTIZED_LINEAR_DIR=/tmp/ \ + WEBGPU_TEST_EMBEDDING_Q4GSW_MODEL="${EMBEDDING_MODEL}" \ + WEBGPU_TEST_EMBEDDING_Q4GSW_INDICES="${EMBEDDING_INDICES}" \ + WEBGPU_TEST_EMBEDDING_Q4GSW_GOLDEN="${EMBEDDING_GOLDEN}" \ + WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_MODEL="${EMBEDDING_LLAMA1B_MODEL}" \ + WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_INDICES="${EMBEDDING_LLAMA1B_INDICES}" \ + WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_GOLDEN="${EMBEDDING_LLAMA1B_GOLDEN}" \ + WEBGPU_TEST_ROPE_MODEL="${ROPE_MODEL}" \ + WEBGPU_TEST_ROPE_XQ_GOLDEN="${ROPE_XQ_GOLDEN}" \ + WEBGPU_TEST_ROPE_XK_GOLDEN="${ROPE_XK_GOLDEN}" \ + WEBGPU_TEST_ROPE_DECODE_MODEL="${ROPE_DECODE_MODEL}" \ + WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN="${ROPE_DECODE_XQ_GOLDEN}" \ + WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN="${ROPE_DECODE_XK_GOLDEN}" \ + WEBGPU_TEST_SYMINT_BLOB="${SYMINT_BLOB}" \ + WEBGPU_TEST_PREPACK_MODEL="${PREPACK_MODEL}" \ + WEBGPU_TEST_PREPACK_GOLDEN="${PREPACK_GOLDEN}" \ + WEBGPU_TEST_PREPACK2_MODEL="${PREPACK2_MODEL}" \ + WEBGPU_TEST_PREPACK2_GOLDEN="${PREPACK2_GOLDEN}" \ + WEBGPU_TEST_PREPACK_TIED_MODEL="${PREPACK_TIED_MODEL}" \ + WEBGPU_TEST_PREPACK_TIED_GOLDEN="${PREPACK_TIED_GOLDEN}" \ + "${BIN_DIR}/webgpu_native_test" +"${BIN_DIR}/webgpu_update_cache_test" "${UPDATE_CACHE_DIR}" +"${BIN_DIR}/webgpu_dispatch_order_test" "${DISPATCH_ORDER_DIR}" +"${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}" +"${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}" +"${BIN_DIR}/webgpu_scratch_buffer_test" +"${BIN_DIR}/webgpu_dispatch_2d_test" +"${BIN_DIR}/webgpu_compute_dispatch_test" +"${BIN_DIR}/webgpu_execution_options_test" +"${BIN_DIR}/webgpu_output_suppression_test" "${OUTPUT_SUPPRESSION_DIR}" +"${BIN_DIR}/webgpu_op_test_util_test" echo "=== WebGPU native tests on Dawn: all run targets passed ===" # ── Op-test codegen framework: generate manifest → build → run (Dawn+SwiftShader) ── -# Reconfigure the SAME build dir adding GTest (EXECUTORCH_BUILD_TESTS=ON), then run -# every op in cases.py against its torch golden. Self-skips if the generator can't run. +# Generate the op-test manifest, build the target from the existing test-enabled +# configuration, then run every op in cases.py against its torch golden. OP_TEST_DIR="/tmp/webgpu_op_tests" -if $PYTHON_EXECUTABLE -m executorch.backends.webgpu.test.op_tests.generate_op_tests \ - --output "${OP_TEST_DIR}"; then - echo "=== Reconfigure with GTest + build/run op-test framework ===" - cmake -DEXECUTORCH_BUILD_TESTS=ON -B "${BUILD_DIR}" "${EXECUTORCH_ROOT}" - OP_DEFINED="$(cmake --build "${BUILD_DIR}" --target help 2>/dev/null || true)" - if printf '%s\n' "${OP_DEFINED}" | grep -qw webgpu_op_test_util_test; then - cmake --build "${BUILD_DIR}" --target webgpu_op_test_util_test -j"${NPROC}" - "${BIN_DIR}/webgpu_op_test_util_test" - fi - if printf '%s\n' "${OP_DEFINED}" | grep -qw webgpu_op_test; then - cmake --build "${BUILD_DIR}" --target webgpu_op_test -j"${NPROC}" - "${BIN_DIR}/webgpu_op_test" --manifest "${OP_TEST_DIR}/manifest.json" - fi - echo "=== WebGPU op-test framework on Dawn: passed ===" -else - echo "WARN: op-test manifest generation failed (needs the executorch wheel); skipping" -fi +$PYTHON_EXECUTABLE -m executorch.backends.webgpu.test.op_tests.generate_op_tests \ + --output "${OP_TEST_DIR}" +cmake --build "${BUILD_DIR}" --target webgpu_op_test -j"${NPROC}" +"${BIN_DIR}/webgpu_op_test" --manifest "${OP_TEST_DIR}/manifest.json" +echo "=== WebGPU op-test framework on Dawn: passed ===" diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index 35032bec983..0963168ccd5 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -6,12 +6,17 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include #include +#include +#include +#include #include +#include #include #include #include @@ -23,9 +28,11 @@ namespace { WGPUDevice g_device = nullptr; // NOLINTEND(cppcoreguidelines-avoid-non-const-global-variables) -struct SigmoidParams { +struct UnaryParams { uint32_t num_elements; - uint32_t padding[3]; + float min; + float max; + uint32_t padding; }; WGPUBuffer create_storage_buffer(size_t nbytes) { @@ -39,6 +46,110 @@ WGPUBuffer create_storage_buffer(size_t nbytes) { return buffer; } +enum class Q4RouteSignal { Static, LegacyGraphMarker, ExplicitOption }; + +void build_q4_route_graph(WebGPUGraph& graph, Q4RouteSignal signal) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](vk::VkDataType dtype, + const std::vector& dims, + int mem_obj_id) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, dtype, &dims, /*constant_id=*/-1, mem_obj_id) + .Union())); + return id; + }; + auto add_int = [&](int64_t value) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, value).Union())); + return id; + }; + + const int input = add_tensor(vk::VkDataType::FLOAT32, {2, 8}, 0); + const int weight = add_tensor(vk::VkDataType::UINT8, {8, 4}, 1); + const int scales = add_tensor(vk::VkDataType::FLOAT32, {1, 8}, 2); + const int group_size = add_int(8); + const int bias = static_cast(values.size()); + values.push_back(vk::CreateVkValue(fbb)); + const int output = add_tensor(vk::VkDataType::FLOAT32, {2, 8}, 3); + + std::vector<::flatbuffers::Offset> chain; + int q4_input = input; + if (signal == Q4RouteSignal::LegacyGraphMarker) { + const int dim = add_int(0); + const int symint = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::SymInt, vk::CreateSymInt(fbb, 0).Union())); + const int alpha = add_int(1); + const int intermediate = add_tensor(vk::VkDataType::FLOAT32, {2, 8}, 4); + const std::vector sym_size_args = {input, dim, symint}; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "sym_size.int", &sym_size_args)); + const std::vector add_args = {input, input, alpha, intermediate}; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 1, "aten.add.Tensor", &add_args)); + q4_input = intermediate; + } + + const std::vector q4_args = { + q4_input, weight, scales, group_size, bias, output}; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, + static_cast(chain.size()), + "et_vk.linear_q4gsw.default", + &q4_args)); + const std::vector input_ids = { + static_cast(input), + static_cast(weight), + static_cast(scales)}; + const std::vector output_ids = {static_cast(output)}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraphConfig config; + if (signal == Q4RouteSignal::ExplicitOption) { + config.record_q4gsw_decode_route = true; + } + graph.build(fbb.GetBufferPointer(), nullptr, nullptr, config); +} + +std::vector q4_dispatches(WebGPUGraph& graph) { + std::vector dispatches; + for (size_t i = 0; i < graph.num_dispatches(); i++) { + const WebGPUDispatch& dispatch = graph.dispatch_at(i); + if (dispatch.kernel_name.rfind("linear_q4gsw", 0) == 0) { + dispatches.push_back(&dispatch); + } + } + return dispatches; +} + +void expect_dual_q4_topology(WebGPUGraph& graph) { + const auto dispatches = q4_dispatches(graph); + ASSERT_EQ(dispatches.size(), 2); + EXPECT_NE(dispatches[0]->pipeline, nullptr); + EXPECT_NE(dispatches[1]->pipeline, nullptr); + EXPECT_NE(dispatches[0]->pipeline, dispatches[1]->pipeline); + EXPECT_NE(dispatches[0]->bind_group, nullptr); + EXPECT_EQ(dispatches[0]->bind_group, dispatches[1]->bind_group); + EXPECT_EQ( + std::count_if( + dispatches.begin(), + dispatches.end(), + [](const WebGPUDispatch* dispatch) { + return dispatch->workgroup_count_x != 0 && + dispatch->workgroup_count_y != 0; + }), + 1); +} + TEST(WebGPUShaderRegistry, FindsKnownShaderAndRejectsUnknownName) { const WebGPUShaderInfo& sigmoid = get_webgpu_shader_info("sigmoid"); EXPECT_EQ(sigmoid.name, "sigmoid"); @@ -48,6 +159,35 @@ TEST(WebGPUShaderRegistry, FindsKnownShaderAndRejectsUnknownName) { get_webgpu_shader_info("not_a_registered_shader"), std::runtime_error); } +TEST(WebGPUQ4RouteSignal, PreservesStaticAndRecordsBothDynamicSignals) { + WebGPUGraph static_graph; + static_graph.set_device(g_device); + build_q4_route_graph(static_graph, Q4RouteSignal::Static); + const auto static_dispatches = q4_dispatches(static_graph); + ASSERT_EQ(static_dispatches.size(), 1); + EXPECT_EQ(static_graph.num_dispatches(), 1); + EXPECT_NE(static_dispatches[0]->pipeline, nullptr); + EXPECT_NE(static_dispatches[0]->bind_group, nullptr); + EXPECT_FALSE(static_graph.has_dynamic_shapes()); + EXPECT_FALSE(static_graph.config().record_q4gsw_decode_route); + + WebGPUGraph legacy_graph; + legacy_graph.set_device(g_device); + build_q4_route_graph(legacy_graph, Q4RouteSignal::LegacyGraphMarker); + EXPECT_TRUE(legacy_graph.has_dynamic_shapes()); + EXPECT_FALSE(legacy_graph.config().record_q4gsw_decode_route); + EXPECT_EQ(legacy_graph.num_dispatches(), 3); + expect_dual_q4_topology(legacy_graph); + + WebGPUGraph explicit_graph; + explicit_graph.set_device(g_device); + build_q4_route_graph(explicit_graph, Q4RouteSignal::ExplicitOption); + EXPECT_FALSE(explicit_graph.has_dynamic_shapes()); + EXPECT_TRUE(explicit_graph.config().record_q4gsw_decode_route); + EXPECT_EQ(explicit_graph.num_dispatches(), 2); + expect_dual_q4_topology(explicit_graph); +} + TEST(WebGPUComputeDispatch, PipelineKeyCanonicalizesConstants) { WebGPUComputeDispatchDescriptor first; first.shader_name = "sigmoid"; @@ -139,32 +279,84 @@ TEST(WebGPUComputeDispatch, ReusesPipelineAndReleasesDawnObjects) { { WebGPUGraph graph; graph.set_device(g_device); - WGPUBuffer params = - graph.create_params_buffer(SigmoidParams{kNumElements, {0, 0, 0}}); + WGPUBuffer clamp_params = + graph.create_params_buffer(UnaryParams{kNumElements, -1.0f, 1.0f, 0}); + WGPUBuffer hardtanh_params = + graph.create_params_buffer(UnaryParams{kNumElements, -2.0f, 2.0f, 0}); WebGPUComputeDispatchDescriptor descriptor; - descriptor.shader_name = "sigmoid"; - descriptor.kernel_name = "sigmoid_test"; + descriptor.shader_name = "clamp"; + descriptor.kernel_name = "clamp_test"; descriptor.bindings = { {input, 0u, kBufferBytes}, {output, 0u, kBufferBytes}, - {params, 0u, sizeof(SigmoidParams)}}; + {clamp_params, 0u, sizeof(UnaryParams)}}; descriptor.constants = {{"wg_size", 64.0}}; - descriptor.grid = {1u, 1u}; + descriptor.grid = {7u, 3u}; - graph.add_compute_dispatch(descriptor); - graph.add_compute_dispatch(descriptor); + const size_t first = graph.add_compute_dispatch(descriptor); + descriptor.kernel_name = "hardtanh_test"; + descriptor.bindings[2].buffer = hardtanh_params; + const size_t second = graph.add_compute_dispatch(descriptor); const WebGPUMemoryStats stats = graph.memory_stats(); EXPECT_EQ(stats.num_dispatches, 2); EXPECT_EQ(stats.num_cached_shaders, 1); EXPECT_EQ(stats.num_cached_pipelines, 1); + EXPECT_EQ(stats.uniform_buffer_bytes, 2 * sizeof(UnaryParams)); + EXPECT_EQ( + graph.dispatch_at(first).pipeline, + graph.dispatch_at(second).pipeline); + EXPECT_EQ(graph.dispatch_at(first).kernel_name, "clamp_test"); + EXPECT_EQ(graph.dispatch_at(second).kernel_name, "hardtanh_test"); + EXPECT_EQ(graph.dispatch_at(first).workgroup_count_x, 7u); + EXPECT_EQ(graph.dispatch_at(first).workgroup_count_y, 3u); + EXPECT_EQ(graph.dispatch_at(second).workgroup_count_x, 7u); + EXPECT_EQ(graph.dispatch_at(second).workgroup_count_y, 3u); } wgpuBufferRelease(output); wgpuBufferRelease(input); } } +TEST(WebGPUComputeDispatch, AlternateShaderReusesLayoutAndBindGroup) { + constexpr size_t kNumElements = 64; + constexpr size_t kBufferBytes = kNumElements * sizeof(float); + WGPUBuffer input = create_storage_buffer(kBufferBytes); + WGPUBuffer output = create_storage_buffer(kBufferBytes); + { + WebGPUGraph graph; + graph.set_device(g_device); + WGPUBuffer params = + graph.create_params_buffer(UnaryParams{kNumElements, 0.0f, 0.0f, 0}); + const std::vector bindings = { + {0, WGPUBufferBindingType_ReadOnlyStorage, input, kBufferBytes}, + {1, WGPUBufferBindingType_Storage, output, kBufferBytes}, + {2, WGPUBufferBindingType_Uniform, params, sizeof(UnaryParams)}, + }; + const WGPUConstantEntry wg_size = utils::make_wg_size_constant(64u); + { + utils::ComputePipelineBundle sigmoid = utils::make_compute_pipeline( + g_device, kSigmoidWGSL, bindings, &wg_size, 1u); + utils::ComputePipelineBundle relu = utils::make_compute_pipeline( + g_device, kReluWGSL, sigmoid, &wg_size, 1u); + + EXPECT_EQ(relu.bind_group, sigmoid.bind_group); + EXPECT_EQ(relu.bind_group_layout, nullptr); + EXPECT_EQ(relu.pipeline_layout, nullptr); + EXPECT_NE(relu.shader, nullptr); + EXPECT_NE(relu.pipeline, nullptr); + EXPECT_NE(relu.pipeline, sigmoid.pipeline); + + graph.add_dispatch( + {sigmoid.pipeline, sigmoid.bind_group, 1u, "sigmoid_test"}); + graph.add_dispatch({relu.pipeline, relu.bind_group, 1u, "relu_test"}); + } + } + wgpuBufferRelease(output); + wgpuBufferRelease(input); +} + TEST(WebGPUComputeDispatch, RejectsBindingRangeBeyondDawnBuffer) { WGPUBuffer buffer = create_storage_buffer(16u); WebGPUComputeDispatchDescriptor descriptor; @@ -176,6 +368,291 @@ TEST(WebGPUComputeDispatch, RejectsBindingRangeBeyondDawnBuffer) { wgpuBufferRelease(buffer); } +constexpr int kResizeQ = 0; +constexpr int kResizeK = 1; +constexpr int kResizeUnrelated = 2; +constexpr int kResizeCascade = 3; +constexpr int kResizeOutput = 4; +constexpr int kResizeSymInt = 5; + +void build_resize_test_graph(WebGPUGraph& graph) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + const std::vector dims = {8, 8}; + for (int mem_obj_id = 0; mem_obj_id < 5; mem_obj_id++) { + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + /*constant_id=*/-1, + mem_obj_id) + .Union())); + } + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::SymInt, vk::CreateSymInt(fbb, 0).Union())); + + const std::vector<::flatbuffers::Offset> chain; + const std::vector input_ids = { + kResizeQ, kResizeK, kResizeUnrelated}; + const std::vector output_ids = {kResizeOutput}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + graph.set_device(g_device); + graph.build(fbb.GetBufferPointer(), nullptr, nullptr); +} + +WebGPUComputeDispatchDescriptor make_dynamic_test_descriptor( + WebGPUGraph& graph, + const char* kernel_name) { + const auto& input = graph.get_tensor(kResizeQ); + const auto& output = graph.get_tensor(kResizeOutput); + WGPUBuffer params = + graph.create_params_buffer(UnaryParams{64u, -1.0f, 1.0f, 0u}); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = "clamp"; + descriptor.kernel_name = kernel_name; + descriptor.bindings = { + {input.buffer, 0u, input.nbytes}, + {output.buffer, 0u, output.nbytes}, + {params, 0u, sizeof(UnaryParams)}}; + descriptor.constants = {{"wg_size", 64.0}}; + descriptor.grid = {99u, 97u}; + return descriptor; +} + +struct ResizeProbeContext { + int marker; + int* observed; + int* calls; +}; + +void record_resize_probe(WebGPUGraph&, const ResizeProbeContext& context) { + *context.observed = context.marker; + ++*context.calls; +} + +struct GridPickerContext { + int tensor_id; + uint32_t x_bias; + uint32_t y_bias; + int* calls; + const bool* fail; +}; + +WebGPUDispatchGrid pick_tensor_grid( + const WebGPUGraph& graph, + const GridPickerContext& context) { + ++*context.calls; + if (context.fail != nullptr && *context.fail) { + throw std::runtime_error("dynamic grid picker failure"); + } + const auto& dims = graph.cur_dims(context.tensor_id); + return { + static_cast(dims.at(0)) + context.x_bias, + static_cast(dims.at(1)) + context.y_bias}; +} + +struct CascadeContext { + int source_id; + int output_id; +}; + +void resize_cascade_output(WebGPUGraph& graph, const CascadeContext& context) { + graph.set_cur_dims(context.output_id, graph.cur_dims(context.source_id)); +} + +void expect_dispatch_grid( + WebGPUGraph& graph, + size_t dispatch_index, + uint32_t x, + uint32_t y) { + const auto& dispatch = graph.dispatch_at(dispatch_index); + EXPECT_EQ(dispatch.workgroup_count_x, x); + EXPECT_EQ(dispatch.workgroup_count_y, y); +} + +TEST(WebGPUResizeHooks, TypedRegistrationOwnsContextCopies) { + WebGPUGraph graph; + build_resize_test_graph(graph); + int tensor_observed = 0; + int tensor_calls = 0; + int symint_observed = 0; + int symint_calls = 0; + { + ResizeProbeContext tensor_context = {17, &tensor_observed, &tensor_calls}; + ResizeProbeContext symint_context = {23, &symint_observed, &symint_calls}; + graph.add_tensor_resize_hook(kResizeQ, record_resize_probe, tensor_context); + graph.add_resize_hook(kResizeSymInt, record_resize_probe, symint_context); + tensor_context.marker = 101; + symint_context.marker = 103; + } + + graph.resize_input(kResizeQ, {7, 8}); + graph.propagate_resize(); + EXPECT_EQ(tensor_observed, 17); + EXPECT_EQ(tensor_calls, 1); + EXPECT_EQ(symint_observed, 0); + EXPECT_EQ(symint_calls, 0); + + graph.set_symint(kResizeSymInt, 9); + graph.propagate_resize(); + EXPECT_EQ(tensor_calls, 1); + EXPECT_EQ(symint_observed, 23); + EXPECT_EQ(symint_calls, 1); +} + +TEST(WebGPUDynamicDispatch, InitializesAndIsolatesTriggeredGrids) { + WebGPUGraph graph; + build_resize_test_graph(graph); + int q_calls = 0; + int k_calls = 0; + GridPickerContext q_context = {kResizeQ, 1u, 2u, &q_calls, nullptr}; + const GridPickerContext k_context = {kResizeK, 3u, 4u, &k_calls, nullptr}; + const size_t q_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_q"), + kResizeQ, + pick_tensor_grid, + q_context); + const size_t k_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_k"), + kResizeK, + pick_tensor_grid, + k_context); + q_context.x_bias = 101u; + q_context.y_bias = 103u; + expect_dispatch_grid(graph, q_dispatch, 9u, 10u); + expect_dispatch_grid(graph, k_dispatch, 11u, 12u); + EXPECT_EQ(q_calls, 1); + EXPECT_EQ(k_calls, 1); + q_calls = 0; + k_calls = 0; + + graph.resize_input(kResizeUnrelated, {7, 7}); + graph.propagate_resize(); + EXPECT_EQ(q_calls, 0); + EXPECT_EQ(k_calls, 0); + + graph.resize_input(kResizeQ, {4, 3}); + graph.propagate_resize(); + expect_dispatch_grid(graph, q_dispatch, 5u, 5u); + expect_dispatch_grid(graph, k_dispatch, 11u, 12u); + EXPECT_EQ(q_calls, 1); + EXPECT_EQ(k_calls, 0); + + graph.resize_input(kResizeQ, {2, 5}); + graph.propagate_resize(); + expect_dispatch_grid(graph, q_dispatch, 3u, 7u); + EXPECT_EQ(q_calls, 2); + graph.resize_input(kResizeQ, {2, 5}); + graph.propagate_resize(); + EXPECT_EQ(q_calls, 2); + + graph.resize_input(kResizeK, {6, 1}); + graph.propagate_resize(); + expect_dispatch_grid(graph, k_dispatch, 9u, 5u); + EXPECT_EQ(k_calls, 1); +} + +TEST(WebGPUDynamicDispatch, HandlesCascadesAndStagesPickerFailures) { + { + WebGPUGraph graph; + build_resize_test_graph(graph); + int same_pass_calls = 0; + int cascade_pass_calls = 0; + graph.add_tensor_resize_hook( + kResizeQ, + resize_cascade_output, + CascadeContext{kResizeQ, kResizeCascade}); + const size_t same_pass_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_same_pass"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeCascade, 5u, 7u, &same_pass_calls, nullptr}); + const size_t cascade_pass_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_cascade_pass"), + kResizeCascade, + pick_tensor_grid, + GridPickerContext{ + kResizeCascade, 9u, 11u, &cascade_pass_calls, nullptr}); + same_pass_calls = 0; + cascade_pass_calls = 0; + graph.resize_input(kResizeQ, {4, 6}); + graph.propagate_resize(); + EXPECT_EQ(same_pass_calls, 1); + EXPECT_EQ(cascade_pass_calls, 1); + expect_dispatch_grid(graph, same_pass_dispatch, 9u, 13u); + expect_dispatch_grid(graph, cascade_pass_dispatch, 13u, 17u); + } + + WebGPUGraph graph; + build_resize_test_graph(graph); + int first_calls = 0; + int second_calls = 0; + bool second_fails = false; + const size_t first = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_first"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 1u, 2u, &first_calls, nullptr}); + const size_t second = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_second"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 3u, 4u, &second_calls, &second_fails}); + expect_dispatch_grid(graph, first, 9u, 10u); + expect_dispatch_grid(graph, second, 11u, 12u); + second_fails = true; + graph.resize_input(kResizeQ, {4, 3}); + EXPECT_THROW(graph.propagate_resize(), std::runtime_error); + expect_dispatch_grid(graph, first, 9u, 10u); + expect_dispatch_grid(graph, second, 11u, 12u); +} + +TEST(WebGPUDynamicDispatch, RejectsRouteOverlapWithoutPoisoningRegistry) { + WebGPUGraph graph; + build_resize_test_graph(graph); + int calls = 0; + bool initial_fails = true; + const size_t dispatches_before_failure = graph.num_dispatches(); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_initial_failure"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, &initial_fails}), + std::runtime_error); + EXPECT_EQ(graph.num_dispatches(), dispatches_before_failure); + + const size_t dynamic = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_route_guard"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}); + ASSERT_EQ(dynamic, 0u); + graph.add_dispatch(WebGPUDispatch{}); + graph.add_dispatch(WebGPUDispatch{}); + + EXPECT_THROW( + graph.register_dispatch_route_group({{0, 1}, {1, 2}}), + std::runtime_error); + calls = 0; + graph.resize_input(kResizeQ, {7, 6}); + graph.propagate_resize(); + EXPECT_EQ(calls, 1); + expect_dispatch_grid(graph, dynamic, 7u, 6u); + const size_t group = graph.register_dispatch_route_group({{1, 2}, {2, 3}}); + EXPECT_EQ(group, 0u); + graph.select_dispatch_route(group, 1, {{13u, 17u}}); + expect_dispatch_grid(graph, 1u, 0u, 0u); + expect_dispatch_grid(graph, 2u, 13u, 17u); +} + TEST(WebGPUExecution, FullySuppressedPlanPerformsNoQueueSubmission) { WebGPUGraph graph; const WebGPUExecutionPlan plan; diff --git a/backends/webgpu/test/native/test_dispatch_2d.cpp b/backends/webgpu/test/native/test_dispatch_2d.cpp index 582caa3c98e..4478c08dbdc 100644 --- a/backends/webgpu/test/native/test_dispatch_2d.cpp +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -9,6 +9,8 @@ // Device-free unit test for the pure 2D workgroup-count fold that lifts the // 65535 per-dim dispatch cap. Exercises the fold arithmetic only — no GPU. +#include +#include #include #include @@ -16,6 +18,10 @@ #include #include +using executorch::backends::webgpu::WebGPUDispatch; +using executorch::backends::webgpu::WebGPUGraph; +using executorch::backends::webgpu::utils::DispatchRange; +using executorch::backends::webgpu::utils::DispatchRouteRegistry; using executorch::backends::webgpu::utils::fold_workgroup_count_2d; using executorch::backends::webgpu::utils::WgCount; @@ -57,4 +63,213 @@ TEST(DispatchFold, ThrowsWhenNeeds3rdDimension) { EXPECT_ANY_THROW(fold_workgroup_count_2d(kMax * kMax + 1u, kMax, "test")); } +void expect_grid(const WgCount& grid, uint32_t x, uint32_t y) { + EXPECT_EQ(grid.x, x); + EXPECT_EQ(grid.y, y); +} + +void expect_grids_equal( + const std::vector& actual, + const std::vector& expected) { + ASSERT_EQ(actual.size(), expected.size()); + for (size_t i = 0; i < actual.size(); i++) { + EXPECT_EQ(actual[i].x, expected[i].x) << "index=" << i; + EXPECT_EQ(actual[i].y, expected[i].y) << "index=" << i; + } +} + +TEST(DispatchRoute, SwitchesUnequalRangesAndRestoresBothDimensions) { + std::vector grids(6, {101, 103}); + const std::vector compute(6, true); + DispatchRouteRegistry registry; + const size_t group = registry.register_group( + grids.size(), {{1, 3}, {4, 5}}, [&](size_t i) { return compute[i]; }); + auto set_grid = [&](size_t i, WgCount grid) { grids[i] = grid; }; + + registry.select(group, 0, {{3, 7}, {5, 11}}, set_grid); + expect_grid(grids[0], 101, 103); + expect_grid(grids[1], 3, 7); + expect_grid(grids[2], 5, 11); + expect_grid(grids[3], 101, 103); + expect_grid(grids[4], 0, 0); + expect_grid(grids[5], 101, 103); + + registry.select(group, 1, {{13, 17}}, set_grid); + expect_grid(grids[1], 0, 0); + expect_grid(grids[2], 0, 0); + expect_grid(grids[4], 13, 17); + + registry.select(group, 0, {{19, 23}, {29, 31}}, set_grid); + expect_grid(grids[1], 19, 23); + expect_grid(grids[2], 29, 31); + expect_grid(grids[4], 0, 0); +} + +TEST(DispatchRoute, InvalidSelectionDoesNotMutateGrids) { + std::vector grids = {{1, 2}, {3, 4}, {5, 6}}; + const std::vector original = grids; + DispatchRouteRegistry registry; + const size_t group = registry.register_group( + grids.size(), {{0, 2}, {2, 3}}, [](size_t) { return true; }); + auto set_grid = [&](size_t i, WgCount grid) { grids[i] = grid; }; + + EXPECT_ANY_THROW(registry.select(group, 2, {{7, 8}}, set_grid)); + expect_grids_equal(grids, original); + EXPECT_ANY_THROW(registry.select(group, 0, {{7, 8}}, set_grid)); + expect_grids_equal(grids, original); + EXPECT_ANY_THROW(registry.select(group, 0, {{0, 8}, {9, 10}}, set_grid)); + expect_grids_equal(grids, original); + EXPECT_ANY_THROW(registry.select(group, 0, {{7, 0}, {9, 10}}, set_grid)); + expect_grids_equal(grids, original); +} + +TEST(DispatchRoute, RejectsInvalidAndMultiplyOwnedRanges) { + const auto all_compute = [](size_t) { return true; }; + DispatchRouteRegistry registry; + // Each pairs the range under test with a valid second range so the group + // clears the size < 2 short-circuit and the range-validation loop is actually + // exercised: {2,1} inverted, {1,1} empty, {0,5} end past dispatch_count. + EXPECT_ANY_THROW(registry.register_group(4, {{2, 1}, {3, 4}}, all_compute)); + EXPECT_ANY_THROW(registry.register_group(4, {{1, 1}, {3, 4}}, all_compute)); + EXPECT_ANY_THROW(registry.register_group(4, {{0, 5}, {3, 4}}, all_compute)); + EXPECT_ANY_THROW(registry.register_group(4, {{0, 2}, {1, 3}}, all_compute)); + + const size_t first = + registry.register_group(4, {{0, 1}, {1, 2}}, all_compute); + EXPECT_EQ(first, 0); + EXPECT_ANY_THROW(registry.register_group(4, {{1, 3}}, all_compute)); + + DispatchRouteRegistry copy_registry; + const std::vector compute = {true, false, true}; + EXPECT_ANY_THROW(copy_registry.register_group( + compute.size(), {{0, 1}, {1, 2}}, [&](size_t i) { return compute[i]; })); +} + +TEST(DispatchRoute, GraphRejectsCopyAndCrossGroupOwnership) { + WebGPUGraph graph; + for (size_t i = 0; i < 6; i++) { + graph.add_dispatch(WebGPUDispatch{}); + } + graph.add_buffer_copy(nullptr, nullptr, 0); + + const size_t group = graph.register_dispatch_route_group({{0, 1}, {1, 2}}); + EXPECT_EQ(group, 0); + EXPECT_ANY_THROW(graph.register_dispatch_route_group({{1, 2}, {2, 3}})); + EXPECT_ANY_THROW(graph.register_dispatch_route_group({{5, 6}, {6, 7}})); + + const size_t second = graph.register_dispatch_route_group({{2, 4}, {4, 5}}); + EXPECT_EQ(second, 1); + + graph.select_dispatch_route(group, 0, {{7, 11}}); + graph.select_dispatch_route(second, 0, {{13, 17}, {19, 23}}); + expect_grid( + {graph.dispatch_at(0).workgroup_count_x, + graph.dispatch_at(0).workgroup_count_y}, + 7, + 11); + expect_grid( + {graph.dispatch_at(1).workgroup_count_x, + graph.dispatch_at(1).workgroup_count_y}, + 0, + 0); + expect_grid( + {graph.dispatch_at(2).workgroup_count_x, + graph.dispatch_at(2).workgroup_count_y}, + 13, + 17); + expect_grid( + {graph.dispatch_at(3).workgroup_count_x, + graph.dispatch_at(3).workgroup_count_y}, + 19, + 23); + expect_grid( + {graph.dispatch_at(4).workgroup_count_x, + graph.dispatch_at(4).workgroup_count_y}, + 0, + 0); + + graph.select_dispatch_route(group, 1, {{29, 31}}); + expect_grid( + {graph.dispatch_at(2).workgroup_count_x, + graph.dispatch_at(2).workgroup_count_y}, + 13, + 17); +} + +TEST(DispatchRoute, ExecuteRejectsHalfZeroGrid) { + WebGPUGraph graph; + WebGPUDispatch dispatch; + dispatch.workgroup_count_x = 0; + dispatch.workgroup_count_y = 1; + graph.add_dispatch(dispatch); + EXPECT_ANY_THROW(graph.make_execution_plan({})); +} + +TEST(DispatchRoute, RecordsEligibleQ4AndDynamicSdpaAlternates) { + using executorch::backends::webgpu::utils::should_record_q4gsw_dual_route; + using executorch::backends::webgpu::utils::should_record_sdpa_dual_route; + + EXPECT_FALSE(should_record_q4gsw_dual_route(1, true, true, true)); + EXPECT_FALSE(should_record_q4gsw_dual_route(32, false, true, true)); + EXPECT_FALSE(should_record_q4gsw_dual_route(32, true, false, false)); + EXPECT_TRUE(should_record_q4gsw_dual_route(32, true, true, false)); + EXPECT_TRUE(should_record_q4gsw_dual_route(32, true, false, true)); + + EXPECT_FALSE(should_record_sdpa_dual_route(true, false)); + EXPECT_FALSE(should_record_sdpa_dual_route(false, true)); + EXPECT_TRUE(should_record_sdpa_dual_route(true, true)); +} + +TEST(WebGPUGraphConfig, ParsesExactBooleanCompileOption) { + using executorch::backends::webgpu::parse_webgpu_graph_config; + using executorch::runtime::CompileSpec; + + auto absent = parse_webgpu_graph_config({}); + ASSERT_TRUE(absent.ok()); + EXPECT_FALSE(absent->record_q4gsw_decode_route); + EXPECT_FALSE(absent->f16_kv_cache); + EXPECT_FALSE(absent->f16_accumulate_gemm); + + uint8_t false_value = 0; + CompileSpec false_spec = { + "webgpu_record_q4gsw_decode_route", {&false_value, 1}}; + auto parsed_false = parse_webgpu_graph_config(false_spec); + ASSERT_TRUE(parsed_false.ok()); + EXPECT_FALSE(parsed_false->record_q4gsw_decode_route); + + uint8_t true_value = 1; + CompileSpec true_spec = { + "webgpu_record_q4gsw_decode_route", {&true_value, 1}}; + auto parsed_true = parse_webgpu_graph_config(true_spec); + ASSERT_TRUE(parsed_true.ok()); + EXPECT_TRUE(parsed_true->record_q4gsw_decode_route); + + CompileSpec unknown_spec = {"unknown_webgpu_option", {nullptr, 0}}; + auto parsed_unknown = parse_webgpu_graph_config(unknown_spec); + ASSERT_TRUE(parsed_unknown.ok()); + EXPECT_FALSE(parsed_unknown->record_q4gsw_decode_route); +} + +TEST(WebGPUGraphConfig, RejectsMalformedRouteCompileOption) { + using executorch::backends::webgpu::parse_webgpu_graph_config; + using executorch::runtime::CompileSpec; + using executorch::runtime::Error; + + uint8_t values[2] = {0, 1}; + CompileSpec empty = {"webgpu_record_q4gsw_decode_route", {values, 0}}; + auto parsed_empty = parse_webgpu_graph_config(empty); + ASSERT_FALSE(parsed_empty.ok()); + EXPECT_EQ(parsed_empty.error(), Error::DelegateInvalidCompatibility); + + CompileSpec oversized = {"webgpu_record_q4gsw_decode_route", {values, 2}}; + auto parsed_oversized = parse_webgpu_graph_config(oversized); + ASSERT_FALSE(parsed_oversized.ok()); + EXPECT_EQ(parsed_oversized.error(), Error::DelegateInvalidCompatibility); + + CompileSpec null_value = {"webgpu_record_q4gsw_decode_route", {nullptr, 1}}; + auto parsed_null = parse_webgpu_graph_config(null_value); + ASSERT_FALSE(parsed_null.ok()); + EXPECT_EQ(parsed_null.error(), Error::DelegateInvalidCompatibility); +} + } // namespace diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index 3752e4ee53f..e271d53b9be 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -24,11 +24,13 @@ // /tmp/dynamic_shape. #include +#include #include #include #include +#include #include #include #include @@ -49,6 +51,26 @@ constexpr int kHidden = 64; // Artifacts directory; set from env/argv in main() before RUN_ALL_TESTS(). std::string g_dir; // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) +#ifdef WGPU_BACKEND_ENABLE_PROFILING +std::vector current_profile_names() { + const auto* context = get_default_webgpu_context(); + if (context == nullptr || context->querypool == nullptr) { + return {}; + } + std::vector names; + for (const auto& duration : context->querypool->results()) { + names.push_back(duration.kernel_name); + } + return names; +} + +bool contains_name( + const std::vector& names, + const std::string& expected) { + return std::find(names.begin(), names.end(), expected) != names.end(); +} +#endif + std::vector read_bin(const std::string& path) { std::ifstream f(path, std::ios::binary | std::ios::ate); if (!f) { @@ -105,16 +127,22 @@ void check_s(Module& m, const std::string& prefix, int s) { // Dynamic quantized linear: input [M, kLinK] -> output [M, n]. kLinN is the // register-tiled/bicol config; kLinNShmem (N>=2048) routes to the shmem GEMM. constexpr int kLinK = 64; +constexpr int kLinAltK = 72; constexpr int kLinN = 128; constexpr int kLinNShmem = 2048; // Run at [m_rows, kLinK] on an already-loaded module (so it can be // reused across M without a fresh load), and compare to the golden. -void run_linear(Module& m, int m_rows, const char* prefix, int n) { +void run_linear( + Module& m, + int m_rows, + const char* prefix, + int n, + int k = kLinK) { const std::string base = g_dir + "/" + prefix + ".S" + std::to_string(m_rows); auto input = read_bin(base + ".input.bin"); auto golden = read_bin(base + ".golden.bin"); ASSERT_FALSE(input.empty()) << "missing " << prefix << ".S" << m_rows; - auto t = make_tensor_ptr({m_rows, kLinK}, std::move(input)); + auto t = make_tensor_ptr({m_rows, k}, std::move(input)); auto r = m.forward({EValue(t)}); ASSERT_TRUE(r.ok() && !r.get().empty() && r.get()[0].isTensor()) << prefix << " M=" << m_rows << " forward failed"; @@ -138,16 +166,60 @@ void check_linear(int m_rows) { void check_linear_shmem(int m_rows) { Module m(g_dir + "/dyn_linear_shmem.pte"); ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_shmem.pte"; - run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem); + run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem, kLinAltK); +} + +void check_linear_tiled(int m_rows) { + Module m(g_dir + "/dyn_linear_tiled.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_tiled.pte"; + run_linear(m, m_rows, "dyn_linear_tiled", kLinN, kLinAltK); +} + +constexpr int kQkvNq = 2048; +constexpr int kQkvNk = 512; +constexpr int kQkvNv = 512; +constexpr int kQkvMaxM = 16; + +void run_qkv_routes(Module& m, int m_rows) { + const std::string base = + g_dir + "/qkv_routes.S" + std::to_string(m_rows) + "."; + auto input = read_bin(base + "input.bin"); + ASSERT_FALSE(input.empty()) << "missing qkv_routes.S" << m_rows; + auto tensor = make_tensor_ptr({m_rows, kLinK}, std::move(input)); + auto result = m.forward({EValue(tensor)}); + ASSERT_TRUE(result.ok()) << "qkv_routes M=" << m_rows << " forward failed"; + ASSERT_EQ(result.get().size(), 3u); + + const int widths[] = {kQkvNq, kQkvNk, kQkvNv}; + const char* names[] = {"q", "k", "v"}; + for (size_t i = 0; i < 3; i++) { + ASSERT_TRUE(result.get()[i].isTensor()); + const auto& output = result.get()[i].toTensor(); + ASSERT_EQ(output.dim(), 2); + ASSERT_EQ(output.size(0), m_rows); + ASSERT_EQ(output.size(1), widths[i]); + const size_t numel = static_cast(m_rows) * widths[i]; + std::vector got( + output.const_data_ptr(), output.const_data_ptr() + numel); + auto golden = read_bin(base + names[i] + ".bin"); + ASSERT_EQ(golden.size(), numel); + EXPECT_LT(max_err(got, golden), 1e-2f) + << "qkv_routes " << names[i] << " M=" << m_rows; + } } // Dynamic SDPA (GQA prefill, input_pos=0): q[1,s,hq,d] k/v[1,s,hkv,d] // caches[1,cmax,hkv,d]; attn output [1,s,hq,d] selected by shape (3 outputs). constexpr int kSdHq = 8, kSdHkv = 2, kSdD = 16, kSdCmax = 64; -void check_sdpa(int s) { - Module m(g_dir + "/sdpa_dyn.pte"); - ASSERT_EQ(m.load_forward(), Error::Ok) << "sdpa_dyn S=" << s << " load"; - const std::string b = g_dir + "/sdpa_dyn.S" + std::to_string(s) + "."; +void run_sdpa_case( + Module& m, + int s, + const char* prefix, + int hq, + int hkv, + int d, + int cmax) { + const std::string b = g_dir + "/" + prefix + ".S" + std::to_string(s) + "."; auto q = read_bin(b + "q.bin"); auto k = read_bin(b + "k.bin"); auto v = read_bin(b + "v.bin"); @@ -158,36 +230,92 @@ void check_sdpa(int s) { q.empty() || k.empty() || v.empty() || kc.empty() || vc.empty() || golden.empty()) << "missing sdpa_dyn.S" << s; - auto tq = make_tensor_ptr({1, s, kSdHq, kSdD}, std::move(q)); - auto tk = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(k)); - auto tv = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(v)); - auto tkc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(kc)); - auto tvc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(vc)); + auto tq = make_tensor_ptr({1, s, hq, d}, std::move(q)); + auto tk = make_tensor_ptr({1, s, hkv, d}, std::move(k)); + auto tv = make_tensor_ptr({1, s, hkv, d}, std::move(v)); + auto tkc = make_tensor_ptr({1, cmax, hkv, d}, std::move(kc)); + auto tvc = make_tensor_ptr({1, cmax, hkv, d}, std::move(vc)); auto r = m.forward({EValue(tq), EValue(tk), EValue(tv), EValue(tkc), EValue(tvc)}); ASSERT_TRUE(r.ok()) << "sdpa S=" << s << " forward failed (err=" << (int)r.error() << ")"; // Select the attn output by full shape [1,s,hq,d] (never numel). const float* attn = nullptr; - const size_t numel = static_cast(s) * kSdHq * kSdD; + const size_t numel = static_cast(s) * hq * d; for (size_t i = 0; i < r.get().size(); i++) { if (!r.get()[i].isTensor()) { continue; } const auto& t = r.get()[i].toTensor(); - if (t.dim() == 4 && t.size(1) == s && t.size(2) == kSdHq && - t.size(3) == kSdD) { + if (t.dim() == 4 && t.size(1) == s && t.size(2) == hq && t.size(3) == d) { attn = t.const_data_ptr(); break; } } ASSERT_NE(attn, nullptr) << "sdpa S=" << s << ": no attn output of shape [1," - << s << "," << kSdHq << "," << kSdD << "]"; + << s << "," << hq << "," << d << "]"; std::vector got(attn, attn + numel); const float e = max_err(got, golden); EXPECT_LT(e, 2e-3f) << "sdpa_dyn S=" << s << " max_err=" << e; } +void run_sdpa(Module& m, int s) { + run_sdpa_case(m, s, "sdpa_dyn", kSdHq, kSdHkv, kSdD, kSdCmax); +} + +void check_sdpa(int s) { + Module m(g_dir + "/sdpa_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "sdpa_dyn S=" << s << " load"; + run_sdpa(m, s); +} + +void run_combined_routes(Module& m, int s) { + const std::string b = g_dir + "/combined_routes.S" + std::to_string(s) + "."; + auto x = read_bin(b + "x.bin"); + auto q = read_bin(b + "q.bin"); + auto k = read_bin(b + "k.bin"); + auto v = read_bin(b + "v.bin"); + auto kc = read_bin(b + "kc.bin"); + auto vc = read_bin(b + "vc.bin"); + auto golden = read_bin(b + "golden.bin"); + ASSERT_FALSE( + x.empty() || q.empty() || k.empty() || v.empty() || kc.empty() || + vc.empty() || golden.empty()) + << "missing combined_routes.S" << s; + auto tx = make_tensor_ptr({s, kLinK}, std::move(x)); + auto tq = make_tensor_ptr({1, s, kSdHq, kSdD}, std::move(q)); + auto tk = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(k)); + auto tv = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(v)); + auto tkc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(kc)); + auto tvc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(vc)); + auto result = m.forward( + {EValue(tx), + EValue(tq), + EValue(tk), + EValue(tv), + EValue(tkc), + EValue(tvc)}); + ASSERT_TRUE(result.ok()) << "combined routes S=" << s + << " forward failed (err=" << (int)result.error() + << ")"; + const float* attn = nullptr; + const size_t numel = static_cast(s) * kSdHq * kSdD; + for (const auto& output : result.get()) { + if (!output.isTensor()) { + continue; + } + const auto& tensor = output.toTensor(); + if (tensor.dim() == 4 && tensor.size(1) == s && tensor.size(2) == kSdHq && + tensor.size(3) == kSdD) { + attn = tensor.const_data_ptr(); + break; + } + } + ASSERT_NE(attn, nullptr); + const std::vector got(attn, attn + numel); + EXPECT_LT(max_err(got, golden), 1e-2f) << "combined_routes S=" << s; +} + // Dynamic embedding: int64 token ids [N] -> [N, kEmbDim] fp32. The int64 host // input exercises copy_inputs' int64->int32 narrow path under dynamic shapes. constexpr int kEmbDim = 64; @@ -407,9 +535,7 @@ TEST(DynamicShape, QuantizedLinearReusedGraph) { } } -// I3: dynamic linear at N=2048 -> the shmem-GEMM route (K>=4096||N>=2048); the -// resize hook recomputes the shmem tile count for the live M on the fixed shmem -// pipeline (M=1 exercises a partial row-tile). +// I3: K=72 disables Steel; N=2048 forces shmem for M>1 and bicol for M=1. TEST(DynamicShape, QuantizedLinearShmem) { for (int m_rows : {128, 32, 1}) { check_linear_shmem(m_rows); @@ -421,26 +547,205 @@ TEST(DynamicShape, QuantizedLinearShmemReusedGraph) { Module m(g_dir + "/dyn_linear_shmem.pte"); ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_shmem.pte"; for (int m_rows : {128, 32, 1, 128}) { - run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem); + run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem, kLinAltK); } } -// J: dynamic SDPA (GQA prefill) at several seq-len S. The whole case skips -// while op coverage is pending (the dynamic-S build throws err 48 until -// registered). -TEST(DynamicShape, Sdpa) { - { - Module probe(g_dir + "/sdpa_dyn.pte"); - if (probe.load_forward() == Error::DelegateInvalidCompatibility) { - GTEST_SKIP() << "sdpa_dyn pending op coverage (err " - << (int)Error::DelegateInvalidCompatibility << ")"; - } +// I5: K=72 disables Steel; N=128 keeps the tiled fallback for M>1. +TEST(DynamicShape, QuantizedLinearTiled) { + for (int m_rows : {128, 32, 1}) { + check_linear_tiled(m_rows); + } +} + +TEST(DynamicShape, QuantizedLinearTiledReusedGraph) { + Module m(g_dir + "/dyn_linear_tiled.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_tiled.pte"; + for (int m_rows : {128, 1, 32, 1, 128}) { + run_linear(m, m_rows, "dyn_linear_tiled", kLinN, kLinAltK); + } +} + +// The first max-shape execution does not invoke resize hooks. This sequence +// therefore checks both initial QKV route selection and max/decode transitions. +TEST(DynamicShape, QkvRoutesReusedGraph) { + Module m(g_dir + "/qkv_routes.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load qkv_routes.pte"; + for (int m_rows : {kQkvMaxM, 1, kQkvMaxM}) { + run_qkv_routes(m, m_rows); + } +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +TEST(DynamicShape, QkvLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported || !context->shader_f16_supported) { + GTEST_SKIP() << "timestamp queries or shader-f16 unavailable"; + } + Module m(g_dir + "/qkv_routes.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load qkv_routes.pte"; + for (int m_rows : {kQkvMaxM, 1, kQkvMaxM}) { + run_qkv_routes(m, m_rows); + const auto names = current_profile_names(); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_qkv_fused"), + m_rows > 1 ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_coop4_bicol"), + m_rows == 1 ? 3 : 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_steel"), 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_shmem"), 0); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_tiled"), 0); + } +} + +TEST(DynamicShape, CombinedLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported) { + GTEST_SKIP() << "timestamp queries unavailable"; + } + Module m(g_dir + "/combined_routes.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load combined_routes.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_combined_routes(m, s); + const auto names = current_profile_names(); + ASSERT_EQ(names.size(), s == 1 ? 6 : 7); + EXPECT_EQ(std::count(names.begin(), names.end(), "add"), 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_coop4_bicol"), + s == 1 ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_steel"), + s != 1 ? 1 : 0); + EXPECT_FALSE(contains_name(names, "linear_q4gsw_shmem")); + EXPECT_FALSE(contains_name(names, "linear_q4gsw_tiled")); + EXPECT_EQ(std::count(names.begin(), names.end(), "update_cache"), 2); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_split"), s == 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_reduce"), s == 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), + s != 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_softmax"), s != 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_out"), s != 1); } +} + +TEST(DynamicShape, StaticRouteProfiles) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported) { + GTEST_SKIP() << "timestamp queries unavailable"; + } + + Module linear_m1(g_dir + "/static_linear_m1.pte"); + ASSERT_EQ(linear_m1.load_forward(), Error::Ok); + run_linear(linear_m1, 1, "static_linear_m1", kLinN); + auto names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_coop4_bicol"), 1); + + Module linear_m32(g_dir + "/static_linear_m32.pte"); + ASSERT_EQ(linear_m32.load_forward(), Error::Ok); + run_linear(linear_m32, 32, "static_linear_m32", kLinN); + names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_steel"), 1); + + Module linear_shmem(g_dir + "/dyn_linear_shmem.pte"); + ASSERT_EQ(linear_shmem.load_forward(), Error::Ok); + run_linear(linear_shmem, 128, "dyn_linear_shmem", kLinNShmem, kLinAltK); + names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_shmem"), 1); + + Module linear_tiled(g_dir + "/dyn_linear_tiled.pte"); + ASSERT_EQ(linear_tiled.load_forward(), Error::Ok); + run_linear(linear_tiled, 128, "dyn_linear_tiled", kLinN, kLinAltK); + names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_tiled"), 1); + + Module sdpa_s1(g_dir + "/static_sdpa_s1.pte"); + ASSERT_EQ(sdpa_s1.load_forward(), Error::Ok); + run_sdpa_case(sdpa_s1, 1, "static_sdpa_s1", kSdHq, kSdHkv, kSdD, kSdCmax); + names = current_profile_names(); + ASSERT_EQ(names.size(), 4); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_split"), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_reduce"), 1); + + Module sdpa_s16(g_dir + "/static_sdpa_s16.pte"); + ASSERT_EQ(sdpa_s16.load_forward(), Error::Ok); + run_sdpa_case(sdpa_s16, 16, "static_sdpa_s16", kSdHq, kSdHkv, kSdD, kSdCmax); + names = current_profile_names(); + ASSERT_EQ(names.size(), 5); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_softmax"), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_compute_out"), 1); +} +#endif + +// J: dynamic SDPA reuses one graph across prefill and FlashDecoding shapes. +TEST(DynamicShape, Sdpa) { for (int s : {64, 16, 1}) { check_sdpa(s); } } +TEST(DynamicShape, SdpaReusedGraph) { + Module m(g_dir + "/sdpa_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load sdpa_dyn.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_sdpa(m, s); + } +} + +TEST(DynamicShape, CombinedLiveRoutes) { + Module m(g_dir + "/combined_routes.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load combined_routes.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_combined_routes(m, s); + } +} + +TEST(DynamicShape, SdpaWideMaterializedOnly) { + Module m(g_dir + "/sdpa_wide.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load sdpa_wide.pte"; + for (int s : {16, 1, 16}) { + run_sdpa_case(m, s, "sdpa_wide", 8, 2, 132, 16); + } +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +TEST(DynamicShape, SdpaLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported) { + GTEST_SKIP() << "timestamp queries unavailable"; + } + Module m(g_dir + "/sdpa_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load sdpa_dyn.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_sdpa(m, s); + const auto names = current_profile_names(); + ASSERT_EQ(names.size(), s == 1 ? 4 : 5); + EXPECT_EQ(std::count(names.begin(), names.end(), "update_cache"), 2); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_split"), s == 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_reduce"), s == 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), + s != 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_softmax"), s != 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_out"), s != 1); + } +} +#endif + // K: dynamic embedding (int64 token ids) at several token counts. TEST(DynamicShape, Embedding) { for (int n : {16, 8, 1}) { diff --git a/backends/webgpu/test/native/test_execution_options.cpp b/backends/webgpu/test/native/test_execution_options.cpp index 3b58c48534b..382bfbf7efa 100644 --- a/backends/webgpu/test/native/test_execution_options.cpp +++ b/backends/webgpu/test/native/test_execution_options.cpp @@ -140,5 +140,28 @@ TEST(WebGPUExecutionPlanTest, CopyOnlyPlanRetainsOneSubmissionChunk) { EXPECT_EQ(plan.copy_outputs, (std::vector{true})); } +TEST(WebGPUExecutionPlanTest, FiltersDisabledDispatchesAcrossChunks) { + const std::vector enabled = {true, false, true, false, true, true}; + const WebGPUExecutionPlan plan = plan_webgpu_execution( + 6, 1, ExecuteConfig{2, 1}, {}, WebGPUGraphExecutionOptions{}, enabled); + + EXPECT_EQ( + plan.dispatch_chunks, + (std::vector>{{0}, {2}, {4}, {5}})); + EXPECT_EQ(plan.copy_outputs, (std::vector{true})); +} + +TEST(WebGPUExecutionPlanTest, RejectsMismatchedEnabledDispatches) { + EXPECT_THROW( + plan_webgpu_execution( + 3, + 1, + ExecuteConfig{}, + {}, + WebGPUGraphExecutionOptions{}, + {true, false}), + std::runtime_error); +} + } // namespace } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index 19cb41451e0..b1705e733d9 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -19,9 +19,11 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes MAXS = 128 # upper bound for the dynamic seq-len dim (within the 1D dispatch cap) HIDDEN = 64 +_Q4_DECODE_ROUTE_SPEC = "webgpu_record_q4gsw_decode_route" def _rms(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: @@ -108,17 +110,51 @@ def _ramp(shape) -> torch.Tensor: return torch.linspace(-1.0, 1.0, n, dtype=torch.float32).reshape(shape) +def _lower_fully_delegated( + ep, + label: str, + compile_options=None, + *, + expect_q4_decode_route: bool = False, +): + edge = to_edge_transform_and_lower( + ep, + partitioner=[VulkanPartitioner(compile_options=compile_options)], + ) + graph = edge.exported_program().graph_module.graph + delegates = get_delegates(graph) + portable = get_non_lowered_nodes(graph) + if len(delegates) != 1: + raise RuntimeError(f"{label}: expected one delegate, got {len(delegates)}") + if portable: + raise RuntimeError(f"{label}: non-lowered nodes: {portable}") + et = edge.to_executorch() + delegate_ids = [ + delegate.id + for plan in et.executorch_program.execution_plan + for delegate in plan.delegates + ] + if delegate_ids != ["VulkanBackend"]: + raise RuntimeError(f"{label}: serialized delegates: {delegate_ids}") + route_values = [ + spec.value + for plan in et.executorch_program.execution_plan + for delegate in plan.delegates + for spec in delegate.compile_specs + if spec.key == _Q4_DECODE_ROUTE_SPEC + ] + expected_route_values = [b"\x01"] if expect_q4_decode_route else [] + if route_values != expected_route_values: + raise RuntimeError( + f"{label}: {_Q4_DECODE_ROUTE_SPEC}: " + f"expected {expected_route_values}, got {route_values}" + ) + return et + + def _export(model, example_inputs, dynamic_shapes, path: str) -> None: ep = torch.export.export(model, example_inputs, dynamic_shapes=dynamic_shapes) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - found = any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ) - assert found, f"Expected VulkanBackend delegate in {path}" + et = _lower_fully_delegated(ep, path) with open(path, "wb") as f: f.write(et.buffer) print(f"Exported {path}") @@ -186,10 +222,29 @@ def export_dynamic_shape_cases(out_dir: str) -> None: # 2d) 4-bit quantized linear with a DYNAMIC rows (M) dim — prefill GEMM # (register-tiled N=128) + a shmem-GEMM-routed variant (N=2048). _export_dynamic_linear(out_dir) - _export_dynamic_linear(out_dir, n=LIN_SHMEM_N, prefix="dyn_linear_shmem") + _export_dynamic_linear( + out_dir, + n=LIN_SHMEM_N, + prefix="dyn_linear_shmem", + k=LIN_ALT_K, + group=LIN_ALT_GROUP, + ) + _export_dynamic_linear( + out_dir, + prefix="dyn_linear_tiled", + k=LIN_ALT_K, + group=LIN_ALT_GROUP, + ) + _export_static_linear(out_dir, 1, "static_linear_m1") + _export_static_linear(out_dir, 32, "static_linear_m32") # 2e) Fused SDPA with a DYNAMIC seq-len S (prefill, input_pos=0). _export_dynamic_sdpa(out_dir) + _export_combined_routes(out_dir) + _export_dynamic_qkv_routes(out_dir) + _export_dynamic_sdpa_wide(out_dir) + _export_static_sdpa(out_dir, 1, "static_sdpa_s1") + _export_static_sdpa(out_dir, 16, "static_sdpa_s16") # 2f) 4-bit embedding with a DYNAMIC token count (int64 indices). _export_dynamic_embedding(out_dir) @@ -224,36 +279,35 @@ def export_dynamic_shape_cases(out_dir: str) -> None: # Quantized linear: K x N weight, dynamic rows M; input [M, K], output [M, N]. LIN_K = 64 LIN_N = 128 -LIN_SHMEM_N = 2048 # N>=2048 routes linear_q4gsw to the shmem-GEMM path +LIN_ALT_K = 72 # K%16 != 0 disables Steel while K%8 keeps bicol eligible. +LIN_ALT_GROUP = 24 +LIN_SHMEM_N = 2048 LIN_GROUP = 32 LIN_MAXM = 128 def _export_dynamic_linear( - out_dir: str, n: int = LIN_N, prefix: str = "dyn_linear" + out_dir: str, + n: int = LIN_N, + prefix: str = "dyn_linear", + k: int = LIN_K, + group: int = LIN_GROUP, ) -> None: - from executorch.backends.webgpu.test.ops.quantized_linear.test_quantized_linear import ( + from executorch.backends.webgpu.test.ops.test_quantized_linear import ( _fp64_golden, _make_quantized_model, ) - model = _make_quantized_model(LIN_K, n, LIN_GROUP) - x = _ramp((LIN_MAXM, LIN_K)) + model = _make_quantized_model(k, n, group) + x = _ramp((LIN_MAXM, k)) m_dim = torch.export.Dim("m", min=1, max=LIN_MAXM) ep = torch.export.export(model, (x,), dynamic_shapes=({0: m_dim},)) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "linear_q4gsw not delegated" + et = _lower_fully_delegated(ep, prefix) with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: f.write(et.buffer) print(f"Exported {prefix}.pte") for m in [LIN_MAXM, 32, 1]: - xm = _ramp((m, LIN_K)) + xm = _ramp((m, k)) g = _fp64_golden(model, xm).astype(" None: + from executorch.backends.webgpu.test.ops.test_quantized_linear import ( + _fp64_golden, + _make_quantized_model, + ) + + model = _make_quantized_model(LIN_K, LIN_N, LIN_GROUP) + x = _ramp((m, LIN_K)) + ep = torch.export.export(model, (x,)) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + x.detach().numpy().astype(" None: - from executorch.backends.webgpu.test.ops.sdpa.test_sdpa import ( + from executorch.backends.webgpu.test.ops.test_sdpa import ( _det_inputs, _golden, SdpaConfig, @@ -286,14 +361,7 @@ def cfg(s: int) -> "SdpaConfig": s_dim = torch.export.Dim("s", min=1, max=SD_MAXS) ds = ({1: s_dim}, {1: s_dim}, {1: s_dim}, None, None) ep = torch.export.export(model, (q, k, v, kc, vc), dynamic_shapes=ds) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "sdpa not delegated" + et = _lower_fully_delegated(ep, "sdpa_dyn") with open(os.path.join(out_dir, "sdpa_dyn.pte"), "wb") as f: f.write(et.buffer) print("Exported sdpa_dyn.pte") @@ -315,6 +383,199 @@ def cfg(s: int) -> "SdpaConfig": print(f" golden sdpa_dyn S={s} (golden shape {tuple(g.shape)})") +def _export_combined_routes(out_dir: str) -> None: + from executorch.backends.webgpu.test.ops.test_quantized_linear import ( + _make_quantized_model, + ) + from executorch.backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + ) + from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 + + class CombinedRoutes(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = _make_quantized_model(LIN_K, SD_HQ * SD_D, LIN_GROUP) + + def forward(self, x, q, k, v, k_cache, v_cache): + projected = self.linear(x).reshape(1, x.shape[0], SD_HQ, SD_D) + return torch.ops.llama.sdpa_with_kv_cache( + q + projected, + k, + v, + k_cache, + v_cache, + 0, + q.shape[1], + None, + 0.0, + True, + None, + ) + + model = CombinedRoutes().eval() + cfg = SdpaConfig("combined", SD_HQ, SD_HKV, SD_D, SD_MAXS, SD_CMAX, 0) + q, k, v, kc, vc = _det_inputs(cfg) + x = _ramp((SD_MAXS, LIN_K)) + s_dim = torch.export.Dim("s", min=1, max=SD_MAXS) + ep = torch.export.export( + model, + (x, q, k, v, kc, vc), + dynamic_shapes=( + {0: s_dim}, + {1: s_dim}, + {1: s_dim}, + {1: s_dim}, + None, + None, + ), + ) + et = _lower_fully_delegated(ep, "combined_routes") + with open(os.path.join(out_dir, "combined_routes.pte"), "wb") as f: + f.write(et.buffer) + print("Exported combined_routes.pte") + + for s in [SD_MAXS, 16, 1]: + live_cfg = SdpaConfig("combined", SD_HQ, SD_HKV, SD_D, s, SD_CMAX, 0) + q, k, v, kc, vc = _det_inputs(live_cfg) + x = _ramp((s, LIN_K)) + with torch.no_grad(): + projected = model.linear(x).reshape(1, s, SD_HQ, SD_D) + golden = _golden(live_cfg, q + projected, k, v, kc, vc) + base = os.path.join(out_dir, f"combined_routes.S{s}.") + for name, tensor in ( + ("x", x), + ("q", q), + ("k", k), + ("v", v), + ("kc", kc), + ("vc", vc), + ("golden", golden), + ): + tensor.detach().numpy().astype(" None: + from executorch.backends.webgpu.test.ops.test_quantized_linear import ( + _make_quantized_model, + ) + + class QkvRoutes(torch.nn.Module): + def __init__(self): + super().__init__() + self.q = _make_quantized_model(LIN_K, QKV_NQ, LIN_GROUP, seed=0) + self.k = _make_quantized_model(LIN_K, QKV_NK, LIN_GROUP, seed=1) + self.v = _make_quantized_model(LIN_K, QKV_NV, LIN_GROUP, seed=2) + + def forward(self, x): + # Keep the linears internal so graph-output copies do not capture + # the buffers that the runtime QKV pass later replaces. + return ( + torch.sigmoid(self.q(x)), + torch.sigmoid(self.k(x)), + torch.sigmoid(self.v(x)), + ) + + model = QkvRoutes().eval() + x = _ramp((QKV_MAXM, LIN_K)) + m_dim = torch.export.Dim("m", min=1, max=QKV_MAXM) + ep = torch.export.export(model, (x,), dynamic_shapes=({0: m_dim},)) + et = _lower_fully_delegated( + ep, + "qkv_routes", + compile_options={_Q4_DECODE_ROUTE_SPEC: True}, + expect_q4_decode_route=True, + ) + with open(os.path.join(out_dir, "qkv_routes.pte"), "wb") as f: + f.write(et.buffer) + print("Exported qkv_routes.pte") + + for m in [QKV_MAXM, 1]: + live_x = _ramp((m, LIN_K)) + with torch.no_grad(): + q, k, v = model(live_x) + base = os.path.join(out_dir, f"qkv_routes.S{m}.") + for name, tensor in (("input", live_x), ("q", q), ("k", k), ("v", v)): + tensor.detach().numpy().astype(" None: + from executorch.backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + SdpaModule, + ) + + hq, hkv, d, cmax, max_s = 8, 2, 132, 16, 16 + + def cfg(s): + return SdpaConfig("wide", hq, hkv, d, s, cmax, 0) + + model = SdpaModule(0) + q, k, v, kc, vc = _det_inputs(cfg(max_s)) + s_dim = torch.export.Dim("s", min=1, max=max_s) + ep = torch.export.export( + model, + (q, k, v, kc, vc), + dynamic_shapes=({1: s_dim}, {1: s_dim}, {1: s_dim}, None, None), + ) + et = _lower_fully_delegated(ep, "sdpa_wide") + with open(os.path.join(out_dir, "sdpa_wide.pte"), "wb") as f: + f.write(et.buffer) + for s in [max_s, 1]: + live = cfg(s) + q, k, v, kc, vc = _det_inputs(live) + golden = _golden(live, q, k, v, kc, vc) + for name, tensor in ( + ("q", q), + ("k", k), + ("v", v), + ("kc", kc), + ("vc", vc), + ("golden", golden), + ): + tensor.detach().numpy().astype(" None: + from executorch.backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + SdpaModule, + ) + + cfg = SdpaConfig(prefix, SD_HQ, SD_HKV, SD_D, s, SD_CMAX, 0) + inputs = _det_inputs(cfg) + ep = torch.export.export(SdpaModule(0), inputs) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + golden = _golden(cfg, *inputs) + for name, tensor in zip(("q", "k", "v", "kc", "vc"), inputs): + tensor.detach().numpy().astype(" [N, EMBED] fp32. EMB_VOCAB = 64 EMB_DIM = 64 @@ -379,14 +640,7 @@ def _export_dynamic_embedding(out_dir: str) -> None: idx_max = torch.arange(EMB_MAXN, dtype=torch.long) n_dim = torch.export.Dim("n", min=1, max=EMB_MAXN) ep = torch.export.export(qm, (idx_max,), dynamic_shapes=({0: n_dim},)) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "embedding_q4gsw not delegated" + et = _lower_fully_delegated(ep, "emb_dyn") with open(os.path.join(out_dir, "emb_dyn.pte"), "wb") as f: f.write(et.buffer) print("Exported emb_dyn.pte") @@ -418,11 +672,7 @@ def _export_dynamic_embedding(out_dir: str) -> None: def _export_dynamic_rope(out_dir: str) -> None: - from executorch.backends.webgpu.test.ops.rope.test_rope import ( - _golden, - _inputs, - Shape, - ) + from executorch.backends.webgpu.test.ops.test_rope import _golden, _inputs, Shape from executorch.examples.models.llama.rope import RotaryEmbedding xq, xk, fc, fs = _inputs(Shape("dyn", 1, ROPE_MAXS, ROPE_NH, ROPE_NKV, ROPE_HD)) @@ -431,14 +681,7 @@ def _export_dynamic_rope(out_dir: str) -> None: ep = torch.export.export( RotaryEmbedding().eval(), (xq, xk, fc, fs), dynamic_shapes=ds ) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "apply_rotary_emb not delegated" + et = _lower_fully_delegated(ep, "rope_dyn") with open(os.path.join(out_dir, "rope_dyn.pte"), "wb") as f: f.write(et.buffer) print("Exported rope_dyn.pte") @@ -471,14 +714,7 @@ def _export_dynamic_select(out_dir: str) -> None: (_ramp((SEL_LEAD, 1, MAXS, HIDDEN)),), dynamic_shapes=({2: s_dim},), ) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "select_copy not delegated" + et = _lower_fully_delegated(ep, "dyn_select") with open(os.path.join(out_dir, "dyn_select.pte"), "wb") as f: f.write(et.buffer) print("Exported dyn_select.pte") @@ -496,6 +732,18 @@ def _export_dynamic_select(out_dir: str) -> None: class TestDynamicShapeExport(unittest.TestCase): + def test_q4_route_compile_specs(self) -> None: + import tempfile + + model = RmsNormModule(HIDDEN).eval() + x = _ramp((1, 1, MAXS, HIDDEN)) + ep = torch.export.export(model, (x,)) + _lower_fully_delegated(ep, "q4_route_negative_control") + + with tempfile.TemporaryDirectory() as d: + _export_dynamic_qkv_routes(d) + self.assertTrue(os.path.exists(os.path.join(d, "qkv_routes.pte"))) + def test_export_dynamic_rms(self) -> None: import tempfile diff --git a/backends/webgpu/test/ops/test_quantized_linear.py b/backends/webgpu/test/ops/test_quantized_linear.py index 451f478d871..c68553c57ca 100644 --- a/backends/webgpu/test/ops/test_quantized_linear.py +++ b/backends/webgpu/test/ops/test_quantized_linear.py @@ -94,8 +94,10 @@ class Q4gswConfig: ] -def _make_quantized_model(k: int, n: int, group_size: int) -> torch.nn.Module: - torch.manual_seed(0) # load-bearing: fixes the weights the golden derives from +def _make_quantized_model( + k: int, n: int, group_size: int, seed: int = 0 +) -> torch.nn.Module: + torch.manual_seed(seed) # load-bearing: fixes the weights used by the golden m = torch.nn.Linear(k, n, bias=False).eval() quantize_( m, diff --git a/backends/webgpu/test/test_native_ci_contract.py b/backends/webgpu/test/test_native_ci_contract.py new file mode 100644 index 00000000000..011354a0e0f --- /dev/null +++ b/backends/webgpu/test/test_native_ci_contract.py @@ -0,0 +1,66 @@ +# 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 pathlib +import re +import shlex +import unittest + + +def _bash_array(source: str, name: str) -> list[str]: + match = re.search( + rf"^{re.escape(name)}=\((.*?)\)", source, re.MULTILINE | re.DOTALL + ) + if match is None: + raise AssertionError(f"{name} Bash array not found") + return shlex.split(match.group(1)) + + +class TestNativeCIContract(unittest.TestCase): + def test_builds_and_runs_every_fixed_target(self) -> None: + backend = pathlib.Path(__file__).parents[1] + cmake = (backend / "CMakeLists.txt").read_text() + script = (backend / "scripts/test_webgpu_native_ci.sh").read_text() + required = { + "webgpu_native_test", + "webgpu_dispatch_order_test", + "webgpu_scratch_buffer_test", + "webgpu_update_cache_test", + "webgpu_index_test", + "webgpu_dynamic_shape_test", + "webgpu_dispatch_2d_test", + "webgpu_compute_dispatch_test", + "webgpu_execution_options_test", + "webgpu_output_suppression_test", + "webgpu_op_test_util_test", + } + + self.assertEqual(set(_bash_array(script, "REQUIRED_TARGETS")), required) + self.assertNotIn("not defined in this tree — skipping", script) + self.assertIn('-DPYTHON_EXECUTABLE="${PYTHON_EXECUTABLE}"', script) + self.assertIn("run_with_required_device env WEBGPU_TEST_SDPA_DIR", script) + self.assertIn( + "if ! grep -q '^WebGPU device acquired (native)$' " '<<<"${output}"; then', + script, + ) + for target in required: + self.assertIn(target, cmake) + self.assertIn(f'"${{BIN_DIR}}/{target}"', script) + + def test_requires_symint_and_suppression_fixtures(self) -> None: + script = ( + pathlib.Path(__file__).parents[1] / "scripts/test_webgpu_native_ci.sh" + ).read_text() + + self.assertIn( + "export_output_suppression_models('${OUTPUT_SUPPRESSION_DIR}')", script + ) + self.assertIn('WEBGPU_TEST_SYMINT_BLOB="${SYMINT_BLOB}"', script) + for fixture in ( + "${SYMINT_BLOB}", + "${OUTPUT_SUPPRESSION_DIR}/input.bin", + ): + self.assertIn(f'require_file "{fixture}"', script) diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index e6c419e4bb3..f64cccec550 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #include #include @@ -1462,15 +1463,137 @@ void test_sdpa_incache_decode( } } -// S1 SymInt round-trip: build a graph directly from a dynamic-input_pos SDPA -// blob; confirm input_pos deserializes as a live SymInt and set/read -// round-trips. +void exercise_symint_host_inputs( + WebGPUGraph& graph, + int symint_id, + int input_tensor_id) { + const auto& input_ids = graph.input_ids(); + std::vector inputs(input_ids.size()); + int64_t host_value = 5; + bool found = false; + for (size_t i = 0; i < input_ids.size(); i++) { + if (input_ids[i] == input_tensor_id) { + inputs[i] = {&host_value, sizeof(host_value), true}; + found = true; + } + } + ASSERT_TRUE(found) << "select_as_symint source is not a graph input"; + + const auto update_from_host = [&](int64_t value) { + host_value = value; + graph.update_symints_from_inputs(inputs); + return graph.read_symint(symint_id); + }; + EXPECT_EQ(update_from_host(5), 5); + EXPECT_EQ( + update_from_host(std::numeric_limits::min()), + std::numeric_limits::min()); + EXPECT_EQ( + update_from_host(std::numeric_limits::max()), + std::numeric_limits::max()); + + const auto expect_out_of_range = [&](int64_t value) { + ASSERT_EQ(update_from_host(17), 17); + host_value = value; + try { + graph.update_symints_from_inputs(inputs); + ADD_FAILURE() << "accepted out-of-range select_as_symint value " << value; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "select_as_symint: selected value is outside int32 range"); + } + EXPECT_EQ(graph.read_symint(symint_id), 17) + << "rejected value changed the live SymInt"; + }; + expect_out_of_range( + int64_t{std::numeric_limits::min()} - int64_t{1}); + expect_out_of_range( + int64_t{std::numeric_limits::max()} + int64_t{1}); + expect_out_of_range(INT64_C(1) << 32); +} + +void test_symint_input_narrowing() { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + const std::vector dims = {1u}; + std::vector<::flatbuffers::Offset> values; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::INT32, + &dims, + /*constant_id=*/-1, + /*mem_obj_id=*/0) + .Union())); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::SymInt, vk::CreateSymInt(fbb, 0).Union())); + const std::vector args = {0, 1, 2, 3}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, 0, "et_vk.select_as_symint.default", &args)); + const std::vector input_ids = {0}; + const std::vector output_ids = {0}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, nullptr)); + ASSERT_EQ(graph.symint_sources().size(), 1u); + const auto& source = graph.symint_sources().front(); + exercise_symint_host_inputs(graph, source.symint_id, source.input_tensor_id); +} + +struct DelegateBlobView { + size_t base_offset; + WebGPUDelegateHeader header; +}; + +std::optional find_delegate_blob( + const std::vector& blob) { + constexpr size_t kHeaderSize = 30; + constexpr size_t kMagicOffset = 4; + constexpr char kMagic[] = {'V', 'H', '0', '0'}; + if (blob.size() < kHeaderSize) { + return std::nullopt; + } + + for (size_t base_offset = 0; base_offset <= blob.size() - kHeaderSize; + base_offset++) { + const uint8_t* base = blob.data() + base_offset; + if (std::memcmp(base + kMagicOffset, kMagic, sizeof(kMagic)) != 0) { + continue; + } + auto header = WebGPUDelegateHeader::parse(base); + if (!header.ok()) { + continue; + } + + const uint64_t available = blob.size() - base_offset; + const auto range_is_in_blob = [available](uint64_t offset, uint64_t size) { + return offset <= available && size <= available - offset; + }; + if (!range_is_in_blob(header->flatbuffer_offset, header->flatbuffer_size) || + !range_is_in_blob(header->bytes_offset, header->bytes_size)) { + continue; + } + return DelegateBlobView{base_offset, *header}; + } + return std::nullopt; +} + +// S1 SymInt round-trip: confirm a dynamic input_pos stays live. void test_symint_roundtrip(const std::string& blob_path) { printf("\n--- Test: symint round-trip (%s) ---\n", blob_path.c_str()); FILE* f = std::fopen(blob_path.c_str(), "rb"); - if (!f) { - GTEST_SKIP() << blob_path << " not present"; - } + ASSERT_NE(f, nullptr) << blob_path << " not present"; std::fseek(f, 0, SEEK_END); long n = std::ftell(f); std::fseek(f, 0, SEEK_SET); @@ -1479,13 +1602,16 @@ void test_symint_roundtrip(const std::string& blob_path) { std::fclose(f); ASSERT_EQ(rd, blob.size()) << "short read of " << blob_path; - auto header = WebGPUDelegateHeader::parse(blob.data()); - ASSERT_TRUE(header.ok()) << "delegate header parse"; - const uint8_t* base = blob.data(); + const auto delegate = find_delegate_blob(blob); + ASSERT_TRUE(delegate.has_value()) + << "no complete VH00 delegate blob found in " << blob_path; + const uint8_t* base = blob.data() + delegate->base_offset; WebGPUGraph graph; try { graph.build( - base + header->flatbuffer_offset, base + header->bytes_offset, nullptr); + base + delegate->header.flatbuffer_offset, + base + delegate->header.bytes_offset, + nullptr); } catch (const std::exception& e) { FAIL() << "graph build: " << e.what(); } @@ -1508,22 +1634,10 @@ void test_symint_roundtrip(const std::string& blob_path) { ASSERT_EQ(graph.read_symint(sid), 7) << "set/read round-trip (got " << graph.read_symint(sid) << ")"; - // Execute-read: feed a fake input_pos=5 via the recorded select_as_symint - // source and confirm update_symints_from_inputs populates the SymInt. const auto& srcs = graph.symint_sources(); ASSERT_FALSE(srcs.empty()) << "no select_as_symint source recorded"; - const auto& in_ids = graph.input_ids(); - std::vector fake_inputs(in_ids.size()); - int64_t fake_pos = 5; - for (size_t i = 0; i < in_ids.size(); i++) { - if (in_ids[i] == srcs[0].input_tensor_id) { - fake_inputs[i] = {&fake_pos, sizeof(int64_t), true}; - } - } - graph.update_symints_from_inputs(fake_inputs); - ASSERT_EQ(graph.read_symint(srcs[0].symint_id), 5) - << "execute-read (got " << graph.read_symint(srcs[0].symint_id) - << ", want 5)"; + exercise_symint_host_inputs( + graph, srcs[0].symint_id, srcs[0].input_tensor_id); printf( "PASS: symint round-trip (SymInt %d: deserialize, live buffer, " @@ -1537,9 +1651,7 @@ void test_symint_roundtrip(const std::string& blob_path) { void test_resize_hook(const std::string& blob_path) { printf("\n--- Test: resize-hook dirty-gating (%s) ---\n", blob_path.c_str()); FILE* f = std::fopen(blob_path.c_str(), "rb"); - if (!f) { - GTEST_SKIP() << blob_path << " not present"; - } + ASSERT_NE(f, nullptr) << blob_path << " not present"; std::fseek(f, 0, SEEK_END); long n = std::ftell(f); std::fseek(f, 0, SEEK_SET); @@ -1547,13 +1659,16 @@ void test_resize_hook(const std::string& blob_path) { size_t rd = std::fread(blob.data(), 1, blob.size(), f); std::fclose(f); ASSERT_EQ(rd, blob.size()) << "short read of " << blob_path; - auto header = WebGPUDelegateHeader::parse(blob.data()); - ASSERT_TRUE(header.ok()) << "delegate header parse"; - const uint8_t* base = blob.data(); + const auto delegate = find_delegate_blob(blob); + ASSERT_TRUE(delegate.has_value()) + << "no complete VH00 delegate blob found in " << blob_path; + const uint8_t* base = blob.data() + delegate->base_offset; WebGPUGraph graph; try { graph.build( - base + header->flatbuffer_offset, base + header->bytes_offset, nullptr); + base + delegate->header.flatbuffer_offset, + base + delegate->header.bytes_offset, + nullptr); } catch (const std::exception& e) { FAIL() << "graph build: " << e.what(); } @@ -2347,7 +2462,8 @@ TEST(WebGPUNative, SdpaAllFamiliesRanWhenDirSet) { TEST(WebGPUNative, SymintRoundtrip) { if (g_symint_blob.empty()) { - GTEST_SKIP() << "WEBGPU_TEST_SYMINT_BLOB not set"; + test_symint_input_narrowing(); + return; } test_symint_roundtrip(g_symint_blob); } diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 46d285aa60b..fe8340d6528 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -9,8 +9,10 @@ Loads the generator by file path (no package/namespace dependency). """ +import contextlib import hashlib import importlib.util +import io import re import tempfile import unittest @@ -72,6 +74,29 @@ def _function_source(text: str, name: str) -> str: class WgslCodegenTest(unittest.TestCase): + def test_registry_entries_match_concrete_headers(self) -> None: + entries = g.registry_entries() + names = [entry.name for entry in entries] + expected = sorted( + header.name[: -len("_wgsl.h")] + for wgsl in g.discover() + for header, _ in g.headers_for_shader(wgsl) + ) + self.assertEqual(names, expected) + self.assertEqual(len(names), len(set(names))) + + def test_registry_render_is_deterministic(self) -> None: + entries = g.registry_entries() + self.assertEqual( + g.render_registry(entries), + g.render_registry(list(reversed(entries))), + ) + + def test_registry_rejects_duplicate_names(self) -> None: + entry = g.registry_entries()[0] + with self.assertRaisesRegex(ValueError, "duplicate shader registry name"): + g.render_registry([entry, entry]) + def test_symbol_base(self) -> None: self.assertEqual(g.symbol_base("binary_add"), "BinaryAdd") self.assertEqual( @@ -179,7 +204,12 @@ def test_check_fails_on_stale_header(self) -> None: orig = g.BACKEND_ROOT g.BACKEND_ROOT = Path(tmp) try: - self.assertEqual(g.main(["--check"]), 1) + output = io.StringIO() + with contextlib.redirect_stdout(output): + self.assertEqual(g.main(["--check"]), 1) + self.assertEqual( + output.getvalue().count("Stale embedded WGSL headers"), 1 + ) finally: g.BACKEND_ROOT = orig