Skip to content

Commit a5c4b47

Browse files
mshanthagitmshanthagit
committed
UCT/ROCM: rebase, address comments
1 parent cc8d39d commit a5c4b47

5 files changed

Lines changed: 29 additions & 22 deletions

File tree

‎src/ucp/wireup/select.c‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2474,9 +2474,13 @@ ucp_wireup_add_device_lanes(const ucp_wireup_select_params_t *select_params,
24742474
const unsigned ep_init_flags = ucp_wireup_ep_init_flags(select_params,
24752475
select_ctx);
24762476
ucp_wireup_select_bw_info_t bw_info = {};
2477-
ucp_tl_bitmap_t mem_type_tl_bitmap;
2478-
ucs_memory_type_t mem_type;
2477+
const uint64_t mem_type_bitmaps[] = {UCS_BIT(UCS_MEMORY_TYPE_CUDA),
2478+
UCS_BIT(UCS_MEMORY_TYPE_CUDA) |
2479+
UCS_BIT(UCS_MEMORY_TYPE_HOST),
2480+
UCS_BIT(UCS_MEMORY_TYPE_ROCM)};
24792481
int found_lane = 0;
2482+
size_t i;
2483+
ucp_tl_bitmap_t mem_type_tl_bitmap;
24802484

24812485
if (!context->config.ext.proto_enable ||
24822486
(ep_init_flags &
@@ -2506,8 +2510,9 @@ ucp_wireup_add_device_lanes(const ucp_wireup_select_params_t *select_params,
25062510
*/
25072511
bw_info.max_lanes = ucp_wireup_bw_max_lanes(select_params);
25082512

2509-
ucs_for_each_bit(mem_type, UCP_DEVICE_MEM_TYPES) {
2510-
ucp_wireup_memaccess_bitmap(context, mem_type, &mem_type_tl_bitmap);
2513+
for (i = 0; i < ucs_static_array_size(mem_type_bitmaps); ++i) {
2514+
ucp_wireup_memaccess_bitmap(context, mem_type_bitmaps[i],
2515+
&mem_type_tl_bitmap);
25112516
found_lane |= ucp_wireup_add_bw_lanes(select_params, &bw_info,
25122517
mem_type_tl_bitmap, UCP_NULL_LANE,
25132518
select_ctx, 0);

‎src/ucs/sys/device_code.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@
3333
/* Number of threads in a warp */
3434
#define UCS_DEVICE_NUM_THREADS_IN_WARP 32
3535

36+
/* Check if _n is aligned to a power-of-2 boundary _p */
37+
#define UCT_IPC_IS_ALIGNED_POW2(_n, _p) (!((_n) & ((_p) - 1)))
38+
3639

3740
/* nvcc does not provide __builtin_ia32_prefetch used by GCC's x86 intrinsic headers.
3841
Redirect to the generic __builtin_prefetch so those headers compile. */

‎src/uct/cuda/cuda_ipc/cuda_ipc.cuh‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
#include <ucs/type/status.h>
1414
#include <cuda/atomic>
1515

16-
#define UCT_CUDA_IPC_IS_ALIGNED_POW2(_n, _p) (!((_n) & ((_p) - 1)))
1716
#define UCT_CUDA_IPC_WARP_SIZE 32
1817
#define UCT_CUDA_IPC_COPY_LOOP_UNROLL 8
1918

@@ -129,8 +128,8 @@ void uct_cuda_ipc_copy_level<UCS_DEVICE_LEVEL_WARP>(void *dst, const void *src,
129128
auto d1 = reinterpret_cast<char *>(dst);
130129

131130
/* 16B-aligned fast path using vec4 */
132-
if (UCT_CUDA_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec4)) &&
133-
UCT_CUDA_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec4))) {
131+
if (UCT_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec4)) &&
132+
UCT_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec4))) {
134133
const vec4 *s4 = reinterpret_cast<const vec4*>(s1);
135134
vec4 *d4 = reinterpret_cast<vec4*>(d1);
136135
size_t n4 = len / sizeof(vec4);
@@ -149,8 +148,8 @@ void uct_cuda_ipc_copy_level<UCS_DEVICE_LEVEL_WARP>(void *dst, const void *src,
149148
}
150149

151150
/* 8B-aligned fast path using vec2 */
152-
if (UCT_CUDA_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec2)) &&
153-
UCT_CUDA_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec2))) {
151+
if (UCT_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec2)) &&
152+
UCT_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec2))) {
154153
const vec2 *s2 = reinterpret_cast<const vec2*>(s1);
155154
vec2 *d2 = reinterpret_cast<vec2*>(d1);
156155
size_t n2 = len / sizeof(vec2);
@@ -186,8 +185,8 @@ void uct_cuda_ipc_copy_level<UCS_DEVICE_LEVEL_BLOCK>(void *dst, const void *src,
186185
int warp, num_warps, idx;
187186
size_t num_lines;
188187

189-
if (UCT_CUDA_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec4)) &&
190-
UCT_CUDA_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec4))) {
188+
if (UCT_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec4)) &&
189+
UCT_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec4))) {
191190
vec4 tmp[UCT_CUDA_IPC_COPY_LOOP_UNROLL];
192191
warp = threadIdx.x / UCT_CUDA_IPC_WARP_SIZE;
193192
num_warps = blockDim.x / UCT_CUDA_IPC_WARP_SIZE;
@@ -232,8 +231,8 @@ void uct_cuda_ipc_copy_level<UCS_DEVICE_LEVEL_BLOCK>(void *dst, const void *src,
232231
}
233232

234233
/* If not 16B-aligned, try 8B-aligned fast path using vec2 */
235-
if (UCT_CUDA_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec2)) &&
236-
UCT_CUDA_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec2))) {
234+
if (UCT_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec2)) &&
235+
UCT_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec2))) {
237236
const vec2 *s2;
238237
vec2 *d2;
239238
vec2 tmp2[UCT_CUDA_IPC_COPY_LOOP_UNROLL];

‎src/uct/rocm/ipc/rocm_ipc.h‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
#include <ucs/sys/device_code.h>
1212
#include <ucs/type/status.h>
1313

14-
#define UCT_ROCM_IPC_IS_ALIGNED_POW2(_n, _p) (!((_n) & ((_p)-1)))
1514

1615
/* Dynamically detect wavefront size using compiler builtin. */
1716
#if __has_builtin(__builtin_amdgcn_wavefrontsize)
@@ -130,8 +129,8 @@ uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_WARP>(void *dst, const void *src,
130129
auto d1 = reinterpret_cast<char*>(dst);
131130

132131
/* 16B-aligned fast path using vec4 */
133-
if (UCT_ROCM_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec4)) &&
134-
UCT_ROCM_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec4))) {
132+
if (UCT_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec4)) &&
133+
UCT_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec4))) {
135134
const vec4 *s4 = reinterpret_cast<const vec4*>(s1);
136135
vec4 *d4 = reinterpret_cast<vec4*>(d1);
137136
size_t n4 = len / sizeof(vec4);
@@ -151,8 +150,8 @@ uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_WARP>(void *dst, const void *src,
151150
}
152151

153152
/* 8B-aligned fast path using vec2 */
154-
if (UCT_ROCM_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec2)) &&
155-
UCT_ROCM_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec2))) {
153+
if (UCT_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec2)) &&
154+
UCT_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec2))) {
156155
const vec2 *s2 = reinterpret_cast<const vec2*>(s1);
157156
vec2 *d2 = reinterpret_cast<vec2*>(d1);
158157
size_t n2 = len / sizeof(vec2);
@@ -187,8 +186,8 @@ uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_BLOCK>(void *dst, const void *src,
187186
auto s1 = reinterpret_cast<const char*>(src);
188187
auto d1 = reinterpret_cast<char*>(dst);
189188

190-
if (UCT_ROCM_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec4)) &&
191-
UCT_ROCM_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec4))) {
189+
if (UCT_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec4)) &&
190+
UCT_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec4))) {
192191
const vec4 *s4 = reinterpret_cast<const vec4*>(s1);
193192
vec4 *d4 = reinterpret_cast<vec4*>(d1);
194193
size_t num_lines = len / sizeof(vec4);
@@ -208,8 +207,8 @@ uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_BLOCK>(void *dst, const void *src,
208207
}
209208

210209
/* 8B-aligned fast path using vec2 */
211-
if (UCT_ROCM_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec2)) &&
212-
UCT_ROCM_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec2))) {
210+
if (UCT_IPC_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec2)) &&
211+
UCT_IPC_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec2))) {
213212
const vec2 *s2 = reinterpret_cast<const vec2*>(s1);
214213
vec2 *d2 = reinterpret_cast<vec2*>(d1);
215214
size_t num_lines = len / sizeof(vec2);

‎src/uct/rocm/ipc/rocm_ipc_iface.c‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ static ucs_status_t uct_rocm_ipc_iface_query(uct_iface_h tl_iface,
125125
UCT_IFACE_FLAG_PUT_ZCOPY | UCT_IFACE_FLAG_PENDING |
126126
UCT_IFACE_FLAG_CONNECT_TO_IFACE |
127127
UCT_IFACE_FLAG_DEVICE_EP;
128+
iface_attr->ctl_device = UCS_SYS_DEVICE_ID_UNKNOWN;
128129

129130
iface_attr->latency = ucs_linear_func_make(iface->config.latency, 0);
130131
iface_attr->bandwidth.dedicated = 0;

0 commit comments

Comments
 (0)