Skip to content

UCS/DATASTRUCT: Add ucs_array_for_each_index #11853

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

Open
guy-ealey-morag wants to merge 3 commits into
base: master
Choose a base branch
from
Open

UCS/DATASTRUCT: Add ucs_array_for_each_index #11853

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
7a41231
UCS/DATASTRUCT: Add `ucs_array_for_each_index`
guy-ealey-morag Aug 28, 2026
7486592
UCS/DATASTRUCT: Add `expected_idx` to unit tests
guy-ealey-morag Aug 28, 2026
f74e7cf
TEST/UCS: Fix for_each tests
guy-ealey-morag Aug 28, 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
2 changes: 2 additions & 0 deletions .clang-format
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ ForEachMacros: ['UCS_STATIC_BITMAP_FOR_EACH_BIT',
'UCS_STATIC_BITMAP_FOR_EACH_BIT',
'FOR_EACH_ENTITY',
'ucs_carray_for_each',
'ucs_carray_for_each_index',
'kh_foreach',
'kh_foreach_key',
'kh_foreach_value',
Expand All @@ -81,6 +82,7 @@ ForEachMacros: ['UCS_STATIC_BITMAP_FOR_EACH_BIT',
'ucp_proto_perf_segment_foreach',
'ucp_proto_perf_segment_foreach_range',
'ucs_array_for_each',
'ucs_array_for_each_index',
'UCS_BITMAP_FOR_EACH_BIT',
'ucs_for_each_bit',
'ucs_for_each_submask',
Expand Down
4 changes: 1 addition & 3 deletions src/ucp/proto/proto_select.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,9 +75,7 @@ static ucs_status_t ucp_proto_thresholds_next_range(
ucs_dynamic_bitmap_reset_all(proto_mask);
ucs_dynamic_bitmap_init(&disabled_proto_mask);

for (proto_idx = 0; proto_idx < ucs_array_length(&proto_init->protocols);
++proto_idx) {
proto = &ucs_array_elem(&proto_init->protocols, proto_idx);
ucs_array_for_each_index(proto, proto_idx, &proto_init->protocols) {
range = ucp_proto_flat_perf_find_lb(proto->flat_perf, msg_length);
if (range == NULL) {
ucs_trace("skipping proto %s for msg_length %zu",
Expand Down
14 changes: 13 additions & 1 deletion src/ucs/datastruct/array.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/**
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020. ALL RIGHTS RESERVED.
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2020-2026. ALL RIGHTS RESERVED.
*
* See file LICENSE for terms.
*/
Expand Down Expand Up @@ -456,6 +456,18 @@ ucs_array_old_buffer_set_null(void **old_buffer_p)
ucs_array_length(_array))


/**
* Iterate over array elements and track the current index
*
* @param _elem Pointer variable to the current array element
* @param _idx Variable containing the current array index
* @param _array Array to iterate over
*/
#define ucs_array_for_each_index(_elem, _idx, _array) \
ucs_carray_for_each_index(_elem, _idx, ucs_array_begin(_array), \

@tvegas1 tvegas1 Aug 28, 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 we not add ucs_carray_for_each_index() macro and here use something like below which would have benefit of only relying on index to get a fresh elem pointer:

#define ucs_array_for_each_index(_elem, _idx, _array) \
    for ((_idx) = 0; \
         ((_idx) < ucs_array_length(_array)) && \
         (((_elem) = ucs_array_elem(_array, _idx)), 1); \
         ++(_idx))

@guy-ealey-morag guy-ealey-morag Aug 28, 2026 •
edited
Loading

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.

Using ucs_array_elem would add an unnecessary assert on every iteration, the fact that it's part of the loop let us skip this check safely.
Also it's simpler to rely directly on the carray variant like it's done in ucs_array_for_each.

@tvegas1 tvegas1 Aug 28, 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.

assert is already there long time and if needed can be removed for that case

@guy-ealey-morag guy-ealey-morag Aug 31, 2026 •
edited
Loading

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.

This implementation is a direct extension to the existing ucs_array_for_each and ucs_carray_for_each. It uses the same patterns, same assumptions and it works in a very similar way.

The idea is that anywhere that can use the existing ucs_array_for_each would be able to use ucs_array_for_each_index to also have the index available.

ucs_array_length(_array))


/* Internal helper function */
ucs_status_t ucs_array_grow(void **buffer_p, size_t *capacity_p,
size_t min_capacity, size_t max_capacity,
Expand Down
6 changes: 2 additions & 4 deletions src/ucs/datastruct/callbackq.c
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/**
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2016. ALL RIGHTS RESERVED.
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
* Copyright (C) ARM Ltd. 2016-2017. ALL RIGHTS RESERVED.
*
* See file LICENSE for terms.
Expand Down Expand Up @@ -295,9 +295,7 @@ static void ucs_callbackq_spill_elems_purge(ucs_callbackq_t *cbq)
* always be equal to dst_idx, so nothing will be actually copied/moved.
*/
dst_idx = 0;
for (src_idx = 0; src_idx < ucs_array_length(&priv->spill_elems);
++src_idx) {
src_elem = &ucs_array_elem(&priv->spill_elems, src_idx);
ucs_array_for_each_index(src_elem, src_idx, &priv->spill_elems) {
if (src_elem->id != UCS_CALLBACKQ_ID_NULL) {
ucs_assert(dst_idx <= src_idx);
if (dst_idx != src_idx) {
Expand Down
8 changes: 7 additions & 1 deletion src/ucs/sys/compiler_def.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/**
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2017. ALL RIGHTS RESERVED.
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
* Copyright (C) UT-Battelle, LLC. 2015. ALL RIGHTS RESERVED.
* Copyright (C) Arm, Ltd. 2021. ALL RIGHTS RESERVED.
* Copyright (C) Advanced Micro Devices, Inc. 2024. ALL RIGHTS RESERVED.
Expand Down Expand Up @@ -250,6 +250,12 @@
#define ucs_carray_for_each(_elem, _array, _length) \
for ((_elem) = (_array); (_elem) < ((_array) + (_length)); ++(_elem))

/*
* Iterate over all elements of a C-array and track the current index
*/
#define ucs_carray_for_each_index(_elem, _idx, _array, _length) \
for ((_idx) = 0, (_elem) = (_array); (_idx) < (_length); ++(_idx), ++(_elem))

/*
* Swap two variables values
*/
Expand Down
49 changes: 48 additions & 1 deletion test/gtest/ucs/test_datatype.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/**
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2014. ALL RIGHTS RESERVED.
* Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
* Copyright (C) UT-Battelle, LLC. 2014. ALL RIGHTS RESERVED.
* Copyright (C) Huawei Technologies Co., Ltd. 2020. ALL RIGHTS RESERVED.
* See file LICENSE for terms.
Expand Down Expand Up @@ -1115,6 +1115,53 @@ class test_array : public test_datatype {
void cleanup_array_of_linked_lists(test_list_links_array_t *test_array);
};

UCS_TEST_F(test_array, carray_for_each_index) {
const int values[] = {3, 5, 7};
const int *elem;
size_t idx, expected_idx;

expected_idx = 0;
ucs_carray_for_each_index(elem, idx, values,
ucs_static_array_size(values)) {
EXPECT_EQ(&values[idx], elem);
EXPECT_EQ(values[idx], *elem);
EXPECT_EQ(expected_idx++, idx);
}
EXPECT_EQ(expected_idx, idx);
EXPECT_EQ(ucs_static_array_size(values), idx);
}

UCS_TEST_F(test_array, array_for_each_index) {
constexpr size_t NUM_ELEMENTS = 3;
test_1int_t test_array;
int *elem;
size_t idx, expected_idx;

ucs_array_init_dynamic(&test_array);

ucs_array_for_each_index(elem, idx, &test_array) {
ADD_FAILURE() << "iterated over an empty array";
}
EXPECT_EQ(0u, idx);

for (idx = 0; idx < NUM_ELEMENTS; ++idx) {
*ucs_array_append(&test_array,
FAIL()) = static_cast<int>(idx * idx + 1);
}
EXPECT_EQ(NUM_ELEMENTS, ucs_array_length(&test_array));

expected_idx = 0;
ucs_array_for_each_index(elem, idx, &test_array) {
EXPECT_EQ(&ucs_array_elem(&test_array, idx), elem);
EXPECT_EQ(idx * idx + 1, static_cast<size_t>(*elem));
EXPECT_EQ(expected_idx++, idx);
}
EXPECT_EQ(expected_idx, idx);
EXPECT_EQ(ucs_array_length(&test_array), idx);

ucs_array_cleanup_dynamic(&test_array);
}

/* generate a list of numbers in a certain size */
void test_array::generate_linked_list(int size, simple_elem_t *head)
{
Expand Down