Skip to content

TOOLS/PERF: Honor configured device id in ZE and ROCm allocators #11993

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

Closed
yafshar wants to merge 3 commits into from
Closed

TOOLS/PERF: Honor configured device id in ZE and ROCm allocators #11993

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
fd4baa9
TOOLS/PERF: Honor configured device id in ZE and ROCm allocators
yafshar Sep 25, 2026
1d5d3b0
TOOLS/PERF: Simplify GPU index helper and ZE index types
yafshar Sep 29, 2026
d9b7556
TOOLS/PERF: Use int GPU index in ZE command list helper
yafshar Sep 29, 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
17 changes: 5 additions & 12 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 @@ -20,26 +20,19 @@

static ucs_status_t ucx_perf_cuda_init(ucx_perf_context_t *perf)
{
ucs_status_t status;
unsigned group_index;
int num_gpus;
int gpu_index;

group_index = rte_call(perf, group_index);

CUDA_CALL_RET(UCS_ERR_NO_DEVICE, cudaGetDeviceCount, &num_gpus);
if (num_gpus == 0) {
ucs_error("no cuda devices available");
return UCS_ERR_NO_DEVICE;
}

gpu_index = (group_index == 0) ? perf->params.recv_device.device_id :
perf->params.send_device.device_id;
if (gpu_index == UCX_PERF_MEM_DEV_DEFAULT) {
gpu_index = group_index % num_gpus;
} else if (gpu_index >= num_gpus) {
ucs_error("Illegal cuda device %d number of devices %d", gpu_index,
num_gpus);
return UCS_ERR_NO_DEVICE;
status = ucx_perf_get_gpu_index(perf, group_index, num_gpus, "cuda",

@tvegas1 tvegas1 Sep 29, 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.

-a currently also switches the run into device api mode: it adds UCP_FEATURE_DEVICE (libperf.c:958), moves thread_count into device_thread_count (perftest_params.c:866), and routes ucp dispatch through ucx_perf_mem_type_device_dispatchers[] (ucp_tests.cc:1209), which is registered for cuda only. So using -a to pick a ze/rocm device also pulls in device mode semantics that those allocators do not support.

Can we add a separate knob to enable device mode, and keep -a for device type and id only, as its help text already says? Then cuda could also select a device in a regular run, which is not possible today.

@yafshar yafshar Sep 29, 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.

Agreed, -a currently mixes device selection with device-API mode, and this is pre-existing: on master, ucx_perftest -t tag_bw -m cuda -a cuda:1 also fails with Invalid test case, so CUDA can't select a device in a regular UCP run either. Adding a separate knob changes CLI semantics that CI relies on (test_jenkins.sh uses -a cuda:0 to enable device mode), so I'd like to keep this PR limited to honoring the id in the allocators and do the knob in a follow-up PR. With this PR alone, the ZE/ROCm device id takes effect for UCT tests.

@tvegas1 tvegas1 Sep 30, 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.

-a today is only usable for the device feature: it implies UCP_FEATURE_DEVICE (libperf.c:958), moves thread_count into device_thread_count (perftest_params.c:866) and routes ucp dispatch through ucx_perf_mem_type_device_dispatchers[] (ucp_tests.cc:1209), which is registered for cuda only. So it would make ze/rocm honor a device id only in a mode those allocators do not support, and no allocator supports placement outside that mode today.

Can we instead move the device id to -m, with the same <type>:<dev-id> syntax, and keep -a without argument to only enable the device feature? -m already selects the allocator that places the buffers, and it avoids the strange -m ze-device -a ze-device:0.

@tvegas1 tvegas1 Sep 30, 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.

actually started #12003 that should cover this PR, if you could confirm?

@yafshar yafshar Sep 30, 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.

Yes, #12003 subsumes this PR: the shared device-index helper, ZE/ROCm honoring an explicit id, the range checks and the ROCm zero-device guard, and it also separates device mode from device selection, which answers the open -a threads and the type-mismatch question here. Since it targets master independently and moves the id from -a <type:id> to -m <allocator:id>, I'll treat it as a replacement and close this PR once #12003 is merged. I left a few notes on #12003; the main one is preserving the GPU-0 pinning in the CUDA device CI job.

&gpu_index);
if (status != UCS_OK) {
return status;
}

CUDA_CALL_RET(UCS_ERR_NO_DEVICE, cudaSetDevice, gpu_index);
Expand Down
31 changes: 31 additions & 0 deletions src/tools/perf/lib/libperf_int.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -337,6 +337,37 @@ ucx_perf_mem_alloc_name(const ucx_perf_params_t *params, int is_send)
return ucs_memory_type_names[mem_type];
}

/*
* Select the GPU to be used by this process: the device id given by the
* '-a <type:dev-id>' option if it was specified, otherwise spread the
* processes over the available devices according to their group index.
*/
static inline ucs_status_t
ucx_perf_get_gpu_index(const ucx_perf_context_t *perf, unsigned group_index,
int num_gpus, const char *dev_type, int *gpu_index_p)
{
int gpu_index;

if (num_gpus <= 0) {
ucs_error("no %s devices available", dev_type);
return UCS_ERR_NO_DEVICE;
}

/* group index 0 is the receiver side */
gpu_index = (group_index == 0) ? perf->params.recv_device.device_id :

@svc-nvidia-pr-review svc-nvidia-pr-review Sep 29, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

minor: should we also check that the -a device type matches the allocator? Only device_id is read here (without looking at send_device.mem_type/recv_device.mem_type), so a device id given for one accelerator type is applied to whichever allocator happens to run init() — e.g. -m rocm -a cuda:1 silently selects rocm device 1 instead of reporting a mismatch. Low impact (mostly a mistyped-argument case) and partly pre-existing for cuda, but this PR is the point where the code becomes shared across three backends. This overlaps the still-open thread about -a semantics for rocm/ze, so it may be simplest to answer it there.

perf->params.send_device.device_id;
if (gpu_index == UCX_PERF_MEM_DEV_DEFAULT) {
gpu_index = group_index % num_gpus;
} else if (gpu_index >= num_gpus) {
ucs_error("illegal %s device %d, number of devices %d", dev_type,
gpu_index, num_gpus);
return UCS_ERR_NO_DEVICE;
}

*gpu_index_p = gpu_index;
return UCS_OK;
}

static UCS_F_ALWAYS_INLINE int ucx_perf_context_done(ucx_perf_context_t *perf)
{
return ucs_unlikely((perf->current.iters >= perf->max_iter) ||
Expand Down
7 changes: 6 additions & 1 deletion src/tools/perf/rocm/rocm_alloc.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
static ucs_status_t ucx_perf_rocm_init(ucx_perf_context_t *perf)
{
hipError_t ret;
ucs_status_t status;
unsigned group_index;
int num_gpus;
int gpu_index;
Expand All @@ -29,7 +30,11 @@ static ucs_status_t ucx_perf_rocm_init(ucx_perf_context_t *perf)
return UCS_ERR_NO_DEVICE;
}

gpu_index = group_index % num_gpus;
status = ucx_perf_get_gpu_index(perf, group_index, num_gpus, "rocm",
Comment thread
yafshar marked this conversation as resolved.
Resolved
&gpu_index);
if (status != UCS_OK) {
return status;
}

ret = hipSetDevice(gpu_index);
if (ret != hipSuccess) {
Expand Down
22 changes: 14 additions & 8 deletions src/tools/perf/ze/ze_alloc.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,9 @@ static pthread_once_t ze_init_once = PTHREAD_ONCE_INIT;
static ze_context_handle_t gpu_context;
static ze_device_handle_t gpu_devices[ZE_PERF_MAX_DEVICES];
static ze_driver_handle_t gpu_driver;
static unsigned gpu_count;
static int gpu_count;

static __thread unsigned tls_gpu_index;
static __thread int tls_gpu_index;
static __thread ze_command_list_handle_t tls_cmdlist;


Expand All @@ -42,7 +42,7 @@ static void ucx_perf_ze_destroy_tls_cmdlist(void)
}

static ze_result_t
ucx_perf_ze_create_tls_cmdlist(unsigned gpu_idx,
ucx_perf_ze_create_tls_cmdlist(int gpu_idx,
const ze_command_queue_desc_t *cmdq_desc)
{
return zeCommandListCreateImmediate(gpu_context, gpu_devices[gpu_idx],
Expand Down Expand Up @@ -109,23 +109,29 @@ static ucs_status_t ucx_perf_ze_init(ucx_perf_context_t *perf)
.priority = ZE_COMMAND_QUEUE_PRIORITY_NORMAL,
};
ze_result_t ret;
unsigned group_index, i;
ucs_status_t status;
unsigned group_index;
int gpu_index;

ret = ze_init_devices();
if (ret != ZE_RESULT_SUCCESS) {
return UCS_ERR_NO_DEVICE;
}
Comment thread
yafshar marked this conversation as resolved.
Resolved

group_index = rte_call(perf, group_index);
i = group_index % gpu_count;
status = ucx_perf_get_gpu_index(perf, group_index, gpu_count, "ze",
&gpu_index);
if (status != UCS_OK) {
return status;
}

if ((tls_cmdlist != NULL) && (tls_gpu_index != i)) {
if ((tls_cmdlist != NULL) && (tls_gpu_index != gpu_index)) {
ucx_perf_ze_destroy_tls_cmdlist();
}

tls_gpu_index = i;
tls_gpu_index = gpu_index;
if (tls_cmdlist == NULL) {
ret = ucx_perf_ze_create_tls_cmdlist(i, &cmdq_desc);
ret = ucx_perf_ze_create_tls_cmdlist(gpu_index, &cmdq_desc);
if (ret != ZE_RESULT_SUCCESS) {
return UCS_ERR_NO_DEVICE;
}
Expand Down