notebooks/pytorch_02_custom_operators.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Warp PyTorch Tutorial: Custom Operators\n",
    "\n",
    "In this example we will determine the minimum value of the Rosenbrock function, a non-convex function that is often used to test optimization algorithms. The function is defined as:\n",
    "\n",
    "$$\n",
    "f(x,y) = (a-x)^2 + b(y-x^2)^2\n",
    "$$\n",
    "\n",
    "where a = 1 and b = 100. The minimum value of the function is 0 at (1, 1).\n",
    "\n",
    "We will make use of PyTorch custom operators, which will allow us to safely incorporate Warp kernel launches in a PyTorch graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install warp-lang torch matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warp as wp\n",
    "import numpy as np\n",
    "\n",
    "# Requires PyTorch 2.4+\n",
    "import torch\n",
    "\n",
    "wp.config.quiet = True\n",
    "\n",
    "# Explicitly initializing Warp is not necessary but\n",
    "# we do it here to ensure everything is good to go.\n",
    "wp.init()\n",
    "\n",
    "# Everything else is solely to visualize the results.\n",
    "import IPython\n",
    "import matplotlib\n",
    "import matplotlib.animation\n",
    "import matplotlib.pyplot\n",
    "\n",
    "matplotlib.pyplot.rc(\"animation\", html=\"jshtml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PyTorch Custom Operators\n",
    "\n",
    "PyTorch custom operators allow you to wrap Python functions (in this case, Warp kernel launches) so that they behave like PyTorch native operators. See the [PyTorch docs](https://pytorch.org/tutorials/advanced/python_custom_ops.html#adding-training-support-for-crop) for more information. This is useful when you have a computational graph that PyTorch manages but want to use Warp kernels in one or more nodes. In the following example, we use Warp to evaluate the Rosenbrock function in the forward pass and use PyTorch's Adam optimizer to determine the function's minimum.\n",
    "\n",
    "Below, we define the Warp kernel `eval_rosenbrock` in the usual way, and wrap its forward implementation with the custom PyTorch operator `warp_rosenbrock` as well its adjoint with `warp_rosenbrock_backward`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the Rosenbrock function\n",
    "@wp.func\n",
    "def rosenbrock(x: float, y: float):\n",
    "    return (1.0 - x) ** 2.0 + 100.0 * (y - x**2.0) ** 2.0\n",
    "\n",
    "\n",
    "@wp.kernel\n",
    "def eval_rosenbrock(\n",
    "    xy: wp.array(dtype=wp.vec2),\n",
    "    # outputs\n",
    "    z: wp.array(dtype=float),\n",
    "):\n",
    "    i = wp.tid()\n",
    "    v = xy[i]\n",
    "    z[i] = rosenbrock(v[0], v[1])\n",
    "\n",
    "\n",
    "@torch.library.custom_op(\"wp::warp_rosenbrock\", mutates_args=())\n",
    "def warp_rosenbrock(xy: torch.Tensor, num_particles: int) -> torch.Tensor:\n",
    "    wp_xy = wp.from_torch(xy, dtype=wp.vec2)\n",
    "    wp_z = wp.zeros(num_particles, dtype=wp.float32)\n",
    "\n",
    "    wp.launch(kernel=eval_rosenbrock, dim=num_particles, inputs=[wp_xy], outputs=[wp_z])\n",
    "\n",
    "    return wp.to_torch(wp_z)\n",
    "\n",
    "\n",
    "@warp_rosenbrock.register_fake\n",
    "def _(xy, num_particles):\n",
    "    return torch.empty(num_particles, dtype=torch.float32)\n",
    "\n",
    "\n",
    "@torch.library.custom_op(\"wp::warp_rosenbrock_backward\", mutates_args=())\n",
    "def warp_rosenbrock_backward(\n",
    "    xy: torch.Tensor, num_particles: int, z: torch.Tensor, adj_z: torch.Tensor\n",
    ") -> torch.Tensor:\n",
    "    wp_xy = wp.from_torch(xy, dtype=wp.vec2)\n",
    "    wp_z = wp.from_torch(z, requires_grad=False)\n",
    "    wp_adj_z = wp.from_torch(adj_z, requires_grad=False)\n",
    "\n",
    "    wp.launch(\n",
    "        kernel=eval_rosenbrock,\n",
    "        dim=num_particles,\n",
    "        inputs=[wp_xy],\n",
    "        outputs=[wp_z],\n",
    "        adj_inputs=[wp_xy.grad],\n",
    "        adj_outputs=[wp_adj_z],\n",
    "        adjoint=True,\n",
    "    )\n",
    "\n",
    "    return wp.to_torch(wp_xy.grad)\n",
    "\n",
    "\n",
    "@warp_rosenbrock_backward.register_fake\n",
    "def _(xy, num_particles, z, adj_z):\n",
    "    return torch.empty_like(xy)\n",
    "\n",
    "\n",
    "def backward(ctx, adj_z):\n",
    "    ctx.xy.grad = warp_rosenbrock_backward(ctx.xy, ctx.num_particles, ctx.z, adj_z)\n",
    "    return ctx.xy.grad, None\n",
    "\n",
    "\n",
    "def setup_context(ctx, inputs, output):\n",
    "    ctx.xy, ctx.num_particles = inputs\n",
    "    ctx.z = output\n",
    "\n",
    "\n",
    "warp_rosenbrock.register_autograd(backward, setup_context=setup_context)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Initialization\"\"\"\n",
    "\n",
    "# Number of initial function queries (these should move towards the function's minimum)\n",
    "num_particles = 1500\n",
    "\n",
    "# Initial particle positions\n",
    "rng = np.random.default_rng(42)\n",
    "xy = torch.tensor(\n",
    "    rng.normal(size=(num_particles, 2)), dtype=torch.float32, requires_grad=True, device=wp.device_to_torch(wp.get_device())\n",
    ")\n",
    "\n",
    "# PyTorch Adam optimizer\n",
    "opt = torch.optim.Adam([xy], lr=5e-2)\n",
    "\n",
    "\n",
    "\"\"\"Plotting\"\"\"\n",
    "\n",
    "# Domain\n",
    "min_x, max_x = -2.0, 2.0\n",
    "min_y, max_y = -2.0, 2.0\n",
    "\n",
    "# Create a grid of points\n",
    "x = np.linspace(min_x, max_x, 100)\n",
    "y = np.linspace(min_y, max_y, 100)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "XY = np.column_stack((X.flatten(), Y.flatten()))\n",
    "N = len(XY)\n",
    "\n",
    "XY = wp.array(XY, dtype=wp.vec2)\n",
    "Z = wp.empty(N, dtype=wp.float32)\n",
    "\n",
    "# Evaluate the function over the domain\n",
    "wp.launch(eval_rosenbrock, dim=N, inputs=[XY], outputs=[Z])\n",
    "Z = Z.numpy().reshape(X.shape)\n",
    "\n",
    "# Plot the function as a heatmap\n",
    "fig = matplotlib.pyplot.figure(figsize=(6, 6))\n",
    "ax = matplotlib.pyplot.gca()\n",
    "matplotlib.pyplot.imshow(\n",
    "    Z,\n",
    "    extent=[min_x, max_x, min_y, max_y],\n",
    "    origin=\"lower\",\n",
    "    interpolation=\"bicubic\",\n",
    "    cmap=\"coolwarm\",\n",
    ")\n",
    "matplotlib.pyplot.contour(\n",
    "    X,\n",
    "    Y,\n",
    "    Z,\n",
    "    extent=[min_x, max_x, min_y, max_y],\n",
    "    levels=150,\n",
    "    colors=\"k\",\n",
    "    alpha=0.5,\n",
    "    linewidths=0.5,\n",
    ")\n",
    "\n",
    "# Plot optimum\n",
    "matplotlib.pyplot.plot(1, 1, \"*\", color=\"r\", markersize=10)\n",
    "\n",
    "matplotlib.pyplot.title(\"Rosenbrock function\")\n",
    "matplotlib.pyplot.xlabel(\"x\")\n",
    "matplotlib.pyplot.ylabel(\"y\")\n",
    "\n",
    "(mean_marker,) = ax.plot([], [], \"o\", color=\"w\", markersize=5)\n",
    "\n",
    "# Create a scatter plot (initially empty)\n",
    "scatter_plot = ax.scatter([], [], c=\"k\", s=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimization with PyTorch\n",
    "\n",
    "Here we set up the optimization procedure. `step()` executes a single optimization pass. Notice that in the `forward()` method, we simply need to call our custom operator `warp_rosenbrock()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward():\n",
    "    z = warp_rosenbrock(xy, num_particles)\n",
    "    return z\n",
    "\n",
    "def step():\n",
    "    opt.zero_grad()\n",
    "    z = forward()\n",
    "    z.backward(torch.ones_like(z))\n",
    "    opt.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to update the scatter plot\n",
    "def render():\n",
    "    # Compute mean\n",
    "    xy_np = xy.numpy(force=True)\n",
    "    mean_pos = np.mean(xy_np, axis=0)\n",
    "    \n",
    "    # Update the scatter plot\n",
    "    scatter_plot.set_offsets(np.c_[xy_np[:, 0], xy_np[:, 1]])\n",
    "    mean_marker.set_data([mean_pos[0]], [mean_pos[1]])\n",
    "\n",
    "# Optimize then render\n",
    "def step_and_render(frame):\n",
    "    for _ in range(200):\n",
    "        step()\n",
    "\n",
    "    render()\n",
    "\n",
    "# Create the animation and visualize in Matplotlib\n",
    "plot_anim = matplotlib.animation.FuncAnimation(\n",
    "    fig,\n",
    "    step_and_render,\n",
    "    frames=30,\n",
    "    interval=100)\n",
    "\n",
    "# Display the result\n",
    "IPython.display.display(plot_anim)\n",
    "matplotlib.pyplot.close()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}