Skip to content

TOOLS/PERF: Add cuda-async memory allocator #11566

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tvegas1 merged 10 commits into from
Aug 27, 2026
Merged

TOOLS/PERF: Add cuda-async memory allocator #11566

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
e3bfc5e
TOOLS/PERF: Allocate cuda-async memory
tvegas1 Jun 17, 2026
ab1968d
TOOLS/PERF: Allocate cuda-async memory
tvegas1 Jun 22, 2026
cf97549
TOOLS/PERF: Allocate cuda-async memory
tvegas1 Jun 22, 2026
886bfb2
TOOLS/PERF: Allocate cuda-async memory
tvegas1 Jun 22, 2026
8d25838
TOOLS/PERF: Allocate cuda-async memory
tvegas1 Jun 23, 2026
b05832a
TOOLS/PERF: Address review comments
tvegas1 Jul 3, 2026
1922402
TOOLS/PERF: Resolve cuda-async mem type from config
tvegas1 Aug 10, 2026
f7f7ce1
TOOLS/PERF: Allocate cuda-async memory
tvegas1 Aug 10, 2026
9f1b0b3
TOOLS/PERF: Use resolve_mem_type for allocator mem type
tvegas1 Aug 10, 2026
20234df
Merge branch 'master' into perftest_cuda_async
tvegas1 Aug 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions contrib/test_jenkins.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -514,6 +514,17 @@ run_ucx_perftest() {

echo "==== Running ucx_perf with cuda memory ===="

if $ucx_perftest -h 2>&1 | grep -q "cuda-async"
then
echo "==== Running ucx_perf with cuda-async memory ===="
cuda_async_test_args="-t tag_lat -D contig,contig"
cuda_async_test_args+=" -m cuda-async,cuda-async -s 8 -n 10 -w 1"
run_client_server_app "$ucx_perftest" "$cuda_async_test_args" \
"$(hostname)" 0 0
else
echo "==== cuda-async memory is not supported, skipping ===="
fi

for memtype_cache in y n
do
for gdr in $gdr_options
Expand Down
11 changes: 7 additions & 4 deletions src/tools/perf/api/libperf.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,9 @@ typedef enum {
} ucx_perf_channel_mode_t;


#define UCX_PERF_ALLOC_NAME_MAX 32 /* Max memory allocator name length */


enum ucx_perf_test_flags {
UCX_PERF_TEST_FLAG_VALIDATE = UCS_BIT(1), /* Validate data. Affects performance. */
UCX_PERF_TEST_FLAG_ONE_SIDED = UCS_BIT(2), /* For tests which involves only one side,
Expand Down Expand Up @@ -272,6 +275,10 @@ typedef struct ucx_perf_params {
ucx_perf_wait_mode_t wait_mode; /* How to wait */
ucs_memory_type_t send_mem_type; /* Send memory type */
ucs_memory_type_t recv_mem_type; /* Recv memory type */
char send_mem_alloc_name[UCX_PERF_ALLOC_NAME_MAX]; /* Send memory allocator
name */
char recv_mem_alloc_name[UCX_PERF_ALLOC_NAME_MAX]; /* Recv memory allocator
name */

@tomerg-nvidia tomerg-nvidia Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alignment/long lines

@tvegas1 tvegas1 Jun 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed

ucx_perf_accel_dev_t send_device; /* Send memory device */
ucx_perf_accel_dev_t recv_device; /* Recv memory device */
ucs_device_level_t device_level; /* Device level */
Expand Down Expand Up @@ -337,10 +344,6 @@ typedef struct {
} ucp_perf_daemon_req_t;


/* Allocators for each memory type */
typedef struct ucx_perf_allocator ucx_perf_allocator_t;
extern const ucx_perf_allocator_t* ucx_perf_mem_type_allocators[];

/* Device test dispatchers for each memory type */
typedef struct ucx_perf_device_dispatcher ucx_perf_device_dispatcher_t;
extern const ucx_perf_device_dispatcher_t* ucx_perf_mem_type_device_dispatchers[];
Expand Down
189 changes: 149 additions & 40 deletions src/tools/perf/cuda/cuda_alloc.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,12 +68,45 @@ static inline ucs_status_t ucx_perf_cuda_alloc(size_t length,
return UCS_OK;
}

static inline ucs_status_t
uct_perf_cuda_alloc_reg_mem(const ucx_perf_context_t *perf,
size_t length,
ucs_memory_type_t mem_type,
unsigned flags,
uct_allocated_memory_t *alloc_mem)
#if CUDART_VERSION >= 11020
static ucs_status_t ucx_perf_cuda_async_mem_alloc(
const ucx_perf_context_t *perf, size_t length, void **address_p)
{
cudaError_t cerr;
void *address;

(void)perf;

CUDA_CALL_RET(UCS_ERR_NO_MEMORY, cudaMallocAsync, &address, length, 0);
cerr = cudaStreamSynchronize(0);
if (cerr != cudaSuccess) {
ucs_error("cudaStreamSynchronize() failed: %d (%s)", (int)cerr,
cudaGetErrorString(cerr));
CUDA_CALL_WARN(cudaFreeAsync, address, 0);
CUDA_CALL_WARN(cudaStreamSynchronize, 0);
return UCS_ERR_IO_ERROR;
}

*address_p = address;
return UCS_OK;
}

static void ucx_perf_cuda_async_mem_free(
const ucx_perf_context_t *perf, void *address)
{
(void)perf;

CUDA_CALL_WARN(cudaStreamSynchronize, 0);
CUDA_CALL_WARN(cudaFreeAsync, address, 0);
CUDA_CALL_WARN(cudaStreamSynchronize, 0);
}
#endif

static ucs_status_t uct_perf_cuda_reg_mem(const ucx_perf_context_t *perf,
size_t length,
ucs_memory_type_t mem_type,
unsigned flags,
uct_allocated_memory_t *alloc_mem)
{
uct_md_attr_v2_t md_attr = {.field_mask = UCT_MD_ATTR_FIELD_REG_ALIGNMENT};
size_t reg_length;
Expand All @@ -86,11 +119,6 @@ uct_perf_cuda_alloc_reg_mem(const ucx_perf_context_t *perf,
return status;
}

status = ucx_perf_cuda_alloc(length, mem_type, &alloc_mem->address);
if (status != UCS_OK) {
return status;
}

/* Register memory respecting MD reg_alignment */
reg_address = alloc_mem->address;
reg_length = length;
Expand All @@ -99,7 +127,6 @@ uct_perf_cuda_alloc_reg_mem(const ucx_perf_context_t *perf,
status = uct_md_mem_reg(perf->uct.md, reg_address, reg_length, flags,
&alloc_mem->memh);
if (status != UCS_OK) {
cudaFree(alloc_mem->address);
ucs_error("failed to register memory");
return status;
}
Expand All @@ -110,6 +137,27 @@ uct_perf_cuda_alloc_reg_mem(const ucx_perf_context_t *perf,
return UCS_OK;
}

static ucs_status_t uct_perf_cuda_alloc_reg_mem(

@tomerg-nvidia tomerg-nvidia Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here and in many other places you use uct_perf prefix, shouldn't it be ucx_perf?

@tvegas1 tvegas1 Jun 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

fixed

const ucx_perf_context_t *perf, size_t length,
ucs_memory_type_t mem_type, unsigned flags,
uct_allocated_memory_t *alloc_mem)
{
ucs_status_t status;

status = ucx_perf_cuda_alloc(length, mem_type, &alloc_mem->address);
if (status != UCS_OK) {
return status;
}

status = uct_perf_cuda_reg_mem(perf, length, mem_type, flags, alloc_mem);
if (status != UCS_OK) {
CUDA_CALL_WARN(cudaFree, alloc_mem->address);
return status;
}

return UCS_OK;
}

static ucs_status_t uct_perf_cuda_alloc(const ucx_perf_context_t *perf,
size_t length, unsigned flags,
uct_allocated_memory_t *alloc_mem)
Expand All @@ -118,16 +166,40 @@ static ucs_status_t uct_perf_cuda_alloc(const ucx_perf_context_t *perf,
flags, alloc_mem);
}

static ucs_status_t uct_perf_cuda_managed_alloc(const ucx_perf_context_t *perf,
size_t length, unsigned flags,
uct_allocated_memory_t *alloc_mem)
static ucs_status_t uct_perf_cuda_managed_alloc(
const ucx_perf_context_t *perf, size_t length, unsigned flags,
uct_allocated_memory_t *alloc_mem)
{
return uct_perf_cuda_alloc_reg_mem(perf, length, UCS_MEMORY_TYPE_CUDA_MANAGED,
flags, alloc_mem);
return uct_perf_cuda_alloc_reg_mem(perf, length,
UCS_MEMORY_TYPE_CUDA_MANAGED, flags,
alloc_mem);
}

static void uct_perf_cuda_free(const ucx_perf_context_t *perf,
uct_allocated_memory_t *alloc_mem)
#if CUDART_VERSION >= 11020
static ucs_status_t uct_perf_cuda_async_alloc(
const ucx_perf_context_t *perf, size_t length, unsigned flags,
uct_allocated_memory_t *alloc_mem)
{
ucs_status_t status;

status = ucx_perf_cuda_async_mem_alloc(perf, length, &alloc_mem->address);
if (status != UCS_OK) {
return status;
}

status = uct_perf_cuda_reg_mem(perf, length, UCS_MEMORY_TYPE_CUDA_MANAGED,

@tomerg-nvidia tomerg-nvidia Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can it use the ASYNC_MEM_TYPE

@tvegas1 tvegas1 Jun 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i would keep it for later pr

flags, alloc_mem);
if (status != UCS_OK) {
ucx_perf_cuda_async_mem_free(perf, alloc_mem->address);
return status;
}

return UCS_OK;
}
#endif

static void uct_perf_cuda_dereg(const ucx_perf_context_t *perf,
uct_allocated_memory_t *alloc_mem)
{
ucs_status_t status;

Expand All @@ -137,10 +209,24 @@ static void uct_perf_cuda_free(const ucx_perf_context_t *perf,
if (status != UCS_OK) {
ucs_error("failed to deregister memory");
}
}

static void uct_perf_cuda_free(const ucx_perf_context_t *perf,
uct_allocated_memory_t *alloc_mem)
{
uct_perf_cuda_dereg(perf, alloc_mem);
CUDA_CALL_WARN(cudaFree, alloc_mem->address);
}

#if CUDART_VERSION >= 11020
static void uct_perf_cuda_async_free(const ucx_perf_context_t *perf,
uct_allocated_memory_t *alloc_mem)
{
uct_perf_cuda_dereg(perf, alloc_mem);
ucx_perf_cuda_async_mem_free(perf, alloc_mem->address);
}
#endif

static void ucx_perf_cuda_memcpy(void *dst, ucs_memory_type_t dst_mem_type,
const void *src, ucs_memory_type_t src_mem_type,
size_t count)
Expand All @@ -156,29 +242,52 @@ static void* ucx_perf_cuda_memset(void *dst, int value, size_t count)
return dst;
}

#if CUDART_VERSION >= 11020
static ucx_perf_allocator_t cuda_async_allocator = {
.name = "cuda-async",
.mem_type = UCS_MEMORY_TYPE_CUDA_MANAGED,

@tomerg-nvidia tomerg-nvidia Jun 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can it use the ASYNC_MEM_TYPE

@tvegas1 tvegas1 Jun 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i would keep it for later pr

@brminich brminich Jul 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why?

@tvegas1 tvegas1 Aug 10, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added it.

.init = ucx_perf_cuda_init,
.uct_alloc = uct_perf_cuda_async_alloc,
.uct_free = uct_perf_cuda_async_free,
.mem_alloc = ucx_perf_cuda_async_mem_alloc,
.mem_free = ucx_perf_cuda_async_mem_free,
.memcpy = ucx_perf_cuda_memcpy,
.memset = ucx_perf_cuda_memset
};
#endif

static ucx_perf_allocator_t cuda_allocator = {
.name = "cuda",
.mem_type = UCS_MEMORY_TYPE_CUDA,
.init = ucx_perf_cuda_init,
.uct_alloc = uct_perf_cuda_alloc,
.uct_free = uct_perf_cuda_free,
.memcpy = ucx_perf_cuda_memcpy,
.memset = ucx_perf_cuda_memset
};

static ucx_perf_allocator_t cuda_managed_allocator = {
.name = "cuda-managed",
.mem_type = UCS_MEMORY_TYPE_CUDA_MANAGED,
.init = ucx_perf_cuda_init,
.uct_alloc = uct_perf_cuda_managed_alloc,
.uct_free = uct_perf_cuda_free,
.memcpy = ucx_perf_cuda_memcpy,
.memset = ucx_perf_cuda_memset
};

UCS_STATIC_INIT {
static ucx_perf_allocator_t cuda_allocator = {
.mem_type = UCS_MEMORY_TYPE_CUDA,
.init = ucx_perf_cuda_init,
.uct_alloc = uct_perf_cuda_alloc,
.uct_free = uct_perf_cuda_free,
.memcpy = ucx_perf_cuda_memcpy,
.memset = ucx_perf_cuda_memset
};
static ucx_perf_allocator_t cuda_managed_allocator = {
.mem_type = UCS_MEMORY_TYPE_CUDA_MANAGED,
.init = ucx_perf_cuda_init,
.uct_alloc = uct_perf_cuda_managed_alloc,
.uct_free = uct_perf_cuda_free,
.memcpy = ucx_perf_cuda_memcpy,
.memset = ucx_perf_cuda_memset
};

ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_CUDA] = &cuda_allocator;
ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_CUDA_MANAGED] = &cuda_managed_allocator;
ucx_perf_allocator_register(&cuda_allocator);
ucx_perf_allocator_register(&cuda_managed_allocator);
#if CUDART_VERSION >= 11020
ucx_perf_allocator_register(&cuda_async_allocator);
#endif
}

UCS_STATIC_CLEANUP {
ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_CUDA] = NULL;
ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_CUDA_MANAGED] = NULL;
#if CUDART_VERSION >= 11020
ucx_perf_allocator_unregister(&cuda_async_allocator);
#endif
ucx_perf_allocator_unregister(&cuda_managed_allocator);
ucx_perf_allocator_unregister(&cuda_allocator);
}
44 changes: 33 additions & 11 deletions src/tools/perf/lib/libperf.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,11 +83,26 @@ typedef struct {
} ucp_perf_flush_context_t;


