[WebGPU] Expose safe graph capture I/O in Python

onnxruntime/core/providers/webgpu/data_transfer.cc

@@ -36,12 +36,24 @@ common::Status DataTransferImpl::CopyTensor(void const* src_data,
   return Status::OK();
 }
 
-bool DataTransfer::CanCopy(const OrtDevice& src_device, const OrtDevice& dst_device) const {
+bool DataTransfer::IsSupportedDevicePair(const OrtDevice& src_device, const OrtDevice& dst_device) {
+  // WebGPU allocations use VendorIds::NONE; reject foreign GPU handles from this transfer.
+  if (src_device.Type() == OrtDevice::GPU && src_device.Vendor() != OrtDevice::VendorIds::NONE) {
+    return false;
+  }
+  if (dst_device.Type() == OrtDevice::GPU && dst_device.Vendor() != OrtDevice::VendorIds::NONE) {
+    return false;
+  }
+
   return (dst_device.Type() == OrtDevice::GPU && src_device.Type() == OrtDevice::CPU) ||
          (dst_device.Type() == OrtDevice::GPU && src_device.Type() == OrtDevice::GPU) ||
          (dst_device.Type() == OrtDevice::CPU && src_device.Type() == OrtDevice::GPU);
 }
 
+bool DataTransfer::CanCopy(const OrtDevice& src_device, const OrtDevice& dst_device) const {
+  return IsSupportedDevicePair(src_device, dst_device);
+}
+
 common::Status DataTransfer::CopyTensor(const Tensor& src, Tensor& dst) const {
   return impl_.CopyTensor(src.DataRaw(),
                           src.Location().device.Type() == OrtDevice::GPU,

onnxruntime/core/providers/webgpu/data_transfer.h

@@ -33,6 +33,10 @@ class DataTransfer : public IDataTransfer {
   DataTransfer(const BufferManager& buffer_manager) : impl_{buffer_manager} {};
   ~DataTransfer() {};
 
+  // Device-compatibility predicate for WebGPU copies. Does not touch the BufferManager, so it can
+  // be exercised without a live device.
+  static bool IsSupportedDevicePair(const OrtDevice& src_device, const OrtDevice& dst_device);
+
   bool CanCopy(const OrtDevice& src_device, const OrtDevice& dst_device) const override;
 
   common::Status CopyTensor(const Tensor& src, Tensor& dst) const override;

onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc

@@ -399,12 +399,38 @@ struct WebGpuDataTransferImpl : OrtDataTransferImpl {
            (src_type == OrtMemoryInfoDeviceType_CPU && dst_type == OrtMemoryInfoDeviceType_GPU);
   }
 
+  // BufferManager reports misuse by throwing (ORT_ENFORCE) rather than by returning a Status: an
+  // aliased src/dst pair, an undersized destination and a still-mapped buffer all throw. This
+  // callback is invoked across the C ABI and is therefore noexcept, so letting the exception escape
+  // would terminate the process instead of surfacing an error. Convert any throw to an OrtStatus.
   static OrtStatus* ORT_API_CALL CopyTensorsImpl(
       OrtDataTransferImpl* this_ptr,
       const OrtValue** src_tensors,
       OrtValue** dst_tensors,
-      OrtSyncStream** /*streams*/,
+      OrtSyncStream** streams,
       size_t num_tensors) noexcept {
+    OrtStatus* result = nullptr;
+    ORT_TRY {
+      result = CopyTensorsOrThrow(this_ptr, src_tensors, dst_tensors, streams, num_tensors);
+    }
+    ORT_CATCH(const std::exception& ex) {
+      ORT_HANDLE_EXCEPTION([&]() {
+        result = OrtApis::CreateStatus(ORT_RUNTIME_EXCEPTION, ex.what());
+      });
+    }
+    ORT_CATCH(...) {
+      result = OrtApis::CreateStatus(ORT_RUNTIME_EXCEPTION, "Unknown exception during WebGPU CopyTensors.");
+    }
+    return result;
+  }
+
+  // Body of CopyTensorsImpl. May throw; the caller is responsible for converting.
+  static OrtStatus* CopyTensorsOrThrow(
+      OrtDataTransferImpl* this_ptr,
+      const OrtValue** src_tensors,
+      OrtValue** dst_tensors,
+      OrtSyncStream** /*streams*/,
+      size_t num_tensors) {
     auto& impl = *static_cast<WebGpuDataTransferImpl*>(this_ptr);
 
     if (num_tensors == 0) {

onnxruntime/python/onnxruntime_pybind_iobinding.cc

@@ -21,6 +21,14 @@
 namespace py = pybind11;
 
 namespace {
+const char* GetBindingAllocatorName(const OrtDevice& device) {
+  if (device.Type() == OrtDevice::GPU &&
+      device.Vendor() == OrtDevice::VendorIds::NONE) {
+    return WEBGPU_BUFFER;
+  }
+  return GetDeviceName(device);
+}
+
 void BindOutput(SessionIOBinding* io_binding, const std::string& name, const OrtDevice& device,
                 MLDataType element_type, const std::vector<int64_t>& shape, int64_t data_ptr) {
   ORT_ENFORCE(data_ptr != 0, "Pointer to data memory is not valid");
@@ -43,7 +51,7 @@
   }
 
   OrtValue ml_value;
-  OrtMemoryInfo info(GetDeviceName(device), OrtDeviceAllocator, device);
+  OrtMemoryInfo info(GetBindingAllocatorName(device), OrtDeviceAllocator, device);
   Tensor::InitOrtValue(element_type, gsl::make_span(shape), reinterpret_cast<void*>(data_ptr), info, ml_value);
 
   auto status = io_binding->Get()->BindOutput(name, ml_value);
@@ -57,9 +65,10 @@
   py::class_<SessionIOBinding> session_io_binding(m, "SessionIOBinding");
   session_io_binding
       .def(py::init([](PyInferenceSession* sess) {
-        auto sess_io_binding = std::make_unique<SessionIOBinding>(sess->GetSessionHandle());
-        return sess_io_binding;
-      }))
+             auto sess_io_binding = std::make_unique<SessionIOBinding>(sess->GetSessionHandle());
+             return sess_io_binding;
+           }),
+           py::keep_alive<1, 2>())
       // May create Tensor/Sequence based OrtValues. Use bind_ortvalue_input for universal binding.
       .def("bind_input", [](SessionIOBinding* io_binding, const std::string& name, py::object& arr_on_cpu) -> void {
         InferenceSession* sess = io_binding->GetInferenceSession();
@@ -100,7 +109,7 @@
         }
         auto ml_type = OnnxTypeToOnnxRuntimeTensorType(element_type);
         OrtValue ml_value;
-        OrtMemoryInfo info(GetDeviceName(device), OrtDeviceAllocator, device);
+        OrtMemoryInfo info(GetBindingAllocatorName(device), OrtDeviceAllocator, device);
         Tensor::InitOrtValue(ml_type, gsl::make_span(shape), reinterpret_cast<void*>(data_ptr), info, ml_value);
 
         auto status = io_binding->Get()->BindInput(name, ml_value);
@@ -117,7 +126,7 @@
         int type_num = dtype->type_num;
         Py_DECREF(dtype);
 
-        OrtMemoryInfo info(GetDeviceName(device), OrtDeviceAllocator, device);
+        OrtMemoryInfo info(GetBindingAllocatorName(device), OrtDeviceAllocator, device);
         auto ml_type = NumpyTypeToOnnxRuntimeTensorType(type_num);
         // See comment in the int32_t element_type overload above: string tensors are not safe
         // to bind via a raw, non-owning pointer because no std::string objects are constructed

onnxruntime/python/onnxruntime_pybind_ortvalue.cc

@@ -19,6 +19,11 @@
 namespace py = pybind11;
 
 namespace {
+bool IsWebGpuBuffer(const OrtValue& ort_value) {
+  return ort_value.IsTensor() &&
+         ort_value.Get<Tensor>().Location().name == WEBGPU_BUFFER;
+}
+
 std::unique_ptr<OrtValue> OrtValueFromShapeAndType(const std::vector<int64_t>& shape,
                                                    MLDataType element_type,
                                                    const OrtDevice& device) {
@@ -376,6 +381,9 @@
       })
       .def("device_name", [](const OrtValue* ort_value) -> std::string {
         if (ort_value->IsTensor()) {
+          if (IsWebGpuBuffer(*ort_value)) {
+            return "webgpu";
+          }
           return std::string(GetDeviceName(ort_value->Get<Tensor>().Location().device));
         }
 #if !defined(DISABLE_SPARSE_TENSORS)
@@ -446,6 +454,7 @@
       .def("is_tensor", [](const OrtValue* ort_value) -> bool { return ort_value->IsTensor(); })
       .def("is_sparse_tensor", [](const OrtValue* ort_value) -> bool { return ort_value->IsSparseTensor(); })
       .def("is_tensor_sequence", [](const OrtValue* ort_value) -> bool { return ort_value->IsTensorSequence(); })
+      .def("_is_webgpu_buffer", [](const OrtValue* ort_value) -> bool { return IsWebGpuBuffer(*ort_value); })
       // Converts Tensor into a numpy array
       .def("numpy", [](const OrtValue* ml_value) -> py::object {
         ORT_ENFORCE(ml_value->IsTensor(), "Only OrtValues that are Tensors are convertible to Numpy objects");
@@ -496,17 +505,22 @@
 #endif
       })
 #if defined(ENABLE_DLPACK)
-      .def("to_dlpack", [](OrtValue* ort_value) -> py::object { return py::reinterpret_steal<py::object>(ToDlpack(*ort_value)); },
+      .def("to_dlpack", [](OrtValue* ort_value) -> py::object {
+        ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues.");
+        return py::reinterpret_steal<py::object>(ToDlpack(*ort_value)); },
            "Returns a DLPack representing the tensor. This method does not copy the pointer shape, "
            "instead, it copies the pointer value. The OrtValue must be persist until the dlpack structure "
            "is consumed.")
       .def_static("from_dlpack", [](py::object data, bool is_bool_tensor) { return FromDlpack(data.ptr(), is_bool_tensor); }, py::arg("data"), py::arg("is_bool_tensor") = false, "Converts a tensor from a external library into an OrtValue by means of the __dlpack__ protocol.")
-      .def("__dlpack__", [](OrtValue* ort_value, py::object /* stream */) -> py::object { return py::reinterpret_steal<py::object>(ToDlpack(*ort_value)); }, py::arg("stream") = py::none(),
+      .def("__dlpack__", [](OrtValue* ort_value, py::object /* stream */) -> py::object {
+        ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues.");
+        return py::reinterpret_steal<py::object>(ToDlpack(*ort_value)); }, py::arg("stream") = py::none(),
            "Returns a DLPack representing the tensor (part of __dlpack__ protocol). "
            "This method does not copy the pointer shape, instead, it copies the pointer value. "
            "The OrtValue must persist until the dlpack structure is consumed.")
       .def("__dlpack_device__", [](const OrtValue* ort_value) -> py::tuple {
             ORT_ENFORCE(ort_value->IsTensor(), "Only tensor type OrtValues are supported");
+            ORT_ENFORCE(!IsWebGpuBuffer(*ort_value), "DLPack export is not supported for WebGPU OrtValues.");
             const onnxruntime::Tensor& tensor = ort_value->Get<Tensor>();
             DLDevice device = onnxruntime::dlpack::GetDlpackDevice(*ort_value, tensor.Location().device.Id());
             return py::make_tuple(static_cast<int>(device.device_type), device.device_id); }, "Returns a tuple of integers, (device, device index) (part of __dlpack__ protocol).")
@@ -516,6 +530,16 @@
   py::class_<std::vector<OrtValue>>(m, "OrtValueVector")
       .def(py::init<>())
       .def("push_back", [](std::vector<OrtValue>* v, const OrtValue& ortvalue) {
+        // A standalone OrtValueVector has no parent object, so unlike the vector returned by
+        // SessionIOBinding::get_outputs it cannot keep the owning session alive. A WebGPU buffer
+        // allocated by a session's EP is freed through GpuBufferAllocator, whose buffer-manager
+        // getter captures that EP, so releasing it after the session is gone is a use-after-free.
+        // TODO: remove once GpuBufferAllocator owns a reference to whatever owns its BufferManager,
+        // at which point a WebGPU OrtValue is safe to hold independently of any session.
+        ORT_ENFORCE(!IsWebGpuBuffer(ortvalue),
+                    "A WebGPU OrtValue cannot be stored in a standalone OrtValueVector because the "
+                    "vector cannot keep the session that allocated the buffer alive. Use IOBinding, "
+                    "or IOBinding.get_outputs_as_ortvaluevector() for device-resident outputs.");
         v->push_back(ortvalue);
       })
 #if defined(ENABLE_DLPACK)
@@ -537,6 +561,14 @@
 
               auto ml_type = NumpyTypeToOnnxRuntimeTensorType(type_num);
               auto device = devices.at(i);
+              // This overload wraps foreign (PyTorch) storage by raw address. A WebGPU allocation is
+              // an opaque WGPUBuffer handle rather than an addressable pointer, so no torch tensor
+              // can back one. The OrtMemoryInfo built below would also be mislabelled, since it uses
+              // GetDeviceName rather than the WEBGPU_BUFFER allocator name the binding paths expect.
+              ORT_ENFORCE(!(device.Type() == OrtDevice::GPU &&
+                            device.Vendor() == OrtDevice::VendorIds::NONE),
+                          "OrtValueVector.push_back_batch cannot wrap WebGPU memory: a WebGPU "
+                          "allocation is an opaque buffer handle, not a raw data pointer.");
               OrtMemoryInfo info(GetDeviceName(device), OrtDeviceAllocator, device);
               OrtValue ml_value;
               Tensor::InitOrtValue(ml_type, gsl::make_span(shape), reinterpret_cast<void*>(data_ptr), info, ml_value);
@@ -547,7 +579,7 @@
       .def("shrink_to_fit", [](std::vector<OrtValue>* v) { v->shrink_to_fit(); })
       .def("__len__", [](const std::vector<OrtValue>& v) { return v.size(); })
       .def("__iter__", [](const std::vector<OrtValue>& v) { return py::make_iterator(v.cbegin(), v.cend()); }, py::keep_alive<0, 1>())
-      .def("__getitem__", [](const std::vector<OrtValue>& v, const size_t idx) { return v.at(idx); })
+      .def("__getitem__", [](const std::vector<OrtValue>& v, const size_t idx) { return v.at(idx); }, py::keep_alive<0, 1>())
       .def("bool_tensor_indices", [](std::vector<OrtValue>* v) -> std::vector<int64_t> {
             std::vector<int64_t> indices;
             for (size_t i = 0; i < v->size(); ++i) {
@@ -561,7 +593,9 @@
            "If torch consumes the dlpack structure, `.to(torch.bool)` must be applied to the torch tensor "
            "to get a boolean tensor.")
 #if defined(ENABLE_DLPACK)
-      .def("dlpack_at", [](std::vector<OrtValue>* v, const size_t idx) { return py::reinterpret_steal<py::object>(ToDlpack(v->at(idx))); })
+      .def("dlpack_at", [](std::vector<OrtValue>* v, const size_t idx) {
+        ORT_ENFORCE(!IsWebGpuBuffer(v->at(idx)), "DLPack export is not supported for WebGPU OrtValues.");
+        return py::reinterpret_steal<py::object>(ToDlpack(v->at(idx))); })
 #endif
       .def("element_type_at", [](std::vector<OrtValue>* v, const size_t idx) -> int32_t { return GetTensorProtoType(v->at(idx)); },
            "Returns an integer equal to the ONNX proto type of the tensor at position i. "
@@ -583,6 +617,7 @@
               DLManagedTensor* dlmanaged_tensor;
 
               for (auto it : v) {
+                ORT_ENFORCE(!IsWebGpuBuffer(it), "DLPack export is not supported for WebGPU OrtValues.");
                 dlmanaged_tensor = dlpack::OrtValueToDlpack(it);
                 py::capsule capsule(dlmanaged_tensor, "dltensor", DlpackCapsuleDestructor);
                 list_dlpacks.append(capsule);
@@ -595,6 +630,7 @@
               for (auto it : v) {
                 // A new instance of dlpack needs to be created. The object which consumes it
                 // is responsible for its deletion.
+                ORT_ENFORCE(!IsWebGpuBuffer(it), "DLPack export is not supported for WebGPU OrtValues.");
                 dlmanaged_tensor = dlpack::OrtValueToDlpack(it);
                 if (capsule == NULL) {
                   capsule = PyCapsule_New(dlmanaged_tensor, "dltensor", NULL);