pyverbs: Add Completion Counters support

Expose the Completion Counters verbs API through pyverbs.

Add CompCntr class with set, set_err, inc, inc_err, read and read_err
methods. Add CompCntrInitAttr and CompCntrAttachAttr helper classes.
Add attach_comp_cntr method to the QP class. Expose supported
capabilities using query_comp_cntr_caps device method.

Additionally add pyverbs interface for creating EFA Completion Counters
with external memory.

Signed-off-by: Michael Margolin <mrgolin@amazon.com>
diff --git a/pyverbs/CMakeLists.txt b/pyverbs/CMakeLists.txt
index 4beb118..6399fea 100644
--- a/pyverbs/CMakeLists.txt
+++ b/pyverbs/CMakeLists.txt
@@ -22,6 +22,7 @@
   addr.pyx
   base.pyx
   cmid.pyx
+  comp_cntr.pyx
   cq.pyx
   device.pyx
   ${DMA_UTIL}
diff --git a/pyverbs/comp_cntr.pxd b/pyverbs/comp_cntr.pxd
new file mode 100644
index 0000000..8699786
--- /dev/null
+++ b/pyverbs/comp_cntr.pxd
@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: (GPL-2.0 OR Linux-OpenIB)
+# Copyright Amazon.com, Inc. or its affiliates. All rights reserved.
+
+#cython: language_level=3
+
+from pyverbs.base cimport PyverbsObject, PyverbsCM
+from pyverbs.device cimport Context
+cimport pyverbs.libibverbs as v
+
+cdef class CompCntrInitAttr(PyverbsObject):
+    cdef v.ibv_comp_cntr_init_attr attr
+
+cdef class CompCntr(PyverbsCM):
+    cdef v.ibv_comp_cntr *comp_cntr
+    cdef Context ctx
+    cpdef close(self)
diff --git a/pyverbs/comp_cntr.pyx b/pyverbs/comp_cntr.pyx
new file mode 100644
index 0000000..12e812c
--- /dev/null
+++ b/pyverbs/comp_cntr.pyx
@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: (GPL-2.0 OR Linux-OpenIB)
+# Copyright Amazon.com, Inc. or its affiliates. All rights reserved.
+
+#cython: language_level=3
+
+from pyverbs.base import PyverbsRDMAErrno, PyverbsRDMAError
+from pyverbs.device cimport Context
+cimport pyverbs.libibverbs as v
+
+
+cdef class CompCntrInitAttr(PyverbsObject):
+    """Represents ibv_comp_cntr_init_attr struct."""
+    def __init__(self, comp_mask=0, cntr_type=0, flags=0):
+        super().__init__()
+        self.attr.comp_mask = comp_mask
+        self.attr.type = cntr_type
+        self.attr.flags = flags
+
+    @property
+    def comp_mask(self):
+        return self.attr.comp_mask
+
+    @property
+    def cntr_type(self):
+        return self.attr.type
+
+    @property
+    def flags(self):
+        return self.attr.flags
+
+
+cdef class CompCntr(PyverbsCM):
+    """Completion Counter object for tracking aggregate completions."""
+    def __init__(self, Context ctx not None, CompCntrInitAttr attr not None):
+        """Create a completion counter.
+
+        :param ctx: Device context to create the counter on.
+        :param attr: Completion counter init attributes.
+        """
+        super().__init__()
+        self.comp_cntr = v.ibv_create_comp_cntr(ctx.context, &attr.attr)
+        if self.comp_cntr == NULL:
+            raise PyverbsRDMAErrno('Failed to create comp_cntr')
+        self.ctx = ctx
+        ctx.add_ref(self)
+
+    def __dealloc__(self):
+        self.close()
+
+    cpdef close(self):
+        """Destroy the completion counter."""
+        if self.comp_cntr != NULL:
+            rc = v.ibv_destroy_comp_cntr(self.comp_cntr)
+            if rc:
+                raise PyverbsRDMAError('Failed to destroy comp_cntr', rc)
+            self.comp_cntr = NULL
+
+    def set(self, value):
+        """Set the completion count to the given value."""
+        rc = v.ibv_set_comp_cntr(self.comp_cntr, value)
+        if rc:
+            raise PyverbsRDMAError('Failed to set comp_cntr', rc)
+
+    def set_err(self, value):
+        """Set the error count to the given value."""
+        rc = v.ibv_set_err_comp_cntr(self.comp_cntr, value)
+        if rc:
+            raise PyverbsRDMAError('Failed to set_err comp_cntr', rc)
+
+    def inc(self, amount):
+        """Increment the completion count by amount."""
+        rc = v.ibv_inc_comp_cntr(self.comp_cntr, amount)
+        if rc:
+            raise PyverbsRDMAError('Failed to inc comp_cntr', rc)
+
+    def inc_err(self, amount):
+        """Increment the error count by amount."""
+        rc = v.ibv_inc_err_comp_cntr(self.comp_cntr, amount)
+        if rc:
+            raise PyverbsRDMAError('Failed to inc_err comp_cntr', rc)
+
+    def read(self):
+        """Read and return the current completion count."""
+        cdef unsigned long value = 0
+        rc = v.ibv_read_comp_cntr(self.comp_cntr, &value)
+        if rc:
+            raise PyverbsRDMAError('Failed to read comp_cntr', rc)
+        return value
+
+    def read_err(self):
+        """Read and return the current error count."""
+        cdef unsigned long value = 0
+        rc = v.ibv_read_err_comp_cntr(self.comp_cntr, &value)
+        if rc:
+            raise PyverbsRDMAError('Failed to read_err comp_cntr', rc)
+        return value
diff --git a/pyverbs/device.pxd b/pyverbs/device.pxd
index 26bbe9a..02446ce 100644
--- a/pyverbs/device.pxd
+++ b/pyverbs/device.pxd
@@ -23,6 +23,7 @@
     cdef object pps
     cdef object sched_nodes
     cdef object sched_leafs
