border_op — Halo Communication Operator for LAMMPS Domain Decomposition

中文 README

border_op is a key communication operator (Halo / border exchange) extracted from the MatterSim / OmniMat AI molecular dynamics (AI-MD) application, open-sourced in a clean, standalone form for reuse in other AI-MD projects.

In MPI domain-decomposed molecular dynamics simulations, each process owns only nlocal local atoms, while the coordinates of nghost ghost atoms are maintained by LAMMPS. End-to-end differentiable AI potentials (GNN potentials such as MatterSim) need atom features to be available in the ghost region during message passing, and during backpropagation the gradients on ghost atoms must be accumulated back onto the local atoms of their owner processes. border_op does exactly this:

• Forward: equivalent to LAMMPS forward communication — following the CommBrick swap tables, sends features of local atoms to neighboring processes and fills the local ghost region;
• Backward: equivalent to LAMMPS reverse communication — sends the gradients on the ghost region back to the owner processes and accumulates them onto the corresponding local atoms via index_add_.

The operator is registered as a torch.autograd.Function, so it can be embedded directly into the PyTorch computation graph; forward/backward communication is transparent to both training and inference.

How It Works

The communication schedule fully reuses the existing swap tables of LAMMPS CommBrick, without introducing any additional topology-partitioning logic:

swap iswap = 0 ... nswap-1
  this rank -> sendproc[iswap] : rows sendlist[iswap][0:sendnum[iswap]] of the features
  this rank <- recvproc[iswap] : append recvnum[iswap] rows into the ghost region in order
  if sendproc[iswap] == this rank (periodic image within the same process), copy locally

make_border_pack_uint64 in lammps_bridge.cc reads these tables (sendlist / sendproc / recvproc / sendnum / recvnum / nswap), together with nlocal / nghost and the MPI_Comm handle, from a running LAMMPS instance pointer, and packs them into 8 CPU tensors (the "pack"). Each forward pass then only needs to hand the pack and the feature tensor to border_op. The backward pass swaps the send/recv roles and traverses the swaps in reverse order, yielding reverse communication.

Directory Layout

border_op/
├── csrc/
│   ├── border_op.cc        # The halo operator itself (forward comm / reverse comm backward)
│   ├── lammps_bridge.cc    # Bridge layer extracting swap tables from LAMMPS CommBrick
│   └── device.h            # CUDA / HIP portability shim
├── python/
│   └── halo.py             # PyTorch-side wrapper: init() + halo()
├── tests/
│   └── test_border_op.py   # Two-process PBC forward + backward correctness test
├── Makefile
├── LICENSE                 # MIT (inherited from MatterSim)
└── README.md

Dependencies

• PyTorch (C++ extension; must match the compile-time ABI — the Makefile auto-detects _GLIBCXX_USE_CXX11_ABI)
• MPI (OpenMPI or MPICH, mpicxx)
• LAMMPS (only required for bridge_op.so: the src/ header directory and liblammps_mpi; the lammps-python package is needed at runtime)
• Optional: CUDA (NVIDIA) or ROCm/DTK (AMD/DCU)

Build

# CPU + MPI
make

# NVIDIA GPU (CUDA-aware MPI, communicates directly on device pointers)
make GPU=cuda CUDA_HOME=/usr/local/cuda

# GPU without CUDA-aware MPI: stage through the host before communicating
make GPU=cuda NO_CUDA_AWARE=1

# AMD / DCU (HIP)
make GPU=hip ROCM_PATH=/opt/rocm

# Log forward/backward communication timings to comm_timing.txt
make TIMING=1

# LAMMPS bridge library (requires LAMMPS headers and library)
make bridge LAMMPS_SRC=/path/to/lammps/src LAMMPS_LIB_DIR=/path/to/lammps/lib

Artifacts are written to build/border_op.so and build/bridge_op.so. make show prints the actual compile flags used, which helps with debugging.

Usage

import sys
sys.path.insert(0, "python")   # or add python/ to PYTHONPATH
import torch
import halo
from lammps import lammps

lmp = lammps()
# ... build the box, create atoms, set the potential ...
lmp.command("run 0")           # let the comm tables get built first (nswap > 0)

halo.init(lmp)                 # grab the CommBrick swap tables

n = halo._nlocal + halo._nghost
g1 = torch.zeros(n, 64, device="cuda")   # atom features (float32/float64 both OK)
g1 = halo.halo(g1)             # forward: fill the ghost region
# backward communication happens automatically on loss.backward()

Notes:

• Only comm_style brick (CommBrick) is supported; make_border_pack_uint64 raises an error otherwise.
• halo.init() must be called after the neighbor/communication tables have been built (e.g., after run 0). If the domain decomposition or the atom count changes during the simulation (repartitioning, or atom migration changing nlocal/nghost), call halo.refresh_pack_pair() to re-grab the tables.
• The input feature tensor has layout (nlocal + nghost, K), with the ghost region occupying the last nghost rows, consistent with the LAMMPS atom array layout. float32 and float64 are supported.
• Shared libraries are loaded from <repo>/build/*.so by default; this can be overridden with the environment variables BORDER_OP_SO / BRIDGE_OP_SO.
• lammps_bridge.cc reads the protected members of CommBrick via #define protected public — a pragmatic hack; if LAMMPS internals change, this may need to be adjusted accordingly.

Testing

make && make bridge LAMMPS_SRC=... LAMMPS_LIB_DIR=...
mpirun -np 2 python tests/test_border_op.py

The test builds a periodic box decomposed into 2 subdomains along the x direction and verifies that:

1. after the forward pass, all ghost rows are correctly filled (cross-process or same-process periodic images);
2. local rows remain unchanged;
3. backpropagation correctly accumulates the gradients on ghost rows back onto the local rows of their owners.

Origin and Acknowledgements

This operator is extracted from the mattersim/Exten implementation in the MatterScale paper (SC26), an AI-MD application built on Microsoft's MatterSim. The operator is also currently used in DeePMD-kit. The code is released under the MIT License (see LICENSE). This README was entirely generated by Kimi-K3.

Special thanks to @nahso for the contributions to this project.