[WebGPU] Batch input uploads for built-in and plugin EPs

onnxruntime/core/providers/webgpu/buffer_manager.cc

@@ -533,6 +533,62 @@ void BufferManager::Upload(CommandRecordingState& recording, void* src, WGPUBuff
   ORT_THROW_IF_ERROR(context_.Flush(*this, recording));
 }
 
+void BufferManager::UploadBatch(CommandRecordingState& recording, gsl::span<const UploadInfo> uploads) const {
+  // Bound staging memory while amortizing buffer creation and queue submission across small inputs.
+  constexpr size_t kBatchBytes = 16 * 1024 * 1024;
+  size_t begin = 0;
+  while (begin < uploads.size()) {
+    size_t end = begin;
+    size_t bytes = 0;
+    while (end < uploads.size()) {
+      const auto aligned_size = NormalizeCopySize(uploads[end].size);
+      if (aligned_size > kBatchBytes - bytes) break;
+      bytes += aligned_size;
+      ++end;
+    }
+    if (end == begin) {
+      const auto& upload = uploads[begin++];
+      Upload(recording, const_cast<void*>(upload.source), upload.target, upload.size);
+      continue;
+    }
+    if (bytes != 0) {
+      wgpu::BufferDescriptor desc{};
+      desc.size = bytes;
+      desc.usage = wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::MapWrite;
+      desc.mappedAtCreation = true;
+      auto staging_buffer = context_.Device().CreateBuffer(&desc);
+      auto* mapped = static_cast<uint8_t*>(staging_buffer.GetMappedRange());
+      size_t offset = 0;
+      for (size_t i = begin; i < end; ++i) {
+        const auto& upload = uploads[i];
+        const auto aligned_size = NormalizeCopySize(upload.size);
+        if (upload.size != 0) {
+          memcpy(mapped + offset, upload.source, upload.size);
+          memset(mapped + offset + upload.size, 0, aligned_size - upload.size);
+        }
+        offset += aligned_size;
+      }
+      staging_buffer.Unmap();
+      ORT_THROW_IF_ERROR(context_.EncodeDeferredDispatches(recording));
+      context_.EndComputePass(recording);
+      auto& encoder = context_.GetCommandEncoder(recording);
+      offset = 0;
+      for (size_t i = begin; i < end; ++i) {
+        const auto& upload = uploads[i];
+        const auto aligned_size = NormalizeCopySize(upload.size);
+        if (aligned_size != 0) {
+          ORT_ENFORCE(aligned_size <= wgpuBufferGetSize(upload.target));
+          EnforceBufferUnmapped(context_, upload.target);
+          encoder.CopyBufferToBuffer(staging_buffer, offset, upload.target, 0, aligned_size);
+        }
+        offset += aligned_size;
+      }
+      ORT_THROW_IF_ERROR(context_.Flush(*this, recording));
+    }
+    begin = end;
+  }
+}
+
 void BufferManager::MemCpy(CommandRecordingState& recording, WGPUBuffer src, WGPUBuffer dst, size_t size) const {
   ORT_ENFORCE(src != dst, "Source and destination buffers must be different.");
   EnforceBufferUnmapped(context_, src);

onnxruntime/core/providers/webgpu/buffer_manager.h

@@ -94,6 +94,12 @@ class BufferManager {
  public:
   BufferManager(WebGpuContext& context, BufferCacheMode storage_buffer_cache_mode, BufferCacheMode uniform_buffer_cache_mode, BufferCacheMode query_resolve_buffer_cache_mode, BufferCacheMode default_buffer_cache_mode);
   void Upload(CommandRecordingState& recording, void* src, WGPUBuffer dst, size_t size) const;
+  struct UploadInfo {
+    const void* source;
+    WGPUBuffer target;
+    size_t size;
+  };
+  void UploadBatch(CommandRecordingState& recording, gsl::span<const UploadInfo> uploads) const;
   void MemCpy(CommandRecordingState& recording, WGPUBuffer src, WGPUBuffer dst, size_t size) const;
   WGPUBuffer Create(CommandRecordingState& recording, size_t size, wgpu::BufferUsage usage,
                     bool initialize_to_zero = false,

onnxruntime/core/providers/webgpu/data_transfer.cc

@@ -71,5 +71,30 @@
                           src.SizeInBytes());
 }
 
+common::Status DataTransfer::CopyTensors(const std::vector<SrcDstPair>& src_dst_pairs) const {
+  if (src_dst_pairs.size() < 2) {
+    return IDataTransfer::CopyTensors(src_dst_pairs);
+  }
+  InlinedVector<BufferManager::UploadInfo> uploads;
+  uploads.reserve(src_dst_pairs.size());
+  for (const auto& pair : src_dst_pairs) {
+    const auto& src = pair.src.get();
+    auto& dst = pair.dst.get();
+    if (pair.src_stream != nullptr || src.Location().device.Type() != OrtDevice::CPU ||
+        dst.Location().device.Type() != OrtDevice::GPU) {
+      return IDataTransfer::CopyTensors(src_dst_pairs);
+    }
+    const size_t size = src.SizeInBytes();
+    auto target = static_cast<WGPUBuffer>(dst.MutableDataRaw());
+    // Preserve Upload's direct write/unmap path without partially copying the list.
+    if (size != 0 && wgpuBufferGetMapState(target) != WGPUBufferMapState_Unmapped) {
+      return IDataTransfer::CopyTensors(src_dst_pairs);
+    }
+    uploads.push_back({src.DataRaw(), target, size});
+  }
+  buffer_manager_.UploadBatch(recording_, uploads);
+  return Status::OK();
+}
+
 }  // namespace webgpu
 }  // namespace onnxruntime

