Repository navigation
Expand file tree
/
Copy pathdiff.patch
More file actions
100 lines (95 loc) · 2.61 KB
/
Copy pathdiff.patch
File metadata and controls
100 lines (95 loc) · 2.61 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
diff --git a/ompi/include/mpi.h.in b/ompi/include/mpi.h.in
index fa4082a37f..daf778e284 100644
--- a/ompi/include/mpi.h.in
+++ b/ompi/include/mpi.h.in
@@ -1320,6 +1320,16 @@ OMPI_DECLSPEC extern struct ompi_predefined_datatype_t ompi_mpi_ub;
* MPI API
*/
+
+#define MPI_EXPOSE_UCP_WORKER
+
+
+#if defined(MPI_EXPOSE_UCP_WORKER)
+OMPI_DECLSPEC extern void* MPI_mca_UCP_get_worker(void);
+OMPI_DECLSPEC extern void MPI_mca_UCP_worker_progress(void *);
+#endif
+
+
OMPI_DECLSPEC int MPI_Abort(MPI_Comm comm, int errorcode);
OMPI_DECLSPEC int MPI_Accumulate(const void *origin_addr, int origin_count, MPI_Datatype origin_datatype,
int target_rank, MPI_Aint target_disp, int target_count,
diff --git a/ompi/mca/pml/ucx/pml_ucx.c b/ompi/mca/pml/ucx/pml_ucx.c
index 5eb621990c..f4f9d63180 100644
--- a/ompi/mca/pml/ucx/pml_ucx.c
+++ b/ompi/mca/pml/ucx/pml_ucx.c
@@ -279,6 +279,26 @@ int mca_pml_ucx_close(void)
return OMPI_SUCCESS;
}
+
+extern void* (*mca_pml_ucx_get_worker_cb)(void) = NULL;
+extern void (*mca_pml_ucx_worker_progress_cb)(void*) = NULL;
+
+void* mca_pml_ucx_get_worker(void)
+{
+ return (void *)&ompi_pml_ucx.ucp_worker;
+}
+
+void mca_pml_ucx_worker_progress(void *worker_ptr)
+{
+ ucp_worker_h *worker = (ucp_worker_h *)worker_ptr;
+ /*fprintf(stderr, "Info: worker address %p\n", worker);*/
+ (void) ucp_worker_progress(*worker);
+ return;
+}
+
+
+
+
int mca_pml_ucx_init(int enable_mpi_threads)
{
ucp_worker_params_t params;
@@ -311,6 +331,10 @@ int mca_pml_ucx_init(int enable_mpi_threads)
goto err;
}
+ // Register custom callbacks
+ mca_pml_ucx_get_worker_cb = mca_pml_ucx_get_worker;
+ mca_pml_ucx_worker_progress_cb = mca_pml_ucx_worker_progress;
+
attr.field_mask = UCP_WORKER_ATTR_FIELD_THREAD_MODE;
status = ucp_worker_query(ompi_pml_ucx.ucp_worker, &attr);
if (UCS_OK != status) {
diff --git a/opal/runtime/opal_init.c b/opal/runtime/opal_init.c
index 855da937ee..eb0094e119 100644
--- a/opal/runtime/opal_init.c
+++ b/opal/runtime/opal_init.c
@@ -512,6 +512,32 @@ opal_init_util(int* pargc, char*** pargv)
}
+
+
+void* (*mca_pml_ucx_get_worker_cb)(void) = NULL;
+void (*mca_pml_ucx_worker_progress_cb)(void *) = NULL;
+
+
+void* MPI_mca_UCP_get_worker(void)
+{
+ if (mca_pml_ucx_get_worker_cb) {
+ return mca_pml_ucx_get_worker_cb();
+ }
+ return NULL;
+}
+
+
+void MPI_mca_UCP_worker_progress(void *worker_ptr)
+{
+ if (mca_pml_ucx_worker_progress_cb) {
+ return mca_pml_ucx_worker_progress_cb(worker_ptr);
+ }
+ return;
+}
+
+
+
+
int
opal_init(int* pargc, char*** pargv)
{