pyverbs: Add DevX and UMEM export/import support

Add cross-process sharing support for UMEM and DevX objects through
export/import APIs.

The export/import mechanism allows a UMEM or DevX object allocated in
one process to be shared with another via an opaque data buffer. The
close() method dispatches to the appropriate cleanup based on whether
the object was locally allocated or imported.

Move devx_objs and devx_umems tracking from Mlx5Context to Context
so that ibv_import_device (which returns a base Context, not
Mlx5Context) can properly track imported DevX objects and UMEMs.

Signed-off-by: Linoy Ganti <lganti@nvidia.com>
Signed-off-by: Shachar Kagan <skagan@nvidia.com>
diff --git a/pyverbs/device.pxd b/pyverbs/device.pxd
index a13dd06..26bbe9a 100644
--- a/pyverbs/device.pxd
+++ b/pyverbs/device.pxd
@@ -29,6 +29,8 @@
     cdef object crypto_logins
     cdef object event_channels
     cdef object dmahs
+    cdef object devx_umems
+    cdef object devx_objs
 
 cdef class DeviceAttr(PyverbsObject):
     cdef v.ibv_device_attr dev_attr
diff --git a/pyverbs/device.pyx b/pyverbs/device.pyx
index 137eb0c..25466f1 100644
--- a/pyverbs/device.pyx
+++ b/pyverbs/device.pyx
@@ -125,6 +125,8 @@
         self.crypto_logins = weakref.WeakSet()
         self.event_channels = weakref.WeakSet()
         self.dmahs = weakref.WeakSet()
+        self.devx_umems = weakref.WeakSet()
+        self.devx_objs = weakref.WeakSet()
 
         self.name = kwargs.get('name')
         provider_attr = kwargs.get('attr')
@@ -179,9 +181,11 @@
         if self.context != NULL:
             if self.logger:
                 self.logger.debug('Closing Context')
-            close_weakrefs([self.qps, self.crypto_logins, self.rwq_ind_tbls, self.wqs, self.ccs, self.cqs,
-                            self.dms, self.pds, self.xrcds, self.vars, self.sched_leafs,
-                            self.sched_nodes, self.dr_domains, self.event_channels, self.dmahs])
+            close_weakrefs([self.devx_objs, self.devx_umems, self.qps,
+                            self.crypto_logins, self.rwq_ind_tbls, self.wqs,
+                            self.ccs, self.cqs, self.dms, self.pds, self.xrcds,
+                            self.vars, self.sched_leafs, self.sched_nodes,
+                            self.dr_domains, self.event_channels, self.dmahs])
             rc = v.ibv_close_device(self.context)
             if rc != 0:
                 raise PyverbsRDMAErrno(f'Failed to close device {self.name}')
diff --git a/pyverbs/providers/mlx5/libmlx5.pxd b/pyverbs/providers/mlx5/libmlx5.pxd
index 48f3ee7..eb10ac4 100644
--- a/pyverbs/providers/mlx5/libmlx5.pxd
+++ b/pyverbs/providers/mlx5/libmlx5.pxd
@@ -585,6 +585,9 @@
     mlx5dv_devx_umem *mlx5dv_devx_umem_reg_ex(v.ibv_context *ctx,
                                               mlx5dv_devx_umem_in *umem_in)
     int mlx5dv_devx_umem_dereg(mlx5dv_devx_umem *umem)
+    int mlx5dv_devx_umem_export(mlx5dv_devx_umem *umem, void *data)
+    mlx5dv_devx_umem *mlx5dv_devx_umem_import(v.ibv_context *context, void *data)
+    void mlx5dv_devx_umem_unimport(mlx5dv_devx_umem *umem)
     int mlx5dv_devx_query_eqn(v.ibv_context *context, uint32_t vector, uint32_t *eqn)
     mlx5dv_devx_obj *mlx5dv_devx_obj_create(v.ibv_context *context, const void *_in,
                                             size_t inlen, void *out, size_t outlen)
@@ -600,6 +603,9 @@
     int mlx5dv_devx_obj_modify(mlx5dv_devx_obj *obj, const void *in_,
                                size_t inlen, void *out, size_t outlen)
     int mlx5dv_devx_obj_destroy(mlx5dv_devx_obj *obj)
