Skip to content

Commit aadb4d2

Browse files
UCT/GDA: Use wqe_idx for rc_gda progress
1 parent e75960c commit aadb4d2

File tree

11 files changed

+134
-140
lines changed

11 files changed

+134
-140
lines changed

src/ucp/api/device/ucp_device_impl.h

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
*/
2525
typedef struct ucp_device_request {
2626
uct_device_completion_t comp;
27+
ucs_status_t status;
2728
uct_device_ep_h device_ep;
2829
} ucp_device_request_t;
2930

@@ -51,9 +52,7 @@ UCS_F_DEVICE void ucp_device_request_init(uct_device_ep_t *device_ep,
5152
if (req != nullptr) {
5253
comp = &req->comp;
5354
req->device_ep = device_ep;
54-
uct_device_completion_init(comp);
5555
/* TODO: Handle multiple device posts with same req? */
56-
++comp->count;
5756
} else {
5857
comp = nullptr;
5958
}
@@ -63,16 +62,20 @@ UCS_F_DEVICE void ucp_device_request_init(uct_device_ep_t *device_ep,
6362
/**
6463
* Macro for device put operations with retry logic
6564
*/
66-
#define UCP_DEVICE_SEND_BLOCKING(_level, _uct_device_ep_send, _device_ep, ...) \
65+
#define UCP_DEVICE_SEND_BLOCKING(_level, _uct_device_ep_send, _device_ep, \
66+
_req, ...) \
6767
({ \
6868
ucs_status_t _status; \
6969
do { \
7070
_status = _uct_device_ep_send<_level>(_device_ep, __VA_ARGS__); \
7171
if (_status != UCS_ERR_NO_RESOURCE) { \
7272
break; \
7373
} \
74-
_status = uct_device_ep_progress<_level>(_device_ep); \
75-
} while (!UCS_STATUS_IS_ERR(_status)); \
74+
uct_device_ep_progress<_level>(_device_ep); \
75+
} while (1); \
76+
if (_req != nullptr) { \
77+
_req->status = _status; \
78+
} \
7679
_status; \
7780
})
7881

@@ -148,8 +151,8 @@ UCS_F_DEVICE ucs_status_t ucp_device_put_single(
148151
}
149152

150153
return UCP_DEVICE_SEND_BLOCKING(level, uct_device_ep_put_single, device_ep,
151-
uct_elem, address, remote_address, length,
152-
flags, comp);
154+
req, uct_elem, address, remote_address,
155+
length, flags, comp);
153156
}
154157

155158

@@ -199,8 +202,8 @@ UCS_F_DEVICE ucs_status_t ucp_device_counter_inc(
199202
}
200203

201204
return UCP_DEVICE_SEND_BLOCKING(level, uct_device_ep_atomic_add, device_ep,
202-
uct_elem, inc_value, remote_address, flags,
203-
comp);
205+
req, uct_elem, inc_value, remote_address,
206+
flags, comp);
204207
}
205208

206209

@@ -263,8 +266,9 @@ UCS_F_DEVICE ucs_status_t ucp_device_put_multi(
263266
}
264267

265268
return UCP_DEVICE_SEND_BLOCKING(level, uct_device_ep_put_multi, device_ep,
266-
uct_mem_list, mem_list_h->mem_list_length,
267-
addresses, remote_addresses, lengths,
269+
req, uct_mem_list,
270+
mem_list_h->mem_list_length, addresses,
271+
remote_addresses, lengths,
268272
counter_inc_value, counter_remote_address,
269273
flags, comp);
270274
}
@@ -338,10 +342,11 @@ UCS_F_DEVICE ucs_status_t ucp_device_put_multi_partial(
338342
}
339343

340344
return UCP_DEVICE_SEND_BLOCKING(level, uct_device_ep_put_multi_partial,
341-
device_ep, uct_mem_list, mem_list_indices,
342-
mem_list_count, addresses, remote_addresses,
343-
lengths, counter_index, counter_inc_value,
344-
counter_remote_address, flags, comp);
345+
device_ep, req, uct_mem_list,
346+
mem_list_indices, mem_list_count, addresses,
347+
remote_addresses, lengths, counter_index,
348+
counter_inc_value, counter_remote_address,
349+
flags, comp);
345350
}
346351