const ucx_perf_allocator_t* ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_LAST];
const ucx_perf_allocator_t*
ucx_perf_allocators[UCX_PERF_ALLOCATOR_MAX];
unsigned ucx_perf_num_allocators;

const ucx_perf_device_dispatcher_t*
ucx_perf_mem_type_device_dispatchers[UCS_MEMORY_TYPE_LAST];

const ucx_perf_allocator_t *ucx_perf_allocator_by_name(const char *name)
{
unsigned i;

for (i = 0; i < ucx_perf_num_allocators; ++i) {
if (!strcmp(name, ucx_perf_allocators[i]->name)) {
return ucx_perf_allocators[i];
}
}

return NULL;
}

static const char *perf_iface_ops[] = {
[ucs_ilog2(UCT_IFACE_FLAG_AM_SHORT)] = "am short",
[ucs_ilog2(UCT_IFACE_FLAG_AM_BCOPY)] = "am bcopy",
Expand Down Expand Up @@ -2163,24 +2178,31 @@ ucx_perf_funcs_t ucx_perf_funcs[] = {
ucp_perf_test_dispatch, ucp_perf_thread_barrier}
};

static const ucx_perf_allocator_t *ucx_perf_params_allocator(
const ucx_perf_params_t *params, int is_send)
{
const char *alloc_name;

alloc_name = ucx_perf_mem_alloc_name(params, is_send);
return ucx_perf_allocator_by_name(alloc_name);
}