+    cdef object comp_cntrs
     cdef object dr_domains
     cdef object wqs
     cdef object rwq_ind_tbls
@@ -63,6 +64,9 @@
 cdef class DeviceAttrEx(PyverbsObject):
     cdef v.ibv_device_attr_ex dev_attr
 
+cdef class CompCntrCaps(PyverbsObject):
+    cdef v.ibv_comp_cntr_caps caps
+
 cdef class AllocDmAttr(PyverbsObject):
     cdef v.ibv_alloc_dm_attr alloc_dm_attr
 
diff --git a/pyverbs/device.pyx b/pyverbs/device.pyx
index 25466f1..ea1f71d 100644
--- a/pyverbs/device.pyx
+++ b/pyverbs/device.pyx
@@ -9,6 +9,7 @@
 import weakref
 
 from .pyverbs_error import PyverbsRDMAError, PyverbsError
+from pyverbs.comp_cntr cimport CompCntr
 from pyverbs.cq cimport CQEX, CQ, CompChannel
 from .pyverbs_error import PyverbsUserError
 from pyverbs.base import PyverbsRDMAErrno
@@ -119,6 +120,7 @@
         self.pps = weakref.WeakSet()
         self.sched_nodes = weakref.WeakSet()
         self.sched_leafs = weakref.WeakSet()
+        self.comp_cntrs = weakref.WeakSet()
         self.dr_domains = weakref.WeakSet()
         self.wqs = weakref.WeakSet()
         self.rwq_ind_tbls = weakref.WeakSet()
@@ -185,7 +187,8 @@
                             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])
+                            self.dr_domains, self.event_channels, self.dmahs,
+                            self.comp_cntrs])
             rc = v.ibv_close_device(self.context)
             if rc != 0:
                 raise PyverbsRDMAErrno(f'Failed to close device {self.name}')
@@ -228,6 +231,17 @@
                                    format(name=self.name), rc)
         return dev_attr_ex
 
+    def query_comp_cntr_caps(self):
+        """
+        Query completion counter capabilities.
+        :return: CompCntrCaps object
+        """
+        caps = CompCntrCaps()
+        rc = v.ibv_query_comp_cntr_caps(self.context, &caps.caps)
+        if rc != 0:
+            raise PyverbsRDMAError('Failed to query comp_cntr caps', rc)
+        return caps
+
     def query_pkey(self, unsigned int port_num, int index):
         cdef uint16_t pkey
         rc = v.ibv_query_pkey(self.context, port_num, index, &pkey)
