diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp index 32a6672593..7363bba4a3 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.cpp @@ -97,6 +97,12 @@ namespace Dml { m_context->Close(); } + + void ExecutionProviderImpl::WaitForOutstandingWork() + { + Flush(); + m_context->GetCurrentCompletionEvent().WaitForSignal(); + } HRESULT __stdcall ExecutionProviderImpl::AllocatePooledResource( size_t size, diff --git a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h index 5c2db8bc70..257f481032 100644 --- a/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h +++ b/onnxruntime/core/providers/dml/DmlExecutionProvider/src/ExecutionProvider.h @@ -130,6 +130,8 @@ namespace Dml // Waits for flushed work, discards unflushed work, and discards associated references to // prevent circular references. Must be the last call on the object before destruction. void Close() override; + + void WaitForOutstandingWork(); // Allocate a resource from pools. Releasing pooledResource returns it to the pool. STDMETHOD(AllocatePooledResource)( @@ -249,6 +251,22 @@ namespace Dml return m_impl->OnSessionInitializationEnd(); } + virtual onnxruntime::Status Sync() const final override + { + // Completely wait until the device has completed all preceding tasks. + // The application could have called SynchronizeBoundOutputs(). + m_impl->WaitForOutstandingWork(); + return Status::OK(); + } + + virtual onnxruntime::Status OnRunEnd(bool /*sync_stream*/) final override + { + // Flush any pending work to the GPU, but don't block for completion, permitting it + // to overlap other work. + m_impl->Flush(); + return Status::OK(); + } + void Flush() { return m_impl->Flush();