-
Notifications
You must be signed in to change notification settings - Fork 32
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #158 from tvegas1/uct_read_plugin
UCT plugin: Add read-based implementation
- Loading branch information
Showing
9 changed files
with
578 additions
and
14 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,103 @@ | ||
| /************************************************************************* | ||
| * Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | ||
| * | ||
| * See LICENSE.txt for license information | ||
| ************************************************************************/ | ||
|
|
||
| #ifndef NCCL_UCX_UCT_RING_H_ | ||
| #define NCCL_UCX_UCT_RING_H_ | ||
|
|
||
| #include "nccl.h" | ||
| #include <assert.h> | ||
|
|
||
| #define NCCL_UCT_RING_SIZE (1 << 7) | ||
| #define NCCL_UCT_RING_MASK (NCCL_UCT_RING_SIZE - 1) | ||
|
|
||
| typedef struct nccl_uct_ring { | ||
| unsigned first; | ||
| unsigned last; | ||
| unsigned size; | ||
| unsigned entry_size; | ||
| int tag[NCCL_UCT_RING_SIZE]; | ||
| void *entry; | ||
| } nccl_uct_ring_t; | ||
|
|
||
| static inline ncclResult_t nccl_uct_ring_init(nccl_uct_ring_t *ring, | ||
| unsigned entry_size) { | ||
| int i; | ||
|
|
||
| ring->first = 0; | ||
| ring->last = 0; | ||
| ring->entry_size = entry_size; | ||
| ring->entry = malloc(entry_size * NCCL_UCT_RING_SIZE); | ||
| if (ring->entry == NULL) { | ||
| free(ring->entry); | ||
| return ncclSystemError; | ||
| } | ||
|
|
||
| for (i = 0; i < NCCL_UCT_RING_SIZE; i++) { | ||
| ring->tag[i] = INT_MAX; | ||
| } | ||
| return ncclSuccess; | ||
| } | ||
|
|
||
| static inline void nccl_uct_ring_deinit(nccl_uct_ring_t *ring) { | ||
| free(ring->entry); | ||
| } | ||
|
|
||
| static inline void *nccl_uct_ring_get_entry(nccl_uct_ring_t *ring, unsigned i) { | ||
| return (uint8_t*)ring->entry + (ring->entry_size * (i & NCCL_UCT_RING_MASK)); | ||
| } | ||
|
|
||
| static inline void nccl_uct_ring_append(nccl_uct_ring_t *ring, int tag, | ||
| void *data, size_t len) { | ||
| int j = ring->last & NCCL_UCT_RING_MASK; | ||
|
|
||
| ring->last++; | ||
|
|
||
| assert((ring->last & NCCL_UCT_RING_MASK) != | ||
| (ring->first & NCCL_UCT_RING_MASK)); | ||
| assert(ring->tag[j] == INT_MAX); | ||
| assert(len == ring->entry_size); | ||
|
|
||
| ring->tag[j] = tag; | ||
| memcpy(nccl_uct_ring_get_entry(ring, j), data, len); | ||
| } | ||
|
|
||
| static inline int nccl_uct_ring_is_empty(const nccl_uct_ring_t *ring) { | ||
| return ring->first == ring->last; | ||
| } | ||
|
|
||
| static inline void nccl_uct_ring_consume(nccl_uct_ring_t *ring, unsigned i) { | ||
| unsigned j = i & NCCL_UCT_RING_MASK; | ||
|
|
||
| assert(ring->tag[j] != INT_MAX); | ||
| ring->tag[j] = INT_MAX; | ||
|
|
||
| /* Cleanup upon tag hit */ | ||
| if (i == ring->first) { | ||
| for (; i != ring->last; i++) { | ||
| j = i & NCCL_UCT_RING_MASK; | ||
| if (ring->tag[j] != INT_MAX) { | ||
| break; | ||
| } | ||
| ring->first = i + 1; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| static inline unsigned nccl_uct_ring_find(nccl_uct_ring_t *ring, int tag) { | ||
| unsigned i; | ||
|
|
||
| assert(tag != INT_MAX); | ||
|
|
||
| for (i = ring->first; i != ring->last; i++) { | ||
| if (ring->tag[i & NCCL_UCT_RING_MASK] == tag) { | ||
| return i; | ||
| } | ||
| } | ||
|
|
||
| return ring->last; | ||
| } | ||
|
|
||
| #endif /* NCCL_UCX_UCT_RING_H_ */ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.