diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp index 90ecc7f87f..a16c00b9d4 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp @@ -946,6 +946,13 @@ namespace Dml assert(!m_closed); m_context->Flush(); } + + void __stdcall ExecutionProviderImpl::FlushAndSync() const + { + Flush(); + m_context->GetCurrentCompletionEvent().WaitForSignal(); + m_context->ReleaseCompletedReferences(); + } void ExecutionProviderImpl::SetDefaultRoundingMode(AllocatorRoundingMode roundingMode) { diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h index b9ac772095..8881347936 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h @@ -125,6 +125,7 @@ namespace Dml STDMETHOD_(D3D12_COMMAND_LIST_TYPE, GetCommandListTypeForQueue)() const override; STDMETHOD_(void, Flush)() const override; + STDMETHOD_(void, FlushAndSync)() const override; void SetDefaultRoundingMode(AllocatorRoundingMode roundingMode); diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.cpp index b7f24d49d1..4f72d1dc12 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/FusedGraphKernel.cpp @@ -72,6 +72,9 @@ namespace Dml m_persistentResourceBinding = DML_BUFFER_BINDING { m_persistentResource.Get(), 0, persistentResourceSize }; } + // For temporary prototyping, allow weight upload to complete before issuing initialization work which reads those weights + m_provider->FlushAndSync(); + ORT_THROW_IF_FAILED(m_provider->InitializeOperator( m_compiledExecutionPlanOperator.Get(), m_persistentResourceBinding ? &*m_persistentResourceBinding : nullptr, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/IExecutionProvider.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/IExecutionProvider.h index d7a0a607cd..f5556ae2ab 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/IExecutionProvider.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/IExecutionProvider.h @@ -64,6 +64,7 @@ namespace Dml STDMETHOD_(D3D12_COMMAND_LIST_TYPE, GetCommandListTypeForQueue)() const noexcept = 0; STDMETHOD_(void, Flush)() const noexcept = 0; + STDMETHOD_(void, FlushAndSync)() const noexcept = 0; STDMETHOD_(ID3D12Resource*, DecodeResource)(void* allocation) const noexcept = 0; STDMETHOD(AllocatePooledResource(size_t size, AllocatorRoundingMode roundingMode, ID3D12Resource **d3dResource, IUnknown* *pooledResource)) const noexcept = 0;