Repository: jax-ml/jax
Pull request #27480 (2025-03-26 17:23:25)

[Pallas] Skip reads/writes from/to slices of kernel input/output buffers when the slices do not change between iterations of the grid loop that interprets kernels on CPU.
[Pallas] Skip reads/writes from/to slices of kernel input/output buffers when the slices do not change between iterations of the grid loop that interprets kernels on CPU.



Patch:
diff --git a/jax/_src/pallas/mosaic/interpret.py b/jax/_src/pallas/mosaic/interpret.py
index 5acbabc67..9d7b03ad5 100644
--- a/jax/_src/pallas/mosaic/interpret.py
+++ b/jax/_src/pallas/mosaic/interpret.py
@@ -1307,8 +1307,13 @@ def _compute_start_indices(
       jaxpr.jaxpr, *jaxpr.consts, *loop_idx, *args,
       compiler_params=compiler_params, interpret_params=interpret_params)
   if isinstance(block_mapping.indexing_mode, pallas_core.Blocked):
-    ret = tuple(i if b is pallas_core.mapped else b * i
-                for b, i in zip(block_mapping.block_shape, block_indices))
+    ret = jnp.array(
+        tuple(
+            i if b is pallas_core.mapped else b * i
+            for b, i in zip(block_mapping.block_shape, block_indices)
+        ),
+        dtype=jnp.int32,
+    )
   elif isinstance(block_mapping.indexing_mode, pallas_core.Unblocked):
     ret = block_indices
   else:
@@ -1534,64 +1539,114 @@ def interpret_pallas_call(
     # Base case is always one iteration when grid is ()
     num_iterations = 1
 
-  def body(carry):
-    # The loop carry: (i, loop_idx) --
-    #  - i:int32 is the interation index
-    #  - loop_idx: tuple[int32] are the program ids for each grid axis
-    i, loop_idx = carry
-
+  def _get_local_grid_env(loop_idx):
     if grid_mapping.local_grid_env is not None:
-      local_grid_env = grid_mapping.local_grid_env(loop_idx, grid)
+      return grid_mapping.local_grid_env(loop_idx, grid)
     else:
-      local_grid_env = tuple(
+      return tuple(
           pallas_core.GridAxis(idx, b)
           for dim, (idx, b) in enumerate(zip(loop_idx, grid))
           if dim not in grid_mapping.vmapped_dims
       )
 
-    with pallas_core.grid_env(local_grid_env):
-      start_indices = [
+  def body(
+      carry: tuple[
+          jnp.int32, tuple[jnp.int32, ...], list[jnp.ndarray], list[jnp.ndarray]
+      ],
+  ):
+    """Performs a single iteration of `jaxpr` in the device grid.
+
+    Execution of `jaxpr` is preceded by reading kernel input buffers and
+    followed by writing kernel output buffers.
+
+    Args:
+      carry: (iteration_idx, loop_idx, prev_start_indices, cur_start_indices).
+        - iteration_idx is the interation index.
+        - loop_idx are the program ids for each grid axis.
+        - prev_start_indices is a rank-1 array that contains the start indices
+          for the slices of inputs and outputs processed in the previous loop
+          iteration.
+        - cur_start_indices is a rank-1 array that contains the start indices
+          for the slices of inputs and outputs processed in the current loop
+          iteration.
+
+        Note that by carrying the previous *and* current start indices between
+        loop iterations, it suffices to compute only one list of start indices,
+        i.e. `next_start_indices` (see below), per iteration.
+
+    Returns:
+      The carry for the next iteration.
+    """
+    iteration_idx, loop_idx, prev_start_indices, cur_start_indices = carry
+
+    with pallas_core.grid_env(_get_local_grid_env(loop_idx)):
+      next_loop_idx = _get_next_indices(grid, loop_idx)
+      next_start_indices = [
           _compute_start_indices(
-              bm, loop_idx, *scalar_buffer_ids, compiler_params=compiler_params,
-              interpret_params=interpret_params)
-          for bm in grid_mapping.block_mappings]
+              bm,
+              next_loop_idx,
+              *scalar_buffer_ids,
+              compiler_params=compiler_params,
+              interpret_params=interpret_params,
+          )
+          for bm in grid_mapping.block_mappings
+      ]
+
       # Copy slices of the input to the kernel buffers.
-      #
-      # TODO(jburnim): Only copy slices when the index mapping has changed?
-      for j, var in enumerate(input_vars):
-        if _is_any(var.aval.memory_space):
-          continue
+
+      def _store_slice_to_kernel_input(index, input_var):
         # Copy from the HBM buffer for the pallas_call input to the kernel
         # input buffer.
         # TODO(jburnim): Just use input_args[j] when the input is not aliased?
         transform = indexing.NDIndexer(
-            indices=tuple(indexing.ds(st, sz) if not iid else st
-                          for st, sz, iid in zip(start_indices[j],
-                                                 block_shapes[j],
-                                                 is_indexing_dim[j])),
-            shape=input_args[j].shape,
-            int_indexer_shape=())
+            indices=tuple(
+                indexing.ds(st, sz) if not iid else st
+                for st, sz, iid in zip(
+                    cur_start_indices[index],
+                    block_shapes[index],
+                    is_indexing_dim[index],
+                )
+            ),
+            shape=input_args[index].shape,
+            int_indexer_shape=(),
+        )
         sliced_val = callback.io_callback(
             # TODO(jburnim): Pass source_info from the pallas_call, in case this
             # read is involved in a data race.
             get,
-            jax.ShapeDtypeStruct(var.aval.shape, var.aval.dtype),
+            jax.ShapeDtypeStruct(input_var.aval.shape, input_var.aval.dtype),
             device_id,
             TPU_MEMORY_SPACE_IDXS[mosaic_core.TPUMemorySpace.ANY],
-            input_buffer_ids[j],
+            input_buffer_ids[index],
             (transform,),
-            ordered=True)
+            ordered=True,
+        )
         callback.io_callback(
             # TODO(jburnim): Pass source_info from the pallas_call, in case this
             # store is involved in a data race.
             store,
             (),
             device_id,
-            TPU_MEMORY_SPACE_IDXS[var.aval.memory_space],
-            input_ids[j],
+            TPU_MEMORY_SPACE_IDXS[input_var.aval.memory_space],
+            input_ids[index],
             (),
             sliced_val,
-            ordered=True)
+            ordered=True,
+        )
+
+      for j, var in enumerate(input_vars):
+        if _is_any(var.aval.memory_space):
+          continue
+        assert len(cur_start_indices[j].shape) == 1
+        assert len(prev_start_indices[j].shape) == 1
+        jax.lax.cond(
+            (iteration_idx == 0)
+            | jax.lax.reduce_or(
+                cur_start_indices[j] != prev_start_indices[j], axes=(0,)
+            ),
+            functools.partial(_store_slice_to_kernel_input, j, var),
+            lambda: None,
+        )
 
       # Invoke the kernel.
       _interpret_jaxpr(jaxpr, *kernel_buffer_ids,
@@ -1599,29 +1654,30 @@ def interpret_pallas_call(
                        interpret_params=interpret_params)
 
       # Copy from the kernel buffers to slices of the output in HBM.
-      #
-      # TODO(jburnim): Only copy if the index mapping will change in the
-      # next iteration (or if this is the last iteration)?
-      for j, var in enumerate(output_vars):
-        if _is_any(var.aval.memory_space):
-          continue
+      def _store_to_output_buffer(index, output_var):
         kernel_output_val = callback.io_callback(
             # TODO(jburnim): Pass source_info from the pallas_call, in case this
             # get is involved in a data race.
             get,
-            var.aval,
+            output_var.aval,
             device_id,
-            TPU_MEMORY_SPACE_IDXS[var.aval.memory_space],
+            TPU_MEMORY_SPACE_IDXS[output_var.aval.memory_space],
             kernel_output_ids[j],
             (),
-            ordered=True)
+            ordered=True,
+        )
         transform = indexing.NDIndexer(
-            indices=tuple(indexing.ds(st, sz) if not iid else st
-                          for st, sz, iid  in zip(start_indices[num_inputs + j],
-                                                  block_shapes[num_inputs + j],
-                                                  is_indexing_dim[num_inputs + j])),
-            shape=output_vals[j].shape,
-            int_indexer_shape=())
+            indices=tuple(
+                indexing.ds(st, sz) if not iid else st
+                for st, sz, iid in zip(
+                    cur_start_indices[num_inputs + index],
+                    block_shapes[num_inputs + index],
+                    is_indexing_dim[num_inputs + index],
+                )
+            ),
+            shape=output_vals[index].shape,
+            int_indexer_shape=(index),
+        )
         callback.io_callback(
             # TODO(jburnim): Pass source_info from the pallas_call, in case this
             # store is involved in a data race.
@@ -1629,18 +1685,52 @@ def interpret_pallas_call(
             (),
             device_id,
             TPU_MEMORY_SPACE_IDXS[mosaic_core.TPUMemorySpace.ANY],
-            output_buffer_ids[j],
+            output_buffer_ids[index],
             (transform,),
             kernel_output_val,
-            ordered=True)
+            ordered=True,
+        )
 
