torchdynamo_dynamic_inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "63da03a1",
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x107c1d7f0>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import functools\n",
    "import itertools\n",
    "import traceback\n",
    "from dataclasses import dataclass, field\n",
    "from enum import Enum\n",
    "from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, Union\n",
    "\n",
    "import torch\n",
    "\n",
    "torch.manual_seed(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12e2ab68",
   "metadata": {},
   "source": [
    "This notebook explains how Jason Ansel's proposal for very simple\n",
    "dynamic shapes in TorchDynamo works in\n",
    "https://github.com/facebookresearch/torchdynamo/issues/38\n",
    "\n",
    "The general model for torchdynamo graphs is that they consist of a\n",
    "set of guards plus a trace.  The guards say whether or not the trace\n",
    "is valid; if it is not, torchdynamo must redo its analysis and\n",
    "recompile the graph in question.\n",
    "\n",
    "In this simplified model, we will model torchdynamo graphs as a\n",
    "dead simple AST (in reality, you need a graph representation to\n",
    "specify ordering of operations, sharing and multiple outputs, but\n",
    "they are not relevant to this example so I've dumped them.)\n",
    "\n",
    "First, we define various operations on the graph.  add and mul\n",
    "do what you expect: they perform a (broadcasting) PyTorch add and\n",
    "mul.  `dynamic_param` and `static_param` both represent inputs\n",
    "to the graph.  The distinction is that `dynamic_param` inputs\n",
    "correspond to inputs which are fully dynamic: their shapes can\n",
    "vary from execution to execution of the graph.  `static_param`\n",
    "inputs, on the other hand, are required to be some specific size.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4d57a5aa",
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "@dataclass(frozen=True)\n",
    "class Op:\n",
    "    name: str\n",
    "\n",
    "    def __str__(self):\n",
    "        return self.name\n",
    "\n",
    "\n",
    "v_dynamic_param = Op(\"v_dynamic_param\")\n",
    "v_static_param = Op(\"v_static_param\")\n",
    "v_add = Op(\"v_add\")\n",
    "v_mul = Op(\"v_mul\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24264adb",
   "metadata": {},
   "source": [
    "We can stitch these operations together in an AST of expressions\n",
    "of operators applied to some other expressions (and possibly some\n",
    "other, static metadata)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "47ccaa14",
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "@dataclass(eq=False)\n",
    "class Node:\n",
    "    op: Op\n",
    "    inputs: List[\"Node\"] = field(default_factory=list)\n",
    "    params: Dict[str, Any] = field(default_factory=dict)\n",
    "\n",
    "    def __repr__(self):\n",
    "        inputs_str = \", \".join(repr(i) for i in self.inputs)\n",
    "        params_str = \"\"\n",
    "        if self.inputs and self.params:\n",
    "            params_str += \", \"\n",
    "        params_str += \", \".join(\n",
    "            f\"{k}={v}\"\n",
    "            for k, v in self.params.items()\n",
    "            if k != \"size\" and self.op is v_dynamic_param\n",
    "        )\n",
    "        return f\"{self.op}({inputs_str}{params_str})\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfcd54fb",
   "metadata": {},
   "source": [
    "And then we can write an interpreter for these inputs.  Notice that\n",
    "we fetch parameters from an environment that's passed into the\n",
    "interpreter; if the parameter is dynamic we pass it in directly,\n",
    "but if it's static, we first check that the size of the parameter\n",
    "is consistent with the saved size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "924b75a6",
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "INTERP_RULES = {}\n",
    "INTERP_RULES[v_add] = lambda x, y: x + y\n",
    "INTERP_RULES[v_mul] = lambda x, y: x * y\n",
    "\n",
    "\n",
    "def interp_node(n: Node, env: Dict[Node, torch.Tensor]):\n",
    "    if n.op is v_dynamic_param:\n",
    "        return env[n.params['name']]\n",
    "    elif n.op is v_static_param:\n",
    "        r = env[n.params['name']]\n",
    "        assert (\n",
    "            r.shape == n.params[\"size\"]\n",
    "        ), f\"static shape mismatch: {r.shape} and {n.params['size']}\"\n",
    "        return r\n",
    "    args = [interp_node(i, env) for i in n.inputs]\n",
    "    return INTERP_RULES[n.op](*args, **n.params)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14a70302",
   "metadata": {},
   "source": [
    "In actual torchdynamo, we can construct our IR directly via\n",
    "bytecode analysis.  But this isn't really necessary for our\n",
    "example here; we can use an ordinary tracer to construct the IR as\n",
    "well.  Our tracer is very simple."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "036777d0",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Variable:\n",
    "    tensor: torch.Tensor\n",
    "    node: Node\n",
    "\n",
    "    # This will be implemented later\n",
    "    def size(self):\n",
    "        return variable_size(self)\n",
    "\n",
    "    @staticmethod\n",
    "    def param(tensor: torch.Tensor, name: str):\n",
    "        # Save the observed shape, but by default dynamic_param won't\n",
    "        # check it!\n",
    "        return Variable(tensor, Node(v_dynamic_param, [], {\"name\": name, \"size\": tensor.shape}))\n",
    "\n",
    "    def __mul__(self, rhs: \"Variable\") -> \"Variable\":\n",
    "        r_tensor = self.tensor * rhs.tensor\n",
    "        r_node = Node(v_mul, [self.node, rhs.node])\n",
    "        return Variable(r_tensor, r_node)\n",
    "\n",
    "    def __add__(self, rhs: \"Variable\") -> \"Variable\":\n",
    "        r_tensor = self.tensor + rhs.tensor\n",
    "        r_node = Node(v_add, [self.node, rhs.node])\n",
    "        return Variable(r_tensor, r_node)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d05d056",
   "metadata": {},
   "source": [
    "With this, we can run a simple example, print out the IR for it,\n",
    "and then rerun it.  By default, we treat the inputs as dynamics,\n",
    "so we are allowed to rerun the IR even though the input sizes have\n",
    "changed (because there is nothing shape specific in the IR.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "42443acd",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = Variable.param(torch.randn(4), \"a\")\n",
    "b = Variable.param(torch.randn(4), \"b\")\n",
    "r = a * b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6b22da5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "v_mul(v_dynamic_param(name=a), v_dynamic_param(name=b))\n"
     ]
    }
   ],
   "source": [
    "print(r.node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c8fb3092",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.7916, -0.4439, -0.6567,  0.2004, -0.9429])\n"
     ]
    }
   ],
   "source": [
    "print(interp_node(r.node, {\"a\": torch.randn(5), \"b\": torch.randn(1)}))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b0b306f",
   "metadata": {},
   "source": [
    "Now, the problem is what happens if a user wants to vary the\n",
    "behavior of their computation based on the size of their input?\n",
    "Then our trace is no longer valid in this situation!\n",
    "\n",
    "torchdynamo deals with this situation by looking for explicit uses\n",
    "of sizes.  If there is an explicit use of a size, it goes ahead\n",
    "and conservatively marks all of the parameters which could have\n",
    "contributed to the size of this tensor as static, indicating that\n",
    "the trace is now only valid for those specific sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c32d30d6",
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "def input_sources(node):\n",
    "    r = set()\n",
    "    for i in node.inputs:\n",
    "        r |= input_sources(i)\n",
    "    if node.op is v_dynamic_param:\n",
    "        r.add(node)\n",
    "    return r\n",
    "\n",
    "def variable_size(self):\n",
    "    for i in input_sources(self.node):\n",
    "        # change it from dynamic to static.  (the parameter\n",
    "        # already saved the size, we don't need to recover it)\n",
    "        i.op = v_static_param\n",
    "    return self.tensor.size()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e151cbc8",
   "metadata": {},
   "source": [
    "Now if we have dependent control flow on an input, we will\n",
    "appropriately fail if you pass in mismatching sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7a5b2a4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "a = Variable.param(torch.randn(4), \"a\")\n",
    "b = Variable.param(torch.randn(4), \"b\")\n",
    "if a.size()[0] == 4:\n",
    "    r = a + b\n",
    "else:\n",
    "    r = a * b\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ca4017fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "v_add(v_static_param(), v_dynamic_param(name=b))\n"
     ]
    }
   ],
   "source": [
    "print(r.node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "22c28e13",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.3506, -0.0163,  0.1710,  0.5453])\n"
     ]
    }
   ],
   "source": [
    "print(interp_node(r.node, {\"a\": torch.randn(4), \"b\": torch.randn(4)}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "cfa558aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Traceback (most recent call last):\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_65823/2739262870.py\", line 2, in <cell line: 1>\n",
      "    print(interp_node(r.node, {\"a\": torch.randn(5), \"b\": torch.randn(1)}))\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_65823/4116253730.py\", line 15, in interp_node\n",
      "    args = [interp_node(i, env) for i in n.inputs]\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_65823/4116253730.py\", line 15, in <listcomp>\n",
      "    args = [interp_node(i, env) for i in n.inputs]\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_65823/4116253730.py\", line 11, in interp_node\n",
      "    assert (\n",
      "AssertionError: static shape mismatch: torch.Size([5]) and torch.Size([4])\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    print(interp_node(r.node, {\"a\": torch.randn(5), \"b\": torch.randn(1)}))\n",
    "except Exception:\n",
    "    traceback.print_exc()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd961281",
   "metadata": {},
   "source": [
    "It will still work even if the shape check is done an intermediate\n",
    "computation (in this case, both a and b are marked as dynamic)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ee8a24c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "a = Variable.param(torch.randn(1), \"a\")\n",
    "b = Variable.param(torch.randn(1), \"b\")\n",
    "c = a + b\n",
    "if c.size()[0] == 1:\n",
    "    r = a + c\n",
    "else:\n",
    "    r = a * c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "17f83008",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "v_add(v_static_param(), v_add(v_static_param(), v_static_param()))\n"
     ]
    }
   ],
   "source": [
    "print(r.node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3eaf8e72",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Traceback (most recent call last):\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_65823/81332487.py\", line 2, in <cell line: 1>\n",
      "    print(interp_node(r.node, {\"a\": torch.randn(1), \"b\": torch.randn(5)}))\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_65823/4116253730.py\", line 15, in interp_node\n",
      "    args = [interp_node(i, env) for i in n.inputs]\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_65823/4116253730.py\", line 15, in <listcomp>\n",
      "    args = [interp_node(i, env) for i in n.inputs]\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_65823/4116253730.py\", line 15, in interp_node\n",
      "    args = [interp_node(i, env) for i in n.inputs]\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_65823/4116253730.py\", line 15, in <listcomp>\n",
      "    args = [interp_node(i, env) for i in n.inputs]\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_65823/4116253730.py\", line 11, in interp_node\n",
      "    assert (\n",
      "AssertionError: static shape mismatch: torch.Size([5]) and torch.Size([1])\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    print(interp_node(r.node, {\"a\": torch.randn(1), \"b\": torch.randn(5)}))\n",
    "except Exception:\n",
    "    traceback.print_exc()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "989c9887",
   "metadata": {},
   "source": [
    "This analysis is VERY conservative.  Although there are some easy\n",
    "improvements you can apply, you are limited in the precision you can\n",
    "have without having shape formulas for operators that can propagate\n",
    "dynamic shapes.  With shape formulas, you can track exact dependencies\n",
    "on a size-by-size basis; if you matrix multiply two tensors C = A @ B,\n",
    "a use of C.shape[0] will only add a guard for A.shape[0], and a use of\n",
    "C.shape[1] will only add a guard for B.shape[1].  The analysis here\n",
    "will just make both A and B static, and we cannot do any better\n",
    "without more knowledge of formulas.  This suggests that an important\n",
    "workstream to improve precision is to get dynamic-aware shape formulas\n",
    "in Python for as many operators as possible."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:light"
  },
  "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.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}