ucs_status_t ucx_perf_allocators_init(ucx_perf_context_t *perf,
const ucx_perf_params_t *params)
{
ucs_debug("set send allocator by send mem type %s",
ucs_memory_type_names[params->send_mem_type]);
perf->send_allocator = ucx_perf_mem_type_allocators[params->send_mem_type];

ucs_debug("set recv allocator by recv mem type %s",
ucs_memory_type_names[params->recv_mem_type]);
perf->recv_allocator = ucx_perf_mem_type_allocators[params->recv_mem_type];
perf->send_allocator = ucx_perf_params_allocator(params, 1);
perf->recv_allocator = ucx_perf_params_allocator(params, 0);

if ((perf->send_allocator == NULL) || (perf->recv_allocator == NULL)) {
ucs_error("Unsupported memory types %s<->%s",
ucs_memory_type_names[params->send_mem_type],
ucs_memory_type_names[params->recv_mem_type]);
ucs_error("Unsupported memory allocators %s<->%s",
ucx_perf_mem_alloc_name(params, 1),
ucx_perf_mem_alloc_name(params, 0));
return UCS_ERR_UNSUPPORTED;
}

ucs_debug("set send allocator %s", perf->send_allocator->name);
ucs_debug("set recv allocator %s", perf->recv_allocator->name);

if (perf->params.api == UCX_PERF_API_UCT) {
if (perf->send_allocator->mem_type != UCS_MEMORY_TYPE_HOST) {
ucs_diag("UCT tests also copy one-byte value from %s memory to "
Expand Down
Loading