onnxruntime/core/providers/webgpu/data_transfer.h

@@ -73,7 +73,7 @@
 class DataTransfer : public IDataTransfer {
  public:
   DataTransfer(const BufferManager& buffer_manager, CommandRecordingState& recording, int context_id)
-      : impl_{buffer_manager, recording}, context_id_{context_id} {}
+      : impl_{buffer_manager, recording}, context_id_{context_id}, buffer_manager_{buffer_manager}, recording_{recording} {}
   ~DataTransfer() {};
 
   // Device-compatibility half of CanCopy, split out because it needs no BufferManager and so can
@@ -83,10 +83,13 @@
   bool CanCopy(const OrtDevice& src_device, const OrtDevice& dst_device) const override;
 
   common::Status CopyTensor(const Tensor& src, Tensor& dst) const override;
+  common::Status CopyTensors(const std::vector<SrcDstPair>& src_dst_pairs) const override;
 
  private:
   DataTransferImpl impl_;
   int context_id_;
+  const BufferManager& buffer_manager_;
+  CommandRecordingState& recording_;
 };
 
 }  // namespace webgpu

onnxruntime/core/providers/webgpu/ep/sync_stream.cc

@@ -87,6 +87,13 @@ common::Status CopyTensorOnWebGpuStream(const OrtSyncStream* stream, const void*
   return transfer.CopyTensor(src_data, src_is_gpu, dst_data, dst_is_gpu, bytes);
 }
 
+void UploadTensorsOnWebGpuStream(const OrtSyncStream* stream, int device_id,
+                                 gsl::span<const BufferManager::UploadInfo> uploads) {
+  auto& ep = WebGpuSyncStream::From(stream).ep_;
+  ORT_ENFORCE(ep.GetDeviceId() == device_id, "WebGPU upload destination does not match the stream device.");
+  ep.BufferManager().UploadBatch(ep.Recording(), uploads);
+}
+
 }  // namespace ep
 }  // namespace webgpu
 }  // namespace onnxruntime

onnxruntime/core/providers/webgpu/ep/sync_stream.h

@@ -6,6 +6,7 @@
 #include <cstddef>
 
 #include "core/common/status.h"
+#include "core/providers/webgpu/buffer_manager.h"
 #include "core/session/onnxruntime_c_api.h"
 
 namespace onnxruntime {
@@ -18,6 +19,8 @@ OrtSyncStreamImpl* CreateWebGpuSyncStream(WebGpuExecutionProvider& ep);
 CommandRecordingState& GetWebGpuStreamCommandState(const OrtSyncStream* stream);
 common::Status CopyTensorOnWebGpuStream(const OrtSyncStream* stream, const void* src_data,
                                         bool src_is_gpu, void* dst_data, bool dst_is_gpu, size_t bytes);
+void UploadTensorsOnWebGpuStream(const OrtSyncStream* stream, int device_id,
+                                 gsl::span<const BufferManager::UploadInfo> uploads);
 
 }  // namespace ep
 }  // namespace webgpu

onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc

@@ -580,6 +580,39 @@ struct WebGpuDataTransferImpl : OrtDataTransferImpl {
       return nullptr;
     }
 
