Repository: microsoft/onnxruntime
File: onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc (C++)

Code before:
    if (dst_type == OrtMemoryInfoDeviceType_GPU) {
      uint32_t dst_vendor = impl.ep_api.MemoryDevice_GetVendorId(dst_memory_device);
      if (dst_vendor != 0) {
        return false;  // Not a WebGPU device
      }
    }

    // WebGPU supports GPU<->GPU, GPU<->CPU copies (where GPU has vendor ID 0)
    return (src_type == OrtMemoryInfoDeviceType_GPU && dst_type == OrtMemoryInfoDeviceType_GPU) ||
           (src_type == OrtMemoryInfoDeviceType_GPU && dst_type == OrtMemoryInfoDeviceType_CPU) ||
           (src_type == OrtMemoryInfoDeviceType_CPU && dst_type == OrtMemoryInfoDeviceType_GPU);
  }

  static OrtStatus* CopyTensorsImpl(OrtDataTransferImpl* this_ptr,
                                    const OrtValue** src_tensors,
                                    OrtValue** dst_tensors,
                                    OrtSyncStream** /*streams*/,
                                    size_t num_tensors) noexcept {
    auto& impl = *static_cast<WebGpuDataTransferImpl*>(this_ptr);

    if (num_tensors == 0) {
      return nullptr;
    }

    // Lazy initialization: Get context_id from the first GPU tensor's device
    int context_id = 0;  // Default to context_id 0
    bool found_gpu_tensor = false;

    // Check both src_tensors and dst_tensors to find the first GPU tensor
    for (size_t idx = 0; idx < num_tensors && !found_gpu_tensor; ++idx) {
      // Check source tensor
      const OrtMemoryDevice* src_device = impl.ep_api.Value_GetMemoryDevice(src_tensors[idx]);
      OrtMemoryInfoDeviceType src_device_type = impl.ep_api.MemoryDevice_GetDeviceType(src_device);
      if (src_device_type == OrtMemoryInfoDeviceType_GPU) {
        context_id = static_cast<int>(impl.ep_api.MemoryDevice_GetDeviceId(src_device));
        found_gpu_tensor = true;
        break;
      }

      // Check destination tensor
      const OrtMemoryDevice* dst_device = impl.ep_api.Value_GetMemoryDevice(dst_tensors[idx]);
      OrtMemoryInfoDeviceType dst_device_type = impl.ep_api.MemoryDevice_GetDeviceType(dst_device);
      if (dst_device_type == OrtMemoryInfoDeviceType_GPU) {
        context_id = static_cast<int>(impl.ep_api.MemoryDevice_GetDeviceId(dst_device));
        found_gpu_tensor = true;
        break;
      }
    }

    // If no GPU tensor found, return an error as this indicates an invalid state
    // CanCopy should have rejected CPU-only tensor copies

Review comment (suggestion):
About this context ID:

in my understanding, when an instance of `WebGpuDataTransferImpl` is created, its context ID is always determined. For the usage of the current PR, it must always be `0`.

Checks in `CanCopyImpl` and `CopyTensorsImpl` should be able to simplified according to this:

in `CanCopyImpl`, any input on GPU must have the same device ID.
in `CopyTensorsImpl` those code can be removed.

Code after:
    if (dst_type == OrtMemoryInfoDeviceType_GPU) {
      uint32_t dst_vendor = impl.ep_api.MemoryDevice_GetVendorId(dst_memory_device);
      if (dst_vendor != 0) {
        return false;  // Not a WebGPU device
      }
    }

    // If both are GPU, they must have the same device ID
    if (src_type == OrtMemoryInfoDeviceType_GPU && dst_type == OrtMemoryInfoDeviceType_GPU) {
      uint64_t src_device_id = impl.ep_api.MemoryDevice_GetDeviceId(src_memory_device);
      uint64_t dst_device_id = impl.ep_api.MemoryDevice_GetDeviceId(dst_memory_device);
      if (src_device_id != dst_device_id) {
        return false;  // Cannot copy between different devices
      }
    }

    // WebGPU supports GPU<->GPU, GPU<->CPU copies (where GPU has vendor ID 0)
    return (src_type == OrtMemoryInfoDeviceType_GPU && dst_type == OrtMemoryInfoDeviceType_GPU) ||
           (src_type == OrtMemoryInfoDeviceType_GPU && dst_type == OrtMemoryInfoDeviceType_CPU) ||
           (src_type == OrtMemoryInfoDeviceType_CPU && dst_type == OrtMemoryInfoDeviceType_GPU);
  }

  static OrtStatus* CopyTensorsImpl(OrtDataTransferImpl* this_ptr,
                                    const OrtValue** src_tensors,
                                    OrtValue** dst_tensors,
                                    OrtSyncStream** /*streams*/,
                                    size_t num_tensors) noexcept {
    auto& impl = *static_cast<WebGpuDataTransferImpl*>(this_ptr);

    if (num_tensors == 0) {
      return nullptr;
    }

    // Lazy initialization: Use double-checked locking to avoid unnecessary lock operations
    if (impl.data_transfer_ == nullptr) {
      std::lock_guard<std::mutex> lock(impl.init_mutex_);
      if (impl.data_transfer_ == nullptr) {
        // Always create a new context with context_id 0
        WebGpuContextParams params = GetDefaultWebGpuContextParams();
        params.context_config.context_id = impl.context_id_;
        auto* context_ptr = &webgpu::WebGpuContextFactory::CreateContext(params.context_config);
        context_ptr->Initialize(params.buffer_cache_config, params.backend_type, params.enable_pix_capture);

        // Create the DataTransfer instance
        // Note: The DataTransfer holds a const reference to BufferManager. The BufferManager's lifecycle
        // is managed by the WebGpuContext, which is stored in a static WebGpuContextFactory and persists
        // for the lifetime of the application, ensuring the reference remains valid.
        impl.data_transfer_ = std::make_unique<webgpu::DataTransfer>(context_ptr->BufferManager());
      }
    }