347352

@@ -409,19 +414,14 @@ UCS_F_DEVICE void ucp_device_counter_write(void *counter_ptr, uint64_t value)
409414
template<ucs_device_level_t level = UCS_DEVICE_LEVEL_THREAD>
410415
UCS_F_DEVICE ucs_status_t ucp_device_progress_req(ucp_device_request_t *req)
411416
{
412-
ucs_status_t status;
413-
414-
if (ucs_likely(req->comp.count == 0)) {
415-
return req->comp.status;
416-
}
417-
418-
status = uct_device_ep_progress<level>(req->device_ep);
419-
if (status != UCS_OK) {
420-
return status;
417+
if (ucs_likely(req->status != UCS_INPROGRESS)) {
418+
return req->status;
421419
}
422420

423-
return (ucs_likely(req->comp.count == 0)) ? req->comp.status :
424-
UCS_INPROGRESS;
421+
uct_device_ep_progress<level>(req->device_ep);
422+
req->status = uct_device_ep_check_completion<level>(req->device_ep,
423+
&req->comp);
424+
return req->status;
425425
}
426426

427427
#endif /* UCP_DEVICE_IMPL_H */

src/uct/api/device/uct_device_impl.h

Lines changed: 23 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,11 @@
1515

1616
#include <uct/ib/mlx5/gdaki/gdaki.cuh>
1717

18+
union uct_device_completion {
19+
uct_rc_gda_completion_t rc;
20+
uct_cuda_ipc_completion_t cuda_ipc;
21+
};
22+
1823

1924
/**
2025
* @ingroup UCT_DEVICE
@@ -242,34 +247,37 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_put_multi_partial(
242247
* @brief Progress all operations on device endpoint @a device_ep.
243248
*
244249
* @param [in] device_ep Device endpoint to be used for the operation.
245-
*
246-
* @return UCS_OK - Some operation was completed.
247-
* @return UCS_INPROGRESS - No progress on the endpoint.
248-
* @return Error code as defined by @ref ucs_status_t
249250
*/
250251
template<ucs_device_level_t level>
251-
UCS_F_DEVICE ucs_status_t uct_device_ep_progress(uct_device_ep_h device_ep)
252+
UCS_F_DEVICE void uct_device_ep_progress(uct_device_ep_h device_ep)
252253
{
253254
if (device_ep->uct_tl_id == UCT_DEVICE_TL_RC_MLX5_GDA) {
254-
return uct_rc_mlx5_gda_ep_progress<level>(device_ep);
255-
} else if (device_ep->uct_tl_id == UCT_DEVICE_TL_CUDA_IPC) {
256-
return UCS_OK;
255+
uct_rc_mlx5_gda_ep_progress<level>(device_ep);
257256
}
258-
259-
return UCS_ERR_UNSUPPORTED;
260257
}
261258

262259

263260
/**
264261
* @ingroup UCT_DEVICE
265-
* @brief Initialize a device completion object.
262+
* @brief Check whther opetation executed on device endpoint @a device_ep was
263+
* completed.
264+
*
265+
* @param [in] device_ep Device endpoint to be used for the operation.
266+
* @param [in] comp Completion object tracking operation progress.
266267
*
267-
* @param [out] comp Device completion object to initialize.
268+
* @return UCS_OK - Some operation was completed.
269+
* @return UCS_INPROGRESS - No progress on the endpoint.
270+
* @return Error code as defined by @ref ucs_status_t
268271
*/
269-
UCS_F_DEVICE void uct_device_completion_init(uct_device_completion_t *comp)
272+
template<ucs_device_level_t level>
273+
UCS_F_DEVICE ucs_status_t uct_device_ep_check_completion(
274+
uct_device_ep_h device_ep, uct_device_completion_t *comp)
270275
{
271-
comp->count = 0;
272-
comp->status = UCS_OK;
276+
if (device_ep->uct_tl_id == UCT_DEVICE_TL_RC_MLX5_GDA) {
277+
return uct_rc_mlx5_gda_ep_check_completion<level>(device_ep, comp);
278+
}
279+
280+
return UCS_ERR_UNSUPPORTED;
273281
}
274282

275283
#endif

src/uct/api/device/uct_device_types.h

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -42,10 +42,7 @@ typedef struct uct_device_ep {
4242

4343

4444
/* Completion object for device operations */
45-
typedef struct uct_device_completion {
46-
uint32_t count; /* How many operations are pending */
47-
ucs_status_t status; /* Status of the operation */
48-
} uct_device_completion_t;
45+
typedef union uct_device_completion uct_device_completion_t;
4946

5047

5148
/* Base structure for all device memory elements */

src/uct/cuda/cuda_ipc/cuda_ipc.cuh

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -304,8 +304,6 @@ uct_cuda_ipc_ep_put_single(uct_device_ep_h device_ep,
304304
mapped_rem_addr = uct_cuda_ipc_map_remote(cuda_ipc_mem_element, remote_address);
305305
uct_cuda_ipc_copy_level<level>(mapped_rem_addr, address, length);
306306
uct_cuda_ipc_level_sync<level>();
307-
--comp->count;
308-
309307
return UCS_OK;
310308
}
311309

@@ -339,10 +337,6 @@ uct_cuda_ipc_ep_put_multi(uct_device_ep_h device_ep,
339337
}
340338

341339
uct_cuda_ipc_level_sync<level>();
342-
if (lane_id == 0) {
343-
--comp->count;
344-
}
345-
346340
return UCS_OK;
347341
}
348342

@@ -376,10 +370,6 @@ uct_cuda_ipc_ep_put_multi_partial(uct_device_ep_h device_ep,
376370
}
377371

378372
uct_cuda_ipc_level_sync<level>();
379-
if (lane_id == 0) {
380-
--comp->count;
381-
}
382-
383373
return UCS_OK;
384374
}
385375

