From 1855df30e29e92c17b8c05f3926790799fb7d75c Mon Sep 17 00:00:00 2001 From: Thomas Vegas Date: Wed, 10 Jun 2026 15:58:03 +0300 Subject: [PATCH 1/6] UCT/IB/MLX5: Add AArch64 ST64B BF copy mode --- config/m4/sysdep.m4 | 5 + src/ucs/arch/aarch64/cpu.c | 35 ++++++ src/ucs/arch/aarch64/cpu.h | 6 +- src/ucs/arch/cpu.h | 3 +- src/ucs/sys/init.c | 1 + src/uct/ib/mlx5/configure.m4 | 26 +++++ src/uct/ib/mlx5/dc/dc_mlx5.c | 3 +- src/uct/ib/mlx5/dv/ib_mlx5_dv.c | 7 ++ src/uct/ib/mlx5/gdaki/gdaki.c | 1 + src/uct/ib/mlx5/ib_mlx5.c | 159 ++++++++++++++++++++++++++++ src/uct/ib/mlx5/ib_mlx5.h | 35 +++++- src/uct/ib/mlx5/ib_mlx5.inl | 8 +- src/uct/ib/mlx5/rc/rc_mlx5_common.c | 5 +- src/uct/ib/mlx5/rc/rc_mlx5_common.h | 1 + src/uct/ib/mlx5/rc/rc_mlx5_iface.c | 4 +- src/uct/ib/mlx5/ud/ud_mlx5.c | 8 +- src/uct/ib/mlx5/ud/ud_mlx5.h | 2 +- test/gtest/uct/ib/test_rc.cc | 100 +++++++++++++++-- 18 files changed, 383 insertions(+), 26 deletions(-) diff --git a/config/m4/sysdep.m4 b/config/m4/sysdep.m4 index d951db1a0e3..bf4c87f73e0 100644 --- a/config/m4/sysdep.m4 +++ b/config/m4/sysdep.m4 @@ -20,6 +20,7 @@ AC_CHECK_LIB([rt], [timer_create], [], AC_MSG_ERROR([librt not found])) # Extended string functions # AC_CHECK_HEADERS([libgen.h]) +AC_CHECK_HEADERS([sys/auxv.h]) AC_CHECK_DECLS([asprintf, basename, fmemopen], [], AC_MSG_ERROR([GNU string extensions not found]), [#define _GNU_SOURCE 1 @@ -29,6 +30,10 @@ AC_CHECK_DECLS([asprintf, basename, fmemopen], [], #include #endif ]) +AC_CHECK_DECLS([getauxval], [], [], + [#if HAVE_SYS_AUXV_H + #include + #endif]) # diff --git a/src/ucs/arch/aarch64/cpu.c b/src/ucs/arch/aarch64/cpu.c index 980bb1e818b..4462e6416b7 100644 --- a/src/ucs/arch/aarch64/cpu.c +++ b/src/ucs/arch/aarch64/cpu.c @@ -13,6 +13,18 @@ #include #include +#if HAVE_SYS_AUXV_H +# include +#endif + +/* 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 static void ucs_aarch64_cpuid_from_proc(ucs_aarch64_cpuid_t *cpuid) { @@ -68,4 +80,27 @@ 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; + static int initialized = 0; + + if (!initialized) { + ucs_cpu_flag_t result = 0; + +#if HAVE_SYS_AUXV_H && HAVE_DECL_GETAUXVAL + if (getauxval(AT_HWCAP3) & HWCAP3_LS64) { + result |= UCS_CPU_FLAG_ST64B; + } +#endif + + cpu_flag = result; + ucs_memory_cpu_store_fence(); + initialized = 1; + } + + ucs_memory_cpu_load_fence(); + return cpu_flag; +} + #endif diff --git a/src/ucs/arch/aarch64/cpu.h b/src/ucs/arch/aarch64/cpu.h index b590a1452f9..41492dc0964 100644 --- a/src/ucs/arch/aarch64/cpu.h +++ b/src/ucs/arch/aarch64/cpu.h @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -169,10 +170,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() UCS_F_NOOPTIMIZE; static inline void ucs_cpu_init() { diff --git a/src/ucs/arch/cpu.h b/src/ucs/arch/cpu.h index f5ab9be9a05..651abd8178d 100644 --- a/src/ucs/arch/cpu.h +++ b/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_ST64B = UCS_BIT(11) } ucs_cpu_flag_t; diff --git a/src/ucs/sys/init.c b/src/ucs/sys/init.c index 6d72121801c..955cd6faab9 100644 --- a/src/ucs/sys/init.c +++ b/src/ucs/sys/init.c @@ -50,6 +50,7 @@ static UCS_F_NOOPTIMIZE void ucs_check_cpu_flags(void) { "sse42", UCS_CPU_FLAG_SSE42 }, { "avx", UCS_CPU_FLAG_AVX }, { "avx2", UCS_CPU_FLAG_AVX2 }, + { "st64b", UCS_CPU_FLAG_ST64B }, { NULL, UCS_CPU_FLAG_UNKNOWN }, }; diff --git a/src/uct/ib/mlx5/configure.m4 b/src/uct/ib/mlx5/configure.m4 index ed62c325f66..8ae89834d67 100644 --- a/src/uct/ib/mlx5/configure.m4 +++ b/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]) diff --git a/src/uct/ib/mlx5/dc/dc_mlx5.c b/src/uct/ib/mlx5/dc/dc_mlx5.c index 46ca61aef00..94d8116a0e2 100644 --- a/src/uct/ib/mlx5/dc/dc_mlx5.c +++ b/src/uct/ib/mlx5/dc/dc_mlx5.c @@ -464,7 +464,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; diff --git a/src/uct/ib/mlx5/dv/ib_mlx5_dv.c b/src/uct/ib/mlx5/dv/ib_mlx5_dv.c index 824719d74c7..f74d2e9dd02 100644 --- a/src/uct/ib/mlx5/dv/ib_mlx5_dv.c +++ b/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, diff --git a/src/uct/ib/mlx5/gdaki/gdaki.c b/src/uct/ib/mlx5/gdaki/gdaki.c index 75356419180..9e1ce08b681 100644 --- a/src/uct/ib/mlx5/gdaki/gdaki.c +++ b/src/uct/ib/mlx5/gdaki/gdaki.c @@ -459,6 +459,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); diff --git a/src/uct/ib/mlx5/ib_mlx5.c b/src/uct/ib/mlx5/ib_mlx5.c index 40cbba23578..3136bac5715 100644 --- a/src/uct/ib/mlx5/ib_mlx5.c +++ b/src/uct/ib/mlx5/ib_mlx5.c @@ -32,6 +32,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_GENERIC] = "generic", + [UCT_IB_MLX5_BF_COPY_MODE_AUTO] = "auto", + [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", @@ -54,6 +61,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" + " generic - Use portable scalar stores.\n" + " st64b - Use AArch64 ST64B store when supported.\n" + " auto - Select best according to runtime CPU capabilities.", + 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" @@ -763,8 +779,146 @@ uct_ib_mlx5_get_mmio_mode(uct_priv_worker_t *worker, return UCS_OK; } +#if defined(__aarch64__) +static UCS_F_ALWAYS_INLINE void +uct_ib_mlx5_bf_copy_bb_generic(void *restrict dst, void *restrict src) +{ + UCS_WORD_COPY(volatile uint64_t, dst, uint64_t, src, MLX5_SEND_WQE_BB); +} + +static UCS_F_ALWAYS_INLINE void * +uct_ib_mlx5_bf_copy_generic(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_generic(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; + } + } + + return src; +} + +#if HAVE_AARCH64_ST64B_ASM +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"); +} + +static UCS_F_ALWAYS_INLINE void * +uct_ib_mlx5_bf_copy_st64b(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_st64b(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; + } + } + + return src; +} +#endif + +static uct_ib_mlx5_bf_copy_func_t +uct_ib_mlx5_select_bf_copy(uct_ib_mlx5_bf_copy_mode_t bf_copy_mode) +{ + static uct_ib_mlx5_bf_copy_func_t cached_bf_copy_func; + uct_ib_mlx5_bf_copy_mode_t mode = bf_copy_mode; + uct_ib_mlx5_bf_copy_func_t bf_copy_func; + int have_st64b = 0; + + ucs_assert(bf_copy_mode < UCT_IB_MLX5_BF_COPY_MODE_LAST); + + if (cached_bf_copy_func != NULL) { + return cached_bf_copy_func; + } + + bf_copy_func = uct_ib_mlx5_bf_copy_generic; + if (mode == UCT_IB_MLX5_BF_COPY_MODE_GENERIC) { + goto out; + } + +#if HAVE_AARCH64_ST64B_ASM + have_st64b = ucs_arch_get_cpu_flag() & UCS_CPU_FLAG_ST64B; +#endif + if (mode == UCT_IB_MLX5_BF_COPY_MODE_AUTO) { + mode = have_st64b ? UCT_IB_MLX5_BF_COPY_MODE_ST64B : + UCT_IB_MLX5_BF_COPY_MODE_GENERIC; + } + +#if HAVE_AARCH64_ST64B_ASM + if (mode == UCT_IB_MLX5_BF_COPY_MODE_ST64B) { + ucs_assert(have_st64b); + bf_copy_func = uct_ib_mlx5_bf_copy_st64b; + } +#else + ucs_assert(mode != UCT_IB_MLX5_BF_COPY_MODE_ST64B); +#endif + +out: + cached_bf_copy_func = bf_copy_func; + return bf_copy_func; +} +#endif + +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) +{ +#if defined(__aarch64__) + if (bf_copy_mode == UCT_IB_MLX5_BF_COPY_MODE_ST64B) { +#if !HAVE_AARCH64_ST64B_ASM + ucs_error("mlx5 BlueFlame ST64B copy was requested but UCX was built " + "without assembler support for ST64B"); + return UCS_ERR_UNSUPPORTED; +#else + if (!(ucs_arch_get_cpu_flag() & UCS_CPU_FLAG_ST64B)) { + ucs_error("mlx5 BlueFlame ST64B copy was requested but CPU does " + "not report ST64B support"); + return UCS_ERR_UNSUPPORTED; + } +#endif + } + + txwq->bf_copy = uct_ib_mlx5_select_bf_copy(bf_copy_mode); +#else + (void)txwq; + + if (bf_copy_mode == UCT_IB_MLX5_BF_COPY_MODE_ST64B) { + ucs_error("mlx5 BlueFlame ST64B copy is supported only on AArch64"); + 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) { @@ -798,6 +952,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, diff --git a/src/uct/ib/mlx5/ib_mlx5.h b/src/uct/ib/mlx5/ib_mlx5.h index d5c2ddd4c3e..b4fe8690ece 100644 --- a/src/uct/ib/mlx5/ib_mlx5.h +++ b/src/uct/ib/mlx5/ib_mlx5.h @@ -468,6 +468,14 @@ typedef enum { } uct_ib_mlx5_mmio_mode_t; +typedef enum { + UCT_IB_MLX5_BF_COPY_MODE_GENERIC, + UCT_IB_MLX5_BF_COPY_MODE_AUTO, + 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 { @@ -475,9 +483,10 @@ typedef struct uct_ib_mlx5_iface_config { 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 +652,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; @@ -674,6 +684,14 @@ typedef struct uct_ib_mlx5_qp { } uct_ib_mlx5_qp_t; /* Send work-queue */ +struct uct_ib_mlx5_txwq; + +#if defined(__aarch64__) +typedef void *(*uct_ib_mlx5_bf_copy_func_t)( + void *dst, void *src, uint16_t num_bb, + const struct uct_ib_mlx5_txwq *wq); +#endif + typedef struct uct_ib_mlx5_txwq { uct_ib_mlx5_qp_t super; uint16_t sw_pi; /* PI for next WQE */ @@ -690,6 +708,9 @@ typedef struct uct_ib_mlx5_txwq { uint8_t flags; /* Debug flags */ #endif uct_ib_fence_info_t fi; +#if defined(__aarch64__) + uct_ib_mlx5_bf_copy_func_t bf_copy; +#endif } uct_ib_mlx5_txwq_t; @@ -892,7 +913,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); diff --git a/src/uct/ib/mlx5/ib_mlx5.inl b/src/uct/ib/mlx5/ib_mlx5.inl index c9dde7a82aa..bdb544a12be 100644 --- a/src/uct/ib/mlx5/ib_mlx5.inl +++ b/src/uct/ib/mlx5/ib_mlx5.inl @@ -536,9 +536,6 @@ size_t uct_ib_mlx5_set_data_seg_iov(uct_ib_mlx5_txwq_t *txwq, static UCS_F_ALWAYS_INLINE void uct_ib_mlx5_bf_copy_bb(void * restrict dst, void * restrict src) { -#if defined(__ARM_NEON) - UCS_WORD_COPY(int16x8_t, dst, int16x8_t, src, MLX5_SEND_WQE_BB); -#else #if defined(__SSE4_2__) typedef __m128i uct_ib_mlx5_send_wqe_bb_block_t; #else @@ -553,13 +550,15 @@ static UCS_F_ALWAYS_INLINE void uct_ib_mlx5_bf_copy_bb(void * restrict dst, UCS_WORD_COPY(uct_ib_mlx5_send_wqe_bb_t, dst, uct_ib_mlx5_send_wqe_bb_t, src, MLX5_SEND_WQE_BB); -#endif } 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) { +#if defined(__aarch64__) + return wq->bf_copy(dst, src, num_bb, wq); +#else uint16_t n; for (n = 0; n < num_bb; ++n) { @@ -571,6 +570,7 @@ void *uct_ib_mlx5_bf_copy(void *dst, void *src, uint16_t num_bb, } } return src; +#endif } static UCS_F_ALWAYS_INLINE uint16_t diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_common.c b/src/uct/ib/mlx5/rc/rc_mlx5_common.c index ed7d5cd1249..a478ada1495 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_common.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_common.c @@ -289,6 +289,7 @@ ucs_status_t uct_rc_mlx5_devx_create_cmd_qp(uct_rc_mlx5_iface_common_t *iface) attr.super.srq_num = iface->rx.srq.srq_num; attr.super.port = dev->first_port; attr.mmio_mode = iface->tx.mmio_mode; + attr.bf_copy_mode = iface->tx.bf_copy_mode; status = uct_ib_mlx5_devx_create_qp(&iface->super.super, &iface->cq[UCT_IB_DIR_RX], &iface->cq[UCT_IB_DIR_RX], @@ -414,6 +415,7 @@ uct_rc_mlx5_get_cmd_qp(uct_rc_mlx5_iface_common_t *iface) iface->tm.cmd_wq.super.super.qp_num = qp->qp_num; return uct_ib_mlx5_txwq_init(iface->super.super.super.worker, iface->tx.mmio_mode, + iface->tx.bf_copy_mode, &iface->tm.cmd_wq.super, qp); } #endif @@ -534,12 +536,13 @@ void uct_rc_mlx5_iface_fill_attr(uct_rc_mlx5_iface_common_t *iface, case UCT_IB_MLX5_OBJ_TYPE_DEVX: case UCT_IB_MLX5_OBJ_TYPE_NULL: uct_rc_iface_fill_attr(&iface->super, &qp_attr->super, max_send_wr, NULL); - qp_attr->mmio_mode = iface->tx.mmio_mode; break; case UCT_IB_MLX5_OBJ_TYPE_LAST: break; } + qp_attr->mmio_mode = iface->tx.mmio_mode; + qp_attr->bf_copy_mode = iface->tx.bf_copy_mode; qp_attr->super.srq_num = srq->srq_num; } diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_common.h b/src/uct/ib/mlx5/rc/rc_mlx5_common.h index e38bfa45f54..6c56382c628 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_common.h +++ b/src/uct/ib/mlx5/rc/rc_mlx5_common.h @@ -347,6 +347,7 @@ typedef struct uct_rc_mlx5_iface_common { struct { ucs_mpool_t atomic_desc_mp; uct_ib_mlx5_mmio_mode_t mmio_mode; + uct_ib_mlx5_bf_copy_mode_t bf_copy_mode; uint16_t bb_max; /* limit number of outstanding WQE BBs */ } tx; struct { diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c index a5120aab5dc..f74015d928f 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_iface.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_iface.c @@ -315,7 +315,8 @@ ucs_status_t uct_rc_mlx5_iface_create_qp(uct_rc_mlx5_iface_common_t *iface, if (attr->super.cap.max_send_wr) { status = uct_ib_mlx5_txwq_init(iface->super.super.super.worker, - iface->tx.mmio_mode, txwq, + iface->tx.mmio_mode, + iface->tx.bf_copy_mode, txwq, qp->verbs.qp); if (status != UCS_OK) { ucs_error("Failed to get mlx5 QP information"); @@ -777,6 +778,7 @@ UCS_CLASS_INIT_FUNC(uct_rc_mlx5_iface_common_t, uct_iface_ops_t *tl_ops, dev = uct_ib_iface_device(&self->super.super); self->tx.mmio_mode = mlx5_config->super.mmio_mode; + self->tx.bf_copy_mode = mlx5_config->super.bf_copy_mode; self->tx.bb_max = ucs_min(mlx5_config->tx_max_bb, UINT16_MAX); self->tm.am_desc.super.cb = uct_rc_mlx5_release_desc; diff --git a/src/uct/ib/mlx5/ud/ud_mlx5.c b/src/uct/ib/mlx5/ud/ud_mlx5.c index 73e6e3ec0fa..85588e934d3 100644 --- a/src/uct/ib/mlx5/ud/ud_mlx5.c +++ b/src/uct/ib/mlx5/ud/ud_mlx5.c @@ -815,7 +815,8 @@ static ucs_status_t uct_ud_mlx5_iface_create_qp(uct_ib_iface_t *ib_iface, } status = uct_ib_mlx5_txwq_init(iface->super.super.super.worker, - iface->tx.mmio_mode, &iface->tx.wq, + iface->tx.mmio_mode, + iface->tx.bf_copy_mode, &iface->tx.wq, qp->verbs.qp); if (status != UCS_OK) { goto err_destroy_qp; @@ -1002,8 +1003,9 @@ static UCS_CLASS_INIT_FUNC(uct_ud_mlx5_iface_t, uct_md_h tl_md, uct_ib_mlx5_parse_cqe_zipping(md, &config->mlx5_common, &init_attr); - self->tx.mmio_mode = config->mlx5_common.mmio_mode; - self->tx.wq.super.type = UCT_IB_MLX5_OBJ_TYPE_LAST; + self->tx.mmio_mode = config->mlx5_common.mmio_mode; + self->tx.bf_copy_mode = config->mlx5_common.bf_copy_mode; + self->tx.wq.super.type = UCT_IB_MLX5_OBJ_TYPE_LAST; UCS_CLASS_CALL_SUPER_INIT(uct_ud_iface_t, &uct_ud_mlx5_iface_ops, &uct_ud_mlx5_iface_tl_ops, tl_md, worker, params, diff --git a/src/uct/ib/mlx5/ud/ud_mlx5.h b/src/uct/ib/mlx5/ud/ud_mlx5.h index 4649b35048d..b05975b2779 100644 --- a/src/uct/ib/mlx5/ud/ud_mlx5.h +++ b/src/uct/ib/mlx5/ud/ud_mlx5.h @@ -37,6 +37,7 @@ typedef struct { struct { uct_ib_mlx5_txwq_t wq; uct_ib_mlx5_mmio_mode_t mmio_mode; + uct_ib_mlx5_bf_copy_mode_t bf_copy_mode; } tx; struct { uct_ib_mlx5_rxwq_t wq; @@ -60,4 +61,3 @@ uct_ud_mlx5_tx_moderation(uct_ud_mlx5_iface_t *iface, uint8_t ce_se) } #endif - diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index 9a3cc64e092..e49a220d2b1 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -9,6 +9,16 @@ #include #include +extern "C" { +#include +} + +#ifdef HAVE_MLX5_DV +extern "C" { +#include +#include +} +#endif void test_rc::init() { @@ -144,6 +154,90 @@ UCS_TEST_SKIP_COND_P(test_rc_max_wr, send_limit, UCT_INSTANTIATE_RC_TEST_CASE(test_rc_max_wr) +#ifdef HAVE_MLX5_DV +/* Real-HCA BF MMIO workload to compare copy modes with external counters. */ +class test_rc_mlx5_bf_mmio : public test_rc { +protected: + enum { + NUM_ITERS = 65536 + }; + + void verify_bf_post() + { + uct_rc_mlx5_base_ep_t *ep; + + ep = ucs_derived_of(m_e1->ep(0), uct_rc_mlx5_base_ep_t); + ASSERT_TRUE(ep->tx.wq.reg != NULL); + ASSERT_EQ(UCT_IB_MLX5_MMIO_MODE_BF_POST, ep->tx.wq.reg->mode); + } + + void run_bf_mmio_workload() + { + const uint64_t payload = 0xdeadbeefULL; + ucs_status_t status; + + ASSERT_NO_FATAL_FAILURE(verify_bf_post()); + + for (unsigned i = 0; i < NUM_ITERS; ++i) { + do { + status = uct_ep_am_short(m_e1->ep(0), 0, i, &payload, + sizeof(payload)); + if (status == UCS_ERR_NO_RESOURCE) { + short_progress_loop(); + } + } while (status == UCS_ERR_NO_RESOURCE); + + ASSERT_UCS_OK(status); + + if ((i & 0x3ff) == 0) { + progress(); + } + } + + flush(); + UCS_TEST_MESSAGE << "posted " << NUM_ITERS + << " rc_mlx5 BlueFlame AM short operations"; + } +}; + +UCS_TEST_SKIP_COND_P(test_rc_mlx5_bf_mmio, bf_copy_generic, + !check_caps(UCT_IFACE_FLAG_AM_SHORT), + "RC_MLX5_MMIO_MODE=bf_post", + "RC_MLX5_BF_COPY_MODE=generic", + "RC_FC_ENABLE?=n") +{ + run_bf_mmio_workload(); +} + +_UCT_INSTANTIATE_TEST_CASE(test_rc_mlx5_bf_mmio, rc_mlx5) + +#if defined(__aarch64__) && HAVE_AARCH64_ST64B_ASM +class test_rc_mlx5_bf_mmio_st64b : public test_rc_mlx5_bf_mmio { +public: + virtual void init() + { + if (!(ucs_arch_get_cpu_flag() & UCS_CPU_FLAG_ST64B)) { + UCS_TEST_SKIP_R("CPU does not report ST64B support"); + } + + test_rc_mlx5_bf_mmio::init(); + } +}; + +UCS_TEST_SKIP_COND_P(test_rc_mlx5_bf_mmio_st64b, bf_copy_st64b, + !check_caps(UCT_IFACE_FLAG_AM_SHORT), + "RC_MLX5_MMIO_MODE=bf_post", + "RC_MLX5_BF_COPY_MODE=st64b", + "RC_FC_ENABLE?=n") +{ + run_bf_mmio_workload(); +} + +_UCT_INSTANTIATE_TEST_CASE(test_rc_mlx5_bf_mmio_st64b, rc_mlx5) +#endif + +#endif + class test_rc_iface_flush_remote : public uct_test { protected: @@ -1026,12 +1120,6 @@ UCT_INSTANTIATE_RC_TEST_CASE(test_rc_flow_control_stats) #endif -#ifdef HAVE_MLX5_DV -extern "C" { -#include -} -#endif - test_uct_iface_attrs::attr_map_t test_rc_iface_attrs::get_num_iov() { if (has_transport("rc_mlx5") || has_transport("gga_mlx5")) { return get_num_iov_mlx5_common(0ul); From 826574b2c570743d43192f1a8a6c4b28e1ccd3ba Mon Sep 17 00:00:00 2001 From: Thomas Vegas Date: Wed, 10 Jun 2026 16:40:14 +0300 Subject: [PATCH 2/6] UCT/IB/MLX5: Add AArch64 ST64B BF copy mode --- src/uct/ib/mlx5/ib_mlx5.c | 48 +++++++++++++++++---------------------- 1 file changed, 21 insertions(+), 27 deletions(-) diff --git a/src/uct/ib/mlx5/ib_mlx5.c b/src/uct/ib/mlx5/ib_mlx5.c index 3136bac5715..5bd5e5161e9 100644 --- a/src/uct/ib/mlx5/ib_mlx5.c +++ b/src/uct/ib/mlx5/ib_mlx5.c @@ -842,8 +842,9 @@ uct_ib_mlx5_bf_copy_st64b(void *dst, void *src, uint16_t num_bb, } #endif -static uct_ib_mlx5_bf_copy_func_t -uct_ib_mlx5_select_bf_copy(uct_ib_mlx5_bf_copy_mode_t bf_copy_mode) +static ucs_status_t +uct_ib_mlx5_select_bf_copy(uct_ib_mlx5_bf_copy_mode_t bf_copy_mode, + uct_ib_mlx5_bf_copy_func_t *bf_copy_func_p) { static uct_ib_mlx5_bf_copy_func_t cached_bf_copy_func; uct_ib_mlx5_bf_copy_mode_t mode = bf_copy_mode; @@ -852,10 +853,6 @@ uct_ib_mlx5_select_bf_copy(uct_ib_mlx5_bf_copy_mode_t bf_copy_mode) ucs_assert(bf_copy_mode < UCT_IB_MLX5_BF_COPY_MODE_LAST); - if (cached_bf_copy_func != NULL) { - return cached_bf_copy_func; - } - bf_copy_func = uct_ib_mlx5_bf_copy_generic; if (mode == UCT_IB_MLX5_BF_COPY_MODE_GENERIC) { goto out; @@ -869,18 +866,29 @@ uct_ib_mlx5_select_bf_copy(uct_ib_mlx5_bf_copy_mode_t bf_copy_mode) UCT_IB_MLX5_BF_COPY_MODE_GENERIC; } -#if HAVE_AARCH64_ST64B_ASM if (mode == UCT_IB_MLX5_BF_COPY_MODE_ST64B) { - ucs_assert(have_st64b); +#if HAVE_AARCH64_ST64B_ASM + if (!have_st64b) { + ucs_error("mlx5 BlueFlame ST64B copy was requested but CPU does " + "not report ST64B support"); + return UCS_ERR_UNSUPPORTED; + } + bf_copy_func = uct_ib_mlx5_bf_copy_st64b; - } #else - ucs_assert(mode != UCT_IB_MLX5_BF_COPY_MODE_ST64B); + ucs_error("mlx5 BlueFlame ST64B copy was requested but UCX was built " + "without assembler support for ST64B"); + return UCS_ERR_UNSUPPORTED; #endif + } out: - cached_bf_copy_func = bf_copy_func; - return bf_copy_func; + if (cached_bf_copy_func == NULL) { + cached_bf_copy_func = bf_copy_func; + } + + *bf_copy_func_p = cached_bf_copy_func; + return UCS_OK; } #endif @@ -889,21 +897,7 @@ uct_ib_mlx5_txwq_init_bf_copy(uct_ib_mlx5_txwq_t *txwq, uct_ib_mlx5_bf_copy_mode_t bf_copy_mode) { #if defined(__aarch64__) - if (bf_copy_mode == UCT_IB_MLX5_BF_COPY_MODE_ST64B) { -#if !HAVE_AARCH64_ST64B_ASM - ucs_error("mlx5 BlueFlame ST64B copy was requested but UCX was built " - "without assembler support for ST64B"); - return UCS_ERR_UNSUPPORTED; -#else - if (!(ucs_arch_get_cpu_flag() & UCS_CPU_FLAG_ST64B)) { - ucs_error("mlx5 BlueFlame ST64B copy was requested but CPU does " - "not report ST64B support"); - return UCS_ERR_UNSUPPORTED; - } -#endif - } - - txwq->bf_copy = uct_ib_mlx5_select_bf_copy(bf_copy_mode); + return uct_ib_mlx5_select_bf_copy(bf_copy_mode, &txwq->bf_copy); #else (void)txwq; From 938781d62e0462f3ccd6306674c8ac6f8f10f4a7 Mon Sep 17 00:00:00 2001 From: Thomas Vegas Date: Tue, 16 Jun 2026 14:35:04 +0300 Subject: [PATCH 3/6] UCT/IB/MLX5: Add AArch64 ST64B BF copy mode --- src/ucs/arch/aarch64/cpu.c | 4 +- src/ucs/arch/cpu.h | 2 +- src/ucs/sys/init.c | 2 +- src/uct/ib/mlx5/ib_mlx5.c | 126 +++++------------------------------ src/uct/ib/mlx5/ib_mlx5.h | 17 ++--- src/uct/ib/mlx5/ib_mlx5.inl | 60 +++++++++++++---- test/gtest/uct/ib/test_rc.cc | 4 +- 7 files changed, 75 insertions(+), 140 deletions(-) diff --git a/src/ucs/arch/aarch64/cpu.c b/src/ucs/arch/aarch64/cpu.c index 4462e6416b7..5bd498701fe 100644 --- a/src/ucs/arch/aarch64/cpu.c +++ b/src/ucs/arch/aarch64/cpu.c @@ -15,7 +15,6 @@ #if HAVE_SYS_AUXV_H # include -#endif /* Older userspace headers may not expose Arm64 HWCAP3 definitions. */ #ifndef AT_HWCAP3 @@ -25,6 +24,7 @@ #ifndef HWCAP3_LS64 # define HWCAP3_LS64 (1UL << 3) #endif +#endif static void ucs_aarch64_cpuid_from_proc(ucs_aarch64_cpuid_t *cpuid) { @@ -90,7 +90,7 @@ ucs_cpu_flag_t ucs_arch_get_cpu_flag() #if HAVE_SYS_AUXV_H && HAVE_DECL_GETAUXVAL if (getauxval(AT_HWCAP3) & HWCAP3_LS64) { - result |= UCS_CPU_FLAG_ST64B; + result |= UCS_CPU_FLAG_LS64; } #endif diff --git a/src/ucs/arch/cpu.h b/src/ucs/arch/cpu.h index 651abd8178d..b8086f4ad06 100644 --- a/src/ucs/arch/cpu.h +++ b/src/ucs/arch/cpu.h @@ -66,7 +66,7 @@ typedef enum ucs_cpu_flag { UCS_CPU_FLAG_SSE42 = UCS_BIT(8), UCS_CPU_FLAG_AVX = UCS_BIT(9), UCS_CPU_FLAG_AVX2 = UCS_BIT(10), - UCS_CPU_FLAG_ST64B = UCS_BIT(11) + UCS_CPU_FLAG_LS64 = UCS_BIT(11) } ucs_cpu_flag_t; diff --git a/src/ucs/sys/init.c b/src/ucs/sys/init.c index 955cd6faab9..7d795eef5e5 100644 --- a/src/ucs/sys/init.c +++ b/src/ucs/sys/init.c @@ -50,7 +50,7 @@ static UCS_F_NOOPTIMIZE void ucs_check_cpu_flags(void) { "sse42", UCS_CPU_FLAG_SSE42 }, { "avx", UCS_CPU_FLAG_AVX }, { "avx2", UCS_CPU_FLAG_AVX2 }, - { "st64b", UCS_CPU_FLAG_ST64B }, + { "ls64", UCS_CPU_FLAG_LS64 }, { NULL, UCS_CPU_FLAG_UNKNOWN }, }; diff --git a/src/uct/ib/mlx5/ib_mlx5.c b/src/uct/ib/mlx5/ib_mlx5.c index 5bd5e5161e9..c72bec0f125 100644 --- a/src/uct/ib/mlx5/ib_mlx5.c +++ b/src/uct/ib/mlx5/ib_mlx5.c @@ -33,8 +33,8 @@ static const char *uct_ib_mlx5_mmio_modes[] = { }; static const char *uct_ib_mlx5_bf_copy_modes[] = { - [UCT_IB_MLX5_BF_COPY_MODE_GENERIC] = "generic", [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 }; @@ -64,9 +64,9 @@ ucs_config_field_t uct_ib_mlx5_iface_config_table[] = { {"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 when supported.\n" - " auto - Select best according to runtime CPU capabilities.", + " st64b - Use AArch64 ST64B store when supported.", ucs_offsetof(uct_ib_mlx5_iface_config_t, bf_copy_mode), UCS_CONFIG_TYPE_ENUM(uct_ib_mlx5_bf_copy_modes)}, @@ -779,125 +779,35 @@ uct_ib_mlx5_get_mmio_mode(uct_priv_worker_t *worker, return UCS_OK; } -#if defined(__aarch64__) -static UCS_F_ALWAYS_INLINE void -uct_ib_mlx5_bf_copy_bb_generic(void *restrict dst, void *restrict src) -{ - UCS_WORD_COPY(volatile uint64_t, dst, uint64_t, src, MLX5_SEND_WQE_BB); -} - -static UCS_F_ALWAYS_INLINE void * -uct_ib_mlx5_bf_copy_generic(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_generic(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; - } - } - - return src; -} - -#if HAVE_AARCH64_ST64B_ASM -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"); -} - -static UCS_F_ALWAYS_INLINE void * -uct_ib_mlx5_bf_copy_st64b(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_st64b(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; - } - } - - return src; -} -#endif - -static ucs_status_t -uct_ib_mlx5_select_bf_copy(uct_ib_mlx5_bf_copy_mode_t bf_copy_mode, - uct_ib_mlx5_bf_copy_func_t *bf_copy_func_p) +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) { - static uct_ib_mlx5_bf_copy_func_t cached_bf_copy_func; - uct_ib_mlx5_bf_copy_mode_t mode = bf_copy_mode; - uct_ib_mlx5_bf_copy_func_t bf_copy_func; - int have_st64b = 0; - ucs_assert(bf_copy_mode < UCT_IB_MLX5_BF_COPY_MODE_LAST); - bf_copy_func = uct_ib_mlx5_bf_copy_generic; - if (mode == UCT_IB_MLX5_BF_COPY_MODE_GENERIC) { - goto out; +#if defined(__aarch64__) + 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 HAVE_AARCH64_ST64B_ASM - have_st64b = ucs_arch_get_cpu_flag() & UCS_CPU_FLAG_ST64B; -#endif - if (mode == UCT_IB_MLX5_BF_COPY_MODE_AUTO) { - mode = have_st64b ? UCT_IB_MLX5_BF_COPY_MODE_ST64B : - UCT_IB_MLX5_BF_COPY_MODE_GENERIC; + if (ucs_arch_get_cpu_flag() & UCS_CPU_FLAG_LS64) { + txwq->bf_copy_mode = UCT_IB_MLX5_BF_COPY_MODE_ST64B; + return UCS_OK; } +#endif - if (mode == UCT_IB_MLX5_BF_COPY_MODE_ST64B) { + if (bf_copy_mode == UCT_IB_MLX5_BF_COPY_MODE_ST64B) { #if HAVE_AARCH64_ST64B_ASM - if (!have_st64b) { - ucs_error("mlx5 BlueFlame ST64B copy was requested but CPU does " - "not report ST64B support"); - return UCS_ERR_UNSUPPORTED; - } - - bf_copy_func = uct_ib_mlx5_bf_copy_st64b; + ucs_error("mlx5 BlueFlame ST64B copy was requested but CPU does " + "not report LS64 support"); #else ucs_error("mlx5 BlueFlame ST64B copy was requested but UCX was built " "without assembler support for ST64B"); - return UCS_ERR_UNSUPPORTED; #endif + return UCS_ERR_UNSUPPORTED; } - -out: - if (cached_bf_copy_func == NULL) { - cached_bf_copy_func = bf_copy_func; - } - - *bf_copy_func_p = cached_bf_copy_func; - return UCS_OK; -} -#endif - -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) -{ -#if defined(__aarch64__) - return uct_ib_mlx5_select_bf_copy(bf_copy_mode, &txwq->bf_copy); #else (void)txwq; diff --git a/src/uct/ib/mlx5/ib_mlx5.h b/src/uct/ib/mlx5/ib_mlx5.h index b4fe8690ece..0a5e68ec8a0 100644 --- a/src/uct/ib/mlx5/ib_mlx5.h +++ b/src/uct/ib/mlx5/ib_mlx5.h @@ -469,8 +469,8 @@ typedef enum { typedef enum { - UCT_IB_MLX5_BF_COPY_MODE_GENERIC, 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; @@ -683,15 +683,6 @@ typedef struct uct_ib_mlx5_qp { }; } uct_ib_mlx5_qp_t; -/* Send work-queue */ -struct uct_ib_mlx5_txwq; - -#if defined(__aarch64__) -typedef void *(*uct_ib_mlx5_bf_copy_func_t)( - void *dst, void *src, uint16_t num_bb, - const struct uct_ib_mlx5_txwq *wq); -#endif - typedef struct uct_ib_mlx5_txwq { uct_ib_mlx5_qp_t super; uint16_t sw_pi; /* PI for next WQE */ @@ -703,14 +694,14 @@ typedef struct uct_ib_mlx5_txwq { void *qend; uint16_t bb_max; uint16_t sig_pi; /* PI for last signaled WQE */ +#if defined(__aarch64__) + 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 */ #endif uct_ib_fence_info_t fi; -#if defined(__aarch64__) - uct_ib_mlx5_bf_copy_func_t bf_copy; -#endif } uct_ib_mlx5_txwq_t; diff --git a/src/uct/ib/mlx5/ib_mlx5.inl b/src/uct/ib/mlx5/ib_mlx5.inl index bdb544a12be..f247c3fe010 100644 --- a/src/uct/ib/mlx5/ib_mlx5.inl +++ b/src/uct/ib/mlx5/ib_mlx5.inl @@ -536,6 +536,9 @@ size_t uct_ib_mlx5_set_data_seg_iov(uct_ib_mlx5_txwq_t *txwq, static UCS_F_ALWAYS_INLINE void uct_ib_mlx5_bf_copy_bb(void * restrict dst, void * restrict src) { +#if defined(__ARM_NEON) + UCS_WORD_COPY(int16x8_t, dst, int16x8_t, src, MLX5_SEND_WQE_BB); +#else #if defined(__SSE4_2__) typedef __m128i uct_ib_mlx5_send_wqe_bb_block_t; #else @@ -550,27 +553,58 @@ static UCS_F_ALWAYS_INLINE void uct_ib_mlx5_bf_copy_bb(void * restrict dst, UCS_WORD_COPY(uct_ib_mlx5_send_wqe_bb_t, dst, uct_ib_mlx5_send_wqe_bb_t, src, MLX5_SEND_WQE_BB); +#endif } +#if defined(__aarch64__) && HAVE_AARCH64_ST64B_ASM +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) { -#if defined(__aarch64__) - return wq->bf_copy(dst, src, num_bb, wq); -#else - 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 defined(__aarch64__) && HAVE_AARCH64_ST64B_ASM + 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 uint16_t diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index e49a220d2b1..01911cb0c5d 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -216,8 +216,8 @@ class test_rc_mlx5_bf_mmio_st64b : public test_rc_mlx5_bf_mmio { public: virtual void init() { - if (!(ucs_arch_get_cpu_flag() & UCS_CPU_FLAG_ST64B)) { - UCS_TEST_SKIP_R("CPU does not report ST64B support"); + if (!(ucs_arch_get_cpu_flag() & UCS_CPU_FLAG_LS64)) { + UCS_TEST_SKIP_R("CPU does not report LS64 support"); } test_rc_mlx5_bf_mmio::init(); From db95d46a4ac68d46ad15d104d7d689da85e0ba55 Mon Sep 17 00:00:00 2001 From: Thomas Vegas Date: Wed, 23 Sep 2026 18:35:13 +0300 Subject: [PATCH 4/6] UCT/IB/MLX5: Improve ST64B BF copy test and CPU flag check Add UCT_IB_MLX5_HAVE_ST64B so the ST64B build condition is defined once. --- config/m4/sysdep.m4 | 11 ++-- src/ucs/arch/cpu.h | 19 +++++++ src/ucs/sys/init.c | 1 - src/uct/ib/mlx5/ib_mlx5.c | 16 +++--- src/uct/ib/mlx5/ib_mlx5.h | 9 ++- src/uct/ib/mlx5/ib_mlx5.inl | 4 +- test/gtest/uct/ib/test_rc.cc | 105 +++++++++++++++++++++++++---------- 7 files changed, 118 insertions(+), 47 deletions(-) diff --git a/config/m4/sysdep.m4 b/config/m4/sysdep.m4 index bf4c87f73e0..a6db54775b4 100644 --- a/config/m4/sysdep.m4 +++ b/config/m4/sysdep.m4 @@ -20,7 +20,6 @@ AC_CHECK_LIB([rt], [timer_create], [], AC_MSG_ERROR([librt not found])) # Extended string functions # AC_CHECK_HEADERS([libgen.h]) -AC_CHECK_HEADERS([sys/auxv.h]) AC_CHECK_DECLS([asprintf, basename, fmemopen], [], AC_MSG_ERROR([GNU string extensions not found]), [#define _GNU_SOURCE 1 @@ -30,10 +29,12 @@ AC_CHECK_DECLS([asprintf, basename, fmemopen], [], #include #endif ]) -AC_CHECK_DECLS([getauxval], [], [], - [#if HAVE_SYS_AUXV_H - #include - #endif]) + + +# +# Auxiliary vector (getauxval() is checked by ucm.m4) +# +AC_CHECK_HEADERS([sys/auxv.h]) # diff --git a/src/ucs/arch/cpu.h b/src/ucs/arch/cpu.h index b8086f4ad06..b445c950caa 100644 --- a/src/ucs/arch/cpu.h +++ b/src/ucs/arch/cpu.h @@ -162,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_prefer_relaxed_order() { ucs_cpu_vendor_t cpu_vendor = ucs_arch_get_cpu_vendor(); diff --git a/src/ucs/sys/init.c b/src/ucs/sys/init.c index 7d795eef5e5..6d72121801c 100644 --- a/src/ucs/sys/init.c +++ b/src/ucs/sys/init.c @@ -50,7 +50,6 @@ static UCS_F_NOOPTIMIZE void ucs_check_cpu_flags(void) { "sse42", UCS_CPU_FLAG_SSE42 }, { "avx", UCS_CPU_FLAG_AVX }, { "avx2", UCS_CPU_FLAG_AVX2 }, - { "ls64", UCS_CPU_FLAG_LS64 }, { NULL, UCS_CPU_FLAG_UNKNOWN }, }; diff --git a/src/uct/ib/mlx5/ib_mlx5.c b/src/uct/ib/mlx5/ib_mlx5.c index c72bec0f125..e44a9ac97f3 100644 --- a/src/uct/ib/mlx5/ib_mlx5.c +++ b/src/uct/ib/mlx5/ib_mlx5.c @@ -785,34 +785,32 @@ uct_ib_mlx5_txwq_init_bf_copy(uct_ib_mlx5_txwq_t *txwq, { ucs_assert(bf_copy_mode < UCT_IB_MLX5_BF_COPY_MODE_LAST); -#if defined(__aarch64__) +#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 HAVE_AARCH64_ST64B_ASM - if (ucs_arch_get_cpu_flag() & UCS_CPU_FLAG_LS64) { + if (ucs_cpu_has_flag(UCS_CPU_FLAG_LS64)) { txwq->bf_copy_mode = UCT_IB_MLX5_BF_COPY_MODE_ST64B; return UCS_OK; } -#endif if (bf_copy_mode == UCT_IB_MLX5_BF_COPY_MODE_ST64B) { -#if HAVE_AARCH64_ST64B_ASM ucs_error("mlx5 BlueFlame ST64B copy was requested but CPU does " "not report LS64 support"); -#else - ucs_error("mlx5 BlueFlame ST64B copy was requested but UCX was built " - "without assembler support for ST64B"); -#endif return UCS_ERR_UNSUPPORTED; } #else (void)txwq; 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 diff --git a/src/uct/ib/mlx5/ib_mlx5.h b/src/uct/ib/mlx5/ib_mlx5.h index 0a5e68ec8a0..6864e73aaf3 100644 --- a/src/uct/ib/mlx5/ib_mlx5.h +++ b/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)) @@ -694,7 +701,7 @@ typedef struct uct_ib_mlx5_txwq { void *qend; uint16_t bb_max; uint16_t sig_pi; /* PI for last signaled WQE */ -#if defined(__aarch64__) +#if UCT_IB_MLX5_HAVE_ST64B uint8_t bf_copy_mode; #endif #if UCS_ENABLE_ASSERT diff --git a/src/uct/ib/mlx5/ib_mlx5.inl b/src/uct/ib/mlx5/ib_mlx5.inl index f247c3fe010..d2af1dc4d5b 100644 --- a/src/uct/ib/mlx5/ib_mlx5.inl +++ b/src/uct/ib/mlx5/ib_mlx5.inl @@ -556,7 +556,7 @@ static UCS_F_ALWAYS_INLINE void uct_ib_mlx5_bf_copy_bb(void * restrict dst, #endif } -#if defined(__aarch64__) && HAVE_AARCH64_ST64B_ASM +#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) { @@ -597,7 +597,7 @@ 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) { -#if defined(__aarch64__) && HAVE_AARCH64_ST64B_ASM +#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); diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index 01911cb0c5d..801651f9e59 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -155,85 +155,132 @@ UCS_TEST_SKIP_COND_P(test_rc_max_wr, send_limit, UCT_INSTANTIATE_RC_TEST_CASE(test_rc_max_wr) #ifdef HAVE_MLX5_DV -/* Real-HCA BF MMIO workload to compare copy modes with external counters. */ -class test_rc_mlx5_bf_mmio : public test_rc { +class test_rc_mlx5_bf_copy : public test_rc { +public: + virtual void init() + { + test_rc::init(); + + m_rx_count = 0; + m_rx_errors = 0; + uct_iface_set_am_handler(m_e2->iface(), AM_ID, am_handler, this, 0); + } + protected: enum { - NUM_ITERS = 65536 + AM_ID = 0, + NUM_ITERS = 512, + MAX_PAYLOAD = UCT_IB_MLX5_BF_REG_SIZE }; - void verify_bf_post() + /* Post AM shorts of varying size, so that single-BB and multi-BB BlueFlame + * copies and the send queue wrap-around are all exercised. */ + void test_bf_copy() { uct_rc_mlx5_base_ep_t *ep; + uint8_t payload[MAX_PAYLOAD]; + size_t max_len, len, i; + ucs_status_t status; + unsigned iter; ep = ucs_derived_of(m_e1->ep(0), uct_rc_mlx5_base_ep_t); ASSERT_TRUE(ep->tx.wq.reg != NULL); ASSERT_EQ(UCT_IB_MLX5_MMIO_MODE_BF_POST, ep->tx.wq.reg->mode); - } - void run_bf_mmio_workload() - { - const uint64_t payload = 0xdeadbeefULL; - ucs_status_t status; + max_len = m_e1->iface_attr().cap.am.max_short - sizeof(uint64_t); + if (max_len > sizeof(payload)) { + max_len = sizeof(payload); + } - ASSERT_NO_FATAL_FAILURE(verify_bf_post()); + for (iter = 0; iter < NUM_ITERS; ++iter) { + len = iter % (max_len + 1); + for (i = 0; i < len; ++i) { + payload[i] = iter + i; + } - for (unsigned i = 0; i < NUM_ITERS; ++i) { do { - status = uct_ep_am_short(m_e1->ep(0), 0, i, &payload, - sizeof(payload)); + status = uct_ep_am_short(m_e1->ep(0), AM_ID, iter, payload, + len); if (status == UCS_ERR_NO_RESOURCE) { - short_progress_loop(); + progress_loop(); } } while (status == UCS_ERR_NO_RESOURCE); ASSERT_UCS_OK(status); + } + + flush(); + wait_for_value(&m_rx_count, (unsigned)NUM_ITERS, true); + EXPECT_EQ((unsigned)NUM_ITERS, m_rx_count); + EXPECT_EQ(0u, m_rx_errors); + } + + static ucs_status_t am_handler(void *arg, void *data, size_t length, + unsigned flags) + { + test_rc_mlx5_bf_copy *self = reinterpret_cast( + arg); + const uint8_t *payload; + uint64_t hdr; + size_t i; + + if (length < sizeof(hdr)) { + ++self->m_rx_errors; + return UCS_OK; + } + + memcpy(&hdr, data, sizeof(hdr)); + payload = reinterpret_cast(data) + sizeof(hdr); - if ((i & 0x3ff) == 0) { - progress(); + for (i = 0; i < (length - sizeof(hdr)); ++i) { + if (payload[i] != (uint8_t)(hdr + i)) { + ++self->m_rx_errors; + break; } } - flush(); - UCS_TEST_MESSAGE << "posted " << NUM_ITERS - << " rc_mlx5 BlueFlame AM short operations"; + ++self->m_rx_count; + return UCS_OK; } + + volatile unsigned m_rx_count; + volatile unsigned m_rx_errors; }; -UCS_TEST_SKIP_COND_P(test_rc_mlx5_bf_mmio, bf_copy_generic, +UCS_TEST_SKIP_COND_P(test_rc_mlx5_bf_copy, generic, !check_caps(UCT_IFACE_FLAG_AM_SHORT), "RC_MLX5_MMIO_MODE=bf_post", "RC_MLX5_BF_COPY_MODE=generic", "RC_FC_ENABLE?=n") { - run_bf_mmio_workload(); + test_bf_copy(); } -_UCT_INSTANTIATE_TEST_CASE(test_rc_mlx5_bf_mmio, rc_mlx5) +_UCT_INSTANTIATE_TEST_CASE(test_rc_mlx5_bf_copy, rc_mlx5) -#if defined(__aarch64__) && HAVE_AARCH64_ST64B_ASM -class test_rc_mlx5_bf_mmio_st64b : public test_rc_mlx5_bf_mmio { +#if UCT_IB_MLX5_HAVE_ST64B +class test_rc_mlx5_bf_copy_st64b : public test_rc_mlx5_bf_copy { public: virtual void init() { - if (!(ucs_arch_get_cpu_flag() & UCS_CPU_FLAG_LS64)) { + if (!ucs_cpu_has_flag(UCS_CPU_FLAG_LS64)) { UCS_TEST_SKIP_R("CPU does not report LS64 support"); } - test_rc_mlx5_bf_mmio::init(); + test_rc_mlx5_bf_copy::init(); } }; -UCS_TEST_SKIP_COND_P(test_rc_mlx5_bf_mmio_st64b, bf_copy_st64b, +UCS_TEST_SKIP_COND_P(test_rc_mlx5_bf_copy_st64b, st64b, !check_caps(UCT_IFACE_FLAG_AM_SHORT), "RC_MLX5_MMIO_MODE=bf_post", "RC_MLX5_BF_COPY_MODE=st64b", "RC_FC_ENABLE?=n") { - run_bf_mmio_workload(); + test_bf_copy(); } -_UCT_INSTANTIATE_TEST_CASE(test_rc_mlx5_bf_mmio_st64b, rc_mlx5) +_UCT_INSTANTIATE_TEST_CASE(test_rc_mlx5_bf_copy_st64b, rc_mlx5) #endif #endif From bf3e7cf13f048e5eb8ae1fd6530fd93a7238605a Mon Sep 17 00:00:00 2001 From: Thomas Vegas Date: Mon, 28 Sep 2026 01:56:42 -0700 Subject: [PATCH 5/6] UCT/IB/MLX5: Address review comments --- config/m4/sysdep.m4 | 6 ------ src/ucs/arch/aarch64/cpu.c | 12 ++++-------- src/ucs/arch/aarch64/cpu.h | 3 +-- src/uct/ib/mlx5/ib_mlx5.c | 2 -- src/uct/ib/mlx5/ib_mlx5.h | 7 ++++--- src/uct/ib/mlx5/rc/rc_mlx5_common.c | 2 +- test/gtest/uct/ib/test_rc.cc | 19 +++++++++++++++++++ 7 files changed, 29 insertions(+), 22 deletions(-) diff --git a/config/m4/sysdep.m4 b/config/m4/sysdep.m4 index a6db54775b4..d951db1a0e3 100644 --- a/config/m4/sysdep.m4 +++ b/config/m4/sysdep.m4 @@ -31,12 +31,6 @@ AC_CHECK_DECLS([asprintf, basename, fmemopen], [], ]) -# -# Auxiliary vector (getauxval() is checked by ucm.m4) -# -AC_CHECK_HEADERS([sys/auxv.h]) - - # # CPU-sets # diff --git a/src/ucs/arch/aarch64/cpu.c b/src/ucs/arch/aarch64/cpu.c index 5bd498701fe..ba90b8655bb 100644 --- a/src/ucs/arch/aarch64/cpu.c +++ b/src/ucs/arch/aarch64/cpu.c @@ -13,7 +13,7 @@ #include #include -#if HAVE_SYS_AUXV_H +#if HAVE_DECL_GETAUXVAL # include /* Older userspace headers may not expose Arm64 HWCAP3 definitions. */ @@ -82,24 +82,20 @@ void ucs_aarch64_cpuid(ucs_aarch64_cpuid_t *cpuid) ucs_cpu_flag_t ucs_arch_get_cpu_flag() { - static ucs_cpu_flag_t cpu_flag; - static int initialized = 0; + static ucs_cpu_flag_t cpu_flag = UCS_CPU_FLAG_UNKNOWN; - if (!initialized) { + if (cpu_flag == UCS_CPU_FLAG_UNKNOWN) { ucs_cpu_flag_t result = 0; -#if HAVE_SYS_AUXV_H && HAVE_DECL_GETAUXVAL +#if HAVE_DECL_GETAUXVAL if (getauxval(AT_HWCAP3) & HWCAP3_LS64) { result |= UCS_CPU_FLAG_LS64; } #endif cpu_flag = result; - ucs_memory_cpu_store_fence(); - initialized = 1; } - ucs_memory_cpu_load_fence(); return cpu_flag; } diff --git a/src/ucs/arch/aarch64/cpu.h b/src/ucs/arch/aarch64/cpu.h index 41492dc0964..7419df95722 100644 --- a/src/ucs/arch/aarch64/cpu.h +++ b/src/ucs/arch/aarch64/cpu.h @@ -15,7 +15,6 @@ #include #include #include -#include #include #include #include @@ -170,7 +169,7 @@ static inline ucs_cpu_model_t ucs_arch_get_cpu_model() return UCS_CPU_MODEL_ARM_AARCH64; } -ucs_cpu_flag_t ucs_arch_get_cpu_flag() UCS_F_NOOPTIMIZE; +ucs_cpu_flag_t ucs_arch_get_cpu_flag(); static inline void ucs_cpu_init() { diff --git a/src/uct/ib/mlx5/ib_mlx5.c b/src/uct/ib/mlx5/ib_mlx5.c index 16913540810..de9f2546e3e 100644 --- a/src/uct/ib/mlx5/ib_mlx5.c +++ b/src/uct/ib/mlx5/ib_mlx5.c @@ -822,8 +822,6 @@ uct_ib_mlx5_txwq_init_bf_copy(uct_ib_mlx5_txwq_t *txwq, return UCS_ERR_UNSUPPORTED; } #else - (void)txwq; - 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 " diff --git a/src/uct/ib/mlx5/ib_mlx5.h b/src/uct/ib/mlx5/ib_mlx5.h index 3feebd41751..5f7877db0b3 100644 --- a/src/uct/ib/mlx5/ib_mlx5.h +++ b/src/uct/ib/mlx5/ib_mlx5.h @@ -700,6 +700,7 @@ typedef struct uct_ib_mlx5_qp { }; } uct_ib_mlx5_qp_t; +/* Send work-queue */ typedef struct uct_ib_mlx5_txwq { uct_ib_mlx5_qp_t super; uint16_t sw_pi; /* PI for next WQE */ @@ -713,13 +714,13 @@ 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 uint16_t hw_ci; /* First BB index of last completed WQE */ uint16_t ft_ci; /* First BB index of last ft completed WQE */ uint16_t path_mtu_mask; /* Path MTU in bytes - 1 */ uint8_t path_mtu_shift; /* log2(path MTU in bytes) */ +#if UCT_IB_MLX5_HAVE_ST64B + uint8_t bf_copy_mode; +#endif #if UCS_ENABLE_ASSERT uint8_t flags; /* Debug flags */ #endif diff --git a/src/uct/ib/mlx5/rc/rc_mlx5_common.c b/src/uct/ib/mlx5/rc/rc_mlx5_common.c index f8c38a58c04..af7c43e7396 100644 --- a/src/uct/ib/mlx5/rc/rc_mlx5_common.c +++ b/src/uct/ib/mlx5/rc/rc_mlx5_common.c @@ -536,12 +536,12 @@ void uct_rc_mlx5_iface_fill_attr(uct_rc_mlx5_iface_common_t *iface, case UCT_IB_MLX5_OBJ_TYPE_DEVX: case UCT_IB_MLX5_OBJ_TYPE_NULL: uct_rc_iface_fill_attr(&iface->super, &qp_attr->super, max_send_wr, NULL); + qp_attr->mmio_mode = iface->tx.mmio_mode; break; case UCT_IB_MLX5_OBJ_TYPE_LAST: break; } - qp_attr->mmio_mode = iface->tx.mmio_mode; qp_attr->bf_copy_mode = iface->tx.bf_copy_mode; qp_attr->super.srq_num = srq->srq_num; } diff --git a/test/gtest/uct/ib/test_rc.cc b/test/gtest/uct/ib/test_rc.cc index 9d977dd0590..bc042047067 100644 --- a/test/gtest/uct/ib/test_rc.cc +++ b/test/gtest/uct/ib/test_rc.cc @@ -322,11 +322,30 @@ class test_rc_mlx5_bf_copy : public test_rc { { test_rc::init(); + if (!check_bf_caps()) { + cleanup(); + UCS_TEST_SKIP_R("write-combining UAR is not supported"); + } + m_rx_count = 0; m_rx_errors = 0; uct_iface_set_am_handler(m_e2->iface(), AM_ID, am_handler, this, 0); } + /* BlueFlame writes the full WQE to the MMIO register, which is only valid + * when the UAR is write-combining. */ + bool check_bf_caps() + { + FOR_EACH_ENTITY(entity) { + uct_ib_mlx5_md_t *md = ucs_derived_of((*entity)->md(), + uct_ib_mlx5_md_t); + if (!(md->flags & UCT_IB_MLX5_MD_FLAG_UAR_USE_WC)) { + return false; + } + } + return true; + } + protected: enum { AM_ID = 0, From b6d85113878434ca9902d4b8f830c0469e53a0ea Mon Sep 17 00:00:00 2001 From: Thomas Vegas Date: Mon, 28 Sep 2026 13:27:21 +0300 Subject: [PATCH 6/6] UCT/IB/MLX5: Improve configuration helper --- src/uct/ib/mlx5/ib_mlx5.c | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/uct/ib/mlx5/ib_mlx5.c b/src/uct/ib/mlx5/ib_mlx5.c index de9f2546e3e..8801aebbdc2 100644 --- a/src/uct/ib/mlx5/ib_mlx5.c +++ b/src/uct/ib/mlx5/ib_mlx5.c @@ -67,7 +67,7 @@ ucs_config_field_t uct_ib_mlx5_iface_config_table[] = { "the following:\n" " auto - Select best according to runtime CPU capabilities.\n" " generic - Use portable scalar stores.\n" - " st64b - Use AArch64 ST64B store when supported.", + " 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)},