10_loop_transformation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loop Transformation\n",
    "\n",
    "在[之前的教程](https://zhuanlan.zhihu.com/p/553703704)第三章节`Schedules`中我们通过一种基于词典序的方法来表示语句执行的顺序, 这种表示方法称为[Kelly表示](https://ieeexplore.ieee.org/document/472180).但是在`Kelly`表示中，每个语向对应的仿射函数的值域维度并不完全相同. 在比较不同语句实例之间的执行顺序时很不方便. 因此后来有人提出使用[Schedule Tree](https://www.researchgate.net/publication/277075424_Polyhedral_AST_Generation_Is_More_Than_Scanning_Polyhedra),是通过一种树状结构来表示语句的执行顺序的形式. 本篇教程主要介绍如何操作调度树来实现《多面体编译理论与深度学习实践》书中所提到的一些循环变换. 在本文中遇到不理解的概念可以参考[Isl Concept](https://zhuanlan.zhihu.com/p/646406752).\n",
    "\n",
    "关于调度树的其他文章可以参考[零基础入门isl的调度树-设定篇](https://zhuanlan.zhihu.com/p/641737915), 以及《多面体编译理论与深度学习实践》.\n",
    "\n",
    "所有教程源码位于[这里](https://github.com/zhen8838/isl_learn).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pet\n",
    "import isl\n",
    "\n",
    "\n",
    "def parse_code(source: str, func_name: str):\n",
    "  with open(\"/tmp/parse_code.c\", \"w\") as f:\n",
    "    f.write(source)\n",
    "  scop = pet.scop.extract_from_C_source(\"/tmp/parse_code.c\", func_name)\n",
    "  context = scop.get_context()\n",
    "  schedule: isl.schedule = scop.get_schedule()\n",
    "  reads = scop.get_may_reads()\n",
    "  writes = scop.get_may_writes()\n",
    "  return (scop, context, schedule, reads, writes)\n",
    "\n",
    "\n",
    "scop, context, schedule, reads, writes = parse_code(\"\"\"\n",
    "void foo()\n",
    "{\n",
    "#pragma scop\n",
    "  const int M = 1024;\n",
    "  const int N = 2048;\n",
    "  int a[M*N];\n",
    "  for (int i = 0; i < M; i++) {\n",
    "    for (int j = 0; j < N; j++) {                                          \n",
    "      a[i*N+j] = i*N+j;\n",
    "    }\n",
    "  }\n",
    "#pragma endscop\n",
    "}\n",
    "\"\"\", \"foo\")\n",
    "\n",
    "\n",
    "class CSource():\n",
    "  def __init__(self, path: str) -> None:\n",
    "    with open(path, 'r') as f:\n",
    "      self.context = f.read()\n",
    "\n",
    "  def _repr_html_(self) -> str:\n",
    "    return \"<pre class='code'><code class=\\\"cpp hljs\\\">\" + self.context + \"</code></pre>\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们可以dump整个schedule tree, 因为默认是按Flow的方式输出, 所以这里将style修改为block:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain: \"{ S_0[]; S_3[]; S_1[]; S_5[i, j] : 0 <= i <= 1023 and 0 <= j <= 2047; S_2[]; S_4[] }\"\n",
      "child:\n",
      "  sequence:\n",
      "  - filter: \"{ S_0[] }\"\n",
      "  - filter: \"{ S_1[] }\"\n",
      "  - filter: \"{ S_2[] }\"\n",
      "  - filter: \"{ S_3[] }\"\n",
      "  - filter: \"{ S_4[] }\"\n",
      "  - filter: \"{ S_5[i, j] }\"\n",
      "    child:\n",
      "      schedule: \"L_0[{ S_5[i, j] -> [(i)] }]\"\n",
      "      child:\n",
      "        schedule: \"L_1[{ S_5[i, j] -> [(j)] }]\"\n",
      "\n"
     ]
    }
   ],
   "source": [
    "def isl_schedule_dump(s: isl.schedule):\n",
    "  printer = isl.printer.from_file(\"/tmp/sche.yaml\")\n",
    "  printer = printer.set_yaml_style(isl.ISL_YAML_STYLE.BLOCK)\n",
    "  printer = printer.print_schedule(s)\n",
    "  printer.flush()\n",
    "  with open(\"/tmp/sche.yaml\", 'r') as f:\n",
    "    content = f.read()\n",
    "  print(content)\n",
    "\n",
    "isl_schedule_dump(schedule)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们可以直接获取到节点, 或者通过`visitor`进行遍历:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "直接获取节点:\n",
      "domain: \"{ S_0[]; S_3[]; S_1[]; S_5[i, j] : 0 <= i <= 1023 and 0 <= j <= 2047; S_2[]; S_4[] }\"\n",
      "child:\n",
      "  sequence:\n",
      "  - # YOU ARE HERE\n",
      "    filter: \"{ S_0[] }\"\n",
      "  - filter: \"{ S_1[] }\"\n",
      "  - filter: \"{ S_2[] }\"\n",
      "  - filter: \"{ S_3[] }\"\n",
      "  - filter: \"{ S_4[] }\"\n",
      "  - filter: \"{ S_5[i, j] }\"\n",
      "    child:\n",
      "      schedule: \"L_0[{ S_5[i, j] -> [(i)] }]\"\n",
      "      child:\n",
      "        schedule: \"L_1[{ S_5[i, j] -> [(j)] }]\"\n",
      "\n",
      "后序遍历:\n",
      "<class 'isl.schedule_node_leaf'>\n",
      "<class 'isl.schedule_node_filter'>\n",
      "<class 'isl.schedule_node_leaf'>\n",
      "<class 'isl.schedule_node_filter'>\n",
      "<class 'isl.schedule_node_leaf'>\n",
      "<class 'isl.schedule_node_filter'>\n",
      "<class 'isl.schedule_node_leaf'>\n",
      "<class 'isl.schedule_node_filter'>\n",
      "<class 'isl.schedule_node_leaf'>\n",
      "<class 'isl.schedule_node_filter'>\n",
      "<class 'isl.schedule_node_leaf'>\n",
      "<class 'isl.schedule_node_band'>\n",
      "<class 'isl.schedule_node_band'>\n",
      "<class 'isl.schedule_node_filter'>\n",
      "<class 'isl.schedule_node_sequence'>\n",
      "<class 'isl.schedule_node_domain'>\n"
     ]
    }
   ],
   "source": [
    "print(\"直接获取节点:\")\n",
    "print(schedule.get_root().child(0).child(0))\n",
    "\n",
    "\n",
    "print(\"后序遍历:\")\n",
    "root_node = schedule.get_root()\n",
    "def tile_band_node(node: isl.schedule_node):\n",
    "  print(node.__class__)\n",
    "  return node\n",
    "# 使用callback遍历每个节点\n",
    "_ = root_node.map_descendant_bottom_up(tile_band_node)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 调度树的操作\n",
    "\n",
    "在赵捷老师的书中介绍的操作方式比较多:\n",
    "\n",
    "1. 插入节点, 引入额外信息\n",
    "\n",
    "2. 修改`band`节点进行循环变换\n",
    "\n",
    "3. 嵌套的`ban`d节点分裂(`split`)/合并(`combine`), 适配硬件并行维度.\n",
    "\n",
    "4. 并列的`band`节点(合并)`fuse`/(分块)`distribute`.\n",
    "\n",
    "虽然在书中第四章中介绍了需要循环变换的方式, 但是缺乏具体的例子来告诉读者怎样将各种变换应用在schedule tree上, 本篇文章主要关注这部分.\n",
    "\n",
    "首先验证一下输出未经修改的源代码:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre class='code'><code class=\"cpp hljs\">{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = 0; c0 <= 1023; c0 += 1)\n",
       "    for (int c1 = 0; c1 <= 2047; c1 += 1)\n",
       "      a[2048 * c0 + c1] = (2048 * c0 + c1);\n",
       "}\n",
       "</code></pre>"
      ],
      "text/plain": [
       "```c\n",
       "{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = 0; c0 <= 1023; c0 += 1)\n",
       "    for (int c1 = 0; c1 <= 2047; c1 += 1)\n",
       "      a[2048 * c0 + c1] = (2048 * c0 + c1);\n",
       "}\n",
       "\n",
       "```"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from common import CodeGenerator, CSource\n",
    "\n",
    "generator = CodeGenerator(scop, schedule)\n",
    "source: CSource = generator.generate()\n",
    "source\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 循环反转\n",
    "\n",
    "`band`节点作为包含循环的节点类型, 需要操作循环也需要通过它所提供的函数. 首先我们可以尝试拿到band节点, 并输出他的`schedule`, 在`schedule node`中, `.schedule`是返回自身的`schedule tree`. 因此band节点的schedule需要使用`get_partial_schedule`来获取`multi_union_pw_aff`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "L_0[{ S_5[i, j] -> [(i)] }]\n"
     ]
    }
   ],
   "source": [
    "band_node: isl.schedule_node_band = schedule.get_root().child(0).child(5).child(0)\n",
    "aff: isl.multi_union_pw_aff = band_node.get_partial_schedule()\n",
    "print(aff)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这里我想对循环顺序进行反转, 那么我们可以去修改这个`aff`表达式. 正好`isl`中提供了将`aff`函数体进行反转的语法糖:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "取负操作: L_0[{ S_5[i, j] -> [(-i)] }]\n"
     ]
    }
   ],
   "source": [
    "print(\"取负操作:\", aff.neg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "有了新的`aff`之后, 那么我们需要修改原始的`schedule tree`, 而在`isl`中默认`schedule tree`是不可变的, 所以我们每次插入和构造节点最终都会返回新的`schedule tree`, 因此不需要考虑原始对象的引用问题, 十分方便.\n",
    "\n",
    "我们只需要删除当前这个循环的band节点, 然后再插入一个由新的`aff`表达式构造的`band`节点即可(由于`band`节点是通过`aff`来构造的, 所以插入的函数为`insert_partial_schedule`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain: \"{ S_0[]; S_3[]; S_1[]; S_5[i, j] : 0 <= i <= 1023 and 0 <= j <= 2047; S_2[]; S_4[] }\"\n",
      "child:\n",
      "  sequence:\n",
      "  - filter: \"{ S_0[] }\"\n",
      "  - filter: \"{ S_1[] }\"\n",
      "  - filter: \"{ S_2[] }\"\n",
      "  - filter: \"{ S_3[] }\"\n",
      "  - filter: \"{ S_4[] }\"\n",
      "  - filter: \"{ S_5[i, j] }\"\n",
      "    child:\n",
      "      schedule: \"L_0[{ S_5[i, j] -> [(-i)] }]\"\n",
      "      child:\n",
      "        schedule: \"L_1[{ S_5[i, j] -> [(j)] }]\"\n",
      "\n"
     ]
    }
   ],
   "source": [
    "new_tree = band_node.delete() \\\n",
    "    .insert_partial_schedule(aff.neg()) \\\n",
    "    .schedule()\n",
    "isl_schedule_dump(new_tree)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来再生成代码:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre class='code'><code class=\"cpp hljs\">{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = -1023; c0 <= 0; c0 += 1)\n",
       "    for (int c1 = 0; c1 <= 2047; c1 += 1)\n",
       "      a[-2048 * c0 + c1] = (-2048 * c0 + c1);\n",
       "}\n",
       "</code></pre>"
      ],
      "text/plain": [
       "```c\n",
       "{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = -1023; c0 <= 0; c0 += 1)\n",
       "    for (int c1 = 0; c1 <= 2047; c1 += 1)\n",
       "      a[-2048 * c0 + c1] = (-2048 * c0 + c1);\n",
       "}\n",
       "\n",
       "```"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "generator = CodeGenerator(scop, new_tree)\n",
    "source: CSource = generator.generate()\n",
    "source"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 循环交换\n",
    "\n",
    "我们学会了删除节点和插入节点, 同时也知道了`aff`表达式控制着循环的迭代, 那么想要交换两层循环自然就需要将`schedule tree`上的两个节点位置互换.\n",
    "\n",
    "首先我们将`child`节点保存下来, 然后从当前节点开始`cut`整个分支之后先插入原本节点,再插入`child`节点:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain: \"{ S_0[]; S_3[]; S_1[]; S_5[i, j] : 0 <= i <= 1023 and 0 <= j <= 2047; S_2[]; S_4[] }\"\n",
      "child:\n",
      "  sequence:\n",
      "  - filter: \"{ S_0[] }\"\n",
      "  - filter: \"{ S_1[] }\"\n",
      "  - filter: \"{ S_2[] }\"\n",
      "  - filter: \"{ S_3[] }\"\n",
      "  - filter: \"{ S_4[] }\"\n",
      "  - filter: \"{ S_5[i, j] }\"\n",
      "    child:\n",
      "      schedule: \"L_1[{ S_5[i, j] -> [(j)] }]\"\n",
      "      child:\n",
      "        schedule: \"L_0[{ S_5[i, j] -> [(i)] }]\"\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre class='code'><code class=\"cpp hljs\">{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = 0; c0 <= 2047; c0 += 1)\n",
       "    for (int c1 = 0; c1 <= 1023; c1 += 1)\n",
       "      a[c0 + 2048 * c1] = (c0 + 2048 * c1);\n",
       "}\n",
       "</code></pre>"
      ],
      "text/plain": [
       "```c\n",
       "{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = 0; c0 <= 2047; c0 += 1)\n",
       "    for (int c1 = 0; c1 <= 1023; c1 += 1)\n",
       "      a[c0 + 2048 * c1] = (c0 + 2048 * c1);\n",
       "}\n",
       "\n",
       "```"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "aff = band_node.get_partial_schedule()\n",
    "band_node_child: isl.schedule_node_band = band_node.child(0)\n",
    "child_aff = band_node_child.get_partial_schedule()\n",
    "new_tree = band_node.cut(). \\\n",
    "  insert_partial_schedule(aff). \\\n",
    "  insert_partial_schedule(child_aff). \\\n",
    "  schedule()\n",
    "\n",
    "isl_schedule_dump(new_tree)\n",
    "\n",
    "generator = CodeGenerator(scop, new_tree)\n",
    "source: CSource = generator.generate()\n",
    "source"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 循环延展\n",
    "\n",
    "循环延展一般是将循环索引变量系数从1改变为一个大于1的整数系数. 其实isl中已经提供了对`band`节点的`aff`表达式操作的接口, 有`scale/scale_down/shift/split/tile`等. \n",
    "\n",
    "因此可以使用`scale`去调整循环的系数从而完成循环延展. 在isl中要求我们输入的系数值和这个节点的表达式拥有同样的space, 才可以进行修改, 因此我们需要获取表达式的space, 构造出`multi_val`系数后进行操作:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{ L_0[4] }\n"
     ]
    }
   ],
   "source": [
    "mv = isl.multi_val.zero(band_node.space())  # 通过aff的space构造multi value, 保证match aff的space.\n",
    "for i in range(band_node.space().dim(isl.ISL_DIM_TYPE.SET)):\n",
    "  mv = mv.set_at(i, isl.val(\"4\")) \n",
    "print(mv)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后调用`scale`函数操作节点,并返回新的`schedule tree`后就可以拿来生成新的代码了, 不过`isl`中应该还有一些特殊处理, 不合理的情况在生成代码时还是被忽略了, 导致`scale`不起作用:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain: \"{ S_0[]; S_3[]; S_1[]; S_5[i, j] : 0 <= i <= 1023 and 0 <= j <= 2047; S_2[]; S_4[] }\"\n",
      "child:\n",
      "  sequence:\n",
      "  - filter: \"{ S_0[] }\"\n",
      "  - filter: \"{ S_1[] }\"\n",
      "  - filter: \"{ S_2[] }\"\n",
      "  - filter: \"{ S_3[] }\"\n",
      "  - filter: \"{ S_4[] }\"\n",
      "  - filter: \"{ S_5[i, j] }\"\n",
      "    child:\n",
      "      schedule: \"L_0[{ S_5[i, j] -> [(4i)] }]\"\n",
      "      child:\n",
      "        schedule: \"L_1[{ S_5[i, j] -> [(j)] }]\"\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre class='code'><code class=\"cpp hljs\">{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = 0; c0 <= 1023; c0 += 1)\n",
       "    for (int c1 = 0; c1 <= 2047; c1 += 1)\n",
       "      a[2048 * c0 + c1] = (2048 * c0 + c1);\n",
       "}\n",
       "</code></pre>"
      ],
      "text/plain": [
       "```c\n",
       "{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = 0; c0 <= 1023; c0 += 1)\n",
       "    for (int c1 = 0; c1 <= 2047; c1 += 1)\n",
       "      a[2048 * c0 + c1] = (2048 * c0 + c1);\n",
       "}\n",
       "\n",
       "```"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_tree = band_node.scale(mv).schedule()\n",
    "isl_schedule_dump(new_tree)\n",
    "generator = CodeGenerator(scop, new_tree)\n",
    "source: CSource = generator.generate()\n",
    "source"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "不过我们还可以使用`scale_down`, 向下进行延展:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain: \"{ S_0[]; S_3[]; S_1[]; S_5[i, j] : 0 <= i <= 1023 and 0 <= j <= 2047; S_2[]; S_4[] }\"\n",
      "child:\n",
      "  sequence:\n",
      "  - filter: \"{ S_0[] }\"\n",
      "  - filter: \"{ S_1[] }\"\n",
      "  - filter: \"{ S_2[] }\"\n",
      "  - filter: \"{ S_3[] }\"\n",
      "  - filter: \"{ S_4[] }\"\n",
      "  - filter: \"{ S_5[i, j] }\"\n",
      "    child:\n",
      "      schedule: \"L_0[{ S_5[i, j] -> [(floor((i)/4))] }]\"\n",
      "      child:\n",
      "        schedule: \"L_1[{ S_5[i, j] -> [(j)] }]\"\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre class='code'><code class=\"cpp hljs\">{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = 0; c0 <= 255; c0 += 1)\n",
       "    for (int c1 = 0; c1 <= 2047; c1 += 1)\n",
       "      for (int c2 = 4 * c0; c2 <= 4 * c0 + 3; c2 += 1)\n",
       "        a[c1 + 2048 * c2] = (c1 + 2048 * c2);\n",
       "}\n",
       "</code></pre>"
      ],
      "text/plain": [
       "```c\n",
       "{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = 0; c0 <= 255; c0 += 1)\n",
       "    for (int c1 = 0; c1 <= 2047; c1 += 1)\n",
       "      for (int c2 = 4 * c0; c2 <= 4 * c0 + 3; c2 += 1)\n",
       "        a[c1 + 2048 * c2] = (c1 + 2048 * c2);\n",
       "}\n",
       "\n",
       "```"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_tree = band_node.scale_down(mv).schedule()\n",
    "isl_schedule_dump(new_tree)\n",
    "generator = CodeGenerator(scop, new_tree)\n",
    "source: CSource = generator.generate()\n",
    "source"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 循环分块\n",
    "\n",
    "分块同样是调用`isl`所提供的接口即可:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain: \"{ S_0[]; S_3[]; S_1[]; S_5[i, j] : 0 <= i <= 1023 and 0 <= j <= 2047; S_2[]; S_4[] }\"\n",
      "child:\n",
      "  sequence:\n",
      "  - filter: \"{ S_0[] }\"\n",
      "  - filter: \"{ S_1[] }\"\n",
      "  - filter: \"{ S_2[] }\"\n",
      "  - filter: \"{ S_3[] }\"\n",
      "  - filter: \"{ S_4[] }\"\n",
      "  - filter: \"{ S_5[i, j] }\"\n",
      "    child:\n",
      "      schedule: \"L_0[{ S_5[i, j] -> [(i - (i) mod 24)] }]\"\n",
      "      child:\n",
      "        schedule: \"L_0[{ S_5[i, j] -> [((i) mod 24)] }]\"\n",
      "        child:\n",
      "          schedule: \"L_1[{ S_5[i, j] -> [(j)] }]\"\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre class='code'><code class=\"cpp hljs\">{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = 0; c0 <= 1023; c0 += 24)\n",
       "    for (int c1 = 0; c1 <= min(23, -c0 + 1023); c1 += 1)\n",
       "      for (int c2 = 0; c2 <= 2047; c2 += 1)\n",
       "        a[2048 * c0 + 2048 * c1 + c2] = (2048 * c0 + 2048 * c1 + c2);\n",
       "}\n",
       "</code></pre>"
      ],
      "text/plain": [
       "```c\n",
       "{\n",
       "  M = 1024;\n",
       "  N = 2048;\n",
       "  for (int c0 = 0; c0 <= 1023; c0 += 24)\n",
       "    for (int c1 = 0; c1 <= min(23, -c0 + 1023); c1 += 1)\n",
       "      for (int c2 = 0; c2 <= 2047; c2 += 1)\n",
       "        a[2048 * c0 + 2048 * c1 + c2] = (2048 * c0 + 2048 * c1 + c2);\n",
       "}\n",
       "\n",
       "```"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mv = isl.multi_val.zero(band_node.space())  # 通过aff的space构造multi value, 保证match aff的space.\n",
    "for i in range(band_node.space().dim(isl.ISL_DIM_TYPE.SET)):\n",
    "  mv = mv.set_at(i, isl.val(\"24\")) \n",
    "\n",
    "new_tree = band_node.tile(mv).schedule()\n",
    "isl_schedule_dump(new_tree)\n",
    "generator = CodeGenerator(scop, new_tree)\n",
    "source: CSource = generator.generate()\n",
    "source"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 循环合并\n",
    "\n",
    "我们给出的这个例子中, 循环`M`和循环`N`实际上是没有依赖的, 所以可以把两个循环合成为一个循环. 我们继续之前修改调度树的思路, 首先我们需要删除一个`band`节点, 并且少了一个循环层级, 为了保证`statement instance`不变, 那么也就需要去修改`domain`节点. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这里我的做法是首先拿到`domain`的`union set`, 然后修改为新的`domain`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "isl.schedule(\"\"\"{ domain: \"{ S_5[i0] : 0 <= i0 <= 2097151 }\" }\"\"\")"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = isl.map(\"{ S_5[i, j] -> S_5[i*2048+j] }\")\n",
    "domain = schedule.get_domain()\n",
    "new_domain = domain.apply(m)\n",
    "new_schedule = isl.schedule.from_domain(new_domain)\n",
    "new_schedule"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "再按原图顺序重新插入各个`filter`节点.  这里我们知道`sequence`节点下都是`filter`节点, 所以收集各个`filter`节点的信息, 重新构造出`union_set_list`然后统一插入."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "isl.schedule_node_sequence(\"\"\"domain: \"{ S_5[i0] : 0 <= i0 <= 2097151 }\"\n",
       "child:\n",
       "  # YOU ARE HERE\n",
       "  sequence:\n",
       "  - filter: \"{ S_0[] }\"\n",
       "  - filter: \"{ S_1[] }\"\n",
       "  - filter: \"{ S_2[] }\"\n",
       "  - filter: \"{ S_3[] }\"\n",
       "  - filter: \"{ S_4[] }\"\n",
       "  - filter: \"{ S_5[i0] }\"\n",
       "\"\"\")"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "old_seq: isl.schedule_node_sequence = schedule.get_root().child(0)\n",
    "new_seq: isl.schedule_node_sequence = new_schedule.get_root().child(0)\n",
    "\n",
    "union_set_list = isl.union_set_list(0)\n",
    "for i in range(old_seq.n_children() - 1):\n",
    "  old_filter: isl.schedule_node_filter = old_seq.child(i)\n",
    "  union_set_list = union_set_list.add(old_filter.get_filter())\n",
    "union_set_list = union_set_list.add(domain.get_space().add_named_tuple(\"S_5\", 1).universe_set())\n",
    "new_seq = new_seq.insert_sequence(union_set_list)\n",
    "new_seq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来再插入`band`节点."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "isl.schedule_node_band(\"\"\"domain: \"{ S_5[i0] : 0 <= i0 <= 2097151 }\"\n",
       "child:\n",
       "  sequence:\n",
       "  - filter: \"{ S_0[] }\"\n",
       "  - filter: \"{ S_1[] }\"\n",
       "  - filter: \"{ S_2[] }\"\n",
       "  - filter: \"{ S_3[] }\"\n",
       "  - filter: \"{ S_4[] }\"\n",
       "  - filter: \"{ S_5[i0] }\"\n",
       "    child:\n",
       "      # YOU ARE HERE\n",
       "      schedule: \"L_0[{ S_5[i0] -> [(i0)] }]\"\n",
       "\"\"\")"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_band = new_seq.child(5).child(0).insert_partial_schedule(isl.multi_union_pw_aff(\"L_0[{ S_5[i0] -> [(i0)] }]\"))\n",
    "new_band"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后也就是最关键的一步, 需要将原本代码`stmt`中2维的索引修改为1维的索引操作. 幸好`pet`的代码生成逻辑是根据原始`stmt`构造出一个称为`index`的`multi_pw_aff`, 然后再利用我们当前`schedule`的反向`aff`进行函数合成,也就是将反向`aff`的输出映射到`index`上. \n",
    "\n",
    "但是由于原来的`stmt`是2维的索引`S_5[i, j] -> a[(2048i + j)]`, 因此我这里强行把它改成`S_5[i] -> a[i]`, 这样和新的`schedule`中的`aff`进行合成的时候就可以获得正确的`index`. 然后新的index操作取出来变成`ast_expr`, 最后代码生成的时候可以直接使用这个`ast_expr`, 从而生成正确的代码."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "domain: \"{ S_5[i0] : 0 <= i0 <= 2097151 }\"\n",
      "child:\n",
      "  sequence:\n",
      "  - filter: \"{ S_0[] }\"\n",
      "  - filter: \"{ S_1[] }\"\n",
      "  - filter: \"{ S_2[] }\"\n",
      "  - filter: \"{ S_3[] }\"\n",
      "  - filter: \"{ S_4[] }\"\n",
      "  - filter: \"{ S_5[i0] }\"\n",
      "    child:\n",
      "      schedule: \"L_0[{ S_5[i0] -> [(i0)] }]\"\n",
      "\n",
      "old index:  { S_5[i, j] -> a[(2048i + j)] }\n",
      "iter_map:  { [[] -> L_0[i0]] -> S_5[(i0)] : 0 <= i0 <= 2097151 }\n",
      "old index:  { S_5[i, j] -> [(2048i + j)] }\n",
      "iter_map:  { [[] -> L_0[i0]] -> S_5[(i0)] : 0 <= i0 <= 2097151 }\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre class='code'><code class=\"cpp hljs\">for (int c0 = 0; c0 <= 2097151; c0 += 1)\n",
       "  a[c0] = (c0);\n",
       "</code></pre>"
      ],
      "text/plain": [
       "```c\n",
       "for (int c0 = 0; c0 <= 2097151; c0 += 1)\n",
       "  a[c0] = (c0);\n",
       "\n",
       "```"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_tree = new_band.schedule()\n",
    "isl_schedule_dump(new_tree)\n",
    "\n",
    "\n",
    "def custom_pullback(index: isl.multi_pw_aff, id, iter_map):\n",
    "  print(\"old index: \", index)\n",
    "  print(\"iter_map: \", iter_map)\n",
    "  return isl.multi_pw_aff(\"{ S_5[i] -> a[i] }\").pullback(iter_map)\n",
    "\n",
    "\n",
    "generator = CodeGenerator(scop, new_tree, custom_pullback)\n",
    "source: CSource = generator.generate()\n",
    "source\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ci",
   "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.17"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}