+#if defined(ORT_USE_EP_API_ADAPTERS)
+    // Batch unmapped inputs on the owning Session timeline, including on UMA devices.
+    // Check the entire list before copying so mixed transfers retain their original ordering.
+    if (num_tensors > 1 && streams != nullptr && streams[0] != nullptr) {
+      InlinedVector<BufferManager::UploadInfo> uploads;
+      uploads.reserve(num_tensors);
+      int device_id = -1;
+      for (size_t idx = 0; idx < num_tensors; ++idx) {
+        if (streams[idx] != streams[0]) break;
+        Ort::ConstValue src_value{src_tensors[idx]};
+        Ort::UnownedValue dst_value{dst_tensors[idx]};
+        auto src_info = src_value.GetTensorMemoryInfo();
+        auto dst_info = dst_value.GetTensorMemoryInfo();
+        if (src_info.GetDeviceType() != OrtMemoryInfoDeviceType_CPU ||
+            dst_info.GetDeviceType() != OrtMemoryInfoDeviceType_GPU) break;
+        const auto* src_device = impl.ep_api.MemoryInfo_GetMemoryDevice(src_info);
+        const auto* dst_device = impl.ep_api.MemoryInfo_GetMemoryDevice(dst_info);
+        if (!CanCopyImpl(this_ptr, src_device, dst_device)) break;
+        const int dst_device_id = impl.ep_api.MemoryDevice_GetDeviceId(dst_device);
+        if (idx == 0) device_id = dst_device_id;
+        if (dst_device_id != device_id) break;
+        const size_t size = src_value.GetTensorSizeInBytes();
+        auto target = static_cast<WGPUBuffer>(dst_value.GetTensorMutableRawData());
+        if (size != 0 && wgpuBufferGetMapState(target) != WGPUBufferMapState_Unmapped) break;
+        uploads.push_back({src_value.GetTensorRawData(), target, size});
+      }
+      if (uploads.size() == num_tensors) {
+        webgpu::ep::UploadTensorsOnWebGpuStream(streams[0], device_id, uploads);
+        return nullptr;
+      }
+    }
+#endif
+
     for (size_t idx = 0; idx < num_tensors; ++idx) {
       const OrtMemoryInfo* src_memory_info = nullptr;
       const OrtMemoryInfo* dst_memory_info = nullptr;

onnxruntime/test/providers/webgpu/data_transfer_test.cc

@@ -1,10 +1,18 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <array>
+
 #include "gtest/gtest.h"
 
 #include "core/framework/ortdevice.h"
 #include "core/providers/webgpu/data_transfer.h"
+#include "core/providers/webgpu/allocator.h"
+#include "core/providers/webgpu/buffer_manager.h"
+#include "core/providers/webgpu/webgpu_context.h"
+#include "core/providers/webgpu/webgpu_execution_provider.h"
+#include "core/providers/webgpu/webgpu_provider_factory_creator.h"
+#include "test/util/include/asserts.h"
 
 namespace onnxruntime {
 namespace test {
@@ -71,5 +79,105 @@ TEST(WebGpuDataTransferTest, RejectsForeignGpuVendors) {
   }
 }
 
+TEST(WebGpuDataTransferTest, BatchedUploadsPreserveDataAcrossChunksAndReuse) {
+  ConfigOptions options;
+  auto provider = WebGpuProviderFactoryCreator::Create(options)->CreateProvider();
+  auto& ep = static_cast<WebGpuExecutionProvider&>(*provider);
+  auto& context = webgpu::WebGpuContextFactory::GetContext(0);
+  auto& recording = ep.Recording();
+  CPUAllocator cpu_allocator;
+  const OrtMemoryInfo gpu_info(WEBGPU_BUFFER, OrtDeviceAllocator, webgpu::WebGpuDevice(0), OrtMemTypeDefault);
+  webgpu::DataTransfer transfer(ep.BufferManager(), recording, 0);
+
+  // Include empty/unaligned tensors, a batch boundary, and an upload larger than the batch limit.
+  const std::vector<int64_t> sizes{0, 1, 3, 4, 31, 65, 4097, 9 * 1024 * 1024 + 1,
+                                   9 * 1024 * 1024 + 3, 17 * 1024 * 1024 + 1};
+  std::vector<std::vector<uint8_t>> data;
+  std::vector<wgpu::Buffer> buffers;
+  std::vector<Tensor> sources, targets;
+  sources.reserve(sizes.size());
+  targets.reserve(sizes.size());
+  for (const auto size : sizes) {
+    data.emplace_back(static_cast<size_t>(size));
+    wgpu::BufferDescriptor desc{};
+    // Match BufferManager's allocation alignment, which Download also relies on.
+    desc.size = std::max<uint64_t>(16, (static_cast<uint64_t>(size) + 15) / 16 * 16);
+    desc.usage = wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst;
+    buffers.push_back(context.Device().CreateBuffer(&desc));
+    sources.emplace_back(DataTypeImpl::GetType<uint8_t>(), TensorShape({size}), data.back().data(), cpu_allocator.Info());
+    targets.emplace_back(DataTypeImpl::GetType<uint8_t>(), TensorShape({size}), buffers.back().Get(), gpu_info);
+  }
+  std::vector<IDataTransfer::SrcDstPair> uploads;
+  for (size_t i = 0; i < sizes.size(); ++i) uploads.push_back({sources[i], targets[i], nullptr});
+  for (size_t pass = 0; pass < 2; ++pass) {
+    for (size_t i = 0; i < sizes.size(); ++i) {
+      for (size_t j = 0; j < data[i].size(); ++j) data[i][j] = static_cast<uint8_t>(i * 17 + j * 7 + pass);
+    }
+    ASSERT_STATUS_OK(transfer.CopyTensors(uploads));
+    for (size_t i = 0; i < sizes.size(); ++i) {
+      std::vector<uint8_t> actual(static_cast<size_t>(sizes[i]));
+      Tensor result(DataTypeImpl::GetType<uint8_t>(), TensorShape({sizes[i]}), actual.data(), cpu_allocator.Info());
+      ASSERT_STATUS_OK(transfer.CopyTensor(targets[i], result));
+      EXPECT_EQ(actual, data[i]) << "tensor " << i << " pass " << pass;
+    }
+  }
+  ASSERT_STATUS_OK(context.FlushAndWaitChecked(ep.BufferManager(), recording));
+}
+
+TEST(WebGpuDataTransferTest, BatchedUploadsRetainMappedDestinationFallback) {
+  ConfigOptions options;
+  auto provider = WebGpuProviderFactoryCreator::Create(options)->CreateProvider();
+  auto& ep = static_cast<WebGpuExecutionProvider&>(*provider);
+  auto& context = webgpu::WebGpuContextFactory::GetContext(0);
+  auto& recording = ep.Recording();
+  CPUAllocator cpu_allocator;
+  const OrtMemoryInfo gpu_info(WEBGPU_BUFFER, OrtDeviceAllocator, webgpu::WebGpuDevice(0), OrtMemTypeDefault);
+  webgpu::DataTransfer transfer(ep.BufferManager(), recording, 0);
+
+  // Exercise a mapped destination first, last, and at both positions, regardless of UMA support.
+  for (const auto mapped_destinations : {1U, 2U, 3U}) {
+    SCOPED_TRACE(mapped_destinations);
+    std::array<std::array<uint8_t, 16>, 2> data{};
+    std::array<wgpu::Buffer, 2> buffers;
+    InlinedVector<Tensor> sources, targets;
+    sources.reserve(data.size());
+    targets.reserve(data.size());
+    for (size_t i = 0; i < data.size(); ++i) {
+      wgpu::BufferDescriptor desc{};
+      desc.size = data[i].size();
+      desc.usage = wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst;
+      desc.mappedAtCreation = (mapped_destinations & (1U << i)) != 0;
+      buffers[i] = context.Device().CreateBuffer(&desc);
+      ASSERT_EQ(buffers[i].GetMapState(),
+                desc.mappedAtCreation ? wgpu::BufferMapState::Mapped : wgpu::BufferMapState::Unmapped);
+      sources.emplace_back(DataTypeImpl::GetType<uint8_t>(), TensorShape({16}), data[i].data(), cpu_allocator.Info());
+      targets.emplace_back(DataTypeImpl::GetType<uint8_t>(), TensorShape({16}), buffers[i].Get(), gpu_info);
+    }
+
+    const std::vector<IDataTransfer::SrcDstPair> uploads{{sources[0], targets[0], nullptr},
+                                                         {sources[1], targets[1], nullptr}};
+    for (size_t pass = 0; pass < 2; ++pass) {
+      SCOPED_TRACE(pass);
+      for (size_t i = 0; i < data.size(); ++i) {
+        for (size_t j = 0; j < data[i].size(); ++j) data[i][j] = static_cast<uint8_t>(i * 17 + j * 7 + pass);
+      }
+
+      context.PushErrorScope();
+      const auto copy_status = transfer.CopyTensors(uploads);
+      const auto validation_status = context.PopErrorScope();
+      ASSERT_STATUS_OK(copy_status);
+      ASSERT_STATUS_OK(validation_status);
+      for (size_t i = 0; i < data.size(); ++i) {
+        ASSERT_EQ(buffers[i].GetMapState(), wgpu::BufferMapState::Unmapped);
+        std::array<uint8_t, 16> actual{};
+        Tensor result(DataTypeImpl::GetType<uint8_t>(), TensorShape({16}), actual.data(), cpu_allocator.Info());
+        ASSERT_STATUS_OK(transfer.CopyTensor(targets[i], result));
+        EXPECT_EQ(actual, data[i]) << "tensor " << i;
+      }
+    }
+  }
+  ASSERT_STATUS_OK(context.FlushAndWaitChecked(ep.BufferManager(), recording));
+}
+
 }  // namespace test
 }  // namespace onnxruntime