[ExecuTorch][WebGPU] Add shared model runtime prerequisites

backends/vulkan/serialization/targets.bzl

@@ -26,6 +26,7 @@ def define_common_targets(is_fbcode = False):
         name = "vk_delegate_schema",
         srcs = [],
         visibility = [
+            "//executorch/backends/webgpu/...",
             "//executorch/backends/vulkan/...",
         ],
         exported_headers = {

backends/webgpu/BUCK

@@ -0,0 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target")
+
+oncall("executorch")
+
+load(":targets.bzl", "define_common_targets")
+
+non_fbcode_target(_kind = define_common_targets)
+
+fbcode_target(_kind = define_common_targets)

backends/webgpu/CMakeLists.txt

@@ -47,6 +47,17 @@ list(APPEND WEBGPU_SRCS ${WEBGPU_OP_SRCS})
 
 add_library(webgpu_backend ${WEBGPU_SRCS})
 
+add_library(webgpu_model_loader runner/webgpu_model_loader.cpp)
+target_include_directories(
+  webgpu_model_loader PUBLIC $<BUILD_INTERFACE:${EXECUTORCH_ROOT}/..>
+)
+target_link_libraries(webgpu_model_loader PUBLIC extension_module_static)
+target_compile_options(webgpu_model_loader PRIVATE -fexceptions)
+target_compile_definitions(
+  webgpu_model_loader PUBLIC C10_USING_CUSTOM_GENERATED_MACROS
+)
+set_property(TARGET webgpu_model_loader PROPERTY CXX_STANDARD 17)
+
 # Verify committed *_wgsl.h match their *.wgsl (drift fails the build).
 resolve_python_executable()
 add_custom_target(
@@ -165,8 +176,8 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST)
       PRIVATE "${EXECUTORCH_ROOT}/third-party/json/single_include"
     )
 
-    # Device-free util unit test: no backend/Dawn link (pure manifest/tolerance
-    # + dispatch-grid-math helpers), so it does NOT use the native-test helper.
+    # Device-free util unit test. WebGPUUtils.h needs Dawn declarations, but the
+    # test calls only pure helpers and does not link the Dawn implementation.
     add_executable(
       webgpu_op_test_util_test
       test/op_tests/test_driver_util.cpp test/op_tests/driver_util.cpp
@@ -178,7 +189,13 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST)
               "${EXECUTORCH_ROOT}/third-party/json/single_include"
     )
     target_link_libraries(
-      webgpu_op_test_util_test PRIVATE GTest::gtest GTest::gtest_main
+      webgpu_op_test_util_test PRIVATE executorch_core GTest::gtest
+                                       GTest::gtest_main
+    )
+    target_include_directories(
+      webgpu_op_test_util_test
+      PRIVATE
+        "$<TARGET_PROPERTY:dawn::webgpu_dawn,INTERFACE_INCLUDE_DIRECTORIES>"
     )
     target_compile_options(webgpu_op_test_util_test PRIVATE -fexceptions)
     set_property(TARGET webgpu_op_test_util_test PROPERTY CXX_STANDARD 17)
@@ -235,5 +252,16 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST)
       webgpu_compute_dispatch_test test/native/test_compute_dispatch.cpp
     )
     target_link_libraries(webgpu_compute_dispatch_test PRIVATE GTest::gtest)
+
+    add_executable(webgpu_model_loader_test test/native/test_model_loader.cpp)
+    target_include_directories(
+      webgpu_model_loader_test PRIVATE $<BUILD_INTERFACE:${EXECUTORCH_ROOT}/..>
+    )
+    target_link_libraries(
+      webgpu_model_loader_test PRIVATE webgpu_model_loader GTest::gtest
+                                       GTest::gtest_main
+    )
+    target_compile_options(webgpu_model_loader_test PRIVATE -fexceptions)
+    set_property(TARGET webgpu_model_loader_test PROPERTY CXX_STANDARD 17)
   endif()
 endif()

backends/webgpu/runner/webgpu_model_loader.cpp

@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runner/webgpu_model_loader.h>
+
+#include <unordered_set>
+#include <utility>
+
+namespace executorch::backends::webgpu {
+
+runtime::Result<std::unique_ptr<extension::Module>> load_webgpu_model(
+    WebGPUModelLoadSpec spec) {
+  if (spec.pte_path.empty() || spec.required_methods.empty()) {
+    return runtime::Error::InvalidArgument;
+  }
+  std::unordered_set<std::string> methods;
+  for (const auto& method : spec.required_methods) {
+    if (method.empty() || !methods.insert(method).second) {
+      return runtime::Error::InvalidArgument;
+    }
+  }
+  std::unordered_set<std::string> data_files;
+  for (const auto& path : spec.ptd_paths) {
+    if (path.empty() || path == spec.pte_path ||
+        !data_files.insert(path).second) {
+      return runtime::Error::InvalidArgument;
+    }
+  }
+
+  auto module = std::make_unique<extension::Module>(
+      spec.pte_path, std::move(spec.ptd_paths), spec.load_mode);
+  for (const auto& method : spec.required_methods) {
+    const runtime::Error error = module->load_method(method);
+    if (error != runtime::Error::Ok) {
+      return error;
+    }
+  }
+  return module;
+}
+
+} // namespace executorch::backends::webgpu

backends/webgpu/runner/webgpu_model_loader.h