-      return i + 1, _get_next_indices(grid, loop_idx)
+      for j, var in enumerate(output_vars):
+        if _is_any(var.aval.memory_space):
+          continue
+        assert len(cur_start_indices[num_inputs + j].shape) == 1
+        assert len(next_start_indices[num_inputs + j].shape) == 1
+        jax.lax.cond(
+            (iteration_idx + 1 == num_iterations)
+            | jax.lax.reduce_or(
+                cur_start_indices[num_inputs + j]
+                != next_start_indices[num_inputs + j],
+                axes=(0,),
+            ),
+            functools.partial(_store_to_output_buffer, j, var),
+            lambda: None,
+        )
 
+      return iteration_idx + 1, next_loop_idx, cur_start_indices, next_start_indices
+
+  initial_loop_idx = (jnp.int32(0),) * len(grid)
+  with pallas_core.grid_env(_get_local_grid_env(initial_loop_idx)):
+    initial_start_indices = [
+        _compute_start_indices(
+            bm,
+            initial_loop_idx,
+            *scalar_buffer_ids,
+            compiler_params=compiler_params,
+            interpret_params=interpret_params,
+        )
+        for bm in grid_mapping.block_mappings
+    ]
   # TODO(jburnim): Handle parallel grid dimensions + megacore.
   _ = lax.while_loop(
       lambda carry: carry[0] < num_iterations,
       body,
-      (jnp.int32(0), (jnp.int32(0),) * len(grid))
+      (
+          jnp.int32(0),
+          initial_loop_idx,
+          initial_start_indices,  # Previous start indices are ignored on the first iteration.
+          initial_start_indices,
+      ),
   )
 
   # Read the output from the allocated output buffers.


