[pallas:mosaic] Removed the TPU prefix from TPUCompilerParams and TPUMemorySpace

All TPU-specific APIs are always used qualified, e.g. `pltpu.TPUCompilerParams`,
so the prefix is redundant.

PiperOrigin-RevId: 765167675

Author: superbobry

docs/pallas/CHANGELOG.md

@@ -18,6 +18,11 @@ Remember to align the itemized text with the first line of an item within a list
   * {class}`jax.experimental.pallas.triton.TritonCompilerParams` has been
     renamed to {class}`jax.experimental.pallas.triton.CompilerParams`. The
     old name is deprecated and will be removed in a future release.
+  * {class}`jax.experimental.pallas.tpu.TPUCompilerParams`
+    and {class}`jax.experimental.pallas.tpu.TPUMemorySpace` have been
+    renamed to {class}`jax.experimental.pallas.tpu.CompilerParams`
+    and {class}`jax.experimental.pallas.tpu.MemorySpace`. The
+    old names are deprecated and will be removed in a future release.
 
 ## Released with jax 0.6.1
 

docs/pallas/quickstart.ipynb

@@ -280,7 +280,7 @@
    "metadata": {},
    "source": [
     "TPUs distinguish between vector and scalar memory spaces and in this case the\n",
-    "output must be placed in scalar memory (`TPUMemorySpace.SMEM`) since `i` is\n",
+    "output must be placed in scalar memory (`MemorySpace.SMEM`) since `i` is\n",
     "a scalar. For more details read {ref}`tpu_and_its_memory_spaces`.\n",
     "To call the above kernel on TPU, run:"
    ]
@@ -297,7 +297,7 @@
     "\n",
     "def iota(size: int):\n",
     "  return pl.pallas_call(iota_kernel,\n",
-    "                        out_specs=pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.SMEM),\n",
+    "                        out_specs=pl.BlockSpec(memory_space=pltpu.MemorySpace.SMEM),\n",
     "                        out_shape=jax.ShapeDtypeStruct((size,), jnp.int32),\n",
     "                        grid=(size,))()\n",
     "iota(8)"

docs/pallas/quickstart.md

@@ -186,7 +186,7 @@ iota(8)
 ```
 
 TPUs distinguish between vector and scalar memory spaces and in this case the
-output must be placed in scalar memory (`TPUMemorySpace.SMEM`) since `i` is
+output must be placed in scalar memory (`MemorySpace.SMEM`) since `i` is
 a scalar. For more details read {ref}`tpu_and_its_memory_spaces`.
 To call the above kernel on TPU, run:
 
@@ -196,7 +196,7 @@ from jax.experimental.pallas import tpu as pltpu
 
 def iota(size: int):
   return pl.pallas_call(iota_kernel,
-                        out_specs=pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.SMEM),
+                        out_specs=pl.BlockSpec(memory_space=pltpu.MemorySpace.SMEM),
                         out_shape=jax.ShapeDtypeStruct((size,), jnp.int32),
                         grid=(size,))()
 iota(8)

docs/pallas/tpu/details.rst

@@ -170,7 +170,7 @@ grid axes over cores. This is an opt-in procedure. To allow that,
 ..
   pallas_call(
       ...,
-      compiler_params=pltpu.TPUCompilerParams(
+      compiler_params=pltpu.CompilerParams(
           dimension_semantics=["parallel", "parallel", "arbitrary"]
       ),
     )

docs/pallas/tpu/distributed.ipynb

@@ -271,11 +271,11 @@
     "out_shape = jax.ShapeDtypeStruct((8, 128), jnp.float32)\n",
     "grid_spec = pltpu.PrefetchScalarGridSpec(\n",
     "    num_scalar_prefetch=0,\n",
-    "    # TPUMemorySpace.ANY will (usually) place the tensor in HBM.\n",
+    "    # MemorySpace.ANY will (usually) place the tensor in HBM.\n",
     "    in_specs=[\n",
-    "        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),\n",
+    "        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),\n",
     "    ],\n",
-    "    out_specs=pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),\n",
+    "    out_specs=pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),\n",
     "    scratch_shapes=(\n",
     "        # We allocate DMA semaphores in scratch memory.\n",
     "        [pltpu.SemaphoreType.DMA] * 2\n",
@@ -420,10 +420,10 @@
     "grid_spec = pltpu.PrefetchScalarGridSpec(\n",
     "            num_scalar_prefetch=0,\n",
     "            in_specs=[\n",
-    "                # TPUMemorySpace.ANY will (usually) place the tensor in HBM.\n",
-    "                pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),\n",
+    "                # MemorySpace.ANY will (usually) place the tensor in HBM.\n",
+    "                pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),\n",
     "            ],\n",
-    "            out_specs=pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),\n",
+    "            out_specs=pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),\n",
     "            scratch_shapes=(\n",
     "              # DMA semaphores are allocated in scratch memory.\n",
     "              # We allocated one semaphore for a local HBM-VMEM copy,\n",
@@ -569,7 +569,7 @@
     "kernel = pl.pallas_call(\n",
     "      example_kernel,\n",
     "      ...,\n",
-    "      compiler_params=pltpu.TPUCompilerParams(collective_id=0),\n",
+    "      compiler_params=pltpu.CompilerParams(collective_id=0),\n",
     ")\n",
     "```"
    ]
