From 01f39666619e4fc5697ca818d4d54d6bd46abb2e Mon Sep 17 00:00:00 2001 From: Jeff Bloomfield Date: Wed, 6 Jul 2022 17:43:23 -0700 Subject: [PATCH 1/2] Prevent unbounded growth of command allocator memory --- .../src/CommandAllocatorRing.h | 27 +++++++++---------- .../src/DmlCommandRecorder.cpp | 7 ++--- 2 files changed, 14 insertions(+), 20 deletions(-) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/CommandAllocatorRing.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/CommandAllocatorRing.h index 26c8a021a008d..8d12fdf77a613 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/CommandAllocatorRing.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/CommandAllocatorRing.h @@ -28,26 +28,23 @@ namespace Dml } } - ID3D12CommandAllocator* GetCurrentAllocator() + ID3D12CommandAllocator* GetNextAllocator(GpuEvent nextCompletionEvent) { - CommandAllocatorInfo& allocatorInfo = m_commandAllocators[m_currentCommandAllocator]; + size_t earliestOtherAllocator = (m_currentCommandAllocator + 1) % AllocatorCount; - // Take the opportunity to reset the command allocator if possible. - if (allocatorInfo.completionEvent.IsSignaled()) - { - ORT_THROW_IF_FAILED(allocatorInfo.Get()->Reset()); - } - - return m_commandAllocators[m_currentCommandAllocator].Get(); - } + assert(!m_commandAllocators[m_currentCommandAllocator].completionEvent.IsSignaled() || + m_commandAllocators[earliestOtherAllocator].completionEvent.IsSignaled()); - void AdvanceAllocator(GpuEvent completionEvent) - { + if (m_commandAllocators[earliestOtherAllocator].completionEvent.IsSignaled()) + { + ORT_THROW_IF_FAILED(m_commandAllocators[earliestOtherAllocator].Get()->Reset()); + m_currentCommandAllocator = earliestOtherAllocator; + } + // Set the completion event for the current allocator so it can be reset eventually. - m_commandAllocators[m_currentCommandAllocator].completionEvent = completionEvent; + m_commandAllocators[m_currentCommandAllocator].completionEvent = nextCompletionEvent; - // Advance to the next allocator. - m_currentCommandAllocator = (m_currentCommandAllocator + 1) % AllocatorCount; + return m_commandAllocators[m_currentCommandAllocator].Get(); } private: diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommandRecorder.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommandRecorder.cpp index bd1a8bdc668e0..59ceecdc884d2 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommandRecorder.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommandRecorder.cpp @@ -314,14 +314,14 @@ void DmlCommandRecorder::Open() { assert(m_currentDescriptorHeap == nullptr); - ID3D12CommandAllocator* allocator = m_commandAllocatorRing.GetCurrentAllocator(); + ID3D12CommandAllocator* allocator = m_commandAllocatorRing.GetNextAllocator(m_queue->GetNextCompletionEvent()); if (m_cachedCommandLists.empty()) { ORT_THROW_IF_FAILED(m_d3dDevice->CreateCommandList( 0, m_queue->GetType(), - m_commandAllocatorRing.GetCurrentAllocator(), + allocator, nullptr, IID_GRAPHICS_PPV_ARGS(m_currentCommandList.ReleaseAndGetAddressOf()))); } @@ -331,9 +331,6 @@ void DmlCommandRecorder::Open() m_cachedCommandLists.pop_front(); ORT_THROW_IF_FAILED(m_currentCommandList->Reset(allocator, nullptr)); } - - // The current command allocator will become eligible for reset once this command list completes execution - m_commandAllocatorRing.AdvanceAllocator(m_queue->GetNextCompletionEvent()); } void DmlCommandRecorder::CloseAndExecute() From 64861e54db1f9af25d276205194cfeee63d659f4 Mon Sep 17 00:00:00 2001 From: Jeff Bloomfield Date: Wed, 6 Jul 2022 18:14:13 -0700 Subject: [PATCH 2/2] Fix trailing whitespace --- .../dml/DmlExecutionProvider/src/CommandAllocatorRing.h | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/CommandAllocatorRing.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/CommandAllocatorRing.h index 8d12fdf77a613..570f62aac8105 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/CommandAllocatorRing.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/CommandAllocatorRing.h @@ -30,17 +30,17 @@ namespace Dml ID3D12CommandAllocator* GetNextAllocator(GpuEvent nextCompletionEvent) { - size_t earliestOtherAllocator = (m_currentCommandAllocator + 1) % AllocatorCount; + size_t earliestOtherAllocator = (m_currentCommandAllocator + 1) % AllocatorCount; assert(!m_commandAllocators[m_currentCommandAllocator].completionEvent.IsSignaled() || m_commandAllocators[earliestOtherAllocator].completionEvent.IsSignaled()); if (m_commandAllocators[earliestOtherAllocator].completionEvent.IsSignaled()) - { + { ORT_THROW_IF_FAILED(m_commandAllocators[earliestOtherAllocator].Get()->Reset()); m_currentCommandAllocator = earliestOtherAllocator; } - + // Set the completion event for the current allocator so it can be reset eventually. m_commandAllocators[m_currentCommandAllocator].completionEvent = nextCompletionEvent;