Repository: jax-ml/jax
Pull request #27911 (2025-04-10 14:42:02)

[Pallas] Interpret dimensions with parallel semantics by traversing the corresponding grid coordinates in randomized order.
[Pallas] Interpret dimensions with parallel semantics by traversing the corresponding grid coordinates in randomized order.

Note that dynamic grid dimensions with 'parallel' semantics are disallowed. This enables the computation of grid points, with randomized coordinates along 'parallel' dimensions, in Jax/on device.
If randomization of grid dimensions with dynamic sizes (i.e. sizes not known at Jax trace time) were allowed, this would require computing these randomizations on the host/on CPU (where one can have arrays of dynamic shape).



Patch:
diff --git a/jax/_src/pallas/mosaic/interpret.py b/jax/_src/pallas/mosaic/interpret.py
index a980590ba..13b321424 100644
--- a/jax/_src/pallas/mosaic/interpret.py
+++ b/jax/_src/pallas/mosaic/interpret.py
@@ -20,7 +20,7 @@ import gc
 import itertools
 import math
 import threading
-from typing import Any, Literal
+from typing import Any, Callable,Literal
 
 import jax
 from jax import lax
@@ -38,6 +38,7 @@ from jax._src import pjit
 from jax._src.state import discharge as state_discharge
 from jax._src.state import indexing
 from jax._src.state import primitives as state_primitives
