UCT/IB/MLX5: Add AArch64 ST64B BF copy mode - v1.23.x

src/ucs/arch/aarch64/cpu.c

@@ -13,6 +13,18 @@
 #include <ucs/arch/cpu.h>
 #include <stdio.h>
 
+#if HAVE_DECL_GETAUXVAL
+#  include <sys/auxv.h>
+
+/* Older userspace headers may not expose Arm64 HWCAP3 definitions. */
+#ifndef AT_HWCAP3
+#  define AT_HWCAP3 29
+#endif
+
+#ifndef HWCAP3_LS64
+#  define HWCAP3_LS64 (1UL << 3)
+#endif
+#endif
 
 static void ucs_aarch64_cpuid_from_proc(ucs_aarch64_cpuid_t *cpuid)
 {
@@ -68,4 +80,23 @@ void ucs_aarch64_cpuid(ucs_aarch64_cpuid_t *cpuid)
     *cpuid = cached_cpuid;
 }
 
+ucs_cpu_flag_t ucs_arch_get_cpu_flag()
+{
+    static ucs_cpu_flag_t cpu_flag = UCS_CPU_FLAG_UNKNOWN;
+
+    if (cpu_flag == UCS_CPU_FLAG_UNKNOWN) {
+        ucs_cpu_flag_t result = 0;
+
+#if HAVE_DECL_GETAUXVAL
+        if (getauxval(AT_HWCAP3) & HWCAP3_LS64) {
+            result |= UCS_CPU_FLAG_LS64;
+        }
+#endif
+
+        cpu_flag = result;
+    }
+
+    return cpu_flag;
+}
+
 #endif

src/ucs/arch/aarch64/cpu.h

@@ -169,10 +169,7 @@ static inline ucs_cpu_model_t ucs_arch_get_cpu_model()
     return UCS_CPU_MODEL_ARM_AARCH64;
 }
 
-static inline int ucs_arch_get_cpu_flag()
-{
-    return UCS_CPU_FLAG_UNKNOWN;
-}
+ucs_cpu_flag_t ucs_arch_get_cpu_flag();
 
 static inline void ucs_cpu_init()
 {

src/ucs/arch/cpu.h

@@ -65,7 +65,8 @@ typedef enum ucs_cpu_flag {
     UCS_CPU_FLAG_SSE41      = UCS_BIT(7),
     UCS_CPU_FLAG_SSE42      = UCS_BIT(8),
     UCS_CPU_FLAG_AVX        = UCS_BIT(9),
-    UCS_CPU_FLAG_AVX2       = UCS_BIT(10)
+    UCS_CPU_FLAG_AVX2       = UCS_BIT(10),
+    UCS_CPU_FLAG_LS64       = UCS_BIT(11)
 } ucs_cpu_flag_t;
 
 
@@ -161,6 +162,25 @@ static inline void ucs_clear_cache(void *start, void *end)
 }
 
 