@@ -809,13 +809,13 @@
     "    num_scalar_prefetch=0,\n",
     "    in_specs=[\n",
     "        # Our input lives in VMEM\n",
-    "        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.VMEM),\n",
+    "        pl.BlockSpec(memory_space=pltpu.MemorySpace.VMEM),\n",
     "    ],\n",
     "    out_specs=[\n",
     "        # Our output lives in VMEM\n",
-    "        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.VMEM),\n",
+    "        pl.BlockSpec(memory_space=pltpu.MemorySpace.VMEM),\n",
     "        # Our double-buffer lives in HBM\n",
-    "        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),\n",
+    "        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),\n",
     "    ],\n",
     "    grid=(num_devices,),\n",
     "    scratch_shapes=(\n",
@@ -829,7 +829,7 @@
     "    all_reduce_kernel,\n",
     "    out_shape=out_shape,\n",
     "    grid_spec=grid_spec,\n",
-    "    compiler_params=pltpu.TPUCompilerParams(collective_id=0),\n",
+    "    compiler_params=pltpu.CompilerParams(collective_id=0),\n",
     ")\n",
     "\n",
     "pallas_result = jax.jit(\n",
@@ -1146,11 +1146,11 @@
     "grid_spec = pltpu.PrefetchScalarGridSpec(\n",
     "    num_scalar_prefetch=0,\n",
     "    in_specs=[\n",
-    "        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.VMEM),\n",
+    "        pl.BlockSpec(memory_space=pltpu.MemorySpace.VMEM),\n",
     "    ],\n",
     "    out_specs=[\n",
-    "        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.VMEM),\n",
-    "        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),\n",
+    "        pl.BlockSpec(memory_space=pltpu.MemorySpace.VMEM),\n",
+    "        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),\n",
     "    ],\n",
     "    grid=(num_devices, 2),\n",
     "    scratch_shapes=(\n",
@@ -1169,7 +1169,7 @@
     "      reduce_scatter_kernel,\n",
     "      out_shape=out_shape,\n",
     "      grid_spec=grid_spec,\n",
-    "      compiler_params=pltpu.TPUCompilerParams(collective_id=0),\n",
+    "      compiler_params=pltpu.CompilerParams(collective_id=0),\n",
     "  )(input_arr)[0]\n",
     "\n",
     "\n",
@@ -1307,7 +1307,7 @@
     "\n",
     "In this next example we will modify our previous reduce-scatter example to utilize a nested inner pipeline. Note that the communication and computation costs of `reduce_scatter` both scale linearly with the size of the input, so we do not necessarily expect to see the operation become compute-bound with larger block sizes. This example is purely for demonstration purposes on how to use the pipeline emitter.\n",
     "\n",
