{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "474f1562",
   "metadata": {},
   "source": [
    "(pallas-tpu-quickstart)=\n",
    "# Quickstart: TPU\n",
    "\n",
    "This quickstart covers the basics of writing Pallas kernels for TPU:\n",
    "memory spaces, `pl.kernel`, and pipelining.\n",
    "\n",
    "```python\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax.experimental import pallas as pl\n",
    "from jax.experimental.pallas import tpu as pltpu\n",
    "```\n",
    "\n",
    "## TPU memory spaces\n",
    "\n",
    "Pallas kernels interact with memory through **Refs** -- JAX's mutable\n",
    "array references. On TPU, each Ref lives in a specific **memory space**:\n",
    "\n",
    "- **HBM**: large but slow. Kernel inputs and outputs live here.\n",
    "- **VMEM**: small but fast. Computation happens here.\n",
    "\n",
    "You can't compute directly on HBM Refs. Data has to be copied to VMEM\n",
    "first, and results copied back.\n",
    "\n",
    "## Your first kernel\n",
    "\n",
    "Here's a kernel that fills an output array with a constant. We allocate a\n",
    "temporary VMEM buffer, write the constant there, and copy it to the output\n",
    "HBM Ref.\n",
    "\n",
    "```python\n",
    "@pl.kernel(\n",
    "    out_type=jax.ShapeDtypeStruct((128,), jnp.float32),\n",
    "    mesh=pltpu.TensorCoreMesh(axis_name='core'),\n",
    "    scratch_types=dict(o_vmem=pltpu.VMEM((128,), jnp.float32)),\n",
    ")\n",
    "def fill_42(o_ref, o_vmem):\n",
    "  # Compute in VMEM\n",
    "  o_vmem[...] = jnp.full_like(o_vmem, 42.0)\n",
    "\n",
    "  # VMEM → HBM (blocks until the transfer completes)\n",
    "  pltpu.sync_copy(o_vmem, o_ref)\n",
    "\n",
    "result = fill_42()  # [42.0, 42.0, ...]\n",
    "```\n",
    "\n",
    "`pl.kernel` wraps a kernel function to make it callable with standard JAX arrays.\n",
    "It handles allocating the necessary Refs and executing the kernel on-device:\n",
    "\n",
    "- `out_type` allocates an output Ref in HBM for the result.\n",
    "- `scratch_types` allocates the VMEM buffers (like `o_vmem`) that get passed\n",
    "  as extra arguments to the kernel.\n",
    "\n",
    "```{note}\n",
    "If your chip has multiple TensorCores (e.g. TPU v5p has 2), the kernel\n",
    "runs on all of them. As written, both cores will redundantly perform the exact\n",
    "same computation. In a real kernel, you would distribute work across cores\n",
    "either manually or by using pipelining.\n",
    "```\n",
    "\n",
    "### Distributing work across TensorCores\n",
    "\n",
    "To process a larger array, we can distribute work across multiple TensorCores\n",
    "manually. Each core runs a separate invocation of the kernel.\n",
    "\n",
    "```python\n",
    "def iota() -> jax.Array:\n",
    "  tpu_info = pltpu.get_tpu_info()\n",
    "\n",
    "  @pl.kernel(\n",
    "      out_type=jax.ShapeDtypeStruct((128 * tpu_info.num_cores,), jnp.float32),\n",
    "      mesh=pltpu.TensorCoreMesh(axis_name='core'),\n",
    "      scratch_types=dict(o_vmem=pltpu.VMEM((128,), jnp.float32)),\n",
    "  )\n",
    "  def kernel(o_ref, o_vmem):\n",
    "    i = jax.lax.axis_index('core')\n",
    "\n",
    "    # Compute our chunk in VMEM\n",
    "    o_vmem[...] = jnp.arange(128, dtype=jnp.float32) + i * 128\n",
    "\n",
    "    # Copy back to our slice of HBM\n",
    "    pltpu.sync_copy(o_vmem, o_ref.at[pl.ds(i * 128, 128)])\n",
    "\n",
    "  return kernel()\n",
    "\n",
    "result = iota()  # [0.0, 1.0, 2.0, ...]\n",
    "```\n",
    "\n",
    "`jax.lax.axis_index('core')` tells you which TensorCore you're on -- use it to\n",
    "determine your slice of the output.\n",
    "\n",
    "## Making it fast: pipelining\n",
    "\n",
    "The examples above work, but are inefficient. They copy data in, compute,\n",
    "then copy out -- each `sync_copy` blocks until the transfer completes.\n",
    "The TensorCores sit idle while memory is moving, and vice versa.\n",
    "\n",
    "**Pipelining** fixes this by breaking the work into blocks and overlapping\n",
    "memory transfers with compute. While the TensorCores process block N, the\n",
    "DMA engine is already fetching block N+1.\n",
    "\n",
    "`pltpu.emit_pipeline` handles the double-buffering machinery for you.\n",
    "It takes:\n",
    "\n",
    "- A **grid**: how many pipeline blocks to process.\n",
    "- **`BlockSpec`s**: how to slice the arrays into blocks.\n",
    "- A **body** function: the computation to run for each block.\n",
    "\n",
    "Here's an elementwise addition kernel using pipelining:\n",
    "\n",
    "```python\n",
    "def add_matrices_pipelined(x: jax.Array, y: jax.Array) -> jax.Array:\n",
    "  @pl.kernel(\n",
    "      out_type=jax.ShapeDtypeStruct.like(x),\n",
    "      mesh=pltpu.TensorCoreMesh(axis_name='core'),\n",
    "  )\n",
    "  def kernel(x_hbm, y_hbm, o_hbm):\n",
    "    def add_body(x_vmem, y_vmem, o_vmem):\n",
    "      o_vmem[...] = x_vmem[...] + y_vmem[...]\n",
    "\n",
    "    pltpu.emit_pipeline(\n",
    "        add_body,\n",
    "        grid=(x_hbm.shape[0] // 128, x_hbm.shape[1] // 128),\n",
    "        in_specs=[\n",
    "            pl.BlockSpec((128, 128), lambda i, j: (i, j)),\n",
    "            pl.BlockSpec((128, 128), lambda i, j: (i, j)),\n",
    "        ],\n",
    "        out_specs=pl.BlockSpec((128, 128), lambda i, j: (i, j)),\n",
    "        core_axis_name='core',\n",
    "        dimension_semantics=(pltpu.PARALLEL, pltpu.ARBITRARY),\n",
    "    )(x_hbm, y_hbm, o_hbm)\n",
    "  return kernel(x, y)\n",
    "\n",
    "x = 2 * jnp.ones((256, 256))\n",
    "y = 3 * jnp.ones((256, 256))\n",
    "add_matrices_pipelined(x, y)\n",
    "# Array([[5., 5., 5., ..., 5., 5., 5.],\n",
    "#        [5., 5., 5., ..., 5., 5., 5.],\n",
    "#        [5., 5., 5., ..., 5., 5., 5.],\n",
    "#        ...,\n",
    "#        [5., 5., 5., ..., 5., 5., 5.],\n",
    "#        [5., 5., 5., ..., 5., 5., 5.],\n",
    "#        [5., 5., 5., ..., 5., 5., 5.]], dtype=float32)\n",
    "```\n",
    "\n",
    "Notice the body is just the compute -- no `sync_copy`, no scratch allocation.\n",
    "`emit_pipeline` takes care of the memory movement and buffer management.\n",
    "\n",
    "We also passed `core_axis_name='core'` and `dimension_semantics`. This tells\n",
    "the pipeline how to map the grid to the hardware:\n",
    "- `pltpu.PARALLEL` means the first dimension of the grid is distributed across\n",
    "  the TensorCores, splitting the work automatically.\n",
    "- `pltpu.ARBITRARY` indicates that Pallas cannot make assumptions about data\n",
    "  independence along this dimension; it therefore cannot be parallelized across\n",
    "  cores and is instead executed sequentially on a single core.\n",
    "\n",
    "## What's next\n",
    "\n",
    "- [TPU Details](details) -- the full picture of TPU memory spaces and supported ops\n",
    "- [TPU Pipelining](pipelining.md) -- deeper dive into pipelining, reductions, and accumulation\n",
    "- [Matrix Multiplication](matmul.md) -- putting it all together in a full kernel"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}