diff --git a/.github/workflows/run-ctests.yml b/.github/workflows/run-ctests.yml index 5a75676..ba09265 100644 --- a/.github/workflows/run-ctests.yml +++ b/.github/workflows/run-ctests.yml @@ -95,10 +95,17 @@ jobs: - name: Update LD_LIBRARY_PATH run: echo "LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${{ steps.install-boost.outputs.BOOST_ROOT }}/lib" >> $GITHUB_ENV - name: Install and Test Python CLI + if: ${{ matrix.usempi != 'ON' }} run: | export BOOST_ROOT=${{steps.install-boost.outputs.BOOST_ROOT}} pip install -U . --verbose pysa-stern ${{github.workspace}}/tests/test_n120_t8_k64_mld.txt --max-iters 5000 + - name: Install and Test Python CLI (MPI) + if: ${{ matrix.usempi == 'ON' }} + run: | + export BOOST_ROOT=${{steps.install-boost.outputs.BOOST_ROOT}} + PYSA_USE_MPI=1 pip install -U . --verbose + mpirun -n 2 pysa-stern ${{github.workspace}}/tests/test_n120_t8_k64_mld.txt --max-iters 5000 - name: Configure CMake run: > cmake -B ${{github.workspace}}/build diff --git a/lib/libmld/include/libmld/mld.h b/lib/libmld/include/libmld/mld.h index ce8a56f..39c6a7c 100644 --- a/lib/libmld/include/libmld/mld.h +++ b/lib/libmld/include/libmld/mld.h @@ -37,8 +37,7 @@ class MLDException : public std::exception { enum MLDType { G, H, P, W }; -class MLDProblem { -public: +struct MLDProblem { MLDProblem() = default; MLDProblem(const MLDProblem &) = default; template @@ -109,7 +108,6 @@ class MLDProblem { return yv; } -private: MLDType prob_type{MLDType::G}; size_t nvars = 0; size_t nrows = 0; diff --git a/pyproject.toml b/pyproject.toml index 918b61b..58248c3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,7 @@ classifiers =[ "Operating System :: Unix", "Operating System :: MacOS" ] +dynamic = ["dependencies"] [project.scripts] pysa-stern = "pysa_stern.pysa_stern:PySA_Stern_main" diff --git a/setup.py b/setup.py index 4fdadd8..f00888b 100644 --- a/setup.py +++ b/setup.py @@ -1,10 +1,21 @@ -import sys -import setuptools +import os from pathlib import Path +import setuptools +import sys from cmake_build_extension import CMakeExtension, BuildExtension +# Check if MPI is needed +PYSA_USE_MPI = 'PYSA_USE_MPI' in os.environ + +if PYSA_USE_MPI: + mpi_conf = ["-DMPI=ON"] + mpi_req = ["mpi4py"] + print("Installing PySA-Stern with MPI.") +else: + mpi_conf = [] + mpi_req = [] + print("Installing PySA-Stern without MPI.") -init_py = "" setuptools.setup( ext_modules=[ @@ -16,11 +27,12 @@ cmake_configure_options=[ "-DCMAKE_BUILD_TYPE=Release", f"-DPython_ROOT_DIR={Path(sys.prefix)}", - f"-DPYMODULES=ON", + "-DPYMODULES=ON", "-DCALL_FROM_SETUP_PY:BOOL=ON", - ], + ] + mpi_conf, cmake_component="bindings", ), ], - cmdclass=dict(build_ext=BuildExtension) + cmdclass=dict(build_ext=BuildExtension), + install_requires=["numpy"] + mpi_req ) \ No newline at end of file diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 844bb4f..981f055 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -41,11 +41,15 @@ endif () target_compile_options(sternx PRIVATE "$<$:-Wall;-Wpedantic>") target_compile_options(mldpt PRIVATE "$<$:-Wall;-Wpedantic>") -## Set include directories +## Set include directories and compile definitions if(USE_SIMDE) include_directories(../lib/libmld/lib/simde/) add_compile_definitions(USE_SIMDE) endif () +if(MPI) + target_compile_definitions(sterncpp PUBLIC USEMPI) + target_compile_definitions(sternx PUBLIC USEMPI) +endif() target_include_directories(sterncpp PUBLIC ../lib/libmld/include ../lib/simde diff --git a/src/bindings/CMakeLists.txt b/src/bindings/CMakeLists.txt index 070a5e4..ee73f98 100644 --- a/src/bindings/CMakeLists.txt +++ b/src/bindings/CMakeLists.txt @@ -23,6 +23,10 @@ find_package(pybind11 REQUIRED) # Build the Python extension bindings module pybind11_add_module(pysa_stern_ext pysa_stern.cpp) +if(MPI) + target_compile_definitions(pysa_stern_ext PUBLIC USEMPI) +endif() + target_include_directories(pysa_stern_ext PRIVATE ../../include/) target_include_directories(pysa_stern_ext PRIVATE ../../lib/libmld/include) target_link_libraries(pysa_stern_ext PUBLIC sterncpp) diff --git a/src/bindings/pysa_stern.cpp b/src/bindings/pysa_stern.cpp index c4a9424..366701e 100644 --- a/src/bindings/pysa_stern.cpp +++ b/src/bindings/pysa_stern.cpp @@ -21,6 +21,10 @@ specific language governing permissions and limitations under the License. #include +#define MEMBER2PYDICT(_D,_CL,_M) _D[#_M] = _CL._M; + +#define PYDICT2MEMBER(_D,_CL,_M) _CL._M = _D[#_M].cast(); +#define PYTUP2MEMBER(_T,_I,_CL,_M) _CL._M = _T[_I].cast(); namespace py = pybind11; // Wrapper class to make solution accessible as a numpy array @@ -74,6 +78,30 @@ void init_mld_h(py::module& m){ .def("Weight", &MLDProblem::Weight) .def("clause_list", &MLDProblem::clause_list) .def("read_problem_str", &MLDProblem::read_problem_string) + .def(py::pickle( + [](const MLDProblem& mldp){ + return py::make_tuple( + mldp.prob_type, + mldp.nvars, + mldp.nrows, + mldp.yarr, + mldp.clauses, + mldp.y, + mldp.w + ); + }, + [](py::tuple t){ + MLDProblem mldp; + PYTUP2MEMBER(t, 0, mldp, prob_type); + PYTUP2MEMBER(t, 1, mldp, nvars); + PYTUP2MEMBER(t, 2, mldp, nrows); + PYTUP2MEMBER(t, 3, mldp, yarr); + PYTUP2MEMBER(t, 4, mldp, clauses); + PYTUP2MEMBER(t, 5, mldp, y); + PYTUP2MEMBER(t, 6, mldp, w); + return mldp; + } + )) .doc() = "Options definding a MLD problem." ; } @@ -106,6 +134,34 @@ PySA-sternx: Stern algorithm for unstructured decoding problems. .def_readwrite("p", &sternc_opts::p) .def_readwrite("m", &sternc_opts::m) .def_readwrite("block_size", &sternc_opts::block_size) + .def(py::pickle( + [](const sternc_opts& opts){ + py::dict d; + d["parcheck"] = opts.parcheck; + d["bench"] = opts.bench; + d["test_hw1"] = opts.test_hw1; + d["t"] = opts.t; + d["max_iters"] = opts.max_iters; + d["l"] = opts.l; + d["p"] = opts.p; + d["m"] = opts.m; + d["block_size"] = opts.block_size; + return d; + }, + [](py::dict d){ + sternc_opts opts; + PYDICT2MEMBER(d, opts, parcheck); + PYDICT2MEMBER(d, opts, bench); + PYDICT2MEMBER(d, opts, test_hw1); + PYDICT2MEMBER(d, opts, t); + PYDICT2MEMBER(d, opts, max_iters); + PYDICT2MEMBER(d, opts, l); + PYDICT2MEMBER(d, opts, p); + PYDICT2MEMBER(d, opts, m); + PYDICT2MEMBER(d, opts, block_size); + return opts; + } + )) .doc() = "Options for the Stern algorithm routine." //.def_readwrite("nvars", &sternc_opts::nvars) ; @@ -123,5 +179,8 @@ PySA-sternx: Stern algorithm for unstructured decoding problems. ; m.def("sterncpp_adjust_opts", &sterncpp_adjust_opts); m.def("sterncpp_main", &sterncpp_main); +#ifdef USEMPI + m.attr("__pysa_mpi__") = 1; +#endif init_libmld_submod(m); } \ No newline at end of file diff --git a/src/pysa_stern/pysa_stern.py b/src/pysa_stern/pysa_stern.py index 8ebb85f..77d62c1 100644 --- a/src/pysa_stern/pysa_stern.py +++ b/src/pysa_stern/pysa_stern.py @@ -1,11 +1,24 @@ import argparse import sys -import pathlib from pysa_stern.mld import MLDProblem from pysa_stern.bindings import SternOpts, sterncpp_adjust_opts from pysa_stern.sterncpp import pysa_sternx +import numpy as np +try: + # Make sure bindings were compiled with MPI + from pysa_stern.bindings import __pysa_mpi__ + # Check that mpi4py is installed + from mpi4py import MPI + comm = MPI.COMM_WORLD + rank = comm.Get_rank() +except: + comm = None + rank = 0 + def PySA_Stern_main(): + adjusted_opts = None + parser = argparse.ArgumentParser("PySA-Stern") parser.add_argument("input", help="MLD format input file") # Stern algorithm parameters @@ -46,16 +59,13 @@ def PySA_Stern_main(): action='store_true', ) + if(comm is not None and rank == 0): + print("--- PySA-Stern (MPI) ---") + elif comm is None: + print("--- PySA-Stern---") args = parser.parse_args() - stern_opts = SternOpts() - problem = MLDProblem() - try: - with open(args.input) as f: - problem.read_problem(f) - except Exception as e: - print(f"Problem reading MLD file ({args.input}).", file=sys.stderr) - print(e, file=sys.stderr) - exit(1) + # Gather and check options + stern_opts = SternOpts() stern_opts.l = args.col_size if args.col_size is not None else -1 stern_opts.m = args.col_sets stern_opts.block_size = args.block_size @@ -67,19 +77,58 @@ def PySA_Stern_main(): except RuntimeError as e: print(f"Problem with passed options.", file=sys.stderr) print(e, file=sys.stderr) + adjusted_opts = None + if comm is not None: + MPI.Finalize() exit(1) - sol_arr = pysa_sternx(problem, adjusted_opts) - if sol_arr is not None: - print(f"[PySA-Stern] Solution found: ") - for i in range(len(sol_arr)): - print(f"{sol_arr[i]} ", end="") - if (i+1)%32 == 0: - print("\n", end="") - elif (i+1)%8 == 0: - print(" ", end="") - print() + # Read in the MLD problem + problem = MLDProblem() + + problem_ok = np.asarray([1]) + if rank == 0: + try: + with open(args.input, 'r') as f: + problem_string = f.read() + except Exception as e: + problem_ok[0] = 0 + print(f"Problem reading MLD file ({args.input}).", file=sys.stderr) + print(e, file=sys.stderr) else: - print(f"[PySA-Stern] Solution not found.") + problem_string = None + if comm is not None: + comm.Bcast(problem_ok, root=0) + if not problem_ok[0]: + if comm is not None: + MPI.Finalize() + exit(1) + if comm is not None: + problem_string = comm.bcast(problem_string, root=0) + try: + problem = MLDProblem() + problem.read_problem_str(problem_string) + except Exception as e: + problem_ok[0] = 0 + if rank==0: + print(f"Problem with MLD file ({args.input}).", file=sys.stderr) + print(e, file=sys.stderr) + exit(1) + # Call the Stern algorithm solver + sol_arr = pysa_sternx(problem, adjusted_opts) + + if rank==0: + if sol_arr is not None: + print(f"[PySA-Stern] Solution found: ") + for i in range(len(sol_arr)): + print(f"{sol_arr[i]} ", end="") + if (i+1)%32 == 0: + print("\n", end="") + elif (i+1)%8 == 0: + print(" ", end="") + print() + else: + print(f"[PySA-Stern] Solution not found.") + # if comm is not None: + # MPI.Finalize() if __name__ == "__main__": PySA_Stern_main() \ No newline at end of file