+/**
+ * Check whether the CPU supports a given feature flag.
+ *
+ * Architectures which cannot detect CPU flags report @ref
+ * UCS_CPU_FLAG_UNKNOWN, which has all bits set, so the result of @ref
+ * ucs_arch_get_cpu_flag must not be tested by a plain bitwise AND.
+ *
+ * @param flag  Feature flag to check.
+ *
+ * @return Nonzero if the CPU is known to support the flag.
+ */
+static inline int ucs_cpu_has_flag(ucs_cpu_flag_t flag)
+{
+    int flags = ucs_arch_get_cpu_flag();
+
+    return (flags != UCS_CPU_FLAG_UNKNOWN) && (flags & flag);
+}
+
+
 static inline int ucs_cpu_model_prefer_relaxed_order(
         ucs_cpu_vendor_t cpu_vendor, ucs_cpu_model_t cpu_model)
 {

src/uct/ib/mlx5/configure.m4

@@ -8,6 +8,32 @@
 # Add IB mlx5 provider support
 #
 
+AC_CACHE_CHECK([for AArch64 ST64B assembler support],
+               [ucx_cv_have_aarch64_st64b_asm],
+               [AC_COMPILE_IFELSE([AC_LANG_SOURCE([[
+#ifndef __aarch64__
+#error "ST64B is only supported on AArch64"
+#endif
+
+void test_st64b(void *dst, void *src)
+{
+    register void *src_reg asm("x8") = src;
+    register void *dst_reg asm("x9") = dst;
+
+    asm volatile(".arch_extension ls64\n"
+                 "st64b x8, [x9]"
+                 :
+                 : "r"(src_reg), "r"(dst_reg)
+                 : "memory");
+}
+]])],
+               [ucx_cv_have_aarch64_st64b_asm=yes],
+               [ucx_cv_have_aarch64_st64b_asm=no])])
+
+AS_IF([test "x$ucx_cv_have_aarch64_st64b_asm" = "xyes"],
+      [AC_DEFINE([HAVE_AARCH64_ST64B_ASM], [1],
+                 [Define to 1 if AArch64 ST64B assembler is supported])])
+
 AC_CONFIG_FILES([src/uct/ib/mlx5/Makefile src/uct/ib/mlx5/ucx-ib-mlx5.pc])
 
 m4_include([src/uct/ib/mlx5/gdaki/configure.m4])

src/uct/ib/mlx5/dc/dc_mlx5.c

@@ -523,7 +523,8 @@ ucs_status_t uct_dc_mlx5_iface_create_dci(uct_dc_mlx5_iface_t *iface,
 
     if (dci->txwq.super.type == UCT_IB_MLX5_OBJ_TYPE_VERBS) {
         status = uct_ib_mlx5_txwq_init(iface->super.super.super.super.worker,
-                                       iface->super.tx.mmio_mode, &dci->txwq,
+                                       iface->super.tx.mmio_mode,
+                                       iface->super.tx.bf_copy_mode, &dci->txwq,
                                        dci->txwq.super.verbs.qp);
         if (status != UCS_OK) {
             goto err;

src/uct/ib/mlx5/dv/ib_mlx5_dv.c

@@ -155,6 +155,13 @@ ucs_status_t uct_ib_mlx5_devx_create_qp_common(uct_ib_iface_t *iface,
         goto err;
     }
 
+    if (tx != NULL) {
+        status = uct_ib_mlx5_txwq_init_bf_copy(tx, attr->bf_copy_mode);
+        if (status != UCS_OK) {
+            goto err;
+        }
+    }
+
     uar = uct_worker_tl_data_get(iface->super.worker,
                                  UCT_IB_MLX5_DEVX_UAR_KEY,
                                  uct_ib_mlx5_devx_uar_t,

src/uct/ib/mlx5/gdaki/gdaki.c

@@ -487,6 +487,7 @@ uct_rc_gdaki_init_channel_chunk(uct_rc_gdaki_iface_t *iface,
 
     cq_attr.flags                           |= UCT_IB_MLX5_CQ_IGNORE_OVERRUN;
     qp_attr.mmio_mode                        = UCT_IB_MLX5_MMIO_MODE_DB;
+    qp_attr.bf_copy_mode                     = UCT_IB_MLX5_BF_COPY_MODE_GENERIC;
     qp_attr.super.srq_num                    = 0;
     qp_attr.super.max_inl_cqe[UCT_IB_DIR_TX] = 0;
     uct_ib_mlx5_wq_calc_sizes(&qp_attr);

src/uct/ib/mlx5/ib_mlx5.c

@@ -33,6 +33,13 @@ static const char *uct_ib_mlx5_mmio_modes[] = {
     [UCT_IB_MLX5_MMIO_MODE_LAST]       = NULL
 };
 
+static const char *uct_ib_mlx5_bf_copy_modes[] = {
+    [UCT_IB_MLX5_BF_COPY_MODE_AUTO]    = "auto",
+    [UCT_IB_MLX5_BF_COPY_MODE_GENERIC] = "generic",
+    [UCT_IB_MLX5_BF_COPY_MODE_ST64B]   = "st64b",
+    [UCT_IB_MLX5_BF_COPY_MODE_LAST]    = NULL
+};
+
 ucs_config_field_t uct_ib_mlx5_iface_config_table[] = {
 #if HAVE_IBV_DM
     {"DM_SIZE", "2k",
@@ -55,6 +62,15 @@ ucs_config_field_t uct_ib_mlx5_iface_config_table[] = {
      ucs_offsetof(uct_ib_mlx5_iface_config_t, mmio_mode),
      UCS_CONFIG_TYPE_ENUM(uct_ib_mlx5_mmio_modes)},
 
+    {"BF_COPY_MODE", "auto",
+     "How to copy WQE building blocks to BlueFlame MMIO register. One of "
+     "the following:\n"
+     " auto    - Select best according to runtime CPU capabilities.\n"
+     " generic - Use portable scalar stores.\n"
+     " st64b   - Use AArch64 ST64B store, fail if LS64 is not supported.",
+     ucs_offsetof(uct_ib_mlx5_iface_config_t, bf_copy_mode),
+     UCS_CONFIG_TYPE_ENUM(uct_ib_mlx5_bf_copy_modes)},
+
     {"AR_ENABLE", "auto",
      "Enable Adaptive Routing (out of order) feature on SL that supports it.\n"
      "On RoCE devices, this is done by modifying the QP directly.\n"
@@ -765,8 +781,46 @@ uct_ib_mlx5_get_mmio_mode(uct_priv_worker_t *worker,
     return UCS_OK;
 }
 
+ucs_status_t
+uct_ib_mlx5_txwq_init_bf_copy(uct_ib_mlx5_txwq_t *txwq,
+                              uct_ib_mlx5_bf_copy_mode_t bf_copy_mode)
+{
+    ucs_assert(bf_copy_mode < UCT_IB_MLX5_BF_COPY_MODE_LAST);
+
+#if UCT_IB_MLX5_HAVE_ST64B
+    txwq->bf_copy_mode = UCT_IB_MLX5_BF_COPY_MODE_GENERIC;
+    if (bf_copy_mode == UCT_IB_MLX5_BF_COPY_MODE_GENERIC) {
+        return UCS_OK;
+    }
+
+    if (ucs_cpu_has_flag(UCS_CPU_FLAG_LS64)) {
+        txwq->bf_copy_mode = UCT_IB_MLX5_BF_COPY_MODE_ST64B;
+        return UCS_OK;
+    }
+
+    if (bf_copy_mode == UCT_IB_MLX5_BF_COPY_MODE_ST64B) {
+        ucs_error("mlx5 BlueFlame ST64B copy was requested but CPU does "
+                  "not report LS64 support");
+        return UCS_ERR_UNSUPPORTED;
+    }
+#else
+    if (bf_copy_mode == UCT_IB_MLX5_BF_COPY_MODE_ST64B) {
+#if defined(__aarch64__)
+        ucs_error("mlx5 BlueFlame ST64B copy was requested but UCX was built "
+                  "without assembler support for ST64B");
+#else
+        ucs_error("mlx5 BlueFlame ST64B copy is supported only on AArch64");
+#endif
+        return UCS_ERR_UNSUPPORTED;
+    }
+#endif
+
+    return UCS_OK;
+}
+
 ucs_status_t uct_ib_mlx5_txwq_init(uct_priv_worker_t *worker,
                                    uct_ib_mlx5_mmio_mode_t cfg_mmio_mode,
+                                   uct_ib_mlx5_bf_copy_mode_t bf_copy_mode,
                                    uct_ib_mlx5_txwq_t *txwq,
                                    struct ibv_qp *verbs_qp)
 {
@@ -800,6 +854,11 @@ ucs_status_t uct_ib_mlx5_txwq_init(uct_priv_worker_t *worker,
         return status;
     }
 
+    status = uct_ib_mlx5_txwq_init_bf_copy(txwq, bf_copy_mode);
+    if (status != UCS_OK) {
+        return status;
+    }
+
     ucs_debug("tx wq %d bytes [bb=%d, nwqe=%d] mmio_mode %s",
               qp_info.dv.sq.stride * qp_info.dv.sq.wqe_cnt,
               qp_info.dv.sq.stride, qp_info.dv.sq.wqe_cnt,

src/uct/ib/mlx5/ib_mlx5.h

@@ -90,6 +90,13 @@
 #  define UCT_IB_MLX5_UAR_ALLOC_TYPE_NC (1U << 31)
 #endif
 
+/* BlueFlame copy by AArch64 ST64B is compiled in */
+#if defined(__aarch64__) && HAVE_AARCH64_ST64B_ASM
+#  define UCT_IB_MLX5_HAVE_ST64B 1
+#else
+#  define UCT_IB_MLX5_HAVE_ST64B 0
+#endif
+
 #define UCT_IB_MLX5_OPMOD_EXT_ATOMIC(_log_arg_size) \
     ((8) | ((_log_arg_size) - 2))
 
@@ -468,16 +475,25 @@ typedef enum {
 } uct_ib_mlx5_mmio_mode_t;
 
 
+typedef enum {
+    UCT_IB_MLX5_BF_COPY_MODE_AUTO,
+    UCT_IB_MLX5_BF_COPY_MODE_GENERIC,
+    UCT_IB_MLX5_BF_COPY_MODE_ST64B,
+    UCT_IB_MLX5_BF_COPY_MODE_LAST
+} uct_ib_mlx5_bf_copy_mode_t;
+
+
 typedef struct uct_ib_mlx5_iface_config {
 #if HAVE_IBV_DM
     struct {
         size_t               seg_len;
         unsigned             count;
     } dm;
 #endif
-    uct_ib_mlx5_mmio_mode_t  mmio_mode;
-    ucs_ternary_auto_value_t ar_enable;
-    int                      cqe_zip_enable[UCT_IB_DIR_LAST];
+    uct_ib_mlx5_mmio_mode_t     mmio_mode;
+    uct_ib_mlx5_bf_copy_mode_t  bf_copy_mode;
+    ucs_ternary_auto_value_t    ar_enable;
+    int                         cqe_zip_enable[UCT_IB_DIR_LAST];
 } uct_ib_mlx5_iface_config_t;
 
 
@@ -643,6 +659,7 @@ typedef struct uct_ib_mlx5_res_domain {
 typedef struct uct_ib_mlx5_qp_attr {
     uct_ib_qp_attr_t            super;
     uct_ib_mlx5_mmio_mode_t     mmio_mode;
+    uct_ib_mlx5_bf_copy_mode_t  bf_copy_mode;
     uint32_t                    uidx;
     int                         full_handshake;
     int                         rdma_wr_disabled;
@@ -693,6 +710,9 @@ typedef struct uct_ib_mlx5_txwq {
     void                        *qend;
     uint16_t                    bb_max;
     uint16_t                    sig_pi;     /* PI for last signaled WQE */
+#if UCT_IB_MLX5_HAVE_ST64B
+    uint8_t                     bf_copy_mode;
+#endif
 #if UCS_ENABLE_ASSERT
     uint16_t                    hw_ci; /* First BB index of last completed WQE */
     uint8_t                     flags; /* Debug flags */
@@ -900,7 +920,13 @@ uct_ib_mlx5_get_mmio_mode(uct_priv_worker_t *worker,
  */
 ucs_status_t uct_ib_mlx5_txwq_init(uct_priv_worker_t *worker,
                                    uct_ib_mlx5_mmio_mode_t cfg_mmio_mode,
-                                   uct_ib_mlx5_txwq_t *txwq, struct ibv_qp *verbs_qp);
+                                   uct_ib_mlx5_bf_copy_mode_t bf_copy_mode,
+                                   uct_ib_mlx5_txwq_t *txwq,
+                                   struct ibv_qp *verbs_qp);
+
+ucs_status_t
+uct_ib_mlx5_txwq_init_bf_copy(uct_ib_mlx5_txwq_t *txwq,
+                              uct_ib_mlx5_bf_copy_mode_t bf_copy_mode);
 
 /* Get pointer to a WQE by producer index */
 void *uct_ib_mlx5_txwq_get_wqe(const uct_ib_mlx5_txwq_t *txwq, uint16_t pi);

src/uct/ib/mlx5/ib_mlx5.inl

@@ -556,21 +556,55 @@ static UCS_F_ALWAYS_INLINE void uct_ib_mlx5_bf_copy_bb(void * restrict dst,
 #endif
 }
 
+#if UCT_IB_MLX5_HAVE_ST64B
+static UCS_F_ALWAYS_INLINE void
+uct_ib_mlx5_bf_copy_bb_st64b(void *restrict dst, void *restrict src)
+{
+    ucs_assert(((uintptr_t)src % MLX5_SEND_WQE_BB) == 0);
+    ucs_assert(((uintptr_t)dst % MLX5_SEND_WQE_BB) == 0);
+
+    asm volatile(".arch_extension ls64\n"
+                 "ldp x8, x9, [%[src], #0]\n"
+                 "ldp x10, x11, [%[src], #16]\n"
+                 "ldp x12, x13, [%[src], #32]\n"
+                 "ldp x14, x15, [%[src], #48]\n"
+                 "st64b x8, [%[dst]]"
+                 :
+                 : [src] "r"(src), [dst] "r"(dst)
+                 : "x8", "x9", "x10", "x11", "x12", "x13", "x14",
+                   "x15", "memory");
+}
+#endif
+
+#define UCT_IB_MLX5_BF_COPY(_dst, _src, _num_bb, _wq, _copy_bb) \
+    ({ \
+        uint16_t _bb_num = (_num_bb); \
+        uint16_t _n; \
+        \
+        for (_n = 0; _n < _bb_num; ++_n) { \
+            _copy_bb(_dst, _src); \
+            (_dst) = UCS_PTR_BYTE_OFFSET(_dst, MLX5_SEND_WQE_BB); \
+            (_src) = UCS_PTR_BYTE_OFFSET(_src, MLX5_SEND_WQE_BB); \
+            if (ucs_unlikely((_src) == (_wq)->qend)) { \
+                (_src) = (_wq)->qstart; \
+            } \
+        } \
+        \
+        (_src); \
+    })
+
 static UCS_F_ALWAYS_INLINE
 void *uct_ib_mlx5_bf_copy(void *dst, void *src, uint16_t num_bb,
                           const uct_ib_mlx5_txwq_t *wq)
 {
-    uint16_t n;
-
-    for (n = 0; n < num_bb; ++n) {
-        uct_ib_mlx5_bf_copy_bb(dst, src);
-        dst = UCS_PTR_BYTE_OFFSET(dst, MLX5_SEND_WQE_BB);
-        src = UCS_PTR_BYTE_OFFSET(src, MLX5_SEND_WQE_BB);
-        if (ucs_unlikely(src == wq->qend)) {
-            src = wq->qstart;
-        }
+#if UCT_IB_MLX5_HAVE_ST64B
+    if (wq->bf_copy_mode == UCT_IB_MLX5_BF_COPY_MODE_ST64B) {
+        return UCT_IB_MLX5_BF_COPY(dst, src, num_bb, wq,
+                                   uct_ib_mlx5_bf_copy_bb_st64b);
     }
-    return src;
+#endif
+
+    return UCT_IB_MLX5_BF_COPY(dst, src, num_bb, wq, uct_ib_mlx5_bf_copy_bb);
 }
 
 static UCS_F_ALWAYS_INLINE void *