Introducing MPI shared memory based weight pre-processing

.github/workflows/documentation.yaml

@@ -12,16 +12,16 @@ jobs:
     if: "!contains(github.event.head_commit.message, 'skip docs')"
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5.0.0
+      - uses: actions/checkout@v7
       - name: Configure Git Credentials
         run: |
           git config user.name github-actions[bot]
           git config user.email 41898282+github-actions[bot]@users.noreply.github.com
-      - uses: actions/setup-python@v6.0.0
+      - uses: actions/setup-python@v7
         with:
           python-version: 3.11
       - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV 
-      - uses: actions/cache@v4.3.0
+      - uses: actions/cache@v6
         with:
           key: mkdocs-material-${{ env.cache_id }}
           path: .cache

.github/workflows/tests.yaml

@@ -10,6 +10,7 @@ jobs:
   test:
     if: "!contains(github.event.head_commit.message, 'skip ci')"
     runs-on: ${{ matrix.os }}
+    timeout-minutes: 30
 
     env:
       MACOSX_DEPLOYMENT_TARGET: "10.15"  # To be removed when support for Python 3.11 is dropped
@@ -23,17 +24,17 @@ jobs:
 
     steps:
       - name: Checkout BrahMap repository
-        uses: actions/checkout@v6.0.2
+        uses: actions/checkout@v7
         with:
           submodules: true
 
       - name: Install MPI - ${{ matrix.mpi }}
-        uses: mpi4py/setup-mpi@v1.4.2
+        uses: mpi4py/setup-mpi@v1.4.4
         with:
           mpi: ${{ matrix.mpi }}
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v6.2.0
+        uses: actions/setup-python@v7
         with:
           python-version: ${{ matrix.python-version }}
 

.pre-commit-config.yaml

@@ -29,6 +29,7 @@ repos:
   hooks:
   - id: check-toml
   - id: check-yaml
+    args: [--unsafe]
   - id: check-merge-conflict
 
 - repo: local
@@ -37,3 +38,5 @@ repos:
     name: Building the markdown documentation with `mkdocs build`
     entry: bash -c 'mkdocs build'
     language: system
+    pass_filenames: false
+    require_serial: true

README.md

