Conversation
Warp 1.16 removed the generic type generators and moved several helpers: - wp.vec(n, dtype=) -> wp.types.vector(length=n, dtype=) (46 sites) - wp.mat((r,c),dtype=)-> wp.types.matrix(shape=(r,c), dtype=) - wp.mat33(a,b,c)/wp.mat33f(a,b,c) (row-vectors) -> wp.matrix_from_rows(...) - in-kernel wp.vec(x, length=1) -> wp.vector(x, length=1) (bc_zouhe) - wp.build.clear_kernel_cache() -> wp.clear_kernel_cache() - set wp.config.enable_backward = False for Neon (forward-only; avoids NVRTC adjoint compile errors from legacy @wp.func patterns) Also point setup.py [neon] extra at neon_gpu 0.5.2a3 (update4 wheels). Co-authored-by: Cursor <cursoragent@cursor.com>
Completes the Warp 1.16 port beyond the xlb/ package: - examples/ (7 files) + tests/install (1 file): wp.vec -> wp.types.vector for Python-scope type generators (dtype=), and -> wp.vector for the in-kernel value constructor (length=). - Restore the wp.Mesh retention fix in mesh_boundary_masker and aabb_close: only the integer mesh.id reaches the kernels, so the Mesh's BVH arrays could be freed before Neon's deferred NVRTC-compiled container runs. NOT NUMERICALLY VALIDATED: multires_flow_past_sphere_3d on this stack gives Cd ~= 0.348 vs ~= 2.680 on the reference stack (warp 1.7.2 + neon 0.5.2a1). Ruled out: this type port (velocity-set constants are bit-identical), Warp BVH /mesh queries (identical across versions), voxelization (identical cell count), the drag operator (identical code), and the Neon C++ core (one benign diff). Cause still unknown. Co-authored-by: Cursor <cursoragent@cursor.com>
The Warp 1.7.2 3-vector matrix constructor was inconsistent between scopes: in Python scope wp.mat33(a,b,c) builds ROWS, but inside a @wp.func/@wp.kernel it builds COLUMNS. Warp 1.16 split these into matrix_from_rows/matrix_from_cols precisely to remove that ambiguity. Both remaining sites are in-kernel, so porting them to matrix_from_rows transposed the triangle vertex/edge matrices feeding triangle_box_intersect. Mesh voxelization then found almost no triangle-voxel overlaps: the sphere was marked with only 10 solid cells instead of 440, so there was effectively no obstacle in the domain. Validated on multires_flow_past_sphere_3d (10k steps, D3Q27, Re=5000) against the reference stack (warp 1.7.2 + neon 0.5.2a1): - solid-cell mask now identical (440 vs 440) - Cd within 0.15% at every reporting step (final 2.6761 vs 2.6802) - max |u| identical, mean |u| within 0.36% This supersedes the 'NOT NUMERICALLY VALIDATED' note in dd0674f. Co-authored-by: Cursor <cursoragent@cursor.com>
Warp generates an adjoint copy of every kernel by default, roughly doubling codegen and compile time. XLB's LBM solvers are forward-only, so this was already switched off for the Neon backend (where the legacy @wp.func patterns additionally fail NVRTC adjoint compilation). The plain Warp backend never got the same treatment and was still paying for adjoints it never launched. Move the setting into _warp_init_and_select_cuda_device, the common chokepoint for Warp startup, so both backends pick it up. It has to run before the first kernel is created, since Warp snapshots enable_backward into a module's options at that point. examples/out_of_core/autodiff_lbm.py does differentiate through the solver with wp.Tape, so add XLB_WARP_ENABLE_BACKWARD to opt back in, and set it there. Verified: QuadraticEquilibrium's kernel module reports enable_backward=False on the Warp backend (was True) and True under the opt-in; Neon unchanged. The 87 Warp-backend tests in tests/kernels, tests/grids, tests/boundary_conditions pass (the 6 test_grid_jax.py failures are a pre-existing JAX PartitionSpec equality change, identical on an unmodified tree). Co-authored-by: Cursor <cursoragent@cursor.com>
Makes Warp's adjoint (backward) codegen switchable per call instead of only via the XLB_WARP_ENABLE_BACKWARD environment variable. Resolution order is explicit argument, then environment variable, then off, so an explicit enable_backward=False overrides the environment and scripts that cannot pass arguments keep working. The resolved value is recorded on DefaultConfig.enable_backward for introspection; it stays None on the JAX backend, which does not use Warp. autodiff_lbm.py now passes enable_backward=True directly rather than setting the environment variable, which reads better for the one case that differentiates through the solver. The parameter is optional and trailing, and all 34 existing init() call sites pass keyword arguments, so nothing else needs updating. Verified all six combinations of argument (unset/True/False) against environment (unset/1) resolve as intended, down to the kernel module's enable_backward option; Neon honors an explicit override too, and JAX is untouched. The 86 Warp-backend tests in tests/kernels, tests/grids/test_grid_warp.py and tests/boundary_conditions pass. Co-authored-by: Cursor <cursoragent@cursor.com>
JAX 0.10.0 made PartitionSpec stop inheriting from tuple and stop comparing
equal to one, so 'f.sharding.spec == ("cardinality", "x", "y")' became False
and the 2D/3D grid tests failed across all six parametrizations. requirements.txt
declares jax[cuda] unpinned, so the environment resolved to 0.10.2 and picked up
the change.
The sharding XLB produces was never wrong: field shape, sharding mesh, and the
spec's axes all match what the tests expect. Only the comparison was stale, which
is why the shape and mesh assertions kept passing and just the spec one failed.
Full suite is now 94 passed, 0 failed.
Co-authored-by: Cursor <cursoragent@cursor.com>
SIM105: the nested try/except/pass around the cuda:0 fallback becomes
contextlib.suppress(Exception). Behavior is unchanged - try the device from
XLB_WARP_DEVICE, fall back to the first GPU, and otherwise leave Warp's own
default in place.
Also collapse the _pi_vec construction in second_moment.py, where the Warp 1.16
port's mechanical 'wp.vec(' -> 'wp.types.vector(length=' rewrite left the call
split as 'wp.types.vector(length=' / '_pi_dim,'. It now reads as one line,
matching _f_vec directly above it. Whitespace only.
'ruff check .' passes and xlb/ plus tests/ are fully formatted; the one
remaining format deviation, examples/cfd/differentiable_lbm.py, is untouched
upstream code. Suite still 94 passed.
Co-authored-by: Cursor <cursoragent@cursor.com>
tests/run_backend_test_envs.py builds a virtualenv per backend, installs XLB editable from the repo root, reports which Warp/Neon/JAX the env resolved to, and runs the suite in each: python tests/run_backend_test_envs.py # warp and neon python tests/run_backend_test_envs.py --backend neon --reuse warp: pip install -e .[warp,test] -> warp-lang from PyPI neon: pip install -e .[neon,test] -> neon_gpu wheel and its bundled Warp fork JAX is in install_requires, so both profiles get it regardless of backend; --jax cuda adds setup.py's cuda extra to swap in a GPU jaxlib. tests/backends/test_neon_backend.py adds the suite's first Neon backend coverage. Previously nothing exercised it: every test used JAX or WARP, and tests/install/flow_past_sphere_3d_test.py skips Neon because the sphere's HalfwayBounceBackBC has no Neon implementation. These four tests cover what needs no boundary condition - neon.init() registering XLB's custom Warp types, grid and field allocation, the device fill plus host read-back, and operator construction - and skip when neon is absent or no CUDA device exists, so the warp profile stays green. The runner also warns when a Neon env still has the warp-lang distribution registered. Both ship files at site-packages/warp and XLB's install_requires pulls warp-lang in regardless of the [neon] extra, so whichever pip unpacks last owns them; tests pass either way, which makes it invisible otherwise. Verified: both profiles build from scratch and pass (98 tests each, ~6 min per profile); the Neon tests pass alone and with Warp kernels built first in the same session; pytest's exit code 5 on an all-skipped selection is reported as 'no tests ran' rather than a failure. Co-authored-by: Cursor <cursoragent@cursor.com>
pip install -e .[neon,test] left both neon_gpu and warp-lang installed. Both
unpack into site-packages/warp, so whichever pip wrote last owned the
directory: a later 'pip uninstall warp-lang' would delete files Neon needs, and
an upgrade would silently replace Neon's fork. The env still imported and the
whole suite still passed, so nothing caught it.
warp-lang was listed in install_requires, which made the collision unavoidable
for [neon]. The InstallWithNeonHooks cmdclass was meant to undo it but cannot:
pip installs dependencies after building the project, and for wheel and PEP 660
editable installs the 'install' command never runs at all. It also detected the
extra by scanning sys.argv, which PEP 517 builds do not preserve.
So warp-lang moves out of install_requires into the [warp] extra, matching what
the README already documented ('Do not install both in the same environment'),
and the dead hook is removed. Warp now comes from exactly one place: [warp] for
PyPI's build or [neon] for the bundled fork.
pip will not remove a warp-lang that predates the install, so xlb/__init__.py
warns when it finds both distributions and says how to recover. It also turns a
missing Warp into a message naming the two extras instead of a bare
ModuleNotFoundError, since the base install no longer supplies one.
test_neon_install_warp_cleanup.py asserted only that 'import neon' and
'import warp' succeed, which held in the mixed state. It now asserts that a
clean [neon] install has no warp-lang distribution and imports the fork, and
that a pre-existing warp-lang produces the warning. Install docs and the
install-matrix scenarios move to the explicit extras.
Verified by rebuilding both envs from scratch: the warp env resolves warp-lang
1.17.0 with 95 passed and the Neon module skipped, and the neon env resolves
warp 1.16.0.dev0 with warp-lang absent and 99 passed.
Co-authored-by: Cursor <cursoragent@cursor.com>
The extra allowed warp-lang>=1.10.0, so a fresh install resolved 1.17.0 while XLB's kernels target the 1.16 API and the [neon] extra gets a 1.16 fork. Pin to >=1.16,<1.17 so both extras present the same API surface. PyPI currently has one release in the series, 1.16.0; the range leaves room for a future 1.16 patch. Verified by rebuilding the warp env on the pin: it resolves warp-lang 1.16.0, every Warp API XLB uses is present (types.vector, types.matrix, matrix_from_cols, clear_kernel_cache, Tape, Mesh), and the suite is 95 passed with the Neon module skipped. Co-authored-by: Cursor <cursoragent@cursor.com>
Narrowing to <1.17 excluded 1.17.0, which the suite passes on (95 passed, same as 1.16.0). Widen to >=1.16,<1.18 so both tested series are allowed while 1.18 stays out until it has been run. A fresh install now resolves 1.17.0 again, the highest allowed; install warp-lang==1.16.* to match the series of the fork neon_gpu bundles. Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Contributing Guidelines
Description
Upgrade XLB to be compatible with the new API in Warp 1.16.0. or 1.17.0
For the Neon backend, we are using an out-of-tree Warp branch based on 1.16.0.dev0. This branch gives us access to Warp's 'external-type' support until it is merged upstream. The 'external-type' mechanism provides an interface to Neon's CUDA runtime.
The PR also disables Warp compilation of backward-pass kernels by default. This feature must be enabled manually when using XLB with automatic differentiation in the Warp backend.
A new script that run all the tests with all the different backends (JAX, Warp, Neon) has been added.
Type of change
How Has This Been Tested?
Linting and Code Formatting
Make sure the code follows the project's linting and formatting standards. This project uses Ruff for linting.
To run Ruff, execute the following command from the root of the repository:
ruff check .