Test patch:
diff --git a/tests/pallas/tpu_pallas_interpret_test.py b/tests/pallas/tpu_pallas_interpret_test.py
index 5b729f0fe..afb573f8c 100644
--- a/tests/pallas/tpu_pallas_interpret_test.py
+++ b/tests/pallas/tpu_pallas_interpret_test.py
@@ -18,24 +18,48 @@ To work around https://github.com/jax-ml/jax/issues/25671 , this file
 contains only tests that do not use shard_map.
 """
 
-from absl.testing import absltest
-from absl.testing import parameterized
 import functools
 
+from absl.testing import absltest
+from absl.testing import parameterized
 import jax
 from jax._src import test_util as jtu
 import jax._src.pallas.mosaic.interpret as mosaic_interpret
 from jax.experimental import pallas as pl
 from jax.experimental.pallas import tpu as pltpu
 import jax.numpy as jnp
-
 import numpy as np
 
 
 jax.config.parse_flags_with_absl()
 
 
+class CountStoreCallbacksContext(object):
+  """Wraps the I/O callback `store` into a callback that counts the number of calls to `store`."""
+
+  def __init__(self):
+    self._num_stores = 0
+    self._saved = mosaic_interpret.store
+
+  def __enter__(self):
+    def _store_callback(self, *args, **kwargs):
+      self._num_stores += 1
+      return self._saved(*args, **kwargs)
+
+    mosaic_interpret.store = functools.partial(_store_callback, self)
+    return self
+
+  def __exit__(self, ty, value, traceback):
+    del ty, value, traceback
+    mosaic_interpret.store = self._saved
+
+  @property
+  def num_stores(self):
+    return self._num_stores
+
+
 class InterpretTest(jtu.JaxTestCase):
+
   def setUp(self):
     super().setUp()
     self.num_devices = jax.device_count()
@@ -50,17 +74,18 @@ class InterpretTest(jtu.JaxTestCase):
     @jax.jit
     def matmul(x: jax.Array, y: jax.Array):
       return pl.pallas_call(
-        matmul_kernel,
-        out_shape=jax.ShapeDtypeStruct((x.shape[0], y.shape[1]), x.dtype),
-        grid=(2, 2),
-        in_specs=[
-            pl.BlockSpec((x.shape[0] // 2, x.shape[1]), lambda i, j: (i, 0)),
-            pl.BlockSpec((y.shape[0], y.shape[1] // 2), lambda i, j: (0, j))
-        ],
-        out_specs=pl.BlockSpec(
-            (x.shape[0] // 2, y.shape[1] // 2), lambda i, j: (i, j),
-        ),
-        interpret=mosaic_interpret.TPUInterpretParams(),
+          matmul_kernel,
+          out_shape=jax.ShapeDtypeStruct((x.shape[0], y.shape[1]), x.dtype),
+          grid=(2, 2),
+          in_specs=[
+              pl.BlockSpec((x.shape[0] // 2, x.shape[1]), lambda i, j: (i, 0)),
+              pl.BlockSpec((y.shape[0], y.shape[1] // 2), lambda i, j: (0, j)),
+          ],
+          out_specs=pl.BlockSpec(
+              (x.shape[0] // 2, y.shape[1] // 2),
+              lambda i, j: (i, j),
+          ),
+          interpret=mosaic_interpret.TPUInterpretParams(),
       )(x, y)
 
     k1, k2 = jax.random.split(jax.random.key(0))
@@ -79,9 +104,11 @@ class InterpretTest(jtu.JaxTestCase):
       grid_spec = pltpu.PrefetchScalarGridSpec(
           num_scalar_prefetch=1,
           grid=(1, 1),
-          in_specs=[pl.BlockSpec(
-              sizes,
-              lambda i, j, block_idx: (block_idx[0], block_idx[1]))],
+          in_specs=[
+              pl.BlockSpec(
+                  sizes, lambda i, j, block_idx: (block_idx[0], block_idx[1])
+              )
+          ],
           out_specs=pl.BlockSpec(sizes, lambda *_: (0, 0)),
       )
 
@@ -96,17 +123,21 @@ class InterpretTest(jtu.JaxTestCase):
 
     shape = (512, 512)
     x = jnp.reshape(jnp.arange(np.prod(shape), dtype=jnp.int32), shape)
-    result = block_dynamic_slice(x, starts=jnp.array([128, 256]), sizes=(128, 128))
-    ref = jax.lax.dynamic_slice(x, start_indices=(128, 256), slice_sizes=(128, 128))
+    result = block_dynamic_slice(
+        x, starts=jnp.array([128, 256]), sizes=(128, 128)
+    )
+    ref = jax.lax.dynamic_slice(
+        x, start_indices=(128, 256), slice_sizes=(128, 128)
+    )
     diff = jnp.max(jnp.abs(result - ref))
     np.testing.assert_allclose(result, ref)
 
   def test_dynamic_grid_and_aliasing(self):
-    self.skipTest('Broken pending fix to extra reads/writes of inputs/outputs')
     def kernel(s_ref, x_ref, o_ref):
       o_ref[...] = x_ref[...] + s_ref[0].astype(x_ref.dtype)
 
     iters = jax.random.randint(jax.random.key(0), (), 10, 20, dtype=jnp.int32)
+
     @jax.jit
     def f(s, x):
       return pl.pallas_call(
@@ -119,11 +150,11 @@ class InterpretTest(jtu.JaxTestCase):
           ],
           out_specs=pl.BlockSpec(x.shape, lambda i: (0, 0)),
           input_output_aliases={1: 0},
-          interpret=mosaic_interpret.TPUInterpretParams()
+          interpret=mosaic_interpret.TPUInterpretParams(),
       )(s, x)
 
     s = jnp.array([1], dtype=jnp.int32)
-    x = jnp.arange(32 * 128.).reshape((32, 128))
+    x = jnp.arange(32 * 128.0).reshape((32, 128))
     y = f(s, x)
     # NOTE: No matter how many times the kernel body is run, the kernel input
     # buffer will only be written once by the pallas_call machinery, just
@@ -136,6 +167,7 @@ class InterpretTest(jtu.JaxTestCase):
       @pl.when((pl.program_id(0) == 0) & (pl.program_id(1) == 0))
       def _():
         s_ref[0] = jnp.int32(0)
+
       s = s_ref[0]
       s_ref[0] = s + 1
       o_ref[:] = x_ref[:] + s.astype(x_ref.dtype)
@@ -149,7 +181,8 @@ class InterpretTest(jtu.JaxTestCase):
             pl.BlockSpec(block_shape=(8, 128), index_map=lambda i, j: (i, j)),
         ],
         out_specs=pl.BlockSpec(
-            block_shape=(8, 128), index_map=lambda i, j: (j, i)),
+            block_shape=(8, 128), index_map=lambda i, j: (j, i)
+        ),
         scratch_shapes=(pltpu.SMEM((1,), jnp.int32),),
         input_output_aliases={0: 0},
         interpret=mosaic_interpret.TPUInterpretParams(),
@@ -184,7 +217,8 @@ class InterpretTest(jtu.JaxTestCase):
       copy.wait()
 
     x = jnp.zeros((8, 128), jnp.float32)
-    y = pl.pallas_call(kernel_without_race,
+    y = pl.pallas_call(
+        kernel_without_race,
         out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
         in_specs=[pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY)],
         scratch_shapes=[
@@ -192,12 +226,14 @@ class InterpretTest(jtu.JaxTestCase):
             pltpu.SemaphoreType.DMA,
         ],
         interpret=mosaic_interpret.TPUInterpretParams(
-            detect_races=True, dma_execution_mode=dma_execution_mode),
+            detect_races=True, dma_execution_mode=dma_execution_mode
+        ),
     )(x).block_until_ready()
     self.assertFalse(mosaic_interpret.races.races_found)
     np.testing.assert_allclose(y, x + 1.0)
 
-    pl.pallas_call(kernel_with_race,
+    pl.pallas_call(
+        kernel_with_race,
         out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
         in_specs=[pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY)],
         scratch_shapes=[
@@ -205,7 +241,8 @@ class InterpretTest(jtu.JaxTestCase):
             pltpu.SemaphoreType.DMA,
         ],
         interpret=mosaic_interpret.TPUInterpretParams(
-            detect_races=True, dma_execution_mode=dma_execution_mode),
+            detect_races=True, dma_execution_mode=dma_execution_mode
+        ),
     )(x).block_until_ready()
     self.assertTrue(mosaic_interpret.races.races_found)
 
@@ -228,8 +265,8 @@ class InterpretTest(jtu.JaxTestCase):
     z = jax.jit(matmul)(x, y)
     np.testing.assert_array_equal(z, jnp.full_like(z, jnp.inf))
 
-    lowered = jax.jit(matmul).lower(x, y).as_text(dialect="stablehlo")
-    self.assertNotIn("dot_general", lowered)
+    lowered = jax.jit(matmul).lower(x, y).as_text(dialect='stablehlo')
+    self.assertNotIn('dot_general', lowered)
 
   @parameterized.parameters('nan', 'zero')
   def test_uninitialized_memory(self, uninitialized_memory):
@@ -250,7 +287,8 @@ class InterpretTest(jtu.JaxTestCase):
             pltpu.VMEM((8, 128), jnp.int16),
         ],
         interpret=mosaic_interpret.TPUInterpretParams(
-            uninitialized_memory=uninitialized_memory),
+            uninitialized_memory=uninitialized_memory
+        ),
     )()
     if uninitialized_memory == 'nan':
       self.assertTrue(jnp.isnan(x).all())
@@ -261,6 +299,33 @@ class InterpretTest(jtu.JaxTestCase):
       np.testing.assert_equal(np.array(y), 0)
       np.testing.assert_equal(np.array(z), 0)
 
+  def test_correct_number_of_stores(self):
+    def kernel(x_ref, s_ref, o_ref):
+      s = s_ref[0]
+      x_ref[:] += jax.lax.full_like(x_ref, s)
+      s_ref[0] = s + 1
+      o_ref[:] = x_ref[:]
+
+    def kernel_call(x, s):
+      return pl.pallas_call(
+          kernel,
+          out_shape=jax.ShapeDtypeStruct((16, 256), jnp.float32),
+          grid=(2, 2),
+          in_specs=[
+              pl.BlockSpec((8, 256), lambda i, j: (i, 0)),
+              pl.BlockSpec(memory_space=pltpu.SMEM),
+          ],
+          out_specs=pl.BlockSpec((8, 256), lambda i, j: (i, 0)),
+          interpret=mosaic_interpret.TPUInterpretParams(),
+      )(x, s)
+
+    with CountStoreCallbacksContext() as store_callbacks_counter:
+      result = jax.jit(kernel_call)(
+          jnp.zeros((16, 256), jnp.float32), jnp.zeros((1,), jnp.int32)
+      )
+      np.testing.assert_allclose(result[::8, ::256], [[1.0], [5.0]])
+      self.assertEqual(store_callbacks_counter.num_stores, 5)
+
 
-if __name__ == "__main__":
+if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())