Mpi (#93)

* draft v1 * clippy * correctness * fmt * clippy * add mpi finalize * clippy * refactor; try make mpi compatible with threading * add finalize * clippy * try ci install * update * try mpi install * try mpi install * try mpi install * try mpi install * refactor * clippy * identify structure info * clippy * main-mpi * mpi_config * mpiexec in ci * intentially fail ci for mpi * now correct * minor in toml * Update ci.yml Signed-off-by: Tiancheng Xie <[email protected]> * allow static mut ref * Update build.yml Signed-off-by: Tiancheng Xie <[email protected]> * Update ci.yml Signed-off-by: Tiancheng Xie <[email protected]> * Merge CI change (#106) * [FIX] m31 exp (#96) * chore: test vectors for M31 * fix: M31 exponentiation --------- Co-authored-by: enpsi <[email protected]> * fix ci (#104) --------- Co-authored-by: 0xwework <[email protected]> Co-authored-by: enpsi <[email protected]> Co-authored-by: zhenfei <[email protected]> --------- Signed-off-by: Tiancheng Xie <[email protected]> Co-authored-by: Tiancheng Xie <[email protected]> Co-authored-by: 0xwework <[email protected]> Co-authored-by: enpsi <[email protected]> Co-authored-by: zhenfei <[email protected]>
PolyhedraZK · Sep 25, 2024 · 3df9ca0 · 3df9ca0
1 parent 10215e1
commit 3df9ca0
Show file tree

Hide file tree

Showing 33 changed files with 1,634 additions and 848 deletions.
diff --git a/.github/workflows/benchmarks.yml b/.github/workflows/benchmarks.yml
@@ -27,6 +27,8 @@ jobs:
         project_id: ${{ secrets.GCP_PROJECT_ID }}
     - name: Install criterion
       run: cargo install cargo-criterion
+    - name: Install mpi
+      run: python3 ./scripts/install.py
     - name: setup
       run: cargo run --bin=dev-setup --release
     - name: Run benchmark

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -11,7 +11,7 @@ jobs:
   build:
     strategy:
       matrix:
-        platform: [macos, linux-avx2]
+        platform: [macos, 7950x3d]
         include:
           - platform: macos
             ci_image: macos-latest
@@ -24,6 +24,8 @@ jobs:
       - uses: actions/checkout@v4
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
+      - name: Install MPI
+        run: python3 ./scripts/install.py
       - name: Prepare binary
         run: ${{ matrix.flags }} cargo build --release --bin expander-exec
       - name: Upload release asset

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -9,15 +9,15 @@ env:
 jobs:
   lint:
     name: Lint
-    runs-on: ubuntu-latest
+    runs-on: 7950x3d
     steps:
       - uses: actions/checkout@v4
       - uses: dtolnay/rust-toolchain@stable
         with:
           components: rustfmt, clippy
+      - run: python3 ./scripts/install.py
       - run: cargo fmt --all -- --check
       - run: cargo clippy
-
   build-and-test:
     name: Build and Test (${{ matrix.os }}${{ matrix.feature != '' && format(', {0}', matrix.feature) || '' }})
     runs-on: ${{ matrix.os }}
@@ -38,6 +38,7 @@ jobs:
         run: echo "RUSTFLAGS=$RUSTFLAGS -C target-feature=+${{ matrix.feature }}" >> $GITHUB_ENV
       - name: Build and Test
         run: |
+          python3 ./scripts/install.py
           cargo run --bin=dev-setup --release
           cargo build --all-features --release
           cargo test --all-features --release      
@@ -83,6 +84,8 @@ jobs:
         run: echo "RUSTFLAGS=$RUSTFLAGS -C target-feature=+${{ matrix.feature }}" >> $GITHUB_ENV
       - name: Run benchmark
         run: |
+          python3 ./scripts/install.py
           cargo run --bin=dev-setup --release
           cargo run --release -- -t ${{ matrix.os == 'macos-latest' && 2 || 16 }} -f ${{ matrix.field }}
 
+
diff --git a/Cargo.toml b/Cargo.toml
@@ -10,6 +10,7 @@ ark-std.workspace = true
 clap.workspace = true
 env_logger.workspace = true
 log.workspace = true
+mpi.workspace = true
 rand.workspace = true
 sha2.workspace = true
 halo2curves.workspace = true
@@ -26,6 +27,10 @@ tiny-keccak.workspace = true
 [dev-dependencies]
 criterion = "0.5.1"
 
+[[bin]]
+name = "expander-rs-mpi"
+path = "src/main_mpi.rs"
+
 [[bin]]
 name = "expander-exec"
 path = "src/exec.rs"
@@ -58,6 +63,7 @@ halo2curves = { git = "https://github.com/PolyhedraZK/halo2curves", default-feat
 ] }
 itertools = "0.13"
 log = "0.4"
+mpi = "0.8.0"
 rand = "0.8.5"
 rayon = "1.10"
 sha2 = "0.10.8"

diff --git a/benches/gkr_hashes.rs b/benches/gkr_hashes.rs
@@ -1,7 +1,7 @@
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use expander_rs::utils::{KECCAK_M31_CIRCUIT, POSEIDON_CIRCUIT};
 use expander_rs::{
-    BN254ConfigSha2, Circuit, Config, GKRConfig, GKRScheme, M31ExtConfigSha2, Prover,
+    BN254ConfigSha2, Circuit, Config, GKRConfig, GKRScheme, M31ExtConfigSha2, MPIConfig, Prover,
 };
 use std::hint::black_box;
 
@@ -12,7 +12,7 @@ fn prover_run<C: GKRConfig>(config: &Config<C>, circuit: &mut Circuit<C>) {
 }
 
 fn benchmark_setup<C: GKRConfig>(scheme: GKRScheme, circuit_file: &str) -> (Config<C>, Circuit<C>) {
-    let config = Config::<C>::new(scheme);
+    let config = Config::<C>::new(scheme, MPIConfig::new());
     let mut circuit = Circuit::<C>::load_circuit(circuit_file);
     circuit.set_random_input_for_test();
     (config, circuit)

diff --git a/scripts/install.py b/scripts/install.py
@@ -0,0 +1,8 @@
+import subprocess
+from sys import platform
+
+if __name__ == "__main__":
+    if platform == "darwin": # mac os
+        subprocess.run(["brew", "install", "openmpi"])
+    else:
+        pass # Do nothing, assuming mpi has already been installed
diff --git a/scripts/run_multiple_mpi.py b/scripts/run_multiple_mpi.py
@@ -0,0 +1,44 @@
+import sys
+import json
+import subprocess
+
+MPI_CONFIG_JSON = '''
+{
+    "field": "gf2ext128",
+    "n_groups": 2,
+    "mpi_size_each_group": 8,
+    "cpu_ids":
+        [
+            [0, 1, 2, 3, 4, 5, 6, 7],
+            [8, 9, 10, 11, 12, 13, 14, 15]
+        ]
+}
+'''
+
+def parse_config(mpi_config):
+    field = mpi_config["field"]
+    n_groups = mpi_config["n_groups"]
+    mpi_size_each_group = mpi_config["mpi_size_each_group"]
+    cpu_ids = mpi_config["cpu_ids"]
+
+    if field not in ["gf2ext128", "m31ext3", "fr"]:
+        sys.exit("Unrecognized field, now only supports gf2ext128, m31ext3 and fr")
+
+    if n_groups != len(cpu_ids):
+        sys.exit("Lack/Too much cpu specifications.")
+
+    for i in range(n_groups):
+        if len(cpu_ids[i]) != mpi_size_each_group:
+            sys.exit(f"Cpu ids are not correct for group {i}")
+
+    return field, n_groups, mpi_size_each_group, cpu_ids
+
+
+# Run two mpi process
+if __name__ == "__main__":
+    mpi_config = json.loads(MPI_CONFIG_JSON)
+    field, n_groups, mpi_size_each_group, cpu_ids = parse_config(mpi_config)
+
+    for i in range(n_groups):
+        cpu_id = ",".join(map(str, cpu_ids[i]))
+        subprocess.Popen(["mpiexec", "-cpu-set", cpu_id, "-n", str(mpi_size_each_group), "./target/release/expander-rs-mpi", "-f", field])