From e1012a8662585b730ba1973450f1c7e0911735d5 Mon Sep 17 00:00:00 2001
From: sumitsays <sumitagarwal330@gmail.com>
Date: Tue, 25 Jan 2022 13:00:44 -0800
Subject: [PATCH] Added OnRunEnd and Sync method in ExecutionProvider (#10362)

Co-authored-by: Sumit Agarwal <sumitagarwal@microsoft.com>
---
 .../src/ExecutionProvider.cpp                  |  6 ++++++
 .../src/ExecutionProvider.h                    | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+)

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();