[MUSA]Feat: musa support

AUTHORS

@@ -149,6 +149,7 @@ Zihao Zhao <zizhao@nvidia.com>
 lzhang2 <cherry.zhang@intel.com>
 michaelzli <michaelzli@tencent.com>
 shasson5 <103439971+shasson5@users.noreply.github.com>
+superleo <xlei.boll@gmail.com>
 wehy <18174858434@163.com>
 ybenabou <ybenabou@nvidia.com>
 yuanwu <yuan.wu@intel.com>

config/m4/musa.m4

@@ -0,0 +1,134 @@
+#
+# Copyright (C) MThreads 2026. ALL RIGHTS RESERVED.
+# See file LICENSE for terms.
+#
+
+AC_DEFUN([UCX_CHECK_MUSA],[
+    AS_IF([test "x$musa_checked" != "xyes"], [
+        AC_ARG_WITH([musa],
+                    [AS_HELP_STRING([--with-musa=(DIR)],
+                                    [Enable MUSA support (default is guess).])],
+                    [],
+                    [with_musa=guess])
+
+        musa_happy=no
+        MUSA_CPPFLAGS=""
+        MUSA_LDFLAGS=""
+        MUSA_LIBS=""
+        MUSA_ROOT=""
+
+        AS_IF([test "x$with_musa" = "xno"], [
+            AC_MSG_WARN([MUSA support was explicitly disabled])
+        ], [
+            AS_IF([test "x$with_musa" = "xyes" ||
+                   test "x$with_musa" = "xguess"], [
+                AS_IF([test -d /usr/local/musa],
+                      [MUSA_ROOT=/usr/local/musa])
+            ], [
+                MUSA_ROOT="$with_musa"
+            ])
+
+            AS_IF([test -n "$MUSA_ROOT"], [
+                MUSA_CPPFLAGS="-I$MUSA_ROOT/include"
+                MUSA_LDFLAGS="-L$MUSA_ROOT/lib -L$MUSA_ROOT/lib64"
+                MUSA_LIBS="-lmusart -lmusa"
+
+                musa_happy=yes
+                save_CPPFLAGS="$CPPFLAGS"
+                save_LDFLAGS="$LDFLAGS"
+                save_LIBS="$LIBS"
+
+                CPPFLAGS="$MUSA_CPPFLAGS $CPPFLAGS"
+                LDFLAGS="$MUSA_LDFLAGS $LDFLAGS"
+                LIBS="$MUSA_LIBS $LIBS"
+
+                AC_CHECK_HEADERS([musa_runtime_api.h],
+                                 [],
+                                 [musa_happy=no])
+                AC_CHECK_HEADERS([musa.h],
+                                 [],
+                                 [musa_happy=no])
+
+                AS_IF([test "x$musa_happy" = "xyes"], [
+                    AC_CHECK_LIB([musart], [musaGetDeviceCount],
+                                 [],
+                                 [musa_happy=no])
+                ])
+                AS_IF([test "x$musa_happy" = "xyes"], [
+                    AC_CHECK_LIB([musart], [musaMalloc],
+                                 [],
+                                 [musa_happy=no])
+                    AC_CHECK_LIB([musart], [musaFree],
+                                 [],
+                                 [musa_happy=no])
+                    AC_CHECK_LIB([musart], [musaMemcpy],
+                                 [],
+                                 [musa_happy=no])
+                    AC_CHECK_LIB([musart], [musaHostRegister],
+                                 [],
+                                 [musa_happy=no])
+                    AC_CHECK_LIB([musart], [musaHostUnregister],
+                                 [],
+                                 [musa_happy=no])
+                    AC_CHECK_LIB([musart], [musaDeviceGetPCIBusId],
+                                 [],
+                                 [musa_happy=no])
+                ])
+                AS_IF([test "x$musa_happy" = "xyes"], [
+                    AC_CHECK_LIB([musa], [muMemGetAddressRange],
+                                 [],
+                                 [musa_happy=no])
+                ])
+                AS_IF([test "x$musa_happy" = "xyes"], [
+                    AC_CHECK_LIB([musa], [muMemGetHandleForAddressRange],
+                                 [AC_DEFINE([HAVE_MUSA_DMABUF], [1],
+                                            [MUSA supports DMA-BUF export])],
+                                 [])
+                ])
+                AS_IF([test "x$musa_happy" = "xyes"], [
+                    AC_LINK_IFELSE(
+                        [AC_LANG_PROGRAM(
+                            [[#include <stddef.h>
+                              #include <musa.h>
+                              #include <musa_runtime_api.h>]],
+                            [[int count = 0;
+                              int device = 0;
+                              void *address = 0;
+                              MUdeviceptr base = 0;
+                              size_t size = 0;
+                              musaGetDeviceCount(&count);
+                              musaGetDevice(&device);
+                              musaSetDevice(device);
+                              musaPointerGetAttributes(NULL, address);
+                              musaGetErrorString(musaSuccess);
+                              muMemGetAddressRange(&base, &size,
+                                                   (MUdeviceptr)address);
+                              return count < 0;]])],
+                        [],
+                        [musa_happy=no])
+                ])
+
+                CPPFLAGS="$save_CPPFLAGS"
+                LDFLAGS="$save_LDFLAGS"
+                LIBS="$save_LIBS"
+            ])
+
+            AS_IF([test "x$musa_happy" = "xyes"], [
+                AC_DEFINE([HAVE_MUSA], [1], [Enable MUSA support])
+                AC_SUBST([MUSA_CPPFLAGS])
+                AC_SUBST([MUSA_LDFLAGS])
+                AC_SUBST([MUSA_LIBS])
+                AC_SUBST([MUSA_ROOT])
+            ], [
+                AS_IF([test "x$with_musa" != "xguess"], [
+                    AC_MSG_ERROR([MUSA support was requested but MUSA headers and libraries were not found])
+                ], [
+                    AC_MSG_WARN([MUSA headers and libraries not found])
+                ])
+            ])
+        ])
+
+        musa_checked=yes
+        AM_CONDITIONAL([HAVE_MUSA], [test "x$musa_happy" = "xyes"])
+    ])
+])

configure.ac

@@ -281,6 +281,7 @@ AS_IF([test "x$with_docs_only" = xyes],
      AM_CONDITIONAL([HAVE_LCOV], [false])
      AM_CONDITIONAL([HAVE_ZE], [false])
      AM_CONDITIONAL([HAVE_GAUDI], [false])
+     AM_CONDITIONAL([HAVE_MUSA], [false])
      AM_CONDITIONAL([HAVE_GDA], [false])
     ],
     [
@@ -295,6 +296,7 @@ AS_IF([test "x$with_docs_only" = xyes],
      m4_include([config/m4/cuda.m4])
      m4_include([config/m4/verbs.m4])
      m4_include([config/m4/rocm.m4])
+     m4_include([config/m4/musa.m4])
      m4_include([config/m4/ze.m4])
      m4_include([config/m4/gdrcopy.m4])
      m4_include([config/m4/mad.m4])

src/ucm/AGENTS.md

@@ -26,6 +26,7 @@ is `api/ucm.h`; everything else is internal.
   `reloc.c` (ELF relocation walking), `log.c`, `sys.c`, `khash_safe.h`.
 - `cuda/` — `cudamem` hooks for `cudaMalloc`/`cudaFree` and friends. Built
   as a module; emits `MEM_TYPE_ALLOC`/`FREE` events.
+- `musa/` — `musamem` hooks for `musaMalloc`/`musaFree`.
 - `rocm/` — analogous hooks for `hipMalloc`/`hipFree`.
 - `ze/` — Intel Level Zero allocations.
 - `ptmalloc286/` — third-party allocator, only built when in-tree malloc

src/ucm/Makefile.am

@@ -6,7 +6,7 @@
 # See file LICENSE for terms.
 #
 
-SUBDIRS = . cuda rocm ze
+SUBDIRS = . cuda musa rocm ze
 
 lib_LTLIBRARIES    = libucm.la
 libucm_ladir       = $(includedir)/ucm
@@ -62,4 +62,3 @@ libucm_la_SOURCES += \
 noinst_HEADERS += \
     ptmalloc286/malloc-2.8.6.h
 endif
-

src/ucm/configure.m4

@@ -31,6 +31,7 @@ LDFLAGS="$SAVE_LDFLAGS"
 
 ucm_modules=""
 m4_include([src/ucm/cuda/configure.m4])
+m4_include([src/ucm/musa/configure.m4])
 m4_include([src/ucm/rocm/configure.m4])
 m4_include([src/ucm/ze/configure.m4])
 AC_DEFINE_UNQUOTED([ucm_MODULES], ["${ucm_modules}"], [UCM loadable modules])

src/ucm/musa/Makefile.am

@@ -0,0 +1,19 @@
+#
+# Copyright (C) MThreads 2026. ALL RIGHTS RESERVED.
+# See file LICENSE for terms.
+#
+
+if HAVE_MUSA
+module_LTLIBRARIES      = libucm_musa.la
+libucm_musa_la_CPPFLAGS = $(BASE_CPPFLAGS) $(MUSA_CPPFLAGS)
+libucm_musa_la_CFLAGS   = $(BASE_CFLAGS) $(MUSA_CFLAGS) $(LT_CFLAGS)
+libucm_musa_la_LIBADD   = ../libucm.la $(MUSA_LIBS)
+libucm_musa_la_LDFLAGS  = $(UCM_MODULE_LDFLAGS) \
+                          $(patsubst %, -Xlinker %, $(MUSA_LDFLAGS)) \
+                          -version-info $(SOVERSION) $(UCX_LT_RELEASE)
+
+libucm_musa_la_SOURCES = \
+	musamem.c
+
+include $(top_srcdir)/config/module.am
+endif

src/ucm/musa/configure.m4

@@ -0,0 +1,9 @@
+#
+# Copyright (C) MThreads 2026. ALL RIGHTS RESERVED.
+# See file LICENSE for terms.
+#
+
+UCX_CHECK_MUSA
+AS_IF([test "x$musa_happy" = "xyes"],
+      [ucm_modules="${ucm_modules}:musa"])
+AC_CONFIG_FILES([src/ucm/musa/Makefile])

src/ucm/musa/musamem.c

@@ -0,0 +1,174 @@
+/*
+ * Copyright (C) MThreads 2026. ALL RIGHTS RESERVED.
+ * See file LICENSE for terms.
+ */
+
+#ifdef HAVE_CONFIG_H
+#  include "config.h"
+#endif
+
+#include <musa.h>
+#include <musa_runtime_api.h>
+
+#include <ucm/event/event.h>
+#include <ucm/util/log.h>
+#include <ucm/util/reloc.h>
+#include <ucm/util/replace.h>
+#include <ucm/util/sys.h>
+#include <ucs/sys/compiler.h>
+#include <ucs/sys/preprocessor.h>
+
+#include <pthread.h>
+#include <string.h>
+#include <sys/mman.h>
+
+static musaError_t ucm_musaMalloc(void **address, size_t size);
+static musaError_t ucm_musaFree(void *address);
+
+UCM_DEFINE_REPLACE_DLSYM_FUNC(musaMalloc, musaError_t, musaErrorUnknown,
+                              void **, size_t)
+UCM_DEFINE_REPLACE_DLSYM_FUNC(musaFree, musaError_t, musaErrorUnknown, void *)
+
+static void
+ucm_dispatch_musa_mem_type_event(ucm_event_type_t event_type, void *address,
+                                 size_t length)
+{
+    ucm_event_t event;
+
+    if (address == NULL) {
+        return;
+    }
+
+    event.mem_type.address   = address;
+    event.mem_type.size      = length;
+    event.mem_type.mem_type  = UCS_MEMORY_TYPE_MUSA;
+    event.mem_type.sys_dev   = UCS_SYS_DEVICE_ID_UNKNOWN;
+    event.mem_type.mem_flags = UCS_MEM_FLAG_REGISTRABLE;
+    ucm_event_dispatch(event_type, &event);
+}
+
+static size_t ucm_musa_mem_size(void *address)
+{
+    MUdeviceptr base;
+    size_t length;
+
+    if ((address == NULL) ||
+        (muMemGetAddressRange(&base, &length, (MUdeviceptr)address) !=
+         MUSA_SUCCESS)) {
+        return 1;
+    }
+
+    return length;
+}
+
+static musaError_t ucm_musaMalloc(void **address, size_t size)
+{
+    musaError_t status;
+
+    ucm_event_enter();
+    status = ucm_orig_musaMalloc(address, size);
+    if ((status == musaSuccess) && (address != NULL) && (*address != NULL)) {
+        ucm_dispatch_musa_mem_type_event(UCM_EVENT_MEM_TYPE_ALLOC, *address,
+                                         size);
+    }
+    ucm_event_leave();
+    return status;
+}
+
+static musaError_t ucm_musaFree(void *address)
+{
+    musaError_t status;
+    size_t length;
+
+    ucm_event_enter();
+    length = ucm_musa_mem_size(address);
+    ucm_dispatch_musa_mem_type_event(UCM_EVENT_MEM_TYPE_FREE, address, length);
+    status = ucm_orig_musaFree(address);
+    ucm_event_leave();
+    return status;
+}
+
+static ucm_reloc_patch_t ucm_musa_patches[] = {
+    {UCS_PP_MAKE_STRING(musaMalloc), ucm_override_musaMalloc},
+    {UCS_PP_MAKE_STRING(musaFree), ucm_override_musaFree},
+    {NULL, NULL}
+};
+
+static ucs_status_t ucm_musamem_install(int events)
+{
+    static int installed;
+    static pthread_mutex_t install_mutex = PTHREAD_MUTEX_INITIALIZER;
+    ucm_reloc_patch_t *patch;
+    ucs_status_t status = UCS_OK;
+
+    if (!(events & (UCM_EVENT_MEM_TYPE_ALLOC | UCM_EVENT_MEM_TYPE_FREE))) {
+        return UCS_OK;
+    }
+
+    pthread_mutex_lock(&install_mutex);
+    if (installed) {
+        goto out;
+    }
+
+    for (patch = ucm_musa_patches; patch->symbol != NULL; ++patch) {
+        status = ucm_reloc_modify(patch);
+        if (status != UCS_OK) {
+            ucm_warn("failed to install MUSA hook for '%s'", patch->symbol);
+            goto out;
+        }
+    }
+
+    installed = 1;
+    ucm_info("MUSA memory hooks are ready");
+
+out:
+    pthread_mutex_unlock(&install_mutex);
+    return status;
+}
+
+static int
+ucm_musamem_scan_regions_cb(void *arg, void *addr, size_t length, int prot,
+                            const char *path)
+{
+    static const char musa_path_pattern[] = "/dev/musa";
+    ucm_event_handler_t *handler = arg;
+    ucm_event_t event;
+
+    if ((prot & (PROT_READ | PROT_WRITE | PROT_EXEC)) &&
+        strncmp(path, musa_path_pattern, sizeof(musa_path_pattern) - 1)) {
+        return 0;
+    }
+
+    event.mem_type.address   = addr;
+    event.mem_type.size      = length;
+    event.mem_type.mem_type  = UCS_MEMORY_TYPE_LAST;
+    event.mem_type.sys_dev   = UCS_SYS_DEVICE_ID_UNKNOWN;
+    event.mem_type.mem_flags = UCS_MEM_FLAG_REGISTRABLE;
+
+    ucm_event_enter();
+    handler->cb(UCM_EVENT_MEM_TYPE_ALLOC, &event, handler->arg);
+    ucm_event_leave();
+    return 0;
+}
+
+static void ucm_musamem_get_existing_alloc(ucm_event_handler_t *handler)
+{
+    if (handler->events & UCM_EVENT_MEM_TYPE_ALLOC) {
+        ucm_parse_proc_self_maps(ucm_musamem_scan_regions_cb, handler);
+    }
+}
+
+static ucm_event_installer_t ucm_musa_initializer = {
+    .install            = ucm_musamem_install,
+    .get_existing_alloc = ucm_musamem_get_existing_alloc
+};
+
+UCS_STATIC_INIT
+{
+    ucs_list_add_tail(&ucm_event_installer_list, &ucm_musa_initializer.list);
+}
+
+UCS_STATIC_CLEANUP
+{
+    ucs_list_del(&ucm_musa_initializer.list);
+}