[pallas] Added API docs for Triton and Mosaic GPU backends

I've left the TPU backend docs a stub for now. Hopefully, someone working
on Pallas TPU can fill them in later.

docs/jax.experimental.pallas.mosaic_gpu.rst

@@ -0,0 +1,41 @@
+``jax.experimental.pallas.mosaic_gpu`` module
+=============================================
+
+.. automodule:: jax.experimental.pallas.mosaic_gpu
+
+Classes
+-------
+
+.. autosummary::
+   :toctree: _autosummary
+
+   Barrier
+   GPUBlockSpec
+   GPUCompilerParams
+   GPUMemorySpace
+   TilingTransform
+   TransposeTransform
+   WGMMAAccumulatorRef
+
+Functions
+---------
+
+.. autosummary::
+   :toctree: _autosummary
+
+   copy_gmem_to_smem
+   copy_smem_to_gmem
+   wait_barrier
+   wait_smem_to_gmem
+   wgmma
+   wgmma_wait
+
+Aliases
+-------
+
+.. autosummary::
+   :toctree: _autosummary
+
+   ACC
+   GMEM
+   SMEM

docs/jax.experimental.pallas.rst

@@ -3,6 +3,16 @@
 
 .. automodule:: jax.experimental.pallas
 
+Backends
+--------
+
+.. toctree::
+    :maxdepth: 1
+
+    jax.experimental.pallas.mosaic_gpu
+    jax.experimental.pallas.triton
+    jax.experimental.pallas.tpu
+
 Classes
 -------
 

docs/jax.experimental.pallas.tpu.rst

@@ -0,0 +1,16 @@
+``jax.experimental.pallas.tpu`` module
+======================================
+
+.. automodule:: jax.experimental.pallas.tpu
+
+Classes
+-------
+
+.. autosummary::
+   :toctree: _autosummary
+
+Functions
+---------
+
+.. autosummary::
+   :toctree: _autosummary

docs/jax.experimental.pallas.triton.rst

@@ -0,0 +1,22 @@
+``jax.experimental.pallas.triton`` module
+=========================================
+
+.. automodule:: jax.experimental.pallas.triton
+
+Classes
+-------
+
+.. autosummary::
+   :toctree: _autosummary
+
+   TritonCompilerParams
+
+Functions
+---------
+
+.. autosummary::
+   :toctree: _autosummary
+
+   approx_tanh
+   debug_barrier
+   elementwise_inline_asm

jax/_src/pallas/mosaic_gpu/core.py

@@ -59,8 +59,11 @@ class GPUCompilerParams(pallas_core.CompilerParams):
 
 
 class GPUMemorySpace(enum.Enum):
+  #: Global memory.
   GMEM = "gmem"
+  #: Shared memory.
   SMEM = "smem"
+  #: Registers.
   REGS = "regs"
 
   def __str__(self) -> str:

jax/experimental/pallas/mosaic_gpu.py

@@ -23,13 +23,17 @@
 from jax._src.pallas.mosaic_gpu.core import GPUMemorySpace
 from jax._src.pallas.mosaic_gpu.core import TilingTransform
 from jax._src.pallas.mosaic_gpu.core import TransposeTransform
-from jax._src.pallas.mosaic_gpu.core import WGMMAAccumulatorRef as ACC
+from jax._src.pallas.mosaic_gpu.core import WGMMAAccumulatorRef
 from jax._src.pallas.mosaic_gpu.primitives import copy_gmem_to_smem
 from jax._src.pallas.mosaic_gpu.primitives import copy_smem_to_gmem
 from jax._src.pallas.mosaic_gpu.primitives import wait_barrier
 from jax._src.pallas.mosaic_gpu.primitives import wait_smem_to_gmem
 from jax._src.pallas.mosaic_gpu.primitives import wgmma
 from jax._src.pallas.mosaic_gpu.primitives import wgmma_wait
 
+#: Alias of :class:`jax.experimental.pallas.mosaic_gpu.WGMMAAccumulatorRef`.
+ACC = WGMMAAccumulatorRef
+#: Alias of :data:`jax.experimental.pallas.mosaic_gpu.GPUMemorySpace.GMEM`.
 GMEM = GPUMemorySpace.GMEM
+#: Alias of :data:`jax.experimental.pallas.mosaic_gpu.GPUMemorySpace.SMEM`.
 SMEM = GPUMemorySpace.SMEM