examples/compiler/03_iet-A.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we will show how to access and navigate the Iteration/Expression Tree (IET) rooted in an `Operator`.\n",
    "\n",
    "\n",
    "# Part I - Top Down\n",
    "\n",
    "Let's start with a fairly trivial example. First of all, we disable all performance-related optimizations, to maximize the simplicity of the created IET as well as the readability of the generated code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from devito import configuration\n",
    "configuration['opt'] = 'noop'\n",
    "configuration['language'] = 'C'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we create a `TimeFunction` with 3 points in each of the space `Dimension`s _x_ and _y_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from devito import Grid, TimeFunction\n",
    "\n",
    "grid = Grid(shape=(3, 3))\n",
    "u = TimeFunction(name='u', grid=grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now create an `Operator` that increments by 1 all points in the computational domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from devito import Eq, Operator\n",
    "\n",
    "eq = Eq(u.forward, u+1)\n",
    "op = Operator(eq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An `Operator` is an IET node that can generate, JIT-compile, and run low-level code (e.g., C). Just like all other types of IET nodes, it's got a number of metadata attached. For example, we can query an `Operator` to retrieve the input/output `Function`s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(u(t, x, y),)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "op.input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(u(t, x, y),)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "op.writes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we print `op`, we can see how the generated code looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#define _POSIX_C_SOURCE 200809L\n",
      "#define START(S) struct timeval start_ ## S , end_ ## S ; gettimeofday(&start_ ## S , NULL);\n",
      "#define STOP(S,T) gettimeofday(&end_ ## S, NULL); T->S += (double)(end_ ## S .tv_sec-start_ ## S.tv_sec)+(double)(end_ ## S .tv_usec-start_ ## S .tv_usec)/1000000;\n",
      "\n",
      "#include \"stdlib.h\"\n",
      "#include \"math.h\"\n",
      "#include \"sys/time.h\"\n",
      "\n",
      "struct dataobj\n",
      "{\n",
      "  void *restrict data;\n",
      "  unsigned long * size;\n",
      "  unsigned long * npsize;\n",
      "  unsigned long * dsize;\n",
      "  int * hsize;\n",
      "  int * hofs;\n",
      "  int * oofs;\n",
      "  void * dmap;\n",
      "} ;\n",
      "\n",
      "struct profiler\n",
      "{\n",
      "  double section0;\n",
      "} ;\n",
      "\n",
      "\n",
      "int Kernel(struct dataobj *restrict u_vec, const int time_M, const int time_m, const int x_M, const int x_m, const int y_M, const int y_m, struct profiler * timers)\n",
      "{\n",
      "  float (*restrict u)[u_vec->size[1]][u_vec->size[2]] __attribute__ ((aligned (64))) = (float (*)[u_vec->size[1]][u_vec->size[2]]) u_vec->data;\n",
      "\n",
      "  for (int time = time_m, t0 = (time)%(2), t1 = (time + 1)%(2); time <= time_M; time += 1, t0 = (time)%(2), t1 = (time + 1)%(2))\n",
      "  {\n",
      "    START(section0)\n",
      "    for (int x = x_m; x <= x_M; x += 1)\n",
      "    {\n",
      "      for (int y = y_m; y <= y_M; y += 1)\n",
      "      {\n",
      "        u[t1][x + 1][y + 1] = u[t0][x + 1][y + 1] + 1;\n",
      "      }\n",
      "    }\n",
      "    STOP(section0,timers)\n",
      "  }\n",
      "\n",
      "  return 0;\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(op)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An `Operator` is the root of an IET that typically consists of several nested `Iteration`s and `Expression`s – two other fundamental IET node types. The user-provided SymPy equations are wrapped within `Expressions`. Loop nest embedding such expressions are constructed by suitably nesting `Iterations`.\n",
    "\n",
    "The Devito compiler constructs the IET from a collection of `Cluster`s, which represent a higher-level intermediate representation (not covered in this tutorial).\n",
    "\n",
    "The Devito compiler also attaches to the IET key computational properties, such as _sequential_, _parallel_, and _affine_, which are derived through data dependence analysis.\n",
    "\n",
    "We can print the IET structure of an `Operator`, as well as the attached computational properties, using the utility function `pprint`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Callable Kernel>\n",
      "  <CallableBody <unpacks=0, allocs=0, casts=1, maps=0, objs=0> <unmaps=0, frees=0>>\n",
      "    <List (0, 1, 0)>\n",
      "\n",
      "      <[affine,sequential] Iteration time::time::(time_m, time_M, 1)>\n",
      "        <Section (section0)>\n",
      "\n",
      "          <TimedList (1, 1, 1)>\n",
      "            <[affine,parallel,parallel=] Iteration x::x::(x_m, x_M, 1)>\n",
      "              <[affine,parallel,parallel=] Iteration y::y::(y_m, y_M, 1)>\n",
      "                <ExpressionBundle (1)>\n",
      "\n",
      "                  <Expression u[t1, x + 1, y + 1] = u[t0, x + 1, y + 1] + 1>\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from devito.tools import pprint\n",
    "pprint(op)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, `op` is represented as a `<Callable Kernel>`. Attached to it are metadata, such as `_headers` and `_includes`, as well as the `body`, which includes the children IET nodes. Here, the body is the concatenation of an `PointerCast` and a `List` object.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedSet([('_POSIX_C_SOURCE', '200809L'), ('START(S)', 'struct timeval start_ ## S , end_ ## S ; gettimeofday(&start_ ## S , NULL);'), ('STOP(S,T)', 'gettimeofday(&end_ ## S, NULL); T->S += (double)(end_ ## S .tv_sec-start_ ## S.tv_sec)+(double)(end_ ## S .tv_usec-start_ ## S .tv_usec)/1000000;')])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "op._headers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedSet(['stdlib.h', 'math.h', 'sys/time.h'])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "op._includes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<CallableBody <unpacks=0, allocs=0, casts=1, maps=0, objs=0> <unmaps=0, frees=0>>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "op.body"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can explicitly traverse the `body` until we locate the user-provided `SymPy` equations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "float (*restrict u)[u_vec->size[1]][u_vec->size[2]] __attribute__ ((aligned (64))) = (float (*)[u_vec->size[1]][u_vec->size[2]]) u_vec->data;\n"
     ]
    }
   ],
   "source": [
    "print(op.body.casts[0])  # Printing the PointerCast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "for (int time = time_m, t0 = (time)%(2), t1 = (time + 1)%(2); time <= time_M; time += 1, t0 = (time)%(2), t1 = (time + 1)%(2))\n",
      "{\n",
      "  START(section0)\n",
      "  for (int x = x_m; x <= x_M; x += 1)\n",
      "  {\n",
      "    for (int y = y_m; y <= y_M; y += 1)\n",
      "    {\n",
      "      u[t1][x + 1][y + 1] = u[t0][x + 1][y + 1] + 1;\n",
      "    }\n",
      "  }\n",
      "  STOP(section0,timers)\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(op.body.body[0])  # Printing the actual body"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we access the `Iteration` representing the time loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<WithProperties[affine,sequential]::Iteration time[t0,t1]; (time_m, time_M, 1)>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t_iter = op.body.body[0].body[0]\n",
    "t_iter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can for example inspect the `Iteration` to discover what its iteration bounds are."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(time_m, time_M, 1)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t_iter.limits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And as we keep going down through the IET, we can eventually reach the `Expression` wrapping the user-provided SymPy equation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<Expression u[t1, x + 1, y + 1] = u[t0, x + 1, y + 1] + 1>'"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expr = t_iter.nodes[0].body[0].body[0].nodes[0].nodes[0].body[0]\n",
    "expr.view"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of course, there are mechanisms in place to, for example, find all `Expression`s in a given IET. The Devito compiler has a number of IET visitors, among which `FindNodes`, usable to retrieve all nodes of a particular type. So we easily \n",
    "can get all `Expression`s within `op` as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<Expression u[t1, x + 1, y + 1] = u[t0, x + 1, y + 1] + 1>'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from devito.ir.iet import Expression, FindNodes\n",
    "exprs = FindNodes(Expression).visit(op)\n",
    "exprs[0].view"
   ]
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}