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

onnxruntime/core/providers/webgpu/data_transfer.cc

@@ -37,6 +37,14 @@ common::Status DataTransferImpl::CopyTensor(void const* src_data,
 }
 
 bool DataTransfer::CanCopy(const OrtDevice& src_device, const OrtDevice& dst_device) const {
+  // 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);

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 python {
 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) {
@@ -364,6 +369,8 @@ void addOrtValueMethods(pybind11::module& m) {
       .def("data_ptr", [](OrtValue* ml_value) -> uintptr_t {
         // TODO: Assumes that the OrtValue is a Tensor, make this generic to handle non-Tensors
         ORT_ENFORCE(ml_value->IsTensor(), "Only OrtValues that are Tensors are currently supported");
+        ORT_ENFORCE(!IsWebGpuBuffer(*ml_value),
+                    "WebGPU OrtValues contain opaque buffer handles and do not expose a data pointer.");
 
         auto* tensor = ml_value->GetMutable<Tensor>();
 
@@ -376,6 +383,9 @@ void addOrtValueMethods(pybind11::module& m) {
       })
       .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,9 +456,12 @@ void addOrtValueMethods(pybind11::module& m) {
       .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");
+        ORT_ENFORCE(!IsWebGpuBuffer(*ml_value),
+                    "WebGPU OrtValues require explicit readback with IOBinding.copy_outputs_to_cpu().");
         [[maybe_unused]] const auto& device = ml_value->Get<Tensor>().Location().device;
 #ifdef _MSC_VER
 // The switch statement may only contain the 'default' label. In such a case, the MSVC compiler
@@ -496,17 +509,22 @@ void addOrtValueMethods(pybind11::module& m) {
 #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 +534,8 @@ void addOrtValueMethods(pybind11::module& m) {
   py::class_<std::vector<OrtValue>>(m, "OrtValueVector")
       .def(py::init<>())
       .def("push_back", [](std::vector<OrtValue>* v, const OrtValue& ortvalue) {
+        ORT_ENFORCE(!IsWebGpuBuffer(ortvalue),
+                    "OrtValueVector cannot retain the session ownership required by WebGPU OrtValues.");
         v->push_back(ortvalue);
       })
 #if defined(ENABLE_DLPACK)
@@ -537,6 +557,9 @@ void addOrtValueMethods(pybind11::module& m) {
 
               auto ml_type = NumpyTypeToOnnxRuntimeTensorType(type_num);
               auto device = devices.at(i);
+              ORT_ENFORCE(!(device.Type() == OrtDevice::GPU &&
+                            device.Vendor() == OrtDevice::VendorIds::NONE),
+                          "OrtValueVector.push_back_batch does not support WebGPU buffers.");
               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 +570,7 @@ void addOrtValueMethods(pybind11::module& m) {
       .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 +584,9 @@ void addOrtValueMethods(pybind11::module& m) {
            "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 +608,7 @@ void addOrtValueMethods(pybind11::module& m) {
               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 +621,7 @@ void addOrtValueMethods(pybind11::module& m) {
               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);

onnxruntime/python/onnxruntime_pybind_state.cc

@@ -26,6 +26,7 @@
 #include "core/framework/provider_options_utils.h"
 #include "core/framework/random_seed.h"
 #include "core/framework/sparse_tensor.h"
+#include "core/framework/tensor.h"
 #include "core/framework/tensorprotoutils.h"
 #include "core/framework/TensorSeq.h"
 #include "core/graph/graph_viewer.h"
@@ -91,6 +92,65 @@ namespace {
 
 constexpr std::string_view kEpCudaProviderOptionPrefix{"ep.cuda."};
 
+std::unique_ptr<OrtValue> CreateSessionOrtValue(PyInferenceSession* session,
+                                                const std::vector<int64_t>& shape,
+                                                MLDataType element_type,
+                                                const OrtDevice& device) {
+  // Accessing session state enforces that initialization has completed.
+  (void)session->GetSessionHandle()->GetSessionState();
+
+  const char* allocator_name = device.Type() == OrtDevice::GPU &&
+                                       device.Vendor() == OrtDevice::VendorIds::NONE
+                                   ? WEBGPU_BUFFER
+                                   : GetDeviceName(device);
+  OrtMemoryInfo memory_info{allocator_name, OrtDeviceAllocator, device};
+  auto allocator = session->GetSessionHandle()->GetAllocator(memory_info);
+  if (!allocator) {
+    throw std::runtime_error("No session allocator found for " + device.ToString());
+  }
+
+  auto ort_value = std::make_unique<OrtValue>();
+  Tensor::InitOrtValue(element_type, gsl::make_span(shape), std::move(allocator), *ort_value);
+  return ort_value;
+}
+
+void CopySessionOrtValue(PyInferenceSession* session, OrtValue& destination, const OrtValue& source) {
+  if (!destination.IsTensor() || !source.IsTensor()) {
+    throw std::runtime_error("Inplace update of OrtValues is only supported for Tensors");
+  }
+
+  const auto& source_tensor = source.Get<Tensor>();
+  auto* destination_tensor = destination.GetMutable<Tensor>();
+  const auto& source_location = source_tensor.Location();
+  const auto& destination_location = destination_tensor->Location();
+  const bool source_is_webgpu = source_location.name == WEBGPU_BUFFER;
+  const bool destination_is_webgpu = destination_location.name == WEBGPU_BUFFER;
+
+  if (source_location.device.Type() == OrtDevice::GPU &&
+      destination_location.device.Type() == OrtDevice::GPU &&
+      source_is_webgpu != destination_is_webgpu) {
+    throw std::runtime_error("WebGPU OrtValues cannot be copied to or from another GPU provider");
+  }
+
+  if (destination_tensor->DataType() != source_tensor.DataType()) {
+    throw std::runtime_error("The source and destination OrtValues must have the same data type");
+  }
+
+  if (destination_tensor->Shape().Size() != source_tensor.Shape().Size()) {
+    throw std::runtime_error("The source and destination OrtValues must have the same size");
+  }
+
+  py::gil_scoped_release release;
+  ORT_THROW_IF_ERROR(session->GetSessionHandle()->GetDataTransferManager().CopyTensor(
+      source_tensor, *destination_tensor));
+}
+
+void UpdateSessionOrtValue(PyInferenceSession* session, OrtValue& destination, const py::array& source) {
+  OrtValue source_value;
+  CreateGenericMLValue(nullptr, GetAllocator(), "", source, &source_value, true);
+  CopySessionOrtValue(session, destination, source_value);
+}
+
 struct AdaptedProviderOptions {
   std::vector<std::string> keys;
   std::vector<std::string> values;
@@ -2985,6 +3045,36 @@ including arg name, arg type (contains both type and shape).)pbdoc")
       })
       .def("get_providers", [](const PyInferenceSession* sess) -> const std::vector<std::string>& { return sess->GetSessionHandle()->GetRegisteredProviderTypes(); }, py::return_value_policy::reference_internal)
       .def("get_provider_options", [](const PyInferenceSession* sess) -> const ProviderOptionsMap& { return sess->GetSessionHandle()->GetAllProviderOptions(); }, py::return_value_policy::reference_internal)
+      .def("is_webgpu_graph_capture_enabled", [](const PyInferenceSession* sess) {
+        const auto* webgpu_ep =
+            sess->GetSessionHandle()->GetExecutionProviders().Get(kWebGpuExecutionProvider);
+        return webgpu_ep != nullptr && webgpu_ep->IsGraphCaptureEnabled(); })
+      .def("create_ortvalue_from_shape_and_type", [](PyInferenceSession* sess, const std::vector<int64_t>& shape, py::object& numpy_element_type, const OrtDevice& device) {
+        PyArray_Descr* dtype;
+        if (!PyArray_DescrConverter(numpy_element_type.ptr(), &dtype)) {
+          throw std::runtime_error("Not a valid numpy type");
+        }
+
+        int type_num = dtype->type_num;
+        Py_DECREF(dtype);
+        if (!IsNumericNumpyType(type_num)) {
+          throw std::runtime_error("Creation of OrtValues is currently only supported for numeric tensor types");
+        }
+
+        py::gil_scoped_release release;
+        return CreateSessionOrtValue(sess, shape, NumpyTypeToOnnxRuntimeTensorType(type_num), device); }, py::keep_alive<0, 1>())
+      .def("create_ortvalue_from_shape_and_onnx_type", [](PyInferenceSession* sess, const std::vector<int64_t>& shape, int32_t onnx_element_type, const OrtDevice& device) {
+        if (onnx_element_type == ONNX_NAMESPACE::TensorProto_DataType_STRING) {
+          throw std::runtime_error("Creation of OrtValues is currently only supported for numeric tensor types");
+        }
+
+        py::gil_scoped_release release;
+        return CreateSessionOrtValue(sess, shape, OnnxTypeToOnnxRuntimeTensorType(onnx_element_type), device); }, py::keep_alive<0, 1>())
+      .def("update_ortvalue_inplace", [](PyInferenceSession* sess, OrtValue& destination, const py::array& source) { UpdateSessionOrtValue(sess, destination, source); })
+      .def("update_ortvalue_inplace", [](PyInferenceSession* sess, OrtValue& destination, const OrtValue& source) { CopySessionOrtValue(sess, destination, source); })
+      .def("release_captured_graph", [](PyInferenceSession* sess, int graph_annotation_id) {
+        py::gil_scoped_release release;
+        OrtPybindThrowIfError(sess->GetSessionHandle()->ReleaseCapturedGraph(graph_annotation_id)); })
       .def("get_provider_graph_assignment_info", [](const PyInferenceSession* sess) -> const std::vector<const OrtEpAssignedSubgraph*>& {
 #if !defined(ORT_MINIMAL_BUILD)
         const auto* inference_session = sess->GetSessionHandle();