diff --git a/.clang-format b/.clang-format index 9f1ff816891..8fa2aa6a80c 100644 --- a/.clang-format +++ b/.clang-format @@ -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', @@ -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', diff --git a/src/ucp/proto/proto_select.c b/src/ucp/proto/proto_select.c index ad19d6fb3f3..ba37f7bab73 100644 --- a/src/ucp/proto/proto_select.c +++ b/src/ucp/proto/proto_select.c @@ -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", diff --git a/src/ucs/datastruct/array.h b/src/ucs/datastruct/array.h index c3e39c9ca38..7986b356e2c 100644 --- a/src/ucs/datastruct/array.h +++ b/src/ucs/datastruct/array.h @@ -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. */ @@ -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), \ + 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, diff --git a/src/ucs/datastruct/callbackq.c b/src/ucs/datastruct/callbackq.c index b97b093c276..1100517e4a9 100644 --- a/src/ucs/datastruct/callbackq.c +++ b/src/ucs/datastruct/callbackq.c @@ -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. @@ -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) { diff --git a/src/ucs/sys/compiler_def.h b/src/ucs/sys/compiler_def.h index 1836dd0a3ed..82907725710 100644 --- a/src/ucs/sys/compiler_def.h +++ b/src/ucs/sys/compiler_def.h @@ -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. @@ -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 */ diff --git a/test/gtest/ucs/test_datatype.cc b/test/gtest/ucs/test_datatype.cc index f1ec148c05b..63df4393776 100644 --- a/test/gtest/ucs/test_datatype.cc +++ b/test/gtest/ucs/test_datatype.cc @@ -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. @@ -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(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(*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) {