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..e0e182c3dc8fc 100644 --- a/onnxruntime/core/providers/webgpu/data_transfer.cc +++ b/onnxruntime/core/providers/webgpu/data_transfer.cc @@ -71,5 +71,30 @@ 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); + } + 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(); +} + } // 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/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/data_transfer_test.cc b/onnxruntime/test/providers/webgpu/data_transfer_test.cc index 23ac26fbfe227..198a1afe1d843 100644 --- a/onnxruntime/test/providers/webgpu/data_transfer_test.cc +++ b/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 + #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(*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)); +} + +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 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