[pallas:mosaic_gpu] ModuleContext.reserve_barrier is now a context manager

This allows unreserving the barrier once it is no longer needed and is consistent
with how resource estimation works, e.g. for `cond`.

PiperOrigin-RevId: 745483567

Author: superbobry

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -311,6 +311,7 @@ class ModuleContext:
   lowering_semantics: mgpu.LoweringSemantics
   primitive_semantics: gpu_core.PrimitiveSemantics
 
+  @contextlib.contextmanager
   def reserve_barrier(self, barrier: mgpu.Barrier) -> mgpu.BarrierRef:
     """Reserves a barrier.
 
@@ -320,7 +321,9 @@ def reserve_barrier(self, barrier: mgpu.Barrier) -> mgpu.BarrierRef:
     available = self.runtime_barriers.get(barrier, [])
     if not available:
       raise RuntimeError(f"Barrier {barrier} is already reserved")
-    return available.pop()
+    barrier = available.pop()
+    yield barrier
+    available.append(barrier)
 
   @contextlib.contextmanager
   def alloc_tmem(
@@ -1965,7 +1968,7 @@ def _run_scoped_lowering_rule(
           input_refs.append(acc)
         should_discharge.append(True)
       elif isinstance(aval.dtype, gpu_core.BarrierType):
-        input_refs.append(
+        barrier_ref = alloc_stack.enter_context(
             ctx.module_ctx.reserve_barrier(
                 mgpu.Barrier(
                     aval.dtype.num_arrivals
@@ -1974,17 +1977,19 @@ def _run_scoped_lowering_rule(
                 )
             )
         )
+        input_refs.append(barrier_ref)
         should_discharge.append(False)
       elif isinstance(aval.dtype, gpu_core.ClusterBarrierType):
         collective_dims = jax.tree.map(
             lambda axis: _resolve_cluster_axis(ctx.module_ctx.axis_names, axis),
             aval.dtype.collective_axes,
         )
-        input_refs.append(
+        barrier_ref = alloc_stack.enter_context(
             ctx.module_ctx.reserve_barrier(
                 mgpu.ClusterBarrier(collective_dims, *aval.shape)
             )
         )
+        input_refs.append(barrier_ref)
         should_discharge.append(False)
       elif aval.memory_space == gpu_core.SMEM:
         [input_ref] = alloc_stack.enter_context(

tests/pallas/mosaic_gpu_test.py

@@ -1046,6 +1046,26 @@ def body(tmp_ref):
     x = np.ones((8, 128), jnp.float32)
     np.testing.assert_array_equal(kernel(x), x + 1.0)
 
+  def test_run_scoped_in_cond(self):
+    @functools.partial(
+        self.pallas_call,
+        out_shape=jax.ShapeDtypeStruct([256], jnp.int32),
+        in_specs=[pl.BlockSpec(memory_space=plgpu.GPUMemorySpace.GMEM)],
+        out_specs=pl.BlockSpec(memory_space=plgpu.GPUMemorySpace.SMEM),
+    )
+    def kernel(x_ref_gmem, o_ref):
+      def scoped_kernel(barrier_ref):
+        plgpu.copy_gmem_to_smem(x_ref_gmem, o_ref, barrier_ref)
+        plgpu.barrier_wait(barrier_ref)
+
+      def branch():
+        pl.run_scoped(scoped_kernel, plgpu.Barrier(num_arrivals=1))
+
+      jax.lax.cond(x_ref_gmem[0] % 2 == 0, branch, branch)
+
+    x = jnp.full((256,), 1234, dtype=jnp.int32)
+    np.testing.assert_array_equal(kernel(x), x)
+
   def test_program_id(self):
     @functools.partial(
         self.pallas_call,