Skip to content

Commit a4ae783

Browse files
mshanthagitmshanthagit
committed
UCT/ROCM: address review comments
1 parent 5f19248 commit a4ae783

6 files changed

Lines changed: 71 additions & 109 deletions

File tree

‎config/hip.am‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ SUFFIXES = .hip
77

88
HIPCC ?= hipcc
99

10-
HIPCC_CMD = $(HIPCC) -DHAVE_CONFIG_H -DUCT_DEVICE_CODE_HIP -fPIE -I$(top_srcdir)/src -I$(top_builddir)/src $(BASE_CXXFLAGS) $(CXXFLAGS) $(HIP_CPPFLAGS) $(HIP_CXXFLAGS) $(HIPCC_EXTRA_FLAGS) -Wno-c++20-extensions -c $< -MT $@ -MF $(DEPDIR)/hip/$@.d -MMD -o $@
10+
HIPCC_CMD = $(HIPCC) $(DEFS) -DUCT_DEVICE_CODE_HIP -fPIE -I$(top_srcdir)/src -I$(top_builddir)/src $(BASE_CXXFLAGS) $(CXXFLAGS) $(HIP_CPPFLAGS) $(HIP_CXXFLAGS) $(HIPCC_EXTRA_FLAGS) -Wno-c++20-extensions -c $< -MT $@ -MF $(DEPDIR)/hip/$@.d -MMD -o $@
1111
HIPCC_LT_CMD = $(LIBTOOL) --tag=CXX --mode=compile $(HIPCC_CMD)
1212

1313
define hipcc-build

‎src/ucp/core/ucp_device.c‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -378,7 +378,7 @@ ucp_device_local_mem_list_create(const ucp_device_mem_list_params_t *params,
378378
ucp_device_local_mem_list_h *mem_list_h)
379379
{
380380
const ucp_worker_h worker = UCS_PARAM_VALUE(UCP_DEVICE_MEM_LIST_PARAMS_FIELD,
381-
params, worker, WORKER, NULL);
381+
params, worker, WORKER, NULL);
382382
ucs_memory_type_t export_mem_type;
383383
ucs_sys_device_t local_sys_dev;
384384
uct_allocated_memory_t mem;
@@ -561,7 +561,7 @@ ucp_device_remote_mem_list_fill(const ucp_device_mem_list_elem_t *ucp_element,
561561
}
562562

563563
static ucs_status_t ucp_device_remote_mem_list_create_handle(
564-
const ucp_device_mem_list_params_t *params, ucp_ep_h ep,
564+
const ucp_device_mem_list_params_t *params, const ucp_ep_h ep,
565565
const ucs_memory_type_t mem_type, uct_allocated_memory_t *mem,
566566
const ucs_sys_device_t local_sys_dev)
567567
{

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

Lines changed: 45 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -76,11 +76,11 @@ uct_rocm_ipc_map_remote(const uct_rocm_ipc_device_mem_element_t *elem,
7676
elem->mapped_offset);
7777
}
7878

79-
/* System-wide atomic increment */
79+
/* System-wide atomic add */
8080
__device__ static inline void
81-
uct_rocm_ipc_atomic_inc(uint64_t *dst, uint64_t inc_value)
81+
uct_rocm_ipc_atomic_add(uint64_t *dst, uint64_t add_value)
8282
{
83-
atomicAdd_system((unsigned long long*)dst, (unsigned long long)inc_value);
83+
atomicAdd_system((unsigned long long*)dst, (unsigned long long)add_value);
8484
__threadfence_system();
8585
}
8686

@@ -101,98 +101,24 @@ __device__ static inline void uct_rocm_ipc_level_sync()
101101
}
102102
}
103103

