saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-06 04:28:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support batched copies from GPU to CPU in DML EP

Browse source
This commit is contained in:
Jeff 2020-04-16 13:17:04 -07:00
parent 31bb1182e6
commit d40fd897dd
4 changed files with 149 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

61
onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp
View file

@ -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<onnxruntime::Tensor*>(&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<ID3D12Resource*> srcDatas;
std::vector<void*> dstDatas;
std::vector<uint32_t> dataSizesInBytes;
assert(!m_closed);
auto provider = const_cast<ExecutionProviderImpl*>(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<onnxruntime::Tensor*>(src[i]),
false,
provider,
true);
TensorWrapper dstWrapper = TensorWrapper(
dst[i],
true,
provider,
true);
dataSizesInBytes.push_back(static_cast<uint32_t>(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);

7
onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h
View file

@ -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
{

92
onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.cpp
View file

@ -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<ptrdiff_t>(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<ptrdiff_t>(m_readbackHeap->GetDesc().Width) >= totalSize);
}
void ReadbackHeap::ReadbackFromGpu(
gsl::span<std::byte> 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<void*> dst,
gsl::span<const uint32_t > dstSizes,
gsl::span<ID3D12Resource*> 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<uint8_t*>(readbackHeapData) + offset, dstSizes[i]);
offset += dstSizes[i];
}
m_readbackHeap->Unmap(0, nullptr);
}
} // namespace Dml

9
onnxruntime/core/providers/dml/DmlExecutionProvider/src/ReadbackHeap.h
View file

@ -21,8 +21,17 @@ namespace Dml
ID3D12Resource* src,
uint64_t srcOffset,
D3D12_RESOURCE_STATES srcState);
// Overload supporting batching
void ReadbackFromGpu(
gsl::span<void*> dst,
gsl::span<const uint32_t > dstSizes,
gsl::span<ID3D12Resource*> src,
D3D12_RESOURCE_STATES srcState);
private:
void EnsureReadbackHeap(size_t size);
static constexpr size_t c_initialCapacity = 1024 * 1024; // 1MB
ComPtr<ID3D12Device> m_device;