notebooks/tutorial_ivf_pq.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# cuVS IVF-PQ tutorial\n",
    "In this tutorial you will learn to build IVF-PQ index and use it to search approximate nearest neighbors (ANN).\n",
    "We will start with a brief overview of the functionality, but then dive into details to gain the understanding of the model parameters.\n",
    "Along the way, we will benchmark the model and give some practical recommendations on how to maximize its performance for various use cases.\n",
    "\n",
    "This tutorial uses the data from [ANN benchmarks website](https://ann-benchmarks.com)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install adjustText h5py matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tempfile\n",
    "import cupy as cp\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import rmm\n",
    "import urllib.request\n",
    "import h5py\n",
    "\n",
    "from rmm.allocators.cupy import rmm_cupy_allocator\n",
    "from cuvs.common import Resources\n",
    "from cuvs.neighbors import ivf_pq, refine\n",
    "from adjustText import adjust_text\n",
    "from utils import calc_recall, load_dataset\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A clumsy helper for inspecting properties of an object\n",
    "def show_properties(obj):\n",
    "    return {\n",
    "        attr: getattr(obj, attr)\n",
    "        for attr in dir(obj)\n",
    "        if type(getattr(type(obj), attr)).__name__ == 'getset_descriptor'\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll need to load store some data in this tutorial\n",
    "WORK_FOLDER = os.path.join(tempfile.gettempdir(), 'cuvs_ivf_pq_tutorial')\n",
    "\n",
    "if not os.path.exists(WORK_FOLDER):\n",
    "   os.makedirs(WORK_FOLDER)\n",
    "print(\"The index and data will be saved in\", WORK_FOLDER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Report the GPU in use to put the measurements into perspective\n",
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use the pool memory resource\n",
    "cuVS uses RMM allocator widely across its algorithms, including the performance-sensitive parts like IVF-PQ search.\n",
    "It's strongly advised to set up the RMM pool memory resource to minimize the overheads of repeated CUDA allocations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pool = rmm.mr.PoolMemoryResource(\n",
    "    rmm.mr.CudaMemoryResource(),\n",
    "    initial_pool_size=2**30\n",
    ")\n",
    "rmm.mr.set_current_device_resource(pool)\n",
    "cp.cuda.set_allocator(rmm_cupy_allocator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get the data\n",
    "The [ANN benchmarks website](https://ann-benchmarks.com) provides the datasets in [HDF5 format](https://www.hdfgroup.org/solutions/hdf5/).\n",
    "\n",
    "The list of prepared datasets can be found at https://github.com/erikbern/ann-benchmarks/#data-sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_URL = \"http://ann-benchmarks.com/sift-128-euclidean.hdf5\"\n",
    "DATASET_NAME = \"SIFT-128\"\n",
    "f = load_dataset(DATASET_URL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = f.attrs['distance']\n",
    "\n",
    "dataset = cp.array(f['train'])\n",
    "queries = cp.array(f['test'])\n",
    "gt_neighbors = cp.array(f['neighbors'])\n",
    "gt_distances = cp.array(f['distances'])\n",
    "\n",
    "print(f\"Loaded dataset of size {dataset.shape}; metric: '{metric}'.\")\n",
    "print(f\"Number of test queries: {queries.shape[0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the index\n",
    "Construction of the index generally consists of two phases: training (building the clusters) and filling-in (extending the index with data).\n",
    "In the first phase, a balanced hierarchical k-means algorithm clusters the training data.\n",
    "In the second phase, the new data is classified and added into the appropriate clusters in the index.\n",
    "Hence, a user should call `ivf_pq.build` once and then possibly `ivf_pq.extend` several times.\n",
    "Though for user convenience `ivf_pq.build` by default adds the whole training set into the index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cuVS' Resources controls the GPU, cuda stream, memory policies etc.\n",
    "# For now, we just create a default instance.\n",
    "resources = Resources()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, we need to initialize the build/indexing parameters.\n",
    "# One of the more important parameters is the product quantisation (PQ) dim.\n",
    "# Effectively, this parameter says\n",
    "#      \"shrink the dataset to this dimensionality to reduce the index size\".\n",
    "# It must be not bigger than the dataset dim,\n",
    "# and it should be divisible by 32 for better GPU performance.\n",
    "pq_dim = 1\n",
    "while pq_dim * 2 < dataset.shape[1]:\n",
    "    pq_dim = pq_dim * 2\n",
    "# We'll use the ANN-benchmarks-provided metric and sensible defaults for the rest of parameters.\n",
    "index_params = ivf_pq.IndexParams(n_lists=1024, metric=metric, pq_dim=pq_dim)\n",
    "\n",
    "show_properties(index_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "## Build the index\n",
    "# This function takes a row-major either numpy or cupy (GPU) array.\n",
    "# Generally, it's a bit faster with GPU inputs, but the CPU version may come in handy\n",
    "# if the whole dataset cannot fit into GPU memory.\n",
    "index = ivf_pq.build(index_params, dataset, resources=resources)\n",
    "# This function is asynchronous so we need to explicitly synchronize the GPU before we can measure the execution time\n",
    "resources.sync()\n",
    "index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Index serialization\n",
    "For bigger datasets, building an index can take some time. To avoid building the index from scratch every time you need it, you can save it to a file. Here is how this works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "index_filepath = os.path.join(WORK_FOLDER, \"ivf_pq.bin\")\n",
    "ivf_pq.save(index_filepath, index) \n",
    "loaded_index = ivf_pq.load(index_filepath)\n",
    "resources.sync()\n",
    "index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Search\n",
    "The search function returns the requested number `k` of (approximate) nearest neighbor in no particular order.\n",
    "Besides the queries and `k`, the function can take a few more parameters to tweak the performance of the algorithm.\n",
    "Again, these are passed via the struct with some sensible defaults."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 10\n",
    "search_params = ivf_pq.SearchParams()\n",
    "show_properties(search_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "distances, neighbors = ivf_pq.search(search_params, index, queries, k, resources=resources)\n",
    "# Sync the GPU to make sure we've got the timing right\n",
    "resources.sync()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Measuring the quality of the predictions\n",
    "We use [recall](https://en.wikipedia.org/wiki/Precision_and_recall) to measure the quality of the prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recall_first_try = calc_recall(neighbors, gt_neighbors)\n",
    "print(f\"Got recall = {recall_first_try} with the default parameters (k = {k}).\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Refine\n",
    "Let's improve our results a little bit!\n",
    "The refinement operation follows an approximate NN search.\n",
    "It recomputes the exact distances for the already selected candidates and selects a subset of them thus improving the recall."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "candidates = ivf_pq.search(search_params, index, queries, k * 2, resources=resources)[1]\n",
    "distances, neighbors = refine(dataset, queries, candidates, k, resources=resources)\n",
    "resources.sync()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recall_refine2x = calc_recall(neighbors, gt_neighbors)\n",
    "print(f\"Got recall = {recall_refine2x} with 2x refinement (k = {k}).\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tweaking search parameters\n",
    "Before diving deep into tweaking the model, let's quickly define the performance metrics.\n",
    "As we've mentioned earlier, we use the recall to measure the quality of prediction.\n",
    "The other important metric is the speed of the search.\n",
    "We measure the speed in terms of queries per second (QPS).\n",
    "\n",
    "Most of the time, by changing the model parameters we balance the trade-off between the QPS and the recall."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Number of neighbors\n",
    "Let's see how QPS depens on `k`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench_k = np.exp2(np.arange(10)).astype(np.int32)\n",
    "bench_avg = np.zeros_like(bench_k, dtype=np.float32)\n",
    "bench_std = np.zeros_like(bench_k, dtype=np.float32)\n",
    "for i, k in enumerate(bench_k):\n",
    "    r = %timeit -o ivf_pq.search(search_params, index, queries, k, resources=resources); resources.sync()\n",
    "    bench_avg[i] = (queries.shape[0] * r.loops / np.array(r.all_runs)).mean()\n",
    "    bench_std[i] = (queries.shape[0] * r.loops / np.array(r.all_runs)).std()\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=plt.figaspect(1/2))\n",
    "ax.errorbar(bench_k, bench_avg, bench_std)\n",
    "ax.set_xscale('log')\n",
    "ax.set_xticks(bench_k, bench_k)\n",
    "ax.set_xlabel('k')\n",
    "ax.grid()\n",
    "ax.set_ylabel('QPS');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Number of probes\n",
    "IVF-PQ search runs in two phases; first it looks for nearest clusters,\n",
    "then it searches for the neighbors in every selected cluster.\n",
    "\n",
    "We can set how many clusters we want to inspect.\n",
    "For this, `ivf_pq.SearchParams` has a parameter `n_probes`.\n",
    "This is the core parameter to control the QPS/recall trade-off."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench_probes = np.exp2(np.arange(11)).astype(np.int32)\n",
    "bench_qps = np.zeros_like(bench_probes, dtype=np.float32)\n",
    "bench_recall = np.zeros_like(bench_probes, dtype=np.float32)\n",
    "k = 100\n",
    "for i, n_probes in enumerate(bench_probes):\n",
    "    sp = ivf_pq.SearchParams(n_probes=n_probes)\n",
    "    r = %timeit -o ivf_pq.search(sp, index, queries, k, resources=resources); resources.sync()\n",
    "    bench_qps[i] = (queries.shape[0] * r.loops / np.array(r.all_runs)).mean()\n",
    "    bench_recall[i] = calc_recall(ivf_pq.search(sp, index, queries, k, resources=resources)[1], gt_neighbors)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's clear that the search time scales almost linearly with the number of probes.\n",
    "This is due to the algorithm spending most of the time in the second phase scanning through individual clusters.\n",
    "Thanks to the balanced nature of the clustering k-means algorithm, the sizes of the clusters are roughly similar;\n",
    "hence the linear relation `n_probes` ~ query time.\n",
    "\n",
    "Let's draw some plots to illustrate how the number of probes affects QPS and recall."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 3, figsize=plt.figaspect(1/4))\n",
    "\n",
    "ax[0].plot(bench_probes, bench_recall)\n",
    "ax[0].set_xscale('log')\n",
    "ax[0].set_xticks(bench_probes, bench_probes)\n",
    "ax[0].set_xlabel('n_probes')\n",
    "ax[0].set_ylabel('recall')\n",
    "ax[0].grid()\n",
    "\n",
    "ax[1].plot(bench_probes, bench_qps)\n",
    "ax[1].set_xscale('log')\n",
    "ax[1].set_xticks(bench_probes, bench_probes)\n",
    "ax[1].set_xlabel('n_probes')\n",
    "ax[1].set_ylabel('QPS')\n",
    "ax[1].set_yscale('log')\n",
    "ax[1].grid()\n",
    "\n",
    "ax[2].plot(bench_recall, bench_qps)\n",
    "ax[2].set_xlabel('recall')\n",
    "ax[2].set_ylabel('QPS')\n",
    "ax[2].set_yscale('log')\n",
    "ax[2].grid();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Internal search types\n",
    "Besides `n_probes`, `ivf_pq.SearchParams` contains a couple more parameters, which affect the internal workings of the algorithm.\n",
    "\n",
    "`internal_distance_dtype` controls the representation of the distance/similarity during the search.\n",
    "By default, it's `np.float32`, but you can change it to `np.float16` when appropriate to save the memory bandwidth.\n",
    "This can be a good idea when the dataset type is low precision anyway (e.g. `np.uint8`),\n",
    "yet it may help with 32-bit float datasets too.\n",
    "\n",
    "`lut_dtype` is the Look-Up Table Data Type.\n",
    "The specifics of the PQ algorithm is that it stores the data in the Product Quantizer (PQ) encoded format,\n",
    "which needs to be decoded during the second-phase (in-cluster) search.\n",
    "Thus, the algorithm constructs a lookup table for each cluster.\n",
    "This is a costly operation, and the table itself can be rather large.\n",
    "By default, the individual elements in the table are stored as 32-bit floats,\n",
    "but you can change this to `np.float16` or `np.uint8` to reduce the table size.\n",
    "\n",
    "The exact size of the table is as follows:\n",
    "\n",
    "$ \\mathtt{lut\\_size} = \\mathtt{pq\\_dim} \\cdot \\mathtt{sizeof(lut\\_dtype) \\cdot 2^{\\mathtt{pq\\_bits}}} $\n",
    "\n",
    "Ideally, the lookup table should fit in the shared memory of a GPU's multiprocessor,\n",
    "but it's not the case for wider datasets.\n",
    "The logic of deciding whether this table should stay in the shared or the global memory of the GPU is somewhat complicated.\n",
    "Yet, you can see the outcome when you gradually change `pq_dim` and observe a sudden drop in QPS after a certain threshold.\n",
    "The shared-memory kernel version is typically 2-5x faster than the global-memory version.\n",
    "\n",
    "However `pq_dim` strongly affects the recall and requires the index to be re-build on change.\n",
    "This is where `lut_dtype` comes in handy: you can halve or quarter the lookup table size by changing it.\n",
    "Though it does affect the recall too.\n",
    "\n",
    "Also note, it does not make sense to set the `lut_dtype` to a more precise type than `internal_distance_dtype`,\n",
    "as the former is converted to the latter internally.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bench_qps_s1 = np.zeros((5,), dtype=np.float32)\n",
    "bench_recall_s1 = np.zeros((5,), dtype=np.float32)\n",
    "k = 10\n",
    "n_probes = 256\n",
    "search_params_32_32 = ivf_pq.SearchParams(n_probes=n_probes, internal_distance_dtype=np.float32, lut_dtype=np.float32)\n",
    "search_params_32_16 = ivf_pq.SearchParams(n_probes=n_probes, internal_distance_dtype=np.float32, lut_dtype=np.float16)\n",
    "search_params_32_08 = ivf_pq.SearchParams(n_probes=n_probes, internal_distance_dtype=np.float32, lut_dtype=np.uint8)\n",
    "search_params_16_16 = ivf_pq.SearchParams(n_probes=n_probes, internal_distance_dtype=np.float16, lut_dtype=np.float16)\n",
    "search_params_16_08 = ivf_pq.SearchParams(n_probes=n_probes, internal_distance_dtype=np.float16, lut_dtype=np.uint8)\n",
    "search_ps = [search_params_32_32, search_params_32_16, search_params_32_08, search_params_16_16, search_params_16_08]\n",
    "bench_names = ['32/32', '32/16', '32/8', '16/16', '16/8']\n",
    "\n",
    "for i, sp in enumerate(search_ps):\n",
    "    r = %timeit -o ivf_pq.search(sp, index, queries, k, resources=resources); resources.sync()\n",
    "    bench_qps_s1[i] = (queries.shape[0] * r.loops / np.array(r.all_runs)).mean()\n",
    "    bench_recall_s1[i] = calc_recall(ivf_pq.search(sp, index, queries, k, resources=resources)[1], gt_neighbors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=plt.figaspect(1/2))\n",
    "fig.suptitle(\n",
    "    f'Effects of search parameters on QPS/recall trade-off ({DATASET_NAME})\\n' + \\\n",
    "    f'k = {k}, n_probes = {n_probes}, pq_dim = {pq_dim}')\n",
    "ax.plot(bench_recall_s1, bench_qps_s1, 'o')\n",
    "ax.set_xlabel('recall')\n",
    "ax.set_ylabel('QPS')\n",
    "ax.grid()\n",
    "annotations = []\n",
    "for i, label in enumerate(bench_names):\n",
    "    annotations.append(ax.text(\n",
    "        bench_recall_s1[i], bench_qps_s1[i],\n",
    "        f\" {label} \",\n",
    "        ha='center', va='center'))\n",
    "clutter = [\n",
    "    ax.text(\n",
    "        0.02, 0.08,\n",
    "        'Labels denote the bitsize of: internal_distance_dtype/lut_dtype',\n",
    "        verticalalignment='top',\n",
    "        bbox={'facecolor': 'white', 'edgecolor': 'grey'},\n",
    "        transform = ax.transAxes)\n",
    "]\n",
    "adjust_text(annotations, objects=clutter);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This figure represents the trade-offs one does by choosing different combintations of the internal search types (the bit sizes of the data types are shown as point labels).\n",
    "Depending on the GPU and the selected dataset, you may see different pictures.\n",
    "With SIFT-128 (`pq_dim = 64`), reducing the `internal_distance_dtype` comes at a huge cost to recall,\n",
    "whereas `lut_dtype` doesn't cost too much while significantly improving QPS.\n",
    "\n",
    "Also, often you may see `16/16` version being faster than `16/8`.\n",
    "This indicates that ALU is the bottleneck in this configuration, and a few extra ALU operations for converting between fp8 and fp16 do more harm than the saved L1 bandwidth does good for the performance.\n",
    "\n",
    "\n",
    "Let's try the same experiment, but with refinement.\n",
    "We'll try ratio 2 and 4 and see how it affects recall and QPS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_refine(ps, ratio):\n",
    "    k_search = k * ratio\n",
    "    candidates = ivf_pq.search(ps, index, queries, k_search, resources=resources)[1]\n",
    "    return candidates if ratio == 1 else refine(dataset, queries, candidates, k, resources=resources)[1]\n",
    "\n",
    "ratios = [1, 2, 4]\n",
    "bench_qps_sr = np.zeros((len(ratios), len(search_ps)), dtype=np.float32)\n",
    "bench_recall_sr = np.zeros((len(ratios), len(search_ps)), dtype=np.float32)\n",
    "\n",
    "for j, ratio in enumerate(ratios): \n",
    "    for i, ps in enumerate(search_ps):\n",
    "        r = %timeit -o search_refine(ps, ratio); resources.sync()\n",
    "        bench_qps_sr[j, i] = (queries.shape[0] * r.loops / np.array(r.all_runs)).mean()\n",
    "        bench_recall_sr[j, i] = calc_recall(search_refine(ps, ratio), gt_neighbors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=plt.figaspect(1/2))\n",
    "fig.suptitle(\n",
    "    f'Effects of search parameters on QPS/recall trade-off ({DATASET_NAME})\\n' + \\\n",
    "    f'k = {k}, n_probes = {n_probes}, pq_dim = {pq_dim}')\n",
    "labels = []\n",
    "for j, ratio in enumerate(ratios):\n",
    "    ax.plot(bench_recall_sr[j, :], bench_qps_sr[j, :], 'o')\n",
    "    labels.append(f\"refine ratio = {ratio}\")\n",
    "ax.legend(labels)\n",
    "ax.set_xlabel('recall')\n",
    "ax.set_ylabel('QPS')\n",
    "ax.grid()\n",
    "colors = plt.rcParams[\"axes.prop_cycle\"].by_key()[\"color\"]\n",
    "annotations = []\n",
    "for j, ratio in enumerate(ratios):\n",
    "    for i, label in enumerate(bench_names):\n",
    "        annotations.append(ax.text(\n",
    "            bench_recall_sr[j, i], bench_qps_sr[j, i],\n",
    "            f\" {label} \",\n",
    "            color=colors[j],\n",
    "            ha='center', va='center'))\n",
    "clutter = [\n",
    "    ax.text(\n",
    "        0.02, 0.08,\n",
    "        'Labels denote the bitsize of: internal_distance_dtype/lut_dtype',\n",
    "        verticalalignment='top',\n",
    "        bbox={'facecolor': 'white', 'edgecolor': 'grey'},\n",
    "        transform = ax.transAxes)\n",
    "]\n",
    "adjust_text(annotations, objects=clutter);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Depending on the dataset, you may see very different pictures here. For SIFT-128, we pick three interesting candidates candidates featuring compromizes between the QPS and the recall:\n",
    "  - `internal_distance_dtype = 16, lut_dtype = 16`\n",
    "  - `internal_distance_dtype = 32, lut_dtype = 8`\n",
    "  - `internal_distance_dtype = 32, lut_dtype = 8, refine_ratio = 2`\n",
    "\n",
    "This is all for the search parameters, but we will come back to the look-up table question in the next section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_refine(internal_distance_dtype, lut_dtype, ratio, n_probes):\n",
    "    k_search = k * ratio\n",
    "    ps = ivf_pq.SearchParams(\n",
    "        n_probes=n_probes,\n",
    "        internal_distance_dtype=internal_distance_dtype,\n",
    "        lut_dtype=lut_dtype)\n",
    "    candidates = ivf_pq.search(ps, index, queries, k_search, resources=resources)[1]\n",
    "    return candidates if ratio == 1 else refine(dataset, queries, candidates, k, resources=resources)[1]\n",
    "\n",
    "search_configs = [\n",
    "    lambda n_probes: search_refine(np.float16, np.float16, 1, n_probes),\n",
    "    lambda n_probes: search_refine(np.float32, np.uint8, 1, n_probes),\n",
    "    lambda n_probes: search_refine(np.float32, np.uint8, 2, n_probes)\n",
    "]\n",
    "search_config_names = [\n",
    "    '16/16', '32/8', '32/8/r2'\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tweaking indexing parameters\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Deciding on the indexing parameters is a bit more involved than on the search parameters. This is obviously because `ivf_pq.IndexParams` has more members than `ivf_pq.SearchParams`, but also because the try-test loop takes longer time when it includes training.\n",
    "Since cuVS' IVF-PQ algorithm uses balanced-hierarchical k-means clustering and efficient logic for encoding, we find significantly improved index build times.\n",
    "\n",
    "First of all, let's pick the parameters we __don't need__ to tweak:\n",
    "\n",
    "  - `metric` - the distance metric often depens on the problem and thus fixed (currently cuVS supports variations of eucliean and inner product distances).\n",
    "  - `conservative_memory_allocation` only affects how data is allocated - does not affect the search performance.\n",
    "  - `add_data_on_build` is a convenience flag. When activated, it automatically adds the training data to the index during `ivf_pq.build`. Otherwise, no data is added during `ivf_pq.build` and vectors need to be explicitly added to the index using `ivf_pq.extend`.\n",
    "  - `force_random_rotation` may slightly affect performance when the data dimensionality is a power of two (see the module docs), but normally you don't need to change the defaults. \n",
    "\n",
    "The rest of the parameters can be divided in two categories: influencing the coarse search (`kmeans_n_iters`, `kmeans_trainset_fraction` , `n_lists`) and the fine search / product quantization (`codebook_kind`, `pq_dim`, `pq_bits`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Indexing parameters affecting the coarse search\n",
    "\n",
    "#### n_lists\n",
    "\n",
    "`n_lists` is the first parameter to look at. It has a profound impact on overall performance during both training and search.\n",
    "`n_lists` defines the number of clusters into which the index data is partitioned; you should keep this in mind when selecting the `n_probes` search parameter.\n",
    "\n",
    "The ratio `n_probes/n_lists` tells how large fraction of the dataset is compared to each query. If `n_lists == n_probes`, that is like a brute force search: we compare all dataset vectors to all query vectors. One would expect the recall is equal to `1` in such a case, but that does not take into account the PQ compression, which is lossy; in reality the recall is always lower unless you refine the search results.\n",
    "\n",
    "As `n_probes` approaches `n_lists`, IVF-PQ becomes slower than brute force because of all the extra work the algorithm does: dimension padding / transform, two-step search, extra PQ compute, etc. In practice searching around 0.1-1% of lists is enough for many datasets. But this depends on how well the input can be clustered. (e.g. for uniform random numbers as inputs, IVF methods don't work well).\n",
    "\n",
    "`n_lists = sqrt(n_samples)` is a good starting point for the balance of coarse/fine search time. To make sure the GPU resources are utilized efficiently, keep in mind:\n",
    "  - The average cluster size (i.e. `n_smaples / n_lists`) should be in the range of at least ~2k records to keep individual SMs busy\n",
    "  - Total amount of search work (`n_queries * n_probes`) should be a good multiple of number of SMs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_list_variants = [100, 500, 1000, 2000, 5000]\n",
    "pl_ratio_variants = [500, 200, 100, 50, 10, 5]\n",
    "selected_search_variant = 1\n",
    "search_fun = search_configs[selected_search_variant]\n",
    "search_label = search_config_names[selected_search_variant]\n",
    "\n",
    "bench_qps_nl = np.zeros((len(n_list_variants), len(pl_ratio_variants)), dtype=np.float32)\n",
    "bench_recall_nl = np.zeros_like(bench_qps_nl, dtype=np.float32)\n",
    "\n",
    "for i, n_lists in enumerate(n_list_variants):\n",
    "    index_params = ivf_pq.IndexParams(n_lists=n_lists, metric=metric, pq_dim=pq_dim)\n",
    "    index = ivf_pq.build(index_params, dataset, resources=resources)\n",
    "    for j, pl_ratio in enumerate(pl_ratio_variants):\n",
    "        n_probes = max(1, n_lists // pl_ratio)\n",
    "        r = %timeit -o search_fun(n_probes);  resources.sync()\n",
    "        bench_qps_nl[i, j] = (queries.shape[0] * r.loops / np.array(r.all_runs)).mean()\n",
    "        bench_recall_nl[i, j] = calc_recall(search_fun(n_probes), gt_neighbors)\n",
    "    del index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=plt.figaspect(1/2))\n",
    "fig.suptitle(\n",
    "    f'Effects of n_list on QPS/recall trade-off ({DATASET_NAME})\\n' + \\\n",
    "    f'k = {k}, pq_dim = {pq_dim}, search = {search_label}')\n",
    "labels = []\n",
    "for i, n_lists in enumerate(n_list_variants):\n",
    "    ax.plot(bench_recall_nl[i, :], bench_qps_nl[i, :])\n",
    "    labels.append(f\"n_lists = {n_lists}\")\n",
    "\n",
    "ax.legend(labels)\n",
    "ax.set_xlabel('recall')\n",
    "ax.set_ylabel('QPS')\n",
    "ax.set_yscale('log')\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This chart demonstrates that for the given data set (SIFT-128) and the selected parameters, the QPS/recall curves are rather close to each other.\n",
    "Yet, two lines, which correspond to 100- and 5000-cluster indices, lag below the others.\n",
    "This suggests that 5000 clusters is probably too many and 100 clusters is probably too few for this dataset. In the range of 500-2000 the algorithm performs very similar though.\n",
    "Hence, you shouldn't worry about finding the exact single best value of `n_lists`, but rather make sure it's within a reasonable range.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### kmeans_trainset_fraction\n",
    "\n",
    "This parameter defines how much of the original data should be fed into training.\n",
    "This is useful when in conjunction with `add_data_on_build = True`.\n",
    "For example, having a 100M-record dataset, it's reasonable to set `kmeans_trainset_fraction = 0.1` to train the index (i.e. run the k-means clustering) using 10M records only (10% of data), and then add the whole dataset to the index.\n",
    "Hence, this parameter directly affects the training speed, but can indirectly affect the search performance (depending on how well the training set represents the full dataset).\n",
    "\n",
    "Note, if `add_data_on_build = False`, setting the trainset fraction less than one is identical to passing a smaller dataset to the `ivf_pq.build`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### kmeans_n_iters\n",
    "\n",
    "This parameter is passed directly to the k-means algorithm during training. It's set to a reasonable default of 20, which works for most datasets. However, once in a while you may see a warning complaining that the trained clusters are imbalanced. You can try to fix that by increasing the number of iterations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Indexing parameters affecting the fine search / product quantization\n",
    "\n",
    "In the IVF-PQ index, a database vector y is approximated with two level quantization:\n",
    "\n",
    "$ y = Q_1(y) + Q_2(y - Q_1(y)) $\n",
    "\n",
    "The first level quantizer ($Q_1$), maps the vector y to the nearest cluster center. The number of\n",
    "clusters is `n_lists`.\n",
    "\n",
    "The second quantizer encodes the residual, and it is defined as a product quantizer\n",
    "(see [_\"Product quantization for nearest neighbor search\" by Herve Jegou, Matthijs Douze, Cordelia Schmid_](https://www.researchgate.net/publication/47815472_Product_Quantization_for_Nearest_Neighbor_Search)).\n",
    "\n",
    "A product quantizer encodes a `dim` dimensional vector with a `pq_dim` dimensional vector.\n",
    "First we split the input vector into `pq_dim` subvectors (denoted by u), where each u vector\n",
    "contains `pq_len` distinct components of y\n",
    "```\n",
    "y_1, y_2, ... y_{pq_len}, y_{pq_len+1}, ... y_{2*pq_len}, ... y_{dim-pq_len+1} ... y_{dim}\n",
    " \\___________________/     \\____________________________/      \\______________________/\n",
    "        u_1                         u_2                          u_{pq_dim}\n",
    "```\n",
    "Then each subvector encoded with a separate quantizer $q_i$, end the results are concatenated\n",
    "\n",
    "$ Q_2(y) = q_1(u_1),q_2(u_2),...,q_\\mathtt{pq\\_dim}(u_\\mathtt{pq\\_dim}) $\n",
    "\n",
    "Each quantizer $q_i$ outputs a code with `pq_bit` bits. The second level quantizers are also defined\n",
    "by k-means clustering in the corresponding sub-space: the reproduction values are the centroids,\n",
    "and the set of reproduction values is the codebook.\n",
    "\n",
    "During the search, for every query and probed list, a look-up table (LUT) is constructed using appropriate codebooks and the query coordinates.\n",
    "The size of the LUT has profound effect on the performance; here it is one more time:\n",
    "\n",
    "$ \\mathtt{lut\\_size} = \\mathtt{pq\\_dim} \\cdot \\mathtt{sizeof(lut\\_dtype) \\cdot 2^{\\mathtt{pq\\_bits}}} $\n",
    "\n",
    "If possible, the LUT is stored fully in GPU L1 (shared) memory during search;\n",
    "otherwise, a slower version of the kernel is used, which stores the LUT in the global memory.\n",
    "\n",
    "\n",
    "#### codebook_kind\n",
    "\n",
    "The second-level quantizers are trained either for each subspace or for each cluster, controlled by parameter `codebook_kind`:\n",
    "\n",
    "  1. \"subspace\" (C++ api: `codebook_gen::PER_SUBSPACE`): \\\n",
    "        creates `pq_dim` second-level quantizers - one for each slice of the data along features;\n",
    "  2. \"cluster\" (C++ api: `codebook_gen::PER_CLUSTER`): \\\n",
    "        creates `n_lists` second-level quantizers - one for each first-level cluster.\n",
    "\n",
    "In either case, the centroids are found using k-means clustering interpreting the data as having `pq_len` dimensions.\n",
    "\n",
    "There's no definitive way to tell in advance, which of the two options yields better performance for a particular use case.\n",
    "A few observations, however, may help:\n",
    "\n",
    "  - A per-cluster codebook tends to take more time to train, since `n_lists` is usually much higher than `pq_dim` - more codebooks to train.\n",
    "  - Search with a per-cluster codebook usually utilizes L1 cache of the GPU better than with a per-subspace codebook; this may result in a faster search when the LUT is big and occupies a large part of the GPU L1 memory.\n",
    "  - However, in practice, the recall is slightly higher with a per-subspace codebook.\n",
    "\n",
    "\n",
    "#### pq_dim, pq_bits\n",
    "\n",
    "`pq_dim` parameter is the main way to control the compression in the database.\n",
    "You should choose it depending on your expectations about the sparsity of the information in the data.\n",
    "As an experiment, you could start with `pq_dim` in the range of the data dimensionality `[dim / 2, dim]`.\n",
    "\n",
    "`pq_bits` is the number of bits in a single PQ code.\n",
    "Hence, it controls the codebook size - $2^{\\mathtt{pq\\_bits}}$ - the number of possible values a code can take.\n",
    "IVF-PQ supports the codebooks sizes from 16 to 256, or the `pq_bits` in the range of `[4, 8]`.\n",
    "\n",
    "`pq_bits` affects the compression: a database with `pq_bits = 4` is twice smaller than with the `pq_bits = 8`.\n",
    "Though much stronger `pq_bits` affects the LUT size, as the LUT size is proportional to $2^{\\mathtt{pq\\_bits}}$ (see the formula above).\n",
    "This also means a drastic effect on the recall.\n",
    "\n",
    "A few observations:\n",
    "\n",
    "  - It's required that `(pq_dim * pq_bits) % 8 == 0`; in general, keeping `pq_dim` in powers of two improves the search performance due to better data alignment.\n",
    "  - Keeping `pq_dim * pq_bits >= 128` and `(pq_dim * pq_bits) % 32 == 0` maximizes the GPU memory bandwidth utilization.\n",
    "  - Generally `pq_bits = 8` is a good starting point.\n",
    "  - The recall loss due to smaller `pq_bits` can be compensated by enabling refinement.\n",
    "  - For high-dimensional data and large `pq_dims`, lowering `pq_bits` can yield a drastic search speedup due to enabling the faster kernel that keeps the LUT in L1.\n",
    "  - Alternatively, setting the search parameter `lut_dtype` to `uint8` may be enough to keep the LUT in L1.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's try a few build configurations.\n",
    "#   Warning: this will take some time\n",
    "\n",
    "k = 10\n",
    "n_probes_variants = [10, 20, 50, 100]\n",
    "n_lists = 1000\n",
    "\n",
    "build_configs = {\n",
    "    '64-8-subspace': ivf_pq.IndexParams(n_lists=n_lists, metric=metric, pq_dim=64, pq_bits=8, codebook_kind=\"subspace\"),\n",
    "    '128-8-subspace': ivf_pq.IndexParams(n_lists=n_lists, metric=metric, pq_dim=128, pq_bits=8, codebook_kind=\"subspace\"),\n",
    "    '128-6-subspace': ivf_pq.IndexParams(n_lists=n_lists, metric=metric, pq_dim=128, pq_bits=6, codebook_kind=\"subspace\"),\n",
    "    '128-6-cluster': ivf_pq.IndexParams(n_lists=n_lists, metric=metric, pq_dim=128, pq_bits=6, codebook_kind=\"cluster\"),\n",
    "}\n",
    "\n",
    "bench_qps_ip = np.zeros((len(build_configs), len(search_configs), len(n_probes_variants)), dtype=np.float32)\n",
    "bench_recall_ip = np.zeros_like(bench_qps_ip, dtype=np.float32)\n",
    "\n",
    "for i, index_params in enumerate(build_configs.values()):\n",
    "    index = ivf_pq.build(index_params, dataset, resources=resources)\n",
    "    for l, search_fun in enumerate(search_configs):\n",
    "        for j, n_probes in enumerate(n_probes_variants):\n",
    "            r = %timeit -o search_fun(n_probes);  resources.sync()\n",
    "            bench_qps_ip[i, l, j] = (queries.shape[0] * r.loops / np.array(r.all_runs)).mean()\n",
    "            bench_recall_ip[i, l, j] = calc_recall(search_fun(n_probes), gt_neighbors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(len(search_config_names), 1, figsize=(16, len(search_config_names)*8))\n",
    "fig.suptitle(\n",
    "    f'Effects of index parameters on QPS/recall trade-off ({DATASET_NAME})\\n' + \\\n",
    "    f'k = {k}, n_lists = {n_lists}')\n",
    "\n",
    "for j, search_label in enumerate(search_config_names):\n",
    "    labels = []\n",
    "    for i, index_label in enumerate(build_configs.keys()):\n",
    "        ax[j].plot(bench_recall_ip[i, j, :], bench_qps_ip[i, j, :])\n",
    "        labels.append(index_label)\n",
    "\n",
    "    ax[j].set_title(f\"search: {search_label}\")\n",
    "    ax[j].legend(labels)\n",
    "    ax[j].set_xlabel('recall')\n",
    "    ax[j].set_ylabel('QPS')\n",
    "    ax[j].set_yscale('log')\n",
    "    ax[j].grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks like `pq_dim = 128`, `pq_bits = 6` is the best parameter set for the `SIFT-128` dataset."
   ]
  }
 ],
 "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.11.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}