104-
/* Copy routines for different parallelism levels */
105-
template<ucs_device_level_t level>
106-
__device__ void uct_rocm_ipc_copy_level(void *dst, const void *src, size_t len);
107-
108-
/* Thread-level copy */
109-
template<>
110-
__device__ inline void
111-
uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_THREAD>(void *dst, const void *src,
112-
size_t len)
113-
{
114-
memcpy(dst, src, len);
115-
}
116-
117-
/* Wavefront-level copy (64 threads) */
118-
template<>
119-
__device__ inline void
120-
uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_WARP>(void *dst, const void *src,
121-
size_t len)
104+
/* Shared strided copy used by warp- and block-level specializations */
105+
__device__ static inline void
106+
uct_rocm_ipc_copy_strided(void *dst, const void *src, size_t len,
107+
unsigned lane_id, unsigned num_lanes)
122108
{
123109
using vec4 = int4;
124110
using vec2 = int2;
125-
unsigned int lane_id, num_lanes;
126-
127-
uct_rocm_ipc_get_lane<UCS_DEVICE_LEVEL_WARP>(lane_id, num_lanes);
128111
auto s1 = reinterpret_cast<const char*>(src);
129112
auto d1 = reinterpret_cast<char*>(dst);
130113

131114
/* 16B-aligned fast path using vec4 */
132-
if (UCS_DEVICE_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec4)) &&
133-
UCS_DEVICE_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec4))) {
134-
const vec4 *s4 = reinterpret_cast<const vec4*>(s1);
135-
vec4 *d4 = reinterpret_cast<vec4*>(d1);
136-
size_t n4 = len / sizeof(vec4);
137-
138-
for (size_t i = lane_id; i < n4; i += num_lanes) {
139-
vec4 v = uct_rocm_ipc_ld_global_cg(s4 + i);
140-
uct_rocm_ipc_st_global_cg(d4 + i, v);
141-
}
142-
143-
len = len - n4 * sizeof(vec4);
144-
if (len == 0) {
145-
return;
146-
}
147-
148-
s1 = reinterpret_cast<const char*>(s4 + n4);
149-
d1 = reinterpret_cast<char*>(d4 + n4);
150-
}
151-
152-
/* 8B-aligned fast path using vec2 */
153-
if (UCS_DEVICE_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec2)) &&
154-
UCS_DEVICE_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec2))) {
155-
const vec2 *s2 = reinterpret_cast<const vec2*>(s1);
156-
vec2 *d2 = reinterpret_cast<vec2*>(d1);
157-
size_t n2 = len / sizeof(vec2);
158-
159-
for (size_t i = lane_id; i < n2; i += num_lanes) {
160-
vec2 v2 = uct_rocm_ipc_ld_global_cg(s2 + i);
161-
uct_rocm_ipc_st_global_cg(d2 + i, v2);
162-
}
163-
164-
len = len - n2 * sizeof(vec2);
165-
if (len == 0) {
166-
return;
167-
}
168-
169-
s1 = reinterpret_cast<const char*>(s2 + n2);
170-
d1 = reinterpret_cast<char*>(d2 + n2);
171-
}
172-
173-
/* Byte tail */
174-
for (size_t i = lane_id; i < len; i += num_lanes) {
175-
d1[i] = s1[i];
176-
}
177-
}
178-
179-
template<>
180-
__device__ inline void
181-
uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_BLOCK>(void *dst, const void *src,
182-
size_t len)
183-
{
184-
using vec4 = int4;
185-
using vec2 = int2;
186-
auto s1 = reinterpret_cast<const char*>(src);
187-
auto d1 = reinterpret_cast<char*>(dst);
188-
189115
if (UCS_DEVICE_IS_ALIGNED_POW2((intptr_t)s1, sizeof(vec4)) &&
190116
UCS_DEVICE_IS_ALIGNED_POW2((intptr_t)d1, sizeof(vec4))) {
191117
const vec4 *s4 = reinterpret_cast<const vec4*>(s1);
192118
vec4 *d4 = reinterpret_cast<vec4*>(d1);
193119
size_t num_lines = len / sizeof(vec4);
194120

195-
for (size_t line = threadIdx.x; line < num_lines; line += blockDim.x) {
121+
for (size_t line = lane_id; line < num_lines; line += num_lanes) {
196122
vec4 v = uct_rocm_ipc_ld_global_cg(s4 + line);
197123
uct_rocm_ipc_st_global_cg(d4 + line, v);
198124
}
@@ -213,7 +139,7 @@ uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_BLOCK>(void *dst, const void *src,
213139
vec2 *d2 = reinterpret_cast<vec2*>(d1);
214140
size_t num_lines = len / sizeof(vec2);
215141

216-
for (size_t line = threadIdx.x; line < num_lines; line += blockDim.x) {
142+
for (size_t line = lane_id; line < num_lines; line += num_lanes) {
217143
vec2 v2 = uct_rocm_ipc_ld_global_cg(s2 + line);
218144
uct_rocm_ipc_st_global_cg(d2 + line, v2);
219145
}
@@ -228,11 +154,45 @@ uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_BLOCK>(void *dst, const void *src,
228154
}
229155

230156
/* Byte tail */
231-
for (size_t line = threadIdx.x; line < len; line += blockDim.x) {
157+
for (size_t line = lane_id; line < len; line += num_lanes) {
232158
d1[line] = s1[line];
233159
}
234160
}
235161

162+
/* Copy routines for different parallelism levels */
163+
template<ucs_device_level_t level>
164+
__device__ void uct_rocm_ipc_copy_level(void *dst, const void *src, size_t len);
165+
166+
/* Thread-level copy */
167+
template<>
168+
__device__ inline void
169+
uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_THREAD>(void *dst, const void *src,
170+
size_t len)
171+
{
172+
memcpy(dst, src, len);
173+
}
174+
175+
/* Warp- and block-level copy: strided iteration across lanes */
176+
template<>
177+
__device__ inline void
178+
uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_WARP>(void *dst, const void *src,
179+
size_t len)
180+
{
181+
unsigned int lane_id, num_lanes;
182+
uct_rocm_ipc_get_lane<UCS_DEVICE_LEVEL_WARP>(lane_id, num_lanes);
183+
uct_rocm_ipc_copy_strided(dst, src, len, lane_id, num_lanes);
184+
}
185+
186+
template<>
187+
__device__ inline void
188+
uct_rocm_ipc_copy_level<UCS_DEVICE_LEVEL_BLOCK>(void *dst, const void *src,
189+
size_t len)
190+
{
191+
unsigned int lane_id, num_lanes;
192+
uct_rocm_ipc_get_lane<UCS_DEVICE_LEVEL_BLOCK>(lane_id, num_lanes);
193+
uct_rocm_ipc_copy_strided(dst, src, len, lane_id, num_lanes);
194+
}
195+
236196
/* Grid-level copy - not implemented */
237197
template<>
238198
__device__ inline void
@@ -278,7 +238,7 @@ __device__ ucs_status_t uct_rocm_ipc_ep_atomic_add(
278238
if (lane_id == 0) {
279239
mapped_rem_addr = reinterpret_cast<uint64_t*>(
280240
uct_rocm_ipc_map_remote(rocm_ipc_mem_element, remote_address));
281-
uct_rocm_ipc_atomic_inc(mapped_rem_addr, inc_value);
241+
uct_rocm_ipc_atomic_add(mapped_rem_addr, inc_value);
282242
}
283243

284244
uct_rocm_ipc_level_sync<level>();

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

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -240,18 +240,15 @@ ucs_status_t uct_rocm_ipc_component_init_cache(void)
240240
ucs_status_t status;
241241

242242
UCS_INIT_ONCE(&cache_init_once) {
243-
if (uct_rocm_ipc_component.ipc_cache == NULL) {
244-
status = uct_rocm_ipc_create_cache(&uct_rocm_ipc_component.ipc_cache,
245-
"rocm_ipc_component");
246-
if (status != UCS_OK) {
247-
ucs_error("Failed to create ROCm IPC component cache: %s",
248-
ucs_status_string(status));
249-
pthread_mutex_unlock(&uct_rocm_ipc_component.lock);
250-
return status;
251-
}
252-
253-
ucs_debug("ROCm IPC component cache initialized");
243+
status = uct_rocm_ipc_create_cache(&uct_rocm_ipc_component.ipc_cache,
244+
"rocm_ipc_component");
245+
if (status != UCS_OK) {
246+
ucs_error("Failed to create ROCm IPC component cache: %s",
247+
ucs_status_string(status));
248+
return status;
254249
}
250+
251+
ucs_debug("ROCm IPC component cache initialized");
255252
}
256253

257254
return UCS_OK;

‎test/gtest/uct/rocm/test_kernels_uct.hip‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,13 @@
1414

1515
#define UCS_DEVICE_LEVEL_EXEC_ID 1
1616

17-
#define UCS_DEVICE_LEVEL_EXEC_SELECT(scope_ok, count, id) \
18-
((scope_ok) ? (((count) > UCS_DEVICE_LEVEL_EXEC_ID) ? \
19-
((id) == UCS_DEVICE_LEVEL_EXEC_ID) : true) : false)
17+
static __device__ inline bool
18+
ucs_device_level_exec_select(bool scope_ok, unsigned count, unsigned id)
19+
{
20+
return scope_ok &&
21+
((count > UCS_DEVICE_LEVEL_EXEC_ID) ? (id == UCS_DEVICE_LEVEL_EXEC_ID)
22+
: true);
23+
}
2024

2125
namespace rocm_uct {
2226

@@ -80,18 +84,18 @@ template<typename T> class device_result_ptr {
8084
{
8185
unsigned int thread_id = threadIdx.x;
8286
unsigned int num_threads = blockDim.x;
83-
unsigned int warp_id = thread_id / 64; // ROCm wavefront size
84-
unsigned int num_warps = num_threads / 64;
87+
unsigned int warp_id = thread_id / UCT_ROCM_IPC_WAVEFRONT_SIZE;
88+
unsigned int num_warps = num_threads / UCT_ROCM_IPC_WAVEFRONT_SIZE;
8589
unsigned int block_id = blockIdx.x;
8690
unsigned int num_blocks = gridDim.x;
8791

8892
switch (level) {
8993
case UCS_DEVICE_LEVEL_THREAD:
90-
return UCS_DEVICE_LEVEL_EXEC_SELECT(block_id == 0, num_threads, thread_id);
94+
return ucs_device_level_exec_select(block_id == 0, num_threads, thread_id);
9195
case UCS_DEVICE_LEVEL_WARP:
92-
return UCS_DEVICE_LEVEL_EXEC_SELECT(block_id == 0, num_warps, warp_id);
96+
return ucs_device_level_exec_select(block_id == 0, num_warps, warp_id);
9397
case UCS_DEVICE_LEVEL_BLOCK:
94-
return UCS_DEVICE_LEVEL_EXEC_SELECT(true, num_blocks, block_id);
98+
return ucs_device_level_exec_select(true, num_blocks, block_id);
9599
case UCS_DEVICE_LEVEL_GRID:
96100
return true;
97101
}

‎test/gtest/uct/rocm/test_rocm_ipc_device.hip‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -180,10 +180,11 @@ UCS_TEST_P(test_rocm_ipc_rma, get_mem_elem_pack)
180180
mapped_buffer sendbuf(length, SEED1, *m_sender, 0, UCS_MEMORY_TYPE_ROCM);
181181
mapped_buffer recvbuf(length, SEED2, *m_receiver, 0, UCS_MEMORY_TYPE_ROCM);
182182

183-
ASSERT_EQ(hipSuccess, hipMalloc((void **)&mem_elem, mem_elem_size));
183+
mem_elem = static_cast<uct_device_mem_elem_t*>(malloc(mem_elem_size));
184+
ASSERT_NE(nullptr, mem_elem);
184185
EXPECT_UCS_OK(uct_md_mem_elem_pack(m_sender->md(), sendbuf.memh(),
185186
recvbuf.rkey(), mem_elem));
186-
hipFree(mem_elem);
187+
free(mem_elem);
187188
}
188189

189190
UCS_TEST_P(test_rocm_ipc_rma, get_device_ep)

0 commit comments

Comments
 (0)