feat: render the Mandelbrot set on CPU and GPU

Replace the add/subtract demo with a more instructional example: the
Mandelbrot set rendered two ways so users can read both implementations
side by side and compare their performance.

- src/mandelbrot_cpu.cpp: plain nested loop over every pixel
- src/mandelbrot.cu: the same logic as a CUDA kernel, one thread per pixel
- src/main.cpp: shared wrapper returning a (height, width) int32 NumPy array,
  releasing the GIL during compute
- both functions take identical arguments and return identical arrays
- add numpy as a runtime dependency; test GPU output matches the CPU
- ignore vim swap files

Assisted-by: ClaudeCode:claude-opus-4.8

Author: henryiii

.gitignore

@@ -143,3 +143,7 @@ cython_debug/
 
 _skbuild/
 .pyodide-xbuildenv/
+
+# Vim swap files
+*.swp
+*.swo

CMakeLists.txt

@@ -14,8 +14,8 @@ project(
 find_package(pybind11 CONFIG REQUIRED)
 
 # Add a library using FindPython's tooling (pybind11 also provides a helper like
-# this), combining the pybind11 bindings with the CUDA kernels.
-python_add_library(_core MODULE src/main.cpp src/add.cu WITH_SOABI)
+# this), combining the pybind11 bindings with the CPU and CUDA implementations.
+python_add_library(_core MODULE src/main.cpp src/mandelbrot_cpu.cpp src/mandelbrot.cu WITH_SOABI)
 target_link_libraries(_core PRIVATE pybind11::headers)
 
 # This is passing in the version as a define just as an example

README.md

@@ -11,13 +11,19 @@
 An example project built with [pybind11][], [CUDA][], and
 [scikit-build-core][]. Python 3.9+.
 
-The extension exposes tiny `add` and `subtract` functions that run on the GPU
-with CUDA kernels (`src/add.cu`). Building **requires** the CUDA Toolkit
-(`nvcc`): the CMake project declares CUDA as a required language, so
-configuration fails without it. The CUDA runtime is linked statically, so the
-resulting wheels do not depend on `libcudart` and stay importable on machines
-without a GPU — running `add`/`subtract` there raises, but `cuda_available()`
-lets you check first.
+The extension renders the [Mandelbrot set][mandelbrot] two ways — once on the
+CPU and once on the GPU — so you can read both side by side and compare their
+performance. The two implementations are written the same way on purpose:
+
+* `src/mandelbrot_cpu.cpp` — a plain nested loop over every pixel
+* `src/mandelbrot.cu` — the same logic as a CUDA kernel, one thread per pixel
+
+Both return a `(height, width)` int32 NumPy array of escape counts. Building
+**requires** the CUDA Toolkit (`nvcc`): the CMake project declares CUDA as a
+required language, so configuration fails without it. The CUDA runtime is linked
+statically, so the resulting wheels do not depend on `libcudart` and stay
+importable on machines without a GPU — calling `mandelbrot_gpu` there raises, but
+`cuda_available()` lets you check first.
 
 
 [gitter-badge]:            https://badges.gitter.im/pybind/Lobby.svg
@@ -39,9 +45,42 @@ The CUDA Toolkit (`nvcc`) must be installed and discoverable by CMake.
 ```python
 import cuda_example
 
+# (height, width) int32 array of escape counts
+image = cuda_example.mandelbrot_cpu(width=800, height=600, max_iterations=100)
+
 if cuda_example.cuda_available():
-    cuda_example.add(1, 2)       # 3, computed on the GPU
-    cuda_example.subtract(1, 2)  # -1, computed on the GPU
+    image = cuda_example.mandelbrot_gpu(width=800, height=600, max_iterations=100)
+```
+
+You can view the result with any plotting library, e.g.:
+
+```python
+import matplotlib.pyplot as plt
+
+plt.imshow(image, extent=(-2, 1, -1.25, 1.25), cmap="twilight_shifted")
+plt.show()
+```
+
+## Comparing CPU and GPU
+
+Because both functions take the same arguments and return identical arrays, you
+can run them back to back and time them (on a machine with a GPU):
+
+```python
+import time
+import cuda_example
+
+size = {"width": 2000, "height": 1500, "max_iterations": 200}
+
+start = time.perf_counter()
+cpu = cuda_example.mandelbrot_cpu(**size)
+print(f"CPU: {time.perf_counter() - start:.3f}s")
+
+start = time.perf_counter()
+gpu = cuda_example.mandelbrot_gpu(**size)
+print(f"GPU: {time.perf_counter() - start:.3f}s")
+
+assert (cpu == gpu).all()  # identical results, very different runtimes
 ```
 
 ## Building CUDA wheels
@@ -86,16 +125,17 @@ docker run --rm \
     PY=/opt/python/cp312-cp312/bin/python
     cp -r /io /tmp/src && cd /tmp/src
     $PY -m pip install --upgrade pip build pytest
-    $PY -m build --wheel --outdir /wheelhouse .   # compiles src/add.cu with nvcc
+    $PY -m build --wheel --outdir /wheelhouse .   # compiles src/mandelbrot.cu with nvcc
     $PY -m pip install /wheelhouse/*.whl
     $PY -m pytest                                  # GPU tests skip (no device)
   '
 ```
 
 The compiled wheel is written to `./wheelhouse/` on the host, so you can inspect
 or install it afterwards. Because the container has no GPU, `cuda_available()`
-returns `False` and the `add`/`subtract` tests are skipped. The same flow runs
-in CI in the `cuda` job of `.github/workflows/pip.yml`.
+returns `False` and the `mandelbrot_gpu` test is skipped (the `mandelbrot_cpu`
+tests still run). The same flow runs in CI in the `cuda` job of
+`.github/workflows/pip.yml`.
 
 ## Files
 
@@ -104,9 +144,10 @@ necessary. The necessary files are:
 
 * `pyproject.toml`: The Python project file
 * `CMakeLists.txt`: The CMake configuration file, which requires the CUDA language
-* `src/main.cpp`: The pybind11 bindings
-* `src/add.cu`: The CUDA kernels (`add`/`subtract`) and runtime device query
-* `src/add.h`: The shared declarations
+* `src/main.cpp`: The pybind11 bindings (turns the results into NumPy arrays)
+* `src/mandelbrot_cpu.cpp`: The CPU implementation
+* `src/mandelbrot.cu`: The CUDA kernel and runtime device query
+* `src/mandelbrot.h`: The shared declarations
 * `src/cuda_example/__init__.py`: The Python portion of the module. The root of the module needs to be `<package_name>`, `src/<package_name>`, or `python/<package_name>` to be auto-discovered.
 
 These files are also expected and highly recommended:
@@ -148,6 +189,7 @@ terms and conditions of this license.
 [cibuildwheel]: https://cibuildwheel.readthedocs.io
 [cibw-cuda]: https://github.com/pypa/cibuildwheel/pull/2896
 [cuda]: https://developer.nvidia.com/cuda-toolkit
+[mandelbrot]: https://en.wikipedia.org/wiki/Mandelbrot_set
 [scientific-python development guide]: https://learn.scientific-python.org/development
 [dependabot]: https://docs.github.com/en/code-security/dependabot
 [github actions]: https://docs.github.com/en/actions

pyproject.toml

@@ -25,10 +25,11 @@ classifiers = [
   "Programming Language :: Python :: 3.14",
   "Private :: Do Not Upload",
 ]
+dependencies = ["numpy"]
 
 
 [dependency-groups]
-test = ["pytest"]
+test = ["pytest", "numpy"]
 dev = [{ include-group = "test" }]
 
 

src/add.cu

@@ -3,15 +3,15 @@
 from ._core import (
     __doc__,
     __version__,
-    add,
     cuda_available,
-    subtract,
+    mandelbrot_cpu,
+    mandelbrot_gpu,
 )
 
 __all__ = [
     "__doc__",
     "__version__",
-    "add",
     "cuda_available",
-    "subtract",
+    "mandelbrot_cpu",
+    "mandelbrot_gpu",
 ]

src/add.h

@@ -1,27 +1,50 @@
 """
-Pybind11 + CUDA example plugin
-------------------------------
+Pybind11 + CUDA Mandelbrot example
+----------------------------------
 
 .. currentmodule:: cuda_example
 
 .. autosummary::
     :toctree: _generate
 
-    add
-    subtract
+    mandelbrot_cpu
+    mandelbrot_gpu
     cuda_available
 """
 
 from __future__ import annotations
 
-def add(i: int, j: int) -> int:
+from numpy import int32
+from numpy.typing import NDArray
+
+def mandelbrot_cpu(
+    width: int = ...,
+    height: int = ...,
+    max_iterations: int = ...,
+    xmin: float = ...,
+    xmax: float = ...,
+    ymin: float = ...,
+    ymax: float = ...,
+) -> NDArray[int32]:
     """
-    Add two numbers on the GPU with a CUDA kernel.
+    Render the Mandelbrot set on the CPU.
+
+    Returns a ``(height, width)`` int32 array of escape counts.
     """
 
-def subtract(i: int, j: int) -> int:
+def mandelbrot_gpu(
+    width: int = ...,
+    height: int = ...,
+    max_iterations: int = ...,
+    xmin: float = ...,
+    xmax: float = ...,
+    ymin: float = ...,
+    ymax: float = ...,
+) -> NDArray[int32]:
     """
-    Subtract two numbers on the GPU with a CUDA kernel.
+    Render the Mandelbrot set on the GPU with CUDA.
+
+    Returns a ``(height, width)`` int32 array of escape counts.
     """
 
 def cuda_available() -> bool:

src/cuda_example/__init__.py

@@ -1,34 +1,85 @@
+#include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
 
-#include "add.h"
+#include <cstdint>
+#include <stdexcept>
+
+#include "mandelbrot.h"
 
 #define STRINGIFY(x) #x
 #define MACRO_STRINGIFY(x) STRINGIFY(x)
 
 namespace py = pybind11;
 
+namespace {
+
+using Image = py::array_t<std::int32_t, py::array::c_style>;
+
+// Shared wrapper: validate the arguments, allocate the output image, and run one
+// of the compute functions with the GIL released. Returns a (height, width)
+// NumPy array of escape counts.
+Image render(void (*compute)(const MandelbrotParams &, std::int32_t *), int width, int height,
+             int max_iterations, double xmin, double xmax, double ymin, double ymax) {
+    if (width <= 0 || height <= 0) {
+        throw std::invalid_argument("width and height must be positive");
+    }
+    if (max_iterations <= 0) {
+        throw std::invalid_argument("max_iterations must be positive");
+    }
+
+    const MandelbrotParams params{width, height, max_iterations, xmin, xmax, ymin, ymax};
+    Image image({height, width});
+    std::int32_t *data = image.mutable_data();
+
+    {
+        py::gil_scoped_release release;
+        compute(params, data);
+    }
+    return image;
+}
+
+}  // namespace
+
 PYBIND11_MODULE(_core, m, py::mod_gil_not_used(), py::multiple_interpreters::per_interpreter_gil()) {
     m.doc() = R"pbdoc(
-        Pybind11 + CUDA example plugin
-        ------------------------------
+        Pybind11 + CUDA Mandelbrot example
+        ----------------------------------
 
         .. currentmodule:: cuda_example
 
         .. autosummary::
            :toctree: _generate
 
-           add
-           subtract
+           mandelbrot_cpu
+           mandelbrot_gpu
            cuda_available
     )pbdoc";
 
-    m.def("add", &add, R"pbdoc(
-        Add two numbers on the GPU with a CUDA kernel.
-    )pbdoc");
+    const char *doc = R"pbdoc(
+        Render the Mandelbrot set.
 
-    m.def("subtract", &subtract, R"pbdoc(
-        Subtract two numbers on the GPU with a CUDA kernel.
-    )pbdoc");
+        Returns a ``(height, width)`` int32 NumPy array; each value is the number
+        of iterations before the point escaped (``max_iterations`` if it never
+        did).
+    )pbdoc";
+
+    m.def("mandelbrot_cpu",
+          [](int width, int height, int max_iterations, double xmin, double xmax, double ymin,
+             double ymax) {
+              return render(&mandelbrot_cpu, width, height, max_iterations, xmin, xmax, ymin, ymax);
+          },
+          py::arg("width") = 800, py::arg("height") = 600, py::arg("max_iterations") = 100,
+          py::arg("xmin") = -2.0, py::arg("xmax") = 1.0, py::arg("ymin") = -1.25,
+          py::arg("ymax") = 1.25, doc);
+
+    m.def("mandelbrot_gpu",
+          [](int width, int height, int max_iterations, double xmin, double xmax, double ymin,
+             double ymax) {
+              return render(&mandelbrot_gpu, width, height, max_iterations, xmin, xmax, ymin, ymax);
+          },
+          py::arg("width") = 800, py::arg("height") = 600, py::arg("max_iterations") = 100,
+          py::arg("xmin") = -2.0, py::arg("xmax") = 1.0, py::arg("ymin") = -1.25,
+          py::arg("ymax") = 1.25, doc);
 
     m.def("cuda_available", &cuda_available, R"pbdoc(
         Return True if a CUDA-capable device is available at runtime.