+    int mlx5dv_devx_obj_export(mlx5dv_devx_obj *obj, void *data)
+    mlx5dv_devx_obj *mlx5dv_devx_obj_import(v.ibv_context *context, void *data)
+    void mlx5dv_devx_obj_unimport(mlx5dv_devx_obj *obj)
     int mlx5dv_init_obj(mlx5dv_obj *obj, uint64_t obj_type)
     mlx5dv_devx_msi_vector *mlx5dv_devx_alloc_msi_vector(v.ibv_context *ibctx)
     int mlx5dv_devx_free_msi_vector(mlx5dv_devx_msi_vector *msi)
diff --git a/pyverbs/providers/mlx5/mlx5dv.pxd b/pyverbs/providers/mlx5/mlx5dv.pxd
index c08f7ad..1ddd189 100644
--- a/pyverbs/providers/mlx5/mlx5dv.pxd
+++ b/pyverbs/providers/mlx5/mlx5dv.pxd
@@ -12,8 +12,6 @@
 
 
 cdef class Mlx5Context(Context):
-    cdef object devx_umems
-    cdef object devx_objs
     cdef object devx_eqs
     cdef object cmd_comps
     cdef add_ref(self, obj)
@@ -88,6 +86,7 @@
     cdef Context context
     cdef void *addr
     cdef object is_user_addr
+    cdef bint imported
 
 cdef class Mlx5DevxObj(PyverbsCM):
     cdef dv.mlx5dv_devx_obj *obj
@@ -95,6 +94,7 @@
     cdef object out_view
     cdef object flow_counter_actions
     cdef object dest_tir_actions
+    cdef bint imported
     cdef add_ref(self, obj)
 
 cdef class Mlx5Cqe64(PyverbsObject):
diff --git a/pyverbs/providers/mlx5/mlx5dv.pyx b/pyverbs/providers/mlx5/mlx5dv.pyx
index 54d4b27..00dd977 100644
--- a/pyverbs/providers/mlx5/mlx5dv.pyx
+++ b/pyverbs/providers/mlx5/mlx5dv.pyx
@@ -194,6 +194,7 @@
         self.context.add_ref(self)
         self.flow_counter_actions = weakref.WeakSet()
         self.dest_tir_actions = weakref.WeakSet()
+        self.imported = False
 
     def query(self, in_, outlen):
         """
@@ -293,14 +294,68 @@
     cpdef close(self):
         if self.obj != NULL:
             if self.logger:
-                self.logger.debug('Closing Mlx5DvexObj')
+                self.logger.debug('Closing Mlx5DevxObj')
             close_weakrefs([self.flow_counter_actions, self.dest_tir_actions])
+            if self.imported:
+                self.unimport_obj()
+            else:
+                self.destroy()
+
+    def destroy(self):
+        """Destroy the DevX object."""
+        if self.obj != NULL:
             rc = dv.mlx5dv_devx_obj_destroy(self.obj)
             if rc:
                 raise PyverbsRDMAError('Failed to destroy a DevX object', rc)
             self.obj = NULL
             self.context = None
 
+    def unimport_obj(self):
+        """
+        Release an imported DevX object without destroying the kernel object.
+        """
+        if self.obj != NULL:
+            dv.mlx5dv_devx_obj_unimport(self.obj)
+            self.obj = NULL
+            self.context = None
+
+    def export(self):
+        """
+        Export DevX object into an opaque data buffer for cross-process sharing.
+        :return: Bytes object containing the exported DevX object data.
+        """
+        cdef uint32_t sz = Mlx5DevxObj.get_export_size()
+        cdef void *buf = malloc(sz)
+        if buf == NULL:
+            raise MemoryError('Failed to allocate export buffer')
+        rc = dv.mlx5dv_devx_obj_export(self.obj, buf)
+        if rc != 0:
+            free(buf)
+            raise PyverbsRDMAError('Failed to export DevX object', rc)
+        cdef bytes data = (<char *>buf)[:sz]
+        free(buf)
+        return data
+
+    @staticmethod
+    def import_obj(Context context not None, bytes data not None):
+        """
+        Import a DevX object from an opaque data buffer (e.g. from export()).
+        :param context: Device context to import the DevX object on
+        :param data: Opaque bytes buffer previously returned by Mlx5DevxObj.export()
+        :return: An Mlx5DevxObj object representing the imported DevX object
+        """
+        cdef Mlx5DevxObj obj = Mlx5DevxObj.__new__(Mlx5DevxObj)
+        cdef char *buf = data
+        obj.obj = dv.mlx5dv_devx_obj_import(context.context, <void *>buf)
+        if obj.obj == NULL:
+            raise PyverbsRDMAErrno('Failed to import DevX object')
+        obj.context = context
+        obj.imported = True
+        obj.flow_counter_actions = weakref.WeakSet()
+        obj.dest_tir_actions = weakref.WeakSet()
+        context.devx_objs.add(obj)
+        return obj
+
 
 cdef class Mlx5Context(Context):
     """
