{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5.6.2 Matrix-free operator application on the GPU\n",
    "\n",
    "For operator application we do not need the assembled sparse matrix: the bilinear form\n",
    "\n",
    "$$ A = B^T D B $$\n",
    "\n",
    "factors into the element-wise differential operator $B$ (universal on the reference element), a point-wise operation $D$ built from the coefficient and the geometry, and the transposed test-function operator $B^T$. The `mf` option of the `BilinearForm` sets up this factorization instead of assembling a matrix, and `CreateDeviceMatrix` compiles one fused GPU kernel for the whole operator - gather, $B$, $D$, $B^T$, scatter.\n",
    "\n",
    "The kernel is written in the common GPU language of [unit 5.6.1](commonGPU.ipynb), so the same code runs on CUDA, Metal and the host reference backend. Block sizes, warps and the launch size are set by `MFOpts`, with defaults that suit both NVIDIA and Apple gpus (the group count adapts to the device's compute units)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve import *\n",
    "from ngsolve.comp import MFOpts\n",
    "from ngsolve.gpu import backend   # registers cuda, metal or the host reference device\n",
    "from time import time\n",
    "print (\"gpu backend:\", backend)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with TaskManager():\n",
    "    mesh = Mesh(unit_cube.GenerateMesh(maxh=0.3))\n",
    "    for l in range(1):     # increase for serious timings\n",
    "        mesh.Refine()\n",
    "    mesh.ngmesh.OrderElements()\n",
    "    mesh = Mesh(mesh.ngmesh)\n",
    "print (\"elements:\", mesh.ne)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Correctness and timing\n",
    "\n",
    "For each space and order we assemble the operator twice: matrix-free (`mf=MFOpts()`) and as a sparse matrix. The matrix-free device operator is applied to a random vector and compared against the sparse product. Then we time\n",
    "\n",
    "* the sparse matrix on the host (with `TaskManager`),\n",
    "* the sparse matrix on the device (cuda only),\n",
    "* the matrix-free device operator.\n",
    "\n",
    "Throughput is reported in GDofs/sec, counting the dofs of the discontinuous version of the space (the actual data volume the operator moves)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Sync (vec):\n",
    "    try: vec.WaitUntilCompleted()     # metal\n",
    "    except AttributeError: s = vec.Norm()   # any backend: forces completion\n",
    "\n",
    "def TimeApply (apply_once, syncvec, runs):\n",
    "    for j in range(min(runs,20)): apply_once()\n",
    "    Sync(syncvec)\n",
    "    ts = time()\n",
    "    for j in range(runs): apply_once()\n",
    "    Sync(syncvec)\n",
    "    return (time()-ts)/runs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Bench (space, order, form, runs=100):\n",
    "    fes = space(mesh, order=order)\n",
    "    u,v = fes.TnT()\n",
    "    equ = form(u,v)\n",
    "    ndof = Discontinuous(fes).ndof\n",
    "\n",
    "    bfmf = BilinearForm(equ, mf=MFOpts()).Assemble()\n",
    "    gpumat = bfmf.mat.CreateDeviceMatrix()\n",
    "    bfsp = BilinearForm(equ).Assemble()\n",
    "\n",
    "    # correctness\n",
    "    gfu = GridFunction(fes)\n",
    "    gfu.vec.SetRandom()\n",
    "    vy = (bfsp.mat * gfu.vec).Evaluate()\n",
    "    xdev = gpumat.CreateRowVector(); xdev.data = gfu.vec\n",
    "    ydev = gpumat.CreateColVector(); ydev.data = gpumat * xdev\n",
    "    yh = bfsp.mat.CreateColVector(); yh.data = ydev\n",
    "    relerr = Norm(vy-yh)/Norm(vy)\n",
    "\n",
    "    # timings\n",
    "    x = gfu.vec.CreateVector(); x.SetRandom()\n",
    "    y = gfu.vec.CreateVector()\n",
    "    def apply_host(): y.data = bfsp.mat * x\n",
    "    with TaskManager():\n",
    "        t_host = TimeApply(apply_host, y, max(runs//5,5))\n",
    "\n",
    "    t_spdev = None\n",
    "    if backend == \"cuda\":\n",
    "        spdev = bfsp.mat.CreateDeviceMatrix()\n",
    "        xs = x.CreateDeviceVector(copy=True)\n",
    "        ys = (spdev*xs).Evaluate()\n",
    "        def apply_spdev(): ys.data = spdev * xs\n",
    "        t_spdev = TimeApply(apply_spdev, ys, runs)\n",
    "\n",
    "    def apply_mf(): ydev.data = gpumat * xdev\n",
    "    t_mf = TimeApply(apply_mf, ydev, runs)\n",
    "\n",
    "    G = lambda t: ndof/t*1e-9 if t else None\n",
    "    return ndof, relerr, G(t_host), G(t_spdev), G(t_mf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def RunTable (space, form, orders=range(1,6)):\n",
    "    print (f\"{space.__name__}:  order      ndof    relerr   sparse-host  sparse-dev   MF-dev   [GDofs/s]\")\n",
    "    for order in orders:\n",
    "        ndof, relerr, gh, gs, gm = Bench(space, order, form)\n",
    "        gs = f\"{gs:9.2f}\" if gs else \"        -\"\n",
    "        print (f\"          {order}  {ndof:9d}   {relerr:.1e}   {gh:9.2f}   {gs}  {gm:8.2f}\")\n",
    "\n",
    "RunTable (H1, lambda u,v: grad(u)*grad(v)*dx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RunTable (L2, lambda u,v: u*v*dx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Measured results\n",
    "\n",
    "Throughput in GDofs/sec on a 3-times refined unit cube (approx. 10^5 elements), 2026-09-01.\n",
    "\n",
    "**Apple M4 Pro (Metal, matrix-free in fp32):**\n",
    "\n",
    "| space | order | sparse host | MF device |\n",
    "|---|---|---|---|\n",
    "| H1 | 1 | 15.2 | 2.7 |\n",
    "| H1 | 2 | 3.2 | 7.4 |\n",
    "| H1 | 3 | 1.2 | 7.8 |\n",
    "| H1 | 4 | 0.58 | 7.8 |\n",
    "| H1 | 5 | 0.29 | 3.9 |\n",
    "| L2 | 1 | 3.4 | 7.0 |\n",
    "| L2 | 2 | 1.4 | 5.9 |\n",
    "| L2 | 3 | 0.75 | 8.3 |\n",
    "| L2 | 4 | 0.44 | 7.0 |\n",
    "| L2 | 5 | 0.30 | 2.2 |\n",
    "\n",
    "**NVIDIA RTX 5090 (CUDA):**\n",
    "\n",
    "| space | order | sparse host | sparse device | MF device fp64 | MF device fp32 |\n",
    "|---|---|---|---|---|---|\n",
    "| H1 | 1 | 4.3 | 16.7 | 11.0 | 26.1 |\n",
    "| H1 | 2 | 0.58 | 25.0 | 11.2 | 37.0 |\n",
    "| H1 | 3 | 0.18 | 11.0 | 5.8 | 29.6 |\n",
    "| H1 | 4 | 0.13 | 5.7 | 4.4 | 22.8 |\n",
    "| H1 | 5 | 0.07 | 3.4 | 2.7 | 15.2 |\n",
    "| L2 | 1 | 1.4 | 17.0 | 15.1 | 30.9 |\n",
    "| L2 | 2 | 0.48 | 10.3 | 13.1 | 35.8 |\n",
    "| L2 | 3 | 0.26 | 6.4 | 12.5 | 42.4 |\n",
    "| L2 | 4 | 0.15 | 3.9 | 7.4 | 28.5 |\n",
    "| L2 | 5 | 0.10 | 2.3 | 2.7 | 9.4 |\n",
    "\n",
    "Observations:\n",
    "\n",
    "* Beyond the lowest order the matrix-free operator clearly beats the assembled matrix - it moves element data instead of the sparse matrix, whose size grows with the square of the local dofs.\n",
    "* On consumer NVIDIA hardware fp64 arithmetic is slow (1/64 of fp32), so there the matrix-free operator wants `MFOpts(fp32=True)`; data-center gpus do not have this restriction. On Metal fp32 is the only choice.\n",
    "* The drop at the highest orders is a shared-memory/occupancy limit of the fused kernel - the block sizes in `MFOpts` are the tuning knobs.\n",
    "* Dof numbering matters for the gather/scatter: `Reorder(fes)` (Morton element order with first-touch dof numbering) gives another 5-15% on the gpu."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.14.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}