From 371e1b85ac5cee8d52447c0cd8c4eb196ba5878d Mon Sep 17 00:00:00 2001
From: Aaron Greig <aaron.greig@codeplay.com>
Date: Wed, 1 Nov 2023 11:46:11 +0000
Subject: [PATCH] [OpenCL] Retain native handle objects when properties
 dictate.

---
 source/adapters/opencl/event.cpp   | 13 ++++++++-----
 source/adapters/opencl/kernel.cpp  |  7 +++++--
 source/adapters/opencl/memory.cpp  | 14 ++++++++------
 source/adapters/opencl/program.cpp |  7 +++++--
 4 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/source/adapters/opencl/event.cpp b/source/adapters/opencl/event.cpp
index 78303a0829..64cf410460 100644
--- a/source/adapters/opencl/event.cpp
+++ b/source/adapters/opencl/event.cpp
@@ -50,12 +50,15 @@ convertURProfilingInfoToCL(const ur_profiling_info_t PropName) {
   }
 }
 
-UR_APIEXPORT ur_result_t UR_APICALL urEventCreateWithNativeHandle(
-    ur_native_handle_t hNativeEvent,
-    [[maybe_unused]] ur_context_handle_t hContext,
-    [[maybe_unused]] const ur_event_native_properties_t *pProperties,
-    ur_event_handle_t *phEvent) {
+UR_APIEXPORT ur_result_t UR_APICALL
+urEventCreateWithNativeHandle(ur_native_handle_t hNativeEvent,
+                              [[maybe_unused]] ur_context_handle_t hContext,
+                              const ur_event_native_properties_t *pProperties,
+                              ur_event_handle_t *phEvent) {
   *phEvent = reinterpret_cast<ur_event_handle_t>(hNativeEvent);
+  if (!pProperties || !pProperties->isNativeHandleOwned) {
+    return urEventRetain(*phEvent);
+  }
   return UR_RESULT_SUCCESS;
 }
 
diff --git a/source/adapters/opencl/kernel.cpp b/source/adapters/opencl/kernel.cpp
index 80b1502854..6c688021d7 100644
--- a/source/adapters/opencl/kernel.cpp
+++ b/source/adapters/opencl/kernel.cpp
@@ -335,9 +335,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetNativeHandle(
 
 UR_APIEXPORT ur_result_t UR_APICALL urKernelCreateWithNativeHandle(
     ur_native_handle_t hNativeKernel, ur_context_handle_t, ur_program_handle_t,
-    const ur_kernel_native_properties_t *, ur_kernel_handle_t *phKernel) {
-
+    const ur_kernel_native_properties_t *pProperties,
+    ur_kernel_handle_t *phKernel) {
   *phKernel = reinterpret_cast<ur_kernel_handle_t>(hNativeKernel);
+  if (!pProperties || !pProperties->isNativeHandleOwned) {
+    return urKernelRetain(*phKernel);
+  }
   return UR_RESULT_SUCCESS;
 }
 
diff --git a/source/adapters/opencl/memory.cpp b/source/adapters/opencl/memory.cpp
index 279faad376..8912d333e3 100644
--- a/source/adapters/opencl/memory.cpp
+++ b/source/adapters/opencl/memory.cpp
@@ -331,10 +331,11 @@ urMemGetNativeHandle(ur_mem_handle_t hMem, ur_native_handle_t *phNativeMem) {
 UR_APIEXPORT ur_result_t UR_APICALL urMemBufferCreateWithNativeHandle(
     ur_native_handle_t hNativeMem,
     [[maybe_unused]] ur_context_handle_t hContext,
-    [[maybe_unused]] const ur_mem_native_properties_t *pProperties,
-    ur_mem_handle_t *phMem) {
-
+    const ur_mem_native_properties_t *pProperties, ur_mem_handle_t *phMem) {
   *phMem = reinterpret_cast<ur_mem_handle_t>(hNativeMem);
+  if (!pProperties || !pProperties->isNativeHandleOwned) {
+    return urMemRetain(*phMem);
+  }
   return UR_RESULT_SUCCESS;
 }
 
@@ -343,10 +344,11 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemImageCreateWithNativeHandle(
     [[maybe_unused]] ur_context_handle_t hContext,
     [[maybe_unused]] const ur_image_format_t *pImageFormat,
     [[maybe_unused]] const ur_image_desc_t *pImageDesc,
-    [[maybe_unused]] const ur_mem_native_properties_t *pProperties,
-    ur_mem_handle_t *phMem) {
-
+    const ur_mem_native_properties_t *pProperties, ur_mem_handle_t *phMem) {
   *phMem = reinterpret_cast<ur_mem_handle_t>(hNativeMem);
+  if (!pProperties || !pProperties->isNativeHandleOwned) {
+    return urMemRetain(*phMem);
+  }
   return UR_RESULT_SUCCESS;
 }
 
diff --git a/source/adapters/opencl/program.cpp b/source/adapters/opencl/program.cpp
index 0beca23dab..954c2dc48f 100644
--- a/source/adapters/opencl/program.cpp
+++ b/source/adapters/opencl/program.cpp
@@ -299,9 +299,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urProgramGetNativeHandle(
 
 UR_APIEXPORT ur_result_t UR_APICALL urProgramCreateWithNativeHandle(
     ur_native_handle_t hNativeProgram, ur_context_handle_t,
-    const ur_program_native_properties_t *, ur_program_handle_t *phProgram) {
-
+    const ur_program_native_properties_t *pProperties,
+    ur_program_handle_t *phProgram) {
   *phProgram = reinterpret_cast<ur_program_handle_t>(hNativeProgram);
+  if (!pProperties || !pProperties->isNativeHandleOwned) {
+    return urProgramRetain(*phProgram);
+  }
   return UR_RESULT_SUCCESS;
 }