-    "We will increase the block sizes of the outer kernel such that they would be undesirable to place inside of VMEM, and allocate all inputs and outputs in HBM (`memory_space=TPUMemorySpace.Any`). The only major change from our previous kernel is the body of the kernel where accumulation is done. Rather than manually copying from HBM to VMEM, accumulating, and copying back to HBM, we use `emit_pipeline` to handle the memory transfers for us. Accumulation is done in an inner kernel with a much smaller, VMEM-friendly block size.\n",
+    "We will increase the block sizes of the outer kernel such that they would be undesirable to place inside of VMEM, and allocate all inputs and outputs in HBM (`memory_space=MemorySpace.ANY`). The only major change from our previous kernel is the body of the kernel where accumulation is done. Rather than manually copying from HBM to VMEM, accumulating, and copying back to HBM, we use `emit_pipeline` to handle the memory transfers for us. Accumulation is done in an inner kernel with a much smaller, VMEM-friendly block size.\n",
     "\n",
     "In our previous kernel we had the following kernel body to copy data from HBM to the VMEM accumulator, increment, and then copy the results back to HBM:\n",
     "\n",
@@ -1408,7 +1408,7 @@
     "inner_block_spec = pl.BlockSpec(\n",
     "    index_map=lambda i, j: (i, j),\n",
     "    block_shape=inner_block_size,\n",
-    "    memory_space=pltpu.TPUMemorySpace.ANY,\n",
+    "    memory_space=pltpu.MemorySpace.ANY,\n",
     ")\n",
     "\n",
     "\n",
@@ -1590,11 +1590,11 @@
     "grid_spec = pltpu.PrefetchScalarGridSpec(\n",
     "    num_scalar_prefetch=0,\n",
     "    in_specs=[\n",
-    "        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),\n",
+    "        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),\n",
     "    ],\n",
     "    out_specs=[\n",
-    "        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),\n",
-    "        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),\n",
+    "        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),\n",
+    "        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),\n",
     "    ],\n",
     "    grid=(num_devices, 2),\n",
     "    scratch_shapes=(\n",
@@ -1612,7 +1612,7 @@
     "      reduce_scatter_kernel,\n",
     "      out_shape=out_shape,\n",
     "      grid_spec=grid_spec,\n",
-    "      compiler_params=pltpu.TPUCompilerParams(collective_id=0),\n",
+    "      compiler_params=pltpu.CompilerParams(collective_id=0),\n",
     "  )(input_arr)[0]\n",
     "\n",
     "\n",

docs/pallas/tpu/distributed.md

@@ -233,11 +233,11 @@ def right_permute_kernel(input_ref, output_ref, send_sem, recv_sem):
 out_shape = jax.ShapeDtypeStruct((8, 128), jnp.float32)
 grid_spec = pltpu.PrefetchScalarGridSpec(
     num_scalar_prefetch=0,
-    # TPUMemorySpace.ANY will (usually) place the tensor in HBM.
+    # MemorySpace.ANY will (usually) place the tensor in HBM.
     in_specs=[
-        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
+        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),
     ],
