From 83ead3e2eb7ad01b6e8889cabab9fab21dc18fc1 Mon Sep 17 00:00:00 2001
From: Tiago Koji Castro Shibata <ticastro@microsoft.com>
Date: Thu, 8 Oct 2020 10:18:38 -0700
Subject: [PATCH] Fix com ptr refcount (#5404)

---
 winml/adapter/winml_adapter_dml.cpp     | 3 ++-
 winml/lib/Api.Ort/OnnxruntimeEngine.cpp | 8 ++++----
 winml/test/adapter/AdapterDmlEpTest.cpp | 1 +
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/winml/adapter/winml_adapter_dml.cpp b/winml/adapter/winml_adapter_dml.cpp
index 10e412e68b..02aaf561a9 100644
--- a/winml/adapter/winml_adapter_dml.cpp
+++ b/winml/adapter/winml_adapter_dml.cpp
@@ -146,6 +146,7 @@ ORT_API_STATUS_IMPL(winmla::DmlGetD3D12ResourceFromAllocation, _In_ OrtExecution
       Dml::GetD3D12ResourceFromAllocation(
           dml_provider_internal->GetAllocator(0, ::OrtMemType::OrtMemTypeDefault).get(),
           allocation);
+  (*d3d_resource)->AddRef();
 #endif  // USE_DML USE_DML
   return nullptr;
   API_IMPL_END
@@ -173,4 +174,4 @@ ORT_API_STATUS_IMPL(winmla::DmlCopyTensor, _In_ OrtExecutionProvider* dml_provid
   return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED, "Out of memory");
 #endif  // USE_DML USE_DML
   API_IMPL_END
-}
\ No newline at end of file
+}
diff --git a/winml/lib/Api.Ort/OnnxruntimeEngine.cpp b/winml/lib/Api.Ort/OnnxruntimeEngine.cpp
index 118dd1f758..66cb317289 100644
--- a/winml/lib/Api.Ort/OnnxruntimeEngine.cpp
+++ b/winml/lib/Api.Ort/OnnxruntimeEngine.cpp
@@ -179,11 +179,11 @@ HRESULT OnnxruntimeValue::GetResource(_winml::Resource& out) {
 
   bool is_cpu = false;
   if (SUCCEEDED(IsCpu(&is_cpu)) && !is_cpu) {
-    void* resource;
+    winrt::com_ptr<ID3D12Resource> resource;
     RETURN_HR_IF_NOT_OK_MSG(winml_adapter_api->DmlGetD3D12ResourceFromAllocation(ort_provider, mutable_data,
-                                                                                 reinterpret_cast<ID3D12Resource**>(&resource)),
+                                                                                 resource.put()),
                             ort_api);
-    out = _winml::Resource(resource, [](void*) { /*do nothing, as this pointer is actually a com pointer! */ });
+    out = _winml::Resource(resource.get(), [](void*) { /*do nothing, as this pointer is actually a com pointer! */ });
   } else {
     int is_tensor;
     RETURN_HR_IF_NOT_OK_MSG(ort_api->IsTensor(value_.get(), &is_tensor),
@@ -1360,4 +1360,4 @@ STDAPI CreateOnnxruntimeEngineFactory(_Out_ _winml::IEngineFactory** engine_fact
   RETURN_IF_FAILED(Microsoft::WRL::MakeAndInitialize<OnnxruntimeEngineFactory>(&onnxruntime_engine_factory));
   RETURN_IF_FAILED(onnxruntime_engine_factory.CopyTo(engine_factory));
   return S_OK;
-}
\ No newline at end of file
+}
diff --git a/winml/test/adapter/AdapterDmlEpTest.cpp b/winml/test/adapter/AdapterDmlEpTest.cpp
index e2453827d0..ca503bad17 100644
--- a/winml/test/adapter/AdapterDmlEpTest.cpp
+++ b/winml/test/adapter/AdapterDmlEpTest.cpp
@@ -135,6 +135,7 @@ winrt::com_ptr<ID3D12Resource> CreateD3D12Resource(ID3D12Device& device) {
   return d3d12_resource;
 }
 
+
 void DmlCreateAndFreeGPUAllocationFromD3DResource() {
   GPUTEST;
   winrt::com_ptr<ID3D12Device> device;