{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8ac3acdc",
   "metadata": {},
   "source": [
    "(pallas-gpu-quickstart)=\n",
    "# Quickstart: GPU\n",
    "\n",
    "This quickstart shows you how to write Pallas kernels for NVIDIA GPUs\n",
    "using Mosaic GPU. The examples target Hopper (H100) GPUs, but the core\n",
    "concepts - memory spaces, grids, pipelining - apply to all supported GPU\n",
    "generations.\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 mosaic_gpu as plgpu\n",
    "```\n",
    "\n",
    "## The programming model\n",
    "\n",
    "In Mosaic GPU, each Pallas \"thread\" corresponds to a warpgroup\n",
    "(4 warps = 128 CUDA threads). You write straight-line code that operates\n",
    "on arrays; the warpgroup executes it in lockstep, so there are no\n",
    "individual CUDA threads to manage.\n",
    "\n",
    "## GPU memory spaces\n",
    "\n",
    "Pallas kernels access memory through **Refs**, JAX's mutable array\n",
    "references. On GPU, each Ref lives in a specific **memory space**:\n",
    "\n",
    "- **GMEM** (global memory / HBM): large, slow. Used for kernel inputs\n",
    "  and outputs.\n",
    "- **SMEM** (shared memory): small, fast, per-SM. Shared between\n",
    "  threads within a block, used to stage data for Tensor Core ops.\n",
    "- **TMEM** (tensor memory): fast, per-SM. Used for Tensor Core ops\n",
    "  (only available on Blackwell GPUs).\n",
    "\n",
    "The typical data flow for Tensor Core work on Hopper is:\n",
    "GMEM → SMEM → Tensor Cores → registers → SMEM → GMEM\n",
    "(on Blackwell, TMEM replaces SMEM for staging Tensor Core inputs and outputs).\n",
    "Simpler ALU work (like elementwise operations) bypasses SMEM and Tensor Cores\n",
    "entirely: GMEM → registers → GMEM. Shared memory is primarily needed if\n",
    "threads within a block must exchange or reuse data.\n",
    "\n",
    "## Your first kernel\n",
    "\n",
    "Here's a kernel that fills an output array with a constant:\n",
    "\n",
    "```python\n",
    "@plgpu.kernel(out_type=jax.ShapeDtypeStruct((128,), jnp.float32))\n",
    "def fill_42(o_ref):\n",
    "  o_ref[...] = jnp.full_like(o_ref, 42.0)\n",
    "\n",
    "result = fill_42()  # [42.0, 42.0, ...]\n",
    "```\n",
    "\n",
    "`plgpu.kernel` allocates the output buffer, runs the computation on-device, and\n",
    "returns a JAX array.\n",
    "\n",
    "### Parallel work with a grid\n",
    "\n",
    "To process a larger array, add a **grid**. Each grid point runs as a\n",
    "separate CUDA block in parallel across SMs:\n",
    "\n",
    "```python\n",
    "@plgpu.kernel(\n",
    "    out_type=jax.ShapeDtypeStruct((1024,), jnp.float32),\n",
    "    grid=(8,),\n",
    "    grid_names=('i',),\n",
    ")\n",
    "def iota(o_ref):\n",
    "  i = jax.lax.axis_index('i')\n",
    "  o_ref[pl.ds(i * 128, 128)] = jnp.arange(128, dtype=jnp.float32) + i * 128\n",
    "\n",
    "result = iota()  # [0.0, 1.0, ..., 1023.0]\n",
    "```\n",
    "\n",
    "`jax.lax.axis_index(name)` tells you which grid block you're in - use it to\n",
    "determine your slice of the output. `pl.ds(start, size)` creates a dynamic\n",
    "slice of the given `size`, and is equivalent to `start:start+size`.\n",
    "\n",
    "## Pipelining: a matmul kernel\n",
    "\n",
    "The examples above don't need pipelining, but anything that hits the\n",
    "Tensor Cores - matmuls, attention, etc. - should overlap GMEM↔SMEM\n",
    "transfers with compute. Otherwise, the Tensor Cores sit idle waiting for\n",
    "data.\n",
    "\n",
    "`plgpu.emit_pipeline` handles this. It takes:\n",
    "\n",
    "- A **sequential grid**: how many pipeline steps to run (typically\n",
    "  over the contraction dimension).\n",
    "- **`BlockSpec`s**: how to slice your inputs per step.\n",
    "- A **body** function: the computation to run at each step.\n",
    "\n",
    "The outer `plgpu.kernel` grid handles the parallel work, and\n",
    "`emit_pipeline` handles the sequential reduction within each block.\n",
    "\n",
    "Here's a matmul for Hopper GPUs:\n",
    "\n",
    "```{note}\n",
    "This example uses `wgmma`, which is specific to Hopper GPUs. For Blackwell, see\n",
    "[Blackwell Matrix Multiplication](blackwell_matmul.md), which uses `tcgen05`\n",
    "instead.\n",
    "```\n",
    "\n",
    "```python\n",
    "def matmul(a, b, tile_m=128, tile_n=128, tile_k=64, out_dtype=jnp.float16):\n",
    "  m, k = a.shape\n",
    "  _, n = b.shape\n",
    "\n",
    "  @plgpu.kernel(\n",
    "      out_type=jax.ShapeDtypeStruct((m, n), out_dtype),\n",
    "      scratch_types=dict(\n",
    "          o_smem=plgpu.SMEM((tile_m, tile_n), out_dtype),\n",
    "          acc=plgpu.ACC((tile_m, tile_n), jnp.float32),\n",
    "      ),\n",
    "      grid=(m // tile_m, n // tile_n),\n",
    "      grid_names=('m', 'n'),\n",
    "  )\n",
    "  def kernel(a_gmem, b_gmem, o_gmem, o_smem, acc):\n",
    "    pid_m = jax.lax.axis_index('m')\n",
    "    pid_n = jax.lax.axis_index('n')\n",
    "\n",
    "    def body(_, a_smem, b_smem):\n",
    "      plgpu.wgmma(acc, a_smem, b_smem)\n",
    "      plgpu.wgmma_wait(1)  # Keep one wgmma in flight.\n",
    "\n",
    "    plgpu.emit_pipeline(\n",
    "        body,\n",
    "        grid=(k // tile_k,),\n",
    "        in_specs=[\n",
    "            plgpu.BlockSpec(\n",
    "                (tile_m, tile_k), lambda ki: (pid_m, ki), delay_release=1\n",
    "            ),\n",
    "            plgpu.BlockSpec(\n",
    "                (tile_k, tile_n), lambda ki: (ki, pid_n), delay_release=1\n",
    "            ),\n",
    "        ],\n",
    "        max_concurrent_steps=2,\n",
    "    )(a_gmem, b_gmem)\n",
    "\n",
    "    # Drain: move the accumulated result to GMEM via SMEM.\n",
    "    o_smem[...] = acc[...].astype(out_dtype)\n",
    "    plgpu.commit_smem()  # Make the SMEM write visible to the TMA engine.\n",
    "    plgpu.copy_smem_to_gmem(\n",
    "        o_smem,\n",
    "        o_gmem.at[pl.ds(pid_m * tile_m, tile_m),\n",
    "                  pl.ds(pid_n * tile_n, tile_n)],\n",
    "    )\n",
    "    plgpu.wait_smem_to_gmem(0)  # Wait for all copies to finish.\n",
    "\n",
    "  return kernel(a, b)\n",
    "```\n",
    "\n",
    "A few things to note:\n",
    "\n",
    "- The **parallel grid** (`plgpu.kernel(..., grid=...)`) maps one CUDA block\n",
    "  per output tile.\n",
    "- The **sequential grid** (`emit_pipeline(..., grid=...)`) is the pipeline\n",
    "  loop over the K dimension.\n",
    "- **`scratch_types`** specifies temporary memory allocations needed for each\n",
    "  parallel grid point, such as shared memory buffers (`SMEM`) or Tensor Core\n",
    "  accumulators (`ACC`). Each key in the dictionary is passed as a keyword\n",
    "  argument to the kernel function.\n",
    "- **`plgpu.ACC`** is the Tensor Core accumulator. `wgmma` accumulates\n",
    "  into it asynchronously.\n",
    "- **`delay_release=1`** tells the pipeline to keep an extra buffer alive.\n",
    "  Without this, the pipeline releases the corresponding buffer immediately,\n",
    "  allowing the next iteration to overwrite input data while `wgmma` is still\n",
    "  reading it.\n",
    "\n",
    "## What's next\n",
    "\n",
    "- [Mosaic GPU Pipelining](pipelining.md) - pipelining in depth, including warp\n",
    "  specialization\n",
    "- [Mosaic GPU Reference](reference.md) - full API reference (memory spaces,\n",
    "  layouts, Tensor Core ops)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}