-    out_specs=pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
+    out_specs=pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),
     scratch_shapes=(
         # We allocate DMA semaphores in scratch memory.
         [pltpu.SemaphoreType.DMA] * 2
@@ -356,10 +356,10 @@ out_shape = jax.ShapeDtypeStruct((num_devices, 8, 128), jnp.float32)
 grid_spec = pltpu.PrefetchScalarGridSpec(
             num_scalar_prefetch=0,
             in_specs=[
-                # TPUMemorySpace.ANY will (usually) place the tensor in HBM.
-                pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
+                # MemorySpace.ANY will (usually) place the tensor in HBM.
+                pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),
             ],
-            out_specs=pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
+            out_specs=pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),
             scratch_shapes=(
               # DMA semaphores are allocated in scratch memory.
               # We allocated one semaphore for a local HBM-VMEM copy,
@@ -491,7 +491,7 @@ When using barrier semaphores, the `collective_id` compiler parameter must be pa
 kernel = pl.pallas_call(
       example_kernel,
       ...,
-      compiler_params=pltpu.TPUCompilerParams(collective_id=0),
+      compiler_params=pltpu.CompilerParams(collective_id=0),
 )
 ```
 
@@ -703,13 +703,13 @@ grid_spec = pltpu.PrefetchScalarGridSpec(
     num_scalar_prefetch=0,
     in_specs=[
         # Our input lives in VMEM
-        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.VMEM),
+        pl.BlockSpec(memory_space=pltpu.MemorySpace.VMEM),
     ],
     out_specs=[
         # Our output lives in VMEM
-        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.VMEM),
+        pl.BlockSpec(memory_space=pltpu.MemorySpace.VMEM),
         # Our double-buffer lives in HBM
-        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
+        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),
     ],
     grid=(num_devices,),
     scratch_shapes=(
@@ -723,7 +723,7 @@ kernel = pl.pallas_call(
     all_reduce_kernel,
     out_shape=out_shape,
     grid_spec=grid_spec,
-    compiler_params=pltpu.TPUCompilerParams(collective_id=0),
+    compiler_params=pltpu.CompilerParams(collective_id=0),
 )
 
 pallas_result = jax.jit(
@@ -1019,11 +1019,11 @@ out_shape = (
 grid_spec = pltpu.PrefetchScalarGridSpec(
     num_scalar_prefetch=0,
     in_specs=[
-        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.VMEM),
+        pl.BlockSpec(memory_space=pltpu.MemorySpace.VMEM),
     ],
     out_specs=[
-        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.VMEM),
-        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
+        pl.BlockSpec(memory_space=pltpu.MemorySpace.VMEM),
+        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),
     ],
     grid=(num_devices, 2),
     scratch_shapes=(
@@ -1042,7 +1042,7 @@ def pallas_reduce_scatter(input_arr):
       reduce_scatter_kernel,
       out_shape=out_shape,
       grid_spec=grid_spec,
-      compiler_params=pltpu.TPUCompilerParams(collective_id=0),
+      compiler_params=pltpu.CompilerParams(collective_id=0),
   )(input_arr)[0]
 
 
@@ -1148,7 +1148,7 @@ pl.pallas_call(
 
 In this next example we will modify our previous reduce-scatter example to utilize a nested inner pipeline. Note that the communication and computation costs of `reduce_scatter` both scale linearly with the size of the input, so we do not necessarily expect to see the operation become compute-bound with larger block sizes. This example is purely for demonstration purposes on how to use the pipeline emitter.
 
-We will increase the block sizes of the outer kernel such that they would be undesirable to place inside of VMEM, and allocate all inputs and outputs in HBM (`memory_space=TPUMemorySpace.Any`). The only major change from our previous kernel is the body of the kernel where accumulation is done. Rather than manually copying from HBM to VMEM, accumulating, and copying back to HBM, we use `emit_pipeline` to handle the memory transfers for us. Accumulation is done in an inner kernel with a much smaller, VMEM-friendly block size.
+We will increase the block sizes of the outer kernel such that they would be undesirable to place inside of VMEM, and allocate all inputs and outputs in HBM (`memory_space=MemorySpace.ANY`). The only major change from our previous kernel is the body of the kernel where accumulation is done. Rather than manually copying from HBM to VMEM, accumulating, and copying back to HBM, we use `emit_pipeline` to handle the memory transfers for us. Accumulation is done in an inner kernel with a much smaller, VMEM-friendly block size.
 
 In our previous kernel we had the following kernel body to copy data from HBM to the VMEM accumulator, increment, and then copy the results back to HBM:
 
@@ -1242,7 +1242,7 @@ inner_grid = (
 inner_block_spec = pl.BlockSpec(
     index_map=lambda i, j: (i, j),
     block_shape=inner_block_size,
-    memory_space=pltpu.TPUMemorySpace.ANY,
+    memory_space=pltpu.MemorySpace.ANY,
 )
 
 
@@ -1424,11 +1424,11 @@ out_shape = (
 grid_spec = pltpu.PrefetchScalarGridSpec(
     num_scalar_prefetch=0,
     in_specs=[
-        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
+        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),
     ],
     out_specs=[
-        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
-        pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.ANY),
+        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),
+        pl.BlockSpec(memory_space=pltpu.MemorySpace.ANY),
     ],
     grid=(num_devices, 2),
     scratch_shapes=(
@@ -1446,7 +1446,7 @@ def pallas_reduce_scatter(input_arr):
       reduce_scatter_kernel,
       out_shape=out_shape,
       grid_spec=grid_spec,
-      compiler_params=pltpu.TPUCompilerParams(collective_id=0),
+      compiler_params=pltpu.CompilerParams(collective_id=0),
   )(input_arr)[0]
 
 

docs/pallas/tpu/matmul.ipynb

@@ -210,7 +210,7 @@
     "                pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))],\n",
     "      out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),\n",
     "      grid=(m // bm, n // bn, k // bk),\n",
-    "      compiler_params=pltpu.TPUCompilerParams(\n",
+    "      compiler_params=pltpu.CompilerParams(\n",
     "          dimension_semantics=(\"parallel\", \"parallel\", \"arbitrary\")),\n",
     "  )(x, y)"
    ]
@@ -466,7 +466,7 @@
     "        grid=(m // bm, n // bn, k // bk),\n",
     "      ),\n",
     "      out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),\n",
-    "      compiler_params=pltpu.TPUCompilerParams(\n",
+    "      compiler_params=pltpu.CompilerParams(\n",
     "          dimension_semantics=(\"parallel\", \"parallel\", \"arbitrary\")),\n",
     "  )(x, y)"
    ]
@@ -741,7 +741,7 @@
     "        grid=(m // bm, n // bn, k // bk),\n",
     "      ),\n",
     "      out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),\n",
-    "      compiler_params=pltpu.TPUCompilerParams(\n",
+    "      compiler_params=pltpu.CompilerParams(\n",
     "          dimension_semantics=(\"parallel\", \"parallel\", \"arbitrary\")),\n",
     "  )(x, y)"
    ]
@@ -929,7 +929,7 @@
     "          grid=(m // bm, n // bn, k // bk),\n",
     "      ),\n",
     "      out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),\n",
-    "      compiler_params=pltpu.TPUCompilerParams(\n",
+    "      compiler_params=pltpu.CompilerParams(\n",
     "          dimension_semantics=(\"parallel\", \"parallel\", \"arbitrary\")),\n",
     "  )(x, y)"
    ]

docs/pallas/tpu/matmul.md

@@ -167,7 +167,7 @@ def matmul(
                 pl.BlockSpec((bk, bn), lambda i, j, k: (k, j))],
       out_specs=pl.BlockSpec((bm, bn), lambda i, j, k: (i, j)),
       grid=(m // bm, n // bn, k // bk),
-      compiler_params=pltpu.TPUCompilerParams(
+      compiler_params=pltpu.CompilerParams(
           dimension_semantics=("parallel", "parallel", "arbitrary")),
   )(x, y)
 ```
@@ -321,7 +321,7 @@ def matmul(
         grid=(m // bm, n // bn, k // bk),
       ),
       out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
-      compiler_params=pltpu.TPUCompilerParams(
+      compiler_params=pltpu.CompilerParams(
           dimension_semantics=("parallel", "parallel", "arbitrary")),
   )(x, y)
 ```
@@ -489,7 +489,7 @@ def matmul(
         grid=(m // bm, n // bn, k // bk),
       ),
       out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
-      compiler_params=pltpu.TPUCompilerParams(
+      compiler_params=pltpu.CompilerParams(
           dimension_semantics=("parallel", "parallel", "arbitrary")),
   )(x, y)
 ```
@@ -613,7 +613,7 @@ def matmul(
           grid=(m // bm, n // bn, k // bk),
       ),
       out_shape=jax.ShapeDtypeStruct((m, n), x.dtype),
-      compiler_params=pltpu.TPUCompilerParams(
+      compiler_params=pltpu.CompilerParams(
           dimension_semantics=("parallel", "parallel", "arbitrary")),
   )(x, y)
 ```