diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp index 04c6ef202c..f43487f497 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp @@ -509,6 +509,12 @@ namespace Dml partitionKernelPrefix ); } + + bool IsGpuTensor(const onnxruntime::Tensor& tensor) + { + return strcmp(tensor.Location().name, onnxruntime::CPU) && + !(tensor.Location().mem_type == ::OrtMemType::OrtMemTypeCPUOutput || tensor.Location().mem_type == ::OrtMemType::OrtMemTypeCPUInput); + } Status ExecutionProviderImpl::CopyTensor(const onnxruntime::Tensor& src, onnxruntime::Tensor& dst) const { @@ -518,13 +524,13 @@ namespace Dml TensorWrapper destInternal( &dst, - strcmp(dst.Location().name, onnxruntime::CPU) && !(dst.Location().mem_type == ::OrtMemType::OrtMemTypeCPUOutput || dst.Location().mem_type == ::OrtMemType::OrtMemTypeCPUInput), + IsGpuTensor(dst), provider, true); TensorWrapper srcInternal( const_cast(&src), - strcmp(src.Location().name, onnxruntime::CPU) && !(src.Location().mem_type == ::OrtMemType::OrtMemTypeCPUOutput || src.Location().mem_type == ::OrtMemType::OrtMemTypeCPUInput), + IsGpuTensor(src), provider, true); @@ -533,6 +539,57 @@ namespace Dml return onnxruntime::common::Status::OK(); } + + Status ExecutionProviderImpl::CopyTensors(const onnxruntime::Tensor** src, onnxruntime::Tensor** dst, uint32_t count) const + { + // Source and destination for batched GPU -> CPU copies + std::vector srcDatas; + std::vector dstDatas; + std::vector dataSizesInBytes; + + assert(!m_closed); + auto provider = const_cast(this); + + for (uint32_t i = 0; i < count; ++i) + { + // This batching implementation only handles GPU -> CPU copies. Other copies do not require synchronization + // and are batched across multiple calls to CopyTensor. + if (!IsGpuTensor(*src[i]) || IsGpuTensor(*dst[i])) + { + ORT_RETURN_IF_ERROR(CopyTensor(*src[i], *dst[i])); + continue; + } + + TensorWrapper srcWrapper = TensorWrapper( + const_cast(src[i]), + false, + provider, + true); + + TensorWrapper dstWrapper = TensorWrapper( + dst[i], + true, + provider, + true); + + dataSizesInBytes.push_back(static_cast(ComputeByteSizeFromTensor(dstWrapper))); + THROW_HR_IF(E_INVALIDARG, dataSizesInBytes[i] != ComputeByteSizeFromTensor(srcWrapper)); // Tensors must be the same size + + dstDatas.push_back(dstWrapper.GetData()); + const AllocationInfo* srcAllocInfo = m_allocator->DecodeDataHandle(MLOperatorTensor(&srcWrapper).GetDataInterface().Get()); + + srcDatas.push_back(srcAllocInfo->GetResource()); + } + + const uint64_t srcOffset = 0; + const auto srcState = D3D12_RESOURCE_STATE_UNORDERED_ACCESS; // GPU resources are always kept in UAV state + + // Performs a blocking call to synchronize and read back data from the GPU into the destination buffer + m_readbackHeap->ReadbackFromGpu(dstDatas, dataSizesInBytes, srcDatas, srcState); + + return onnxruntime::common::Status::OK(); + } + Status ExecutionProviderImpl::WaitForGpuCompletion() { assert(!m_closed); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h index 58f73f62ba..48798ca2c8 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h @@ -92,6 +92,7 @@ namespace Dml uint32_t GetSuppportedDeviceDataTypeMask() const; onnxruntime::common::Status CopyTensor(const onnxruntime::Tensor& src, onnxruntime::Tensor& dst) const; + onnxruntime::common::Status CopyTensors(const onnxruntime::Tensor** src, onnxruntime::Tensor** dst, uint32_t count) const; onnxruntime::common::Status WaitForGpuCompletion(); // IWinmlExecutionProvider methods @@ -198,6 +199,12 @@ namespace Dml assert(exec_queue_id == 0); return m_impl->CopyTensor(src, dst); } + + onnxruntime::common::Status CopyTensors(const onnxruntime::Tensor** src, onnxruntime::Tensor** dst, uint32_t count, int exec_queue_id) const + { + assert(exec_queue_id == 0); + return m_impl->CopyTensors(src, dst, count); + } bool CanCopy(const OrtDevice& srcDevice, const OrtDevice& dstDevice) const final { diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.cpp index e7118d37d4..e590a96703 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.cpp @@ -45,6 +45,27 @@ namespace Dml return newCapacity; } + void ReadbackHeap::EnsureReadbackHeap(size_t size) + { + if (!m_readbackHeap) + { + // Initialize the readback heap for the first time + assert(m_capacity == 0); + m_capacity = ComputeNewCapacity(c_initialCapacity, totalSize); + m_readbackHeap = CreateReadbackHeap(m_device.Get(), m_capacity); + } + else if (gsl::narrow_cast(m_capacity) < totalSize) + { + // Ensure there's sufficient capacity + m_capacity = ComputeNewCapacity(m_capacity, totalSize); + + m_readbackHeap = nullptr; + m_readbackHeap = CreateReadbackHeap(m_device.Get(), m_capacity); + } + + assert(gsl::narrow_cast(m_readbackHeap->GetDesc().Width) >= totalSize); + } + void ReadbackHeap::ReadbackFromGpu( gsl::span dst, ID3D12Resource* src, @@ -52,24 +73,8 @@ namespace Dml D3D12_RESOURCE_STATES srcState) { assert(!dst.empty()); - - if (!m_readbackHeap) - { - // Initialize the readback heap for the first time - assert(m_capacity == 0); - m_capacity = ComputeNewCapacity(c_initialCapacity, dst.size()); - m_readbackHeap = CreateReadbackHeap(m_device.Get(), m_capacity); - } - else if (m_capacity < dst.size()) - { - // Ensure there's sufficient capacity - m_capacity = ComputeNewCapacity(m_capacity, dst.size()); - - m_readbackHeap = nullptr; - m_readbackHeap = CreateReadbackHeap(m_device.Get(), m_capacity); - } - - assert(m_readbackHeap->GetDesc().Width >= dst.size()); + + EnsureReadbackHeap(dst.size()); // Copy from the source resource into the readback heap m_executionContext->CopyBufferRegion( @@ -92,5 +97,56 @@ namespace Dml memcpy(dst.data(), readbackHeapData, dst.size()); m_readbackHeap->Unmap(0, nullptr); } + + void ReadbackHeap::ReadbackFromGpu( + gsl::span dst, + gsl::span dstSizes, + gsl::span src, + D3D12_RESOURCE_STATES srcState) + { + assert(!dst.empty()); + uint32_t totalSize = 0; + for (auto size : dstSizes) + { + totalSize += size; + } + + EnsureReadbackHeap(totalSize); + + // Copy from the source resource into the readback heap + uint32_t offset = 0; + for (uint32_t i = 0; i < dst.size(); ++i) + { + m_executionContext->CopyBufferRegion( + m_readbackHeap.Get(), + offset, + D3D12_RESOURCE_STATE_COPY_DEST, + src[i], + 0, + srcState, + dstSizes[i]); + + offset += dstSizes[i]; + } + + // Wait for completion and map the result + m_executionContext->Flush(); + m_executionContext->GetCurrentCompletionEvent().WaitForSignal(); + m_executionContext->ReleaseCompletedReferences(); + + // Map the readback heap and copy it into the destination + void* readbackHeapData = nullptr; + THROW_IF_FAILED(m_readbackHeap->Map(0, nullptr, &readbackHeapData)); + + // Copy from the source resource into the readback heap + offset = 0; + for (uint32_t i = 0; i < dst.size(); ++i) + { + memcpy(dst[i], static_cast(readbackHeapData) + offset, dstSizes[i]); + offset += dstSizes[i]; + } + + m_readbackHeap->Unmap(0, nullptr); + } } // namespace Dml diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.h index 59425f860f..c596d982b7 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.h @@ -21,8 +21,17 @@ namespace Dml ID3D12Resource* src, uint64_t srcOffset, D3D12_RESOURCE_STATES srcState); + + // Overload supporting batching + void ReadbackFromGpu( + gsl::span dst, + gsl::span dstSizes, + gsl::span src, + D3D12_RESOURCE_STATES srcState); private: + void EnsureReadbackHeap(size_t size); + static constexpr size_t c_initialCapacity = 1024 * 1024; // 1MB ComPtr m_device;