[pallas:mgpu] Support per-input delay_release in emit_pipeline

This change allows specifying different `delay_release` values for each input in
`emit_pipeline` (`emit_pipeline_warp_specialized` left for a followup cl).

This is achieved by stratifying the inputs based on their `delay_release` values
and adjusting the pipeline logic to handle the different effective number of
stages for each stratum.

PiperOrigin-RevId: 801093645

Author: cperivol

jax/_src/pallas/mosaic_gpu/pipeline.py

@@ -69,7 +69,7 @@ def _get_block_shape(spec: pallas_core.BlockSpec):
 @jax.tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True)
 class BufferedRef:
-  spec: pallas_core.BlockSpec = dataclasses.field(metadata={"static": True})
+  spec: gpu_core.BlockSpec = dataclasses.field(metadata={"static": True})
   is_index_invariant: bool = dataclasses.field(metadata={"static": True})
   gmem_ref: state.AbstractRef
   # ``None`` if the ref is pinned to GMEM; otherwise, has shape
@@ -184,6 +184,20 @@ def __eq__(self, other: _Slice) -> jax.Array:  # type: ignore
 )
 
 
+def _downcast_spec(
+    spec: gpu_core.BlockSpec | pallas_core.BlockSpec,
+) -> gpu_core.BlockSpec:
+  if isinstance(spec, gpu_core.BlockSpec):
+    return spec
+
+  return gpu_core.BlockSpec(
+      block_shape=spec.block_shape,
+      index_map=spec.index_map,
+      memory_space=spec.memory_space,
+      pipeline_mode=spec.pipeline_mode,
+  )
+
+
 def emit_pipeline(
     body: Callable[..., T],
     *,
@@ -220,25 +234,16 @@ def emit_pipeline(
     pipeline and returns the final carry value (if ``init_carry`` was used),
     otherwise it returns None.
   """
+
+  in_specs = tuple(map(_downcast_spec, in_specs))
+  out_specs = tuple(map(_downcast_spec, out_specs))
   # TODO(justinfu): Factor out common code between warp-specialized and
   # normal pipelines.
-  delay_release = None
-  for in_spec in in_specs:
-    if not isinstance(in_spec, gpu_core.BlockSpec):
-      delay_release = 0
-      continue
-    delay_release = in_spec.delay_release
-    if in_spec.delay_release != delay_release:
-      raise NotImplementedError(
-          "All inputs must have the same delay_release, but"
-          f" {in_spec.delay_release=} != {delay_release=}"
-      )
-
-  delay_release = delay_release or 0
-  if max_concurrent_steps <= delay_release:
+  delay_release_levels = sorted({s.delay_release for s in in_specs}) or [0]
+  if delay_release_levels and max_concurrent_steps <= delay_release_levels[0]:
     raise ValueError(
-        "max_concurrent_steps must be greater than delay_release, but"
-        f" {max_concurrent_steps=}, {delay_release=}"
+        "max_concurrent_steps must be greater than all delay_release values,"
+        f" but {max_concurrent_steps=} and {delay_release_levels=}."
     )
 
   num_steps = math.prod(grid)
@@ -316,7 +321,7 @@ def scoped_pipeline(
 
     def loop_body(step, carry):
       slot = lax.rem(step, max_concurrent_steps)
-      indices, fetch_indices, last_store_slices, prev_body_carry = carry
+      indices, fetch_index_levels, last_store_slices, prev_body_carry = carry
 
       if barrier_ref is not None:
         # Wait for the current GMEM->SMEM copy to complete, if any.
@@ -368,31 +373,43 @@ def loop_body(step, carry):
       if copies_out_in_loop:
         gpu_primitives.commit_smem_to_gmem_group()
 
-      fetch_step = step + (max_concurrent_steps - delay_release)
-      fetch_slot = lax.rem(fetch_step, max_concurrent_steps)
-
-      def do_fetch():
-        for bref in in_brefs:
-          bref.copy_in(fetch_slot, fetch_indices, barrier_ref)
-
-      jax.lax.cond(
-          lax.bitwise_and(step >= delay_release, fetch_step < num_steps),
-          do_fetch,
-          lambda: None,
-      )
+      for delay_release, fetch_indices in zip(
+          delay_release_levels, fetch_index_levels
+      ):
+        fetch_step = step + (max_concurrent_steps - delay_release)
+        fetch_slot = lax.rem(fetch_step, max_concurrent_steps)
+
+        # pylint: disable=cell-var-from-loop
+        def do_fetch():
+          for bref in in_brefs:
+            if bref.spec.delay_release == delay_release:
+              bref.copy_in(fetch_slot, fetch_indices, barrier_ref)
+        # pylint: enable=cell-var-from-loop
+
+        jax.lax.cond(
+            lax.bitwise_and(step >= delay_release, fetch_step < num_steps),
+            do_fetch,
+            lambda: None,
+        )
 
+      next_fetch_indices_levels = [
+          _inc_grid_by_1(fetch_indices, grid)
+          for fetch_indices in fetch_index_levels
+      ]
       return (
           _inc_grid_by_1(indices, grid),
-          _inc_grid_by_1(fetch_indices, grid),
+          next_fetch_indices_levels,
           new_store_slices,
           next_body_carry if init_carry is not None else None,
       )
 
-    # Invariant: ``indices`` and ``fetch_indices`` are always
-    # ``max_concurrent_steps-delay_release`` apart.
-    fetch_indices = indices
-    for _ in range(max_concurrent_steps-delay_release):
-      fetch_indices = _inc_grid_by_1(fetch_indices, grid)
+    fetch_index_levels = []
+    for delay_release in delay_release_levels:
+      fetch_indices = indices
+      for _ in range(max_concurrent_steps - delay_release):
+        fetch_indices = _inc_grid_by_1(fetch_indices, grid)
+      fetch_index_levels.append(fetch_indices)
+
     # TODO(justinfu): Only store base pointer instead of all indices.
     last_store_slices = [
         None
@@ -404,7 +421,7 @@ def do_fetch():
         0,
         num_steps,
         loop_body,
-        (indices, fetch_indices, last_store_slices, init_carry),
+        (indices, fetch_index_levels, last_store_slices, init_carry),
     )
 
     # Outputs invariant to the sequential axis are never written from inside the

tests/pallas/mosaic_gpu_test.py

@@ -4711,6 +4711,47 @@ def pipeline(*gmem_refs):
     y = jax.random.uniform(jax.random.key(1), (m, n), dtype=jnp.float32)
     np.testing.assert_allclose(kernel(x, y), x + y, atol=1e-4)
 
+  def test_different_delay_release(self):
+    self.skip_if_wg_semantics()  # Crashes!
+    m, n = 128, 64
+    blk_m, blk_n = 32, 64
+    in_specs = [
+        plgpu.BlockSpec(
+            block_shape=(blk_m, blk_n),
+            index_map=lambda i, j: (i, j),
+            delay_release=delay,
+        )
+        for delay in range(3)
+    ]
+    out_spec = pl.BlockSpec(
+        block_shape=(blk_m, blk_n),
+        index_map=lambda i, j: (i, j),
+    )
+
+    def tiled_add_kernel(_, x_smem, y_smem, z_smem, o_smem):
+      o_smem[...] = x_smem[...] + y_smem[...] + z_smem[...]
+
+    def pipeline(*gmem_refs):
+      grid = (m // blk_m, n // blk_n)
+      return mgpu_pipeline.emit_pipeline(
+          tiled_add_kernel,
+          grid=grid,
+          max_concurrent_steps=4,
+          in_specs=in_specs,
+          out_specs=[out_spec],
+      )(*gmem_refs)
+
+    kernel = self.kernel(
+        pipeline,
+        out_shape=jax.ShapeDtypeStruct((m, n), jnp.float32),
+        grid=(1,),
+        grid_names=("_",)
+    )
+    x = jax.random.uniform(jax.random.key(0), (m, n), dtype=jnp.float32)
+    y = jax.random.uniform(jax.random.key(1), (m, n), dtype=jnp.float32)
+    z = jax.random.uniform(jax.random.key(3), (m, n), dtype=jnp.float32)
+    np.testing.assert_allclose(kernel(x, y, z), x + y + z)
+
   @parameterized.product(
       delay_release=[0, 1],
   )