diff options
| author | kguerda-idris <84066930+kguerda-idris@users.noreply.github.com> | 2021-09-29 15:29:31 +0200 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2021-09-29 15:29:31 +0200 |
| commit | 1c7e7ce2da8bb362c184fb6eae71fe7e36356494 (patch) | |
| tree | 92fdc31870b6d5384c8ba83ff72d85a2d5a1eee6 /ot/lp/emd_wrap.pyx | |
| parent | 7dde9e8e4b6aae756e103d49198caaa4f24150e3 (diff) | |
[MRG] OpenMP support (#260)
* Added : OpenMP support
Restored : Epsilon and Debug mode
Replaced : parmap => multiprocessing is now replace by multithreading
* Commit clean up
* Number of CPUs correctly calculated on SLURM clusters
* Corrected number of processes for cluster slurm
* Mistake corrected
* parmap is now deprecated
* Now a different solver is used depending on the requested number of threads
* Tiny mistake corrected
* Folders are now in the ot library instead of at the root
* Helpers is now correctly placed
* Attempt to make compilation work smoothly
* OS compatible path
* NumThreads now defaults to 1
* Better flags
* Mistake corrected in case of OpenMP unavailability
* Revert OpenMP flags modification, which do not compile on Windows
* Test helper functions
* Helpers comments
* Documentation update
* File title corrected
* Warning no longer using print
* Last attempt for macos compilation
* pls work
* atempt
* solving a type error
* TypeError OpenMP
* Compilation finally working on Windows
* Bug solve, number of threads now correctly selected
* 64 bits solver to avoid overflows for bigger problems
* 64 bits EMD corrected
Co-authored-by: kguerda-idris <ssos023@jean-zay3.idris.fr>
Co-authored-by: ncassereau-idris <84033440+ncassereau-idris@users.noreply.github.com>
Co-authored-by: ncassereau <nathan.cassereau@idris.fr>
Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>
Diffstat (limited to 'ot/lp/emd_wrap.pyx')
| -rw-r--r-- | ot/lp/emd_wrap.pyx | 9 |
1 files changed, 6 insertions, 3 deletions
diff --git a/ot/lp/emd_wrap.pyx b/ot/lp/emd_wrap.pyx index de9a700..42e08f4 100644 --- a/ot/lp/emd_wrap.pyx +++ b/ot/lp/emd_wrap.pyx @@ -20,6 +20,7 @@ import warnings cdef extern from "EMD.h": int EMD_wrap(int n1,int n2, double *X, double *Y,double *D, double *G, double* alpha, double* beta, double *cost, int maxIter) nogil + int EMD_wrap_omp(int n1,int n2, double *X, double *Y,double *D, double *G, double* alpha, double* beta, double *cost, int maxIter, int numThreads) nogil cdef enum ProblemType: INFEASIBLE, OPTIMAL, UNBOUNDED, MAX_ITER_REACHED @@ -38,7 +39,7 @@ def check_result(result_code): @cython.boundscheck(False) @cython.wraparound(False) -def emd_c(np.ndarray[double, ndim=1, mode="c"] a, np.ndarray[double, ndim=1, mode="c"] b, np.ndarray[double, ndim=2, mode="c"] M, int max_iter): +def emd_c(np.ndarray[double, ndim=1, mode="c"] a, np.ndarray[double, ndim=1, mode="c"] b, np.ndarray[double, ndim=2, mode="c"] M, int max_iter, int numThreads): """ Solves the Earth Movers distance problem and returns the optimal transport matrix @@ -109,8 +110,10 @@ def emd_c(np.ndarray[double, ndim=1, mode="c"] a, np.ndarray[double, ndim=1, mod # calling the function with nogil: - result_code = EMD_wrap(n1, n2, <double*> a.data, <double*> b.data, <double*> M.data, <double*> G.data, <double*> alpha.data, <double*> beta.data, <double*> &cost, max_iter) - + if numThreads == 1: + result_code = EMD_wrap(n1, n2, <double*> a.data, <double*> b.data, <double*> M.data, <double*> G.data, <double*> alpha.data, <double*> beta.data, <double*> &cost, max_iter) + else: + result_code = EMD_wrap_omp(n1, n2, <double*> a.data, <double*> b.data, <double*> M.data, <double*> G.data, <double*> alpha.data, <double*> beta.data, <double*> &cost, max_iter, numThreads) return G, cost, alpha, beta, result_code |