[UR] Add extra param to urMemGetNativeHandle

include/ur_api.h

@@ -2723,13 +2723,15 @@ urMemBufferPartition(
 ///     - ::UR_RESULT_ERROR_ADAPTER_SPECIFIC
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_HANDLE
 ///         + `NULL == hMem`
+///         + `NULL == hDevice`
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
 ///         + `NULL == phNativeMem`
 ///     - ::UR_RESULT_ERROR_UNSUPPORTED_FEATURE
 ///         + If the adapter has no underlying equivalent handle.
 UR_APIEXPORT ur_result_t UR_APICALL
 urMemGetNativeHandle(
     ur_mem_handle_t hMem,           ///< [in] handle of the mem.
+    ur_device_handle_t hDevice,     ///< [in] handle of the device that the native handle will be resident on.
     ur_native_handle_t *phNativeMem ///< [out] a pointer to the native handle of the mem.
 );
 
@@ -9488,6 +9490,7 @@ typedef struct ur_mem_buffer_partition_params_t {
 ///     allowing the callback the ability to modify the parameter's value
 typedef struct ur_mem_get_native_handle_params_t {
     ur_mem_handle_t *phMem;
+    ur_device_handle_t *phDevice;
     ur_native_handle_t **pphNativeMem;
 } ur_mem_get_native_handle_params_t;
 

include/ur_ddi.h

@@ -770,6 +770,7 @@ typedef ur_result_t(UR_APICALL *ur_pfnMemBufferPartition_t)(
 /// @brief Function-pointer for urMemGetNativeHandle
 typedef ur_result_t(UR_APICALL *ur_pfnMemGetNativeHandle_t)(
     ur_mem_handle_t,
+    ur_device_handle_t,
     ur_native_handle_t *);
 
 ///////////////////////////////////////////////////////////////////////////////

include/ur_print.hpp

@@ -11174,6 +11174,12 @@ inline std::ostream &operator<<(std::ostream &os, [[maybe_unused]] const struct
     ur::details::printPtr(os,
                           *(params->phMem));
 
+    os << ", ";
+    os << ".hDevice = ";
+
+    ur::details::printPtr(os,
+                          *(params->phDevice));
+
     os << ", ";
     os << ".phNativeMem = ";
 

scripts/core/memory.yml

@@ -432,6 +432,10 @@ params:
       name: hMem
       desc: |
             [in] handle of the mem.
+    - type: $x_device_handle_t
+      name: hDevice
+      desc: |
+            [in] handle of the device that the native handle will be resident on.
     - type: $x_native_handle_t*
       name: phNativeMem
       desc: |

source/adapters/cuda/memory.cpp

@@ -161,8 +161,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemRelease(ur_mem_handle_t hMem) {
 /// \param[out] phNativeMem Set to the native handle of the UR mem object.
 ///
 /// \return UR_RESULT_SUCCESS
-UR_APIEXPORT ur_result_t UR_APICALL
-urMemGetNativeHandle(ur_mem_handle_t hMem, ur_native_handle_t *phNativeMem) {
+UR_APIEXPORT ur_result_t UR_APICALL urMemGetNativeHandle(
+    ur_mem_handle_t hMem, ur_device_handle_t, ur_native_handle_t *phNativeMem) {
   *phNativeMem = reinterpret_cast<ur_native_handle_t>(
       std::get<BufferMem>(hMem->Mem).get());
   return UR_RESULT_SUCCESS;

source/adapters/hip/memory.cpp

@@ -279,16 +279,32 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemGetInfo(ur_mem_handle_t hMemory,
 /// \param[out] phNativeMem Set to the native handle of the UR mem object.
 ///
 /// \return UR_RESULT_SUCCESS
-UR_APIEXPORT ur_result_t UR_APICALL urMemGetNativeHandle(ur_mem_handle_t,
-                                                         ur_native_handle_t *) {
-  // FIXME: there is no good way of doing this with a multi device context.
-  // If we return a single pointer, how would we know which device's allocation
-  // it should be?
-  // If we return a vector of pointers, this is OK for read only access but if
-  // we write to a buffer, how would we know which one had been written to?
-  // Should unused allocations be updated afterwards? We have no way of knowing
-  // any of these things in the current API design.
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+UR_APIEXPORT ur_result_t UR_APICALL
+urMemGetNativeHandle(ur_mem_handle_t hMem, ur_device_handle_t Device,
+                     ur_native_handle_t *phNativeMem) {
+#if defined(__HIP_PLATFORM_NVIDIA__)
+  if (sizeof(BufferMem::native_type) > sizeof(ur_native_handle_t)) {
+    // Check that all the upper bits that cannot be represented by
+    // ur_native_handle_t are empty.
+    // NOTE: The following shift might trigger a warning, but the check in the
+    // if above makes sure that this does not underflow.
+    BufferMem::native_type UpperBits =
+        std::get<BufferMem>(hMem->Mem).getPtr(Device) >>
+        (sizeof(ur_native_handle_t) * CHAR_BIT);
+    if (UpperBits) {
+      // Return an error if any of the remaining bits is non-zero.
+      return UR_RESULT_ERROR_INVALID_MEM_OBJECT;
+    }
+  }
+  *phNativeMem = reinterpret_cast<ur_native_handle_t>(
+      std::get<BufferMem>(hMem->Mem).getPtr(Device));
+#elif defined(__HIP_PLATFORM_AMD__)
+  *phNativeMem = reinterpret_cast<ur_native_handle_t>(
+      std::get<BufferMem>(hMem->Mem).getPtr(Device));
+#else
+#error("Must define exactly one of __HIP_PLATFORM_AMD__ or __HIP_PLATFORM_NVIDIA__");
+#endif
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urMemBufferCreateWithNativeHandle(

source/adapters/level_zero/memory.cpp

@@ -1856,6 +1856,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemBufferPartition(
 
 UR_APIEXPORT ur_result_t UR_APICALL urMemGetNativeHandle(
     ur_mem_handle_t Mem, ///< [in] handle of the mem.
+    ur_device_handle_t,  ///< [in] handle of the device.
     ur_native_handle_t
         *NativeMem ///< [out] a pointer to the native handle of the mem.
 ) {

source/adapters/native_cpu/memory.cpp

@@ -111,8 +111,10 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemBufferPartition(
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL
-urMemGetNativeHandle(ur_mem_handle_t hMem, ur_native_handle_t *phNativeMem) {
+urMemGetNativeHandle(ur_mem_handle_t hMem, ur_device_handle_t hDevice,
+                     ur_native_handle_t *phNativeMem) {
   std::ignore = hMem;
+  std::ignore = hDevice;
   std::ignore = phNativeMem;
 
   DIE_NO_IMPLEMENTATION

source/adapters/null/ur_nullddi.cpp

@@ -916,6 +916,8 @@ __urdlllocal ur_result_t UR_APICALL urMemBufferPartition(
 /// @brief Intercept function for urMemGetNativeHandle
 __urdlllocal ur_result_t UR_APICALL urMemGetNativeHandle(
     ur_mem_handle_t hMem, ///< [in] handle of the mem.
+    ur_device_handle_t
+        hDevice, ///< [in] handle of the device that the native handle will be resident on.
     ur_native_handle_t
         *phNativeMem ///< [out] a pointer to the native handle of the mem.
     ) try {
@@ -924,7 +926,7 @@ __urdlllocal ur_result_t UR_APICALL urMemGetNativeHandle(
     // if the driver has created a custom function, then call it instead of using the generic path
     auto pfnGetNativeHandle = d_context.urDdiTable.Mem.pfnGetNativeHandle;
     if (nullptr != pfnGetNativeHandle) {
-        result = pfnGetNativeHandle(hMem, phNativeMem);
+        result = pfnGetNativeHandle(hMem, hDevice, phNativeMem);
     } else {
         // generic implementation
         *phNativeMem = reinterpret_cast<ur_native_handle_t>(d_context.get());

source/adapters/opencl/memory.cpp

@@ -331,8 +331,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemBufferPartition(
   return mapCLErrorToUR(RetErr);
 }
 
-UR_APIEXPORT ur_result_t UR_APICALL
-urMemGetNativeHandle(ur_mem_handle_t hMem, ur_native_handle_t *phNativeMem) {
+UR_APIEXPORT ur_result_t UR_APICALL urMemGetNativeHandle(
+    ur_mem_handle_t hMem, ur_device_handle_t, ur_native_handle_t *phNativeMem) {
   return getNativeHandle(hMem, phNativeMem);
 }
 

source/loader/layers/tracing/ur_trcddi.cpp

@@ -986,6 +986,8 @@ __urdlllocal ur_result_t UR_APICALL urMemBufferPartition(
 /// @brief Intercept function for urMemGetNativeHandle
 __urdlllocal ur_result_t UR_APICALL urMemGetNativeHandle(
     ur_mem_handle_t hMem, ///< [in] handle of the mem.
+    ur_device_handle_t
+        hDevice, ///< [in] handle of the device that the native handle will be resident on.
     ur_native_handle_t
         *phNativeMem ///< [out] a pointer to the native handle of the mem.
 ) {
@@ -995,11 +997,11 @@ __urdlllocal ur_result_t UR_APICALL urMemGetNativeHandle(
         return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
     }
 
-    ur_mem_get_native_handle_params_t params = {&hMem, &phNativeMem};
+    ur_mem_get_native_handle_params_t params = {&hMem, &hDevice, &phNativeMem};
     uint64_t instance = context.notify_begin(UR_FUNCTION_MEM_GET_NATIVE_HANDLE,
                                              "urMemGetNativeHandle", &params);
 
-    ur_result_t result = pfnGetNativeHandle(hMem, phNativeMem);
+    ur_result_t result = pfnGetNativeHandle(hMem, hDevice, phNativeMem);
 
     context.notify_end(UR_FUNCTION_MEM_GET_NATIVE_HANDLE,
                        "urMemGetNativeHandle", &params, &result, instance);

source/loader/layers/validation/ur_valddi.cpp

@@ -1188,6 +1188,8 @@ __urdlllocal ur_result_t UR_APICALL urMemBufferPartition(
 /// @brief Intercept function for urMemGetNativeHandle
 __urdlllocal ur_result_t UR_APICALL urMemGetNativeHandle(
     ur_mem_handle_t hMem, ///< [in] handle of the mem.
+    ur_device_handle_t
+        hDevice, ///< [in] handle of the device that the native handle will be resident on.
     ur_native_handle_t
         *phNativeMem ///< [out] a pointer to the native handle of the mem.
 ) {
@@ -1202,12 +1204,16 @@ __urdlllocal ur_result_t UR_APICALL urMemGetNativeHandle(
             return UR_RESULT_ERROR_INVALID_NULL_HANDLE;
         }
 
+        if (NULL == hDevice) {
+            return UR_RESULT_ERROR_INVALID_NULL_HANDLE;
+        }
+
         if (NULL == phNativeMem) {
             return UR_RESULT_ERROR_INVALID_NULL_POINTER;
         }
     }
 
-    ur_result_t result = pfnGetNativeHandle(hMem, phNativeMem);
+    ur_result_t result = pfnGetNativeHandle(hMem, hDevice, phNativeMem);
 
     return result;
 }

source/loader/ur_ldrddi.cpp

@@ -1231,6 +1231,8 @@ __urdlllocal ur_result_t UR_APICALL urMemBufferPartition(
 /// @brief Intercept function for urMemGetNativeHandle
 __urdlllocal ur_result_t UR_APICALL urMemGetNativeHandle(
     ur_mem_handle_t hMem, ///< [in] handle of the mem.
+    ur_device_handle_t
+        hDevice, ///< [in] handle of the device that the native handle will be resident on.
     ur_native_handle_t
         *phNativeMem ///< [out] a pointer to the native handle of the mem.
 ) {
@@ -1246,8 +1248,11 @@ __urdlllocal ur_result_t UR_APICALL urMemGetNativeHandle(
     // convert loader handle to platform handle
     hMem = reinterpret_cast<ur_mem_object_t *>(hMem)->handle;
 
+    // convert loader handle to platform handle
+    hDevice = reinterpret_cast<ur_device_object_t *>(hDevice)->handle;
+
     // forward to device-platform
-    result = pfnGetNativeHandle(hMem, phNativeMem);
+    result = pfnGetNativeHandle(hMem, hDevice, phNativeMem);
 
     if (UR_RESULT_SUCCESS != result) {
         return result;

source/loader/ur_libapi.cpp

@@ -1654,12 +1654,15 @@ ur_result_t UR_APICALL urMemBufferPartition(
 ///     - ::UR_RESULT_ERROR_ADAPTER_SPECIFIC
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_HANDLE
 ///         + `NULL == hMem`
+///         + `NULL == hDevice`
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
 ///         + `NULL == phNativeMem`
 ///     - ::UR_RESULT_ERROR_UNSUPPORTED_FEATURE
 ///         + If the adapter has no underlying equivalent handle.
 ur_result_t UR_APICALL urMemGetNativeHandle(
     ur_mem_handle_t hMem, ///< [in] handle of the mem.
+    ur_device_handle_t
+        hDevice, ///< [in] handle of the device that the native handle will be resident on.
     ur_native_handle_t
         *phNativeMem ///< [out] a pointer to the native handle of the mem.
     ) try {
@@ -1669,7 +1672,7 @@ ur_result_t UR_APICALL urMemGetNativeHandle(
         return UR_RESULT_ERROR_UNINITIALIZED;
     }
 
-    return pfnGetNativeHandle(hMem, phNativeMem);
+    return pfnGetNativeHandle(hMem, hDevice, phNativeMem);
 } catch (...) {
     return exceptionToResult(std::current_exception());
 }

source/ur_api.cpp

@@ -1419,12 +1419,15 @@ ur_result_t UR_APICALL urMemBufferPartition(
 ///     - ::UR_RESULT_ERROR_ADAPTER_SPECIFIC
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_HANDLE
 ///         + `NULL == hMem`
+///         + `NULL == hDevice`
 ///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
 ///         + `NULL == phNativeMem`
 ///     - ::UR_RESULT_ERROR_UNSUPPORTED_FEATURE
 ///         + If the adapter has no underlying equivalent handle.
 ur_result_t UR_APICALL urMemGetNativeHandle(
     ur_mem_handle_t hMem, ///< [in] handle of the mem.
+    ur_device_handle_t
+        hDevice, ///< [in] handle of the device that the native handle will be resident on.
     ur_native_handle_t
         *phNativeMem ///< [out] a pointer to the native handle of the mem.
 ) {

test/conformance/memory/memory_adapter_hip.match

@@ -1,7 +1,4 @@
-{{OPT}}urMemGetInfoTest.InvalidNullPointerParamValue/AMD_HIP_BACKEND___{{.*}}
-{{OPT}}urMemGetInfoTest.InvalidNullPointerParamValue/AMD_HIP_BACKEND___{{.*}}
-{{OPT}}urMemGetInfoTest.InvalidNullPointerPropSizeRet/AMD_HIP_BACKEND___{{.*}}
-{{OPT}}urMemGetInfoTest.InvalidNullPointerPropSizeRet/AMD_HIP_BACKEND___{{.*}}
+{{OPT}}urMemBufferCreateWithNativeHandleTest.Success/AMD_HIP_BACKEND___{{.*}}
 {{OPT}}urMemImageCreateTest.InvalidSize/AMD_HIP_BACKEND___{{.*}}
 {{OPT}}urMemImageGetInfoTest.Success/AMD_HIP_BACKEND___{{.*}}
 {{OPT}}urMemImageGetInfoTest.Success/AMD_HIP_BACKEND___{{.*}}

test/conformance/memory/urMemBufferCreateWithNativeHandle.cpp

@@ -10,7 +10,7 @@ UUR_INSTANTIATE_DEVICE_TEST_SUITE_P(urMemBufferCreateWithNativeHandleTest);
 
 TEST_P(urMemBufferCreateWithNativeHandleTest, Success) {
     ur_native_handle_t hNativeMem = nullptr;
-    if (urMemGetNativeHandle(buffer, &hNativeMem)) {
+    if (urMemGetNativeHandle(buffer, device, &hNativeMem)) {
         GTEST_SKIP();
     }
 

test/conformance/memory/urMemGetNativeHandle.cpp

@@ -9,18 +9,24 @@ UUR_INSTANTIATE_DEVICE_TEST_SUITE_P(urMemGetNativeHandleTest);
 
 TEST_P(urMemGetNativeHandleTest, Success) {
     ur_native_handle_t hNativeMem = nullptr;
-    if (auto error = urMemGetNativeHandle(buffer, &hNativeMem)) {
+    if (auto error = urMemGetNativeHandle(buffer, device, &hNativeMem)) {
         ASSERT_EQ_RESULT(UR_RESULT_ERROR_UNSUPPORTED_FEATURE, error);
     }
 }
 
 TEST_P(urMemGetNativeHandleTest, InvalidNullHandleMem) {
     ur_native_handle_t phNativeMem;
     ASSERT_EQ_RESULT(UR_RESULT_ERROR_INVALID_NULL_HANDLE,
-                     urMemGetNativeHandle(nullptr, &phNativeMem));
+                     urMemGetNativeHandle(nullptr, device, &phNativeMem));
+}
+
+TEST_P(urMemGetNativeHandleTest, InvalidNullHandleDevice) {
+    ur_native_handle_t phNativeMem;
+    ASSERT_EQ_RESULT(UR_RESULT_ERROR_INVALID_NULL_HANDLE,
+                     urMemGetNativeHandle(buffer, nullptr, &phNativeMem));
 }
 
 TEST_P(urMemGetNativeHandleTest, InvalidNullPointerNativeMem) {
     ASSERT_EQ_RESULT(UR_RESULT_ERROR_INVALID_NULL_POINTER,
-                     urMemGetNativeHandle(buffer, nullptr));
+                     urMemGetNativeHandle(buffer, device, nullptr));
 }

test/conformance/memory/urMemImageCreateWithNativeHandle.cpp

@@ -10,7 +10,7 @@ UUR_INSTANTIATE_DEVICE_TEST_SUITE_P(urMemImageCreateWithNativeHandleTest);
 
 TEST_P(urMemImageCreateWithNativeHandleTest, Success) {
     ur_native_handle_t native_handle = nullptr;
-    if (urMemGetNativeHandle(image, &native_handle)) {
+    if (urMemGetNativeHandle(image, device, &native_handle)) {
         GTEST_SKIP();
     }