@@ -403,10 +393,6 @@ uct_cuda_ipc_ep_atomic_add(uct_device_ep_h device_ep,
403393
}
404394

405395
uct_cuda_ipc_level_sync<level>();
406-
if (lane_id == 0) {
407-
--comp->count;
408-
}
409-
410396
return UCS_OK;
411397
}
412398

src/uct/cuda/cuda_ipc/cuda_ipc_device.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,4 +12,8 @@ typedef struct {
1212
ptrdiff_t mapped_offset;
1313
} uct_cuda_ipc_device_mem_element_t;
1414

15+
16+
typedef struct {
17+
} uct_cuda_ipc_completion_t;
18+
1519
#endif

src/uct/ib/mlx5/gdaki/gdaki.c

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -92,10 +92,8 @@ static UCS_CLASS_INIT_FUNC(uct_rc_gdaki_ep_t, const uct_ep_params_t *params)
9292
uct_ib_mlx5_wq_calc_sizes(&qp_attr);
9393

9494
cq_attr.flags |= UCT_IB_MLX5_CQ_IGNORE_OVERRUN;
95-
cq_attr.umem_offset = ucs_align_up_pow2(
96-
sizeof(uct_rc_gdaki_dev_ep_t) +
97-
qp_attr.max_tx * sizeof(uct_rc_gdaki_op_t),
98-
ucs_get_page_size());
95+
cq_attr.umem_offset = ucs_align_up_pow2(sizeof(uct_rc_gdaki_dev_ep_t),
96+
ucs_get_page_size());
9997

10098
qp_attr.mmio_mode = UCT_IB_MLX5_MMIO_MODE_DB;
10199
qp_attr.super.srq_num = 0;
@@ -109,9 +107,9 @@ static UCS_CLASS_INIT_FUNC(uct_rc_gdaki_ep_t, const uct_ep_params_t *params)
109107
dev_ep_size = qp_attr.umem_offset + qp_attr.len;
110108
/*
111109
* dev_ep layout:
112-
* +---------------------+-------+---------+---------+
113-
* | counters, dbr | ops | cq buff | wq buff |
114-
* +---------------------+-------+---------+---------+
110+
* +---------------------+---------+---------+
111+
* | counters, dbr | cq buff | wq buff |
112+
* +---------------------+---------+---------+
115113
*/
116114
status = uct_rc_gdaki_alloc(dev_ep_size, ucs_get_page_size(),
117115
(void**)&self->ep_gpu, &self->ep_raw);

0 commit comments

Comments
 (0)