dynamic_strides.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "91bfc7a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy as sp\n",
    "from typing import List\n",
    "from dataclasses import dataclass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "526ec319",
   "metadata": {},
   "source": [
    "In this notebook, we explore how to model contiguity and strides in a\n",
    "universe where we support dynamic shapes.  We don't care about dynamic\n",
    "strides/contiguity per se (we'd be OK with specializing on the input\n",
    "being contiguous, channels-last, etc), but strides and contiguity\n",
    "are *derived* from shapes, so if you have dynamic shapes, you\n",
    "also end up with dynamic strides and contiguity.\n",
    "\n",
    "Let's take a concrete look at this phenomenon in the simplest possible\n",
    "context: a contiguous tensor.  Here is the C++ code which implements\n",
    "computation of contiguous strides for a tensor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0bd29fe4",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"\\n// From c10/util/strides.h\\n// Computes the contiguous strides of a tensor, given its sizes.\\nstatic inline std::vector<typename IntArrayRef::value_type> contiguous_strides(\\n    const IntArrayRef sizes) {\\n  using Int = IntArrayRef::value_type;\\n  const Int dims = static_cast<Int>(sizes.size());\\n\\n  std::vector<Int> strides;\\n\\n  if (dims > 0) {\\n    strides.assign(dims, 0);\\n    // Start by populating the last dimension: its strides is always 1.\\n    strides[dims - 1] = 1;\\n    for (auto i = dims - 2; i >= 0; --i) {\\n      // Strides can't be 0 even if sizes are 0.\\n      strides[i] = strides[i + 1] * std::max(sizes[i + 1], Int{1});\\n    }\\n  }\\n\\n  return strides;\\n}\\n\""
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "// From c10/util/strides.h\n",
    "// Computes the contiguous strides of a tensor, given its sizes.\n",
    "static inline std::vector<typename IntArrayRef::value_type> contiguous_strides(\n",
    "    const IntArrayRef sizes) {\n",
    "  using Int = IntArrayRef::value_type;\n",
    "  const Int dims = static_cast<Int>(sizes.size());\n",
    "\n",
    "  std::vector<Int> strides;\n",
    "\n",
    "  if (dims > 0) {\n",
    "    strides.assign(dims, 0);\n",
    "    // Start by populating the last dimension: its strides is always 1.\n",
    "    strides[dims - 1] = 1;\n",
    "    for (auto i = dims - 2; i >= 0; --i) {\n",
    "      // Strides can't be 0 even if sizes are 0.\n",
    "      strides[i] = strides[i + 1] * std::max(sizes[i + 1], Int{1});\n",
    "    }\n",
    "  }\n",
    "\n",
    "  return strides;\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c40f3277",
   "metadata": {},
   "source": [
    "And a port to Python:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5e9eab47",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def contiguous_strides(sizes: List[int]):\n",
    "    dims = len(sizes)\n",
    "    strides = []\n",
    "    if dims > 0:\n",
    "        strides = [0] * dims\n",
    "        strides[dims - 1] = 1\n",
    "        for i in range(dims - 2, -1, -1):\n",
    "            strides[i] = strides[i + 1] * sp.Max(sizes[i + 1], 1)\n",
    "    return strides"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "98350f5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[15, 5, 1]\n"
     ]
    }
   ],
   "source": [
    "print(contiguous_strides([2, 3, 5]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b0ae57f",
   "metadata": {},
   "source": [
    "Let's look at the symbolic output of this function.  When only the batch\n",
    "dimension is dynamic, things are pretty simple:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e5b2e645",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[15, 5, 1]\n"
     ]
    }
   ],
   "source": [
    "x = sp.symbols('x')\n",
    "print(contiguous_strides([x, 3, 5]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e926e5c6",
   "metadata": {},
   "source": [
    "However, if an inner dimension is dynamic, the dynamic shape variable\n",
    "shows up in the stride calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d07e23d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[5*Max(1, x), 5, 1]\n"
     ]
    }
   ],
   "source": [
    "print(contiguous_strides([2, x, 5]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0f51c99",
   "metadata": {},
   "source": [
    "The set of strides returned by contiguous is guaranteed to be\n",
    "contiguous, but the inverse is not true: there are some degrees of\n",
    "freedom in the definition of strides when sizes are one or zero.\n",
    "Here is our definition of \"when something is contiguous\" (not accounting\n",
    "for overflow):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ac9205c3",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\n// In c10/core/TensorImpl.h\\ninline bool is_empty() const {\\n  return numel() == 0;\\n}\\n\\n// In c10/core/TensorImpl.cpp\\nbool TensorImpl::compute_contiguous() const {\\n  bool is_contiguous = true;\\n  if (is_empty())\\n    return is_contiguous;\\n  int64_t z = 1;\\n  for (int64_t d = dim() - 1; d >= 0; d--) {\\n    const auto size_d = sizes_and_strides_.size_at_unchecked(d);\\n    if (size_d != 1) {\\n      if (sizes_and_strides_.stride_at_unchecked(d) == z) {\\n        z *= size_d;\\n      } else {\\n        is_contiguous = false;\\n        break;\\n      }\\n    }\\n  }\\n  return is_contiguous;\\n}\\n'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "// In c10/core/TensorImpl.h\n",
    "inline bool is_empty() const {\n",
    "  return numel() == 0;\n",
    "}\n",
    "\n",
    "// In c10/core/TensorImpl.cpp\n",
    "bool TensorImpl::compute_contiguous() const {\n",
    "  bool is_contiguous = true;\n",
    "  if (is_empty())\n",
    "    return is_contiguous;\n",
    "  int64_t z = 1;\n",
    "  for (int64_t d = dim() - 1; d >= 0; d--) {\n",
    "    const auto size_d = sizes_and_strides_.size_at_unchecked(d);\n",
    "    if (size_d != 1) {\n",
    "      if (sizes_and_strides_.stride_at_unchecked(d) == z) {\n",
    "        z *= size_d;\n",
    "      } else {\n",
    "        is_contiguous = false;\n",
    "        break;\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "  return is_contiguous;\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a680f77",
   "metadata": {},
   "source": [
    "In Python (note that we will use the suffix branchy to refer\n",
    "to code which branches on the concrete value of sizes/strides):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e2fbcbf7",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def compute_numel(sizes: List[int]):\n",
    "    numel = 1\n",
    "    for s in sizes:\n",
    "        numel *= s\n",
    "    return numel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5c33f0b9",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def compute_contiguous_branchy(sizes: List[int], strides: List[int]):\n",
    "    is_contiguous = True\n",
    "    if compute_numel(sizes) == 0:\n",
    "        return is_contiguous\n",
    "    z = 1\n",
    "    for d in range(len(sizes) - 1, -1, -1):\n",
    "        if sizes[d] != 1:\n",
    "            if strides[d] == z:\n",
    "                z *= sizes[d]\n",
    "            else:\n",
    "                is_contiguous = False\n",
    "                break\n",
    "    return is_contiguous"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5ac31e9",
   "metadata": {},
   "source": [
    "When a dimension has size 1, we are indifferent to the stride at that\n",
    "dimension:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f81e2e36",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[5, 5, 1]\n"
     ]
    }
   ],
   "source": [
    "print(contiguous_strides([3, 1, 5]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "aabd727c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "print(compute_contiguous_branchy([3, 1, 5], [5, 5, 1]))\n",
    "print(compute_contiguous_branchy([3, 1, 5], [5, 999999, 1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c71a8bfb",
   "metadata": {},
   "source": [
    "When a tensor contains zero elements, we are indifferent to all the\n",
    "strides"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "33533558",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[5, 5, 1]\n"
     ]
    }
   ],
   "source": [
    "print(contiguous_strides([3, 0, 5]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4efb2274",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "print(compute_contiguous_branchy([3, 0, 5], [5, 5, 1]))\n",
    "print(compute_contiguous_branchy([3, 0, 5], [123456, 999999, 424242]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12a6dc9f",
   "metadata": {},
   "source": [
    "Can we compute_contiguous symbolically?  Unfortunately, the \"branchy\"\n",
    "implementation, as written above cannot be run directly on SymPy\n",
    "integers, as in several points in the code we condition on the\n",
    "concrete values of various comparisons on integers.  Fortunately,\n",
    "we can introduce a SymInt/SymBool abstraction (as done in previous\n",
    "notebooks) to provide concrete values and record guards expressing\n",
    "what is required to be true for the computation to be correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f7bb62b0",
   "metadata": {
    "lines_to_end_of_cell_marker": 0,
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "\n",
    "GUARDS = []\n",
    "\n",
    "def is_constant(e):\n",
    "    if hasattr(e, 'is_constant'):\n",
    "        return e.is_constant()\n",
    "    elif e is sp.true or e is sp.false:\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "class SymObject:\n",
    "    def __post_init__(self):\n",
    "        if self.expr is None:\n",
    "            self.expr = sp.sympify(self.val)\n",
    "        elif not isinstance(self.expr, sp.Expr):\n",
    "            self.expr = sp.sympify(self.expr)\n",
    "\n",
    "@dataclass\n",
    "class SymBool(SymObject):\n",
    "    val: bool\n",
    "    expr: sp.Expr = None\n",
    "    guarded: bool = False\n",
    "\n",
    "    def __bool__(self):\n",
    "        if not self.guarded:\n",
    "            self.guarded = True\n",
    "            if not is_constant(self.expr):\n",
    "                if self.val:\n",
    "                    GUARDS.append(self.expr)\n",
    "                else:\n",
    "                    GUARDS.append(sp.Not(self.expr))\n",
    "        return self.val\n",
    "\n",
    "def logical_and(self: bool, other: bool):\n",
    "    if isinstance(self, SymBool) and isinstance(other, SymBool):\n",
    "        return SymBool(self.val and other.val, sp.And(self.expr, other.expr))\n",
    "    return sp.And(self, other)\n",
    "\n",
    "def logical_or(self: bool, other: bool):\n",
    "    if isinstance(self, SymBool) and isinstance(other, SymBool):\n",
    "        return SymBool(self.val or other.val, sp.Or(self.expr, other.expr))\n",
    "    return sp.Or(self, other)\n",
    "\n",
    "@dataclass\n",
    "class SymInt(SymObject):\n",
    "    val: int\n",
    "    expr: sp.Expr = None\n",
    "    guarded: bool = False\n",
    "\n",
    "    def __int__(self):\n",
    "        if not self.guarded:\n",
    "            self.guarded = True\n",
    "            if not is_constant(self.expr):\n",
    "                GUARDS.append(self.Eq(self.expr, self.val).simplify())\n",
    "        return self.val\n",
    "\n",
    "    def __eq__(self, other):\n",
    "        if not isinstance(other, SymInt):\n",
    "            other = SymInt(other)\n",
    "        return SymBool(self.val == other.val, sp.Eq(self.expr, other.expr))\n",
    "\n",
    "    def __ne__(self, other):\n",
    "        if not isinstance(other, SymInt):\n",
    "            other = SymInt(other)\n",
    "        return SymBool(self.val != other.val, sp.Ne(self.expr, other.expr))\n",
    "\n",
    "    def __mul__(self, other):\n",
    "        if not isinstance(other, SymInt):\n",
    "            other = SymInt(other)\n",
    "        return SymInt(self.val * other.val, sp.Mul(self.expr, other.expr))\n",
    "\n",
    "    def __rmul__(self, other):\n",
    "        if not isinstance(other, SymInt):\n",
    "            other = SymInt(other)\n",
    "        return SymInt(self.val * other.val, sp.Mul(self.expr, other.expr))\n",
    "\n",
    "def I(val, expr=None):\n",
    "    return SymInt(val, expr)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6f4db58",
   "metadata": {},
   "source": [
    "Let's run our example.  Under the guards model, we must provide\n",
    "concrete values for every symbolic integer, so we can resolve\n",
    "conditionals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ead6141f",
   "metadata": {},
   "outputs": [],
   "source": [
    "x1, x2, x3, y1, y2, y3 = sp.symbols(\"x1 x2 x3 y1 y2 y3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "df4eb309",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "GUARDS.clear()\n",
    "print(compute_contiguous_branchy(\n",
    "    [I(3, x1), I(1, x2), I(5, x3)],\n",
    "    [I(5, y1), I(99999, y2), I(1, y3)]\n",
    "))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "621b3252",
   "metadata": {},
   "source": [
    "We see that this tensor is contiguous..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ffa5c60c",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Ne(x1*x2*x3, 0), Ne(x3, 1), Eq(y3, 1), Eq(x2, 1), Ne(x1, 1), Eq(y1, x3)]\n"
     ]
    }
   ],
   "source": [
    "print(GUARDS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6710018d",
   "metadata": {},
   "source": [
    "...subject to these conditions.  These conditions say which particular\n",
    "path through the loop we took: we require the sizes to be nonzero,\n",
    "there are number of size one equalities/disequalities, and the\n",
    "equality requirement between y1 and x3 is the \"true\" contiguity\n",
    "requirement."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3ba4e98",
   "metadata": {},
   "source": [
    "If we are willing to rewrite the definition of compute contiguous, we\n",
    "can eliminate the branches, giving a symbolic expression with no\n",
    "guards."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2e66f468",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "def compute_contiguous(sizes, strides):\n",
    "    is_contiguous = True\n",
    "    z = 1\n",
    "    for d in range(len(sizes) - 1, -1, -1):\n",
    "        is_contiguous = logical_and(\n",
    "            is_contiguous,\n",
    "            logical_or(\n",
    "                sp.Eq(sizes[d], 1),\n",
    "                sp.Eq(strides[d], z)\n",
    "            )\n",
    "        )\n",
    "        z *= sizes[d]\n",
    "    return logical_or(sp.Eq(compute_numel(sizes), 0), is_contiguous)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09b49776",
   "metadata": {},
   "source": [
    "TODO: prove these two implementations are equivalent, somehow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cc20966",
   "metadata": {},
   "source": [
    "We can see that no matter the choice of the stride for a size one\n",
    "dimension, the result is always contiguous:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "c0874461",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "print(compute_contiguous([3, 1, 5], [5, x, 1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7e6ab8d",
   "metadata": {},
   "source": [
    "And we can see the unflattened contiguity requirement for a completely\n",
    "general size/stride tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "39e58372",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Eq(x1*x2*x3, 0) | ((Eq(x2, 1) | Eq(y2, x3)) & (Eq(x3, 1) | Eq(y3, 1)) & (Eq(x1, 1) | Eq(y1, x2*x3)))\n"
     ]
    }
   ],
   "source": [
    "print(compute_contiguous([x1, x2, x3], [y1, y2, y3]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b04a2af0",
   "metadata": {},
   "source": [
    "There's other stuff too:\n",
    "\n",
    "  - We are not \"just\" compute_contiguous; we also have have variations\n",
    "    of this for every memory layout we support.  So the same exercise\n",
    "    needs to apply everywhere.\n",
    "\n",
    "  - We also have non_overlapping_and_dense which which involves a sort\n",
    "    which is very annoying."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6952548c",
   "metadata": {},
   "source": [
    "In conclusion:\n",
    "\n",
    "  - We have an explicit choice whether or not to branch inside\n",
    "    implementations of code that may be traced.  More trace friendly\n",
    "    code is not as good for eager execution (because you can't do\n",
    "    things like short circuit).\n",
    "\n",
    "  - If we store SymInt inside TensorImpl, we need to make a call about\n",
    "    how we represent the contiguity bits inside Tensor.  These bits\n",
    "    are literally a single bit, so we cannot store a symbolic boolean\n",
    "    in them.  It seems the easiest fix is to ensure the\n",
    "    is_contiguous() is virtualized (it is), and then internally run\n",
    "    (and cache) the symbolic formula done here."
   ]
  }
 ],
 "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
}