feat: Convert a subset of Mesa backend to Rust

[rht]

Git branch: https://github.com/rht/mesa/tree/rust

This is an experiment in using a Rust backend for Mesa. The tooling is very straightforward. Initial implementation of get_neighborhood2 shows that it is about 2x faster than pure Python, while a Cython version with a NumPy grid (the grid data type does not matter for get_neighborhood) is about 7x faster than pure Python. According to ChatGPT, there are still lots of room for improvement.

Benchmark parameter

grid_size = "100, 100"
pos = "(50, 50)"
radius = 50
setup = """import pure_python"""
stmt = f"pure_python.get_neighborhood({pos}, True, True, {radius}, False, {grid_size})"
elapsed_default = time_elapsed("default (pure python)", setup, stmt)

Measured using https://github.com/rht/mesa-perf/blob/main/experiments/get_neighborhood/benchmark.py.

To quickly test this branch, you need:

1. rustc
2. maturin (via pip install maturin; haven't tested uv pip install yet)
3. run maturin develop, alternatively, run maturin develop --release for faster code
4. In a Python interpreter:

In [2]: import mesa

In [3]: mesa.get_neighborhood((5, 5), True, True, 2, False, 10, 10)
Out[3]:
[(3, 3),
 (3, 4),
 (3, 5),
 (3, 6),
 (3, 7),
 (4, 3),
 (4, 4),
 (4, 5),
 (4, 6),
 (4, 7),
 (5, 3),
 (5, 4),
 (5, 6),
 (5, 7),
 (6, 3),
 (6, 4),
 (6, 5),
 (6, 6),
 (6, 7),
 (7, 3),
 (7, 4),
 (7, 5),
 (7, 6),
 (7, 7)]

[rht]

Update: an apple to apple comparison of get_neighborhood shows that Rust is faster.

default (pure python) 4205.103 μs
python empty 0.096 μs

cython vector 1122.303 μs
  faster 3.75 X
cython old less-readable 580.982 μs
  faster 7.24 X

rust hashmap 1360.873 μs
  faster 3.09 X
rust vector 662.559 μs
  faster 6.35 X
rust old less-readable 627.504 μs
  faster 6.7 X

I find "rust vector" syntax to be more expressive and pleasant than "cython vector":

#[pyfunction]
fn get_neighborhood_vector(
    pos: (i32, i32),
    moore: bool,
    include_center: bool,
    radius: i32,
    torus: bool,
    width: i32,
    height: i32,
) -> PyResult<Vec<(i32, i32)>> {
    let (x, y) = pos;
    // Estimate the maximum possible size of the neighborhood
    let capacity = (2 * radius + 1).pow(2) as usize;
    let mut neighborhood = Vec::with_capacity(capacity);

    let check_bounds = |x: i32, y: i32| -> bool {
        !(x < 0 || x >= width || y < 0 || y >= height)
    };

    for dx in -radius..=radius {
        for dy in -radius..=radius {
            if !moore && dx.abs() + dy.abs() > radius {
                continue;
            }

            let mut new_x = x + dx;
            let mut new_y = y + dy;

            if torus {
                new_x = ((new_x % width) + width) % width;
                new_y = ((new_y % height) + height) % height;
            } else if !check_bounds(new_x, new_y) {
                continue;
            }

            if !(!include_center && dx == 0 && dy == 0) {
                neighborhood.push((new_x, new_y));
            }
        }
    }

    Ok(neighborhood)
}

vs Cython

    @cython.boundscheck(False)
    @cython.wraparound(False)
    cpdef list get_neighborhood(
        self,
        object pos,
        bint moore,
        bint include_center,
        int radius,
        bint torus,
    ):

        cdef int x, y
        x, y = pos

        cdef int max_neighborhood_count
        max_neighborhood_count = (2 * radius + 1) ** 2
        cdef long[:, :] neighborhood
        neighborhood = np.empty((max_neighborhood_count, 2), long)

        cdef int count
        count = 0
        cdef int new_x, new_y
        cdef int dx, dy
        for dx in range(-radius, radius + 1):
            for dy in range(-radius, radius + 1):
                if not moore and abs(dx) + abs(dy) > radius:
                    continue
                new_x = x + dx
                new_y = y + dy

                if torus:
                    new_x %= self.width
                    new_y %= self.height
                elif self.out_of_bounds((new_x, new_y)):
                    continue

                if not (not include_center and dx == 0 and dy == 0):
                    neighborhood[count, 0] = new_x
                    neighborhood[count, 1] = new_y
                    count += 1

        # Convert to list
        cdef list neighborhood_list
        neighborhood_list = [0] * count
        for i in range(count):
            neighborhood_list[i] = (neighborhood[i, 0], neighborhood[i, 1])

        return neighborhood_list

Where in Rust:

1. conversion to Python object is straightforward.
2. no need to declare variables first, and so is more concise.
3. has better debugging tools than Cython, which outputs opaque C code and is hard to debug.
4. (edit) has hybrid Python-Rust profiling tool via py-spy. Whereas you can't profile the inside of the Cython code itself.

Note: It's not fully apple to apple comparison either, because the Cython version cheats by not doing bounds checking.

@quaquel what do you think of the syntax so far? Or would you delay judgement until I translate get_cell_list_contents and see how the syntax looks like for a Rust extension Python object?

[rht]

Cython with boundscheck and wraparound enabled, with radius of 50

default (pure python) 4201.759 μs
python empty 0.100 μs

cython vector 1128.796 μs
  faster 3.72 X
cython old less-readable 588.826 μs
  faster 7.14 X

rust hashmap 1353.057 μs
  faster 3.11 X
rust vector 646.587 μs
  faster 6.5 X
rust old less-readable 622.935 μs
  faster 6.75 X

with radius of 3

default (pure python) 8.505 μs
python empty 0.058 μs

cython vector 6.663 μs
  faster 1.28 X
cython old less-readable 3.573 μs
  faster 2.38 X

rust vector 2.739 μs
  faster 3.11 X
rust old less-readable 2.537 μs
  faster 3.35 X

[rht]

Here is a benchmark result of get_cell_list_contents:

default 36.563 μs
cython np.ndarray 8.219 μs
cython list-of-list 20.414 μs
cython memoryview-only 2.313 μs
rust 19.482 μs

The current Rust version is not as fast as cython memoryview-only because Cython doesn't require an extra allocation for a NumPy array. If the cell list is a NumPy array returned by get_neighborhood, the entire operation can be very fast. Whereas in Rust, it needs to convert the Python list to Rust Vec, and hence it is as slow as Cython list-of-list.
A solution is to use NumPy array everywhere via rust-numpy.

[rht]

The underlying Rust class code is once again, simpler than its Cython cdef class counterpart:

#[pyclass]
struct RustGrid {
    width: usize,
    height: usize,
    torus: bool,
    num_cells: usize,
    _grid: Vec<Vec<GridContent>>,
}

#[pymethods]
impl RustGrid {
    #[new]
    fn new(width: usize, height: usize, torus: bool) -> Self {
        let _grid = vec![vec![None; height]; width];
        RustGrid {
            width,
            height,
            torus,
            num_cells: width * height,
            _grid,
        }
    }

    fn get_cell_list_contents(&self, py: Python, cell_list: Vec<Coordinate>) -> PyResult<Vec<PyObject>> {
        // Reserve capacity to minimize reallocations
        let mut contents = Vec::with_capacity(cell_list.len()); 
        for cell in cell_list {
            let (x, y) = cell;
            if let Some(agent) = &self._grid[x][y] {
                contents.push(agent.clone_ref(py));
            }
        }
        Ok(contents)
    }

[quaquel]

1. I started looking into Rust last week. I really like what I have seen so far. However, I don't know if I can commit the required time to properly learn the language and be able to contribute to porting parts of MESA to Rust.
2. I would like to agree on a timeline. There are so many new experimental features being discussed and developed. Those need to be completed first, in my view. Only after those have been completed can we properly profile MESA and identify the key bottlenecks. For example, with the new DiscreteSpace, the function being tested here does not exist anymore.
3. Can we extend in Python a class defined in Rust? This is, for me, a critical question.
4. At face value, Rust seems superior to Cython. It has better readability, a superior memory model, and many other excellent features. Moreover, if we decide to port the core of MESA to Rust, we can do it in such a way that it is possible to also provide bindings in other languages (if anyone wants to do it).
5. What would the CI process look like for building and releasing?

[rht]

1

I can be the go-to person to review and answer questions, to speed up learning. Some of the speedup tricks I found came from blog posts (e.g. this one (2023)) instead of ChatGPT. But ChatGPT has been very helpful in automating the translation and contextualizing the reasons why it wrote the code in some specific ways.

2

The main purpose of this discussion thread is to characterize implementing Rust as a backend, orthogonal to the project of porting the new API's to Cython/Rust. The result so far has been beyond my expectation. Tooling and the Cargo inline debugging error messages have been pleasant. I expected a somewhat tedious port because of Rust ownership constraint, and so had a rather low bar to begin with.
The schedule to porting the new API's can be discussed once we have decided on the Cython vs Rust question.

3

I tried, and it worked, via a subclass macro: PyO3/pyo3#176. The missing feature is to extend a Python class from Rust: PyO3/pyo3#991

5

It doesn't require additional setup on GH Actions. Polars has been doing it with vanilla Ubuntu OS, which has both Python and Rust: https://github.com/pola-rs/polars/blob/bf5310dcc7c2dfe4a19269aa2e70b98235cd7683/.github/workflows/test-python.yml. Installation is done via maturin develop instead of pip install.

[quaquel]

1. Great, although let's also try to avoid the bus problem here. It would be nice if there were several contributors and maintainers who are comfortable with either Rust or Cython. Also, I am working my way through the Rust programming language. I personally prefer getting a broad understanding of a language before relying on tools like ChatGPT.
2. Great and I agree
3. That is very good to know. This makes it possible, in principle, to port a lot of the MESA framework to Rust while maintaining the current design.
4. ...
5. And what about building and distributing?

I am slowly coming around to the idea that Rust might be the more promising avenue.

[rht]

4. Yes, indeed I want a Mesa binding in other high-level languages for technical computing / data analysis, in particular: JavaScript and R. Julia already has Agents.jl. With a Cython backend, this would be impossible.
5. maturin build and maturin publish cover these use cases.