Preserve the original handle ptr

Signed-off-by: Larsen, Steffen <[email protected]>

Author: steffenlarsen

sycl/doc/extensions/experimental/sycl_ext_oneapi_inter_process_communication.asciidoc

@@ -117,25 +117,57 @@ functions specified in the following section.
 
 === Inter-process communicable memory
 
-
-This extension adds new free functions under the `ipc_memory` experimental
-namespace.
+This extension adds a new type alias `handle_data_t` and a new IPC memory handle
+type `handle` under the `ipc_memory` experimental namespace.
 
 ```
 namespace sycl::ext::oneapi::experimental::ipc_memory {
 
-using handle_data_t = std::vector<std::byte>;
+using handle_data_t = sycl::span<const std::byte, dynamic_extent>;
+
+struct handle {
+  handle_data_t data() const;
+};
+
+}
+```
+
+|====
+a|
+[frame=all,grid=none]
+!====
+a!
+[source]
+----
+handle_data_t data() const;
+----
+!====
+
+_Preconditions:_ The `put` function has not previously been called with this
+handle and the USM device memory associated with this handle has not been freed.
+
+_Returns:_ The handle data associated with the IPC handle object. This data can
+be transferred to other processes, but cannot be used to recreate a `handle`
+object.
+
+!====
+
+Additionally, this extension adds new free functions under the `ipc_memory`
+experimental namespace.
+
+```
+namespace sycl::ext::oneapi::experimental::ipc_memory {
 
-handle_data_t get(void *ptr, const sycl::context &ctx);
+handle get(void *ptr, const sycl::context &ctx);
 
-handle_data_t get(void *ptr);
+handle get(void *ptr);
 
-void put(const handle_data_t &handle_data, const sycl::context &ctx);
+void put(handle &ipc_handle, const sycl::context &ctx);
 
-void put(const handle_data_t &handle_data);
+void put(handle &ipc_handle);
 
 void *open(const handle_data_t &handle_data, const sycl::context &ctx,
-                  const sycl::device &dev);
+           const sycl::device &dev);
 
 void *open(const handle_data_t &handle_data, const sycl::device &dev);
 
@@ -155,7 +187,7 @@ a|
 a!
 [source]
 ----
-handle_data_t get(void *ptr, const sycl::context &ctx)
+handle get(void *ptr, const sycl::context &ctx)
 ----
 !====
 
@@ -175,7 +207,7 @@ device _D_ does not have `aspect::ext_oneapi_ipc_memory`.
 a!
 [source]
 ----
-handle_data_t get(void *ptr)
+handle get(void *ptr)
 ----
 !====
 
@@ -192,13 +224,14 @@ return ipc_memory::get(ptr, ctxt);
 a!
 [source]
 ----
-void put(const handle_data_t &handle_data, const sycl::context &ctx)
+void put(handle &ipc_handle, const sycl::context &ctx)
 ----
 !====
 
-_Preconditions:_ `handle_data` is the IPC "handle" to USM device memory that was
+_Preconditions:_ `ipc_handle` is the IPC "handle" to USM device memory that was
 returned from a call to `get` either in this process or in some other process on
-the same system. The USM device memory has not yet been freed in this process.
+the same system. The `put` function has not been previously called on the
+handle.
 
 _Effects:_ Deallocates resources associated with the handle. These resources are
 automatically deallocated when the USM device memory is freed, so it is not
@@ -216,7 +249,7 @@ _{endnote}_]
 a!
 [source]
 ----
-void put(const handle_data_t &handle_data)
+void put(handle &ipc_handle)
 ----
 !====
 

sycl/include/sycl/ext/oneapi/experimental/ipc_memory.hpp

@@ -15,43 +15,58 @@
 #include <sycl/detail/export.hpp>
 #include <sycl/device.hpp>
 #include <sycl/platform.hpp>
+#include <sycl/sycl_span.hpp>
 
 #include <cstddef>
-#include <vector>
 
 namespace sycl {
 inline namespace _V1 {
 
 namespace ext::oneapi::experimental::ipc_memory {
 
-using handle_data_t = std::vector<std::byte>;
+using handle_data_t = sycl::span<const std::byte, dynamic_extent>;
 
-__SYCL_EXPORT handle_data_t get(void *Ptr, const sycl::context &Ctx);
+struct handle {
+public:
+  handle_data_t data() const {
+    return {reinterpret_cast<std::byte *>(MData), MSize};
+  }
 
-inline handle_data_t get(void *Ptr) {
+private:
+  handle(void *Data, size_t Size) : MData{Data}, MSize{Size} {}
+
+  void *MData;
+  size_t MSize;
+
+  friend handle get(void *Ptr, const sycl::context &Ctx);
+  friend void put(handle &HandleData, const sycl::context &Ctx);
+};
+
+__SYCL_EXPORT handle get(void *Ptr, const sycl::context &Ctx);
+
+inline handle get(void *Ptr) {
   sycl::device Dev;
   sycl::context Ctx = Dev.get_platform().khr_get_default_context();
   return ipc_memory::get(Ptr, Ctx);
 }
 
-__SYCL_EXPORT void put(const handle_data_t &HandleData,
-                       const sycl::context &Ctx);
+__SYCL_EXPORT void put(handle &HandleData, const sycl::context &Ctx);
 
-inline void put(const handle_data_t &HandleData) {
+inline void put(handle &HandleData) {
   sycl::device Dev;
   sycl::context Ctx = Dev.get_platform().khr_get_default_context();
   ipc_memory::put(HandleData, Ctx);
 }
 
-__SYCL_EXPORT void *open(const handle_data_t &HandleData,
-                         const sycl::context &Ctx, const sycl::device &Dev);
+__SYCL_EXPORT void *open(handle_data_t HandleData, const sycl::context &Ctx,
+                         const sycl::device &Dev);
 
-inline void *open(const handle_data_t &HandleData, const sycl::device &Dev) {
+inline void *open(handle_data_t HandleData, const sycl::device &Dev) {
   sycl::context Ctx = Dev.get_platform().khr_get_default_context();
   return ipc_memory::open(HandleData, Ctx, Dev);
 }
 
-inline void *open(const handle_data_t &HandleData) {
+inline void *open(handle_data_t HandleData) {
   sycl::device Dev;
   sycl::context Ctx = Dev.get_platform().khr_get_default_context();
   return ipc_memory::open(HandleData, Ctx, Dev);

sycl/source/ipc_memory.cpp

@@ -16,7 +16,7 @@ namespace sycl {
 inline namespace _V1 {
 namespace ext::oneapi::experimental::ipc_memory {
 
-__SYCL_EXPORT handle_data_t get(void *Ptr, const sycl::context &Ctx) {
+__SYCL_EXPORT handle get(void *Ptr, const sycl::context &Ctx) {
   auto CtxImpl = sycl::detail::getSyclObjImpl(Ctx);
   sycl::detail::adapter_impl &Adapter = CtxImpl->getAdapter();
 
@@ -30,42 +30,26 @@ __SYCL_EXPORT handle_data_t get(void *Ptr, const sycl::context &Ctx) {
           "Device does not support aspect::ext_oneapi_ipc_memory.");
   };
 
+  void *HandlePtr = nullptr;
   size_t HandleSize = 0;
   auto UrRes =
       Adapter.call_nocheck<sycl::detail::UrApiKind::urIPCGetMemHandleExp>(
-          CtxImpl->getHandleRef(), Ptr, nullptr, &HandleSize);
+          CtxImpl->getHandleRef(), Ptr, &HandlePtr, &HandleSize);
   if (UrRes != UR_RESULT_SUCCESS) {
     CheckDeviceSupport();
     Adapter.checkUrResult(UrRes);
   }
-
-  handle_data_t Res(HandleSize);
-  UrRes = Adapter.call_nocheck<sycl::detail::UrApiKind::urIPCGetMemHandleExp>(
-      CtxImpl->getHandleRef(), Ptr, Res.data(), nullptr);
-  if (UrRes != UR_RESULT_SUCCESS) {
-    CheckDeviceSupport();
-    Adapter.checkUrResult(UrRes);
-  }
-  return Res;
+  return {HandlePtr, HandleSize};
 }
 
-__SYCL_EXPORT void put(const handle_data_t &HandleData,
-                       const sycl::context &Ctx) {
-  // TODO: UMF and UR currently requires the handle data to be non-const, so we
-  //       need to make a copy of the data. Once this has been changed, the copy
-  //       can be removed.
-  //       CMPLRLLVM-71181
-  //       https://github.com/oneapi-src/unified-memory-framework/issues/1536
-  handle_data_t HandleDataCopy = HandleData;
-
+__SYCL_EXPORT void put(handle &Handle, const sycl::context &Ctx) {
   auto CtxImpl = sycl::detail::getSyclObjImpl(Ctx);
-  sycl::detail::adapter_impl &Adapter = CtxImpl->getAdapter();
-  Adapter.call<sycl::detail::UrApiKind::urIPCPutMemHandleExp>(
-      CtxImpl->getHandleRef(), HandleDataCopy.data());
+  CtxImpl->getAdapter().call<sycl::detail::UrApiKind::urIPCPutMemHandleExp>(
+      CtxImpl->getHandleRef(), Handle.MData);
 }
 
-__SYCL_EXPORT void *open(const handle_data_t &HandleData,
-                         const sycl::context &Ctx, const sycl::device &Dev) {
+__SYCL_EXPORT void *open(handle_data_t HandleData, const sycl::context &Ctx,
+                         const sycl::device &Dev) {
   if (!Dev.has(aspect::ext_oneapi_ipc_memory))
     throw sycl::exception(
         sycl::make_error_code(errc::feature_not_supported),
@@ -75,17 +59,17 @@ __SYCL_EXPORT void *open(const handle_data_t &HandleData,
   sycl::detail::adapter_impl &Adapter = CtxImpl->getAdapter();
 
   // TODO: UMF and UR currently requires the handle data to be non-const, so we
-  //       need to make a copy of the data. Once this has been changed, the copy
-  //       can be removed.
+  //       need const-cast the data pointer. Once this has been changed, the
+  //       const-cast can be removed.
   //       CMPLRLLVM-71181
   //       https://github.com/oneapi-src/unified-memory-framework/issues/1536
-  handle_data_t HandleDataCopy = HandleData;
+  std::byte *HandleDataPtr = const_cast<std::byte *>(HandleData.data());
 
   void *Ptr = nullptr;
   ur_result_t UrRes =
       Adapter.call_nocheck<sycl::detail::UrApiKind::urIPCOpenMemHandleExp>(
           CtxImpl->getHandleRef(), getSyclObjImpl(Dev)->getHandleRef(),
-          HandleDataCopy.data(), HandleDataCopy.size(), &Ptr);
+          HandleDataPtr, HandleData.size(), &Ptr);
   if (UrRes == UR_RESULT_ERROR_INVALID_VALUE)
     throw sycl::exception(sycl::make_error_code(errc::invalid),
                           "HandleData data size does not correspond "
@@ -96,8 +80,7 @@ __SYCL_EXPORT void *open(const handle_data_t &HandleData,
 
 __SYCL_EXPORT void close(void *Ptr, const sycl::context &Ctx) {
   auto CtxImpl = sycl::detail::getSyclObjImpl(Ctx);
-  sycl::detail::adapter_impl &Adapter = CtxImpl->getAdapter();
-  Adapter.call<sycl::detail::UrApiKind::urIPCCloseMemHandleExp>(
+  CtxImpl->getAdapter().call<sycl::detail::UrApiKind::urIPCCloseMemHandleExp>(
       CtxImpl->getHandleRef(), Ptr);
 }
 

sycl/test-e2e/Experimental/ipc_memory.cpp

@@ -48,8 +48,9 @@ int spawner(int argc, char *argv[]) {
   {
     // Write handle data to file.
     {
-      syclexp::ipc_memory::handle_data_t HandleData =
+      syclexp::ipc_memory::handle Handle =
           syclexp::ipc_memory::get(DataPtr, Q.get_context());
+      syclexp::ipc_memory::handle_data_t HandleData = Handle.data();
       size_t HandleDataSize = HandleData.size();
       std::fstream FS(CommsFile, std::ios_base::out | std::ios_base::binary);
       FS.write(reinterpret_cast<const char *>(&HandleDataSize), sizeof(size_t));

sycl/test-e2e/Experimental/ipc_put_after_free.cpp

@@ -0,0 +1,49 @@
+// REQUIRES: aspect-usm_device_allocations && aspect-ext_oneapi_ipc_memory
+
+// UNSUPPORTED: level_zero && windows
+// UNSUPPORTED-TRACKER: UMFW-348
+
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+
+#include <sycl/detail/core.hpp>
+#include <sycl/ext/oneapi/experimental/ipc_memory.hpp>
+#include <sycl/usm.hpp>
+
+#include <cstdio>
+#include <cstdlib>
+
+#if defined(__linux__)
+#include <linux/prctl.h>
+#include <sys/prctl.h>
+#include <unistd.h>
+#endif // defined(__linux__)
+
+namespace syclexp = sycl::ext::oneapi::experimental;
+
+int main() {
+  sycl::queue Q;
+
+#if defined(__linux__)
+  // UMF currently requires ptrace permissions to be set for the spawner. As
+  // such we need to set it until this limitation has been addressed.
+  // https://github.com/oneapi-src/unified-memory-framework/tree/main?tab=readme-ov-file#level-zero-memory-provider
+  if (Q.get_backend() == sycl::backend::ext_oneapi_level_zero &&
+      prctl(PR_SET_PTRACER, getppid()) == -1) {
+    std::cout << "Failed to set ptracer permissions!" << std::endl;
+    return 1;
+  }
+#endif // defined(__linux__)
+
+  int *DataPtr = sycl::malloc_device<int>(32, Q);
+  syclexp::ipc_memory::handle Handle =
+      syclexp::ipc_memory::get(DataPtr, Q.get_context());
+
+  // Free data before put.
+  sycl::free(DataPtr, Q);
+
+  // Try calling put after free.
+  syclexp::ipc_memory::put(Handle, Q.get_context());
+
+  return 0;
+}