From df7d5ace08ac3416b0311a2b49f0f9c599d5a5ce Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 16 Jan 2026 15:28:13 +0800 Subject: [PATCH 1/7] [webgpu] Enable profiling for graph capture --- .../core/providers/webgpu/webgpu_context.cc | 21 +++++++++++++++---- .../core/providers/webgpu/webgpu_context.h | 15 ++++++++++++- .../webgpu/webgpu_execution_provider.cc | 6 ++++++ 3 files changed, 37 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 7cb6a852e8d7e..aaff4d4611db8 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -792,11 +792,20 @@ void WebGpuContext::LaunchComputePipeline(const wgpu::ComputePassEncoder& comput if (indirect_dispatch_tensor != nullptr) { indirect_buffer = reinterpret_cast(const_cast(indirect_dispatch_tensor->DataRaw())); } + + // Store profiling info if profiling is enabled + std::optional, std::vector>> profiling_data; + if (is_profiling_ && !pending_kernels_.empty()) { + const auto& kernel_info = pending_kernels_.back(); + profiling_data = std::make_tuple(kernel_info.name, kernel_info.cache_key, kernel_info.input_shapes, kernel_info.output_shapes); + } + external_captured_commands_->push_back({program_artifact.compute_pipeline, bind_group, bind_group_layout, {x, y, z}, - indirect_buffer}); + indirect_buffer, + profiling_data}); } else { compute_pass_encoder.SetPipeline(program_artifact.compute_pipeline); wgpuComputePassEncoderSetBindGroup(compute_pass_encoder.Get(), 0, bind_group, 0, nullptr); @@ -827,9 +836,6 @@ void WebGpuContext::CaptureBegin(std::vector* captu external_captured_commands_->clear(); } - // TODO: support profiling with graph capture. - ORT_ENFORCE(!is_profiling_, "profiling is not supported yet under graph capture mode"); - graph_capture_state_ = GraphCaptureState::Capturing; } @@ -842,6 +848,13 @@ void WebGpuContext::Replay(const std::vector& captu auto& command = captured_commands[i]; const auto& compute_pass_encoder = GetComputePassEncoder(); WriteTimestamp(num_pending_dispatches_ * 2); + + // Restore profiling info if available and profiling is enabled + if (is_profiling_ && command.pending_kernel_info.has_value()) { + const auto& [name, cache_key, input_shapes, output_shapes] = command.pending_kernel_info.value(); + pending_kernels_.emplace_back(name, cache_key, input_shapes, output_shapes); + } + compute_pass_encoder.SetPipeline(command.compute_pipeline); wgpuComputePassEncoderSetBindGroup(compute_pass_encoder.Get(), 0, command.bind_group, 0, nullptr); diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index 9feea69f3702b..c8d7275858237 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -5,6 +5,8 @@ #include #include +#include +#include #include "core/providers/webgpu/webgpu_external_header.h" @@ -31,7 +33,8 @@ struct CapturedCommandInfo { WGPUBindGroup bind_group; WGPUBindGroupLayout bind_group_layout; std::array dispatch_group; - WGPUBuffer indirect_buffer; // WGPUBuffer for indirect dispatch, nullptr if not using indirect dispatch + WGPUBuffer indirect_buffer; // WGPUBuffer for indirect dispatch, nullptr if not using indirect dispatch + std::optional, std::vector>> pending_kernel_info; // Optional profiling data: (name, cache_key, input_shapes, output_shapes) }; struct WebGpuBufferCacheConfig { @@ -280,6 +283,16 @@ class WebGpuContext final { } } + // Constructor for replay - takes shapes directly + PendingKernelInfo(std::string name_in, + std::string cache_key_in, + std::vector input_shapes_in, + std::vector output_shapes_in) + : name{std::move(name_in)}, + cache_key{std::move(cache_key_in)}, + input_shapes{std::move(input_shapes_in)}, + output_shapes{std::move(output_shapes_in)} {} + PendingKernelInfo(PendingKernelInfo&&) = default; PendingKernelInfo& operator=(PendingKernelInfo&&) = default; ORT_DISALLOW_COPY_AND_ASSIGNMENT(PendingKernelInfo); diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index ee6b7707384e2..426c1b6f4f123 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -1050,7 +1050,13 @@ bool WebGpuExecutionProvider::IsGraphCaptured(int graph_annotation_id) const { Status WebGpuExecutionProvider::ReplayGraph(int graph_annotation_id) { ORT_ENFORCE(IsGraphCaptured(graph_annotation_id)); + if (profiler_->Enabled()) { + context_.StartProfiling(); + } context_.Replay(captured_commands_, *graph_buffer_mgr_); + if (profiler_->Enabled()) { + context_.CollectProfilingData(profiler_->Events()); + } return Status::OK(); } From 4cdc68b37f9bbead0d3bfeb25cffe928ae4bb394 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 20 Jan 2026 14:52:42 +0800 Subject: [PATCH 2/7] address Copilot's comments --- onnxruntime/core/providers/webgpu/webgpu_context.h | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index c8d7275858237..cb1020d27ba05 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -33,8 +33,10 @@ struct CapturedCommandInfo { WGPUBindGroup bind_group; WGPUBindGroupLayout bind_group_layout; std::array dispatch_group; - WGPUBuffer indirect_buffer; // WGPUBuffer for indirect dispatch, nullptr if not using indirect dispatch - std::optional, std::vector>> pending_kernel_info; // Optional profiling data: (name, cache_key, input_shapes, output_shapes) + // WGPUBuffer for indirect dispatch, nullptr if not using indirect dispatch + WGPUBuffer indirect_buffer; + // Optional profiling data: (name, cache_key, input_shapes, output_shapes) + std::optional, std::vector>> pending_kernel_info; }; struct WebGpuBufferCacheConfig { From ba18d6d70e4a28ec4ddf02dd723706b348d887a1 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 20 Jan 2026 17:06:14 +0800 Subject: [PATCH 3/7] address copilot's comments --- .../core/providers/webgpu/webgpu_context.cc | 13 ++-- .../core/providers/webgpu/webgpu_context.h | 75 ++++++++----------- 2 files changed, 39 insertions(+), 49 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index aaff4d4611db8..bb3ace497d4a7 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -3,6 +3,7 @@ #include #include +#include #if defined(__GNUC__) #pragma GCC diagnostic push @@ -793,11 +794,10 @@ void WebGpuContext::LaunchComputePipeline(const wgpu::ComputePassEncoder& comput indirect_buffer = reinterpret_cast(const_cast(indirect_dispatch_tensor->DataRaw())); } - // Store profiling info if profiling is enabled - std::optional, std::vector>> profiling_data; - if (is_profiling_ && !pending_kernels_.empty()) { - const auto& kernel_info = pending_kernels_.back(); - profiling_data = std::make_tuple(kernel_info.name, kernel_info.cache_key, kernel_info.input_shapes, kernel_info.output_shapes); + // Always store profiling metadata to support profiling during replay regardless of current profiling state + std::optional profiling_data; + if (!pending_kernels_.empty()) { + profiling_data = pending_kernels_.back(); } external_captured_commands_->push_back({program_artifact.compute_pipeline, @@ -851,8 +851,7 @@ void WebGpuContext::Replay(const std::vector& captu // Restore profiling info if available and profiling is enabled if (is_profiling_ && command.pending_kernel_info.has_value()) { - const auto& [name, cache_key, input_shapes, output_shapes] = command.pending_kernel_info.value(); - pending_kernels_.emplace_back(name, cache_key, input_shapes, output_shapes); + pending_kernels_.emplace_back(command.pending_kernel_info.value()); } compute_pass_encoder.SetPipeline(command.compute_pipeline); diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index cb1020d27ba05..e5b98bec58571 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -6,7 +6,6 @@ #include #include #include -#include #include "core/providers/webgpu/webgpu_external_header.h" @@ -27,6 +26,37 @@ class WebGpuContext; class ComputeContextBase; class ProgramBase; +// PendingKernelInfo stores profiling information for a kernel execution +struct PendingKernelInfo { + PendingKernelInfo(std::string_view kernel_name, + std::string_view kernel_type, + std::string_view program_name, + std::string_view cache_key, + const std::vector& inputs, + const std::vector& outputs) + : name{absl::StrJoin({kernel_name, kernel_type, program_name}, "&")}, cache_key{cache_key} { + // Store shape information instead of tensor pointers to avoid accessing released tensors + input_shapes.reserve(inputs.size()); + for (const auto& input : inputs) { + input_shapes.emplace_back(input.use_override_shape ? input.override_shape : input.tensor->Shape()); + } + output_shapes.reserve(outputs.size()); + for (const auto& output : outputs) { + output_shapes.emplace_back(output.use_override_shape ? output.override_shape : output.tensor->Shape()); + } + } + + PendingKernelInfo(const PendingKernelInfo&) = default; + PendingKernelInfo& operator=(const PendingKernelInfo&) = default; + PendingKernelInfo(PendingKernelInfo&&) = default; + PendingKernelInfo& operator=(PendingKernelInfo&&) = default; + + std::string name; + std::string cache_key; + std::vector input_shapes; + std::vector output_shapes; +}; + // Definition for CapturedCommandInfo in the webgpu namespace struct CapturedCommandInfo { wgpu::ComputePipeline compute_pipeline; @@ -35,8 +65,8 @@ struct CapturedCommandInfo { std::array dispatch_group; // WGPUBuffer for indirect dispatch, nullptr if not using indirect dispatch WGPUBuffer indirect_buffer; - // Optional profiling data: (name, cache_key, input_shapes, output_shapes) - std::optional, std::vector>> pending_kernel_info; + // Optional profiling data + std::optional pending_kernel_info; }; struct WebGpuBufferCacheConfig { @@ -266,45 +296,6 @@ class WebGpuContext final { wgpu::Limits GetRequiredLimits(const wgpu::Adapter& adapter) const; void WriteTimestamp(uint32_t query_index); - struct PendingKernelInfo { - PendingKernelInfo(std::string_view kernel_name, - std::string_view kernel_type, - std::string_view program_name, - std::string_view cache_key, - const std::vector& inputs, - const std::vector& outputs) - : name{absl::StrJoin({kernel_name, kernel_type, program_name}, "&")}, cache_key{cache_key} { - // Store shape information instead of tensor pointers to avoid accessing released tensors - input_shapes.reserve(inputs.size()); - for (const auto& input : inputs) { - input_shapes.emplace_back(input.use_override_shape ? input.override_shape : input.tensor->Shape()); - } - output_shapes.reserve(outputs.size()); - for (const auto& output : outputs) { - output_shapes.emplace_back(output.use_override_shape ? output.override_shape : output.tensor->Shape()); - } - } - - // Constructor for replay - takes shapes directly - PendingKernelInfo(std::string name_in, - std::string cache_key_in, - std::vector input_shapes_in, - std::vector output_shapes_in) - : name{std::move(name_in)}, - cache_key{std::move(cache_key_in)}, - input_shapes{std::move(input_shapes_in)}, - output_shapes{std::move(output_shapes_in)} {} - - PendingKernelInfo(PendingKernelInfo&&) = default; - PendingKernelInfo& operator=(PendingKernelInfo&&) = default; - ORT_DISALLOW_COPY_AND_ASSIGNMENT(PendingKernelInfo); - - std::string name; - std::string cache_key; - std::vector input_shapes; - std::vector output_shapes; - }; - struct PendingQueryInfo { PendingQueryInfo(std::vector&& kernels, wgpu::Buffer query_buffer) : kernels{std::move(kernels)}, query_buffer{query_buffer} {} From 5f92ab1150e07669ba90a7c6fc84ab222227926c Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 21 Jan 2026 17:17:24 +0800 Subject: [PATCH 4/7] address comments --- .../core/providers/webgpu/webgpu_context.cc | 45 ++++++++++++------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index bb3ace497d4a7..5bf81bf13d0e3 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -282,16 +282,6 @@ Status WebGpuContext::Run(ComputeContextBase& context, const ProgramBase& progra auto key = CalculateProgramCacheKey(program, inputs_segments, outputs_segments, is_1d_dispatch); - if (is_profiling_) { - PendingKernelInfo pending_kernel_info(context.NodeName(), - context.OpType(), - program.Name(), - key, - inputs, - outputs); - pending_kernels_.emplace_back(std::move(pending_kernel_info)); - } - LOGS(context.Logger(), INFO) << "Starting program \"" << key << "\" (" << x << ", " << y << ", " << z << ")"; const auto* program_artifact = program_mgr_->Get(key); @@ -480,6 +470,26 @@ Status WebGpuContext::Run(ComputeContextBase& context, const ProgramBase& progra WriteTimestamp(num_pending_dispatches_ * 2 + 1); ++num_pending_dispatches_; + // Update profiling data after LaunchComputePipeline + if (is_profiling_ || graph_capture_state_ == GraphCaptureState::Capturing) { + PendingKernelInfo pending_kernel_info(context.NodeName(), + context.OpType(), + program.Name(), + key, + inputs, + outputs); + + if (graph_capture_state_ == GraphCaptureState::Capturing) { + // Update the last captured command's profiling info + if (external_captured_commands_ && !external_captured_commands_->empty()) { + external_captured_commands_->back().pending_kernel_info = std::move(pending_kernel_info); + } + } else { + // Add to pending kernels for current run profiling + pending_kernels_.emplace_back(std::move(pending_kernel_info)); + } + } + if (num_pending_dispatches_ >= max_num_pending_dispatches_ || (is_profiling_ && query_type_ == TimestampQueryType::AtPasses)) { EndComputePass(); @@ -715,6 +725,10 @@ void WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr) { EndComputePass(); if (is_profiling_ && num_pending_dispatches_ > 0) { + ORT_ENFORCE(num_pending_dispatches_ == pending_kernels_.size(), + "Number of pending dispatches (", num_pending_dispatches_, + ") does not match pending kernels size (", pending_kernels_.size(), ")"); + uint32_t query_count = num_pending_dispatches_ * 2; current_command_encoder_.ResolveQuerySet( query_set_, @@ -794,18 +808,13 @@ void WebGpuContext::LaunchComputePipeline(const wgpu::ComputePassEncoder& comput indirect_buffer = reinterpret_cast(const_cast(indirect_dispatch_tensor->DataRaw())); } - // Always store profiling metadata to support profiling during replay regardless of current profiling state - std::optional profiling_data; - if (!pending_kernels_.empty()) { - profiling_data = pending_kernels_.back(); - } - + // Profiling data will be populated after this call in Run() external_captured_commands_->push_back({program_artifact.compute_pipeline, bind_group, bind_group_layout, {x, y, z}, indirect_buffer, - profiling_data}); + std::nullopt}); } else { compute_pass_encoder.SetPipeline(program_artifact.compute_pipeline); wgpuComputePassEncoderSetBindGroup(compute_pass_encoder.Get(), 0, bind_group, 0, nullptr); @@ -836,6 +845,8 @@ void WebGpuContext::CaptureBegin(std::vector* captu external_captured_commands_->clear(); } + // Disable profiling during capture mode + is_profiling_ = false; graph_capture_state_ = GraphCaptureState::Capturing; } From 5c4c5bef556dd778b2bada334379e9cd7ef79746 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 21 Jan 2026 18:05:14 +0800 Subject: [PATCH 5/7] address comments --- .../core/providers/webgpu/webgpu_context.cc | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 5bf81bf13d0e3..3ad61810639cc 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -808,7 +808,7 @@ void WebGpuContext::LaunchComputePipeline(const wgpu::ComputePassEncoder& comput indirect_buffer = reinterpret_cast(const_cast(indirect_dispatch_tensor->DataRaw())); } - // Profiling data will be populated after this call in Run() + // Profiling data will be populated in Run() after this call returns. external_captured_commands_->push_back({program_artifact.compute_pipeline, bind_group, bind_group_layout, @@ -845,7 +845,7 @@ void WebGpuContext::CaptureBegin(std::vector* captu external_captured_commands_->clear(); } - // Disable profiling during capture mode + // Temporarily disable active profiling during capture; profiling data is deferred and used during replay is_profiling_ = false; graph_capture_state_ = GraphCaptureState::Capturing; } @@ -860,9 +860,15 @@ void WebGpuContext::Replay(const std::vector& captu const auto& compute_pass_encoder = GetComputePassEncoder(); WriteTimestamp(num_pending_dispatches_ * 2); - // Restore profiling info if available and profiling is enabled - if (is_profiling_ && command.pending_kernel_info.has_value()) { - pending_kernels_.emplace_back(command.pending_kernel_info.value()); + // Restore profiling info when profiling is enabled. All commands are expected + // to have profiling data in this mode to keep pending_kernels_ consistent + // with num_pending_dispatches_. + if (is_profiling_) { + ORT_ENFORCE(command.pending_kernel_info.has_value(), + "WebGpuContext::Replay: profiling is enabled but captured command at index ", + i, + " is missing pending_kernel_info."); + pending_kernels_.emplace_back(*command.pending_kernel_info); } compute_pass_encoder.SetPipeline(command.compute_pipeline); From fb702a7c27f94dbda6b9ac5f8c221651ee7663a1 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 27 Jan 2026 16:03:03 +0800 Subject: [PATCH 6/7] align with the latest code --- .../core/providers/webgpu/webgpu_execution_provider.cc | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 9ed90a813115c..844591a930c0c 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -1075,12 +1075,13 @@ bool WebGpuExecutionProvider::IsGraphCaptured(int graph_annotation_id) const { Status WebGpuExecutionProvider::ReplayGraph(int graph_annotation_id) { ORT_ENFORCE(IsGraphCaptured(graph_annotation_id)); - if (profiler_->Enabled()) { + // TODO: enable profiling in run level + if (session_profiler_ && session_profiler_->Enabled()) { context_.StartProfiling(); } context_.Replay(captured_commands_, *graph_buffer_mgr_); - if (profiler_->Enabled()) { - context_.CollectProfilingData(profiler_->Events()); + if (session_profiler_ && session_profiler_->Enabled()) { + context_.CollectProfilingData(); } return Status::OK(); } From 7c1b7f9cfba11d4d0411a54ae48f78079c676404 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 2 Feb 2026 18:40:57 +0800 Subject: [PATCH 7/7] Don't always store profiling info during capturing --- onnxruntime/core/providers/webgpu/webgpu_context.cc | 8 +++----- onnxruntime/core/providers/webgpu/webgpu_context.h | 2 +- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index b825a9f26634d..318e815255ff7 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -472,7 +472,7 @@ Status WebGpuContext::Run(ComputeContextBase& context, const ProgramBase& progra ++num_pending_dispatches_; // Update profiling data after LaunchComputePipeline - if (is_profiling_ || graph_capture_state_ == GraphCaptureState::Capturing) { + if (is_profiling_) { PendingKernelInfo pending_kernel_info(context.NodeName(), context.OpType(), program.Name(), @@ -588,7 +588,7 @@ wgpu::Limits WebGpuContext::GetRequiredLimits(const wgpu::Adapter& adapter) cons } void WebGpuContext::WriteTimestamp(uint32_t query_index) { - if (!is_profiling_ || query_type_ != TimestampQueryType::InsidePasses) { + if (!is_profiling_ || graph_capture_state_ == GraphCaptureState::Capturing || query_type_ != TimestampQueryType::InsidePasses) { return; } @@ -725,7 +725,7 @@ void WebGpuContext::Flush(const webgpu::BufferManager& buffer_mgr) { EndComputePass(); - if (is_profiling_ && num_pending_dispatches_ > 0) { + if (is_profiling_ && num_pending_dispatches_ > 0 && graph_capture_state_ != GraphCaptureState::Capturing) { ORT_ENFORCE(num_pending_dispatches_ == pending_kernels_.size(), "Number of pending dispatches (", num_pending_dispatches_, ") does not match pending kernels size (", pending_kernels_.size(), ")"); @@ -846,8 +846,6 @@ void WebGpuContext::CaptureBegin(std::vector* captu external_captured_commands_->clear(); } - // Temporarily disable active profiling during capture; profiling data is deferred and used during replay - is_profiling_ = false; graph_capture_state_ = GraphCaptureState::Capturing; } diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index e280cb58963a6..7645f1e6b2482 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -180,7 +180,7 @@ class WebGpuContext final { wgpu::ComputePassDescriptor compute_pass_desc{}; - if (is_profiling_ && query_type_ == TimestampQueryType::AtPasses) { + if (is_profiling_ && query_type_ == TimestampQueryType::AtPasses && graph_capture_state_ != GraphCaptureState::Capturing) { wgpu::PassTimestampWrites timestampWrites = { nullptr, query_set_,