Multi-res grid refinement + Neon backend support

[hsalehipour]

Contributing Guidelines

• I have read and understood the CONTRIBUTING.md guidelines

Description

Grid refinement capability is now supported in XLB through the Neon backend. The Neon backend provides full support for dense grids on multi-GPU systems, as well as multi-resolution grids on single GPUs. All newly introduced functionalities have been carefully tested and optimized. This represents a major enhancement to the library and involves substantial additions and improvements to the codebase.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update

How Has This Been Tested?

• All pytest tests pass

============================================= test session starts ==============================================
platform linux -- Python 3.12.3, pytest-9.0.2, pluggy-1.6.0
rootdir: /home/max/repo/test/XLB
collected 93 items                                                                                             

tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_jax.py ....                                 [  4%]
tests/boundary_conditions/bc_equilibrium/test_bc_equilibrium_warp.py ....                                [  8%]
tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_jax.py ......               [ 15%]
tests/boundary_conditions/bc_fullway_bounce_back/test_bc_fullway_bounce_back_warp.py ......              [ 21%]
tests/boundary_conditions/mask/test_bc_indices_masker_jax.py .......                                     [ 29%]
tests/boundary_conditions/mask/test_bc_indices_masker_warp.py ......                                     [ 35%]
tests/grids/test_grid_jax.py .......                                                                     [ 43%]
tests/grids/test_grid_warp.py ....                                                                       [ 47%]
tests/kernels/collision/test_bgk_collision_jax.py ......                                                 [ 53%]
tests/kernels/collision/test_bgk_collision_warp.py ......                                                [ 60%]
tests/kernels/equilibrium/test_equilibrium_jax.py ......                                                 [ 66%]
tests/kernels/equilibrium/test_equilibrium_warp.py ......                                                [ 73%]
tests/kernels/macroscopic/test_macroscopic_jax.py ......                                                 [ 79%]
tests/kernels/macroscopic/test_macroscopic_warp.py .......                                               [ 87%]
tests/kernels/stream/test_stream_jax.py ......                                                           [ 93%]
tests/kernels/stream/test_stream_warp.py ......                                                          [100%]

======================================== 93 passed in 248.34s (0:04:08) ========================================

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 .

• Ruff passes

[mehdiataei]

Thank you @massimim, this looks like a strong contribution. I left code-level comments in this round, but I also wanted to share a few higher-level suggestions.

1. I think framing Neon as a new compute "backend" may be slightly misleading. Conceptually, Neon here does not seem fully parallel to JAX. The implementation largely reuses Warp functionals and then executes them through Neon handles, containers, and skeletons. In that sense, Neon feels more like an execution/runtime layer on top of Warp code generation than a standalone compute backend.

I think a different framing could make the design clearer:

• If Neon is fundamentally "Warp math + Neon execution", I would be hesitant to model it as a third peer backend throughout the operator hierarchy.
• Instead, I would consider splitting the abstraction into:
  • kernel / math backend: JAX vs Warp
  • execution runtime: direct Warp launch vs Neon container/skeleton launch

I think this would make Neon more generic and would better highlight its real strength: the execution model and skeleton abstraction, rather than presenting it as a bespoke backend. It may also make the integration easier to extend and adopt. Several of the current issues feel like symptoms of the abstraction boundary being one layer off. This likely needs some careful design thought, but I would strongly encourage it. To me, the more compelling framing is that Neon provides a skeleton/runtime that Warp kernels can target.

2. The multires implementation also feels too monolithic. Kernels, schedule planning, state ownership, and runtime graph compilation all live in roughly the same layer, which makes the system harder to reason about, test, and extend.

One possible improvement would be to introduce a typed MultiresPlan / Schedule layer that represents the recursive timestep as explicit operations, then have a separate Neon graph builder that lowers that plan into containers/skeletons. I would also keep simulation state in a manager and keep kernels separate from schedule construction.

3. The topology and coordinate model feels too implicit at the moment, which makes it harder to debug and reuse. An explicit MultiresTopology or LevelInfo abstraction could help a lot, with methods such as:

• level_shape(level)
• global_bounds(level)
• to_global(level, coords)
• face_indices(level, side)
• active_indices(level)

I think making those concepts explicit would improve both clarity and correctness.

4. For the new BC, I suggest clearly clarifying the Re that it has been validated for.

[mehdiataei]

you can delete WarpGrid import

[massimim]

Hi @mehdiataei , thank you for the review. This PR is the result of quite a long journey with @hsalehipour and others too.

I'll go through points 1 to 3.

1. I see your point on the compute "backends" as Neon reuses directly most of the functionals defined by the warp backend as the python version of Neon is built on top of warp. It is different however if we consider their programming models. Neon follows a structured parallel pattern approach as it raises the abstraction level by fully hiding the complexity of multi-GPU systems (for the dense grid). I think the critical point for determining the best structure would be until we get the automatic differentiation for Neon too. I'll open an issue to track this.
2. I agree on the structure to be too monolithic. There are some changes to the Neon API that are coming that will also need to be considered. Once these are in we can consider a restructuring.
3. I agree here too. There are still some mechanisms that Neon does not expose yet, but will be there hopefully soon along with a new graph mechanism to improve on independent kernels.

@hsalehipour and I are addressing the remaining comments and will send a PR revision.

[hsalehipour]

Thanks a lot, @mehdiataei , for the thorough review. We have now addressed the remaining issues in the latest commits, and the changes are ready to be merged.

[mehdiataei]

Hi @hsalehipour , maybe I’m missing something, but are there any benchmarks reported for this implementation? For example, DrivAer at the reference Reynolds number of 4.87M.

I think this is important, since people may use this for industrial applications. If those benchmarks cannot be passed, it would be helpful to clearly document the limits of the method in an MD file inside the example folder, along with the multi-resolution example and the relevant explanations.

Also, there are a number of WIP/debugging commits that should be squashed. Could you please clean those up?

[hsalehipour]

@mehdiataei Agreed. The new example multires_windtunnel_3d already serves that purpose and relies on ahmed.json for reference data.

We attempted to squash just the WIP commits via interactive rebase. The squash itself worked, but replaying the ~190 commits that sit above the squash range produced merge conflicts in the upper commits, and we aborted to keep the branch intact. Resolving those conflicts is feasible but non-trivial, so we're proposing simple merge instead as we did for the OOC PR.

[mehdiataei]

What is the Reynolds number range within which the BC is intended to operate? I believe this should be explicitly clarified, particularly since the BC was previously presented as a novel design. In addition, I suggest either add a README file explaining the validity and add the tables etc or cite the appropriate literature on the method.

Merging such a large number of commits (269 commits while the full library history is just 369 commits) into one PR is fairly unconventional and also risks reducing the visibility of substantial prior contributions by pushing them further back in the project history. We can also fix the OCC merge if that's your concern.

Let me know what you think.

[mehdiataei]

If the concern is primarily around conflict resolution or preserving the existing branch state, I would be happy to help with the squashing/rebase process to make the PR history cleaner and more manageable.

Please let me know!

[hsalehipour]

I don’t think it would be appropriate to provide empirical bounds on the stability or accuracy of HybridBC, as these characteristics are highly problem- and geometry-dependent. Based on the experiments conducted so far, the method has demonstrated strong stability. Any discrepancies with reference data are likely influenced by other important modeling components that are not yet incorporated into the library, such as wall models.

Regarding the commit history, Max and I, as maintainers of the library, have decided to follow a squash-and-merge strategy for these integrations in order to keep the history concise and manageable. It would be appreciated if the same approach could also be applied retroactively to the previous OOC merge.