Merge pull request #1226 from hdelan/get-native-mem-on-device2

[UR] Add extra param to urMemGetNativeHandle

Author: kbenzie

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