Skip to content

Commit 1e6885c

Browse files
austin-hooveraustin-hoover
committed
Merge branch 'main' into env-tracker
2 parents c90c12c + 305a783 commit 1e6885c

4 files changed

Lines changed: 285 additions & 108 deletions

File tree

‎src/mpi/orbit_mpi.cc‎

Lines changed: 36 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -6,44 +6,38 @@
66
#include <ctime>
77

88
#if !defined(DOXYGEN_SHOULD_SKIP_THIS)
9+
10+
#if USE_MPI == 0
11+
static std::size_t ORBIT_MPI_Type_size(MPI_Datatype data) {
12+
switch(data){
13+
case MPI_CHAR: return sizeof(char);
14+
case MPI_UNSIGNED_CHAR: return sizeof(unsigned char);
15+
case MPI_BYTE: return sizeof(char);
16+
case MPI_SHORT: return sizeof(short);
17+
case MPI_UNSIGNED_SHORT: return sizeof(unsigned short);
18+
case MPI_INT: return sizeof(int);
19+
case MPI_UNSIGNED: return sizeof(unsigned);
20+
case MPI_LONG: return sizeof(long);
21+
case MPI_UNSIGNED_LONG: return sizeof(unsigned long);
22+
case MPI_FLOAT: return sizeof(float);
23+
case MPI_DOUBLE: return sizeof(double);
24+
case MPI_LONG_DOUBLE: return sizeof(long double);
25+
case MPI_LONG_LONG_INT: return sizeof(long long);
26+
default: return 0;
27+
}
28+
}
29+
#endif
30+
931
/** A C wrapper around MPI_Init. */
1032
int ORBIT_MPI_Init(){
11-
int res = 0;
1233
#if USE_MPI > 0
13-
int len=0;
14-
char** ch = NULL;
15-
// Getting arguments from sys.argv
16-
PyObject* sys_module = PyImport_ImportModule("sys");
17-
PyObject* argv_list = PyObject_GetAttrString(sys_module, "argv");
18-
19-
// Check if argv_list is a list
20-
if (PyList_Check(argv_list)) {
21-
// Access individual command-line arguments
22-
len = PyList_Size(argv_list);
23-
ch = (char**) malloc(sizeof(char*) * len);
24-
for (Py_ssize_t i = 0; i < len; ++i) {
25-
PyObject* item = PyList_GetItem(argv_list, i);
26-
if (item && PyUnicode_Check(item)) {
27-
ch[i] = const_cast<char*>(PyUnicode_AsUTF8(item));
28-
}
29-
}
30-
}
31-
32-
// Release references
33-
Py_XDECREF(argv_list);
34-
Py_XDECREF(sys_module);
35-
36-
res = MPI_Init(&len,&ch);
37-
38-
free(ch);
39-
ch = NULL;
34+
// Ignoring result; if it fails, the proc is doomed anyway.
35+
MPI_Init(NULL, NULL);
4036

4137
// Registering MPI finalize method at cleanup stage
4238
Py_AtExit(ORBIT_MPI_Finalize);
43-
#else
44-
res = MPI_SUCCESS;
4539
#endif
46-
return res;
40+
return MPI_SUCCESS;
4741
}
4842

