Multi-Dimensional Mesh DMA in Pallas Interpret Mode

[R-Oliver]

Hey all,

I'm prototyping an all-gather kernel using Pallas with shard_map on a 2D mesh. When using make_async_remote_copy with DeviceIdType.MESH, interpret mode raises:

NotImplementedError: Meshes with more than 1 named dimension not implemented in dma_start_p

jax/jax/_src/pallas/mosaic/primitives.py

Lines 619 to 621 in fc1b321

	if device_id_len > 1 or len(nonempty_axes) > 1:
	raise NotImplementedError("Meshes with more than 1 named dimension not "
	"implemented in dma_start_p")

The same code runs correctly on the TPU v5e-4, and I wanted to see if there was a common workaround here that doesn't involve flattening the mesh axes.

Reproducer:

import jax
from jax.experimental import pallas as pl
from jax.experimental.pallas import tpu as pltpu
import jax.numpy as jnp
from jax.sharding import PartitionSpec as P

mesh = jax.make_mesh((2, 2), ("x", "y"))
x = jax.device_put(
    jnp.ones((2048, 1024), dtype=jnp.bfloat16),
    jax.NamedSharding(mesh, P("x", None)),
)

def kernel(input_ref, output_ref, send_sem, recv_sem):
    this_x = jax.lax.axis_index("x")
    this_y = jax.lax.axis_index("y")
    neighbor_x = (this_x + 1) % 2

    dma = pltpu.make_async_remote_copy(
        src_ref=input_ref,
        dst_ref=output_ref,
        send_sem=send_sem,
        recv_sem=recv_sem,
        device_id=(neighbor_x, this_y),
        device_id_type=pltpu.DeviceIdType.MESH,
    )
    dma.start()
    dma.wait()

def pallas_call(x):
    return pl.pallas_call(
        kernel,
        grid_spec=pltpu.PrefetchScalarGridSpec(
            num_scalar_prefetch=0,
            grid=(1,),
            in_specs=[pl.BlockSpec(memory_space=pl.ANY)],
            out_specs=pl.BlockSpec(memory_space=pl.ANY),
            scratch_shapes=[pltpu.SemaphoreType.DMA] * 2,
        ),
        out_shape=jax.ShapeDtypeStruct((1024, 1024), jnp.bfloat16),
        interpret=True,
    )(x)

jax.jit(jax.shard_map(pallas_call, mesh=mesh, in_specs=P("x", None), out_specs=P("x", None), check_vma=False))(x)

It looks like the TPU lowering path linearizes ND mesh coordinates to a scalar logical device ID

jax/jax/_src/pallas/primitives.py

Lines 1431 to 1478 in fc1b321

	def device_id_to_logical(
	mesh_context: pallas_utils.MeshInfo | None,
	device_id: ir.Value | tuple[ir.Value, ...] | dict[Any, ir.Value],
	device_id_type: DeviceIdType,
	get_axis_index,
	) -> tuple[ir.Value | None, dict[Any, ir.Value]]:
	"""Normalizes a device id into a logical device id and axes that don't correspond to JAX mesh axes.
	The indexing implied by the returned axis dict should be handled by the
	caller. If there are no cross-device operations, then the returned logical
	device id will be None.
	"""
	non_mesh_axes = {}
	if isinstance(device_id, dict):
	if device_id_type is not DeviceIdType.MESH:
	raise ValueError(
	"`device_id_type` must be MESH if `device_id` is a dict,"
	f" got: {device_id_type = }."
	)
	device_id, non_mesh_axes = _device_id_dict_to_mesh(mesh_context, device_id, get_axis_index)
	if device_id_type is DeviceIdType.MESH:
	# Mesh means we are passed the mesh coordinates for the device
	device_ids = tree_util.tree_leaves(device_id)
	mesh_strides: tuple[int, ...]
	if mesh_context is None:
	mesh_strides = ()
	else:
	mesh_strides = mesh_context.mesh_strides
	if len(device_ids) != len(mesh_strides):
	raise ValueError(
	"Number of device ids must match the number of mesh axes, but got"
	f" {len(device_ids)} ids for a {len(mesh_strides)}D mesh."
	)
	i32 = ir.IntegerType.get_signless(32)
	if not device_ids:
	# If there are no device ids, then it is purely local communication.
	return None, non_mesh_axes
	return functools.reduce(
	arith.addi,
	(
	arith.muli(a, arith.constant(i32, b))
	for a, b in zip(device_ids, mesh_strides)
	),
	), non_mesh_axes
	elif device_id_type is DeviceIdType.LOGICAL:
	return device_id, non_mesh_axes
	raise NotImplementedError(f"Unsupported device id type: {device_id_type}")

and interpret/utils.py contains device_coords_to_logical_id() with equivalent logic

jax/jax/_src/pallas/mosaic/interpret/utils.py

Lines 206 to 220 in fc1b321

	def device_coords_to_logical_id(device_coords, axis_sizes, axis_indices):
	if isinstance(device_coords, dict):
	device_coords, non_mesh_axes = _device_id_dict_to_mesh(
	device_coords, axis_sizes, axis_indices
	)
	if non_mesh_axes:
	raise NotImplementedError(non_mesh_axes)
	if not isinstance(device_coords, tuple):
	device_coords = (device_coords,)
	assert len(device_coords) == len(axis_sizes)
	sizes = list(axis_sizes.values())
	ret = 0
	for i in range(len(device_coords)):
	ret += device_coords[i] * math.prod(sizes[i + 1 :])
	return ret

Questions:

• Would it be possible to replace the NotImplementedError with linearization logic mirroring the TPU path?
• Is this the right approach? Or is there a reason the interpret mode was kept single-axis?