+from jax._src.typing import Array
 from jax._src.util import (
     safe_map,
     safe_zip,
@@ -74,10 +75,10 @@ class TPUInterpretParams:
       is waiting on a DMA semaphore that will be signaled when the read or write
       is complete.
       Default: "on_wait".
-    detect_races: If True, a dynamic, happens-before race detector will be
-      used to detect data races during kernel interpretation.  If any races are
-      detected, a message will be printed and `races.races_found` will be set
-      to True.
+    detect_races: If True, a dynamic, happens-before race detector will be used
+      to detect data races during kernel interpretation.  If any races are
+      detected, a message will be printed and `races.races_found` will be set to
+      True.
       Default: False.
     skip_floating_point_ops: If True, operations that produce only floating
       point values will not be interpreted; instead, their results will be
@@ -85,14 +86,25 @@ class TPUInterpretParams:
       operands to any operation will be replaced with (arrays of) `jnp.inf`.
       Default: False.
     uninitialized_memory: If "nan", allocated buffers are initialized to
-      to contain all NaNs (or to their maximum possible value for integers).
-      If "zero", allocated buffers are initialized to all zeros.
+      contain all NaNs (or to their maximum possible value for integers). If
+      "zero", allocated buffers are initialized to all zeros.
       Default: "nan".
+    random_seed: Seed for random number generator used during interpretation.
+      Currently random numbers are used to randomize the grid coordinates along
+      dimensions with 'parallel' semantics.
+      Default: None.
+    grid_point_recorder: Callback that is invoked by the interpreter for each
+      grid point in the order in which the grid points are traversed. This is
+      intended for inspecting the randomization of coordinates along grid
+      dimensions with 'parallel' semantics.
+      Default: None.
   """
   dma_execution_mode: Literal["eager", "on_wait"] = "on_wait"
   detect_races: bool = False
   skip_floating_point_ops: bool = False
   uninitialized_memory: Literal["nan", "zero"] = "nan"
+  random_seed: int | None = None
+  grid_point_recorder: Callable[[tuple[jnp.int32, ...]], None] | None = None
 
 
 VectorClock = np.ndarray
@@ -1358,6 +1370,96 @@ def _get_next_indices(grid, indices):
     next_indices.append(jnp.where(carry, 0, i))
   return tuple(reversed(next_indices))
 
+def _get_parallel_dim_semantics(
+    compiler_params: dict[str, Any], grid: tuple[int, ...]
+) -> tuple[bool, ...]:
+  """Returns a tuple of booleans indicating whether the corresponding dimension in `grid` is parallel."""
+  dimension_semantics = compiler_params.get('mosaic', {}).get(
+      'dimension_semantics', None
+  )
+  if dimension_semantics is None:
+    return (False,) * len(grid)
+  return tuple(ds == 'parallel' for ds in dimension_semantics)
+
+_GridPointCoordinatesPerDim = tuple[Array, ...]
+
+def _get_randomized_grid_coordinates(
+    grid: tuple[int, ...],
+    compiler_params: dict[str, Any],
+    random_seed: int | None,
+) -> _GridPointCoordinatesPerDim:
+  """Returns a tuple of randomized coordinates for each 'parallel' dimension in `grid`.
+
+  For a dimension with 'parallel' semantics at position `d` in the grid, the
+  returned tuple contains a random permutation of the sequence `[0,...,
+  grid[d] - 1]` at index `d`. For each dimension with 'arbitrary' semantics,
+  the resulting tuple contains an empty array. (Inserting an empty arry for an
+  'arbitrary' dimension at position `d` in the grid, instead of the sequence
+  `[0,..., grid[d] - 1]`, allows `grid[d]` to be a dynamic value, i.e. a value
+  not known at Jax trace time.)
+
+  Args:
+    grid: Tuple of sizes of the dimensions in the grid.
+    compiler_params: Representation of a `mosaic_core.TPUCompilerParams` object
+      as a dictionary.
+    parallel_semantics_per_dim: A tuple of booleans indicating whether the
+      corresponding dimension in the grid has parallel semantics.
+    random_seed: The seed to use for randomizing coordinates in parallel
+      dimensions.
+  """
+  parallel_semantics_per_dim = _get_parallel_dim_semantics(
+      compiler_params, grid
+  )
+
+  key = jax.random.key(random_seed or 0)
+  grid_point_coordinates = []
+  for dim_size, parallel_dim in zip(grid, parallel_semantics_per_dim):
+    if parallel_dim:
+      # The size of a dimension with `parallel` semantics must be known at Jax
+      # trace time. This ensures that the arguments to `jnp.arange` and
+      # `jax.random.permutation` below are valid.
+      dim_size = jax_core.concrete_or_error(None, dim_size)
+
+      coordindates_along_dim = jnp.arange(dim_size, dtype=jnp.int32)
+      key, subkey = jax.random.split(key)
+      coordindates_along_dim = jax.random.permutation(
+          subkey, coordindates_along_dim
+      )
+      grid_point_coordinates.append(coordindates_along_dim)
+    else:
+      grid_point_coordinates.append(jnp.array((), dtype=jnp.int32))
+
+  return tuple(grid_point_coordinates)
+
+
+def _get_grid_point(
+    loop_indices: tuple[Array, ...],
+    grid_point_coordinates: _GridPointCoordinatesPerDim,
+) -> Array:
+  """Indexes each entry in `grid_point_coordinates` with the corresponding entry in `loop_indices`.
+
+  If an entry in `grid_point_coordinates` is an empty array, the corresponding
+  entry in the returned array is the corresponding entry in `loop_indices`.
+  Otherwise, the returned array contains the entry in `grid_point_coordinates`
+  indexed with the corresponding entry in `loop_indices`.
+
+  Args:
+    loop_indices: A tuple of loop indices.
+    grid_point_coordinates: A tuple of coordinate arrays for each dimension in
+      the grid. Dimensions with 'arbitrary' semantics are represented by empty
+      arrays. Dimensions with 'parallel' semantics are represented by arrays of
+      randomized coordinates.
+
+  Returns:
+    A 1-dimensional array containing the coordinates for the grid point
+    corresponding to the specified `loop_indices`.
+  """
+  grid_point = []
+  for li, coords in zip(loop_indices, grid_point_coordinates):
+    grid_point.append(li if jnp.size(coords) == 0 else coords[li])
+  return jnp.array(grid_point, dtype=np.int32)
+
+
 def _maybe_dynamic_slice(start_idx, block_shape, value, is_indexing):
   start_idx = tuple(jnp.array(s, dtype=jnp.int32) for s in start_idx)
   output = lax.dynamic_slice(value, start_idx, slice_sizes=block_shape)
@@ -1411,7 +1513,7 @@ def interpret_pallas_call(
     input_output_aliases: tuple[tuple[int, int], ...],
     grid_mapping: GridMapping,
     mesh: pallas_core.Mesh | None,
-    compiler_params: mosaic_core.TPUCompilerParams,
+    compiler_params: dict[str, Any],
     cost_estimate: CostEstimate,
     out_avals: tuple[jax_core.AbstractValue, ...],
     interpret_params: TPUInterpretParams,
@@ -1568,6 +1670,10 @@ def interpret_pallas_call(
     # Base case is always one iteration when grid is ()
     num_iterations = 1
 
+  randomized_grid_coordinates = _get_randomized_grid_coordinates(
+      grid, compiler_params, interpret_params.random_seed  # type: ignore[arg-type]
+  )
+
   def _get_local_grid_env(loop_idx):
     if grid_mapping.local_grid_env is not None:
       return grid_mapping.local_grid_env(loop_idx, grid)
@@ -1607,13 +1713,19 @@ def interpret_pallas_call(
       The carry for the next iteration.
     """
     iteration_idx, loop_idx, prev_start_indices, cur_start_indices = carry
+    if interpret_params.grid_point_recorder is not None:
+      grid_point = _get_grid_point(loop_idx, randomized_grid_coordinates)
+      callback.io_callback(interpret_params.grid_point_recorder, (), grid_point)
 
     with pallas_core.grid_env(_get_local_grid_env(loop_idx)):
       next_loop_idx = _get_next_indices(grid, loop_idx)
+      next_grid_point = _get_grid_point(
+          next_loop_idx, randomized_grid_coordinates
+      )
       next_start_indices = [
           _compute_start_indices(
               bm,
-              next_loop_idx,
+              next_grid_point,
               *scalar_buffer_ids,
               compiler_params=compiler_params,
               interpret_params=interpret_params,
@@ -1739,11 +1851,14 @@ def interpret_pallas_call(
       return iteration_idx + 1, next_loop_idx, cur_start_indices, next_start_indices
 
   initial_loop_idx = (jnp.int32(0),) * len(grid)
+  initial_grid_point = _get_grid_point(
+      initial_loop_idx, randomized_grid_coordinates
+  )
   with pallas_core.grid_env(_get_local_grid_env(initial_loop_idx)):
     initial_start_indices = [
         _compute_start_indices(
             bm,
-            initial_loop_idx,
+            initial_grid_point,
             *scalar_buffer_ids,
             compiler_params=compiler_params,
             interpret_params=interpret_params,


Test patch:
diff --git a/tests/pallas/tpu_pallas_interpret_test.py b/tests/pallas/tpu_pallas_interpret_test.py
index afb573f8c..c4bf07f39 100644
--- a/tests/pallas/tpu_pallas_interpret_test.py
+++ b/tests/pallas/tpu_pallas_interpret_test.py
@@ -32,6 +32,7 @@ import numpy as np
 
 
 jax.config.parse_flags_with_absl()
+jax.config.update('jax_threefry_partitionable', True)
 
 
 class CountStoreCallbacksContext(object):
@@ -58,6 +59,29 @@ class CountStoreCallbacksContext(object):
     return self._num_stores
 
 
+class GridPointRecorderContext(object):
+  """Records grid points in the order in which they are traversed."""
+
+  def __init__(self):
+    self._grid_points = []
+
+  def __enter__(self):
+    return self
+
+  def __exit__(self, ty, value, traceback):
+    ...
+
+  def get_recorder(self):
+    def _recorder(grid_point):
+      self._grid_points.append(grid_point)
+
+    return _recorder
+
+  @property
+  def grid_points(self):
+    return self._grid_points
+
+
 class InterpretTest(jtu.JaxTestCase):
 
   def setUp(self):
@@ -326,6 +350,142 @@ class InterpretTest(jtu.JaxTestCase):
       np.testing.assert_allclose(result[::8, ::256], [[1.0], [5.0]])
       self.assertEqual(store_callbacks_counter.num_stores, 5)
 
+  def test_randomization_of_parallel_dimensions(self):
+    def kernel(s_ref, o_ref):
+      s = s_ref[0]
+      s_ref[0] = s + 1
+      o_ref[:] = jax.lax.full_like(o_ref, s)
+
+    def kernel_call_dimensions_arbitrary_parallel(s, grid_point_recorder):
+      return pl.pallas_call(
+          kernel,
+          out_shape=jax.ShapeDtypeStruct((32, 512), jnp.float32),
+          grid=(4, 4),
+          in_specs=[pl.BlockSpec(memory_space=pltpu.SMEM)],
+          out_specs=pl.BlockSpec((8, 128), lambda i, j: (i, j)),
+          interpret=mosaic_interpret.TPUInterpretParams(
+              random_seed=12345, grid_point_recorder=grid_point_recorder
+          ),
+          compiler_params=pltpu.TPUCompilerParams(
+              dimension_semantics=('arbitrary', 'parallel')
+          ),
+      )(s)
+
+    with GridPointRecorderContext() as grid_point_recorder:
+      result = jax.jit(
+          kernel_call_dimensions_arbitrary_parallel, static_argnums=1
+      )(
+          jnp.zeros((1,), jnp.int32),
+          grid_point_recorder.get_recorder(),
+      )
+      np.testing.assert_allclose(
+          result[::8, ::128],
+          [
+              [ 2.0,  3.0,  0.0,  1.0],
+              [ 6.0,  7.0,  4.0,  5.0],
+              [10.0, 11.0,  8.0,  9.0],
+              [14.0, 15.0, 12.0, 13.0],
+          ],
+      )
+      np.testing.assert_array_equal(
+          grid_point_recorder.grid_points,
+          [
+              [0, 2],
+              [0, 3],
+              [0, 0],
+              [0, 1],
+              [1, 2],
+              [1, 3],
+              [1, 0],
+              [1, 1],
+              [2, 2],
+              [2, 3],
+              [2, 0],
+              [2, 1],
+              [3, 2],
+              [3, 3],
+              [3, 0],
+              [3, 1],
+          ],
+      )
+
+    def kernel_call_dimensions_parallel_arbitrary(s, grid_point_recorder):
+      return pl.pallas_call(
+          kernel,
+          out_shape=jax.ShapeDtypeStruct((32, 512), jnp.float32),
+          grid=(4, 4),
+          in_specs=[pl.BlockSpec(memory_space=pltpu.SMEM)],
+          out_specs=pl.BlockSpec((8, 128), lambda i, j: (i, j)),
+          interpret=mosaic_interpret.TPUInterpretParams(
+              random_seed=12345, grid_point_recorder=grid_point_recorder
+          ),
+          compiler_params=pltpu.TPUCompilerParams(
+              dimension_semantics=('parallel', 'arbitrary')
+          ),
+      )(s)
+
+    with GridPointRecorderContext() as grid_point_recorder:
+      result = jax.jit(
+          kernel_call_dimensions_parallel_arbitrary, static_argnums=1
+      )(
+          jnp.zeros((1,), jnp.int32),
+          grid_point_recorder.get_recorder(),
+      )
+      np.testing.assert_allclose(
+          result[::8, ::128],
+          [
+              [ 8.0,  9.0, 10.0, 11.0],
+              [12.0, 13.0, 14.0, 15.0],
+              [ 0.0,  1.0,  2.0,  3.0],
+              [ 4.0,  5.0,  6.0,  7.0],
+          ],
+      )
+      np.testing.assert_array_equal(
+          grid_point_recorder.grid_points,
+          [
+              [2, 0],
+              [2, 1],
+              [2, 2],
+              [2, 3],
+              [3, 0],
+              [3, 1],
+              [3, 2],
+              [3, 3],
+              [0, 0],
+              [0, 1],
+              [0, 2],
+              [0, 3],
+              [1, 0],
+              [1, 1],
+              [1, 2],
+              [1, 3],
+          ],
+      )
+
+  def test_dynamic_parallel_dimension_raises(self):
+    def kernel(o_ref):
+      o_ref[0] = 42.0
+
+    @jax.jit
+    def kernel_call_dynamic_parallel_dimension():
+      dim_size = jax.random.randint(
+          jax.random.key(0), (), 10, 20, dtype=jnp.int32
+      )
+      return pl.pallas_call(
+          kernel,
+          out_shape=jax.ShapeDtypeStruct((1,), jnp.float32),
+          grid=(dim_size,),
+          in_specs=[],
+          out_specs=pl.BlockSpec((1,), lambda _: (0,)),
+          interpret=mosaic_interpret.TPUInterpretParams(),
+          compiler_params=pltpu.TPUCompilerParams(
+              dimension_semantics=('parallel',)
+          ),
+      )()
+
+    with self.assertRaises(jax.errors.ConcretizationTypeError):
+      kernel_call_dynamic_parallel_dimension()
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())