From d40fd897dda9840c9b4ed09b2a5ad8c86de0f8dd Mon Sep 17 00:00:00 2001 From: Jeff Date: Thu, 16 Apr 2020 13:17:04 -0700 Subject: [PATCH 1/4] Support batched copies from GPU to CPU in DML EP --- .../src/ExecutionProvider.cpp | 61 +++++++++++- .../src/ExecutionProvider.h | 7 ++ .../DmlExecutionProvider/src/ReadbackHeap.cpp | 92 +++++++++++++++---- .../DmlExecutionProvider/src/ReadbackHeap.h | 9 ++ 4 files changed, 149 insertions(+), 20 deletions(-) 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; From c925156fd21129bd858055850508650592f5549e Mon Sep 17 00:00:00 2001 From: Jeff Date: Thu, 16 Apr 2020 13:42:09 -0700 Subject: [PATCH 2/4] Fix build error --- .../providers/dml/DmlExecutionProvider/src/ReadbackHeap.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.cpp index e590a96703..a3517e12e1 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.cpp @@ -51,13 +51,13 @@ namespace Dml { // Initialize the readback heap for the first time assert(m_capacity == 0); - m_capacity = ComputeNewCapacity(c_initialCapacity, totalSize); + m_capacity = ComputeNewCapacity(c_initialCapacity, size); m_readbackHeap = CreateReadbackHeap(m_device.Get(), m_capacity); } - else if (gsl::narrow_cast(m_capacity) < totalSize) + else if (gsl::narrow_cast(m_capacity) < size) { // Ensure there's sufficient capacity - m_capacity = ComputeNewCapacity(m_capacity, totalSize); + m_capacity = ComputeNewCapacity(m_capacity, size); m_readbackHeap = nullptr; m_readbackHeap = CreateReadbackHeap(m_device.Get(), m_capacity); From 027b0cb3f3da2b250e6a9ef3de64bec6ccc486d6 Mon Sep 17 00:00:00 2001 From: Jeff Date: Fri, 17 Apr 2020 15:04:58 -0700 Subject: [PATCH 3/4] Update to match ORT signature --- .../DmlExecutionProvider/src/ExecutionProvider.cpp | 12 ++++++------ .../dml/DmlExecutionProvider/src/ExecutionProvider.h | 7 +++---- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp index f43487f497..f01f9de534 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp @@ -540,7 +540,7 @@ namespace Dml } - Status ExecutionProviderImpl::CopyTensors(const onnxruntime::Tensor** src, onnxruntime::Tensor** dst, uint32_t count) const + Status ExecutionProviderImpl::CopyTensors(const std::vector& src_dst_pairs) const { // Source and destination for batched GPU -> CPU copies std::vector srcDatas; @@ -550,24 +550,24 @@ namespace Dml assert(!m_closed); auto provider = const_cast(this); - for (uint32_t i = 0; i < count; ++i) + for (uint32_t i = 0; i < src_dst_pairs.size(); ++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])) + if (!IsGpuTensor(src_dst_pairs[i].src) || IsGpuTensor(src_dst_pairs[i].dst)) { - ORT_RETURN_IF_ERROR(CopyTensor(*src[i], *dst[i])); + ORT_RETURN_IF_ERROR(CopyTensor(src_dst_pairs[i].src, src_dst_pairs[i].dst)); continue; } TensorWrapper srcWrapper = TensorWrapper( - const_cast(src[i]), + const_cast(&src_dst_pairs[i].src.get()), false, provider, true); TensorWrapper dstWrapper = TensorWrapper( - dst[i], + &src_dst_pairs[i].dst.get(), true, provider, true); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h index 48798ca2c8..8e85171e54 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h @@ -92,7 +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 CopyTensors(const std::vector& src_dst_pairs) const; onnxruntime::common::Status WaitForGpuCompletion(); // IWinmlExecutionProvider methods @@ -200,10 +200,9 @@ namespace Dml 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 + onnxruntime::common::Status CopyTensors(const std::vector& src_dst_pairs) const { - assert(exec_queue_id == 0); - return m_impl->CopyTensors(src, dst, count); + return m_impl->CopyTensors(src_dst_pairs); } bool CanCopy(const OrtDevice& srcDevice, const OrtDevice& dstDevice) const final From c47490ab317a7db10a3eb7f5912768c712784dd5 Mon Sep 17 00:00:00 2001 From: Jeff Date: Fri, 17 Apr 2020 17:16:14 -0700 Subject: [PATCH 4/4] Bug fix --- .../dml/DmlExecutionProvider/src/ExecutionProvider.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp index f01f9de534..ca3f2475a1 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp @@ -562,13 +562,13 @@ namespace Dml TensorWrapper srcWrapper = TensorWrapper( const_cast(&src_dst_pairs[i].src.get()), - false, + true, provider, true); TensorWrapper dstWrapper = TensorWrapper( &src_dst_pairs[i].dst.get(), - true, + false, provider, true);