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

Code before:
      }

      // 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
    if (!found_gpu_tensor) {
      return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT,
                                   "No GPU tensor found in CopyTensors call - all tensors are CPU-only. "
                                   "This indicates an invalid call as CanCopy should have rejected this.");
    }

    // Initialize data_transfer if not already done or if context_id changed
    // Use mutex to ensure thread-safe lazy initialization
    {
      std::lock_guard<std::mutex> lock(impl.init_mutex_);

      if (impl.data_transfer_ == nullptr || impl.context_id_ != context_id) {
        impl.context_id_ = context_id;

        // Check if context exists, create a default one if it doesn't
        webgpu::WebGpuContext* context_ptr = nullptr;
        if (webgpu::WebGpuContextFactory::HasContext(context_id)) {
          context_ptr = &webgpu::WebGpuContextFactory::GetContext(context_id);
        } else {
          WebGpuContextParams params = GetDefaultWebGpuContextParams();
          params.context_config.context_id = context_id;
          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());
      }
    }  // Release lock

    // Now perform the actual tensor copy
    for (size_t idx = 0; idx < num_tensors; ++idx) {
      const OrtValue* src_tensor = src_tensors[idx];
      OrtValue* dst_tensor = dst_tensors[idx];

Review comment (bug):
I understand that the code are designed for situations that assume different context ID may be used for the same `WebGpuDataTransferImpl` instance. 

This causes the lock behavior to be complicated. for now, I think it is safe to assume context ID won't change: it should always be 0.

based on this, code may be simplified:

- use patterns like this to avoid unnecessary lock operation:
  ```c++
  if (impl.data_transfer_ == nullptr) {
    std::lock_guard<std::mutex> lock(impl.init_mutex_);
    if (impl.data_transfer_ == nullptr) {
      ...

      impl.data_transfer_ = ...;
    }
  }
  ```

- always create new context:
  ```c++
  context_ptr = &webgpu::WebGpuContextFactory::CreateContext(params.context_config);
  ```

  and in ReleaseImpl release it
  ```c++
  static void ReleaseImpl(OrtDataTransferImpl* this_ptr) noexcept {
    WebGpuDataTransferImpl* p_impl = static_cast<WebGpuDataTransferImpl*>(this_ptr);
    int context_id = p_impl->context_id_;
    bool data_transfer_initialized = false;
    {
      std::lock_guard<std::mutex> lock(p_impl->init_mutex_);
      data_transfer_initialized = p_impl->data_transfer_ == nullptr;
    }
    delete p_impl;
    if (data_transfer_initialized) {
      WebGpuContextFactory::ReleaseContext(context_id);
    }
  }
  ```

Code after:
        // 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());
      }
    }

    // Now perform the actual tensor copy
    for (size_t idx = 0; idx < num_tensors; ++idx) {
      const OrtValue* src_tensor = src_tensors[idx];
      OrtValue* dst_tensor = dst_tensors[idx];
      auto status = impl.data_transfer_->CopyTensor(src_tensor->Get<Tensor>(), *dst_tensor->GetMutable<Tensor>());
      if (!status.IsOK()) {
        return OrtApis::CreateStatus(ORT_RUNTIME_EXCEPTION, status.ErrorMessage().c_str());
      }
    }
    return nullptr;
  }

  static void ReleaseImpl(OrtDataTransferImpl* this_ptr) noexcept {
    auto* p_impl = static_cast<WebGpuDataTransferImpl*>(this_ptr);
    int context_id = p_impl->context_id_;
    bool data_transfer_initialized = false;
    {
      std::lock_guard<std::mutex> lock(p_impl->init_mutex_);
      data_transfer_initialized = (p_impl->data_transfer_ != nullptr);
    }
    delete p_impl;
    if (data_transfer_initialized) {
      webgpu::WebGpuContextFactory::ReleaseContext(context_id);
    }
  }

  const OrtApi& ort_api;
  const OrtEpApi& ep_api;
  std::unique_ptr<webgpu::DataTransfer> data_transfer_;  // Lazy-initialized
  int context_id_;                                       // Track which context we're using
  std::mutex init_mutex_;                                // Protects lazy initialization
};

OrtDataTransferImpl* OrtWebGpuCreateDataTransfer() {
  // Validate API version is supported
  const OrtApi* api = OrtApis::GetApi(ORT_API_VERSION);
  if (!api) {