{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# 5.6 Common GPU programming from Python\n",
    "\n",
    "One kernel source, written in a small macro vocabulary, compiled at run-time for whichever\n",
    "backend is loaded: **CUDA**, **Metal**, or a **host reference backend** that is always\n",
    "available and compiles the kernel with the system C++ compiler.\n",
    "\n",
    "* [Getting started](#Getting-started)\n",
    "* [Reference: the kernel syntax](#Reference:-the-kernel-syntax)\n",
    "* [Using tinybla](#Using-tinybla)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Getting started\n",
    "\n",
    "`ngsolve.gpu` picks up whichever backend this build has - Metal, CUDA, or the host\n",
    "reference device if neither is there - so `GetGPUDevice()` never returns `None`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from ngsolve.gpu import *\n",
    "\n",
    "dev = GetGPUDevice()\n",
    "print(dev, \" backend =\", backend, \" fp64 =\", dev.has_float64,\n",
    "      \" max group =\", dev.max_threads_per_group)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "A kernel. `KERNEL` declares the entry point, `GLOBAL_IN`/`GLOBAL` are device buffers,\n",
    "`VALUE` a scalar passed by value, and `GLOBAL_ID_X` is the work-item index.\n",
    "Prepend `GPUKernelPrelude`, which defines these for the current backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "src = '''\n",
    "KERNEL(saxpy, GLOBAL_IN(float,x), GLOBAL(float,y), VALUE(float,a), VALUE(int,n))\n",
    "{\n",
    "  int i = GLOBAL_ID_X;\n",
    "  if (i < n) y[i] = a*x[i] + y[i];\n",
    "}\n",
    "'''\n",
    "\n",
    "lib = dev.CompileSource(GPUKernelPrelude + src)\n",
    "q   = dev.DefaultQueue()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "`NewBuffer` allocates device memory, `H2D` uploads a numpy array, `D2H(n)` reads values back.\n",
    "Buffers and scalars are single precision here: `H2D`/`D2H` use `float32`, and a Python\n",
    "`float` argument arrives as `float`, an `int` as `int`.\n",
    "\n",
    "The launch geometry is as in CUDA - `groups` threadgroups of `groupsize` work items.\n",
    "Launches are asynchronous, `Finish()` waits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "n, tg = 1000, 256\n",
    "ngroups = (n+tg-1)//tg\n",
    "hx = np.arange(n, dtype=np.float32)\n",
    "\n",
    "x, y = dev.NewBuffer(4*n), dev.NewBuffer(4*n)\n",
    "x.H2D(hx)\n",
    "y.H2D(np.ones(n, dtype=np.float32))\n",
    "\n",
    "q.Launch(lib.GetKernel(\"saxpy\"), groups=[ngroups], groupsize=[tg], args=[x, y, 2.0, n])\n",
    "q.Finish()\n",
    "\n",
    "print(\"device :\", y.D2H(5))\n",
    "print(\"numpy  :\", 2*hx[:5]+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "The point of the common syntax: the *same* source runs on the host reference backend, so a\n",
    "GPU result can be checked against a reference that shares no code with it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run(device):\n",
    "    lib = device.CompileSource(GPUKernelPrelude + src)\n",
    "    q   = device.DefaultQueue()\n",
    "    xb, yb = device.NewBuffer(4*n), device.NewBuffer(4*n)\n",
    "    xb.H2D(hx)\n",
    "    yb.H2D(np.ones(n, dtype=np.float32))\n",
    "    q.Launch(lib.GetKernel(\"saxpy\"), groups=[ngroups], groupsize=[tg], args=[xb, yb, 2.0, n])\n",
    "    q.Finish()\n",
    "    return yb.D2H(n)\n",
    "\n",
    "on_device = run(dev)\n",
    "on_host   = run(GetCPUDevice())\n",
    "print(\"max difference :\", np.abs(on_device-on_host).max())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Reference: the kernel syntax\n",
    "\n",
    "`GPUKernelPrelude` defines one vocabulary that expands to CUDA, Metal or plain C++.\n",
    "\n",
    "**Entry point and arguments.** Arguments are positional: the i-th declared argument is the\n",
    "i-th entry of the `args` list of `Launch`.\n",
    "\n",
    "| | |\n",
    "|---|---|\n",
    "| `KERNEL(name, args...)` | declares an entry point |\n",
    "| `GLOBAL(T,x)` | device buffer, read and write |\n",
    "| `GLOBAL_IN(T,x)` | device buffer, read only |\n",
    "| `GLOBAL_ATOMIC(T,x)` | device buffer written with `ATOMIC_ADD` |\n",
    "| `VALUE(T,a)` | scalar passed by value |\n",
    "\n",
    "**Indices.** All exist as `_X`, `_Y` and `_Z`.\n",
    "\n",
    "| | |\n",
    "|---|---|\n",
    "| `GLOBAL_ID_X` | index of this work item in the whole grid |\n",
    "| `LOCAL_ID_X` | index inside its threadgroup |\n",
    "| `GROUP_ID_X` | index of the threadgroup |\n",
    "| `GROUP_SIZE_X` | work items per threadgroup |\n",
    "| `NUM_GROUPS_X` | number of threadgroups |\n",
    "\n",
    "**Group-shared memory and synchronisation.**\n",
    "\n",
    "| | |\n",
    "|---|---|\n",
    "| `SHARED(T,s,N)` | array shared by the threadgroup |\n",
    "| `SHARED_2D(T,s,N,M)` | two-dimensional variant |\n",
    "| `BARRIER()` | wait until the whole group arrives |\n",
    "| `ATOMIC_ADD(p,val)` | atomic accumulate into a `GLOBAL_ATOMIC` buffer |\n",
    "\n",
    "**Types for helper functions.**\n",
    "\n",
    "| | |\n",
    "|---|---|\n",
    "| `GPU_FUNC` | qualifier for a function called from a kernel |\n",
    "| `GLOBAL_PTR(T)`, `LOCAL_PTR(T)` | pointer into global / group-shared memory |\n",
    "| `ATOMIC_PTR(T)` | pointer into a `GLOBAL_ATOMIC` buffer |\n",
    "\n",
    "Note that `uint` and `float4` are *not* portable spellings - use `unsigned`, and index\n",
    "buffers element by element."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "A kernel using group-shared memory and a barrier: each group reduces its slice of the\n",
    "input to one partial sum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "reduce_src = '''\n",
    "KERNEL(blocksum, GLOBAL_IN(float,x), GLOBAL(float,out), VALUE(int,n))\n",
    "{\n",
    "  SHARED(float, s, 256);\n",
    "  int lid = LOCAL_ID_X;\n",
    "  int gid = GLOBAL_ID_X;\n",
    "  s[lid] = (gid < n) ? x[gid] : 0.0f;\n",
    "  BARRIER();\n",
    "  for (int stride = 128; stride > 0; stride >>= 1)\n",
    "    { if (lid < stride) s[lid] += s[lid+stride]; BARRIER(); }\n",
    "  if (lid == 0) out[GROUP_ID_X] = s[0];\n",
    "}\n",
    "'''\n",
    "\n",
    "libr = dev.CompileSource(GPUKernelPrelude + reduce_src)\n",
    "out  = dev.NewBuffer(4*ngroups)\n",
    "q.Launch(libr.GetKernel(\"blocksum\"), groups=[ngroups], groupsize=[tg], args=[x, out, n])\n",
    "q.Finish()\n",
    "\n",
    "partial = out.D2H(ngroups)\n",
    "print(\"partial sums :\", partial)\n",
    "print(\"device total :\", partial.sum(), \"  numpy :\", hx.sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "And one where many work items accumulate into the same locations, so the target is\n",
    "declared `GLOBAL_ATOMIC` and written with `ATOMIC_ADD`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "hist_src = '''\n",
    "KERNEL(histogram, GLOBAL_IN(float,x), GLOBAL_ATOMIC(float,bins),\n",
    "                  VALUE(int,n), VALUE(int,nbins))\n",
    "{\n",
    "  int i = GLOBAL_ID_X;\n",
    "  if (i < n) ATOMIC_ADD(&bins[i % nbins], x[i]);\n",
    "}\n",
    "'''\n",
    "\n",
    "nbins = 4\n",
    "libh  = dev.CompileSource(GPUKernelPrelude + hist_src)\n",
    "bins  = dev.NewBuffer(4*nbins)\n",
    "bins.H2D(np.zeros(nbins, dtype=np.float32))\n",
    "\n",
    "q.Launch(libh.GetKernel(\"histogram\"), groups=[ngroups], groupsize=[tg],\n",
    "         args=[x, bins, n, nbins])\n",
    "q.Finish()\n",
    "\n",
    "print(\"device :\", bins.D2H(nbins))\n",
    "print(\"numpy  :\", np.array([hx[i::nbins].sum() for i in range(nbins)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "### Buffers\n",
    "\n",
    "`MemType.Shared` is addressable from host and device, `MemType.Device` lives only on the\n",
    "accelerator. On unified memory shared costs nothing; on a discrete card a device-only\n",
    "buffer avoids keeping a host copy. `D2H` works for both."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "for mt in [MemType.Shared, MemType.Device]:\n",
    "    b = dev.NewBuffer(4*n, mt)\n",
    "    print(f\"{str(mt):18s} host_visible = {b.host_visible}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16",
   "metadata": {},
   "source": [
    "## Using tinybla\n",
    "\n",
    "`tinybla` is a small linear algebra library that lives *inside* the kernel: fixed size\n",
    "vectors and matrices in registers, views on buffers, and tile products across a\n",
    "threadgroup. Prepend `TinyBlaPrelude` after `GPUKernelPrelude` and open the namespace.\n",
    "\n",
    "**Per work item.** `Vec<S,T>` and `Mat<H,W,T>` are held in registers, one per work item.\n",
    "\n",
    "| | |\n",
    "|---|---|\n",
    "| `Vec<S,T>`, `Mat<H,W,T>` | fixed size, `v(i)` and `m(i,j)` |\n",
    "| `+`, `-`, `*` | with a scalar, a vector or a matrix |\n",
    "| `Trans`, `Det`, `Cof`, `Inv` | `Det`/`Cof`/`Inv` for sizes 1, 2, 3 |\n",
    "| `ToVec`, `ToMat<H,W>` | reinterpret between the two |\n",
    "| `v.Range<FIRST,NEXT>()`, `v.SetRange<FIRST,NEXT>(w)` | sub-vectors |\n",
    "\n",
    "**Views on memory.** `BareMatrix` is a pointer plus a leading dimension - no storage of its\n",
    "own - over a buffer or a `SHARED_2D` array.\n",
    "\n",
    "| | |\n",
    "|---|---|\n",
    "| `MakeBareMatrix<RowMajor>(ptr, ld)` | view with explicit leading dimension |\n",
    "| `MakeBareMatrix<RowMajor>(shared2d)` | view of a `SHARED_2D` array |\n",
    "| `m.SubMatrix(r,c)` | view starting at that entry |\n",
    "| `m.Transpose()` | swaps `RowMajor` and `ColMajor`, no data moved |\n",
    "| `m.ShiftRows(n)`, `m.ShiftCols(n)` | move the origin |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "A kernel with one small matrix per work item - invert a 3x3 and apply it, the shape of a\n",
    "geometry transformation in an FE kernel. Checked against the host backend and numpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "tb_src = '''\n",
    "using namespace tinybla;\n",
    "\n",
    "KERNEL(transform, GLOBAL_IN(float,jac), GLOBAL_IN(float,v), GLOBAL(float,out), VALUE(int,n))\n",
    "{\n",
    "  int i = GLOBAL_ID_X;\n",
    "  if (i >= n) return;\n",
    "\n",
    "  Mat<3,3,float> F;\n",
    "  for (int r = 0; r < 3; r++)\n",
    "    for (int c = 0; c < 3; c++)\n",
    "      F(r,c) = jac[9*i + 3*r + c];\n",
    "\n",
    "  Vec<3,float> u;\n",
    "  for (int k = 0; k < 3; k++) u(k) = v[3*i+k];\n",
    "\n",
    "  Vec<3,float> res = Inv(F) * u;\n",
    "  for (int k = 0; k < 3; k++) out[3*i+k] = res(k);\n",
    "  out[3*n+i] = Det(F);\n",
    "}\n",
    "'''\n",
    "\n",
    "m = 64\n",
    "rng = np.random.default_rng(0)\n",
    "J = rng.random((m,3,3)).astype(np.float32) + np.eye(3, dtype=np.float32)\n",
    "V = rng.random((m,3)).astype(np.float32)\n",
    "\n",
    "def run_tb(device):\n",
    "    lib = device.CompileSource(GPUKernelPrelude + TinyBlaPrelude + tb_src)\n",
    "    qq  = device.DefaultQueue()\n",
    "    bj, bv, bo = device.NewBuffer(4*9*m), device.NewBuffer(4*3*m), device.NewBuffer(4*4*m)\n",
    "    bj.H2D(J.reshape(-1))\n",
    "    bv.H2D(V.reshape(-1))\n",
    "    qq.Launch(lib.GetKernel(\"transform\"), groups=[1], groupsize=[64], args=[bj, bv, bo, m])\n",
    "    qq.Finish()\n",
    "    r = bo.D2H(4*m)\n",
    "    return r[:3*m].reshape(m,3), r[3*m:]\n",
    "\n",
    "sol, det = run_tb(dev)\n",
    "sol_host, det_host = run_tb(GetCPUDevice())\n",
    "\n",
    "print(\"Inv(F)*u   vs host :\", np.abs(sol-sol_host).max(),\n",
    "      \"  vs numpy :\", np.abs(sol-np.linalg.solve(J, V[...,None])[...,0]).max())\n",
    "print(\"Det(F)     vs host :\", np.abs(det-det_host).max(),\n",
    "      \"  vs numpy :\", np.abs(det-np.linalg.det(J)).max())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "### Tiles across a threadgroup\n",
    "\n",
    "`WarpMatrix<H,W,T>` is one tile held *collectively* by a simdgroup (32 lanes), which is how\n",
    "the matrix-free operators reach the tensor hardware. It is built from a `BareMatrix` and\n",
    "the work-item index, accumulates products with `AddMM<K>`, and is written back with `Store`.\n",
    "\n",
    "| | |\n",
    "|---|---|\n",
    "| `WarpMatrix<H,W,T> t = 0.0f;` | a zero tile, `H%8==0`, `W%4==0` |\n",
    "| `WarpMatrix<H,W,T> t(baremat, tid);` | load a tile |\n",
    "| `t.AddMM<K>(a, b, tid)` | `t += a*b`, with `a` H x K and `b` K x W |\n",
    "| `t.Store(baremat, tid)` | write the tile back |\n",
    "\n",
    "This one needs a real GPU: the shuffle primitives it is built on are implemented for Metal\n",
    "only, so the host reference backend cannot compile it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "warp_src = '''\n",
    "using namespace tinybla;\n",
    "\n",
    "KERNEL(tilemul, GLOBAL_IN(float,a), GLOBAL_IN(float,b), GLOBAL(float,c))\n",
    "{\n",
    "  unsigned tid = LOCAL_ID_X;\n",
    "  auto ma = MakeBareMatrix<RowMajor>(a, 16u);   // 8 x 16\n",
    "  auto mb = MakeBareMatrix<RowMajor>(b,  8u);   // 16 x 8\n",
    "  auto mc = MakeBareMatrix<RowMajor>(c,  8u);   // 8 x 8\n",
    "\n",
    "  WarpMatrix<8,8,float> sum = 0.0f;\n",
    "  sum.AddMM<16>(ma, mb, tid);\n",
    "  sum.Store(mc, tid);\n",
    "}\n",
    "'''\n",
    "\n",
    "K = 16\n",
    "A = rng.random((8,K)).astype(np.float32)\n",
    "B = rng.random((K,8)).astype(np.float32)\n",
    "\n",
    "try:\n",
    "    libw = dev.CompileSource(GPUKernelPrelude + TinyBlaPrelude + warp_src)\n",
    "    ba, bb, bc = dev.NewBuffer(4*8*K), dev.NewBuffer(4*K*8), dev.NewBuffer(4*64)\n",
    "    ba.H2D(A.reshape(-1))\n",
    "    bb.H2D(B.reshape(-1))\n",
    "    q.Launch(libw.GetKernel(\"tilemul\"), groups=[1], groupsize=[32], args=[ba, bb, bc])\n",
    "    q.Finish()\n",
    "    print(\"max error vs numpy :\", np.abs(bc.D2H(64).reshape(8,8) - A@B).max())\n",
    "except Exception as e:\n",
    "    print(\"WarpMatrix not available on this backend:\", str(e).splitlines()[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "The kernel sources used here are also available as a standalone script in\n",
    "`ngstd/examples/gpukernel.py`."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}