@@ -320,8 +375,7 @@
         if self.context == NULL:
             raise PyverbsRDMAErrno('Failed to open mlx5 context on {dev}'
                                    .format(dev=self.name))
-        self.devx_umems = weakref.WeakSet()
-        self.devx_objs = weakref.WeakSet()
+
         self.devx_eqs = weakref.WeakSet()
         self.cmd_comps = weakref.WeakSet()
 
@@ -512,8 +566,8 @@
 
     cpdef close(self):
         if self.context != NULL:
-            close_weakrefs([self.pps, self.devx_objs, self.devx_umems, self.devx_eqs,
-                           self.cmd_comps])
+            close_weakrefs([self.pps, self.devx_objs, self.devx_umems,
+                            self.devx_eqs, self.cmd_comps])
             super(Mlx5Context, self).close()
 
 
@@ -1849,6 +1903,7 @@
             raise PyverbsRDMAErrno("Failed to register a UMEM.")
         self.context = context
         self.context.add_ref(self)
+        self.imported = False
 
     def __dealloc__(self):
         self.close()
@@ -1857,6 +1912,14 @@
         if self.umem != NULL:
             if self.logger:
                 self.logger.debug('Closing Mlx5UMEM')
+            if self.imported:
+                self.unimport_umem()
+            else:
+                self.dereg()
+
+    def dereg(self):
+        """Deregister the UMEM."""
+        if self.umem != NULL:
             rc = dv.mlx5dv_devx_umem_dereg(self.umem)
             try:
                 if rc:
@@ -1864,9 +1927,54 @@
             finally:
                 if not self.is_user_addr:
                     free(self.addr)
+                    self.addr = NULL
             self.umem = NULL
             self.context = None
 
+    def unimport_umem(self):
+        """Release an imported UMEM without deregistering the kernel object."""
+        if self.umem != NULL:
+            dv.mlx5dv_devx_umem_unimport(self.umem)
+            self.umem = NULL
+            self.context = None
+
+    def export(self):
+        """
+        Export DEVX UMEM into an opaque data buffer for cross-process sharing.
+        :return: Bytes object containing the exported UMEM data.
+        """
+        cdef uint32_t sz = Mlx5UMEM.get_export_size()
+        cdef void *buf = malloc(sz)
+        if buf == NULL:
+            raise MemoryError('Failed to allocate export buffer')
+        rc = dv.mlx5dv_devx_umem_export(self.umem, buf)
+        if rc != 0:
+            free(buf)
+            raise PyverbsRDMAError('Failed to export UMEM', rc)
+        cdef bytes data = (<char *>buf)[:sz]
+        free(buf)
+        return data
+
+    @staticmethod
+    def import_umem(Context context not None, bytes data not None):
+        """
+        Import a DEVX UMEM from an opaque data buffer (e.g. from export()).
+        :param context: Device context to import the UMEM on
+        :param data: Opaque bytes buffer previously returned by Mlx5UMEM.export()
+        :return: An Mlx5UMEM object representing the imported UMEM
+        """
+        cdef Mlx5UMEM umem_obj = Mlx5UMEM.__new__(Mlx5UMEM)
+        cdef char *buf = data
+        umem_obj.umem = dv.mlx5dv_devx_umem_import(context.context, <void *>buf)
+        if umem_obj.umem == NULL:
+            raise PyverbsRDMAErrno('Failed to import UMEM')
+        umem_obj.context = context
+        umem_obj.imported = True
+        umem_obj.addr = NULL
+        umem_obj.is_user_addr = True
+        context.devx_umems.add(umem_obj)
+        return umem_obj
+
     def __str__(self):
         print_format = '{:20}: {:<20}\n'
         return print_format.format('umem id', self.umem_id) + \