[ExecuTorch][WebGPU] Add certified scoped output suppression

backends/webgpu/CMakeLists.txt

@@ -26,9 +26,13 @@ if(NOT TARGET vulkan_schema)
 endif()
 
 set(WEBGPU_SRCS
-    runtime/WebGPUBackend.cpp runtime/WebGPUGraph.cpp
-    runtime/WebGPUDelegateHeader.cpp runtime/WebGPUDevice.cpp
-    runtime/WebGPUQueryPool.cpp runtime/ops/OperatorRegistry.cpp
+    runtime/WebGPUBackend.cpp
+    runtime/WebGPUExecutionOptions.cpp
+    runtime/WebGPUGraph.cpp
+    runtime/WebGPUDelegateHeader.cpp
+    runtime/WebGPUDevice.cpp
+    runtime/WebGPUQueryPool.cpp
+    runtime/ops/OperatorRegistry.cpp
 )
 
 # Op handlers: glob so adding an op needs no CMakeLists edit. CONFIGURE_DEPENDS
@@ -201,5 +205,22 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST)
     target_link_libraries(
       webgpu_dispatch_2d_test PRIVATE GTest::gtest GTest::gtest_main
     )
+    add_executable(
+      webgpu_execution_options_test test/native/test_execution_options.cpp
+                                    runtime/WebGPUExecutionOptions.cpp
+    )
+    target_include_directories(
+      webgpu_execution_options_test
+      PRIVATE $<BUILD_INTERFACE:${EXECUTORCH_ROOT}/..>
+    )
+    target_link_libraries(
+      webgpu_execution_options_test PRIVATE GTest::gtest GTest::gtest_main
+    )
+    target_compile_options(webgpu_execution_options_test PRIVATE -fexceptions)
+    set_property(TARGET webgpu_execution_options_test PROPERTY CXX_STANDARD 17)
+    add_webgpu_native_test(
+      webgpu_output_suppression_test test/native/test_output_suppression.cpp
+    )
+    target_link_libraries(webgpu_output_suppression_test PRIVATE GTest::gtest)
   endif()
 endif()

backends/webgpu/runtime/WebGPUBackend.cpp

@@ -8,6 +8,7 @@
 
 #include <executorch/backends/webgpu/runtime/WebGPUBackend.h>
 #include <executorch/backends/webgpu/runtime/WebGPUDelegateHeader.h>
+#include <executorch/backends/webgpu/runtime/WebGPUExecutionOptions.h>
 #include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
 
 #include <executorch/backends/vulkan/serialization/schema_generated.h>
@@ -121,9 +122,11 @@ Error WebGPUBackend::execute(
     DelegateHandle* handle,
     Span<EValue*> args) const {
   WebGPUGraph* graph = static_cast<WebGPUGraph*>(handle);
+  const WebGPUExecutionOptions options = current_webgpu_execution_options();
 
   const size_t num_inputs = graph->input_ids().size();
   const size_t num_outputs = graph->output_ids().size();
+  WebGPUGraphExecutionOptions graph_options;
 
   // Copy inputs from EValue tensors to GPU buffers
   std::vector<InputData> inputs;
@@ -159,6 +162,14 @@ Error WebGPUBackend::execute(
         return Error::Internal;
       }
     }
