Upgrading to Warp 1.16

.gitignore

@@ -133,6 +133,7 @@ env.bak/
 venv.bak/
 .xlb-env/
 .xlb_install_test_venvs/
+.xlb_backend_test_venvs/
 
 # Spyder project settings
 .spyderproject

README.md

@@ -9,29 +9,36 @@
 XLB is a fully differentiable 2D/3D Lattice Boltzmann Method (LBM) library that leverages hardware acceleration. It supports [JAX](https://github.com/google/jax), [NVIDIA Warp](https://github.com/NVIDIA/warp), and [Neon](https://github.com/Autodesk/Neon) backends, and is specifically designed to solve fluid dynamics problems in a computationally efficient and differentiable manner. Its unique combination of features positions it as an exceptionally suitable tool for applications in physics-based machine learning. With the Warp backend, XLB offers state-of-the-art single-GPU performance, and with the new Neon backend it extends to multi-GPU (single-resolution). More importantly, the Neon backend provides grid refinement capabilities for multi-resolution simulations.
 
 ## Getting Started
-To get started with XLB, you can install it using pip. There are different installation options depending on your hardware and needs:
+To get started with XLB, you can install it using pip. XLB imports Warp at import time and Warp must come from exactly one of two extras, so every install picks either `xlb[warp]` or `xlb[neon]`:
+
+| Extra | Source of Warp | Use for |
+| --- | --- | --- |
+| `xlb[warp]` | `warp-lang` from PyPI | JAX and Warp backends (single-GPU) |
+| `xlb[neon]` | the fork bundled in `neon_gpu` | Neon backend (multi-GPU / multi-resolution) |
+
+Both unpack into the same `site-packages/warp`, so installing both leaves whichever pip wrote last in charge; XLB warns on import if it finds both.
 
 ### Basic Installation (CPU-only)
 ```bash
-pip install xlb
+pip install "xlb[warp]"
 ```
 
 ### Installation with Warp support (single-GPU)
-For the NVIDIA Warp backend (single-GPU, state-of-the-art performance):
+For the NVIDIA Warp backend (single-GPU, state-of-the-art performance) the same extra applies:
 ```bash
 pip install "xlb[warp]"
 ```
 
 ### Installation with CUDA support (for NVIDIA GPUs)
 This installation is for the JAX backend with CUDA support:
 ```bash
-pip install "xlb[cuda]"
+pip install "xlb[cuda,warp]"
 ```
 
 ### Installation with TPU support
 This installation is for the JAX backend with TPU support:
 ```bash
-pip install "xlb[tpu]"
+pip install "xlb[tpu,warp]"
 ```
 
 ### Installation with Neon support
@@ -47,7 +54,7 @@ pip install '.[neon]'
 
 **Requirements:** The Neon wheel supports **Python 3.11** to **Python 3.14** on **Linux x86_64** and **Linux ARM**. 
 
-**Note:** Neon uses a custom fork of warp.
+**Note:** Neon uses a custom fork of warp, which is why `xlb[neon]` does not install `warp-lang`. If the environment already has `warp-lang`, remove it first with `pip uninstall warp-lang`, then reinstall `.[neon]`.
 
 ### Notes:
 - For Mac users: Use the basic CPU installation command as JAX's GPU support is not available on MacOS
@@ -57,7 +64,15 @@ pip install '.[neon]'
 To install the latest development version from source:
 
 ```bash
-pip install git+https://github.com/Autodesk/XLB.git
+pip install "xlb[warp] @ git+https://github.com/Autodesk/XLB.git"
+```
+
+### Running the tests
+`tests/run_backend_test_envs.py` builds a virtualenv per backend, installs XLB from the repository root, and runs the suite in each:
+
+```bash
+python tests/run_backend_test_envs.py                 # warp and neon
+python tests/run_backend_test_envs.py --backend neon  # one backend
 ```
 
 The changelog for the releases can be found [here](https://github.com/Autodesk/XLB/blob/main/CHANGELOG.md).

docs/index.md

@@ -7,20 +7,20 @@
 # XLB: A Differentiable Massively Parallel Lattice Boltzmann Library in Python for Physics-Based Machine Learning
 
 🎉 **Exciting News!** 🎉 XLB version 0.2.0 has been released, featuring a complete rewrite of the library and introducing support for the NVIDIA Warp backend! 
-XLB can now be installed via pip: `pip install xlb`.
+XLB can now be installed via pip: `pip install "xlb[warp]"`.
 
 XLB is a fully differentiable 2D/3D Lattice Boltzmann Method (LBM) library that leverages hardware acceleration. It supports [JAX](https://github.com/google/jax) and [NVIDIA Warp](https://github.com/NVIDIA/warp) backends, and is specifically designed to solve fluid dynamics problems in a computationally efficient and differentiable manner. Its unique combination of features positions it as an exceptionally suitable tool for applications in physics-based machine learning. With the new Warp backend, XLB now offers state-of-the-art performance for even faster simulations.
 
 ## Getting Started
-To get started with XLB, you can install it using pip:
+To get started with XLB, you can install it using pip. Warp comes from either the `warp` extra (PyPI `warp-lang`) or the `neon` extra (the fork bundled with `neon_gpu`); pick exactly one:
 ```bash
-pip install xlb
+pip install "xlb[warp]"
 ```
 
 To install the latest development version from source:
 
 ```bash
-pip install git+https://github.com/Autodesk/XLB.git
+pip install "xlb[warp] @ git+https://github.com/Autodesk/XLB.git"
 ```
 
 The changelog for the releases can be found [here](https://github.com/Autodesk/XLB/blob/main/CHANGELOG.md).

examples/cfd/flow_past_sphere_3d.py

@@ -98,7 +98,7 @@ def bc_profile_warp(index: wp.vec3i):
             r_squared = (two * y_center / H_y) ** two + (two * z_center / H_z) ** two
 
             # Parabolic profile: u = u_max * (1 - r²)
-            return wp.vec(wp_dtype(u_max) * wp.max(wp_dtype(0.0), wp_dtype(1.0) - r_squared), length=1)
+            return wp.vector(wp_dtype(u_max) * wp.max(wp_dtype(0.0), wp_dtype(1.0) - r_squared), length=1)
 
         return bc_profile_warp
 

examples/cfd/multires_flow_past_sphere_3d.py

@@ -171,7 +171,7 @@ def bc_profile():
     one = dtype(1.0)
     zero = dtype(0.0)
     u_max_wp = dtype(u_max)
-    _u_vec = wp.vec(velocity_set.d, dtype=dtype)
+    _u_vec = wp.types.vector(length=velocity_set.d, dtype=dtype)
 
     @wp.func
     def bc_profile_warp(index: wp.vec3i):
@@ -191,7 +191,7 @@ def bc_profile_warp(index: wp.vec3i):
         # return _u_vec(u_max_wp * wp.max(zero, one - r_squared), zero, zero)
 
         # For Regularized and ZouHe
-        return wp.vec(u_max_wp * wp.max(zero, one - r_squared), length=1)
+        return wp.vector(u_max_wp * wp.max(zero, one - r_squared), length=1)
 
     return bc_profile_warp
 

examples/cfd/rotating_sphere_3d.py

@@ -115,7 +115,7 @@
 def bc_profile():
     """Build a Warp function returning the rotational wall velocity at a voxel."""
     dtype = precision_policy.compute_precision.wp_dtype
-    _u_vec = wp.vec(velocity_set.d, dtype=dtype)
+    _u_vec = wp.types.vector(length=velocity_set.d, dtype=dtype)
     angular_velocity = _u_vec(0.0, rot_rate, 0.0)
     origin_np = shift + diam / 2
     origin_wp = _u_vec(origin_np[0], origin_np[1], origin_np[2])

examples/ibm/airfoil_ibm.py

@@ -70,7 +70,7 @@ def bc_profile(precision_policy, grid_shape, u_max):
 
     @wp.func
     def bc_profile_warp(index: wp.vec3i):
-        return wp.vec(dtype(u_max_d), length=1)
+        return wp.vector(dtype(u_max_d), length=1)
 
     return bc_profile_warp
 

examples/ibm/sphere_ibm.py

@@ -57,7 +57,7 @@ def bc_profile(precision_policy, grid_shape, u_max):
 
     @wp.func
     def bc_profile_warp(index: wp.vec3i):
-        return wp.vec(_dtype(u_max_d), length=1)
+        return wp.vector(_dtype(u_max_d), length=1)
 
     return bc_profile_warp
 

examples/ibm/wind_turbine_ibm.py

@@ -43,7 +43,7 @@ def bc_profile(precision_policy, grid_shape, u_max):
 
     @wp.func
     def bc_profile_warp(index: wp.vec3i):
-        return wp.vec(u_max_d, length=1)
+        return wp.vector(u_max_d, length=1)
 
     return bc_profile_warp
 

examples/ibm/windtunnel_ibm.py

@@ -239,7 +239,7 @@ def bc_profile_warp(index: wp.vec3i):
         z_center = z - (L_z / _dtype(2.0))
         r_sq = ((_dtype(2.0) * y_center) / L_y) ** _dtype(2.0) + ((_dtype(2.0) * z_center) / L_z) ** _dtype(2.0)
         velocity_x = u_max_d * wp.max(_dtype(0.0), _dtype(1.0) - r_sq)
-        return wp.vec(velocity_x, length=1)
+        return wp.vector(velocity_x, length=1)
 
     return bc_profile_warp
 
@@ -559,11 +559,11 @@ def save_drag_coefficient(cd_values, filename):
 cy_list = [c[1] for c in wheel_centers]
 cz_list = [c[2] for c in wheel_centers]
 
-_wheel_starts = wp.constant(wp.vec(len(starts_list), dtype=int)(starts_list))
-_wheel_ends = wp.constant(wp.vec(len(ends_list), dtype=int)(ends_list))
-_wheel_centers_x = wp.constant(wp.vec(len(cx_list), dtype=float)(cx_list))
-_wheel_centers_y = wp.constant(wp.vec(len(cy_list), dtype=float)(cy_list))
-_wheel_centers_z = wp.constant(wp.vec(len(cz_list), dtype=float)(cz_list))
+_wheel_starts = wp.constant(wp.types.vector(length=len(starts_list), dtype=int)(starts_list))
+_wheel_ends = wp.constant(wp.types.vector(length=len(ends_list), dtype=int)(ends_list))
+_wheel_centers_x = wp.constant(wp.types.vector(length=len(cx_list), dtype=float)(cx_list))
+_wheel_centers_y = wp.constant(wp.types.vector(length=len(cy_list), dtype=float)(cy_list))
+_wheel_centers_z = wp.constant(wp.types.vector(length=len(cz_list), dtype=float)(cz_list))
 
 bc_list = setup_boundary_conditions(grid, velocity_set, precision_policy, grid_shape, u_max)
 stepper = setup_stepper(grid, bc_list, omega)

examples/out_of_core/autodiff_lbm.py

@@ -194,11 +194,13 @@ def backward(
     else:
         raise ValueError("Invalid velocity set")
 
-    # Initialize XLB
+    # Initialize XLB. This example differentiates through the solver with
+    # wp.Tape, so it needs the adjoint kernels that XLB otherwise skips.
     xlb.init(
         velocity_set=velocity_set,
         default_backend=compute_backend,
         default_precision_policy=precision_policy,
+        enable_backward=True,
     )
 
     # Make grid for constructing stepper

setup.py

@@ -1,36 +1,21 @@
-import os
 import platform
-import subprocess
 import sys
 
 from setuptools import setup, find_packages
-from setuptools.command.install import install
 
-
-def _neon_extra_requested():
-    """Best-effort detection of [neon] extra from install invocation."""
-    for arg in sys.argv:
-        if "neon" in arg and ("[" in arg or "xlb" in arg):
-            return True
-    return False
-
-
-def _uninstall_warp_lang(*, reason: str) -> None:
-    """Uninstall the ``warp-lang`` distribution so Neon's bundled warp fork is used."""
-    if os.environ.get("XLB_NEON_SKIP_UNINSTALL_WARP", "").lower() in ("1", "true", "yes"):
-        return
-    print(f"[xlb] {reason}")
-    try:
-        subprocess.run(
-            [sys.executable, "-m", "pip", "uninstall", "warp-lang", "-y"],
-            check=False,
-            capture_output=True,
-        )
-    except Exception as exc:  # noqa: BLE001
-        print(f"[xlb] Warning: failed to uninstall warp-lang: {exc}")
-
-
-_NEON_VERSION = "0.5.2a1"
+# Warp is required at import time, but it is deliberately not in
+# ``install_requires``: it must come from exactly one of the two extras.
+# ``[warp]`` installs ``warp-lang`` from PyPI, while ``[neon]`` installs
+# ``neon_gpu``, which bundles its own Warp fork. Both unpack into
+# ``site-packages/warp``, so if pip installs both, whichever lands last owns the
+# directory and later uninstalling ``warp-lang`` deletes files Neon needs.
+#
+# Listing ``warp-lang`` as a core dependency made that collision unavoidable for
+# ``pip install xlb[neon]``. A setup.py hook cannot undo it either: pip installs
+# dependencies after building this project, and for wheel and PEP 660 editable
+# installs the ``install`` command never runs at all. Keeping it out of the core
+# dependencies is the only way ``[neon]`` reliably gets one Warp.
+_NEON_VERSION = "0.5.2a5"
 _NEON_RELEASE_URL = f"https://github.com/Autodesk/Neon/releases/download/v{_NEON_VERSION}"
 
 
@@ -47,33 +32,6 @@ def _neon_wheel_requirement():
     return req
 
 
-class InstallWithNeonHooks(install):
-    """Uninstall ``warp-lang`` before and after install when ``[neon]`` is requested.
-
-    * **Before** ``pip``/setuptools install dependencies: removes any previously
-      installed ``warp-lang`` so an older or PyPI build does not linger next to
-      Neon's fork (``neon_gpu`` ships its own warp).
-    * **After** install: removes the ``warp-lang`` pulled in by ``install_requires``,
-      leaving Neon's warp as the one on the path.
-
-    Only runs when installing from source (e.g. sdist or git). Wheel installs
-    do not run setup.py, so for ``pip install xlb[neon]`` from PyPI you may
-    need to run ``pip uninstall warp-lang`` first if it is already installed.
-    Set XLB_NEON_SKIP_UNINSTALL_WARP=1 to disable this behaviour.
-    """
-
-    def run(self):
-        if _neon_extra_requested():
-            _uninstall_warp_lang(
-                reason=("Removing any existing warp-lang before Neon install (neon_gpu provides its own warp fork)."),
-            )
-        install.run(self)
-        if _neon_extra_requested():
-            _uninstall_warp_lang(
-                reason=("Removing PyPI warp-lang after install (core deps); use the warp bundled with neon_gpu."),
-            )
-
-
 setup(
     name="xlb",
     version="0.3.2",
@@ -93,11 +51,14 @@ def run(self):
         "pydantic>=2.9.1",
         "ruff>=0.14.1",
         "jax>=0.8.2",  # Base JAX CPU-only requirement
-        "warp-lang>=1.10.0",  # Required at import time (core modules import warp)
         "nvtx>=0.2.0",  # NVTX ranges (e.g. nse_multires_stepper); listed in requirements.txt
     ],
     extras_require={
-        "warp": ["warp-lang>=1.10.0"],  # Kept for explicit `pip install xlb[warp]` / Neon uninstall hook docs
+        # Exactly one of [warp] and [neon] is required; see the note above.
+        # Bounded to the versions the suite has been run against: 1.16, the API
+        # XLB's kernels target and the series of the fork neon_gpu bundles, and
+        # 1.17. A fresh install takes 1.17; pin warp-lang==1.16.* to match Neon.
+        "warp": ["warp-lang>=1.16,<1.18"],
         "cuda": ["jax[cuda13]>=0.8.2"],  # For CUDA installations (pip install -U "jax[cuda13]")
         "tpu": ["jax[tpu]>=0.8.2"],  # For TPU installations
         # h5py: MultiresIO / Neon multi-resolution export to HDF5 (see xlb.utils.mesher).
@@ -106,5 +67,4 @@ def run(self):
     },
     python_requires=">=3.11",
     dependency_links=["https://storage.googleapis.com/jax-releases/libtpu_releases.html"],
-    cmdclass={"install": InstallWithNeonHooks},
 )

tests/backends/test_neon_backend.py

@@ -0,0 +1,81 @@
+"""
+Smoke coverage for the Neon compute backend.
+
+The rest of the suite only exercises the JAX and WARP backends, and
+``tests/install/flow_past_sphere_3d_test.py`` skips Neon because the sphere's
+``HalfwayBounceBackBC`` has no Neon implementation. These tests cover the parts
+of the Neon path that need no boundary condition: ``neon.init()`` registering
+XLB's custom Warp types for NVRTC, grid and field allocation, the device-side
+fill and host read-back, and operator kernel construction.
+
+Skipped when ``neon`` is missing (an env built with the ``[warp]`` extra rather
+than ``[neon]``) or when no CUDA device is available.
+"""
+
+import pytest
+
+neon = pytest.importorskip("neon", reason="neon_gpu is not installed; install XLB with the [neon] extra")
+
+import warp as wp  # noqa: E402
+
+import xlb  # noqa: E402
+from xlb.compute_backend import ComputeBackend  # noqa: E402
+from xlb.grid import grid_factory  # noqa: E402
+from xlb.precision_policy import PrecisionPolicy  # noqa: E402
+
+GRID_SHAPE = (16, 16, 16)
+CARDINALITY = 19
+FILL_VALUE = 2.5
+
+
+@pytest.fixture(scope="module", autouse=True)
+def neon_env():
+    """Initialize XLB on the Neon backend once for the module.
+
+    ``xlb.init`` calls ``neon.init()``, which must run before any kernel using
+    Neon's custom types is built, so every test here shares one initialization.
+    """
+    wp.init()
+    if wp.get_cuda_device_count() == 0:
+        pytest.skip("the Neon backend requires a CUDA device")
+
+    precision_policy = PrecisionPolicy.FP32FP32
+    xlb.init(
+        velocity_set=xlb.velocity_set.D3Q19(precision_policy=precision_policy, compute_backend=ComputeBackend.NEON),
+        default_backend=ComputeBackend.NEON,
+        default_precision_policy=precision_policy,
+    )
+
+
+def test_neon_grid_shape():
+    grid = grid_factory(GRID_SHAPE, compute_backend=ComputeBackend.NEON)
+    assert grid.shape == GRID_SHAPE, "Neon grid shape is incorrect"
+
+
+def test_neon_field_cardinality():
+    grid = grid_factory(GRID_SHAPE, compute_backend=ComputeBackend.NEON)
+    f = grid.create_field(cardinality=CARDINALITY)
+    assert f.cardinality == CARDINALITY, "Neon field cardinality is incorrect"
+
+
+def test_neon_field_fill_value():
+    """Fill on the device, then read back on the host.
+
+    Exercises allocation, the fill kernel, and the device-to-host transfer.
+    """
+    grid = grid_factory(GRID_SHAPE, compute_backend=ComputeBackend.NEON)
+    f = grid.create_field(cardinality=CARDINALITY, fill_value=FILL_VALUE)
+    f.update_host(0)
+
+    for idx in (neon.Index_3d(0, 0, 0), neon.Index_3d(1, 2, 3), neon.Index_3d(*(d - 1 for d in GRID_SHAPE))):
+        for card in (0, CARDINALITY // 2, CARDINALITY - 1):
+            assert f.read(idx, card) == pytest.approx(FILL_VALUE), "Neon field not initialized with fill_value"
+
+
+def test_neon_operator_kernel_construction():
+    """Operators must produce a Neon functional and container, not Warp ones."""
+    from xlb.operator.equilibrium import QuadraticEquilibrium
+
+    op = QuadraticEquilibrium(compute_backend=ComputeBackend.NEON)
+    assert op.neon_functional is not None, "Neon functional was not constructed"
+    assert op.neon_container is not None, "Neon container was not constructed"