From dc70d87bb6afbc4ba8e014ef1f91061dd56e3ba2 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 20 Oct 2026 22:18:05 +0800 Subject: [PATCH 1/3] webgpu: Batch built-in input uploads with bounded staging Extract the shared uploader and built-in CopyTensors optimization from the QLoRA work so input-transfer performance can be reviewed independently. Group CPU-to-GPU inputs in staging buffers capped at 16 MiB, including on UMA, while retaining single-copy and mixed-transfer fallbacks. Include the existing round-trip test for empty and unaligned tensors, chunk boundaries, oversized inputs, and repeated writes. All five extracted files match the source branch exactly. Scoped lint and whitespace checks pass; standalone build and runtime tests have not been rerun on this branch. Source commits: 8eab658292 and 034acaaf84; final source tree: 2971861c9054525b1b16422749ae352de9537fb9. --- .../core/providers/webgpu/buffer_manager.cc | 56 +++++++++++++++++++ .../core/providers/webgpu/buffer_manager.h | 6 ++ .../core/providers/webgpu/data_transfer.cc | 19 +++++++ .../core/providers/webgpu/data_transfer.h | 5 +- .../providers/webgpu/data_transfer_test.cc | 51 +++++++++++++++++ 5 files changed, 136 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.cc b/onnxruntime/core/providers/webgpu/buffer_manager.cc index 220ec4942f42e..38a9ef450ead1 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.cc +++ b/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 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(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(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); diff --git a/onnxruntime/core/providers/webgpu/buffer_manager.h b/onnxruntime/core/providers/webgpu/buffer_manager.h index f7e1fb033c816..acb3f0f262402 100644 --- a/onnxruntime/core/providers/webgpu/buffer_manager.h +++ b/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 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, diff --git a/onnxruntime/core/providers/webgpu/data_transfer.cc b/onnxruntime/core/providers/webgpu/data_transfer.cc index 4411c922a6fc3..7e0998b605621 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.cc +++ b/onnxruntime/core/providers/webgpu/data_transfer.cc @@ -71,5 +71,24 @@ common::Status DataTransfer::CopyTensor(const Tensor& src, Tensor& dst) const { src.SizeInBytes()); } +common::Status DataTransfer::CopyTensors(const std::vector& src_dst_pairs) const { + if (src_dst_pairs.size() < 2) { + return IDataTransfer::CopyTensors(src_dst_pairs); + } + InlinedVector 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); + } + uploads.push_back({src.DataRaw(), static_cast(dst.MutableDataRaw()), src.SizeInBytes()}); + } + buffer_manager_.UploadBatch(recording_, uploads); + return Status::OK(); +} + } // namespace webgpu } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/data_transfer.h b/onnxruntime/core/providers/webgpu/data_transfer.h index 1ebc869f66d39..7951ff032ddd1 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.h +++ b/onnxruntime/core/providers/webgpu/data_transfer.h @@ -73,7 +73,7 @@ class DataTransferImpl { 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 @@ class DataTransfer : public IDataTransfer { 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& src_dst_pairs) const override; private: DataTransferImpl impl_; int context_id_; + const BufferManager& buffer_manager_; + CommandRecordingState& recording_; }; } // namespace webgpu diff --git a/onnxruntime/test/providers/webgpu/data_transfer_test.cc b/onnxruntime/test/providers/webgpu/data_transfer_test.cc index 23ac26fbfe227..050ca1206634f 100644 --- a/onnxruntime/test/providers/webgpu/data_transfer_test.cc +++ b/onnxruntime/test/providers/webgpu/data_transfer_test.cc @@ -5,6 +5,12 @@ #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 +77,50 @@ TEST(WebGpuDataTransferTest, RejectsForeignGpuVendors) { } } +TEST(WebGpuDataTransferTest, BatchedUploadsPreserveDataAcrossChunksAndReuse) { + ConfigOptions options; + auto provider = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& ep = static_cast(*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 sizes{0, 1, 3, 4, 31, 65, 4097, 9 * 1024 * 1024 + 1, + 9 * 1024 * 1024 + 3, 17 * 1024 * 1024 + 1}; + std::vector> data; + std::vector buffers; + std::vector sources, targets; + sources.reserve(sizes.size()); + targets.reserve(sizes.size()); + for (const auto size : sizes) { + data.emplace_back(static_cast(size)); + wgpu::BufferDescriptor desc{}; + // Match BufferManager's allocation alignment, which Download also relies on. + desc.size = std::max(16, (static_cast(size) + 15) / 16 * 16); + desc.usage = wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + buffers.push_back(context.Device().CreateBuffer(&desc)); + sources.emplace_back(DataTypeImpl::GetType(), TensorShape({size}), data.back().data(), cpu_allocator.Info()); + targets.emplace_back(DataTypeImpl::GetType(), TensorShape({size}), buffers.back().Get(), gpu_info); + } + std::vector 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(i * 17 + j * 7 + pass); + } + ASSERT_STATUS_OK(transfer.CopyTensors(uploads)); + for (size_t i = 0; i < sizes.size(); ++i) { + std::vector actual(static_cast(sizes[i])); + Tensor result(DataTypeImpl::GetType(), 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)); +} + } // namespace test } // namespace onnxruntime From 7e9ebd912886000d2461884762a66a0502abd512 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 20 Oct 2026 22:18:42 +0800 Subject: [PATCH 2/3] webgpu: Batch plugin input uploads on the owning session stream Extract plugin batching from a49f0196ec. Group same-device CPU-to-GPU inputs on a shared non-null WebGPU session stream using the bounded staging uploader, including on UMA. Preserve mapped-buffer, mixed-transfer, single-copy, and streamless fallbacks and existing command ordering. Include five existing submission-count and data-integrity regression tests covering repeated uploads, staging boundaries, empty inputs, oversized inputs, and single-input fallback. All four extracted files match source commit 2971861c9054525b1b16422749ae352de9537fb9. Scoped lint and whitespace checks pass; standalone build and runtime tests have not been rerun on this branch. --- .../core/providers/webgpu/ep/sync_stream.cc | 7 + .../core/providers/webgpu/ep/sync_stream.h | 3 + .../webgpu/webgpu_provider_factory.cc | 33 ++++ .../webgpu/plugin/test_webgpu_allocators.cc | 165 ++++++++++++++++++ 4 files changed, 208 insertions(+) diff --git a/onnxruntime/core/providers/webgpu/ep/sync_stream.cc b/onnxruntime/core/providers/webgpu/ep/sync_stream.cc index f7831672483b3..396084276095d 100644 --- a/onnxruntime/core/providers/webgpu/ep/sync_stream.cc +++ b/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 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 diff --git a/onnxruntime/core/providers/webgpu/ep/sync_stream.h b/onnxruntime/core/providers/webgpu/ep/sync_stream.h index 043f0655ec074..aec71587aebae 100644 --- a/onnxruntime/core/providers/webgpu/ep/sync_stream.h +++ b/onnxruntime/core/providers/webgpu/ep/sync_stream.h @@ -6,6 +6,7 @@ #include #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 uploads); } // namespace ep } // namespace webgpu diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index 520103e687241..3af2a76ff361d 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/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 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(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; diff --git a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc index 2bf2e6b3ab7d1..107446a80b07f 100644 --- a/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc +++ b/onnxruntime/test/providers/webgpu/plugin/test_webgpu_allocators.cc @@ -506,6 +506,119 @@ int VerifyInstanceNormScratchAllocationsStayBatched(Ort::Env& env, Ort::ConstEpD return 0; } +int VerifySessionInputUploadBatching(Ort::Env& env, Ort::ConstEpDevice ep_device, + gsl::span input_sizes, size_t expected_submissions) { + ONNX_NAMESPACE::ModelProto model; + model.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); + model.add_opset_import()->set_version(18); + auto* graph = model.mutable_graph(); + graph->set_name("webgpu_batched_session_inputs"); + auto* node = graph->add_node(); + node->set_op_type("Concat"); + node->add_output("Y"); + auto* axis = node->add_attribute(); + axis->set_name("axis"); + axis->set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + axis->set_i(0); + InlinedVector names; + names.reserve(input_sizes.size()); + int64_t output_size = 0; + for (size_t i = 0; i < input_sizes.size(); ++i) { + names.push_back("X" + std::to_string(i)); + node->add_input(names.back()); + auto* input = graph->add_input(); + input->set_name(names.back()); + auto* type = input->mutable_type()->mutable_tensor_type(); + type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + type->mutable_shape()->add_dim()->set_dim_value(input_sizes[i]); + output_size += input_sizes[i]; + } + auto* output = graph->add_output(); + output->set_name("Y"); + auto* output_type = output->mutable_type()->mutable_tensor_type(); + output_type->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + output_type->mutable_shape()->add_dim()->set_dim_value(output_size); + + struct SubmissionProbe { + SubmissionProbe() = default; + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(SubmissionProbe); + + bool armed{false}; + bool saw_program{false}; + size_t baseline{0}; + size_t input_submissions{0}; + + static void ORT_API_CALL Log(void* param, OrtLoggingLevel, const char*, const char*, const char*, + const char* message) { + auto& probe = *static_cast(param); + if (probe.armed && !probe.saw_program && + std::string_view{message}.find("Starting program") != std::string_view::npos) { + probe.saw_program = true; + probe.input_submissions = nonempty_queue_submissions.load(std::memory_order_relaxed) - probe.baseline; + } + } + } probe; + auto procs = dawn::native::GetProcs(); + procs.queueSubmit = CountNonemptyQueueSubmissions; + dawnProcSetProcs(&procs); + Ort::SessionOptions options; + options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL); + options.AddConfigEntry(kOrtSessionOptionsDisableCPUEPFallback, "1"); + options.SetLogSeverityLevel(ORT_LOGGING_LEVEL_INFO); + Ort::ThrowOnError(Ort::GetApi().SetUserLoggingFunction(options, SubmissionProbe::Log, &probe)); + options.AppendExecutionProvider_V2( + env, {ep_device}, {{"dawnProcTable", std::to_string(reinterpret_cast(&procs))}}); + const auto model_bytes = model.SerializeAsString(); + Ort::Session session(env, model_bytes.data(), model_bytes.size(), options); + const auto cpu_memory = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + InlinedVector> input_data; + InlinedVector inputs; + InlinedVector input_names; + input_data.reserve(input_sizes.size()); + inputs.reserve(input_sizes.size()); + input_names.reserve(input_sizes.size()); + for (size_t i = 0; i < input_sizes.size(); ++i) { + input_data.emplace_back(static_cast(input_sizes[i])); + inputs.push_back(Ort::Value::CreateTensor(cpu_memory, input_data.back().data(), input_data.back().size(), + &input_sizes[i], 1)); + input_names.push_back(names[i].c_str()); + } + const char* output_names[] = {"Y"}; + // Warm pipelines/caches, then measure actual input submissions before the first program. + // Changing every input also catches stale uploads and clears submitted after the copy. + for (int iteration = 0; iteration < 6; ++iteration) { + for (size_t i = 0; i < input_data.size(); ++i) { + for (size_t j = 0; j < input_data[i].size(); ++j) { + input_data[i][j] = static_cast(iteration * 100 + i * 10 + j % 7); + } + } + probe.baseline = nonempty_queue_submissions.load(std::memory_order_relaxed); + probe.saw_program = false; + probe.armed = true; + auto outputs = session.Run(Ort::RunOptions{nullptr}, input_names.data(), inputs.data(), inputs.size(), + output_names, 1); + probe.armed = false; + if (iteration >= 3 && (!probe.saw_program || probe.input_submissions != expected_submissions)) { + std::fprintf(stderr, "Input batching failed: saw_program=%d, submissions=%zu, expected=%zu\n", + static_cast(probe.saw_program), probe.input_submissions, expected_submissions); + return 11; + } + if (outputs[0].GetTensorTypeAndShapeInfo().GetElementCount() != static_cast(output_size)) return 12; + const float* actual = outputs[0].GetTensorData(); + size_t offset = 0; + for (const auto& values : input_data) { + for (float expected : values) { + if (actual[offset++] != expected) { + std::fputs("Batched Session input data was corrupted or stale\n", stderr); + return 13; + } + } + } + } + std::fputs("Session input uploads preserved data and expected submission count\n", stderr); + return 0; +} + WGPUBufferMapState InjectMappedBufferOnFree(WGPUBuffer buffer) { if (buffer == free_failure_buffer) { ++free_failure_probes; @@ -627,6 +740,58 @@ TEST_F(WebGpuSessionAllocatorDeathTest, WarmedInstanceNormalizationScratchAlloca #endif } +TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsAreBatched) { +#if defined(BUILD_DAWN_SHARED_LIBRARY) + GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; +#else + ASSERT_EXIT( + std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{16, 16, 16, 16}, 1)), + ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); +#endif +} + +TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsSpanMultipleChunks) { +#if defined(BUILD_DAWN_SHARED_LIBRARY) + GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; +#else + // Four 5 MiB + 4 byte inputs require two staging chunks at the 16 MiB limit. + ASSERT_EXIT( + std::_Exit(VerifySessionInputUploadBatching( + Env(), EpDevice(), std::array{1310721, 1310721, 1310721, 1310721}, 2)), + ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); +#endif +} + +TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsIncludeEmptyTensors) { +#if defined(BUILD_DAWN_SHARED_LIBRARY) + GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; +#else + ASSERT_EXIT( + std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{0, 16, 0, 16}, 1)), + ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); +#endif +} + +TEST_F(WebGpuSessionAllocatorDeathTest, SessionInputUploadsHandleOversizedTensor) { +#if defined(BUILD_DAWN_SHARED_LIBRARY) + GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; +#else + ASSERT_EXIT( + std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{4194305, 16}, 2)), + ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); +#endif +} + +TEST_F(WebGpuSessionAllocatorDeathTest, SessionSingleInputUploadRetainsFallback) { +#if defined(BUILD_DAWN_SHARED_LIBRARY) + GTEST_SKIP() << "Shared Dawn calls bypass the replaceable proc table required for submission counting."; +#else + ASSERT_EXIT( + std::_Exit(VerifySessionInputUploadBatching(Env(), EpDevice(), std::array{16}, 1)), + ::testing::ExitedWithCode(0), "Session input uploads preserved data and expected submission count"); +#endif +} + #endif } // namespace From c31f7f26ed3ce78ec8ccc2c716b45b83e260f255 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 21 Oct 2026 15:24:23 +0800 Subject: [PATCH 3/3] webgpu: Preserve mapped input upload fallback Check destination map state during the built-in batch preflight and fall back for the entire list before copying. This preserves the existing direct memcpy/unmap behavior for mapped inputs while retaining batching for unmapped inputs. Add coverage for mapped destinations first, last, and in both positions, plus repeated uploads after unmapping. The regression fails before the fix and passes after it. Verified built-in and plugin Vulkan builds, 6 built-in transfer tests, 62 enabled plugin tests, and scoped lint. --- .../core/providers/webgpu/data_transfer.cc | 8 ++- .../providers/webgpu/data_transfer_test.cc | 57 +++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/webgpu/data_transfer.cc b/onnxruntime/core/providers/webgpu/data_transfer.cc index 7e0998b605621..e0e182c3dc8fc 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.cc +++ b/onnxruntime/core/providers/webgpu/data_transfer.cc @@ -84,7 +84,13 @@ common::Status DataTransfer::CopyTensors(const std::vector& src_dst_ dst.Location().device.Type() != OrtDevice::GPU) { return IDataTransfer::CopyTensors(src_dst_pairs); } - uploads.push_back({src.DataRaw(), static_cast(dst.MutableDataRaw()), src.SizeInBytes()}); + const size_t size = src.SizeInBytes(); + auto target = static_cast(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(); diff --git a/onnxruntime/test/providers/webgpu/data_transfer_test.cc b/onnxruntime/test/providers/webgpu/data_transfer_test.cc index 050ca1206634f..198a1afe1d843 100644 --- a/onnxruntime/test/providers/webgpu/data_transfer_test.cc +++ b/onnxruntime/test/providers/webgpu/data_transfer_test.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "gtest/gtest.h" #include "core/framework/ortdevice.h" @@ -122,5 +124,60 @@ TEST(WebGpuDataTransferTest, BatchedUploadsPreserveDataAcrossChunksAndReuse) { ASSERT_STATUS_OK(context.FlushAndWaitChecked(ep.BufferManager(), recording)); } +TEST(WebGpuDataTransferTest, BatchedUploadsRetainMappedDestinationFallback) { + ConfigOptions options; + auto provider = WebGpuProviderFactoryCreator::Create(options)->CreateProvider(); + auto& ep = static_cast(*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, 2> data{}; + std::array buffers; + InlinedVector 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(), TensorShape({16}), data[i].data(), cpu_allocator.Info()); + targets.emplace_back(DataTypeImpl::GetType(), TensorShape({16}), buffers[i].Get(), gpu_info); + } + + const std::vector 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(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 actual{}; + Tensor result(DataTypeImpl::GetType(), 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