Skip to content

Commit fced151

Browse files
edgargabrieledgargabriel
authored
Merge pull request #12986 from edgargabriel/topic/coll-accelerator-extension
coll/accelerator initial enhancements
2 parents cfb4d38 + 118c1b0 commit fced151

File tree

6 files changed

+68
-36
lines changed

6 files changed

+68
-36
lines changed

‎ompi/mca/coll/accelerator/coll_accelerator.h‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
* reserved.
66
* Copyright (c) 2014-2024 NVIDIA Corporation. All rights reserved.
77
* Copyright (c) 2024 Triad National Security, LLC. All rights reserved.
8+
* Copyright (c) 2024 Advanced Micro Devices, Inc. All Rights reserved.
89
* $COPYRIGHT$
910
*
1011
* Additional copyrights may follow
@@ -87,22 +88,24 @@ mca_coll_accelerator_reduce_scatter_block(const void *sbuf, void *rbuf, size_t r
8788
* @retval >0 The buffer belongs to a managed buffer in
8889
* device memory.
8990
*/
90-
static inline int mca_coll_accelerator_check_buf(void *addr)
91+
static inline int mca_coll_accelerator_check_buf(void *addr, int *dev_id)
9192
{
9293
uint64_t flags;
93-
int dev_id;
94+
9495
if (OPAL_LIKELY(NULL != addr)) {
95-
return opal_accelerator.check_addr(addr, &dev_id, &flags);
96+
return opal_accelerator.check_addr(addr, dev_id, &flags);
9697
} else {
98+
*dev_id = MCA_ACCELERATOR_NO_DEVICE_ID;
9799
return 0;
98100
}
99101
}
100102

101-
static inline void *mca_coll_accelerator_memcpy(void *dest, const void *src, size_t size)
103+
static inline void *mca_coll_accelerator_memcpy(void *dest, int dest_dev, const void *src, int src_dev, size_t size,
104+
opal_accelerator_transfer_type_t type)
102105
{
103106
int res;
104-
res = opal_accelerator.mem_copy(MCA_ACCELERATOR_NO_DEVICE_ID, MCA_ACCELERATOR_NO_DEVICE_ID,
105-
dest, src, size, MCA_ACCELERATOR_TRANSFER_UNSPEC);
107+
108+
res = opal_accelerator.mem_copy(dest_dev, src_dev, dest, src, size, type);
106109
if (res != 0) {
107110
opal_output(0, "coll/accelerator: Error in mem_copy: res=%d, dest=%p, src=%p, size=%d", res, dest, src,
108111
(int) size);

‎ompi/mca/coll/accelerator/coll_accelerator_allreduce.c‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
* Copyright (c) 2014-2015 NVIDIA Corporation. All rights reserved.
66
* Copyright (c) 2022 Amazon.com, Inc. or its affiliates. All Rights reserved.
77
* Copyright (c) 2024 Triad National Security, LLC. All rights reserved.
8+
* Copyright (c) 2024 Advanced Micro Devices, Inc. All Rights reserved.
89
* $COPYRIGHT$
910
*
1011
* Additional copyrights may follow
@@ -37,11 +38,12 @@ mca_coll_accelerator_allreduce(const void *sbuf, void *rbuf, size_t count,
3738
mca_coll_accelerator_module_t *s = (mca_coll_accelerator_module_t*) module;
3839
ptrdiff_t gap;
3940
char *rbuf1 = NULL, *sbuf1 = NULL, *rbuf2 = NULL;
41+
int sbuf_dev, rbuf_dev;
4042
size_t bufsize;
4143
int rc;
4244

4345
bufsize = opal_datatype_span(&dtype->super, count, &gap);
44-
rc = mca_coll_accelerator_check_buf((void *)sbuf);
46+
rc = mca_coll_accelerator_check_buf((void *)sbuf, &sbuf_dev);
4547
if (rc < 0) {
4648
return rc;
4749
}
@@ -50,10 +52,11 @@ mca_coll_accelerator_allreduce(const void *sbuf, void *rbuf, size_t count,
5052
if (NULL == sbuf1) {
5153
return OMPI_ERR_OUT_OF_RESOURCE;
5254
}
53-
mca_coll_accelerator_memcpy(sbuf1, sbuf, bufsize);
55+
mca_coll_accelerator_memcpy(sbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, sbuf, sbuf_dev,
56+
bufsize, MCA_ACCELERATOR_TRANSFER_DTOH);
5457
sbuf = sbuf1 - gap;
5558
}
56-
rc = mca_coll_accelerator_check_buf(rbuf);
59+
rc = mca_coll_accelerator_check_buf(rbuf, &rbuf_dev);
5760
if (rc < 0) {
5861
return rc;
5962
}
@@ -63,7 +66,8 @@ mca_coll_accelerator_allreduce(const void *sbuf, void *rbuf, size_t count,
6366
if (NULL != sbuf1) free(sbuf1);
6467
return OMPI_ERR_OUT_OF_RESOURCE;
6568
}
66-
mca_coll_accelerator_memcpy(rbuf1, rbuf, bufsize);
69+
mca_coll_accelerator_memcpy(rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, rbuf, rbuf_dev,
70+
bufsize, MCA_ACCELERATOR_TRANSFER_DTOH);
6771
rbuf2 = rbuf; /* save away original buffer */
6872
rbuf = rbuf1 - gap;
6973
}
@@ -73,7 +77,8 @@ mca_coll_accelerator_allreduce(const void *sbuf, void *rbuf, size_t count,
7377
}
7478
if (NULL != rbuf1) {
7579
rbuf = rbuf2;
76-
mca_coll_accelerator_memcpy(rbuf, rbuf1, bufsize);
80+
mca_coll_accelerator_memcpy(rbuf, rbuf_dev, rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, bufsize,
81+
MCA_ACCELERATOR_TRANSFER_HTOD);
7782
free(rbuf1);
7883
}
7984
return rc;

‎ompi/mca/coll/accelerator/coll_accelerator_exscan.c‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
* Copyright (c) 2014-2015 NVIDIA Corporation. All rights reserved.
66
* Copyright (c) 2022 Amazon.com, Inc. or its affiliates. All Rights reserved.
77
* Copyright (c) 2024 Triad National Security, LLC. All rights reserved.
8+
* Copyright (c) 2024 Advanced Micro Devices, Inc. All Rights reserved.
89
* $COPYRIGHT$
910
*
1011
* Additional copyrights may follow
@@ -29,11 +30,12 @@ int mca_coll_accelerator_exscan(const void *sbuf, void *rbuf, size_t count,
2930
mca_coll_accelerator_module_t *s = (mca_coll_accelerator_module_t*) module;
3031
ptrdiff_t gap;
3132
char *rbuf1 = NULL, *sbuf1 = NULL, *rbuf2 = NULL;
33+
int sbuf_dev, rbuf_dev;
3234
size_t bufsize;
3335
int rc;
3436

3537
bufsize = opal_datatype_span(&dtype->super, count, &gap);
36-
rc = mca_coll_accelerator_check_buf((void *)sbuf);
38+
rc = mca_coll_accelerator_check_buf((void *)sbuf, &sbuf_dev);
3739
if (rc < 0) {
3840
return rc;
3941
}
@@ -43,10 +45,11 @@ int mca_coll_accelerator_exscan(const void *sbuf, void *rbuf, size_t count,
4345
if (NULL == sbuf1) {
4446
return OMPI_ERR_OUT_OF_RESOURCE;
4547
}
46-
mca_coll_accelerator_memcpy(sbuf1, sbuf, bufsize);
48+
mca_coll_accelerator_memcpy(sbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, sbuf, sbuf_dev, bufsize,
49+
MCA_ACCELERATOR_TRANSFER_DTOH);
4750
sbuf = sbuf1 - gap;
4851
}
49-
rc = mca_coll_accelerator_check_buf(rbuf);
52+
rc = mca_coll_accelerator_check_buf(rbuf, &rbuf_dev);
5053
if (rc < 0) {
5154
return rc;
5255
}
@@ -56,7 +59,8 @@ int mca_coll_accelerator_exscan(const void *sbuf, void *rbuf, size_t count,
5659
if (NULL != sbuf1) free(sbuf1);
5760
return OMPI_ERR_OUT_OF_RESOURCE;
5861
}
59-
mca_coll_accelerator_memcpy(rbuf1, rbuf, bufsize);
62+
mca_coll_accelerator_memcpy(rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, rbuf, rbuf_dev, bufsize,
63+
MCA_ACCELERATOR_TRANSFER_DTOH);
6064
rbuf2 = rbuf; /* save away original buffer */
6165
rbuf = rbuf1 - gap;
6266
}
@@ -68,7 +72,8 @@ int mca_coll_accelerator_exscan(const void *sbuf, void *rbuf, size_t count,
6872
}
6973
if (NULL != rbuf1) {
7074
rbuf = rbuf2;
71-
mca_coll_accelerator_memcpy(rbuf, rbuf1, bufsize);
75+
mca_coll_accelerator_memcpy(rbuf, rbuf_dev, rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, bufsize,
76+
MCA_ACCELERATOR_TRANSFER_HTOD);
7277
free(rbuf1);
7378
}
7479
return rc;

‎ompi/mca/coll/accelerator/coll_accelerator_reduce.c‎

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
* Copyright (c) 2014-2015 NVIDIA Corporation. All rights reserved.
77
* Copyright (c) 2022 Amazon.com, Inc. or its affiliates. All Rights reserved.
88
* Copyright (c) 2024 Triad National Security, LLC. All rights reserved.
9+
* Copyright (c) 2024 Advanced Micro Devices, Inc. All Rights reserved.
910
* $COPYRIGHT$
1011
*
1112
* Additional copyrights may follow
@@ -39,12 +40,13 @@ mca_coll_accelerator_reduce(const void *sbuf, void *rbuf, size_t count,
3940
int rank = ompi_comm_rank(comm);
4041
ptrdiff_t gap;
4142
char *rbuf1 = NULL, *sbuf1 = NULL, *rbuf2 = NULL;
43+
int rbuf_dev, sbuf_dev;
4244
size_t bufsize;
4345
int rc;
4446

4547
bufsize = opal_datatype_span(&dtype->super, count, &gap);
4648

47-
rc = mca_coll_accelerator_check_buf((void *)sbuf);
49+
rc = mca_coll_accelerator_check_buf((void *)sbuf, &sbuf_dev);
4850
if (rc < 0) {
4951
return rc;
5052
}
@@ -53,11 +55,12 @@ mca_coll_accelerator_reduce(const void *sbuf, void *rbuf, size_t count,
5355
if (NULL == sbuf1) {
5456
return OMPI_ERR_OUT_OF_RESOURCE;
5557
}
56-
mca_coll_accelerator_memcpy(sbuf1, sbuf, bufsize);
58+
mca_coll_accelerator_memcpy(sbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, sbuf, sbuf_dev, bufsize,
59+
MCA_ACCELERATOR_TRANSFER_DTOH);
5760
sbuf = sbuf1 - gap;
5861
}
5962

60-
rc = mca_coll_accelerator_check_buf(rbuf);
63+
rc = mca_coll_accelerator_check_buf(rbuf, &rbuf_dev);
6164
if (rc < 0) {
6265
return rc;
6366
}
@@ -67,7 +70,8 @@ mca_coll_accelerator_reduce(const void *sbuf, void *rbuf, size_t count,
6770
if (NULL != sbuf1) free(sbuf1);
6871
return OMPI_ERR_OUT_OF_RESOURCE;
6972
}
70-
mca_coll_accelerator_memcpy(rbuf1, rbuf, bufsize);
73+
mca_coll_accelerator_memcpy(rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, rbuf, rbuf_dev, bufsize,
74+
MCA_ACCELERATOR_TRANSFER_DTOH);
7175
rbuf2 = rbuf; /* save away original buffer */
7276
rbuf = rbuf1 - gap;
7377
}
@@ -80,7 +84,8 @@ mca_coll_accelerator_reduce(const void *sbuf, void *rbuf, size_t count,
8084
}
8185
if (NULL != rbuf1) {
8286
rbuf = rbuf2;
83-
mca_coll_accelerator_memcpy(rbuf, rbuf1, bufsize);
87+
mca_coll_accelerator_memcpy(rbuf, rbuf_dev, rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, bufsize,
88+
MCA_ACCELERATOR_TRANSFER_HTOD);
8489
free(rbuf1);
8590
}
8691
return rc;
@@ -94,12 +99,13 @@ mca_coll_accelerator_reduce_local(const void *sbuf, void *rbuf, size_t count,
9499
{
95100
ptrdiff_t gap;
96101
char *rbuf1 = NULL, *sbuf1 = NULL, *rbuf2 = NULL;
102+
int sbuf_dev, rbuf_dev;
97103
size_t bufsize;
98104
int rc;
99105

100106
bufsize = opal_datatype_span(&dtype->super, count, &gap);
101107

102-
rc = mca_coll_accelerator_check_buf((void *)sbuf);
108+
rc = mca_coll_accelerator_check_buf((void *)sbuf, &sbuf_dev);
103109
if (rc < 0) {
104110
return rc;
105111
}
@@ -109,11 +115,12 @@ mca_coll_accelerator_reduce_local(const void *sbuf, void *rbuf, size_t count,
109115
if (NULL == sbuf1) {
110116
return OMPI_ERR_OUT_OF_RESOURCE;
111117
}
112-
mca_coll_accelerator_memcpy(sbuf1, sbuf, bufsize);
118+
mca_coll_accelerator_memcpy(sbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, sbuf, sbuf_dev, bufsize,
119+
MCA_ACCELERATOR_TRANSFER_DTOH);
113120
sbuf = sbuf1 - gap;
114121
}
115122

116-
rc = mca_coll_accelerator_check_buf(rbuf);
123+
rc = mca_coll_accelerator_check_buf(rbuf, &rbuf_dev);
117124
if (rc < 0) {
118125
return rc;
119126
}
@@ -124,7 +131,8 @@ mca_coll_accelerator_reduce_local(const void *sbuf, void *rbuf, size_t count,
124131
if (NULL != sbuf1) free(sbuf1);
125132
return OMPI_ERR_OUT_OF_RESOURCE;
126133
}
127-
mca_coll_accelerator_memcpy(rbuf1, rbuf, bufsize);
134+
mca_coll_accelerator_memcpy(rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, rbuf, rbuf_dev, bufsize,
135+
MCA_ACCELERATOR_TRANSFER_DTOH);
128136
rbuf2 = rbuf; /* save away original buffer */
129137
rbuf = rbuf1 - gap;
130138
}
@@ -137,7 +145,8 @@ mca_coll_accelerator_reduce_local(const void *sbuf, void *rbuf, size_t count,
137145
}
138146
if (NULL != rbuf1) {
139147
rbuf = rbuf2;
140-
mca_coll_accelerator_memcpy(rbuf, rbuf1, bufsize);
148+
mca_coll_accelerator_memcpy(rbuf, rbuf_dev, rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, bufsize,
149+
MCA_ACCELERATOR_TRANSFER_HTOD);
141150
free(rbuf1);
142151
}
143152
return rc;

‎ompi/mca/coll/accelerator/coll_accelerator_reduce_scatter_block.c‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
* Copyright (c) 2014-2015 NVIDIA Corporation. All rights reserved.
66
* Copyright (c) 2022 Amazon.com, Inc. or its affiliates. All Rights reserved.
77
* Copyright (c) 2024 Triad National Security, LLC. All rights reserved.
8+
* Copyright (c) 2024 Advanced Micro Devices, Inc. All Rights reserved.
89
* $COPYRIGHT$
910
*
1011
* Additional copyrights may follow
@@ -41,13 +42,14 @@ mca_coll_accelerator_reduce_scatter_block(const void *sbuf, void *rbuf, size_t r
4142
mca_coll_accelerator_module_t *s = (mca_coll_accelerator_module_t*) module;
4243
ptrdiff_t gap;
4344
char *rbuf1 = NULL, *sbuf1 = NULL, *rbuf2 = NULL;
45+
int sbuf_dev, rbuf_dev;
4446
size_t sbufsize, rbufsize;
4547
int rc;
4648

4749
rbufsize = opal_datatype_span(&dtype->super, rcount, &gap);
4850

4951
sbufsize = rbufsize * ompi_comm_size(comm);
50-
rc = mca_coll_accelerator_check_buf((void *)sbuf);
52+
rc = mca_coll_accelerator_check_buf((void *)sbuf, &sbuf_dev);
5153
if (rc < 0) {
5254
return rc;
5355
}
@@ -56,10 +58,11 @@ mca_coll_accelerator_reduce_scatter_block(const void *sbuf, void *rbuf, size_t r
5658
if (NULL == sbuf1) {
5759
return OMPI_ERR_OUT_OF_RESOURCE;
5860
}
59-
mca_coll_accelerator_memcpy(sbuf1, sbuf, sbufsize);
61+
mca_coll_accelerator_memcpy(sbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, sbuf, sbuf_dev, sbufsize,
62+
MCA_ACCELERATOR_TRANSFER_DTOH);
6063
sbuf = sbuf1 - gap;
6164
}
62-
rc = mca_coll_accelerator_check_buf(rbuf);
65+
rc = mca_coll_accelerator_check_buf(rbuf, &rbuf_dev);
6366
if (rc < 0) {
6467
return rc;
6568
}
@@ -69,7 +72,8 @@ mca_coll_accelerator_reduce_scatter_block(const void *sbuf, void *rbuf, size_t r
6972
if (NULL != sbuf1) free(sbuf1);
7073
return OMPI_ERR_OUT_OF_RESOURCE;
7174
}
72-
mca_coll_accelerator_memcpy(rbuf1, rbuf, rbufsize);
75+
mca_coll_accelerator_memcpy(rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, rbuf, rbuf_dev, rbufsize,
76+
MCA_ACCELERATOR_TRANSFER_DTOH);
7377
rbuf2 = rbuf; /* save away original buffer */
7478
rbuf = rbuf1 - gap;
7579
}
@@ -80,7 +84,8 @@ mca_coll_accelerator_reduce_scatter_block(const void *sbuf, void *rbuf, size_t r
8084
}
8185
if (NULL != rbuf1) {
8286
rbuf = rbuf2;
83-
mca_coll_accelerator_memcpy(rbuf, rbuf1, rbufsize);
87+
mca_coll_accelerator_memcpy(rbuf, rbuf_dev, rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, rbufsize,
88+
MCA_ACCELERATOR_TRANSFER_HTOD);
8489
free(rbuf1);
8590
}
8691
return rc;

‎ompi/mca/coll/accelerator/coll_accelerator_scan.c‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
* Copyright (c) 2014-2015 NVIDIA Corporation. All rights reserved.
66
* Copyright (c) 2022 Amazon.com, Inc. or its affiliates. All Rights reserved.
77
* Copyright (c) 2024 Triad National Security, LLC. All rights reserved.
8+
* Copyright (c) 2024 Advanced Micro Devices, Inc. All Rights reserved.
89
* $COPYRIGHT$
910
*
1011
* Additional copyrights may follow
@@ -36,11 +37,12 @@ int mca_coll_accelerator_scan(const void *sbuf, void *rbuf, size_t count,
3637
mca_coll_accelerator_module_t *s = (mca_coll_accelerator_module_t*) module;
3738
ptrdiff_t gap;
3839
char *rbuf1 = NULL, *sbuf1 = NULL, *rbuf2 = NULL;
40+
int sbuf_dev, rbuf_dev;
3941
size_t bufsize;
4042
int rc;
4143

4244
bufsize = opal_datatype_span(&dtype->super, count, &gap);
43-
rc = mca_coll_accelerator_check_buf((void *)sbuf);
45+
rc = mca_coll_accelerator_check_buf((void *)sbuf, &sbuf_dev);
4446
if (rc < 0) {
4547
return rc;
4648
}
@@ -49,10 +51,11 @@ int mca_coll_accelerator_scan(const void *sbuf, void *rbuf, size_t count,
4951
if (NULL == sbuf1) {
5052
return OMPI_ERR_OUT_OF_RESOURCE;
5153
}
52-
mca_coll_accelerator_memcpy(sbuf1, sbuf, bufsize);
54+
mca_coll_accelerator_memcpy(sbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, sbuf, sbuf_dev, bufsize,
55+
MCA_ACCELERATOR_TRANSFER_DTOH);
5356
sbuf = sbuf1 - gap;
5457
}
55-
rc = mca_coll_accelerator_check_buf(rbuf);
58+
rc = mca_coll_accelerator_check_buf(rbuf, &rbuf_dev);
5659
if (rc < 0) {
5760
return rc;
5861
}
@@ -62,7 +65,8 @@ int mca_coll_accelerator_scan(const void *sbuf, void *rbuf, size_t count,
6265
if (NULL != sbuf1) free(sbuf1);
6366
return OMPI_ERR_OUT_OF_RESOURCE;
6467
}
65-
mca_coll_accelerator_memcpy(rbuf1, rbuf, bufsize);
68+
mca_coll_accelerator_memcpy(rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, rbuf, rbuf_dev, bufsize,
69+
MCA_ACCELERATOR_TRANSFER_DTOH);
6670
rbuf2 = rbuf; /* save away original buffer */
6771
rbuf = rbuf1 - gap;
6872
}
@@ -73,7 +77,8 @@ int mca_coll_accelerator_scan(const void *sbuf, void *rbuf, size_t count,
7377
}
7478
if (NULL != rbuf1) {
7579
rbuf = rbuf2;
76-
mca_coll_accelerator_memcpy(rbuf, rbuf1, bufsize);
80+
mca_coll_accelerator_memcpy(rbuf, rbuf_dev, rbuf1, MCA_ACCELERATOR_NO_DEVICE_ID, bufsize,
81+
MCA_ACCELERATOR_TRANSFER_HTOD);
7782
free(rbuf1);
7883
}
7984
return rc;

0 commit comments

Comments
 (0)