Taco Notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to TACO\n",
    "\n",
    "TACO, which stands for Tensor Algebra Compiler, is a library for performing sparse and dense linear algebra and tensor algebra computations. \n",
    "\n",
    "The computations can range from relatively simple ones like sparse matrix-vector multiplication to more complex ones like matricized tensor times Khatri-Rao product. All these computations can be performed on any mix of dense and sparse tensors. Under the hood, TACO automatically generates efficient code to perform these computations.\n",
    "\n",
    "This notebook provides a brief introduction to the TACO Python Library. For a more comprehensive overview, please see the documentation linked [here](http://tensor-compiler.org/symphony-docs/index.html). We will also link to relevant pages as we progress."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Table of Contents:\n",
    "* [Getting Started](#first-section)\n",
    "* [Defining Tensor Formats](#second-section)\n",
    "* [NumPy and SciPy I/O: SpMV](#third-section)\n",
    "* [Example Application: SpMV](#fourth-section)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting Started <a class=\"anchor\" id=\"first-section\"></a>\n",
    "\n",
    "First, let's import TACO. Press Shift + Enter to run the code below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytaco as pt\n",
    "from pytaco import dense, compressed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the above, `dense` and `compressed` are [mode (dimension) formats](http://tensor-compiler.org/symphony-docs/reference/rst_files/mode_format.html). We can think of tensors as a multi-dimensional array, and the mode formats allow us to specify how we would like to store the data in each dimension: \n",
    "* If a dimension is `dense`, then all of the elements in that dimension are stored. \n",
    "* And if a dimension is `compressed`, then only nonzeros are stored.\n",
    "\n",
    "For example, we can declare a $512 \\times 64 \\times 2048$ [tensor](http://tensor-compiler.org/symphony-docs/reference/rst_files/tensor_class.html) whose first dimension is dense and second and third dimensions are compressed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = pt.tensor([512, 64, 2048], pt.format([dense, compressed, compressed]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can initialize $T$ by calling its `insert` [method](http://tensor-compiler.org/symphony-docs/reference/rst_files/functions/pytaco.tensor.insert.html) to add a nonzero element to the tensor. The `insert` method takes two arguments: a list of coordinates and the value to be inserted at that coordinate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set T(0, 1, 0) = 42.0\n",
    "T.insert([0, 1, 0], 42.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If multiple elements are inserted at the same coordinates, they are summed together:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set T(0, 0, 1) = 12.0 + 24.0 = 36.0\n",
    "T.insert([0, 0, 1], 12.0)\n",
    "T.insert([0, 0, 1], 24.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then iterate over the nonzero elements of the tensor as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for coordinate, value in T: \n",
    "    print(\"Coordinate: {}, Value: {}\".format(coordinate, value))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining Tensor Formats <a class=\"anchor\" id=\"second-section\"></a>\n",
    "\n",
    "Consider a matrix $M$ (aka a two-dimensional tensor) containing the following values:\n",
    "\n",
    "$$M = \\begin{bmatrix}\n",
    "    6 & \\cdot & 9 & 8 \\\\\n",
    "    \\cdot & \\cdot & \\cdot & \\cdot \\\\\n",
    "    5 & \\cdot & \\cdot & 7  \n",
    "\\end{bmatrix}$$\n",
    "\n",
    "Denote the rows and columns as dimensions $d_1$ and $d_2$, respectively. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We look at how $M$ is represented differently in different formats. For convenience, let's define a helper function to initialize $M$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_example_matrix(format):\n",
    "    M = pt.tensor([3, 4], format)\n",
    "    M.insert([0, 0], 6)\n",
    "    M.insert([0, 2], 9)\n",
    "    M.insert([0, 3], 8)\n",
    "    M.insert([2, 0], 5)\n",
    "    M.insert([2, 3], 7)\n",
    "    return M"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (dense $d_1$, dense $d_2$)\n",
    "Note that passing in `dense` makes all of the dimensions dense. This is equivalent to `pt.format(dense, dense)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_example_matrix(dense)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this example, we focus on the last line of the output, the `vals` array: since all values are stored, it is simply a flattened $3 \\times 4$ matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (dense $d_1$, compressed $d_2$)\n",
    "This is called compressed sparse row (CSR) format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "csr = pt.format([dense, compressed])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "make_example_matrix(csr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since $d_1$ is dense, we need only store the `size` of the dimension; values, both zero and nonzero, are stored for every coordinate in dimension $d_1$.  \n",
    "\n",
    "Since $d_2$ is compressed, we store a `pos` array and an `idx` array; these together form a segmented vector with one segment per entry in the previous dimension. The `idx` array stores all the indices with nonzero values in the dimension, while the `pos` array stores the location in the `idx` array where each segment begins. In particular, segment $i$ is stored in locations `pos[i]:pos[i+1]` in the `idx` array.\n",
    "\n",
    "The below animation visualizes the format. Hover over any non-empty entry of the matrix on the left to see how the value is stored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run ./animation/animation_2.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (dense $d_2$, compressed $d_1$)\n",
    "\n",
    "We switch the order of the dimensions by passing in `[1, 0]` for the optional parameter `mode_ordering`. This results in a column-major (rather than row-major) format called compressed sparse column (CSC)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "csc = pt.format([dense, compressed], [1, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "make_example_matrix(csc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this format, $d_2$ has only `size`, while $d_1$ has a `pos` and `idx` array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%run ./animation/animation_3.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (compressed $d_1$, compressed $d_2$)\n",
    "\n",
    "This is called doubly compressed sparse row (DCSR) format.\n",
    "\n",
    "Similarly as above, passing in `compressed` is equivalent to `pt.format(compressed, compressed)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_example_matrix(compressed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run ./animation/animation_4.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NumPy and SciPy I/O <a class=\"anchor\" id=\"third-section\"></a>\n",
    "\n",
    "We can also initialize tensors with NumPy arrays or SciPy sparse (CSR or CSC) matrices. Let's start by importing the packages we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given a NumPy array such as the one randomly generated below, we can convert it to a TACO tensor using the `pytaco.from_array` [function](http://tensor-compiler.org/symphony-docs/reference/rst_files/functions/pytaco.from_array.html), which creates a dense tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array = np.random.uniform(size=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "pt.from_array(array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, given a SciPy sparse matrix, we can convert it to a TACO tensor. For a CSR matrix like the one below, we use the `pytaco.from_sp_csr` [function](http://tensor-compiler.org/symphony-docs/reference/rst_files/functions/pytaco.from_sp_csr.html), which creates a CSR tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix = sp.sparse.rand(m=10, n=10, density=0.1, format = 'csr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt.from_sp_csr(matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example Application: SpMV <a class=\"anchor\" id=\"fourth-section\"></a>\n",
    "\n",
    "The following example demonstrates how computations on tensors are performed using TACO."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sparse matrix-vector multiplication (SpMV) is a bottleneck computation in many scientific and engineering computations. Mathematically, SpMV can be expressed as $$y = Ax + z,$$ \n",
    "\n",
    "where $A$ is a sparse matrix and $x$, $y$, and $z$ are dense vectors. The computation can also be expressed in [index notation](http://tensor-compiler.org/symphony-docs/pycomputations/index.html#specifying-tensor-algebra-computations) as \n",
    "\n",
    "$$y_i = A_{ij} \\cdot x_j + z_i.$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We generate $A$ randomly with SciPy. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 100\n",
    "density = 0.1\n",
    "sparse_matrix = sp.sparse.rand(size, size, density, format = 'csr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = pt.from_sp_csr(sparse_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the dimensions of $A$ using its `shape` attribute, which we expect to be `[size, size]`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Examining the formula above, we need to define a vector $x$ whose length is the number of columns of $A$, and a vector $z$ whose length is the number of rows. We generate $x$ and $z$ randomly with NumPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "x = pt.from_array(np.random.uniform(size=A.shape[1]))\n",
    "z = pt.from_array(np.random.uniform(size=A.shape[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Expressing the Computation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can express the result $y$ as a dense vector. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y = pt.tensor([A.shape[0]], dense)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The syntax for TACO computations closely mirrors index notation, with the caveat that we also have to explicitly declare the index variables beforehand:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i, j = pt.get_index_vars(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y[i] = A[i, j] * x[j] + z[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performing the Computation\n",
    "\n",
    "Once a tensor algebra computation has been defined, we can simply invoke the result tensor's `evaluate` method to perform the actual computation.\n",
    "\n",
    "[Under the hood](http://tensor-compiler.org/symphony-docs/pycomputations/index.html#performing-the-computation), TACO will first invoke the result tensor's `compile` method to generate code that performs the computation. TACO will then perform the actual computation by first invoking `assemble` to compute the sparsity structure of the result and subsequently invoking `compute` to compute the values of the result's nonzero elements.\n",
    "\n",
    "We can manually invoke these methods as well. In order to accurately measure TACO's computational performance, only the time it takes to actually perform a computation should be measured. Below, we compile $y$, and then benchmark the computation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "y.assemble()\n",
    "y.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we define a computation and then access the result without first manually invoking `evaluate` or `compile`/`assemble`/`compute`, TACO will automatically invoke the computation immediately before the result is accessed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can write the result to a file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pt.write(\"y.tns\", y)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}