Implement DML execution provider functions from adapter (#2846)

* Implement DML execution provider functions from adapter * Use functions in OnnxruntimeEngine.cpp
microsoft · Jan 16, 2020 · 8fb6419 · 8fb6419
1 parent 272d829
commit 8fb6419
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 10 deletions.
diff --git a/cmake/external/protobuf b/cmake/external/protobuf
diff --git a/winml/adapter/winml_adapter_apis.h b/winml/adapter/winml_adapter_apis.h
@@ -73,9 +73,9 @@ ORT_API_STATUS(SessionRegisterCustomRegistry, _In_ OrtSession* session, _In_ IML
 
 // Dml methods (TODO need to figure out how these need to move to session somehow...)
 // ORT_API_STATUS(DmlExecutionProviderSetDefaultRoundingMode, _In_ const OrtExecutionProvider* dml_provider, _In_ bool is_enabled);
-// ORT_API_STATUS(DmlExecutionProviderFlushContext, _In_ OrtExecutionProvider * dml_provider);
-// ORT_API_STATUS(DmlExecutionProviderTrimUploadHeap, _In_ const OrtExecutionProvider* dml_provider);
-// ORT_API_STATUS(DmlExecutionProviderReleaseCompletedReferences, _In_ const OrtExecutionProvider* dml_provider);
+ORT_API_STATUS(DmlExecutionProviderFlushContext, _In_ OrtExecutionProvider * dml_provider);
+ORT_API_STATUS(DmlExecutionProviderTrimUploadHeap, _In_ OrtExecutionProvider* dml_provider);
+ORT_API_STATUS(DmlExecutionProviderReleaseCompletedReferences, _In_ OrtExecutionProvider* dml_provider);
 // ORT_API_STATUS(DmlCreateGPUAllocationFromD3DResource, _In_ ID3D12Resource* pResource, _Out_ void* dml_resource);
 // ORT_API_STATUS(DmlFreeGPUAllocation, _In_ void* ptr);
 

diff --git a/winml/adapter/winml_adapter_c_api.cpp b/winml/adapter/winml_adapter_c_api.cpp
@@ -77,9 +77,9 @@ static constexpr WinmlAdapterApi winml_adapter_api_1 = {
 
   // Dml methods (TODO need to figure out how these need to move to session somehow...)
   nullptr, //OrtStatus*(ORT_API_CALL* DmlExecutionProviderSetDefaultRoundingMode)(_In_ bool is_enabled)NO_EXCEPTION;
-  nullptr, // OrtStatus*(ORT_API_CALL* DmlExecutionProviderFlushContext(onnxruntime::IExecutionProvider * dml_provider)NO_EXCEPTION;
-  nullptr, // OrtStatus*(ORT_API_CALL* DmlExecutionProviderTrimUploadHeap(onnxruntime::IExecutionProvider* dml_provider)NO_EXCEPTION;
-  nullptr, // OrtStatus*(ORT_API_CALL* DmlExecutionProviderReleaseCompletedReferences(onnxruntime::IExecutionProvider* dml_provider)NO_EXCEPTION;
+  &winmla::DmlExecutionProviderFlushContext,
+  &winmla::DmlExecutionProviderTrimUploadHeap,
+  &winmla::DmlExecutionProviderReleaseCompletedReferences,
 
   nullptr, // OrtStatus*(ORT_API_CALL* DmlCreateGPUAllocationFromD3DResource)(_In_ ID3D12Resource* pResource, _Out_ void* dml_resource)NO_EXCEPTION;
   nullptr, // OrtStatus*(ORT_API_CALL* DmlFreeGPUAllocation)(_In_ void* ptr)NO_EXCEPTION;

diff --git a/winml/adapter/winml_adapter_dml.cpp b/winml/adapter/winml_adapter_dml.cpp
@@ -9,6 +9,7 @@
 #include "core/framework/error_code_helper.h"
 
 #include "core/providers/dml/dml_provider_factory.h"
+#include "core/providers/dml/DmlExecutionProvider/inc/DmlExecutionProvider.h"
 
 namespace winmla = Windows::AI::MachineLearning::Adapter;
 
@@ -45,4 +46,28 @@ ORT_API_STATUS_IMPL(winmla::OrtSessionOptionsAppendExecutionProviderEx_DML, _In_
                     ID3D12Device* d3d_device, ID3D12CommandQueue* queue) {
   auto dml_device = CreateDmlDevice(d3d_device);
   return OrtSessionOptionsAppendExecutionProviderEx_DML(options, dml_device.Get(), queue);
+}
+
+ORT_API_STATUS_IMPL(winmla::DmlExecutionProviderFlushContext, _In_ OrtExecutionProvider* dml_provider) {
+  API_IMPL_BEGIN
+  auto dml_provider_internal = reinterpret_cast<::onnxruntime::IExecutionProvider*>(dml_provider);
+  Dml::FlushContext(dml_provider_internal);
+  return nullptr;
+  API_IMPL_END
+}
+
+ORT_API_STATUS_IMPL(winmla::DmlExecutionProviderTrimUploadHeap, _In_ OrtExecutionProvider* dml_provider) {
+  API_IMPL_BEGIN
+  auto dml_provider_internal = reinterpret_cast<::onnxruntime::IExecutionProvider*>(dml_provider);
+  Dml::TrimUploadHeap(dml_provider_internal);
+  return nullptr;
+  API_IMPL_END
+}
+
+ORT_API_STATUS_IMPL(winmla::DmlExecutionProviderReleaseCompletedReferences, _In_ OrtExecutionProvider* dml_provider) {
+  API_IMPL_BEGIN
+  auto dml_provider_internal = reinterpret_cast<::onnxruntime::IExecutionProvider*>(dml_provider);
+  Dml::ReleaseCompletedReferences(dml_provider_internal);
+  return nullptr;
+  API_IMPL_END
 }
diff --git a/winml/lib/Api.Ort/OnnxruntimeEngine.cpp b/winml/lib/Api.Ort/OnnxruntimeEngine.cpp
@@ -73,15 +73,34 @@ HRESULT OnnxruntimeEngine::StartProfiling() {
 }
 
 HRESULT OnnxruntimeEngine::FlushContext() {
-  return E_NOTIMPL;
+  auto winml_adapter_api = engine_factory_->UseWinmlAdapterApi();
+
+  const OrtExecutionProvider* ort_provider;
+  winml_adapter_api->SessionGetExecutionProvider(session_.get(), 0, &ort_provider);
+
+  winml_adapter_api->DmlExecutionProviderFlushContext(ort_provider);
+  return S_OK;
 }
 
 HRESULT OnnxruntimeEngine::TrimUploadHeap() {
-  return E_NOTIMPL;
+  auto winml_adapter_api = engine_factory_->UseWinmlAdapterApi();
+
+  const OrtExecutionProvider* ort_provider;
+  winml_adapter_api->SessionGetExecutionProvider(session_.get(), 0, &ort_provider);
+
+  winml_adapter_api->DmlExecutionProviderTrimUploadHeap(ort_provider);
+  return S_OK;
+
 }
 
 HRESULT OnnxruntimeEngine::ReleaseCompletedReferences() {
-  return E_NOTIMPL;
+  auto winml_adapter_api = engine_factory_->UseWinmlAdapterApi();
+
+  const OrtExecutionProvider* ort_provider;
+  winml_adapter_api->SessionGetExecutionProvider(session_.get(), 0, &ort_provider);
+
+  winml_adapter_api->DmlExecutionProviderReleaseCompletedReferences(ort_provider);
+  return S_OK;
 }
 
 HRESULT OnnxruntimeEngine::CopyOneInputAcrossDevices(const char* input_name, const IValue* src, IValue** dest) {