pyverbs: Add ibv_buf provider-aware buffer support Wrap the provider-aware buffer API (ibv_alloc_buf, ibv_free_buf, ibv_reg_buf_mr) with new Buf and BufMR classes. Buf owns a buffer allocated through a PD and is tracked in a per-PD weakset, so it is torn down before the PD it belongs to. BufMR registers an MR over a (sub)range of a Buf and deregisters it on close without freeing the underlying buffer. Signed-off-by: Jiri Pirko <jiri@nvidia.com>
diff --git a/pyverbs/libibverbs.pxd b/pyverbs/libibverbs.pxd index 96b4b09..566089a 100644 --- a/pyverbs/libibverbs.pxd +++ b/pyverbs/libibverbs.pxd
@@ -83,6 +83,9 @@ unsigned int lkey unsigned int rkey + cdef struct ibv_buf: + pass + cdef struct ibv_query_device_ex_input: unsigned int comp_mask @@ -663,6 +666,7 @@ int fd uint64_t fd_offset ibv_dmah *dmah + ibv_buf *buf ibv_device **ibv_get_device_list(int *n) int ibv_get_device_index(ibv_device *device); @@ -685,6 +689,10 @@ ibv_mr *ibv_reg_mr(ibv_pd *pd, void *addr, size_t length, int access) ibv_mr *ibv_reg_dmabuf_mr(ibv_pd *pd, uint64_t offset, size_t length, uint64_t iova, int fd, int access) + void *ibv_alloc_buf(ibv_pd *pd, size_t size, ibv_buf **buf) + void ibv_free_buf(ibv_buf *buf) + ibv_mr *ibv_reg_buf_mr(ibv_pd *pd, ibv_buf *buf, void *addr, + size_t length, int access) int ibv_rereg_mr(ibv_mr *mr, int flags, ibv_pd *pd, void *addr, size_t length, int access) int ibv_dereg_mr(ibv_mr *mr)
diff --git a/pyverbs/libibverbs_enums.pxd b/pyverbs/libibverbs_enums.pxd index fea916b..9219a36 100644 --- a/pyverbs/libibverbs_enums.pxd +++ b/pyverbs/libibverbs_enums.pxd
@@ -522,6 +522,7 @@ IBV_REG_MR_MASK_FD IBV_REG_MR_MASK_FD_OFFSET IBV_REG_MR_MASK_DMAH + IBV_REG_MR_MASK_BUF cdef extern from "<infiniband/verbs_api.h>":
diff --git a/pyverbs/mr.pxd b/pyverbs/mr.pxd index 94ed2c2..acfcea3 100644 --- a/pyverbs/mr.pxd +++ b/pyverbs/mr.pxd
@@ -24,6 +24,18 @@ cdef class MREx(MR): cdef object dmah + cdef object backing_buf + +cdef class Buf(PyverbsCM): + cdef v.ibv_buf *bufh + cdef void *addr + cdef size_t size + cdef object pd + cdef object mrs + cdef add_ref(self, obj) + +cdef class BufMR(MR): + cdef object backing_buf cdef class MWBindInfo(PyverbsCM): cdef v.ibv_mw_bind_info info
diff --git a/pyverbs/mr.pyx b/pyverbs/mr.pyx index 706653b..727c7ab 100644 --- a/pyverbs/mr.pyx +++ b/pyverbs/mr.pyx
@@ -623,18 +623,21 @@ This class provides more flexibility in memory registration compared to the basic MR class. """ def __init__(self, PD pd not None, length=0, access=0, address=None, - iova=None, fd=None, fd_offset=0, dmah=None, implicit=False, **kwargs): + iova=None, fd=None, fd_offset=0, dmah=None, implicit=False, + Buf buf=None, **kwargs): """ Register a memory region using the extended API ibv_reg_mr_ex. :param pd: A PD object :param length: Length (in bytes) of MR's buffer :param access: Access flags, see ibv_access_flags enum - :param address: Memory address to register (Optional) + :param address: Memory address to register; defaults to the start of + buf when buf is given (Optional) :param iova: IOVA address to register (Optional) :param fd: File descriptor for dma-buf based registration (Optional) :param fd_offset: Offset in the dma-buf (Optional) :param dmah: DMA handle for registration (Optional) :param implicit: If True, register implicit MR + :param buf: A Buf object to register the MR on (Optional). :param kwargs: Additional arguments :return: The newly created MREx on success """ @@ -645,8 +648,14 @@ self.is_user_addr = False self.mmap_length = 0 + if buf is not None: + self.is_user_addr = True + if address is None: + self.buf = buf.addr + else: + self.buf = <void*><uintptr_t>address # Handle memory allocation if no address is provided - if not address and length > 0 and fd is None: + elif not address and length > 0 and fd is None: self._allocate_buffer(<size_t>length, self.is_huge, &mmap_len) if self.buf == NULL: raise PyverbsError(f'Failed to allocate MR buffer of size {length}') @@ -667,6 +676,9 @@ in_.addr = self.buf in_.access = access + if buf is not None: + in_.comp_mask |= e.IBV_REG_MR_MASK_BUF + in_.buf = buf.bufh if iova is not None: in_.comp_mask |= e.IBV_REG_MR_MASK_IOVA in_.iova = iova @@ -687,6 +699,9 @@ self.pd = pd pd.add_ref(self) + if buf is not None: + self.backing_buf = buf + buf.add_ref(self) if dmah is not None: (<DMAHandle>dmah).add_ref(self) self.dmah = dmah @@ -707,3 +722,100 @@ if self.mr != NULL: super(MREx, self).close() self.dmah = None + +cdef class Buf(PyverbsCM): + """ + Represents an ibv_buf buffer allocated through a PD. + The device provider selects the backing memory for the given PD. + """ + def __init__(self, PD pd not None, size): + """ + Allocate a buffer of the given size from the provider associated with + the given protection domain (or parent domain). + :param pd: A PD/ParentDomain object used for the allocation + :param size: Size (in bytes) of the buffer to allocate + :return: The newly created Buf on success + """ + super().__init__() + self.mrs = weakref.WeakSet() + self.addr = v.ibv_alloc_buf(pd.pd, size, &self.bufh) + if self.addr == NULL: + raise PyverbsRDMAErrno(f'Failed to allocate ibv_buf of size {size}') + self.size = size + self.pd = pd + pd.add_ref(self) + self.logger.debug(f'Allocated ibv_buf of size {size}') + + def __dealloc__(self): + self.close() + + cpdef close(self): + """ + Frees the underlying buffer using ibv_free_buf(). + :return: None + """ + if self.bufh != NULL: + if self.logger: + self.logger.debug('Closing Buf') + close_weakrefs([self.mrs]) + v.ibv_free_buf(self.bufh) + self.bufh = NULL + self.addr = NULL + self.pd = None + + cdef add_ref(self, obj): + if isinstance(obj, (BufMR, MREx)): + self.mrs.add(obj) + else: + raise PyverbsError('Unrecognized object type') + + @property + def addr(self): + return <uintptr_t>self.addr + + @property + def size(self): + return self.size + + def __str__(self): + print_format = '{:22}: {:<20}\n' + return 'Buf:\n' + \ + print_format.format('addr', <uintptr_t>self.addr) + \ + print_format.format('size', self.size) + + +cdef class BufMR(MR): + """ + BufMR represents a memory region registered for a Buf via ibv_reg_buf_mr(). + Unlike MR, the backing memory is owned by the Buf, so closing a BufMR only + deregisters the MR and never frees the buffer. The IBV_REG_MR_MASK_BUF path + of ibv_reg_mr_ex() is exercised through the MREx class instead. + """ + def __init__(self, PD pd not None, Buf buf not None, length=0, access=0, + offset=0): + """ + Register a memory region for (a subrange of) the given Buf. + :param pd: The same PD/ParentDomain used to allocate the Buf + :param buf: A Buf object allocated with ibv_alloc_buf() + :param length: Length (in bytes) to register + :param access: Access flags, see ibv_access_flags enum + :param offset: Byte offset within the Buf to start the registration + :return: The newly created BufMR on success + """ + self.logger = logging.getLogger(self.__class__.__name__) + cdef void *addr = <void*>(<uintptr_t>buf.addr + <uintptr_t>offset) + self.mr = v.ibv_reg_buf_mr(pd.pd, buf.bufh, addr, length, access) + if self.mr == NULL: + raise PyverbsRDMAErrno(f'Failed to register a buf MR. length: ' + f'{length}, access flags: {access}') + self.buf = addr + super().__init__(pd, length, access) + self.is_user_addr = True + self.is_huge = False + self.mmap_length = 0 + self.pd = pd + self.backing_buf = buf + pd.add_ref(self) + buf.add_ref(self) + self.logger.debug(f'Registered buf ibv_mr. Length: {length}, access ' + f'flags {access}')
diff --git a/pyverbs/pd.pxd b/pyverbs/pd.pxd index 04e4453..585afc5 100644 --- a/pyverbs/pd.pxd +++ b/pyverbs/pd.pxd
@@ -16,6 +16,7 @@ cdef remove_ref(self, obj) cdef object srqs cdef object mrs + cdef object bufs cdef object mws cdef object ahs cdef object qps
diff --git a/pyverbs/pd.pyx b/pyverbs/pd.pyx index 1dc2bb0..0fd05f7 100644 --- a/pyverbs/pd.pyx +++ b/pyverbs/pd.pyx
@@ -15,7 +15,7 @@ from pyverbs.wr cimport copy_sg_array from pyverbs.device cimport Context from pyverbs.cmid cimport CMID -from .mr cimport MR, MW, DMMR +from .mr cimport MR, MW, DMMR, Buf from pyverbs.srq cimport SRQ from pyverbs.addr cimport AH from pyverbs.cq cimport CQEX @@ -63,6 +63,7 @@ self.logger.debug('Created PD') self.srqs = weakref.WeakSet() self.mrs = weakref.WeakSet() + self.bufs = weakref.WeakSet() self.mws = weakref.WeakSet() self.ahs = weakref.WeakSet() self.qps = weakref.WeakSet() @@ -113,7 +114,8 @@ if self.logger: self.logger.debug('Closing PD') close_weakrefs([self.deks, self.mkeys, self.parent_domains, self.qps, - self.wqs, self.ahs, self.mws, self.mrs, self.srqs]) + self.wqs, self.ahs, self.mws, self.bufs, self.mrs, + self.srqs]) if not self._is_imported: rc = v.ibv_dealloc_pd(self.pd) if rc != 0: @@ -124,6 +126,8 @@ cdef add_ref(self, obj): if isinstance(obj, MR) or isinstance(obj, DMMR): self.mrs.add(obj) + elif isinstance(obj, Buf): + self.bufs.add(obj) elif isinstance(obj, MW): self.mws.add(obj) elif isinstance(obj, AH):