From 02b9b121272a5cd4f57a245eb15a2ca67c9d97ea Mon Sep 17 00:00:00 2001 From: Jeff Bloomfield <38966965+jeffbloo@users.noreply.github.com> Date: Fri, 1 Jul 2022 09:49:23 -0700 Subject: [PATCH] Fix DML custom operators which set descriptor heap to command list (#12059) --- .../dml/DmlExecutionProvider/inc/IWinmlExecutionProvider.h | 2 +- .../dml/DmlExecutionProvider/src/DmlCommandRecorder.h | 6 ++++++ .../dml/DmlExecutionProvider/src/ExecutionContext.cpp | 6 +++++- .../dml/DmlExecutionProvider/src/ExecutionContext.h | 2 +- .../dml/DmlExecutionProvider/src/ExecutionProvider.cpp | 4 ++-- .../dml/DmlExecutionProvider/src/ExecutionProvider.h | 2 +- .../dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp | 4 ++-- 7 files changed, 18 insertions(+), 8 deletions(-) diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/inc/IWinmlExecutionProvider.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/inc/IWinmlExecutionProvider.h index 073c0ebf62..14572cbb82 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/inc/IWinmlExecutionProvider.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/inc/IWinmlExecutionProvider.h @@ -45,7 +45,7 @@ namespace Windows::AI::MachineLearning::Adapter IUnknown* data, bool isInternalOperator) = 0; - virtual void GetABIExecutionInterface( + virtual void GetABIExecutionInterfaceAndInvalidateState( bool isInternalOperator, IUnknown** abiExecutionObject) const = 0; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommandRecorder.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommandRecorder.h index 9becf3ad81..7ad7032317 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommandRecorder.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/DmlCommandRecorder.h @@ -61,6 +61,12 @@ namespace Dml return m_operationsRecordedInCurrentCommandList || !m_pendingCommandLists.empty(); } + // Forces the descriptor heap to be reset to D3D before executing future operations + void InvalidateDescriptorHeap() + { + m_currentDescriptorHeap = nullptr; + } + private: std::shared_ptr m_queue; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionContext.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionContext.cpp index b5f070dd3f..a894d0660d 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionContext.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionContext.cpp @@ -125,10 +125,14 @@ namespace Dml m_dmlRecorder.ResourceBarrier(barriers); } - void ExecutionContext::GetCommandListForRecording(ID3D12GraphicsCommandList** commandList) + void ExecutionContext::GetCommandListForRecordingAndInvalidateState(ID3D12GraphicsCommandList** commandList) { assert(!m_closed); SetCommandRecorder(&m_dmlRecorder); + + // Ensure the descriptor heap is reset to D3D as something external may change it before recording + m_dmlRecorder.InvalidateDescriptorHeap(); + m_dmlRecorder.GetCommandList().CopyTo(commandList); } diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionContext.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionContext.h index f70ee13572..b06f11a5ef 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionContext.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionContext.h @@ -66,7 +66,7 @@ namespace Dml void AddUAVBarrier(); void ResourceBarrier(gsl::span barriers); - void GetCommandListForRecording(ID3D12GraphicsCommandList** commandList); + void GetCommandListForRecordingAndInvalidateState(ID3D12GraphicsCommandList** commandList); // Forces all queued work to begin executing on the GPU. This method returns immediately and does not wait // for the submitted work to complete execution on the GPU. diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp index 1e4f10b2cd..b9b00d91f4 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp @@ -688,7 +688,7 @@ namespace Dml return m_allocator->DecodeDataHandle(data)->GetPooledResourceId(); } - void ExecutionProviderImpl::GetABIExecutionInterface( + void ExecutionProviderImpl::GetABIExecutionInterfaceAndInvalidateState( bool isInternalOperator, IUnknown** abiExecutionObject) const { @@ -702,7 +702,7 @@ namespace Dml else { ComPtr commandList; - m_context->GetCommandListForRecording(commandList.GetAddressOf()); + m_context->GetCommandListForRecordingAndInvalidateState(commandList.GetAddressOf()); #ifdef _GAMING_XBOX ComPtr wrappedCommandList = Microsoft::WRL::Make(commandList.Get()); *abiExecutionObject = wrappedCommandList.Detach(); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h index 257f481032..8024657e13 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h @@ -108,7 +108,7 @@ namespace Dml IUnknown* data, bool isInternalOperator) override; - void GetABIExecutionInterface( + void GetABIExecutionInterfaceAndInvalidateState( bool isInternalOperator, IUnknown** abiExecutionObject) const override; diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp index cfa5ba8a2b..e977900427 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/MLOperatorAuthorImpl.cpp @@ -1467,7 +1467,7 @@ namespace Windows::AI::MachineLearning::Adapter m_abiExecutionObject = m_providerExecutionObject; if (m_winmlProvider) { - m_winmlProvider->GetABIExecutionInterface(isInternalOperator, m_abiExecutionObject.ReleaseAndGetAddressOf()); + m_winmlProvider->GetABIExecutionInterfaceAndInvalidateState(isInternalOperator, m_abiExecutionObject.ReleaseAndGetAddressOf()); } TransitionResourcesForOperatorIfRequired(true); @@ -1737,7 +1737,7 @@ namespace Windows::AI::MachineLearning::Adapter if (m_winmlProvider) { // Get the particular object to return to a isInternalOperator based on the registration of that kernel. - m_winmlProvider->GetABIExecutionInterface(isInternalOperator, m_abiExecutionObject.ReleaseAndGetAddressOf()); + m_winmlProvider->GetABIExecutionInterfaceAndInvalidateState(isInternalOperator, m_abiExecutionObject.ReleaseAndGetAddressOf()); } }