[Pallas SC] Allow semaphores to be returned by SCS kernels

PiperOrigin-RevId: 841972235

Author: sharadmv

jax/_src/pallas/mosaic/pallas_call_registration.py

@@ -62,7 +62,7 @@ def _maybe_cast_to_int(x: jax.Array | jax_core.AbstractValue):
 
 
 def _get_memory_space_from_aval(
-    out_aval: jax_core.AbstractValue,
+    out_aval: jax_core.AbstractValue, kernel_type: tpu_core.KernelType
 ) -> tpu_custom_call.MemorySpace | None:
   if not isinstance(out_aval, jax_core.ShapedArray):
     raise ValueError("Memory spaces not defined for non-ShapedArrays")
@@ -84,20 +84,29 @@ def _get_memory_space_from_aval(
     case tpu_core.MemorySpace.SMEM:
       return tpu_custom_call.MemorySpace.SMEM
     case tpu_core.MemorySpace.SEMAPHORE:
-      return tpu_custom_call.MemorySpace.SEMAPHORE_MEM
+      match kernel_type:
+        case tpu_core.KernelType.SC_SCALAR_SUBCORE:
+          return tpu_custom_call.MemorySpace.SC_SCALAR_SEMAPHORE_MEM
+        case tpu_core.KernelType.TC:
+          return tpu_custom_call.MemorySpace.SEMAPHORE_MEM
+        case _:
+          raise ValueError(f"Invalid kernel type for semaphore: {kernel_type}")
     case tpu_core.MemorySpace.HOST:
       return tpu_custom_call.MemorySpace.HOST
   return None
 
 
 def _get_memory_spaces_from_avals(
-    avals: Sequence[jax_core.AbstractValue],
+    avals: Sequence[jax_core.AbstractValue], kernel_type: tpu_core.KernelType
 ) -> tuple[tpu_custom_call.MemorySpace | None, ...] | None:
   memory_spaces = None
   if any(
       isinstance(aval, pallas_core.ShapedArrayWithMemorySpace) for aval in avals
   ):
-    memory_spaces = tuple(map(_get_memory_space_from_aval, avals))
+    memory_spaces = tuple(
+        _get_memory_space_from_aval(aval, kernel_type=kernel_type)
+        for aval in avals
+    )
   return memory_spaces
 
 
@@ -140,7 +149,7 @@ def pallas_call_tpu_lowering_rule(
   mlir_ctx.load_all_available_dialects()
   tpu.register_dialect(mlir_ctx)
 
-  match mosaic_params.kernel_type:
+  match (kernel_type := mosaic_params.kernel_type):
     case tpu_core.KernelType.TC:
       lower_jaxpr_to_module = lowering.lower_jaxpr_to_module
     case tpu_core.KernelType.SC_SCALAR_SUBCORE | tpu_core.KernelType.SC_VECTOR_SUBCORE:
@@ -191,7 +200,9 @@ def _maybe_cast_inputs(*args):
   # Dynamic grid bounds have to go at the front.
   dynamic_grid_args, args = in_nodes[:num_dyn_bounds], in_nodes[num_dyn_bounds:]
   kernel_ctx = ctx.replace(avals_in=kernel_in_avals, avals_out=kernel_out_avals)
-  output_memory_spaces = _get_memory_spaces_from_avals(out_avals)
+  output_memory_spaces = _get_memory_spaces_from_avals(
+      out_avals, kernel_type=kernel_type
+  )
   input_memory_spaces = None
   if any(
       isinstance(aval, pallas_core.ShapedArrayWithMemorySpace)
@@ -202,7 +213,9 @@ def _maybe_cast_inputs(*args):
       raise NotImplementedError(
           "Dynamic grid bounds are not supported when specifying memory spaces for inputs."
       )
-    input_memory_spaces = _get_memory_spaces_from_avals(ctx.avals_in)
+    input_memory_spaces = _get_memory_spaces_from_avals(
+        ctx.avals_in, kernel_type=kernel_type
+    )
   if cost_estimate is not None:
     mosaic_cost_estimate = cast(
         tpu_custom_call.CostEstimate, dataclasses.asdict(cost_estimate)

jax/_src/pallas/mosaic/sc_lowering.py

@@ -330,7 +330,10 @@ def body_func(*args: ir.Value):
       mosaic_grid_mapping.block_mappings,
   ):
     d = {}
-    if str(arg.type.memory_space) == "#tpu.memory_space<hbm>":
+    if (
+        str(arg.type.memory_space) == "#tpu.memory_space<hbm>"
+        or str(arg.type.memory_space) == "#tpu.memory_space<semaphore_mem>"
+    ):
       d["sc.persistent"] = ir.UnitAttr.get()
     if isinstance(bm, sc_core.BlockMapping) and bm.indexed_by is not None:
       d["sc.indexed_by"] = mlir.i32_attr(bm.indexed_by)

jax/_src/tpu_custom_call.py

@@ -101,6 +101,7 @@ class MemorySpace(enum.Enum):
   SEMAPHORE_MEM = enum.auto()
   SMEM = enum.auto()
   HOST = enum.auto()
+  SC_SCALAR_SEMAPHORE_MEM = enum.auto()
 
   @property
   def color(self) -> int:
@@ -110,6 +111,8 @@ def color(self) -> int:
       return 1
     elif self == MemorySpace.SEMAPHORE_MEM:
       return 2
+    elif self == MemorySpace.SC_SCALAR_SEMAPHORE_MEM:
+      return 8
     elif self == MemorySpace.SMEM:
       return 4
     elif self == MemorySpace.HOST:

tests/pallas/tpu_sparsecore_pallas_test.py

@@ -1923,5 +1923,67 @@ class PipelineTestWithTCTiling(TCTilingMixin, PipelineTest):
   pass
 
 
+class PallasSparsecoreAsyncTest(PallasSCTest):
+
+  @parameterized.product(
+      shape=[
+          (8, 128),
+          (8, 256),
+          (8, 512),
+          (8, 1024),
+          (16, 128),
+          (16, 256),
+          (16, 512),
+          (16, 1024),
+          # TODO(sharadmv): These shapes fail right now.
+          # (64, 8),
+      ],
+      dtype=[jnp.int32, jnp.float32, jnp.bfloat16],
+  )
+  def test_basic_async_kernel(self, shape, dtype):
+    if not jtu.is_cloud_tpu_at_least(2025, 12, 8):
+      self.skipTest("Need newer libtpu")
+    x = jnp.arange(shape[0] * shape[1], dtype=dtype).reshape(shape)
+
+    @jax.jit
+    def foo(x):
+      sc_mesh = plsc.ScalarSubcoreMesh(axis_name="core", num_cores=1)
+
+      sem = pl.pallas_call(
+          lambda _: None,
+          out_shape=pltpu.SemaphoreType.DMA(()),
+          out_specs=pl.BlockSpec(memory_space=pltpu.SEMAPHORE),
+          compiler_params=pltpu.CompilerParams(
+              dimension_semantics=["core_parallel"],
+              kernel_type=pltpu.KernelType.SC_SCALAR_SUBCORE,
+          ),
+      )()
+
+      sem_ref = jax.new_ref(sem, memory_space=pltpu.SEMAPHORE)
+      y_ref = pl.empty_ref_like(pltpu.HBM(x.shape, x.dtype))
+      x_ref = jax.new_ref(x)
+
+      run_kernel = pl.core_map(mesh=sc_mesh)
+
+      @run_kernel
+      def _():
+        pltpu.make_async_copy(x_ref, y_ref, sem_ref).start()
+
+      @run_kernel
+      def _():
+        pltpu.make_async_copy(x_ref, y_ref, sem_ref).wait()
+
+      return y_ref[...]
+
+    o = jax.block_until_ready(foo(x))
+    np.testing.assert_array_equal(o, x)
+
+
+class PallasSparsecoreAsyncTestWithTCTiling(
+    TCTilingMixin, PallasSparsecoreAsyncTest
+):
+  pass
+
+
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())