From e8a7e33b53d0cfa79de2ac384117f5cd84db8c2c Mon Sep 17 00:00:00 2001
From: Wei Wang <[email protected]>
Date: Thu, 28 May 2026 10:50:01 -0700
Subject: [PATCH] [WebNN] Guard TensorImplOrt::AsSpan() against non-CPU device tensors

For non-CPU execution providers (e.g. WebGPU EP), ORT's
GetTensorMutableData() returns an opaque device handle rather than a
CPU-dereferenceable address. Dereferencing it as a host pointer is
undefined behavior.

Thia CL adds a `can_access_on_cpu_` member to TensorImplOrt and CHECK
it in AsSpan() to prevent accidental CPU access to device tensor
memory.

Bug: 513177237
Change-Id: I19b9465dce57775f01fe01f2fbaaefbac5ffa9e5
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/7882779
Commit-Queue: Hu, Ningxin <[email protected]>
Reviewed-by: Hu, Ningxin <[email protected]>
Reviewed-by: Reilly Grant <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1637810}
---

diff --git a/services/webnn/ort/tensor_impl_ort.cc b/services/webnn/ort/tensor_impl_ort.cc
index e79a0f7..64e482ec 100644
--- a/services/webnn/ort/tensor_impl_ort.cc
+++ b/services/webnn/ort/tensor_impl_ort.cc
@@ -26,7 +26,8 @@
     : WebNNTensorImpl(std::move(receiver), context, std::move(tensor_info)),
       device_allocator_((std::move(device_allocator))),
       tensor_(std::move(tensor)),
-      size_(size) {
+      size_(size),
+      can_access_on_cpu_(can_access_on_cpu) {
   // Initialize the tensor with zeros, otherwise, reading uninitialized memory
   // will get random values.
   // TODO(crbug.com/461303833): check whether fast HW clears can be used
@@ -54,6 +55,10 @@
 
 base::span<uint8_t> TensorImplOrt::AsSpan() const {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
+  // For non-CPU device tensors (e.g. WebGPU EP), `GetTensorMutableData()`
+  // returns an opaque device handle, not a CPU-dereferenceable address.
+  // Any read or write through it is undefined behavior.
+  CHECK(can_access_on_cpu_);
 
   void* ort_tensor_raw_data = nullptr;
   CHECK_STATUS(
diff --git a/services/webnn/ort/tensor_impl_ort.h b/services/webnn/ort/tensor_impl_ort.h
index 9d47ac0..1e0d7f6 100644
--- a/services/webnn/ort/tensor_impl_ort.h
+++ b/services/webnn/ort/tensor_impl_ort.h
@@ -58,6 +58,10 @@
   scoped_refptr<DeviceAllocator> device_allocator_;
   const ScopedOrtValue tensor_ GUARDED_BY_CONTEXT(sequence_checker_);
   const size_t size_;
+  // Whether `tensor_`'s backing memory is CPU-accessible. When false (e.g.
+  // WebGPU EP device tensors), `AsSpan()` must not be called because
+  // `GetTensorMutableData()` returns a device handle, not a host pointer.
+  const bool can_access_on_cpu_ = true;
 };
 
 }  // namespace webnn::ort