{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6a95031f",
   "metadata": {},
   "source": [
    "# Heuristics: picking a likely-fast operator\n",
    "\n",
    "`get_operators` can return **many** operators that are all *correct* for your\n",
    "problem. Without a heuristic, that list is in **discovery order** — not\n",
    "\"fastest first.\"\n",
    "\n",
    "A **heuristic** can optionally rank the estimated performance of candidate operators,\n",
    "and *reorders and prunes* the candidates accordingly.\n",
    "\n",
    "CUTLASS Operator API natively supports [NVIDIA Matmul Heuristics](https://docs.nvidia.com/cuda/nvidia-matmul-heuristics/) (`nvMatmulHeuristics`)\n",
    "to do such ranking directly in `get_operators()`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc611145",
   "metadata": {},
   "source": [
    "## Prerequisite: Install `nvidia-matmul-heuristics`\n",
    "\n",
    "NVIDIA Matmul Heuristics (nvMatmulHeuristics) is an optional dependency of CUTLASS\n",
    "Operator API, and must be present to run heuristics.\n",
    "\n",
    "**Install: `pip install 'nvidia-cutlass-operators[heuristics]'`**\n",
    "\n",
    "\n",
    "Currently, heuristics integration is only supported for non-blockscaled GEMMs on\n",
    "SM100.\n",
    "\n",
    "If you request a heuristic and ranking cannot run (missing optional package,\n",
    "unsupported args), an error is raised."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "41e10954",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-21T03:04:18.895778Z",
     "iopub.status.busy": "2026-08-21T03:04:18.895652Z",
     "iopub.status.idle": "2026-08-21T03:04:26.189457Z",
     "shell.execute_reply": "2026-08-21T03:04:26.188194Z"
    }
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "import torch\n",
    "\n",
    "import cutlass.operators as ops\n",
    "from cutlass.operators.heuristics.nvmatmul import is_available as nvmatmul_available\n",
    "\n",
    "if not (status := ops.utils.device.device_or_env_supports(\"100\")):\n",
    "    print(f\"This notebook expects an SM100-class GPU.\\n{status.error}\")\n",
    "    sys.exit(0)\n",
    "\n",
    "if not nvmatmul_available():\n",
    "    print(\n",
    "        \"nvmatmul heuristic is unavailable. Install with:\\n\"\n",
    "        \"  pip install 'nvidia-cutlass-operators[heuristics]'\"\n",
    "    )\n",
    "    sys.exit(0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e409cba",
   "metadata": {},
   "source": [
    "## Minimal usage\n",
    "\n",
    "Operator API can create a nvMatmulHeuristics instance using `ops.get_heuristic(\"nvmatmul\")(gpu=\"B200\")`.\n",
    "`gpu=` names the exact GPU SKU to model and defaults to `\"B200\"` if omitted.\n",
    "\n",
    "`get_operators()` then natively supports using this `heuristic` returning candidate\n",
    "operators for given arguments.\n",
    "Additionally, we can use `limit=N` to limit the results to the top-N heuristic-recommended Operators.\n",
    "\n",
    "Currently, it supports only supports non-blockscaled GEMMs for Blackwell (SM100) GPUs.\n",
    "\n",
    "When supported, `get_operators(args, heuristic)` will:\n",
    "1. Query the supported, unsorted candidate operators\n",
    "2. Query the heuristic for recommended kernel configurations for\n",
    "   your given argument (problem size, dtypes, ...).\n",
    "3. Matches those recommendations to the candidate Operators. Operators not matched by the heuristic are excluded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "67ba90c4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-21T03:04:26.191960Z",
     "iopub.status.busy": "2026-08-21T03:04:26.191803Z",
     "iopub.status.idle": "2026-08-21T03:04:27.036200Z",
     "shell.execute_reply": "2026-08-21T03:04:27.035328Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Returned 5 operator(s) (limit=5)\n"
     ]
    }
   ],
   "source": [
    "M, N, K = 4096, 4096, 4096\n",
    "A = torch.randn(M, K, device=\"cuda\", dtype=torch.float16)\n",
    "B = torch.randn(K, N, device=\"cuda\", dtype=torch.float16)\n",
    "out = torch.empty(M, N, device=\"cuda\", dtype=torch.float16)\n",
    "args = ops.GemmArguments(A, B, out, accumulator_type=torch.float32)\n",
    "\n",
    "heuristic = ops.get_heuristic(\"nvmatmul\")(gpu=\"B200\")\n",
    "# or, equivalently:\n",
    "heuristic =  ops.heuristics.NvMatmulHeuristics(gpu=\"B200\")\n",
    "\n",
    "operators = ops.get_operators(\n",
    "    args,\n",
    "    target_sm=\"100a\",\n",
    "    providers=[ops.CuTeDSLProvider],\n",
    "    heuristic=heuristic,\n",
    "    limit=5,\n",
    ")\n",
    "print(f\"Returned {len(operators)} operator(s) (limit=5)\")\n",
    "operators[0].run(args)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e756d736",
   "metadata": {},
   "source": [
    "An error is raised if the ranking cannot run at all, e.g. for\n",
    "missing package, unsupported args, and unsupported GPU (currently only SM100 is supported)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ec483637",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-21T03:04:27.038473Z",
     "iopub.status.busy": "2026-08-21T03:04:27.038373Z",
     "iopub.status.idle": "2026-08-21T03:04:27.041054Z",
     "shell.execute_reply": "2026-08-21T03:04:27.040275Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unsupported GPU: nvmatmul only supports SM100 today; 'H100_SXM' isn't a recognized/supported device (e.g. \"B200\", \"GB200_NVL\", \"GB300_NVL\").\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    ops.get_heuristic(\"nvmatmul\")(gpu=\"H100_SXM\")\n",
    "except ValueError as e:\n",
    "    print(f\"Unsupported GPU: {e}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fdedd80",
   "metadata": {},
   "source": [
    "## Comparing the result\n",
    "\n",
    "We do a quick benchmark below to compare the operators returned by heuristics.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "380febd3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-21T03:04:27.043272Z",
     "iopub.status.busy": "2026-08-21T03:04:27.043185Z",
     "iopub.status.idle": "2026-08-21T03:04:28.409498Z",
     "shell.execute_reply": "2026-08-21T03:04:28.408688Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "892 total operator(s) for this problem\n",
      "\n",
      "After applying heuristics, this was sorted and pruned to 26 operator(s).\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fastest heuristic-recommended operator: 0.0894 ms \t (name: cutedsl.PersistentDenseGemmOperator_sm100_ttt_AFloat16_BFloat16_outFloat16_accFloat32_2cta_cluster2x2x1_tile256x256x64_schedulerCLC_tma_store)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "slowest heuristic-recommended operator: 0.2378 ms \t (name: cutedsl.PersistentDenseGemmEFCOperator_sm100_ttt_AFloat16_BFloat16_outFloat16_accFloat32_1cta_cluster1x2x1_tile64x32x64_tma_store)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "arbitrary operator: 0.2745 ms \t (name: cutedsl.PersistentDenseGemmOperator_sm100_ttt_AFloat16_BFloat16_outFloat16_accFloat32_2cta_cluster2x1x1_tile128x32x64_schedulerCLC_tma_store)\n"
     ]
    }
   ],
   "source": [
    "def benchmark_operator(op: ops.Operator, args: ops.GemmArguments, warmup=10, iters=50):\n",
    "    \"\"\"Return the median GPU time (in ms) for `op.run(args)`.\"\"\"\n",
    "    compiled = op.compile(args)\n",
    "    for _ in range(warmup):\n",
    "        op.run(args, compiled_artifact=compiled, assume_supported_args=True)\n",
    "    torch.cuda.synchronize()\n",
    "\n",
    "    start = torch.cuda.Event(enable_timing=True)\n",
    "    end = torch.cuda.Event(enable_timing=True)\n",
    "    times_ms = []\n",
    "    for _ in range(iters):\n",
    "        start.record()\n",
    "        op.run(args, compiled_artifact=compiled, assume_supported_args=True)\n",
    "        end.record()\n",
    "        torch.cuda.synchronize()\n",
    "        times_ms.append(start.elapsed_time(end))\n",
    "    return sorted(times_ms)[len(times_ms) // 2]\n",
    "\n",
    "\n",
    "unsorted_operators = ops.get_operators(args, target_sm=\"100a\", providers=[ops.providers.CuTeDSLProvider])\n",
    "sorted_operators = ops.get_operators(args, target_sm=\"100a\", providers=[ops.providers.CuTeDSLProvider], heuristic=heuristic)\n",
    "\n",
    "\n",
    "print(f\"{len(unsorted_operators)} total operator(s) for this problem\\n\")\n",
    "print(f\"After applying heuristics, this was sorted and pruned to {len(sorted_operators)} operator(s).\\n\")\n",
    "\n",
    "candidates = {\n",
    "    \"fastest heuristic-recommended operator\": sorted_operators[0],\n",
    "    \"slowest heuristic-recommended operator\": sorted_operators[-1],\n",
    "    \"arbitrary operator\": unsorted_operators[0],\n",
    "}\n",
    "\n",
    "\n",
    "for label, op in candidates.items():\n",
    "    median_ms = benchmark_operator(op, args)\n",
    "    print(f\"{label}: {median_ms:.4f} ms \\t (name: {op.metadata.operator_name})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8a8c93b",
   "metadata": {},
   "source": [
    "Please note that heuristics-based ranking is an estimate. It is helpful to limit and guide the search to promising candidates to top-N recommendations, and may not yield a definitive single winner."
   ]
  }
 ],
 "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}