Pallas pipeline API tweaks for more advanced pipelining patterns.

PiperOrigin-RevId: 678426679

Author: epiqueras

jax/_src/pallas/mosaic/BUILD

@@ -107,6 +107,7 @@ py_library(
         ":primitives",
         "//jax",
         "//jax:api_util",
+        "//jax:pallas",
         "//jax:util",
         "//jax/_src/pallas",
     ] + py_deps("numpy"),

jax/_src/pallas/mosaic/pipeline.py

@@ -178,6 +178,8 @@ class BufferType(enum.Enum):
   ACCUMULATOR = 3
   INPUT_OUTPUT = 4
 
+  MANUAL = 5
+
 
 @tree_util.register_pytree_node_class
 @dataclasses.dataclass(frozen=True)
@@ -234,6 +236,10 @@ def tree_flatten(self):
   def tree_unflatten(cls, meta, data):
     return cls(*meta, *data)
 
+  @staticmethod
+  def buffer_types() -> type[BufferType]:
+    return BufferType
+
   @classmethod
   def create(cls, spec, dtype, buffer_type) -> BufferedRef:
     """Create a BufferedRef.
@@ -1034,6 +1040,7 @@ def pipeline(
     prefetch=None,
     postyeet=None,
     schedule=None,
+    body_prologue=None,
   ):
     """
     Run the pipeline.
@@ -1056,6 +1063,8 @@ def pipeline(
         Called during the outputs phase in the first inner step.
       schedule: manually specified pipeline schedules for brefs, None indicates
         default schedule.
+      body_prologue: For running code within the grid environment before the
+        body is run. Useful for updating manual refs.
     """
     if scratches is None:
       scratches = ()
@@ -1119,6 +1128,9 @@ def loop_body(step, _):
                   lambda: None)
 
       # run the kernel!
+      if body_prologue is not None:
+        with scheduler.grid_env():
+          body_prologue()
       current_refs = map_brefs(lambda x: x.current_ref, brefs)
       with scheduler._named_scope("ep_run_kernel"):
         with scheduler.grid_env():