Repository: jax-ml/jax
Pull request #28628 (2025-05-08 21:37:40)

[Pallas/TPU]
[Pallas/TPU]
* Ensure we wrap scheduler initialize/finalize with grid_env
* Add support for basic swaps/gets to pull_block_spec



Patch:
diff --git a/jax/_src/pallas/fuser/block_spec.py b/jax/_src/pallas/fuser/block_spec.py
index ba2c18201..3d4df5499 100644
--- a/jax/_src/pallas/fuser/block_spec.py
+++ b/jax/_src/pallas/fuser/block_spec.py
@@ -29,6 +29,7 @@ from jax._src import ad_util
 from jax._src import core
 from jax._src import custom_derivatives
 from jax._src import pjit
+from jax._src import state
 from jax._src import tree_util
 from jax._src import util
 from jax._src.interpreters import partial_eval as pe
@@ -451,6 +452,8 @@ def make_kernel_function(
   _no_aval = object()
 
   def _get_block_aval(bs, aval):
+    if isinstance(aval, state.AbstractRef):
+      return aval
     if bs is pallas_core.no_block_spec or bs is None:
       return _no_aval
     return aval.update(shape=_remove_nones(bs.block_shape))  # pytype: disable=attribute-error
@@ -1065,6 +1068,171 @@ def _dynamic_slice_rule(
       len(ctx.avals_in) - 1
   )
 
+@register_pull_block_spec_rule(state_primitives.swap_p)
+def _swap_pull_rule(
+    ctx: PullRuleContext,
+    block_spec: pallas_core.BlockSpec,
+    **kwargs,
+):
+  del ctx, kwargs
+  # The output and val block spec are the same.
+  return [block_spec, block_spec]
+
+@register_eval_rule(state_primitives.swap_p)
+def _swap_eval_rule(
+    ctx: KernelEvalContext,
+    ref,
+    val,
+    *idx,
+    tree
+):
+  indexers = tree_util.tree_unflatten(tree, idx)
+  ref_aval, _ = ctx.avals_in[:2]
+  indexers_avals = tree_util.tree_unflatten(tree, ctx.avals_in[2:])
+  assert hasattr(ref_aval, 'shape')
+  if len(indexers) > 1:
+    raise NotImplementedError('swap not supported yet')
+  indexer_aval = indexers_avals[0]
+  for idx_aval, size in zip(indexer_aval.indices, ref_aval.shape, strict=True):
+    if not isinstance(idx_aval, indexing.Slice):
+      raise NotImplementedError('swap not supported yet')
+    if not isinstance(idx_aval.start, int):
+      raise NotImplementedError('swap not supported yet')
+    if not isinstance(idx_aval.size, int):
+      raise NotImplementedError('swap not supported yet')
+    if idx_aval.stride != 1:
+      raise NotImplementedError('swap not supported yet')
+    if idx_aval.start != 0:
+      raise NotImplementedError('swap not supported yet')
+    if idx_aval.size != size:
+      raise NotImplementedError('swap not supported yet')
+  # We have a pure slice so now we can just re-index the ref according to the
+  # block indices.
+  block_spec = ctx.out_block_specs[0]
+  block_idx = ctx.get_out_block_indices()[0]
+
+  def _slice(i, b):
+    if not isinstance(b, int):
+      raise NotImplementedError('swap not supported yet')
+    return i if b is None else indexing.ds(i * b, b)
+
+  indexer = tuple(
+      _slice(i, b) for i, b in zip(block_idx, block_spec.block_shape,
+                                   strict=True)
+  )
+  return ref.swap(val, idx=indexer)
+
+@register_pull_block_spec_rule(state_primitives.get_p)
+def _get_pull_rule(
+    ctx: PullRuleContext,
+    block_spec: pallas_core.BlockSpec,
+    *,
+    tree
+):
+  ref_aval = ctx.avals_in[0]
+  assert hasattr(ref_aval, 'shape')
+  indexers_avals = tree_util.tree_unflatten(tree, ctx.avals_in[1:])
+  if len(indexers_avals) > 1:
+    raise NotImplementedError('get not supported yet')
+  indexer_aval = indexers_avals[0]
+  block_shape_iter = iter(block_spec.block_shape)
+  block_shape = []
+  if not all(
+      isinstance(bd, (int, pallas_core.Blocked, pallas_core.Squeezed, None))
+      for bd in block_spec.block_shape
+  ):
+    raise NotImplementedError('get not supported yet')
+  for idx_aval, size in zip(indexer_aval.indices, ref_aval.shape, strict=True):
+    if not isinstance(idx_aval, indexing.Slice):
+      assert hasattr(idx_aval, 'shape') and not idx_aval.shape
+      block_shape.append(pallas_core.Squeezed())
+      continue
+    if not isinstance(idx_aval.start, int):
+      raise NotImplementedError('get not supported yet')
+    if not isinstance(idx_aval.size, int):
+      raise NotImplementedError('get not supported yet')
+    if idx_aval.stride != 1:
+      raise NotImplementedError('get not supported yet')
+    if idx_aval.start != 0:
+      raise NotImplementedError('get not supported yet')
+    if idx_aval.size != size:
+      raise NotImplementedError('get not supported yet')
+    bd = next(block_shape_iter)
+    block_shape.append(_block_size(bd))
+  assert next(block_shape_iter, None) is None
+  def new_index_map(*args):
+    idx = block_spec.index_map(*args)
+    idx_iter = iter(idx)
+    indices = tuple(
+        0
+        if (bd is None or isinstance(bd, pallas_core.Squeezed))
+        else next(idx_iter)
+        for bd in range(len(block_shape))
+    )
+    assert next(idx_iter, None) is None
+    return indices
+  block_spec = pallas_core.BlockSpec(block_shape, new_index_map)
+  return [block_spec] + [pallas_core.no_block_spec] * (len(ctx.avals_in) - 1)
+
+@register_eval_rule(state_primitives.get_p)
+def _get_eval_rule(
+    ctx: KernelEvalContext,
+    ref,
+    *idx,
+    tree
+):
+  indexers = tree_util.tree_unflatten(tree, idx)
+  ref_aval = ctx.avals_in[0]
+  indexers_avals = tree_util.tree_unflatten(tree, ctx.avals_in[1:])
+  ref_block_spec = ctx.in_block_specs[0]
+  assert hasattr(ref_aval, 'shape')
+  if len(indexers) > 1:
+    raise NotImplementedError('get not supported yet')
+  indexer = indexers[0]
+  indexer_aval = indexers_avals[0]
+  block_indexer = []
+
+  def _slice(i, b):
+    match b:
+      case int():
+        return indexing.ds(i * b, b)
+      case pallas_core.Blocked(bs):
+        return indexing.ds(i * bs, bs)
+      case pallas_core.Squeezed() | None:
+        return i
+      case _:
+        raise NotImplementedError('get not supported yet')
+
+  if ref_block_spec is pallas_core.no_block_spec:
+    # Short-circuit if the ref is not blocked.
+    return state_primitives.get_p.bind(ref, *idx, tree=tree)
+  block_idx_iter = iter(ctx.get_out_block_indices()[0])
+  for idx_aval, size, idx, bd in zip(
+      indexer_aval.indices,
+      ref_aval.shape,
+      indexer.indices,
+      ref_block_spec.block_shape,
+      strict=True,
+  ):
+    if not isinstance(idx_aval, indexing.Slice):
+      assert hasattr(idx_aval, 'shape') and not idx_aval.shape, idx_aval
+      assert bd is None or isinstance(bd, pallas_core.Squeezed)
+      block_indexer.append(idx)
+      continue
+    if not isinstance(idx_aval.start, int):
+      raise NotImplementedError('get not supported yet')
+    if not isinstance(idx_aval.size, int):
+      raise NotImplementedError('get not supported yet')
+    if idx_aval.stride != 1:
+      raise NotImplementedError('get not supported yet')
+    if idx_aval.start != 0:
+      raise NotImplementedError('get not supported yet')
+    if idx_aval.size != size:
+      raise NotImplementedError('get not supported yet')
+    bidx = next(block_idx_iter)
+    block_indexer.append(_slice(bidx, bd))
+  assert next(block_idx_iter, None) is None
+  return ref.get(idx=tuple(block_indexer))
 
 @register_eval_rule(lax.concatenate_p)
 def _concatenate_eval_rule(ctx: KernelEvalContext, *args, dimension):
diff --git a/jax/_src/pallas/mosaic/pipeline.py b/jax/_src/pallas/mosaic/pipeline.py
index df7be297c..4ef221792 100644
--- a/jax/_src/pallas/mosaic/pipeline.py
+++ b/jax/_src/pallas/mosaic/pipeline.py
@@ -1369,7 +1369,8 @@ def emit_pipeline(
       initial_indices = (0,) * len(grid)
       scheduler = make_scheduler(0, initial_indices)
       brefs = map_brefs(scheduler.alias_local_refs, allocations, refs)
-      map_brefs(scheduler.initialize, brefs, refs, schedule)
+      with scheduler.grid_env():
+        map_brefs(scheduler.initialize, brefs, refs, schedule)
 
       # pipeline loop
       next_indices = lax.fori_loop(0, num_steps, loop_body, initial_indices)
@@ -1378,7 +1379,8 @@ def emit_pipeline(
       final_indices = _prev_index(next_indices, grid)
       scheduler = make_scheduler(num_steps - 1, final_indices)
       brefs = map_brefs(scheduler.alias_local_refs, allocations, refs)
-      map_brefs(scheduler.finalize, brefs, refs, schedule)
+      with scheduler.grid_env():
+        map_brefs(scheduler.finalize, brefs, refs, schedule)
 
   return pipeline
 


Test patch:
diff --git a/tests/pallas/fuser_block_spec_test.py b/tests/pallas/fuser_block_spec_test.py
index 665cdfb1d..f7e70ec1d 100644
--- a/tests/pallas/fuser_block_spec_test.py
+++ b/tests/pallas/fuser_block_spec_test.py
@@ -761,6 +761,101 @@ class PullBlockSpecTest(jtu.JaxTestCase):
     y = kernel_fn((0, 1, 2), scalar_prefetch_values, (), x)
     np.testing.assert_array_equal(y, x.reshape((256, 1024)))
 
+  def test_basic_swap(self):
+    value = jnp.arange((512 * 1024), dtype=jnp.int32).reshape((512, 1024)) * 2
+    x = jnp.zeros((256, 512), dtype=jnp.int32)
+    def outer(refs):
+      ref, y_ref = refs
+      def f(x):
+        return ref.swap(x)
+      in_type = jax.ShapeDtypeStruct((512, 1024), jnp.int32)
+      f2, new_values, scalar_prefetch_values = block_spec_lib.get_fusion_values(
+          f, in_type
+      )
+      self.assertLen(new_values, 1)  # Captures Ref
+      self.assertEmpty(scalar_prefetch_values)
+
+      block_spec = pl.BlockSpec((256, 512), lambda i, j, k: (i, k))
+      kernel_fn, (value_block_specs, x_block_spec), _ = (
+          block_spec_lib.pull_block_spec(
+              f2,
+              block_spec,
+              grid=(2, 3, 4),
+              scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(),
+          )(new_values, in_type)
+      )
+      self.assertLen(value_block_specs, 1)
+      self.assertEqual(x_block_spec.index_map(0, 1, 2), (0, 2))
+      self.assertEqual(x_block_spec.index_map(3, 2, 1), (3, 1))
+
+      y_ref[...] = kernel_fn((0, 1, 1), scalar_prefetch_values, (ref,), x)
+    y = jnp.zeros((256, 512), jnp.int32)
+    _, y = pl.run_state(outer)((value, y))
+    np.testing.assert_array_equal(y, value[:256, 512:1024])
+
+  def test_basic_get(self):
+    value = jnp.arange((512 * 1024), dtype=jnp.int32).reshape((512, 1024)) * 2
+    def outer(refs):
+      ref, y_ref = refs
+      def f():
+        return ref.get()
+
+      block_spec = pl.BlockSpec((256, 512), lambda i, j, k: (i, k))
+      kernel_fn, (), _ = (
+          block_spec_lib.pull_block_spec(
+              f,
+              block_spec,
+              grid=(2, 3, 4),
+              scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(),
+          )()
+      )
+      y_ref[...] = kernel_fn((0, 1, 1), ())
+    y = jnp.zeros((256, 512), jnp.int32)
+    _, y = pl.run_state(outer)((value, y))
+    np.testing.assert_array_equal(y, value[:256, 512:1024])
+
+  def test_get_with_squeezed_block_spec(self):
+    value = jnp.arange((4 * 512 * 1024), dtype=jnp.int32).reshape((4, 512, 1024)) * 2
+    def outer(refs):
+      ref, y_ref = refs
+      def f():
+        return ref.get()
+
+      block_spec = pl.BlockSpec((pl.Squeezed(), 256, 512), lambda i, j, k: (j, i, k))
+      kernel_fn, (), _ = (
+          block_spec_lib.pull_block_spec(
+              f,
+              block_spec,
+              grid=(2, 3, 4),
+              scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(),
+          )()
+      )
+      y_ref[...] = kernel_fn((0, 3, 1), ())
+    y = jnp.zeros((256, 512), jnp.int32)
+    _, y = pl.run_state(outer)((value, y))
+    np.testing.assert_array_equal(y, value[3, :256, 512:1024])
+
+  def test_get_with_squeezed_indexer(self):
+    value = jnp.arange((4 * 512 * 1024), dtype=jnp.int32).reshape((4, 512, 1024)) * 2
+    def outer(refs):
+      ref, y_ref = refs
+      def f():
+        return ref[3]
+
+      block_spec = pl.BlockSpec((256, 512), lambda i, j, k: (i, k))
+      kernel_fn, (), _ = (
+          block_spec_lib.pull_block_spec(
+              f,
+              block_spec,
+              grid=(2, 3, 4),
+              scalar_prefetch_handler=block_spec_lib.make_scalar_prefetch_handler(),
+          )()
+      )
+      y_ref[...] = kernel_fn((0, 2, 1), ())
+    y = jnp.zeros((256, 512), jnp.int32)
+    _, y = pl.run_state(outer)((value, y))
+    np.testing.assert_array_equal(y, value[3, :256, 512:1024])
+
 
 class PullBlockSpecHOPTest(jtu.JaxTestCase):