{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Working with Loki's internal representation\n",
    "\n",
    "The objective of this notebook is to get an impression how Loki's internal representation (IR) can be traversed, searched and manipulated using the provided visitor utilities.\n",
    "\n",
    "We are again going to work with the `phys_kernel_LITE_LOOP` routine. Let's start by parsing the source file and extracting the routine from it. Note, that we can also directly access the routine using its name, although it is wrapped inside a `Module` object as we have seen in the previous notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Subroutine:: phys_kernel_LITE_LOOP"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from loki import Sourcefile\n",
    "source = Sourcefile.from_file('src/phys_mod.F90')\n",
    "routine_lite_loop = source['phys_kernel_LITE_LOOP']\n",
    "routine_lite_loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to manipulate this routine and want to try two different ways of doing that, so we start by creating a copy. That way, we don't change the original object in the subsequent steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SUBROUTINE phys_kernel_LITE_LOOP (dim1, dim2, i1, i2, in1, in2, in3, in4, in5, in6, in7, in8, in9, in10, out1)\n",
      "  INTEGER(KIND=ip), INTENT(IN) :: dim1, dim2, i1, i2\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: in1, in2, in3, in4, in5, in6, in7, in8, in9, in10\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: out1\n",
      "  \n",
      "  INTEGER(KIND=ip) :: i, k\n",
      "  DO k=1,dim2\n",
      "    DO i=i1,i2\n",
      "      out1(i, k) = (in1(i, k) + in2(i, k) + in3(i, k) + in4(i, k) + in5(i, k) + in6(i, k) + in7(i, k) + in8(i, k) + in9(i, k) +  &\n",
      "      & in10(i, k))*0.1\n",
      "      in1(i, k) = out1(i, k)\n",
      "    END DO\n",
      "  END DO\n",
      "END SUBROUTINE phys_kernel_LITE_LOOP\n"
     ]
    }
   ],
   "source": [
    "routine = routine_lite_loop.clone()\n",
    "print(routine.to_fortran())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The routine body consists of two nested loops. What we want to try first is to change the order of the loops (i.e., have the `i` loop outermost and the `k` loop innermost) but leave the loop body untouched.\n",
    "\n",
    "For that, we first need to find the loops in the IR, which can be done using the [_FindNodes_](https://sites.ecmwf.int/docs/loki/main/loki.visitors.find.html#loki.visitors.find.FindNodes) visitor. As argument to the constructor we provide the node type (or a tuple of multiple types) that we want to look for and call the `visit` method with the tree to search.\n",
    "The visitor traverses the IR and collects all matching nodes into a list that is returned. For our purposes we are interested in the [_Loop_](https://sites.ecmwf.int/docs/loki/main/loki.ir.html#loki.ir.Loop) nodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Loop:: k=1:dim2, Loop:: i=i1:i2]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from loki import FindNodes, Loop\n",
    "loops = FindNodes(Loop).visit(routine.body)\n",
    "loops"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the visitor has found both loops. Next, we create a substitution map - essentially a dictionary that maps the original node to its replacement. To exchange the two loops, we use the outer loop but with the inner loop's body and make it the body of the inner loop:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{Loop:: k=1:dim2: Loop:: i=i1:i2}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outer_loop, inner_loop = loops\n",
    "new_inner_loop = outer_loop.clone(body=inner_loop.body)\n",
    "loop_map = {outer_loop: inner_loop.clone(body=(new_inner_loop,))}\n",
    "loop_map"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the substitution map in place, we can call the [_Transformer_](https://sites.ecmwf.int/docs/loki/main/loki.visitors.transform.html#loki.visitors.transform.Transformer). It takes the map as argument to the constructor and applies it to the control flow tree given to the `visit` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from loki import Transformer\n",
    "routine.body = Transformer(loop_map).visit(routine.body)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is the original routine with the exchanged loop order."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SUBROUTINE phys_kernel_LITE_LOOP (dim1, dim2, i1, i2, in1, in2, in3, in4, in5, in6, in7, in8, in9, in10, out1)\n",
      "  INTEGER(KIND=ip), INTENT(IN) :: dim1, dim2, i1, i2\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: in1, in2, in3, in4, in5, in6, in7, in8, in9, in10\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: out1\n",
      "  \n",
      "  INTEGER(KIND=ip) :: i, k\n",
      "  DO i=i1,i2\n",
      "    DO k=1,dim2\n",
      "      out1(i, k) = (in1(i, k) + in2(i, k) + in3(i, k) + in4(i, k) + in5(i, k) + in6(i, k) + in7(i, k) + in8(i, k) + in9(i, k) +  &\n",
      "      & in10(i, k))*0.1\n",
      "      in1(i, k) = out1(i, k)\n",
      "    END DO\n",
      "  END DO\n",
      "END SUBROUTINE phys_kernel_LITE_LOOP\n"
     ]
    }
   ],
   "source": [
    "reordered_loops = FindNodes(Loop).visit(routine.body)\n",
    "assert len(reordered_loops) == 2\n",
    "assert reordered_loops[0].variable == 'i' and reordered_loops[1].variable == 'k'\n",
    "print(routine.to_fortran())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Next, we want to start again with the original routine and this time keep the loop order as is but reverse the memory layout of all arrays. We start by creating another copy of the original routine and verify that it is indeed the original version without the above transformations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SUBROUTINE phys_kernel_LITE_LOOP (dim1, dim2, i1, i2, in1, in2, in3, in4, in5, in6, in7, in8, in9, in10, out1)\n",
      "  INTEGER(KIND=ip), INTENT(IN) :: dim1, dim2, i1, i2\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: in1, in2, in3, in4, in5, in6, in7, in8, in9, in10\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: out1\n",
      "  \n",
      "  INTEGER(KIND=ip) :: i, k\n",
      "  DO k=1,dim2\n",
      "    DO i=i1,i2\n",
      "      out1(i, k) = (in1(i, k) + in2(i, k) + in3(i, k) + in4(i, k) + in5(i, k) + in6(i, k) + in7(i, k) + in8(i, k) + in9(i, k) +  &\n",
      "      & in10(i, k))*0.1\n",
      "      in1(i, k) = out1(i, k)\n",
      "    END DO\n",
      "  END DO\n",
      "END SUBROUTINE phys_kernel_LITE_LOOP\n"
     ]
    }
   ],
   "source": [
    "routine = routine_lite_loop.clone()\n",
    "print(routine.to_fortran())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time, we want to modify variables instead of loops. Loki uses a two-level internal representation that separates expressions from control flow. This means, the IR that we have worked with so far, is in fact the control flow tree and, nested inside, we have a second tree level as property of certain control flow nodes. For example, the loop bounds of the `Loop` node or the left and right hand side expressions in [_Assignment_](https://sites.ecmwf.int/docs/loki/main/loki.ir.html#loki.ir.Assignment) nodes are such expression trees. The advantage of this is that it makes traversing the control flow tree a lot faster and allows to recurse into expressions only when required.\n",
    "\n",
    "Since we are now looking for variables we need to actually search the expression trees and therefore have to use a different visitor [_FindVariables_](https://sites.ecmwf.int/docs/loki/main/loki.expression.expr_visitors.htloki.expression.expr_visitorssitors.FindVariables). Here, we are only interested in arrays and can further restrict ourselves to [_Array_](https://sites.ecmwf.int/docs/loki/main/loki.expression.symbols.html#loki.expression.symbols.Array) expression nodes. We build again a substitution map with the subscript `dimensions` of the arrays reversed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "in10(i, k) -> in10(k, i)\n",
      "in2(i, k) -> in2(k, i)\n",
      "in3(i, k) -> in3(k, i)\n",
      "in5(i, k) -> in5(k, i)\n",
      "in9(i, k) -> in9(k, i)\n",
      "in7(i, k) -> in7(k, i)\n",
      "out1(i, k) -> out1(k, i)\n",
      "in4(i, k) -> in4(k, i)\n",
      "in1(i, k) -> in1(k, i)\n",
      "in8(i, k) -> in8(k, i)\n",
      "in6(i, k) -> in6(k, i)\n"
     ]
    }
   ],
   "source": [
    "from loki import FindVariables, Array\n",
    "variable_map = {}\n",
    "for var in FindVariables().visit(routine.body):\n",
    "    if isinstance(var, Array) and var.dimensions:\n",
    "        variable_map[var] = var.clone(dimensions=var.dimensions[::-1])\n",
    "print('\\n'.join(f'{a!s} -> {b!s}' for a, b in variable_map.items()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just like we have a separate find utility for expression trees there is a separate transformer [_SubstituteExpressions_](https://sites.ecmwf.int/docs/loki/main/loki.expression.expr_visitors.html#loki.expression.expr_visitors.SubstituteExpressions). Applying this to the routine's body we obtain the following result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SUBROUTINE phys_kernel_LITE_LOOP (dim1, dim2, i1, i2, in1, in2, in3, in4, in5, in6, in7, in8, in9, in10, out1)\n",
      "  INTEGER(KIND=ip), INTENT(IN) :: dim1, dim2, i1, i2\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: in1, in2, in3, in4, in5, in6, in7, in8, in9, in10\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: out1\n",
      "  \n",
      "  INTEGER(KIND=ip) :: i, k\n",
      "  DO k=1,dim2\n",
      "    DO i=i1,i2\n",
      "      out1(k, i) = (in1(k, i) + in2(k, i) + in3(k, i) + in4(k, i) + in5(k, i) + in6(k, i) + in7(k, i) + in8(k, i) + in9(k, i) +  &\n",
      "      & in10(k, i))*0.1\n",
      "      in1(k, i) = out1(k, i)\n",
      "    END DO\n",
      "  END DO\n",
      "END SUBROUTINE phys_kernel_LITE_LOOP\n"
     ]
    }
   ],
   "source": [
    "from loki import SubstituteExpressions\n",
    "routine.body = SubstituteExpressions(variable_map).visit(routine.body)\n",
    "print(routine.to_fortran())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The routine's body is correctly modified, with the array subscript dimensions reversed, but the declarations are still unchanged. For that, we need to change the `shape` of the variables as well as the `dimensions` property of the variable nodes inside the declarations.\n",
    "\n",
    "There are two ways of achieving this: The first and easier way would be to modify the [_variables_ property](https://sites.ecmwf.int/docs/loki/main/loki.subroutine.html#loki.subroutine.Subroutine.variables) of the `Subroutine` object and update all array dimensions and shapes. This automatically recreates declarations for modified variables but inserts separate new declarations for each. Let's try this approach for a copy of the routine:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SUBROUTINE phys_kernel_LITE_LOOP (dim1, dim2, i1, i2, in1, in2, in3, in4, in5, in6, in7, in8, in9, in10, out1)\n",
      "  INTEGER(KIND=ip), INTENT(IN) :: dim1, dim2, i1, i2\n",
      "  \n",
      "  INTEGER(KIND=ip) :: i, k\n",
      "  REAL(KIND=lp), INTENT(INOUT) :: in1(1:dim2, 1:dim1)\n",
      "  REAL(KIND=lp), INTENT(INOUT) :: in2(1:dim2, 1:dim1)\n",
      "  REAL(KIND=lp), INTENT(INOUT) :: in3(1:dim2, 1:dim1)\n",
      "  REAL(KIND=lp), INTENT(INOUT) :: in4(1:dim2, 1:dim1)\n",
      "  REAL(KIND=lp), INTENT(INOUT) :: in5(1:dim2, 1:dim1)\n",
      "  REAL(KIND=lp), INTENT(INOUT) :: in6(1:dim2, 1:dim1)\n",
      "  REAL(KIND=lp), INTENT(INOUT) :: in7(1:dim2, 1:dim1)\n",
      "  REAL(KIND=lp), INTENT(INOUT) :: in8(1:dim2, 1:dim1)\n",
      "  REAL(KIND=lp), INTENT(INOUT) :: in9(1:dim2, 1:dim1)\n",
      "  REAL(KIND=lp), INTENT(INOUT) :: in10(1:dim2, 1:dim1)\n",
      "  REAL(KIND=lp), INTENT(INOUT) :: out1(1:dim2, 1:dim1)\n",
      "  DO k=1,dim2\n",
      "    DO i=i1,i2\n",
      "      out1(k, i) = (in1(k, i) + in2(k, i) + in3(k, i) + in4(k, i) + in5(k, i) + in6(k, i) + in7(k, i) + in8(k, i) + in9(k, i) +  &\n",
      "      & in10(k, i))*0.1\n",
      "      in1(k, i) = out1(k, i)\n",
      "    END DO\n",
      "  END DO\n",
      "END SUBROUTINE phys_kernel_LITE_LOOP\n"
     ]
    }
   ],
   "source": [
    "routine_variant1 = routine.clone()\n",
    "variables = []\n",
    "for var in routine_variant1.variables:\n",
    "    if isinstance(var, Array):\n",
    "        shape = var.shape[::-1]\n",
    "        variables += [var.clone(dimensions=shape, type=var.type.clone(shape=shape))]\n",
    "    else:\n",
    "        variables += [var]\n",
    "routine_variant1.variables = variables\n",
    "print(routine_variant1.to_fortran())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a finger exercise we demonstrate also a second approach that avoids recreating the declarations but modifies them directly. For that, we search for all [_VariableDeclaration_](https://sites.ecmwf.int/docs/loki/main/loki.ir.html#loki.ir.VariableDeclaration) nodes in the routine's specification part (`spec`) and build a substitution map with updated declarations where necessary. This involves updating the list of variables declared in a `VariableDeclaration` node and making sure that only `Array` nodes are modified.\n",
    "\n",
    "Fortran allows to specify array dimensions either using the `DIMENSION` attribute or as dimensions in brackets after the declared symbol's name (e.g., `var(dim1, dim2)`). Loki's default behaviour is the latter (as visible from the auto-generated declarations above). To accommodate both variants in Loki's IR, we allow an optional property `dimensions` on `VariableDeclaration` nodes to produce the syntax of the first. Importantly, in both cases Loki stores the `dimensions` property also on the declared variable nodes to make sure they are always accessible in a uniform way.\n",
    "\n",
    "When building the substitution map for the declaration nodes, we honour both versions and adapt our behaviour accordingly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SUBROUTINE phys_kernel_LITE_LOOP (dim1, dim2, i1, i2, in1, in2, in3, in4, in5, in6, in7, in8, in9, in10, out1)\n",
      "  INTEGER(KIND=ip), INTENT(IN) :: dim1, dim2, i1, i2\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim2, 1:dim1) :: in1, in2, in3, in4, in5, in6, in7, in8, in9, in10\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim2, 1:dim1) :: out1\n",
      "  \n",
      "  INTEGER(KIND=ip) :: i, k\n",
      "  DO k=1,dim2\n",
      "    DO i=i1,i2\n",
      "      out1(k, i) = (in1(k, i) + in2(k, i) + in3(k, i) + in4(k, i) + in5(k, i) + in6(k, i) + in7(k, i) + in8(k, i) + in9(k, i) +  &\n",
      "      & in10(k, i))*0.1\n",
      "      in1(k, i) = out1(k, i)\n",
      "    END DO\n",
      "  END DO\n",
      "END SUBROUTINE phys_kernel_LITE_LOOP\n"
     ]
    }
   ],
   "source": [
    "from loki import VariableDeclaration\n",
    "decl_map = {}\n",
    "for decl in FindNodes(VariableDeclaration).visit(routine.spec):\n",
    "    if decl.dimensions:\n",
    "        shape = decl.dimensions[::-1]\n",
    "        symbols = [var.clone(dimensions=shape, type=var.type.clone(shape=shape)) for var in decl.symbols]\n",
    "        decl_map[decl] = decl.clone(dimensions=shape, symbols=symbols)\n",
    "    elif any(isinstance(var, Array) for var in decl.symbols):\n",
    "        symbols = []\n",
    "        for var in decl.symbols:\n",
    "            if isinstance(var, Array):\n",
    "                shape = var.shape[::-1]\n",
    "                symbols += [var.clone(dimensions=shape, type=var.type.clone(shape=shape))]\n",
    "            else:\n",
    "                symbols += [var]\n",
    "        decl_map[decl] = decl.clone(symbols=symbols)\n",
    "routine.spec = Transformer(decl_map).visit(routine.spec)\n",
    "print(routine.to_fortran())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And with that we have achieved the same result while retaining the compacted notation for declarations. Notably, the body is the same for both variants:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from loki import fgen\n",
    "fcode = fgen(routine.spec)\n",
    "assert '(1:dim2, 1:dim1)' in fcode\n",
    "assert '(1:dim1, 1:dim2)' not in fcode\n",
    "fcode = fgen(routine_variant1.spec)\n",
    "assert '(1:dim2, 1:dim1)' in fcode\n",
    "assert '(1:dim1, 1:dim2)' not in fcode\n",
    "assert fgen(routine.body) == fgen(routine_variant1.body)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For further details on how to work with Loki's internal representation, have a look at the [relevant section in the documentation](https://sites.ecmwf.int/docs/loki/main/visitors.html)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.8 64-bit",
   "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.8"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "49643b39beb1b0a7ebd0b57318d9385a5a724f398f0bc0540a61bbc4360c8e5d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
