{
 "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data([[[0., 0., 0.],\n",
       "       [0., 0., 0.],\n",
       "       [0., 0., 0.]],\n",
       "\n",
       "      [[0., 0., 0.],\n",
       "       [0., 0., 0.],\n",
       "       [0., 0., 0.]]], dtype=float32)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from devito import Eq, Grid, TimeFunction, Operator\n",
    "\n",
    "grid = Grid(shape=(3, 3))\n",
    "u = TimeFunction(name='u', grid=grid)\n",
    "u.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we just defined a time-varying grid with 3x3 points in each of the space `Dimension`s _x_ and _y_. `u.data` gives us access to the actual data values. In particular, `u.data[0, :, :]` holds the data values at timestep `t=0`, whereas `u.data[1, :, :]` holds the data values at timestep `t=1`.\n",
    "\n",
    "We now create an `Operator` that increments by 1 all points in the computational domain.\n",
    "Note: we disable all loop-level optimizations to avoid obfuscating the IET structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from devito import configuration\n",
    "configuration['openmp'] = 0  # We don't care about OpenMP in this example\n",
    "\n",
    "eq = Eq(u.forward, u+1)\n",
    "op = Operator(eq, dle='noop')"
   ]
  },
  {
   "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 user-provided SymPy expressions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle u{\\left(t + dt,x,y \\right)} = u{\\left(t,x,y \\right)} + 1$"
      ],
      "text/plain": [
       "Eq(u(t + dt, x, y), u(t, x, y) + 1)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "op.args['expressions']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we print `op`, we can see how the generated code looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#define _POSIX_C_SOURCE 200809L\n",
      "#include \"stdlib.h\"\n",
      "#include \"math.h\"\n",
      "#include \"sys/time.h\"\n",
      "\n",
      "struct dataobj\n",
      "{\n",
      "  void *restrict data;\n",
      "  int * size;\n",
      "  int * npsize;\n",
      "  int * dsize;\n",
      "  int * hsize;\n",
      "  int * hofs;\n",
      "  int * oofs;\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, struct profiler * timers, const int x_M, const int x_m, const int y_M, const int y_m)\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",
      "  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",
      "    struct timeval start_section0, end_section0;\n",
      "    gettimeofday(&start_section0, NULL);\n",
      "    /* Begin 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",
      "    /* End section0 */\n",
      "    gettimeofday(&end_section0, NULL);\n",
      "    timers->section0 += (double)(end_section0.tv_sec-start_section0.tv_sec)+(double)(end_section0.tv_usec-start_section0.tv_usec)/1000000;\n",
      "  }\n",
      "  return 0;\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(op)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As shown in other tutorials, we can JIT-compile and run `op` by executing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Operator `Kernel` run in 0.00 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Data([[[2., 2., 2.],\n",
       "       [2., 2., 2.],\n",
       "       [2., 2., 2.]],\n",
       "\n",
       "      [[3., 3., 3.],\n",
       "       [3., 3., 3.],\n",
       "       [3., 3., 3.]]], dtype=float32)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "op.apply(time=2)\n",
    "u.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the first time we are calling `apply` on `op`, so the `Operator` code gets written in a `.c` file and compiled into a library (a `.so` shared object if on Linux).\n",
    "\n",
    "`op` runs with the user-provided `time_M=2` and with the default arguments `u=u, x_m=0, x_M=2, y_m=0, y_M=2,` and `time_m=0`. The indices `*_m` and `*_M` represent the minimum and maximum iteration points along `Dimension` `*`.\n",
    "\n",
    "`op` can be used for a completely different run, for example providing a new `TimeFunction`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Data([[[0., 0., 0.],\n",
       "       [0., 0., 0.],\n",
       "       [0., 0., 0.]],\n",
       "\n",
       "      [[0., 0., 0.],\n",
       "       [0., 0., 0.],\n",
       "       [0., 0., 0.]]], dtype=float32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u2 = TimeFunction(name='u', grid=grid)\n",
    "u2.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Any of the `Operator` default arguments may be replaced by passing suitable key-value parameters.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Operator `Kernel` run in 0.00 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Data([[[0., 0., 0.],\n",
       "       [4., 4., 0.],\n",
       "       [4., 4., 0.]],\n",
       "\n",
       "      [[0., 0., 0.],\n",
       "       [3., 3., 0.],\n",
       "       [3., 3., 0.]]], dtype=float32)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "op.apply(u=u2, x_m=1, x_M=2, y_m=0, y_M=1, time_M=3)\n",
    "u2.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note, however, that there is no need for recompilation. JIT compilation occurs only once, triggered by the first call to `op.apply()`.\n",
    "\n",
    "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 analyses the IET to determine key computational properties, such as _sequential_, _parallel_, and _vectorizable_. These properties are attached directly to the relevant IET nodes.\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": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Callable Kernel>\n",
      "  <List (0, 2, 0)>\n",
      "\n",
      "    <ArrayCast>\n",
      "    <List (0, 1, 0)>\n",
      "\n",
      "      <[affine,sequential,wrappable] Iteration time::time::(time_m, time_M, 1)::(0, 0)>\n",
      "        <TimedList (2, 1, 2)>\n",
      "          <C.Statement struct timeval start_section0, end_section0;>\n",
      "          <C.Statement gettimeofday(&start_section0, NULL);>\n",
      "          <Section (1)>\n",
      "\n",
      "            <[affine,parallel] Iteration x::x::(x_m, x_M, 1)::(0, 0)>\n",
      "              <[affine,parallel,vector-dim] Iteration y::y::(y_m, y_M, 1)::(0, 0)>\n",
      "                <ExpressionBundle (1)>\n",
      "\n",
      "                  <Expression u[t1, x + 1, y + 1] = u[t0, x + 1, y + 1] + 1>\n",
      "\n",
      "\n",
      "          <C.Statement gettimeofday(&end_section0, NULL);>\n",
      "          <C.Statement timers->section0 += (double)(end_section0.tv_sec-start_section0.tv_sec)+(double)(end_section0.tv_usec-start_section0.tv_usec)/1000000;>\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 a `body`, which includes the children IET nodes. Here, the body is a `List` object.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['#define _POSIX_C_SOURCE 200809L']"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "op._headers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['stdlib.h', 'math.h', 'sys/time.h']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "op._includes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<List (0, 2, 0)>,)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "op.body"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can inspect the `List` to observe that its immediate children are an `<ArrayCast>` and another `<List>`. We can then proceed with the IET traversal until we locate the user-provided `SymPy` equations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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[0].body[0])  # Printing the ArrayCast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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",
      "  struct timeval start_section0, end_section0;\n",
      "  gettimeofday(&start_section0, NULL);\n",
      "  /* Begin 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",
      "  /* End section0 */\n",
      "  gettimeofday(&end_section0, NULL);\n",
      "  timers->section0 += (double)(end_section0.tv_sec-start_section0.tv_sec)+(double)(end_section0.tv_usec-start_section0.tv_usec)/1000000;\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(op.body[0].body[1][0])  # Printing the List"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we access the `Iteration` representing the time loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<WithProperties[affine,sequential,wrappable]::Iteration time[t0,t1]; (time_m, time_M, 1)>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t_iter = op.body[0].body[1][0].body[0]\n",
    "t_iter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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",
      "  struct timeval start_section0, end_section0;\n",
      "  gettimeofday(&start_section0, NULL);\n",
      "  /* Begin 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",
      "  /* End section0 */\n",
      "  gettimeofday(&end_section0, NULL);\n",
      "  timers->section0 += (double)(end_section0.tv_sec-start_section0.tv_sec)+(double)(end_section0.tv_usec-start_section0.tv_usec)/1000000;\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(t_iter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can inspect the `Iteration` to discover what its iteration bounds are."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(time_m, time_M, 1)"
      ]
     },
     "execution_count": 16,
     "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": 17,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<Expression u[t1, x + 1, y + 1] = u[t0, x + 1, y + 1] + 1>'"
      ]
     },
     "execution_count": 17,
     "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": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<Expression u[t1, x + 1, y + 1] = u[t0, x + 1, y + 1] + 1>'"
      ]
     },
     "execution_count": 18,
     "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",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}