@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <executorch/extension/module/module.h>
+#include <executorch/runtime/core/result.h>
+
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+struct WebGPUModelLoadSpec {
+  std::string pte_path;
+  std::vector<std::string> ptd_paths;
+  std::vector<std::string> required_methods;
+  extension::Module::LoadMode load_mode = extension::Module::LoadMode::File;
+};
+
+runtime::Result<std::unique_ptr<extension::Module>> load_webgpu_model(
+    WebGPUModelLoadSpec spec);
+
+} // namespace executorch::backends::webgpu

backends/webgpu/runtime/WebGPUBackend.cpp

@@ -44,6 +44,16 @@ using executorch::runtime::resize_tensor;
 using executorch::runtime::Result;
 using executorch::runtime::Span;
 
+namespace {
+thread_local WebGPUGraph* last_execution_graph = nullptr;
+} // namespace
+
+std::string webgpu_backend_execution_attestation_json() {
+  return last_execution_graph == nullptr
+      ? "{\"schemaVersion\":1,\"unavailable\":true}"
+      : last_execution_graph->execution_attestation_json();
+}
+
 Result<WebGPUGraphConfig> parse_webgpu_graph_config(
     ArrayRef<CompileSpec> compile_specs) {
   WebGPUGraphConfig config;
@@ -228,6 +238,7 @@ Error WebGPUBackend::execute(
   // the backend boundary.
   try {
     const WebGPUExecutionPlan plan = graph->make_execution_plan(graph_options);
+    last_execution_graph = graph;
     graph->execute(plan);
 
     // Copy outputs from GPU staging buffers to EValue tensor data pointers
@@ -242,7 +253,9 @@ Error WebGPUBackend::execute(
           {tensor.mutable_data_ptr(), tensor.nbytes(), host_is_fp32});
     }
     graph->copy_outputs(outputs, plan);
+    graph->complete_execution_attestation();
   } catch (const std::exception& e) {
+    graph->fail_execution_attestation(e.what());
     ET_LOG(Error, "WebGPU execute / output copy failed: %s", e.what());
     return Error::Internal;
   }
@@ -253,6 +266,9 @@ Error WebGPUBackend::execute(
 void WebGPUBackend::destroy(DelegateHandle* handle) const {
   if (handle != nullptr) {
     WebGPUGraph* graph = static_cast<WebGPUGraph*>(handle);
+    if (last_execution_graph == graph) {
+      last_execution_graph = nullptr;
+    }
     graph->~WebGPUGraph();
   }
 }

backends/webgpu/runtime/WebGPUBackend.h

@@ -19,6 +19,8 @@ executorch::runtime::Result<WebGPUGraphConfig> parse_webgpu_graph_config(
     executorch::runtime::ArrayRef<executorch::runtime::CompileSpec>
         compile_specs);
 
+std::string webgpu_backend_execution_attestation_json();
+
 class WebGPUBackend final : public ::executorch::runtime::BackendInterface {
  public:
   ~WebGPUBackend() override = default;

backends/webgpu/runtime/WebGPUDevice.cpp

@@ -1,3 +1,3 @@
 /*
  * Copyright (c) Meta Platforms, Inc. and affiliates.
  * All rights reserved.
@@ -9,6 +9,7 @@
 #include <executorch/backends/webgpu/runtime/WebGPUCompat.h>
 #include <executorch/backends/webgpu/runtime/WebGPUDevice.h>
 
+#include <atomic>
 #include <cstdio>
 #include <cstdlib>
 #include <memory>
@@ -184,16 +185,31 @@
 }
 
 namespace {
-WebGPUContext* g_default_context = nullptr;
+std::atomic<WebGPUContext*> g_default_context{nullptr};
 } // namespace
 
 void set_default_webgpu_context(WebGPUContext* ctx) {
-  g_default_context = ctx;
+  g_default_context.store(ctx, std::memory_order_release);
+}
+
+WebGPUContext* get_explicit_default_webgpu_context() {
+  return g_default_context.load(std::memory_order_acquire);
+}
+
+bool compare_and_set_default_webgpu_context(
+    WebGPUContext* expected,
+    WebGPUContext* desired) {
+  return g_default_context.compare_exchange_strong(
+      expected,
+      desired,
+      std::memory_order_acq_rel,
+      std::memory_order_acquire);
 }
 
 WebGPUContext* get_default_webgpu_context() {
-  if (g_default_context) {
-    return g_default_context;
+  if (WebGPUContext* explicit_context =
+          get_explicit_default_webgpu_context()) {
+    return explicit_context;
   }
 #if !defined(__EMSCRIPTEN__)
   // Native-only lazy process-wide context, mirroring Vulkan api::context().

backends/webgpu/runtime/WebGPUDevice.h

@@ -41,6 +41,13 @@ void destroy_webgpu_context(WebGPUContext& ctx);
 
 // Global context used by WebGPUGraph::build() when no device is pre-set.
 void set_default_webgpu_context(WebGPUContext* ctx);
+// Returns only a caller-installed context, never the native lazy fallback.
+WebGPUContext* get_explicit_default_webgpu_context();
+// Replaces the explicit context only when its current pointer equals expected.
+// Registration is non-owning; the caller keeps the installed context alive.
+bool compare_and_set_default_webgpu_context(
+    WebGPUContext* expected,
+    WebGPUContext* desired);
 WebGPUContext* get_default_webgpu_context();
 
 } // namespace webgpu