4943
/** A C wrapper around MPI_Initialized. */
@@ -484,7 +478,7 @@ int ORBIT_MPI_Graphdims_get(MPI_Comm comm, int *nnodes, int *nedges){
484478
int ORBIT_MPI_Graph_get(MPI_Comm comm, int maxindex, int maxedges, int *index, int *edges){
485479
int res = 0;
486480
#if USE_MPI > 0
487-
res = ORBIT_MPI_Graph_get(comm, maxindex, maxedges, index, edges);
481+
res = MPI_Graph_get(comm, maxindex, maxedges, index, edges);
488482
#else
489483
res = MPI_SUCCESS;
490484
#endif
@@ -552,15 +546,19 @@ int ORBIT_MPI_Wait(MPI_Request *request, MPI_Status *status){
552546
}
553547

554548
/** A C wrapper around MPI_Allreduce. */
555-
int ORBIT_MPI_Allreduce(void* ar1, void* ar2, int n, MPI_Datatype data, MPI_Op op, MPI_Comm comm){
556-
int res = 0;
549+
int ORBIT_MPI_Allreduce(void* ar1, void* ar2, int n, MPI_Datatype data, MPI_Op op, MPI_Comm comm) {
557550
#if USE_MPI > 0
558-
res = MPI_Allreduce(ar1, ar2, n, data, op, comm);
551+
return MPI_Allreduce(ar1, ar2, n, data, op, comm);
559552
#else
560-
memcpy(ar2, ar1, n*sizeof(ar1));
561-
res = MPI_SUCCESS;
553+
if (ar1 == ORBIT_MPI_IN_PLACE || ar1 == ar2) return MPI_SUCCESS;
554+
555+
std::size_t nbytes = (std::size_t)n * ORBIT_MPI_Type_size(data);
556+
if (nbytes > 0) {
557+
std::memcpy(ar2, ar1, nbytes);
558+
}
559+
560+
return MPI_SUCCESS;
562561
#endif
563-
return res;
564562
}
565563

566564
/** A C wrapper around MPI_Bcast. */

‎src/mpi/orbit_mpi.hh‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,8 @@
88
#endif
99

1010
#if USE_MPI > 0
11-
#include "mpi.h"
11+
#include "mpi.h"
12+
#define ORBIT_MPI_IN_PLACE MPI_IN_PLACE
1213
#else
1314
//---------------------------------------------------------------
1415
//START the case when USE_MPI is not defined.
@@ -93,6 +94,9 @@
9394
#define MPI_ANY_SOURCE (-2)
9495
#define MPI_ANY_TAG (-1)
9596

97+
/* MPI in-place operations */
98+
#define ORBIT_MPI_IN_PLACE ((void*)1)
99+
96100
#endif
97101
//-------------------------------------------------------------
98102
//END the case when USE_MPI is defined or not.

‎src/mpi/wrap_orbit_mpi_send_receive_functions.hh‎

Lines changed: 131 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -28,79 +28,141 @@ static PyObject* mpi_wait(PyObject *self, PyObject *args){
2828
return Py_BuildValue("i",res);
2929
}
3030

31-
static PyObject* mpi_allreduce(PyObject *self, PyObject *args){
32-
PyObject* pyO_arr; PyObject* pyO_datatype; PyObject* pyO_op; PyObject* pyO_comm;
33-
if(!PyArg_ParseTuple( args,"OOOO:mpi_allreduce",&pyO_arr,&pyO_datatype,&pyO_op,&pyO_comm)){
34-
error("MPI_Allreduce([...data],MPI_Datatype type,MPI_Op op,MPI_Comm out) - needs 4 params.");
35-
}
36-
pyORBIT_MPI_Comm* pyComm = (pyORBIT_MPI_Comm*) pyO_comm;
37-
pyORBIT_MPI_Datatype* pyDatatype = (pyORBIT_MPI_Datatype*) pyO_datatype;
38-
pyORBIT_MPI_Op* pyOp = (pyORBIT_MPI_Op*) pyO_op;
39-
//check the data type
40-
if(pyDatatype->datatype != MPI_INT && pyDatatype->datatype != MPI_DOUBLE){
41-
error("MPI_Allreduce(...) data type could be INT or DOUBLE. STOP.");
42-
}
43-
//check if it is not sequence
44-
int is_seq = 0;
45-
if(PySequence_Check(pyO_arr) == 1){
46-
is_seq = 1;
31+
static PyObject* allreduce_scalar_int(PyObject* obj, pyORBIT_MPI_Op* pyOp, pyORBIT_MPI_Comm* pyComm) {
32+
long v = PyLong_AsLong(obj);
33+
if (PyErr_Occurred()) return NULL;
34+
35+
int val = (int)v;
36+
ORBIT_MPI_Allreduce(ORBIT_MPI_IN_PLACE, &val, 1, MPI_INT, pyOp->op, pyComm->comm);
37+
return Py_BuildValue("i", val);
38+
}
39+
40+
static PyObject* allreduce_scalar_double(PyObject* obj, pyORBIT_MPI_Op* pyOp, pyORBIT_MPI_Comm* pyComm) {
41+
double val = PyFloat_AsDouble(obj);
42+
if (PyErr_Occurred()) return NULL;
43+
44+
ORBIT_MPI_Allreduce(ORBIT_MPI_IN_PLACE, &val, 1, MPI_DOUBLE, pyOp->op, pyComm->comm);
45+
return Py_BuildValue("d", val);
46+
}
47+
48+
static PyObject* allreduce_sequence_int(PyObject* obj, pyORBIT_MPI_Op* pyOp, pyORBIT_MPI_Comm* pyComm) {
49+
PyObject *seq = PySequence_Fast(obj, "MPI_Allreduce(...) expected a sequence of ints");
50+
if (seq == NULL) return NULL;
51+
52+
Py_ssize_t size = PySequence_Fast_GET_SIZE(seq);
53+
PyObject *pyRes = PyTuple_New(size);
54+
55+
if (pyRes == NULL) {
56+
Py_DECREF(seq);
57+
return NULL;
58+
}
59+
60+
int buff_index = 0;
61+
int *buf = BufferStore::getBufferStore()->getFreeIntArr(buff_index, (int)size);
62+
63+
for (Py_ssize_t i = 0; i < size; ++i) {
64+
long v = PyLong_AsLong(PySequence_Fast_GET_ITEM(seq, i));
65+
if (PyErr_Occurred()) {
66+
BufferStore::getBufferStore()->setUnusedIntArr(buff_index);
67+
Py_DECREF(pyRes);
68+
Py_DECREF(seq);
69+
return NULL;
4770
}
48-
//it is NOT SEQUENCE
49-
if(is_seq == 0){
50-
if(pyDatatype->datatype == MPI_INT){
51-
int val = (int) PyLong_AsLong(pyO_arr);
52-
int val_out = 0;
53-
ORBIT_MPI_Allreduce(&val,&val_out,1,MPI_INT,pyOp->op,pyComm->comm);
54-
return Py_BuildValue("i",val_out);
55-
}
56-
if(pyDatatype->datatype == MPI_DOUBLE){
57-
double val = PyFloat_AsDouble(pyO_arr);
58-
double val_out = 0.;
59-
ORBIT_MPI_Allreduce(&val,&val_out,1,MPI_DOUBLE,pyOp->op,pyComm->comm);
60-
return Py_BuildValue("d",val_out);
61-
}
62-
error("MPI_Allreduce(...) - use only INT or DOUBLE data types");
71+
buf[i] = (int)v;
72+
}
73+
74+
ORBIT_MPI_Allreduce(ORBIT_MPI_IN_PLACE, buf, (int)size, MPI_INT, pyOp->op, pyComm->comm);
75+
76+
for (Py_ssize_t i = 0; i < size; ++i) {
77+
PyObject *item = Py_BuildValue("i", buf[i]);
78+
if (item == NULL || PyTuple_SetItem(pyRes, i, item) != 0) {
79+
Py_XDECREF(item);
80+
BufferStore::getBufferStore()->setUnusedIntArr(buff_index);
81+
Py_DECREF(pyRes);
82+
Py_DECREF(seq);
83+
return NULL;
6384
}
64-
//it IS A SEQUENCE
65-
int size = PySequence_Size(pyO_arr);
66-
PyObject* pyRes = PyTuple_New(size);
67-
//data is an INT array
68-
if(pyDatatype->datatype == MPI_INT){
69-
int buff_index0 = 0;
70-
int buff_index1 = 0;
71-
int* arr = BufferStore::getBufferStore()->getFreeIntArr(buff_index0,size);
72-
int* arr_out = BufferStore::getBufferStore()->getFreeIntArr(buff_index1,size);
73-
for(int i = 0; i < size; i++){
74-
arr[i]= (int) PyLong_AsLong(PySequence_Fast_GET_ITEM(pyO_arr, i));
75-
}
76-
ORBIT_MPI_Allreduce(arr,arr_out,size,MPI_INT,pyOp->op,pyComm->comm);
77-
for(int i = 0; i < size; i++){
78-
if(PyTuple_SetItem(pyRes,i,Py_BuildValue("i",arr_out[i])) != 0){
79-
error("MPI_Allreduce(...) cannot create a resulting tuple.");
80-
}
81-
}
82-
BufferStore::getBufferStore()->setUnusedIntArr(buff_index0);
83-
BufferStore::getBufferStore()->setUnusedIntArr(buff_index1);
85+
}
86+
87+
BufferStore::getBufferStore()->setUnusedIntArr(buff_index);
88+
Py_DECREF(seq);
89+
return pyRes;
90+
}
91+
92+
static PyObject* allreduce_sequence_double(PyObject* obj, pyORBIT_MPI_Op* pyOp, pyORBIT_MPI_Comm* pyComm) {
93+
PyObject *seq = PySequence_Fast(obj, "MPI_Allreduce(...) expected a sequence of ints");
94+
if (seq == NULL) return NULL;
95+
96+
Py_ssize_t size = PySequence_Fast_GET_SIZE(seq);
97+
PyObject *pyRes = PyTuple_New(size);
98+
99+
if (pyRes == NULL) {
100+
Py_DECREF(seq);
101+
return NULL;
102+
}
103+
104+
int buff_index = 0;
105+
double *buf = BufferStore::getBufferStore()->getFreeDoubleArr(buff_index, (int)size);
106+
107+
for (Py_ssize_t i = 0; i < size; ++i) {
108+
double v = PyFloat_AsDouble(PySequence_Fast_GET_ITEM(seq, i));
109+
if (PyErr_Occurred()) {
110+
BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index);
111+
Py_DECREF(pyRes);
112+
Py_DECREF(seq);
113+
return NULL;
84114
}
85-
//data is an DOUBLE array
86-
if(pyDatatype->datatype == MPI_DOUBLE){
87-
int buff_index0 = 0;
88-
int buff_index1 = 0;
89-
double* arr = BufferStore::getBufferStore()->getFreeDoubleArr(buff_index0,size);
90-
double* arr_out = BufferStore::getBufferStore()->getFreeDoubleArr(buff_index1,size);
91-
for(int i = 0; i < size; i++){
92-
arr[i]= PyFloat_AsDouble(PySequence_Fast_GET_ITEM(pyO_arr, i));
93-
}
94-
ORBIT_MPI_Allreduce(arr,arr_out,size,MPI_DOUBLE,pyOp->op,pyComm->comm);
95-
for(int i = 0; i < size; i++){
96-
if(PyTuple_SetItem(pyRes,i,Py_BuildValue("d",arr_out[i])) != 0){
97-
error("MPI_Allreduce(...) cannot create a resulting tuple.");
98-
}
99-
}
100-
BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index0);
101-
BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index1);
115+
buf[i] = v;
116+
}
117+
118+
ORBIT_MPI_Allreduce(ORBIT_MPI_IN_PLACE, buf, (int)size, MPI_DOUBLE, pyOp->op, pyComm->comm);
119+
120+
for (Py_ssize_t i = 0; i < size; ++i) {
121+
PyObject *item = Py_BuildValue("d", buf[i]);
122+
if (item == NULL || PyTuple_SetItem(pyRes, i, item) != 0) {
123+
Py_XDECREF(item);
124+
BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index);
125+
Py_DECREF(pyRes);
126+
Py_DECREF(seq);
127+
return NULL;
102128
}
103-
return pyRes;
129+
}
130+
131+
BufferStore::getBufferStore()->setUnusedDoubleArr(buff_index);
132+
Py_DECREF(seq);
133+
return pyRes;
134+
}
135+
136+
static PyObject* mpi_allreduce(PyObject *self, PyObject *args){
137+
PyObject *pyO_arr, *pyO_datatype, *pyO_op, *pyO_comm;
138+
139+
if(!PyArg_ParseTuple(args,"OOOO:mpi_allreduce", &pyO_arr, &pyO_datatype, &pyO_op, &pyO_comm)) {
140+
error("MPI_Allreduce([...data],MPI_Datatype type,MPI_Op op,MPI_Comm out) - needs 4 params.");
141+
}
142+
143+
pyORBIT_MPI_Comm* pyComm = (pyORBIT_MPI_Comm*) pyO_comm;
144+
pyORBIT_MPI_Datatype* pyDatatype = (pyORBIT_MPI_Datatype*) pyO_datatype;
145+
pyORBIT_MPI_Op* pyOp = (pyORBIT_MPI_Op*) pyO_op;
146+
147+
//check the data type
148+
MPI_Datatype dtype = pyDatatype->datatype;
149+
150+
if(dtype != MPI_INT && dtype != MPI_DOUBLE){
151+
error("MPI_Allreduce(...) data type should be INT or DOUBLE. STOP.");
152+
}
153+
154+
//check if it is a sequence or a scalar
155+
if (PySequence_Check(pyO_arr) == 0) {
156+
if (dtype == MPI_INT) {
157+
return allreduce_scalar_int(pyO_arr, pyOp, pyComm);
158+
}
159+
return allreduce_scalar_double(pyO_arr, pyOp, pyComm);
160+
} else {
161+
if (dtype == MPI_INT) {
162+
return allreduce_sequence_int(pyO_arr, pyOp, pyComm);
163+
}
164+
return allreduce_sequence_double(pyO_arr, pyOp, pyComm);
165+
}
104166
}
105167

106168
static PyObject* mpi_bcast(PyObject *self, PyObject *args){

0 commit comments

Comments
 (0)