@@ -23,10 +23,6 @@
 ![BrahMap documentation build status](https://github.com/anand-avinash/BrahMap/actions/workflows/documentation.yaml/badge.svg)
 <!-- markdownlint-enable MD013 -->
 
-<!-- markdownlint-disable MD033 -->
-<font color="red"> **This project is currently under active development!!!** </font>
-<!-- markdownlint-enable MD033 -->
-
 `BrahMap` is a scalable and modular map-making framework for CMB experiments.
 It features a user-friendly Python interface for the linear operators used in
 map-making. The Python interface seamlessly handles the workflow while
@@ -43,8 +39,9 @@ For a complete reference to the `BrahMap` API, refer to the
 Complete example notebooks and scripts can be
 [found here](./examples).
 
-You can find detailed information on the implementation and features of
-`BrahMap` at [arXiv:2501.16122](https://arxiv.org/abs/2501.16122).
+For detailed information on the implementation and features of `BrahMap`,
+please refer to the paper
+[arXiv:2501.16122](https://arxiv.org/abs/2501.16122).
 
 ## Installation
 
@@ -91,7 +88,7 @@ of its numerical routines and linear operators. The benchmarks are based on
 `pytest-benchmark` and support both serial and MPI-parallel execution.
 
 For detailed instructions on how to run and customize the benchmarks, refer to
-the [benchmarks documentation](./benchmarks/README.md).
+the [benchmarks documentation](./benchmarks).
 
 ## Citation
 
@@ -100,7 +97,7 @@ This work can be cited with:
 <!-- markdownlint-disable MD013 -->
 ```text
 @misc{anand2025brahmap,
-      title={\texttt{BrahMap}: A scalable and modular map-making framework for the CMB experiments}, 
+      title={BrahMap: A scalable and modular map-making framework for the CMB experiments}, 
       author={Avinash Anand and Giuseppe Puglisi},
       year={2025},
       eprint={2501.16122},

benchmarks/README.md

@@ -1,3 +1,5 @@
+<!-- markdownlint-disable MD013 -->
+
 # Benchmarking BrahMap
 
 This directory contains scripts for benchmarking the performance of BrahMap components. While the benchmark scripts are primarily designed for serial execution, they can also be used for parallel execution with `mpirun/mpiexec`, though one has to be careful while saving the benchmark results.
@@ -55,12 +57,16 @@ This allows for quick sanity checks with `small` or more robust performance meas
 
 You can override specific parameters regardless of the `--size` flag:
 
-| Option          | Values               | Description                                               |
-| :-------------- | :------------------- | :-------------------------------------------------------- |
-| `--nside`       | e.g. `128`           | Sets $N_{side}$ (overrides npix)                        |
-| `--nsamples`    | e.g. `1000000`       | Sets global number of samples (overrides global nsamples) |
-| `--dtype-float` | `float32`, `float64` | Sets floating point precision (default: float64)          |
-| `--dtype-int`   | `int32`, `int64`     | Sets integer precision (default: int64)                   |
+| Option             | Values               | Description                                               |
+| :----------------- | :------------------- | :-------------------------------------------------------- |
+| `--nside`          | e.g. `128`           | Sets $N_{side}$ (overrides npix)                        |
+| `--nsamples`       | e.g. `1000000`       | Sets global number of samples (overrides global nsamples) |
+| `--dtype-float`    | `float32`, `float64` | Sets floating point precision (default: float64)          |
+| `--dtype-int`      | `int32`, `int64`     | Sets integer precision (default: int64)                   |
+| `--mpi-rounds`        | e.g. `20`            | Sets number of rounds for `mpi_benchmark` (default: 20)      |
+| `--mpi-iterations`    | e.g. `1`             | Sets iterations per round for `mpi_benchmark` (default: 1)   |
+| `--mpi-warmup-rounds` | e.g. `0`             | Sets number of warmup rounds for `mpi_benchmark` (default: 0)|
+| `--nproc-reduce`      | e.g. `1`             | Sets the size of each sub-communicator group witin the node-level communicator (default: 1)  |
 
 Example usage:
 
@@ -84,6 +90,15 @@ pytest-benchmark compare results_v1.json results_v2.json
 
 This will show a detailed comparison table with percentage differences. You can also use the option `--sort=...` to sort the comparison results.
 
+### Rounds vs. Iterations
+
+- **Iterations (per round):** The number of times the function is called consecutively **within a single timed measurement round**. The iteration count should be set according to the function being benchmarked:
+  - For fast functions (taking microseconds), a higher value (e.g., `100` or `1000`) is more suitable to amortize timing overhead.
+  - For computationally heavy/slower functions (taking milliseconds/seconds), a lower value (e.g., `1` or `5`) is sufficient.
+  - For MPI-related benchmarks, the default number of iterations is `1`. It can be overridden using `--mpi-iterations=<num_iterations>`.
+- **Rounds:** The number of independent measurements taken to compute statistics. A round **constitutes a single timed measurement event**. This should be set to achieve a balance between total run time and the statistical stability of the measurements.
+  - For MPI-related tests, the default number of rounds is `20`. It can be overridden using `--mpi-rounds=<num_rounds>`.
+
 ## Parallel Execution (with MPI)
 
 To run benchmarks in parallel:
@@ -102,22 +117,36 @@ For collective operations (like `*_rmult`), it is recommended to call the indivi
 
 New benchmarks should follow the class-based structure as in `test_bench_extensions.py`. This groups related benchmarks in the final report.
 
+### 1. Choosing the Right Benchmark Fixture
+
+When writing a benchmark, request the appropriate fixture depending on the operations being measured:
+
+- **`benchmark` (Standard):** For non-MPI / local operations (e.g. math functions, local matrix multiplications). It uses dynamic calibration to scale rounds and iterations automatically.
+
+- **`mpi_benchmark` (Custom):** For collective MPI or shared-memory operations (e.g., projection operators, weight computation). It uses fixed numbers of rounds and iterations (default: 20 rounds of 1 iteration) to prevent rank desynchronization.
+
+### 2. Writing the Benchmark Code
+
 ```python
+# Case A: Standard local benchmark using 'benchmark'
 @pytest.mark.benchmark(group="extensions::my_group")
 class TestMyComponent:
-    def test_bench_my_function_1(self, benchmark, data):
-        # Setup randomized buffers using the module-level data fixture
-        rng = data["rng"]
-        vec = rng.random(data["npix"]).astype(data["dtype_float"])
-
-        # Run benchmark
-        benchmark(my_function1, vec)
+    def test_bench_my_function(self, benchmark, data):
+        vec = data["rng"].random(data["npix"]).astype(data["dtype_float"])
+        benchmark(my_function, vec)
 
-    def test_bench_my_function_2(self, benchmark, data):
-        # Setup randomized buffers using the module-level data fixture
-        rng = data["rng"]
-        vec = rng.random(data["npix"]).astype(data["dtype_float"])
+# Case B: MPI collective benchmark using 'mpi_benchmark' with setup reset
+@pytest.mark.benchmark(group="extensions::my_group")
+class TestMyMPIComponent:
+    def test_bench_my_mpi_function(self, mpi_benchmark, data):
+        prod = np.zeros(data["npix"])
+        vec = data["rng"].random(data["npix"]).astype(data["dtype_float"])
 
-        # Run benchmark
-        benchmark(my_function2, vec)
+        # Setup runs before each round to reset arrays and prevent numerical 
+        # overflow
+        def setup():
+            prod.fill(0)
+            return (data["npix"], vec, prod), {}
+
+        mpi_benchmark(my_mpi_function, setup=setup)
 ```

benchmarks/conftest.py

@@ -35,6 +35,39 @@ def pytest_addoption(parser):
         type=int,
         help="Nside for the healpix map",
     )
+    parser.addoption(
+        "--mpi-rounds",
+        action="store",
+        type=int,
+        default=20,
+        help="Number of rounds for mpi_benchmark",
+    )
+    parser.addoption(
+        "--mpi-iterations",
+        action="store",
+        type=int,
+        default=1,
+        help="Number of iterations for mpi_benchmark",
+    )
+    parser.addoption(
+        "--mpi-warmup-rounds",
+        action="store",
+        type=int,
+        default=0,
+        help="Number of warmup rounds for mpi_benchmark",
+    )
+    parser.addoption(
+        "--nproc-reduce",
+        action="store",
+        type=int,
+        default=1,
+        help="Number of processes used in parallel Reduction within nodes for shared memory mode",
+    )
+
+
+@pytest.fixture(scope="session")
+def nproc_reduce(request):
+    return request.config.getoption("--nproc-reduce")
 
 
 @pytest.fixture(scope="session")
@@ -95,3 +128,48 @@ def pytest_report_header(config):
 def pytest_benchmark_update_machine_info(config, machine_info):
     """Add resolved benchmark parameters to the JSON report."""
     machine_info["benchmark_params"] = _resolve_params(config)
+
+
+@pytest.fixture
+def mpi_benchmark(benchmark, request):
+    """A fixture to enable pedantic benchmark in order to enforce fixed
+    iterations and rounds across all MPI ranks.
+    """
+    # For a generic benchmark, the rounds and the iterations are determined
+    # individually by each MPI process. If the number of iterations or the
+    # rounds are not equal across all MPI ranks, the standard benchmark runs
+    # in segmentation faults. Setting the rounds and the iterations to the
+    # same value for all MPI ranks avoids this issue.
+
+    # Retrieve default CLI overrides
+    cli_rounds = request.config.getoption("--mpi-rounds")
+    cli_iterations = request.config.getoption("--mpi-iterations")
+    cli_warmup = request.config.getoption("--mpi-warmup-rounds")
+
+    def _run(func, *args, **kwargs):
+        # Extract benchmark configuration, defaulting to CLI values (or their defaults)
+        rounds = kwargs.pop("rounds", cli_rounds)
+        iterations = kwargs.pop("iterations", cli_iterations)
+        warmup_rounds = kwargs.pop("warmup_rounds", cli_warmup)
+        setup = kwargs.pop("setup", None)
+        teardown = kwargs.pop("teardown", None)
+
+        # Setup function when used in benchmark.pedantic, can also be used to
+        # supply the benchmark parameters. The setup functions is called at
+        # the first iteration of every round. In some of the cases, we need
+        # to supply the zero-ed arrays as the function/class arguments every
+        # once in a while to prevent overflow/underflow. There, we can use
+        # the setup argument to do so (for example, in `test_bench_extensions.py`)
+        # See <https://pytest-benchmark.readthedocs.io/en/latest/pedantic.html#reference>
+        return benchmark.pedantic(
+            func,
+            args=args,
+            kwargs=kwargs,
+            setup=setup,
+            teardown=teardown,
+            iterations=iterations,
+            rounds=rounds,
+            warmup_rounds=warmup_rounds,
+        )
+
+    return _run