tracing_guards.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8ff35933",
   "metadata": {},
   "source": [
    "The purpose of this notebook explain validity hazards that apply\n",
    "to traces we collect when we are tracing code with real inputs,\n",
    "i.e., we secretly know what all the true values of our tensors/shapes\n",
    "are, and we are willing to let the user perform control flow on these\n",
    "values.  The basic problem we have to solve, is that for a tracing\n",
    "system based on overloading (as well as compiler systems that simplify\n",
    "their internals by only dealing with straight-line traces, aka tracing\n",
    "JITs--and torchdynamo falls into this category), we will have resolved\n",
    "control flow into taking a specific path, and if on a subsequent\n",
    "execution we would have taken a different path, the trace we have\n",
    "recorded is no longer valid!  So we have to detect this situation.\n",
    "\n",
    "The way this notebook will proceed, is we will implement a very simple\n",
    "set of validity tests based on \"bailouts\"; that is to say, we will\n",
    "sprinkle our trace with boolean tests (called bailouts) which specify\n",
    "when we had Python control flow at some point in time and relied on a\n",
    "value in our network having some specific value, so that later when we\n",
    "reexecute the trace, we can check that we would have done the same\n",
    "thing on the other path.\n",
    "\n",
    "Bailouts directly embedded in the trace are of limited use, as many\n",
    "trace compilers do not support bailout--if you think about a CUDA\n",
    "kernel fusion engine, you're just going to want to directly compile\n",
    "all of your trace into a CUDA kernel, you don't want to get halfway\n",
    "through your code and then realize \"oops, I need to bail out of\n",
    "executing this kernel and do something else\" (and indeed, you can't\n",
    "even do this in CUDA).  However, we can take the bailouts from our\n",
    "traces and transform them into standalone guards (thus the name of\n",
    "this notebook) which we can use to check for validity on entry.  We\n",
    "can also take a bailout and split the graph at that point, which also\n",
    "ensures validity.  Finally, if the user is willing to modify their\n",
    "code, they can make changes which make it less likely they will\n",
    "accidentally condition on a dynamic value (channeling the approach\n",
    "taken in the [Dynamic\n",
    "Shapes](https://github.com/albanD/subclass_zoo/blob/main/dynamic_shapes.ipynb)\n",
    "notebook)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a4703b6",
   "metadata": {},
   "source": [
    "First, some imports and some utilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a91287ed",
   "metadata": {
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "\n",
    "import itertools\n",
    "import traceback\n",
    "from dataclasses import dataclass, field\n",
    "from typing import Any, Dict, List, NoReturn\n",
    "\n",
    "import torch\n",
    "\n",
    "\n",
    "def assert_never(x: NoReturn) -> NoReturn:\n",
    "    raise AssertionError(\"Unhandled type: {}\".format(type(x).__name__))\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class FreshSupply:\n",
    "    prefix: str\n",
    "    fresh: int = 0\n",
    "\n",
    "    def __call__(self):\n",
    "        r = f\"{self.prefix}{self.fresh}\"\n",
    "        self.fresh += 1\n",
    "        return r\n",
    "\n",
    "\n",
    "fresh_var = FreshSupply(\"v\")\n",
    "fresh_int = FreshSupply(\"i\")\n",
    "fresh_bool = FreshSupply(\"b\")\n",
    "fresh_size = FreshSupply(\"s\")\n",
    "\n",
    "\n",
    "def reset():\n",
    "    global CURRENT_GRAPH\n",
    "    fresh_var.fresh = 0\n",
    "    fresh_int.fresh = 0\n",
    "    fresh_bool.fresh = 0\n",
    "    fresh_size.fresh = 0\n",
    "    CURRENT_GRAPH = Graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d803d60e",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "Once again, we need to define an IR for our traces.  Like in\n",
    "our previous notebook, we will represent each valid operation\n",
    "as an instance of the Op class, which is just a little wrapper\n",
    "around strings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d51571a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass(frozen=True)\n",
    "class Op:\n",
    "    name: str\n",
    "\n",
    "    def __str__(self):\n",
    "        return self.name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e27f9c1b",
   "metadata": {},
   "source": [
    "Our operations are going to look slightly different this time.\n",
    "In particular, I explicitly model booleans and size tuples (in\n",
    "the other notebook, I allowed for tuples and integer literals to be\n",
    "directly expressed in the IR.)  This presentation will result in\n",
    "slightly uglier syntax, but it will easier for us to write\n",
    "interpreters and other analyses on the IR due to its uniform\n",
    "representation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fb1cc06f",
   "metadata": {},
   "outputs": [],
   "source": [
    "bool_bailout = Op(\"bool_bailout\")  # b, *, expect\n",
    "bool_const = Op(\"bool_const\")  # *, val\n",
    "int_eq = Op(\"int_eq\")  # a, b\n",
    "int_const = Op(\"int_const\")  # *, val\n",
    "int_placeholder = Op(\"int_placeholder\")  # *, name\n",
    "size_index = Op(\"size_index\")  # s, *, index\n",
    "size_ctor = Op(\"size_ctor\")  # *args\n",
    "var_placeholder = Op(\"var_placeholder\")  # *, name\n",
    "var_add = Op(\"var_add\")  # a, b\n",
    "var_mul = Op(\"var_mul\")  # a, b\n",
    "var_size = Op(\"var_size\")  # a\n",
    "var_expand = Op(\"var_expand\")  # a, size"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21a44f09",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "Given a dictionary of parameters (for handling `var_placeholder`), we\n",
    "can say what the concrete interpretation of all of these operators\n",
    "should be."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9f4c039e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def assert_(b):\n",
    "    assert b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a27de441",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Unlike conventional add/mul in PyTorch, these operators do not\n",
    "# broadcast (we will insert explicit expands)\n",
    "def prim_add(a, b):\n",
    "    assert a.shape == b.shape\n",
    "    return torch.add(a, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7f192240",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prim_mul(a, b):\n",
    "    assert a.shape == b.shape\n",
    "    return torch.mul(a, b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f822a21e",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def concrete_rules(**params):\n",
    "    # bool -> bool\n",
    "    # int -> int\n",
    "    # size -> Tuple[int, ...]\n",
    "    # var -> torch.Tensor\n",
    "    return {\n",
    "        bool_bailout: lambda b, *, expect: assert_(b == expect),\n",
    "        bool_const: lambda *, val: val,\n",
    "        int_eq: lambda a, b: a == b,\n",
    "        int_const: lambda *, val: val,\n",
    "        int_placeholder: lambda *, name: params[name],\n",
    "        size_index: lambda s, *, index: s[index],\n",
    "        size_ctor: lambda *, args: args,\n",
    "        var_placeholder: lambda *, name: params[name],\n",
    "        var_add: prim_add,\n",
    "        var_mul: prim_mul,\n",
    "        var_size: lambda a: tuple(a.shape),\n",
    "        var_expand: lambda a, size: a.expand(size),\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d80a162",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "Let's fill out the rest of our IR and an actual implementation\n",
    "of the interpreter (given some rules).  Unlike the dynamic shapes\n",
    "notebook, our nodes are truly in ANF (so inputs is just a list of\n",
    "string names of other nodes.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9c69907a",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Node:\n",
    "    op: Op\n",
    "    inputs: List[str]\n",
    "    outputs: List[str]\n",
    "    params: Dict[str, Any] = field(default_factory=dict)\n",
    "\n",
    "    def __str__(self):\n",
    "        outputs_str = \", \".join(self.outputs)\n",
    "        outputs_str += \" = \" if self.outputs else \"\"\n",
    "        inputs_str = \", \".join(self.inputs)\n",
    "        params_str = \", \" if self.inputs and self.params else \"\"\n",
    "        params_str += \", \".join(f\"{k}={v}\" for k, v in self.params.items())\n",
    "        return f\"{outputs_str}{self.op}({inputs_str}{params_str})\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "891f0fdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Graph:\n",
    "    nodes: List[Node] = field(default_factory=list)\n",
    "\n",
    "    def __str__(self):\n",
    "        return \"\\n\".join(str(n) for n in self.nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6b849b61",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tuplify(outs):\n",
    "    if outs is None:\n",
    "        return tuple()\n",
    "    else:\n",
    "        return (outs,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "26bf43f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def interp_node(rules, n: Node, env: Dict[str, Any]):\n",
    "    try:\n",
    "        args = tuple(env[i] for i in n.inputs)\n",
    "        outs = tuplify(rules[n.op](*args, **n.params))\n",
    "        assert len(outs) == len(n.outputs)\n",
    "    except Exception:\n",
    "        print(f\"Failed during: {n}\")\n",
    "        print(\"\\n\".join(f\"{k} = {v}\" for k, v in env.items()))\n",
    "        raise\n",
    "    for k, v in zip(n.outputs, outs):\n",
    "        env[k] = v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7c524aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def interp(graph, rules, print_env=False):\n",
    "    env = {}\n",
    "    for n in graph.nodes:\n",
    "        interp_node(rules, n, env)\n",
    "    if print_env:\n",
    "        for k, v in env.items():\n",
    "            print(f\"{k} = {v}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98f01fa7",
   "metadata": {},
   "source": [
    "We can show it all works with a tiny example constructing the graph\n",
    "manually by ourself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "75fd409b",
   "metadata": {},
   "outputs": [],
   "source": [
    "g = Graph()\n",
    "g.nodes.append(Node(var_placeholder, [], [\"a\"], {\"name\": \"a\"}))\n",
    "g.nodes.append(Node(var_placeholder, [], [\"b\"], {\"name\": \"b\"}))\n",
    "g.nodes.append(Node(var_add, [\"a\", \"b\"], [\"r\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "935bb3c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a = var_placeholder(name=a)\n",
      "b = var_placeholder(name=b)\n",
      "r = var_add(a, b)\n"
     ]
    }
   ],
   "source": [
    "print(g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5aa5f1af",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a = tensor([-0.4152,  1.0221,  0.1116, -0.5071])\n",
      "b = tensor([ 0.2329, -1.1877, -1.0423,  0.0415])\n",
      "r = tensor([-0.1823, -0.1656, -0.9307, -0.4656])\n"
     ]
    }
   ],
   "source": [
    "interp(g, concrete_rules(a=torch.randn(4), b=torch.randn(4)), print_env=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cc4470a",
   "metadata": {},
   "source": [
    "OK, let's write a tracer.  Like before, we will maintain a global\n",
    "graph we are tracing into, and write nodes into this graph (printing\n",
    "them as we go, because why not.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "df923cae",
   "metadata": {},
   "outputs": [],
   "source": [
    "CURRENT_GRAPH = Graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f76b32cb",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def record(r, op, *args, **kwargs):\n",
    "    n = Node(op, [a.name for a in args], [a.name for a in tuplify(r)], kwargs)\n",
    "    print(n)\n",
    "    CURRENT_GRAPH.nodes.append(n)\n",
    "    return r"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "476a621e",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "The objects that we will perform tracing with, however, are going to\n",
    "operate a bit differently.  Because we are in the context of a\n",
    "concrete execution of a PyTorch program,  all of the \"proxy\" objects\n",
    "we are tracing with will actually have real values that correspond to\n",
    "what the untraced PyTorch program would have produced.  Our job\n",
    "is to \"guard\" accesses to the real values, so that we never let\n",
    "the Python program observe that an int was actually 1 unless we record\n",
    "that observance (with a `bool_bailout`).\n",
    "\n",
    "All of our proxy objects will have this structure, so we will call\n",
    "this a \"Guarded\" object.  Guarded objects have an actual concrete value,\n",
    "as well as a name saying how to reference them in the current graph\n",
    "trace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1437f518",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "class Guarded:\n",
    "    name: str\n",
    "    value: Any\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"{self.name}~{self.value}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d480daa1",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "And then we will subclass Guarded for each type we support tracing in\n",
    "our system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "cb67846c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GuardedBool(Guarded):\n",
    "    value: bool\n",
    "\n",
    "    def __init__(self, value, name=None):\n",
    "        self.value = value\n",
    "        self.name = name or fresh_int()\n",
    "\n",
    "    # The conversion to actual Python bool is when a user is actually\n",
    "    # going to observe a value (ostensibly because they're doing a\n",
    "    # condition on it).  So we must record a bailout here, saying that\n",
    "    # on subsequent executions of this trace, the value of this boolean\n",
    "    # node in the graph must match the expected value we saw initially.\n",
    "    def __bool__(self):\n",
    "        record(None, bool_bailout, self, expect=self.value)\n",
    "        return self.value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "626dbdd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GuardedInt(Guarded):\n",
    "    name: str\n",
    "    value: int\n",
    "\n",
    "    def __init__(self, value, name=None):\n",
    "        self.value = value\n",
    "        self.name = name or fresh_int()\n",
    "\n",
    "    def __eq__(self, other):\n",
    "        if isinstance(other, int):\n",
    "            # Promote the int into a node\n",
    "            other = record(GuardedInt(other), int_const, val=other)\n",
    "        # Peephole optimization\n",
    "        # if self.name == other.name:\n",
    "        #     return record(GuardedBool(True), bool_const, val=True)\n",
    "        return record(GuardedBool(self.value == other.value), int_eq, self, other)\n",
    "\n",
    "    @staticmethod\n",
    "    def placeholder(value: int, name: str = None):\n",
    "        r = GuardedInt(value, name)\n",
    "        return record(r, int_placeholder, name=r.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "a4ff6c68",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GuardedSize(Guarded):\n",
    "    name: str\n",
    "    value: List[int]\n",
    "\n",
    "    def __init__(self, value, name=None):\n",
    "        self.value = value\n",
    "        self.name = name or fresh_size()\n",
    "\n",
    "    @staticmethod\n",
    "    def make(value):\n",
    "        return record(GuardedSize([v.value for v in value]), size_ctor, *value)\n",
    "\n",
    "    def __eq__(self, other):\n",
    "        if len(self) != len(other):\n",
    "            return False\n",
    "        for s, o in zip(self, other):\n",
    "            if s != o:\n",
    "                return False\n",
    "        return True\n",
    "\n",
    "    def __len__(self):\n",
    "        # For simplicity, we have kept everything rank specialized, so\n",
    "        # we are allowed to return a raw integer here.  However, if this\n",
    "        # was not OK, we could also return a GuardedInt here (and this\n",
    "        # is in fact what FX does.) in fact what FX does.\n",
    "        return len(self.value)\n",
    "\n",
    "    def __getitem__(self, index: int):\n",
    "        return record(GuardedInt(self.value[index]), size_index, self, index=index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ca3f0e88",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "class Variable(Guarded):\n",
    "    name: str\n",
    "    value: torch.Tensor\n",
    "    _shape = None\n",
    "\n",
    "    @property\n",
    "    def shape(self):\n",
    "        # I don't want to spam the graph with repeated retrievals of the\n",
    "        # size from a tensor, so we will only ever record this retrieval\n",
    "        # once (and it is done lazily, on the first time you access\n",
    "        # shape.)\n",
    "        if self._shape is None:\n",
    "            self._shape = record(GuardedSize(self.value.shape), var_size, self)\n",
    "        return self._shape\n",
    "\n",
    "    def __init__(self, value: torch.Tensor, name: str = None):\n",
    "        self.value = value\n",
    "        self.name = name or fresh_var()\n",
    "\n",
    "    @staticmethod\n",
    "    def placeholder(value: torch.Tensor, name: str = None):\n",
    "        r = Variable(value, name)\n",
    "        return record(r, var_placeholder, name=r.name)\n",
    "\n",
    "    def dim(self):\n",
    "        # choice!\n",
    "        return self.value.dim()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "401b47a6",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "With this change, we can write normal looking code, including\n",
    "conditions on shapes, which we expect to be able to trace through."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "adbe5fc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def broadcast(lhs: List[GuardedInt], rhs: List[GuardedInt]) -> List[GuardedInt]:\n",
    "    r = []\n",
    "    for x, y in itertools.zip_longest(\n",
    "        reversed(lhs.shape), reversed(rhs.shape), fillvalue=1\n",
    "    ):\n",
    "        if x == 1:\n",
    "            r.append(y)\n",
    "        elif y == 1:\n",
    "            r.append(x)\n",
    "        else:\n",
    "            assert x == y\n",
    "            r.append(x)\n",
    "    return GuardedSize.make(list(reversed(r)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "841369e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Elide expands when the sizes match\n",
    "def expand(self: Variable, size: GuardedSize) -> Variable:\n",
    "    if self.shape == size:\n",
    "        return self\n",
    "    else:\n",
    "        return record(Variable(self.value.expand(size.value)), var_expand, self, size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "a0c53791",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add(self: Variable, rhs: Variable) -> Variable:\n",
    "    shape = broadcast(self, rhs)\n",
    "    self_expanded = expand(self, shape)\n",
    "    rhs_expanded = expand(rhs, shape)\n",
    "    return record(\n",
    "        Variable(self_expanded.value + rhs_expanded.value),\n",
    "        var_add,\n",
    "        self_expanded,\n",
    "        rhs_expanded,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd51e237",
   "metadata": {},
   "source": [
    "Let's take a look at the trace produced by this code, and see\n",
    "in particular what bailouts got produced."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "88580dc2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a = var_placeholder(name=a)\n",
      "b = var_placeholder(name=b)\n",
      "s0 = var_size(a)\n",
      "s1 = var_size(b)\n",
      "i0 = size_index(s0, index=0)\n",
      "i1 = size_index(s1, index=0)\n",
      "i2 = int_const(val=1)\n",
      "i3 = int_eq(i0, i2)\n",
      "bool_bailout(i3, expect=False)\n",
      "i4 = int_const(val=1)\n",
      "i5 = int_eq(i1, i4)\n",
      "bool_bailout(i5, expect=False)\n",
      "i6 = int_eq(i0, i1)\n",
      "bool_bailout(i6, expect=True)\n",
      "s2 = size_ctor(i0)\n",
      "i7 = size_index(s0, index=0)\n",
      "i8 = size_index(s2, index=0)\n",
      "i9 = int_eq(i7, i8)\n",
      "bool_bailout(i9, expect=True)\n",
      "i10 = size_index(s1, index=0)\n",
      "i11 = size_index(s2, index=0)\n",
      "i12 = int_eq(i10, i11)\n",
      "bool_bailout(i12, expect=True)\n",
      "v0 = var_add(a, b)\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "a = Variable.placeholder(torch.randn(4), \"a\")\n",
    "b = Variable.placeholder(torch.randn(4), \"b\")\n",
    "v0 = add(a, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a98087ea",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "There is a lot of code generated and a lot of bailouts, and it's\n",
    "kind of hard to see what is going on.  The first three bailouts result\n",
    "from the broadcasting test: we have to check if each of the sizes\n",
    "are one (that's two bailouts), and then we check if the sizes are\n",
    "equal (the third bailout).  Then, when we run expand on the shape,\n",
    "there is another equality test between the computed broadcasted\n",
    "shape and the size, which results in the last two bailouts.\n",
    "\n",
    "The graph here is quite ugly, so let's try to clean it up a little.\n",
    "Because this graph has no data-dependent control flow, we can\n",
    "recompute the bailouts in terms of integer computations ONLY, by\n",
    "running a slightly different interpreter which maps tensors to\n",
    "their (possibly dynamic) shapes, but otherwise works the same way\n",
    "as before (we'll also assume that everything is well typed, which\n",
    "it is, assuming the bailouts in the original program are sufficient)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "bd015837",
   "metadata": {},
   "outputs": [],
   "source": [
    "def shape_rules(**params):\n",
    "    # bool -> GuardedBool\n",
    "    # int -> GuardedInt\n",
    "    # size -> GuardedSize\n",
    "    # var -> GuardedSize !!!\n",
    "    return {\n",
    "        bool_bailout: lambda b, *, expect: record(\n",
    "            None, bool_bailout, b, expect=b.value\n",
    "        ),\n",
    "        bool_const: lambda *, val: val,\n",
    "        int_eq: lambda a, b: a == b,\n",
    "        int_const: lambda *, val: val,\n",
    "        int_placeholder: lambda *, name: params[name],\n",
    "        size_index: lambda s, *, index: s[index],\n",
    "        size_ctor: lambda *args: args,\n",
    "        var_placeholder: lambda *, name: params[name],\n",
    "        var_add: lambda a, b: a,\n",
    "        var_mul: lambda a, b: a,\n",
    "        var_size: lambda a: a,\n",
    "        var_expand: lambda a, size: size,\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98980b69",
   "metadata": {},
   "source": [
    "Let's save our current trace, and reset the context for the new\n",
    "trace we are construct by interpret our original trace with\n",
    "the shape rules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "cc962d2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = CURRENT_GRAPH\n",
    "reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17281755",
   "metadata": {},
   "source": [
    "Now we can see that we get a graph with only integer/bool operations\n",
    "in it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "6871faa6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a0 = int_placeholder(name=a0)\n",
      "b0 = int_placeholder(name=b0)\n",
      "i0 = int_const(val=1)\n",
      "i1 = int_eq(a0, i0)\n",
      "bool_bailout(i1, expect=False)\n",
      "i2 = int_const(val=1)\n",
      "i3 = int_eq(b0, i2)\n",
      "bool_bailout(i3, expect=False)\n",
      "i4 = int_eq(a0, b0)\n",
      "bool_bailout(i4, expect=True)\n",
      "i5 = int_eq(a0, a0)\n",
      "bool_bailout(i5, expect=True)\n",
      "i6 = int_eq(b0, a0)\n",
      "bool_bailout(i6, expect=True)\n"
     ]
    }
   ],
   "source": [
    "interp(\n",
    "    graph,\n",
    "    shape_rules(\n",
    "        a=[GuardedInt.placeholder(4, \"a0\")], b=[GuardedInt.placeholder(4, \"b0\")]\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e5743d1",
   "metadata": {},
   "source": [
    "We can run this graph, asking if a new set of concrete sizes\n",
    "is valid or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "d448bba2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Failed during: bool_bailout(i4, expect=True)\n",
      "a0 = 4\n",
      "b0 = 8\n",
      "i0 = 1\n",
      "i1 = False\n",
      "i2 = 1\n",
      "i3 = False\n",
      "i4 = False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Traceback (most recent call last):\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_89442/649004965.py\", line 2, in <cell line: 1>\n",
      "    interp(CURRENT_GRAPH, concrete_rules(a0=4, b0=8))\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_89442/791977129.py\", line 4, in interp\n",
      "    interp_node(rules, n, env)\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_89442/4070304308.py\", line 4, in interp_node\n",
      "    outs = tuplify(rules[n.op](*args, **n.params))\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_89442/2500127150.py\", line 7, in <lambda>\n",
      "    bool_bailout: lambda b, *, expect: assert_(b == expect),\n",
      "  File \"/var/folders/11/bcmcs8d57q7dxbysb4w_h1ym0000gn/T/ipykernel_89442/1101229770.py\", line 2, in assert_\n",
      "    assert b\n",
      "AssertionError\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    interp(CURRENT_GRAPH, concrete_rules(a0=4, b0=8))\n",
    "except Exception:\n",
    "    traceback.print_exc()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19406a70",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "We can also write a little printer for our trace to say what our\n",
    "guards should be."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "c428c5b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_rules():\n",
    "    return {\n",
    "        bool_bailout: lambda b, *, expect: print(\n",
    "            f\"assert {b}\" if expect else f\"assert not {b}\"\n",
    "        ),\n",
    "        bool_const: lambda *, val: str(val),\n",
    "        int_eq: lambda a, b: f\"({a} == {b})\",\n",
    "        int_const: lambda *, val: str(val),\n",
    "        int_placeholder: lambda *, name: name,\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "e1ab515a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assert not (a0 == 1)\n",
      "assert not (b0 == 1)\n",
      "assert (a0 == b0)\n",
      "assert (a0 == a0)\n",
      "assert (b0 == a0)\n"
     ]
    }
   ],
   "source": [
    "interp(CURRENT_GRAPH, print_rules())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f44524ca",
   "metadata": {},
   "source": [
    "Obviously there is some redundancy, but you can write a little\n",
    "optimizer to clean it up, or send it to your favorite symbolic\n",
    "reasoning engine.  For example, you can see that the constraints\n",
    "here are only equalities on one, and equalities between items;\n",
    "so we could easily use unification to get these into canonical form.\n",
    "\n",
    "TODO: Relate this to symbolic dynamic sizes which prevent constraints\n",
    "from being generated."
   ]
  }
 ],
 "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
}