Passing Pytrees to Pallas Kernell

[cdelv]

Hi, I have the following Pallas Kernell that implements a sweep and prune force calculation:

def sap_kernel(state: State, system: System, m, M, forces):
    """
    Currently 1 thread per particle. One could use more threads for bigger particles for work balancing but 
    Pallas does not allow dynamic thread counts.
    """
    i = pl.program_id(0) # Particle i-th
    n, dim = state.r.shape
    M_i = pl.load(M, (i,)) # Upper bound bounding box projection
    pos_i = pl.load(state.r, (i, slice(None)))
    aabb_i = pl.load(state.aabb, (i, slice(None)))

    pl.store(forces, (i, slice(None)), jnp.zeros_like(pos_i)) # Initialize forces at 0

    def cond(j):
        return (j < n) * (pl.load(m, (j,)) <= M_i) # m[j] is traversed linearly due to the sorting. (Good for performance)

    def body(j):
        pos_j = pl.load(state.r, (j, slice(None)))
        aabb_j = pl.load(state.aabb, (j, slice(None)))
        r_ij = rij(pos_i, pos_j, system)
        overlap = jnp.sum(jnp.abs(r_ij) <= (aabb_i + aabb_j)) == dim # check bounding box intersection

        # This if is less performant than computing the force and multiplying by the intersection
        # Should be worth it with more complicated force laws. 
        f = jax.lax.cond(
            overlap,
            lambda _: force_between(state, system, i, j),
            lambda _: jnp.zeros_like(pos_i),
            operand = None
        )
        # Newton's 3rd law
        pl.atomic_add(forces, (i, slice(None)), f)
        pl.atomic_add(forces, (j, slice(None)), -f)
        return j + 1

    jax.lax.while_loop(cond, body, i + 1) # j = i + 1 (avoid self interaction)

And this is how I call the kernel:

    state.a = pl.pallas_call(
        sap_kernel,
        out_shape=state.a,
        grid=(n,),
        interpret=False
    )(state, system, m, M)

This is the definition of state and system:

@jax.tree_util.register_dataclass
@dataclass
class State:
    r: jnp.ndarray      # Position
    v: jnp.ndarray      # Velocity
    a: jnp.ndarray      # Acceleration
    aabb: jnp.ndarray   # AABB half-extents (bounding box)
    rad: jnp.ndarray    # Radius
    mass: jnp.ndarray   # Mass
    ID: jnp.ndarray     # Particle IDs

@jax.tree_util.register_dataclass
@dataclass
class System:
    dt: float
    box_anchor: jnp.ndarray
    box_length: jnp.ndarray 

The problem I have is the usage of the system object. When I try to use it, for example, in the rij() function, the program fails with

ValueError: Attempting to pass a Ref MemRef<None>{float32[2]} to a primitive: sub - did you forget to unpack ([...]) the ref?

However, I have no problem at all using the state Pytree. I tried diverse ways of loading the system object and the data from the system I need, but I had a very similar error.

I also tried to use the in_specs field with BlockSpec, but because I'm passing Pytrees, I get this error:

     * at [0], `in_specs` is a pytree leaf but `inputs` is a <class '__main__.State'>, so their Python types differ
    * at [1], `in_specs` is a pytree leaf but `inputs` is a <class '__main__.System'>, so their Python types differ

I appreciate any recommendations with this. Or if there is an obvious way of loading the data into the system. Here I attach a minimal working example: sim_minimal.txt

[justinjfu]

What is in the contents of the rij() function? You are getting this error because the leaves of system when passed into the kernels are Refs, and you need to dereference/load them as the error message says. You can probably easily load all values in the entire object by using a one-liner such as system = jax.tree.map(lambda x: x[...], system).

However, I have no problem at all using the state Pytree.

You aren't seeing this problem with state because you are using pl.load to load values out of the Refs.

[cdelv]

Hi @justinjfu, thank you for the reply.

The contents of the rij() are just the way of computing the displacement vector. The reason for it to be a function is that the computation changes with periodic boundaries.

@partial(jax.jit, inline=True)
def rij(ri, rj, system):
    rij = (ri - system.box_anchor) - (rj - system.box_anchor)
    return rij  - system.box_length * jnp.round(rij / system.box_length)

Thank you for `system = jax.tree. map (lambda x: x [...], system). ' I didn't know that one could do this, and it's exactly what I've been trying to do all this time. I tried it in the kernel, and it worked. I noticed that it's not super fast. Do you have any tips on performance?

Thanks.

[justinjfu]

I noticed that it's not super fast. Do you have any tips on performance?

At a high level, your code doesn't look like it's vectorized very well? In triton the program id corresponds to a block rather than a single thread as you suggest in your comment. So you're using an entire block to process a single pair of forces at a time. There's other nitpicks (for example, atomic adds are slow) but I think structurally your code should somehow process an entire chunk of objects for each program instance.