PySA-Stern Python MPI support

.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

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<typename IS_t>
@@ -109,7 +108,6 @@ class MLDProblem {
     return yv;
   }
 
-private:
   MLDType prob_type{MLDType::G};
   size_t nvars = 0;
   size_t nrows = 0;

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"

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
 )

src/CMakeLists.txt

@@ -41,11 +41,15 @@ endif ()
 target_compile_options(sternx PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:-Wall;-Wpedantic>")
 target_compile_options(mldpt PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:-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

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)

src/bindings/pysa_stern.cpp

@@ -21,6 +21,10 @@ specific language governing permissions and limitations under the License.
 
 #include <sternx/stern.h>
 
+#define MEMBER2PYDICT(_D,_CL,_M) _D[#_M] = _CL._M;
+
+#define PYDICT2MEMBER(_D,_CL,_M) _CL._M = _D[#_M].cast<decltype(_CL._M)>();
+#define PYTUP2MEMBER(_T,_I,_CL,_M) _CL._M = _T[_I].cast<decltype(_CL._M)>();
 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);
 }

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()