{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {
    "deletable": true,
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# 5.6.2 Solving the Poisson Equation on the common GPU layer\n",
    "\n",
    "* After finite element discretization we obtain a (linear) system of equations.\n",
    "* The common GPU layer (`ngsolve.gpu`) moves the linear operators to whichever accelerator this build has: an NVIDIA card through **ngscuda**, an Apple GPU through **ngsmetal**, or the host reference backend if neither is available.\n",
    "* The host is steering, data stays on the device.\n",
    "\n",
    "The kernels behind the device operators are written once in the common syntax of [unit 5.6.1](commonGPU.ipynb) and compiled at run-time for the backend in use. This notebook is the counterpart of [5.5.1](../unit-5.5-cuda/poisson_cuda.ipynb), which uses the CUDA-only classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve import *\n",
    "from ngsolve.webgui import Draw\n",
    "from time import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = Mesh(unit_square.GenerateMesh(maxh=0.1))\n",
    "for l in range(5): mesh.Refine()\n",
    "fes = H1(mesh, order=2, dirichlet=\".*\")\n",
    "print (\"ndof =\", fes.ndof)\n",
    "\n",
    "u, v = fes.TnT()\n",
    "with TaskManager():\n",
    "    a = BilinearForm(grad(u)*grad(v)*dx+u*v*dx).Assemble()\n",
    "    f = LinearForm(x*v*dx).Assemble()\n",
    "\n",
    "gfu = GridFunction(fes)\n",
    "\n",
    "jac = a.mat.CreateSmoother(fes.FreeDofs())\n",
    "\n",
    "with TaskManager():\n",
    "    inv_host = CGSolver(a.mat, jac, maxiter=2000)\n",
    "    ts = time()\n",
    "    gfu.vec.data = inv_host * f.vec\n",
    "    te = time()\n",
    "    print (\"steps =\", inv_host.GetSteps(), \", time =\", te-ts)\n",
    "\n",
    "# Draw (gfu);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "Now we import the common GPU layer. It picks the backend, and tells us which one it took and whether the device computes in double precision. Apple GPUs do not, there the device operators are created in single precision.\n",
    "\n",
    "* NGSolve matrices create their counterparts on the device with `CreateDeviceMatrix`. A sparse matrix becomes a `DeviceSparseMatrix`, a Jacobi smoother a `DeviceDiagonalMatrix`, a block smoother a `DeviceBlockJacobi`, a sparse Cholesky factorization a `DeviceSparseCholesky`. Sums, products and embeddings are translated recursively.\n",
    "* Device operators create device vectors, which live in device memory and are transferred to the host only on demand.\n",
    "* In the following, the conjugate gradients iteration runs on the host, but all operations involving big data are performed on the accelerator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve.gpu import *\n",
    "\n",
    "dev = GetGPUDevice()\n",
    "print (\"backend =\", backend, \", device =\", dev.name, \", fp64 =\", dev.has_float64)\n",
    "ngsglobals.msg_level=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "adev = a.mat.CreateDeviceMatrix()\n",
    "jacdev = jac.CreateDeviceMatrix()\n",
    "print (adev.GetOperatorInfo())\n",
    "print (jacdev.GetOperatorInfo())\n",
    "\n",
    "fdev = adev.CreateColVector()\n",
    "fdev.data = f.vec\n",
    "\n",
    "inv = CGSolver(adev, jacdev, maxiter=2000, printrates=False)\n",
    "\n",
    "ts = time()\n",
    "res = (inv * fdev).Evaluate()\n",
    "te = time()\n",
    "\n",
    "print (\"Time on device:\", te-ts)\n",
    "diff = Norm(gfu.vec - res)\n",
    "print (\"diff = \", diff)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "The difference to the host solution reflects the precision of the device: about $10^{-14}$ in double precision, about $10^{-5}$ relative to the solution in single precision."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "`CGSolver` computes its inner products on the host: every iteration waits for the device and reads two numbers back. `DeviceCGSolver` keeps the coefficients in device scalars and computes them with device kernels, so an iteration is a sequence of kernel launches without synchronisation; the residual is read back only every `check` iterations. On unified-memory machines, where the host is a strong competitor for sparse matrix-vector products, this is what makes the device iteration pay off.\n",
    "\n",
    "With `record=True` the launches of one iteration are recorded once (`ngsolve.gpu.Recording`) and replayed afterwards: a CUDA graph on NVIDIA, a single command buffer on Metal. This removes the per-launch Python and driver cost, which dominates for small problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "inv = DeviceCGSolver(adev, jacdev, maxiter=2000, check=20, record=True, printrates=False)\n",
    "\n",
    "ts = time()\n",
    "res = (inv * fdev).Evaluate()\n",
    "te = time()\n",
    "\n",
    "print (\"Time on device:\", te-ts, \", iterations =\", inv.GetSteps())\n",
    "print (\"diff = \", Norm(gfu.vec - res))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## CG Solver with Block-Jacobi and exact low-order solver:\n",
    "\n",
    "$$\n",
    "A = \\left( \\begin{array}{cc}\n",
    "    A_{cc} & A_{cf} \\\\\n",
    "    A_{fc} & A_{ff}\n",
    "        \\end{array} \\right)\n",
    "$$\n",
    "\n",
    "Additive Schwarz preconditioner:\n",
    "\n",
    "$$\n",
    "C^{-1} = P A_{cc}^{-1} P^T + \\sum_i E_i A_i^{-1} E_i^T\n",
    "$$\n",
    "\n",
    "with\n",
    "* $P$ .. embedding of low-order space\n",
    "* $A_i$ .. blocks on edges/faces/cells\n",
    "* $E_i$ .. embedding matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "fes = H1(mesh, order=5, dirichlet=\".*\")\n",
    "print (\"ndof =\", fes.ndof)\n",
    "\n",
    "u, v = fes.TnT()\n",
    "with TaskManager():\n",
    "    a = BilinearForm(grad(u)*grad(v)*dx+u*v*dx).Assemble()\n",
    "    f = LinearForm(x*v*dx).Assemble()\n",
    "\n",
    "gfu = GridFunction(fes)\n",
    "\n",
    "jac = a.mat.CreateBlockSmoother(fes.CreateSmoothingBlocks())\n",
    "lospace = fes.lospace\n",
    "loinv = a.loform.mat.Inverse(inverse=\"sparsecholesky\", freedofs=lospace.FreeDofs())\n",
    "loemb = fes.loembedding\n",
    "\n",
    "pre = jac + loemb@loinv@loemb.T\n",
    "print (\"mat\", a.mat.GetOperatorInfo())\n",
    "print (\"preconditioner:\")\n",
    "print(pre.GetOperatorInfo())\n",
    "\n",
    "with TaskManager():\n",
    "    inv = CGSolver(a.mat, pre, maxiter=2000, printrates=False)\n",
    "    ts = time()\n",
    "    gfu.vec.data = inv * f.vec\n",
    "    te = time()\n",
    "    print (\"iterations =\", inv.GetSteps(), \"time =\", te-ts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "The whole preconditioner is translated by one call. The block inverses and the Cholesky factorization are computed on the host, the device applies them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "adev = a.mat.CreateDeviceMatrix()\n",
    "predev = pre.CreateDeviceMatrix()\n",
    "print (predev.GetOperatorInfo())\n",
    "\n",
    "fdev = adev.CreateColVector()\n",
    "fdev.data = f.vec\n",
    "\n",
    "inv = CGSolver(adev, predev, maxiter=2000, printrates=False)\n",
    "ts = time()\n",
    "res = (inv * fdev).Evaluate()\n",
    "te = time()\n",
    "print (\"iterations =\", inv.GetSteps(), \"time =\", te-ts)\n",
    "print (\"diff =\", Norm(gfu.vec-res) / Norm(gfu.vec))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Using the BDDC preconditioner:\n",
    "\n",
    "For the BDDC (balancing domain decomposition with constraints) preconditioning, we build a FEM system with relaxed connectivity:\n",
    "\n",
    "<img src=\"../unit-5.5-cuda/pictures/auxspace.png\" width=\"500\">\n",
    "\n",
    "This allows for static condensation of all local and interface degrees of freedom, only the wirebasket dofs enter the global solver. The resulting matrix $\\tilde A$ is much cheaper to invert.\n",
    "\n",
    "The preconditioner is\n",
    "\n",
    "$$\n",
    "P = R {\\tilde A}^{-1} R^T\n",
    "$$\n",
    "\n",
    "with an averaging operator $R$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = Mesh(unit_square.GenerateMesh(maxh=0.1))\n",
    "for l in range(3): mesh.Refine()\n",
    "fes = H1(mesh, order=10, dirichlet=\".*\")\n",
    "print (\"ndof =\", fes.ndof)\n",
    "\n",
    "u, v = fes.TnT()\n",
    "with TaskManager():\n",
    "    a = BilinearForm(grad(u)*grad(v)*dx+u*v*dx)\n",
    "    pre = Preconditioner(a, \"bddc\")\n",
    "    a.Assemble()\n",
    "    f = LinearForm(x*v*dx).Assemble()\n",
    "\n",
    "gfu = GridFunction(fes)\n",
    "with TaskManager():\n",
    "    inv = CGSolver(a.mat, pre, maxiter=2000, printrates=False)\n",
    "    ts = time()\n",
    "    gfu.vec.data = (inv * f.vec).Evaluate()\n",
    "    te = time()\n",
    "    print (\"iterations =\", inv.GetSteps(), \"time =\", te-ts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "predev = pre.mat.CreateDeviceMatrix()\n",
    "print (predev.GetOperatorInfo())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "adev = a.mat.CreateDeviceMatrix()\n",
    "fdev = adev.CreateColVector()\n",
    "fdev.data = f.vec\n",
    "\n",
    "inv = CGSolver(adev, predev, maxiter=2000, printrates=False)\n",
    "ts = time()\n",
    "res = (inv * fdev).Evaluate()\n",
    "te = time()\n",
    "print (\"iterations =\", inv.GetSteps(), \"time =\", te-ts)\n",
    "print (\"diff =\", Norm(gfu.vec-res) / Norm(gfu.vec))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "The same notebook runs on the host reference backend, where the kernels are compiled with the system C++ compiler. This is slow, but a useful check that a device result is not an artifact of the backend:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# SetGPUDevice(GetCPUDevice())\n",
    "# adev = a.mat.CreateDeviceMatrix(); predev = pre.mat.CreateDeviceMatrix()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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"
  },
  "nbsphinx": {
   "execute": "never"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
