notebooks/matmul/01_introduction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "88073684-ba4e-42eb-9d9e-f7541473ce4f",
   "metadata": {},
   "source": [
    "Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES\n",
    "\n",
    "SPDX-License-Identifier: BSD-3-Clause"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1388882b-0c36-4ac2-ac0b-9e909f543605",
   "metadata": {},
   "source": [
    "# Introduction to GEMM with nvmath-python"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "292c2024-bd73-48de-8320-62c0ec0df645",
   "metadata": {},
   "source": [
    "In this tutorial we will demonstrate how to perform GEMM (General Matrix Multiply) with nvmath-python library.\n",
    "\n",
    "We will demonstrate two APIs to execute matrix multiplication with nvmath-python:\n",
    "- `matmul` function (*stateless API*), which performs a single GEMM on its arguments and returns the result.\n",
    "- `Matmul` class (*stateful API*), which can perform multiple GEMMs on different input data, allowing to amortize the cost of initialization and planning."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f8a78ce-ae83-4a25-bd21-ebffdb9d1edd",
   "metadata": {},
   "source": [
    "## Stateless API"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e5782b5-17c3-42e7-8c9f-234fd2ef3f8c",
   "metadata": {},
   "source": [
    "Let us demonstrate the usage of `matmul` function. We will perform the computations on CuPy arrays, but nvmath-python supports NumPy arrays and PyTorch tensors as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "59c42527-6c4a-4359-beae-26ccc9148696",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy\n",
    "from nvmath.linalg.advanced import matmul"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bcffe5e-f125-421d-bf28-f34bd62f003b",
   "metadata": {},
   "source": [
    "### Basic matrix multiplication"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6177b46b-ab4d-4271-91f1-9737bdb4e35e",
   "metadata": {},
   "source": [
    "In its most basic use-case, `matmul` can be used to just multiply two matrices:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "58db13d3-5566-4334-bd18-d93bff6d442f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 4.  7.]\n",
      " [ 8. 15.]]\n",
      "Result type: <class 'cupy.ndarray'>\n"
     ]
    }
   ],
   "source": [
    "a = cupy.asarray([[1, 2], [3, 4]], dtype=cupy.float32)\n",
    "b = cupy.asarray([[0, 1], [2, 3]], dtype=cupy.float32)\n",
    "result = matmul(a, b)\n",
    "print(result)\n",
    "print(\"Result type:\", type(result))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bdf9a94-b9c8-46ef-b999-b87b14f79aed",
   "metadata": {},
   "source": [
    "Note that the result is CuPy array as well. Nvmath-python always returns the result of the same type as its inputs: if the inputs were PyTorch tensors, the result will be a PyTorch tensor as well."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fe690e2-acb9-408d-ae53-e630a3e3d98b",
   "metadata": {},
   "source": [
    "### GEMM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68449c41-0b3e-4169-895a-c9f021c4ec30",
   "metadata": {},
   "source": [
    "We can as well use it to perform a GEMM, which is defined as: $$\\alpha A@B + \\beta C$$ for matrices $A$, $B$ and $C$ and constants $\\alpha$ and $\\beta$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ea8079de-8ed5-40a3-acdf-0654e2e0998e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[402.  503.5]\n",
      " [604.  707.5]\n",
      " [806.  911.5]]\n"
     ]
    }
   ],
   "source": [
    "a = cupy.asarray([[1, 2], [3, 4], [5, 6]], dtype=cupy.float32)\n",
    "b = cupy.asarray([[0, 1], [2, 3]], dtype=cupy.float32)\n",
    "c = cupy.asarray([[4, 5], [6, 7], [8, 9]], dtype=cupy.float32)\n",
    "\n",
    "result = matmul(a, b, c=c, alpha=0.5, beta=100)\n",
    "print(result)\n",
    "\n",
    "# Check with cupy\n",
    "assert cupy.allclose(result, 0.5 * a @ b + 100 * c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1f7c5f2-dfdb-4102-a065-f04b2891c7ac",
   "metadata": {},
   "source": [
    "## Stateful API"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8528ccf-aed2-486c-9a5a-536eeb7295ab",
   "metadata": {},
   "source": [
    "The stateless `matmul` function, which we demonstrated above, plans and executes the matrix multiplication at once.\n",
    "\n",
    "The *stateful API* allows you to first create a `Matmul` object, then plan the multiplication, optionally fine-tune it and then execute (possibly multiple times). This stateful API is recommended for use cases where similar multiplication (with same shapes and types of tensors, but with different numbers inside) is performed many times. This amortizes the planning cost and results in significant reduction of the overhead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4b40c4ff-b45e-4ee9-98f9-151b94903e1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nvmath.linalg.advanced import Matmul"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b353284c-c820-457d-984c-dbb2bc008bb2",
   "metadata": {},
   "source": [
    "### Basic usage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa8ed9fa-f693-42a0-92a6-5f74cc98a75f",
   "metadata": {},
   "source": [
    "Let us show an example of how `Matmul` object can be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "40f722c6-2d4d-41fc-bfcf-b60ca4f42c00",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = cupy.asarray([[1, 2], [3, -4]], dtype=cupy.float32)\n",
    "b = cupy.asarray([[5, -6], [7, -8]], dtype=cupy.float32)\n",
    "\n",
    "mm = Matmul(a, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc9a15d2-bbab-4dab-b872-b38d197fb0db",
   "metadata": {},
   "source": [
    "With our `Matmul` object created, let us execute the planning. In this phase the underlying cuBLAS library will choose the best algorithms based on input shapes and types. The proposed algorithms will be stored in `mm.algorithms` and ordered starting from the best one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c9227a3e-d0b0-4d0c-a9de-3ee844029131",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7 algorithms were proposed in the planning phase.\n"
     ]
    }
   ],
   "source": [
    "mm.plan()\n",
    "print(f\"{len(mm.algorithms)} algorithms were proposed in the planning phase.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2ade4bb-17b3-4828-9d0a-84f23c2838b0",
   "metadata": {},
   "source": [
    "Optionally, you can *autotune* the multiplications. Autotuning runs benchmarks on each of the algorithms in `mm.algorithms` and shuffles the list according to the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0f9c9be7-dd77-4956-acc2-54900e5a5e91",
   "metadata": {},
   "outputs": [],
   "source": [
    "mm.autotune()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a68b1987-1e6b-4e2c-ab13-15e8d6d9d323",
   "metadata": {},
   "source": [
    "With the planning complete, you can execute the multiplication:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f6404c21-6934-44ba-847a-adc0d52deab7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 19. -22.]\n",
      " [-13.  14.]]\n"
     ]
    }
   ],
   "source": [
    "result1 = mm.execute()\n",
    "print(result1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18c8ced6-7806-416f-935f-1394650b8029",
   "metadata": {},
   "source": [
    "### Resetting the operands"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25774b98-6073-4797-8918-78a9a210be69",
   "metadata": {},
   "source": [
    "To perform the next multiplication, you can use `mm.reset_operands` to change all or some of the inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ea370d91-2eb3-4b40-8c37-cbe7973404be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 7. -8.]\n",
      " [ 0.  0.]]\n"
     ]
    }
   ],
   "source": [
    "mm.reset_operands(a=cupy.asarray([[0, 1], [0, 0]], dtype=cupy.float32))\n",
    "result2 = mm.execute()\n",
    "print(result2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b735435d-954e-424d-a20e-1f71ea81650a",
   "metadata": {},
   "source": [
    "The new inputs must be of the same shape and type as the original. Otherwise, an error will be raised:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1dc8ed7f-d3e8-4a8d-baf1-25cd62c03f43",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: The extents of the new operand must match the extents of the original operand.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    mm.reset_operands(a=cupy.asarray([[7]], dtype=cupy.float32))  # This is 1x1 matrix!\n",
    "except ValueError as e:\n",
    "    print(\"Error:\", e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9edf8279-349c-4d5b-8329-d3dff7e7325e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: The data type of the new operand must match the data type of the original operand.\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    mm.reset_operands(a=cupy.asarray([[0, 1], [0, 0]], dtype=cupy.float64))  # This is float64!\n",
    "except ValueError as e:\n",
    "    print(\"Error:\", e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5ee6f18-6560-42f2-a76a-b3de3487df96",
   "metadata": {},
   "source": [
    "### Managing the resources"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55d31e97-5dcb-4ea4-95ce-a10674839c87",
   "metadata": {},
   "source": [
    "Finally, we should release the resources of `mm`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f25815b0-5fa0-4ffd-84ca-38364791a219",
   "metadata": {},
   "outputs": [],
   "source": [
    "mm.free()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cd0b9a5-8899-46ce-bd37-3b4d042a746c",
   "metadata": {},
   "source": [
    "Alternatively, we can use `with` to manage the resources automatically:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2444b8d9-473c-40cf-ba0c-10002598d57e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 45. -37.]\n",
      " [-12.  49.]]\n"
     ]
    }
   ],
   "source": [
    "a = cupy.asarray([[1, 2], [3, -4]], dtype=cupy.float32)\n",
    "b = cupy.asarray([[5, -6], [7, -8]], dtype=cupy.float32)\n",
    "c = cupy.asarray([[10, 10], [20, 30]], dtype=cupy.float32)\n",
    "\n",
    "with Matmul(a, b, c=c, alpha=2, beta=0.7) as mm:\n",
    "    mm.plan()\n",
    "    result = mm.execute()\n",
    "    # mm.free() is no longer needed, the resources are freed by the context manager.\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c63568c-00fe-4394-90ed-80d2ab3fc04b",
   "metadata": {},
   "source": [
    "In fact, `matmul` is just a thin wrapper around `Matmul`, which under the hood creates a `Matmul` object, and then calls `.plan` and `.execute` just as we did above."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c760afa2-a3fd-4f89-805e-309dccc8a229",
   "metadata": {},
   "source": [
    "## Learning more"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b881243-b501-4edd-b823-2ddfde138003",
   "metadata": {},
   "source": [
    "To learn more, we encourage you to visit our [documentation pages for linear algebra](https://docs.nvidia.com/cuda/nvmath-python/latest/host-apis/linalg/index.html#specialized-linear-algebra-apis-nvmath-linalg-advanced)."
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}