JAX unable to optimize reduce operation, leading to OOM issues or really slow jax.lax.scan / jax.lax.map

[jeertmans]

Hi! First, thanks for this great package!

I use JAX heavily for my research, for which I implemented ray tracing utilities.

One critical operation among many is the ray-triangle intersection test. We usually have millions of rays and millions of triangles, and we would like to test, for each ray, if it intersects with any triangle. A common approach to check for this is to use the Möller-Trumbore algorithm, which I implemented using JAX.

However, I am struggling to improve the performance of this function, and I am reaching out to see if anyone has any suggestions. I am considering moving to Mitsuba (or using DrJit directly) or implementing a C++ (or Rust) version, but I hope I can find a way to improve the code without leaving JAX.
After reading some related discussions and issues (https://stackoverflow.com/questions/77527847/jax-vmap-limit-memory, #11319, https://stackoverflow.com/questions/77659069/how-to-understand-and-debug-memory-usage-with-jax), I realized that JIT compilation could optimize away the need for a large allocation.

Here is the code to reproduce (uv run file.py works):

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "jax[cuda]>=0.7.0",
#     "jaxtyping",
# ]
# ///

import jax
import jax.numpy as jnp
from jaxtyping import Array, Bool, Float


def ray_intersect_triangle(
    ray_origin: Float[Array, "3"],
    ray_direction: Float[Array, "3"],
    triangle_vertices: Float[Array, "3 3"],
) -> Bool[Array, " "]:
    # Placeholder computations for the Möller-Trumbore algorithm
    t = (ray_origin * ray_direction).sum()
    hit = (ray_origin * ray_direction * triangle_vertices[0, :]).sum() > 0
    # Actual ray-intersection boolean value
    return (t <= 1.0) & hit


# Intermediate function, used to show that we perform a reduction on 'num_triangles'
@jax.jit
def _ray_intersect_any_triangle(
    ray_origin: Float[Array, "3"],
    ray_direction: Float[Array, "3"],
    triangle_vertices: Float[Array, "num_triangles 3 3"],
) -> Bool[Array, " "]:
    return jax.vmap(ray_intersect_triangle, in_axes=(None, None, 0))(
        ray_origin, ray_direction, triangle_vertices
    ).any()


# Vectorized function, that is the bottleneck
@jax.jit
def rays_intersect_any_triangle(
    ray_origins: Float[Array, "num_rays 3"],
    ray_directions: Float[Array, "num_rays 3"],
    triangle_vertices: Float[Array, "num_triangles 3 3"],
) -> Bool[Array, "num_rays"]:
    # Solution 1: OOM
    return jnp.vectorize(
        _ray_intersect_any_triangle, excluded={2}, signature="(3),(3)->()"
    )(ray_origins, ray_directions, triangle_vertices)
    # Solution 2: OOM
    return jax.vmap(_ray_intersect_any_triangle, in_axes=(0, 0, None))(
        ray_origins, ray_directions, triangle_vertices
    )
    # Solution 3: takes ages to compute
    return jax.lax.map(
        lambda xs: _ray_intersect_any_triangle(*xs, triangle_vertices),
        (ray_origins, ray_directions),
    )

    # Solution 4: OOM (but uses way less memory), and super slow if we reduce 'num_rays'
    def scan_fun(
        intersect: Bool[Array, "num_rays"], triangle_vertices: Float[Array, "3 3"]
    ) -> tuple[Bool[Array, "num_rays"], None]:
        intersect |= jax.vmap(ray_intersect_triangle, in_axes=(0, 0, None))(
            ray_origins, ray_directions, triangle_vertices
        )
        return intersect, None

    return jax.lax.scan(
        scan_fun,
        init=jnp.zeros(ray_origins.shape[0], dtype=bool),
        xs=triangle_vertices,
    )


if __name__ == "__main__":
    key = jax.random.key(1234)

    num_rays = 1_000_000
    num_triangles = 1_000_000

    key_ray_origins, key_ray_directions, key_triangle_vertices = jax.random.split(
        key, 3
    )
    ray_origins = jax.random.normal(key_ray_origins, (num_rays, 3))
    ray_directions = jax.random.normal(key_ray_directions, (num_rays, 3))
    triangle_vertices = jax.random.normal(key_triangle_vertices, (num_triangles, 3, 3))

    # Boom
    rays_intersect_any_triangle(ray_origins, ray_directions, triangle_vertices)

FYI, I have an NVIDIA 3070 with 8192 MiB of memory, which is not enormous but should do the job as it should not need to allocate a matrix that is too large to fit this.

N.B.: I am aware that I could use a hierarchical representation of the scene to accelerate the intersection-check, but I couldn't find any JAX-compatible library for that, so I would prefer avoiding this solution :'-)

Thanks in advance for your help!

---

EDIT: other (unsuccessful) things I have considered:

• swapping the order of vmaps: inner is on all rays and outer (reduced) is on all triangles;
• same as above, but with the outer vmap have out_axes=-1 (and thus reduction on last axis);
• look for XLA compiler flags that could indicate the compiler that it must find a way to reduce memory usage;
• specify inline=True with jax.jit.

[justinjfu]

A simple suggestion based on what you already have: You can try using the batch_size argument to jax.lax.map to apply vmap to chunks of the input instead to save on memory compared to using vmap directly.

[jeertmans]

Hi, thanks for the suggestion! I tried this, but it only temporarily delays the problem. You just reduce the allocated size by the factor of len(array)/batch_size. Another problem with this solution is deciding on the batch size, given that num_rays and num_triangles can vary by orders of magnitude and the available size on the GPU is unknown a priori. Perhaps there is a good heuristic, of which I am unaware.

[jeertmans]

Some updates

After many attempts, and inspired by #4968, I managed to improve the performance by creating a custom batched version of jax.lax.scan, using jax.lax.fori_loop and indexing batches with jax.lax.dynamic_slices_in_dim. This significantly speeds-up the computation for medium-sized arrays, and avoid OOM issues for large arrays (but still takes > 20 seconds to compute on a GPU, at least).

Pseudocode:

@jax.jit
def rays_intersect_any_triangle(
    ray_origins: Float[Array, "num_rays 3"],
    ray_directions: Float[Array, "num_rays 3"],
    triangle_vertices: Float[Array, "num_triangles 3 3"],
) -> Bool[Array, "num_rays"]:
    batch_size = 1024
    num_batches = triangle_vertices.shape[-3] // batch_size
    # ... make sure to pad arrays to be a multiple of batch_size
    def body_fun(batch_index, intersect):
        batch_of_triangle_vertices = jax.lax.dynamic_slice_in_dim(
            triangles_vertices, batch_index * batch_size, batch_size, axis=-3)
        return intersect | jax.vmap(ray_intersect_triangle)(...).any(axis=-1)
    return jax.lax.fori_loop(0, num_batches, body_fun, init_val=jnp.zeros(..., dtype=bool))

The actual code (if you are curious): https://github.com/jeertmans/DiffeRT/pull/300/files.