+    std::vector<const void*> delegate_outputs;
+    delegate_outputs.reserve(num_outputs);
+    for (size_t i = 0; i < num_outputs; i++) {
+      delegate_outputs.push_back(
+          args[num_inputs + i]->toTensor().mutable_data_ptr());
+    }
+    graph_options =
+        resolve_webgpu_graph_execution_options(delegate_outputs, options);
   } catch (const std::exception& e) {
     ET_LOG(Error, "WebGPU input/output resize / copy failed: %s", e.what());
     return Error::Internal;
@@ -167,15 +178,18 @@ Error WebGPUBackend::execute(
   // Execute + read back; fail loud as a runtime Error so a throw never crosses
   // the backend boundary.
   try {
-    graph->execute();
+    const WebGPUExecutionPlan plan = graph->make_execution_plan(graph_options);
+    graph->execute(plan);
+
+    // Copy outputs from GPU staging buffers to EValue tensor data pointers
     std::vector<std::pair<void*, size_t>> outputs;
     outputs.reserve(num_outputs);
     for (size_t i = 0; i < num_outputs; i++) {
       const size_t arg_idx = num_inputs + i;
       auto& tensor = args[arg_idx]->toTensor();
       outputs.emplace_back(tensor.mutable_data_ptr(), tensor.nbytes());
     }
-    graph->copy_outputs(outputs);
+    graph->copy_outputs(outputs, plan);
   } catch (const std::exception& e) {
     ET_LOG(Error, "WebGPU execute / output copy failed: %s", e.what());
     return Error::Internal;

backends/webgpu/runtime/WebGPUExecutionOptions.cpp

@@ -0,0 +1,137 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUExecutionOptions.h>
+
+#include <algorithm>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+namespace {
+
+thread_local WebGPUExecutionOptions execution_options;
+
+} // namespace
+
+WebGPUExecutionOptions current_webgpu_execution_options() {
+  return execution_options;
+}
+
+ScopedWebGPUExecutionOptions::ScopedWebGPUExecutionOptions(
+    WebGPUExecutionOptions options)
+    : previous_(execution_options) {
+  execution_options = options;
+}
+
+ScopedWebGPUExecutionOptions::~ScopedWebGPUExecutionOptions() {
+  execution_options = previous_;
+}
+
+WebGPUExecutionPlan plan_webgpu_execution(
+    size_t dispatch_count,
+    size_t output_count,
+    ExecuteConfig config,
+    const std::vector<SuppressibleOutput>& suppressible_outputs,
+    WebGPUGraphExecutionOptions options) {
+  std::vector<bool> suppressed_dispatches(dispatch_count, false);
+  std::vector<bool> copy_outputs(output_count, true);
+  std::vector<bool> seen_output_ordinals(output_count, false);
+
+  for (const auto& output : suppressible_outputs) {
+    if (output.output_ordinal >= output_count ||
+        output.dispatch_begin >= output.dispatch_end ||
+        output.dispatch_end > dispatch_count) {
+      throw std::runtime_error(
+          "WebGPU: invalid suppressible output range (output_id " +
+          std::to_string(output.output_id) + ")");
+    }
+    if (seen_output_ordinals[output.output_ordinal]) {
+      throw std::runtime_error(
+          "WebGPU: duplicate suppressible output (output_id " +
+          std::to_string(output.output_id) + ")");
+    }
+    seen_output_ordinals[output.output_ordinal] = true;
+    if (output.output_ordinal != options.suppress_output_ordinal) {
+      continue;
+    }
+    copy_outputs[output.output_ordinal] = false;
+    // Only the one ordinal matching suppress_output_ordinal reaches here (the
+    // duplicate check above rejects a repeat), so its dispatch range is
+    // disjoint by construction — mark it suppressed without a redundant overlap
+    // check.
+    for (size_t i = output.dispatch_begin; i < output.dispatch_end; i++) {
+      suppressed_dispatches[i] = true;
+    }
+  }
+
+  WebGPUExecutionPlan plan;
+  plan.copy_outputs = std::move(copy_outputs);
+
+  auto append_chunk = [&](size_t begin, size_t end) {
+    std::vector<size_t> indices;
+    indices.reserve(end - begin);
+    for (size_t i = begin; i < end; i++) {
+      if (!suppressed_dispatches[i]) {
+        indices.push_back(i);
+      }
+    }
+    if (!indices.empty()) {
+      plan.dispatch_chunks.push_back(std::move(indices));
+    }
+  };
+
+  if (config.chunk_size == 0 || dispatch_count <= config.chunk_size) {
+    append_chunk(0, dispatch_count);
+  } else {
+    size_t start = 0;
+    size_t current_chunk = config.initial_chunk_size > 0
+        ? config.initial_chunk_size
+        : config.chunk_size;
+    while (start < dispatch_count) {
+      const size_t end = std::min(start + current_chunk, dispatch_count);
+      append_chunk(start, end);
+      start = end;
+      current_chunk = config.chunk_size;
+    }
+  }
+  if (plan.dispatch_chunks.empty() &&
+      std::any_of(
+          plan.copy_outputs.begin(), plan.copy_outputs.end(), [](bool copy) {
+            return copy;
+          })) {
+    plan.dispatch_chunks.emplace_back();
+  }
+  return plan;
+}
+
+WebGPUGraphExecutionOptions resolve_webgpu_graph_execution_options(
+    const std::vector<const void*>& delegate_outputs,
+    WebGPUExecutionOptions options) {
+  if (options.discardable_output_data == nullptr) {
+    return {};
+  }
+  if (!options.exact_method_certificate_verified) {
+    return {};
+  }
+
+  size_t match = kNoOutputOrdinal;
+  for (size_t i = 0; i < delegate_outputs.size(); i++) {
+    if (delegate_outputs[i] != options.discardable_output_data) {
+      continue;
+    }
+    if (match != kNoOutputOrdinal) {
+      return {};
+    }
+    match = i;
+  }
+  return {match};
+}
+
+} // namespace executorch::backends::webgpu

backends/webgpu/runtime/WebGPUExecutionOptions.h

@@ -0,0 +1,86 @@
+/*
+ * 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 <cstddef>
+#include <utility>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+constexpr size_t kNoOutputOrdinal = static_cast<size_t>(-1);
+
+struct WebGPUExecutionOptions {
+  // The certificate must bind the exact PTE and method and prove one delegate,
+  // no portable nodes, and a unique leaf method output at this data pointer.
+  // The caller must keep this pointer valid and unchanged for the complete
+  // synchronous backend invocation in which these options are scoped.
+  const void* discardable_output_data = nullptr;
+  bool exact_method_certificate_verified = false;
+};
+
+struct WebGPUGraphExecutionOptions {
+  size_t suppress_output_ordinal = kNoOutputOrdinal;
+};
+
+struct ExecuteConfig {
+  size_t chunk_size = 0;
+  size_t initial_chunk_size = 0;
+};
+
+struct SuppressibleOutput {
+  int output_id = -1;
+  size_t output_ordinal = 0;
+  size_t dispatch_begin = 0;
+  size_t dispatch_end = 0;
+};
+
+struct WebGPUExecutionPlan {
+  std::vector<std::vector<size_t>> dispatch_chunks;
+  std::vector<bool> copy_outputs;
+};
+
+WebGPUExecutionPlan plan_webgpu_execution(
+    size_t dispatch_count,
+    size_t output_count,
+    ExecuteConfig config,
+    const std::vector<SuppressibleOutput>& suppressible_outputs,
+    WebGPUGraphExecutionOptions options);
+
+WebGPUGraphExecutionOptions resolve_webgpu_graph_execution_options(
+    const std::vector<const void*>& delegate_outputs,
+    WebGPUExecutionOptions options);
+
+WebGPUExecutionOptions current_webgpu_execution_options();
+
+class ScopedWebGPUExecutionOptions final {
+ public:
+  explicit ScopedWebGPUExecutionOptions(WebGPUExecutionOptions options);
+  ~ScopedWebGPUExecutionOptions();
+
+  ScopedWebGPUExecutionOptions(const ScopedWebGPUExecutionOptions&) = delete;
+  ScopedWebGPUExecutionOptions& operator=(const ScopedWebGPUExecutionOptions&) =
+      delete;
+  ScopedWebGPUExecutionOptions(ScopedWebGPUExecutionOptions&&) = delete;
+  ScopedWebGPUExecutionOptions& operator=(ScopedWebGPUExecutionOptions&&) =
+      delete;
+
+ private:
+  WebGPUExecutionOptions previous_;
+};
+
+template <typename Fn>
+decltype(auto) with_webgpu_execution_options(
+    WebGPUExecutionOptions options,
+    Fn&& fn) {
+  ScopedWebGPUExecutionOptions scope(options);
+  return std::forward<Fn>(fn)();
+}
+
+} // namespace executorch::backends::webgpu