[Pallas][NFC] Move the remainder of Semaphore-related extended dtypes to Pallas core

This completes the move started in jax-ml#26673.

PiperOrigin-RevId: 741487331

Author: apaszke

jax/_src/pallas/core.py

@@ -67,9 +67,55 @@ def __repr__(self):
 SEMAPHORE_INTERPRET_DTYPE = jnp.int16
 SEMAPHORE_MAX_VALUE = jnp.iinfo(SEMAPHORE_INTERPRET_DTYPE).max
 
-class semaphore_dtype(dtypes.extended): pass
-class semaphore(semaphore_dtype): pass
-class barrier_semaphore(semaphore_dtype): pass
+class AbstractSemaphoreTyRules:
+  @staticmethod
+  def pallas_interpret_element_aval(_) -> jax_core.ShapedArray:
+    return jax_core.ShapedArray((), SEMAPHORE_INTERPRET_DTYPE)
+
+  @staticmethod
+  def physical_element_aval(_) -> jax_core.ShapedArray:
+    return jax_core.ShapedArray((), jnp.int32)
+
+# TODO(sharadmv): implement dtype rules for AbstractSemaphoreTy
+class AbstractSemaphoreTy(dtypes.ExtendedDType):
+  name: str
+  _rules = AbstractSemaphoreTyRules
+
+  def __repr__(self) -> str:
+    return self.name
+
+  def __eq__(self, other):
+    return self.__class__ == other.__class__
+
+  def __hash__(self) -> int:
+    return hash(self.__class__)
+
+class semaphore_dtype(dtypes.extended):
+  """Common dtype for all kinds of semaphore dtypes.
+
+  This is an abstract class that should never be instantiated, but rather
+  exists for the sake of `jnp.issubdtype`.
+  """
+
+class semaphore(semaphore_dtype):
+  """Regular semaphore dtype.
+
+  Like its superclass, this class should never be instantiated.
+  """
+
+class Semaphore(AbstractSemaphoreTy):
+  name = "semaphore"
+  type = semaphore
+
+class barrier_semaphore(semaphore_dtype):
+  """Barrier semaphore dtype.
+
+  Like its superclass, this class should never be instantiated.
+  """
+
+class BarrierSemaphore(AbstractSemaphoreTy):
+  name = "barrier_semaphore"
+  type = barrier_semaphore
 
 @runtime_checkable
 class CompilerParams(Protocol):

jax/_src/pallas/mosaic/core.py

@@ -25,7 +25,6 @@
 import jax
 from jax._src import config
 from jax._src import core as jax_core
-from jax._src import dtypes
 from jax._src import util
 from jax._src.pallas import core as pallas_core
 import jax.numpy as jnp
@@ -114,42 +113,10 @@ def __call__(self, shape: tuple[int, ...], dtype: jnp.dtype):
 
 class dma_semaphore(pallas_core.semaphore_dtype): pass
 
-class AbstractSemaphoreTyRules:
-  @staticmethod
-  def pallas_interpret_element_aval(_) -> jax_core.ShapedArray:
-    return jax_core.ShapedArray((), pallas_core.SEMAPHORE_INTERPRET_DTYPE)
-
-  @staticmethod
-  def physical_element_aval(_) -> jax_core.ShapedArray:
-    return jax_core.ShapedArray((), jnp.int32)
-
-class AbstractSemaphoreTy(dtypes.ExtendedDType):
-  name: str
-  _rules = AbstractSemaphoreTyRules
-
-  def __repr__(self) -> str:
-    return self.name
-
-  def __eq__(self, other):
-    return self.__class__ == other.__class__
-
-  def __hash__(self) -> int:
-    return hash(self.__class__)
-
-# TODO(sharadmv): implement dtype rules for AbstractSemaphoreTy
-
-class SemaphoreTy(AbstractSemaphoreTy):
-  type = pallas_core.semaphore
-  name = "sem"
-
-class DmaSemaphoreTy(AbstractSemaphoreTy):
+class DMASemaphore(pallas_core.AbstractSemaphoreTy):
   type = dma_semaphore
   name = "dma_sem"
 
-class BarrierSemaphoreTy(AbstractSemaphoreTy):
-  type = pallas_core.barrier_semaphore
-  name = "barrier_sem"
-
 class SemaphoreType(enum.Enum):
   REGULAR = "regular"
   DMA = "dma"
@@ -158,11 +125,11 @@ class SemaphoreType(enum.Enum):
   def __call__(self, shape: tuple[int, ...]):
     dtype: Any
     if self == SemaphoreType.DMA:
-      dtype = DmaSemaphoreTy()
+      dtype = DMASemaphore()
     elif self == SemaphoreType.BARRIER:
-      dtype = BarrierSemaphoreTy()
+      dtype = pallas_core.BarrierSemaphore()
     else:
-      dtype = SemaphoreTy()
+      dtype = pallas_core.Semaphore()
     return pallas_core.MemoryRef(shape, dtype, TPUMemorySpace.SEMAPHORE)
 
   def get_array_aval(self) -> pallas_core.ShapedArrayWithMemorySpace:

jax/_src/pallas/mosaic/primitives.py

@@ -623,7 +623,7 @@ def async_remote_copy(src_ref, dst_ref, send_sem, recv_sem, device_id,
 @get_barrier_semaphore_p.def_abstract_eval
 def _get_barrier_semaphore_abstract_eval():
   return pl_core.AbstractMemoryRef(
-      jax_core.ShapedArray((), tpu_core.BarrierSemaphoreTy()),
+      jax_core.ShapedArray((), pl_core.BarrierSemaphore()),
       tpu_core.TPUMemorySpace.SEMAPHORE,
   )