@@ -366,6 +380,8 @@
             self.rwq_ind_tbls.add(obj)
         elif isinstance(obj, DMAHandle):
             self.dmahs.add(obj)
+        elif isinstance(obj, CompCntr):
+            self.comp_cntrs.add(obj)
         else:
             raise PyverbsError('Unrecognized object type')
 
@@ -749,6 +765,21 @@
         return self.dev_attr.phys_port_cnt_ex
 
 
+cdef class CompCntrCaps(PyverbsObject):
+    """Completion counter capabilities."""
+    @property
+    def max_counters(self):
+        return self.caps.max_counters
+
+    @property
+    def max_value(self):
+        return self.caps.max_value
+
+    @property
+    def supported_qp_attach_ops(self):
+        return self.caps.supported_qp_attach_ops
+
+
 cdef class AllocDmAttr(PyverbsObject):
     def __init__(self, length, log_align_req = 0, comp_mask = 0):
         """
diff --git a/pyverbs/libibverbs.pxd b/pyverbs/libibverbs.pxd
index a8e942a..82bfa86 100644
--- a/pyverbs/libibverbs.pxd
+++ b/pyverbs/libibverbs.pxd
@@ -247,6 +247,22 @@
     cdef struct ibv_poll_cq_attr:
         unsigned int    comp_mask
 
+    cdef struct ibv_comp_cntr_caps:
+        unsigned long   max_value
+        unsigned int    max_counters
+        unsigned int    supported_qp_attach_ops
+
+    cdef struct ibv_comp_cntr:
+        ibv_context     *context
+        unsigned int    handle
+        unsigned long   comp_count_max_value
+        unsigned long   err_count_max_value
+
+    cdef struct ibv_comp_cntr_init_attr:
+        unsigned int    comp_mask
+        unsigned int    type
+        unsigned int    flags
+
     cdef struct ibv_wc_tm_info:
         unsigned long   tag
         unsigned int    priv
@@ -512,6 +528,10 @@
         ibv_qp_type     qp_type;
         unsigned int    events_completed;
 
+    cdef struct ibv_qp_attach_comp_cntr_attr:
+        unsigned int    comp_mask
+        unsigned int    op_mask
+
     cdef struct ibv_parent_domain_init_attr:
         ibv_pd          *pd;
         uint32_t        comp_mask;
@@ -756,6 +776,17 @@
     unsigned int ibv_wc_read_flow_tag(ibv_cq_ex *cq)
     void ibv_wc_read_tm_info(ibv_cq_ex *cq, ibv_wc_tm_info *tm_info)
     unsigned long ibv_wc_read_completion_wallclock_ns(ibv_cq_ex *cq)
+    int ibv_query_comp_cntr_caps(ibv_context *context,
+                                 ibv_comp_cntr_caps *caps)
+    ibv_comp_cntr *ibv_create_comp_cntr(ibv_context *context,
+                                        ibv_comp_cntr_init_attr *attr)
+    int ibv_destroy_comp_cntr(ibv_comp_cntr *comp_cntr)
+    int ibv_set_comp_cntr(ibv_comp_cntr *comp_cntr, unsigned long value)
+    int ibv_set_err_comp_cntr(ibv_comp_cntr *comp_cntr, unsigned long value)
+    int ibv_inc_comp_cntr(ibv_comp_cntr *comp_cntr, unsigned long amount)
+    int ibv_inc_err_comp_cntr(ibv_comp_cntr *comp_cntr, unsigned long amount)
+    int ibv_read_comp_cntr(ibv_comp_cntr *comp_cntr, unsigned long *value)
+    int ibv_read_err_comp_cntr(ibv_comp_cntr *comp_cntr, unsigned long *value)
     ibv_ah *ibv_create_ah(ibv_pd *pd, ibv_ah_attr *attr)
     int ibv_init_ah_from_wc(ibv_context *context, uint8_t port_num,
                             ibv_wc *wc, ibv_grh *grh, ibv_ah_attr *ah_attr)
@@ -770,6 +801,8 @@
     int ibv_query_qp(ibv_qp *qp, ibv_qp_attr *attr, int attr_mask,
                      ibv_qp_init_attr *init_attr)
     int ibv_destroy_qp(ibv_qp *qp)
+    int ibv_qp_attach_comp_cntr(ibv_qp *qp, ibv_comp_cntr *comp_cntr,
+                                ibv_qp_attach_comp_cntr_attr *attr)
     int ibv_post_recv(ibv_qp *qp, ibv_recv_wr *wr, ibv_recv_wr **bad_wr)
     int ibv_post_send(ibv_qp *qp, ibv_send_wr *wr, ibv_send_wr **bad_wr)
     int ibv_bind_mw(ibv_qp *qp, ibv_mw *mw, ibv_mw_bind *mw_bind)
diff --git a/pyverbs/libibverbs_enums.pxd b/pyverbs/libibverbs_enums.pxd
index 9219a36..c37f9fe 100644
--- a/pyverbs/libibverbs_enums.pxd
+++ b/pyverbs/libibverbs_enums.pxd
@@ -524,6 +524,14 @@
         IBV_REG_MR_MASK_DMAH
         IBV_REG_MR_MASK_BUF
 
+    cpdef enum ibv_qp_attach_comp_cntr_op:
+        IBV_QP_ATTACH_COMP_CNTR_OP_SEND
+        IBV_QP_ATTACH_COMP_CNTR_OP_RECV
+        IBV_QP_ATTACH_COMP_CNTR_OP_RDMA_READ
+        IBV_QP_ATTACH_COMP_CNTR_OP_REMOTE_RDMA_READ
+        IBV_QP_ATTACH_COMP_CNTR_OP_RDMA_WRITE
+        IBV_QP_ATTACH_COMP_CNTR_OP_REMOTE_RDMA_WRITE
+
 
 cdef extern from "<infiniband/verbs_api.h>":
     cdef unsigned long long IBV_ADVISE_MR_ADVICE_PREFETCH
diff --git a/pyverbs/providers/efa/efa_enums.pxd b/pyverbs/providers/efa/efa_enums.pxd
index 179f30c..725e486 100644
--- a/pyverbs/providers/efa/efa_enums.pxd
+++ b/pyverbs/providers/efa/efa_enums.pxd
@@ -11,6 +11,8 @@
         EFADV_DEVICE_ATTR_CAPS_CQ_WITH_SGID
         EFADV_DEVICE_ATTR_CAPS_RDMA_WRITE
         EFADV_DEVICE_ATTR_CAPS_UNSOLICITED_WRITE_RECV
+        EFADV_DEVICE_ATTR_CAPS_CQ_WITH_EXT_MEM_DMABUF
+        EFADV_DEVICE_ATTR_CAPS_COMP_CNTR
 
     cpdef enum:
         EFADV_QP_DRIVER_TYPE_SRD
@@ -30,6 +32,14 @@
         EFADV_WC_EX_WITH_IS_UNSOLICITED
 
     cpdef enum:
+        EFADV_MEMORY_LOCATION_VA
+        EFADV_MEMORY_LOCATION_DMABUF
+
+    cpdef enum:
+        EFADV_COMP_CNTR_INIT_WITH_COMP_EXTERNAL_MEM
+        EFADV_COMP_CNTR_INIT_WITH_ERR_EXTERNAL_MEM
+
+    cpdef enum:
         EFADV_MR_ATTR_VALIDITY_RECV_IC_ID
         EFADV_MR_ATTR_VALIDITY_RDMA_READ_IC_ID
         EFADV_MR_ATTR_VALIDITY_RDMA_RECV_IC_ID
diff --git a/pyverbs/providers/efa/efadv.pyx b/pyverbs/providers/efa/efadv.pyx
index ab56f38..791a3cb 100644
--- a/pyverbs/providers/efa/efadv.pyx
+++ b/pyverbs/providers/efa/efadv.pyx
@@ -6,12 +6,17 @@
 
 from pyverbs.addr cimport GID
 from pyverbs.base import PyverbsRDMAErrno, PyverbsRDMAError
+from pyverbs.base cimport PyverbsCM
+from pyverbs.comp_cntr cimport CompCntrInitAttr, CompCntr
 from pyverbs.cq cimport CQEX, CqInitAttrEx
+from pyverbs.device cimport Context
 from pyverbs.libibverbs_enums import ibv_qp_attr_mask
 cimport pyverbs.libibverbs as v
+from pyverbs.mr cimport MR
 from pyverbs.pd cimport PD
 from pyverbs.qp cimport QP, QPEx, QPInitAttr, QPInitAttrEx
-from pyverbs.mr cimport MR
+from libc.string cimport memset
+from libc.stdint cimport uintptr_t, uint8_t
 
 
 def dev_cap_to_str(flags):
@@ -21,6 +26,8 @@
             dve.EFADV_DEVICE_ATTR_CAPS_CQ_WITH_SGID: 'CQ entries with source GID',
             dve.EFADV_DEVICE_ATTR_CAPS_RDMA_WRITE: 'RDMA Write',
             dve.EFADV_DEVICE_ATTR_CAPS_UNSOLICITED_WRITE_RECV: 'Unsolicited RDMA Write receive',
+            dve.EFADV_DEVICE_ATTR_CAPS_CQ_WITH_EXT_MEM_DMABUF: 'CQ with external memory',
+            dve.EFADV_DEVICE_ATTR_CAPS_COMP_CNTR: 'Completion Counters',
     }
     return bitmask_to_str(flags, l)
 
@@ -408,3 +415,56 @@
     @max_recv_sge.setter
     def max_recv_sge(self, val):
         self.rq_depth_attr.max_recv_sge = val
+
+
+cdef class EfaCompCntrInitAttr(PyverbsObject):
+    """Represents efadv_comp_cntr_init_attr struct."""
+    cdef dv.efadv_comp_cntr_init_attr attr
+
+    def __init__(self, comp_ext_mem_ptr=None, err_ext_mem_ptr=None):
+        super().__init__()
+        memset(&self.attr, 0, sizeof(self.attr))
+        if comp_ext_mem_ptr is not None:
+            self.attr.flags |= dve.EFADV_COMP_CNTR_INIT_WITH_COMP_EXTERNAL_MEM
+            self.attr.comp_cntr_ext_mem.type = dve.EFADV_MEMORY_LOCATION_VA
+            self.attr.comp_cntr_ext_mem.ptr = <uint8_t *><uintptr_t>comp_ext_mem_ptr
+        if err_ext_mem_ptr is not None:
+            self.attr.flags |= dve.EFADV_COMP_CNTR_INIT_WITH_ERR_EXTERNAL_MEM
+            self.attr.err_cntr_ext_mem.type = dve.EFADV_MEMORY_LOCATION_VA
+            self.attr.err_cntr_ext_mem.ptr = <uint8_t *><uintptr_t>err_ext_mem_ptr
+
+    def set_comp_ext_mem_va(self, uintptr_t ptr):
+        """Set completion counter external memory VA."""
+        self.attr.flags |= dve.EFADV_COMP_CNTR_INIT_WITH_COMP_EXTERNAL_MEM
+        self.attr.comp_cntr_ext_mem.type = dve.EFADV_MEMORY_LOCATION_VA
+        self.attr.comp_cntr_ext_mem.ptr = <uint8_t *>ptr
+
+    def set_err_ext_mem_va(self, uintptr_t ptr):
+        """Set error counter external memory VA."""
+        self.attr.flags |= dve.EFADV_COMP_CNTR_INIT_WITH_ERR_EXTERNAL_MEM
+        self.attr.err_cntr_ext_mem.type = dve.EFADV_MEMORY_LOCATION_VA
+        self.attr.err_cntr_ext_mem.ptr = <uint8_t *>ptr
+
+
+cdef class EfaCompCntr(CompCntr):
+    """EFA-specific Completion Counter with external memory support."""
+    def __init__(self, Context ctx not None, CompCntrInitAttr attr not None,
+                 EfaCompCntrInitAttr efa_attr=None):
+        """Create an EFA completion counter.
+
+        :param ctx: Device context.
+        :param attr: Completion counter init attributes.
+        :param efa_attr: EFA-specific init attributes (optional).
+        """
+        PyverbsCM.__init__(self)
+        if efa_attr is None:
+            efa_attr = EfaCompCntrInitAttr()
+        self.comp_cntr = dv.efadv_create_comp_cntr(
+            ctx.context, &attr.attr, &efa_attr.attr, sizeof(efa_attr.attr))
+        if self.comp_cntr == NULL:
+            raise PyverbsRDMAErrno('Failed to create EFA comp_cntr')
+        self.ctx = ctx
+        ctx.add_ref(self)
+
+    def __dealloc__(self):
+        self.close()
diff --git a/pyverbs/providers/efa/libefa.pxd b/pyverbs/providers/efa/libefa.pxd
index 5cd4d1b..878cc5b 100644
--- a/pyverbs/providers/efa/libefa.pxd
+++ b/pyverbs/providers/efa/libefa.pxd
@@ -86,3 +86,20 @@
                                uint32_t inlen)
     int efadv_get_max_rq_depth(v.ibv_context *ibvctx, efadv_rq_depth_attr *attr,
                                uint32_t inlen);
+
+    cdef struct efadv_memory_location:
+        uint8_t *ptr
+        uint8_t type
+        uint8_t reserved[7]
+
+    cdef struct efadv_comp_cntr_init_attr:
+        uint64_t comp_mask
+        uint32_t flags
+        uint32_t reserved
+        efadv_memory_location comp_cntr_ext_mem
+        efadv_memory_location err_cntr_ext_mem
+
+    v.ibv_comp_cntr *efadv_create_comp_cntr(v.ibv_context *ibvctx,
+                                            v.ibv_comp_cntr_init_attr *attr,
+                                            efadv_comp_cntr_init_attr *efa_attr,
+                                            uint32_t inlen)
diff --git a/pyverbs/qp.pxd b/pyverbs/qp.pxd
index 25afd32..050d59d 100644
--- a/pyverbs/qp.pxd
+++ b/pyverbs/qp.pxd
@@ -30,6 +30,9 @@
 cdef class QPRateLimitAttr(PyverbsObject):
     cdef v.ibv_qp_rate_limit_attr attr
 
+cdef class QPAttachCompCntrAttr(PyverbsObject):
+    cdef v.ibv_qp_attach_comp_cntr_attr attr
+
 cdef class QP(PyverbsCM):
     cdef v.ibv_qp *qp
     cdef int type
diff --git a/pyverbs/qp.pyx b/pyverbs/qp.pyx
index b4f8403..a14fd93 100644
--- a/pyverbs/qp.pyx
+++ b/pyverbs/qp.pyx
@@ -21,6 +21,7 @@
 from pyverbs.device cimport Context
 from cpython.ref cimport PyObject
 from pyverbs.cq cimport CQ, CQEX
+from pyverbs.comp_cntr cimport CompCntr
 cimport pyverbs.libibverbs as v
 from pyverbs.xrcd cimport XRCD
 from pyverbs.srq cimport SRQ
@@ -978,6 +979,17 @@
                print_format.format('Comp mask', self.attr.comp_mask)
 
 
+cdef class QPAttachCompCntrAttr(PyverbsObject):
+    """Attributes for attaching a completion counter to a QP."""
+    def __init__(self, op_mask=0):
+        super().__init__()
+        self.attr.op_mask = op_mask
+
+    @property
+    def op_mask(self):
+        return self.attr.op_mask
+
+
 cdef class QP(PyverbsCM):
     def __init__(self, object creator not None, object init_attr not None,
                  QPAttr qp_attr=None):
@@ -1235,6 +1247,18 @@
         if rc != 0:
             raise PyverbsRDMAError('Failed to modify QP rate limit', rc)
 
+    def attach_comp_cntr(self, CompCntr comp_cntr not None,
+                         QPAttachCompCntrAttr attr not None):
+        """
+        Attach a completion counter to this QP.
+        :param comp_cntr: The completion counter to attach
+        :param attr: Attach attributes including op_mask
+        """
+        rc = v.ibv_qp_attach_comp_cntr(self.qp, comp_cntr.comp_cntr,
+                                        &attr.attr)
+        if rc != 0:
+            raise PyverbsRDMAError('Failed to attach comp cntr to QP', rc)
+
     def post_recv(self, RecvWR wr not None, RecvWR bad_wr=None):
         """
         Post a receive WR on the QP.