1.tensor_program_abstraction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[张量程序抽象](#张量程序抽象)\n",
    "\n",
    "[TensorIR：张量程序抽象案例研究](#TensorIR：张量程序抽象案例研究)\n",
    "\n",
    "[转换张量程序](#转换张量程序)\n",
    "\n",
    "[构建和运行](#构建和运行)\n",
    "\n",
    "[创建TensorIR并与之交互的方法](#创建tensorir并与之交互的方法)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import tir as T\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    " ## 元张量函数\n",
    "*****\n",
    "- MLC的过程可以看作张量函数间的变换。\n",
    "\n",
    "- <font color=red>元张量函数</font>\n",
    "    - 机器学习模型的执行包含将输入张量转换到输出预测的多步计算，其中的每一步都为元张量函数。\n",
    "    - 同一个元张量函数可以由多种不同的抽象表示和实现。\n",
    "        - 如：调用预编译框架库(torch.add numpy.add)\n",
    "        - 目标：将元张量函数转换为更专门和针对特定工作和部署环境的函数。\n",
    "\n",
    "![元张量函数间的变换](img/tensor_func_transformation.png)\n",
    "- 左图中的add被转换为右侧的一段表示可能的组合优化的add(被拆分为长度为4的单元)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span id=\"张量程序抽象\"></span>\n",
    "\n",
    "# 张量程序抽象\n",
    "- 针对元张量函数变换的需要，需要一个有效的抽象来表示，从而有效地变换元张量函数\n",
    "- 一个典型抽象包含：\n",
    "    - 存储数据的多维数组\n",
    "    - 驱动张量计算的循环嵌套\n",
    "    - 计算部分语句\n",
    "- 张量程序抽象的一个重要性质：能够被一系列有效的程序变换所改变\n",
    "*****\n",
    "- 我们不能随意对程序进行变换，因为一些计算会依赖循环间的顺序。\n",
    "- 大多数元张量函数都有良好的属性，如循环迭代间的独立性。\n",
    "\n",
    "下图中包含额外的T.axis.spatial标注，表明 vi 这个特定的变量被映射到循环变量 i，并且所有的迭代都是独立的。\n",
    "- 这个信息会使得我们在变换这个程序时更加方便。我们可以安全地并行或者重新排序所有与 vi 有关的循环，只要实际执行中 vi 的值按照从 0 到 128 的顺序变化。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# from tvm.script import tir as T\n",
      "@tvm.script.ir_module\n",
      "class Module:\n",
      "    @T.prim_func\n",
      "    def main(A: T.Buffer[128, \"float32\"], B: T.Buffer[128, \"float32\"], C: T.Buffer[128, \"float32\"]) -> None:\n",
      "        # function attr dict\n",
      "        T.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
      "        # body\n",
      "        # with T.block(\"root\")\n",
      "        for i in T.serial(128):\n",
      "            with T.block(\"C\"):\n",
      "                vi = T.axis.spatial(128, i)\n",
      "                T.reads(A[vi], B[vi])\n",
      "                T.writes(C[vi])\n",
      "                C[vi] = A[vi] + B[vi]\n",
      "    \n",
      "\n"
     ]
    }
   ],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModule:\n",
    "\n",
    "    # 张量程序抽象\n",
    "    @T.prim_func\n",
    "    def main(A: T.Buffer[128, \"float32\"],   # 多维张量缓冲区：用来存放输入，输出，中间结果\n",
    "             B: T.Buffer[128, \"float32\"], \n",
    "             C: T.Buffer[128, \"float32\"]):\n",
    "        # extra annotations for the function\n",
    "        T.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
    "        for i in range(128):                # 驱动张量计算的循环嵌套\n",
    "            with T.block(\"C\"):\n",
    "                # declare a data parallel iterator on spatial domain\n",
    "                vi = T.axis.spatial(128, i)\n",
    "                C[vi] = A[vi] + B[vi]       # 计算语句\n",
    "\n",
    "# TVMScript 用于在python ast中表达张量程序。\n",
    "# 虽然这段代码不等价于python程序，但可以在MLC过程中使用。\n",
    "# 该语言旨在与python语法一致，且有便于分析和转换的额外结构。\n",
    "\n",
    "type(MyModule)\n",
    "# MyMudule的数据结构为IRModule：用于保存一群张量函数的集合\n",
    "# 这一数据结构可以用script函数获取字符串形式的表达，这是一个用于检查每步转换后模块的重要方法。\n",
    "print(MyModule.script())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 通过tvm.build将IRModule转化为可运行函数\n",
    "- build后的mod包含一群可运行函数的集合，可以通过函数名取回函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'tvm.driver.build_module.OperatorModule'>\n",
      "<tvm.runtime.packed_func.PackedFunc object at 0x7f536b627380>\n"
     ]
    }
   ],
   "source": [
    "rt_mod = tvm.build(MyModule, target=\"llvm\")  # The module for CPU backends.\n",
    "print(type(rt_mod))\n",
    "\n",
    "func = rt_mod[\"main\"]\n",
    "print(func)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 在tvm runtime创建三个NDArray\n",
    "a = tvm.nd.array(np.arange(128, dtype=\"float32\"))\n",
    "b = tvm.nd.array(np.ones(128, dtype=\"float32\"))\n",
    "c = tvm.nd.empty((128,), dtype=\"float32\") \n",
    "func(a,b,c)\n",
    "print(a)\n",
    "print(b)\n",
    "print(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 分解循环\n",
    "sch = tvm.tir.Schedule(MyModule)\n",
    "print(type(sch))\n",
    "\n",
    "# Get block by its name\n",
    "block_c = sch.get_block(\"C\")\n",
    "# Get loops surronding the block\n",
    "(i,) = sch.get_loops(block_c)\n",
    "# Tile the loop nesting.\n",
    "i_0, i_1, i_2 = sch.split(i, factors=[None, 4, 4])\n",
    "print(sch.mod.script())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span id=\"TensorIR：张量程序抽象案例研究\"></span>\n",
    "\n",
    "# TensorIR：张量程序抽象案例研究\n",
    "TensorIR 是标准机器学习编译框架 Apache TVM 中使用的张量程序抽象。\n",
    "\n",
    "使用张量程序抽象的主要目的是表示循环和相关的硬件加速选择，如多线程、特殊硬件指令的使用和内存访问。\n",
    "\n",
    "以两个128×128大小的矩阵A和B进行两步计算： matmul->relu 线性层和激活层"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tvm\n",
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import tir as T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 简单的numpy实现\n",
    "\n",
    "在底层，NumPy 调用库（例如 OpenBLAS）和它自己在低级 C 语言中的一些实现来执行这些计算。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype = \"float32\"\n",
    "a_np = np.random.rand(128, 128).astype(dtype)\n",
    "b_np = np.random.rand(128, 128).astype(dtype)\n",
    "# a @ b is equivalent to np.matmul(a, b)\n",
    "c_mm_relu = np.maximum(a_np @ b_np, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "底层细节实现\n",
    "\n",
    "为了说明底层细节，我们将在 NumPy API 的一个受限子集中编写示例 —— 我们称之为 低级 NumPy。它使用以下的约定：\n",
    "\n",
    "- 我们将在必要时使用循环而不是数组函数来展示可能的循环计算。\n",
    "- 如果可能，我们总是通过 numpy.empty 显式地分配数组并传递它们。\n",
    "\n",
    "需要注意的是，这不是人们通常编写 NumPy 程序的方式。不过，它们仍然与幕后发生的事情非常相似 —— 大多数现实世界的部署解决方案都将分配与计算分开处理。特定的库使用不同形式的循环和算术计算来执行计算。当然首先它们是使用诸如 C 之类的低级语言实现的。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):\n",
    "    # 首先分配中间存储A，存放矩阵乘结果\n",
    "    Y = np.empty((128, 128), dtype=\"float32\")\n",
    "    for i in range(128):\n",
    "        for j in range(128):\n",
    "            for k in range(128):\n",
    "                if k == 0:\n",
    "                    Y[i, j] = 0\n",
    "                Y[i, j] = Y[i, j] + A[i, k] * B[k, j]\n",
    "    # 计算RELU\n",
    "    for i in range(128):\n",
    "        for j in range(128):\n",
    "            C[i, j] = max(Y[i, j], 0)\n",
    "\n",
    "c_np = np.empty((128, 128), dtype=dtype)\n",
    "lnumpy_mm_relu(a_np, b_np, c_np)\n",
    "# 对比验证结果正确性\n",
    "np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TensorIR\n",
    "使用TVMScript语言实现\n",
    "\n",
    "- 主要区别\n",
    "    - 计算块中T.block额外结构。块是TensoIR中的基本计算单位。\n",
    "        - 一个块包含一组块轴(vi, vj, vk)和围绕他们定义的计算。\n",
    "            - 块轴的<font color=red>关键性质</font>用T.axis.[axis_type]声明\n",
    "            - 语法：[block_axis] = T.axis.[axis_type]([axis_range], [mapped_value])     \n",
    "                - [axis_range]提供了[block_axis]的预期范围 如vi = T.axis.spatial(128.i)表面vi值应该在(0,128)中。\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModule:\n",
    "    @T.prim_func\n",
    "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"],\n",
    "                B: T.Buffer[(128, 128), \"float32\"],\n",
    "                C: T.Buffer[(128, 128), \"float32\"]):\n",
    "        \n",
    "        T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
    "        # 首先分配中间存储A，存放矩阵乘结果\n",
    "        Y = T.alloc_buffer((128, 128), dtype=\"float32\")\n",
    "        # 循环迭代\n",
    "        for i, j, k in T.grid(128, 128, 128):\n",
    "            # 计算块\n",
    "            with T.block(\"Y\"):\n",
    "                vi = T.axis.spatial(128, i)\n",
    "                vj = T.axis.spatial(128, j)\n",
    "                vk = T.axis.reduce(128, k)\n",
    "                with T.init():\n",
    "                    Y[vi, vj] = T.float32(0)\n",
    "                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
    "        for i, j in T.grid(128, 128):\n",
    "            with T.block(\"C\"):\n",
    "                vi = T.axis.spatial(128, i)\n",
    "                vj = T.axis.spatial(128, j)\n",
    "                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 块轴的属性\n",
    "\n",
    "这些轴属性标记了轴与正在执行的计算之间的关系。\n",
    "\n",
    "下图总结了块（迭代）轴和块 Y 的读写关系。注意到，严格来说这个块正在对缓冲区 Y 进行（规约）更新，我们暂时将其标记为写入，因为我们不需要来自另一个块的缓冲区 Y 的值。\n",
    "\n",
    "![块轴与计算的关系](img/tensor_ir_block_axis.png)\n",
    "\n",
    "块 Y 通过读取来自 A[vi, vk] 和 B[vk, vj] 的值来计算结果 Y[vi, vj]，并对所有可能的 vk 执行求和。 在这个特定示例中，如果我们将 vi、vj 固定为 (0, 1)，并对 vk in range(0, 128) 执行块 Y，我们可以独立于其他可能的位置（具有不同 vi, vj 值的位置）有效地计算 C[0, 1]。\n",
    "\n",
    "值得注意的是，对于一组固定的 vi 和 vj，计算块在 Y 的空间位置 (Y[vi, vj]) 处生成一个点值，该点值独立于 Y 中的其他位置（具有不同的vi, vj 值的位置）。我们可以称 vi、vj 为<font color=red>空间轴</font>，因为它们直接对应于块写入的缓冲区空间区域的开始。 涉及归约（vk）的轴被命名为<font color=red>归约轴</font>。\n",
    "\n",
    "- 这些附加信息有助于我们进行机器学习编译分析。例如，虽然我们总是可以在空间轴上做并行化，在规约轴上进行并行化将需要特定的策略。\n",
    "\n",
    "块轴绑定的语法糖\n",
    "\n",
    "- 在每个块轴直接映射到外部循环迭代器的情况下，我们可以使用 T.axis.remap 在一行中声明所有块轴。\n",
    "\n",
    "- SSR means the properties of each axes are \"spatial\", \"spatial\", \"reduce\"\n",
    "\n",
    "- vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 函数属性和装饰器\n",
    "T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
    "\n",
    "- 这里的 global_symbol 对应函数名，tir.noalias 是一个属性，表示所有的缓冲存储器不重叠。你现在可以放心地跳过这些属性，因为它们不会影响对概念的整体理解。\n",
    "\n",
    "@tvm.script.ir_module 和 @T.prim_func 这两个装饰器用于表示对应部分的类型。\n",
    "\n",
    "- @tvm.script.ir_module 表示 MyModule 是一个 IRModule。IRModule 是在机器学习编译中保存张量函数集合的容器对象。\n",
    "\n",
    "机器学习编译过程中的一个 IRModule 可以包含多个张量函数。 以下代码块显示了具有两个函数的 IRModule 示例。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span id=\"转换张量程序\"></span>\n",
    "\n",
    "# 转换张量程序\n",
    "- 用辅助数据结构schedule来转换张量程序。\n",
    "- 核心问题：\n",
    "    - 什么是表示张量函数可能的抽象？\n",
    "    - 什么是张量函数间可能的变换？"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">])</span> <span class=\"o\">-&gt;</span> <span class=\"kc\">None</span><span class=\"p\">:</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">,</span> <span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">,</span> <span class=\"n\">k</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">,</span> <span class=\"n\">vk</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SSR&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">init</span><span class=\"p\">():</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SS&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)} \\PY{o}{\\PYZhy{}}\\PY{o}{\\PYZgt{}} \\PY{k+kc}{None}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{,} \\PY{n}{k} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{,} \\PY{n}{vk} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SSR}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{,} \\PY{n}{k}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{init}\\PY{p}{(}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "                \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SS}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]) -> None:\n",
       "        # function attr dict\n",
       "        T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i, j, k in T.grid(128, 128, 128):\n",
       "            with T.block(\"Y\"):\n",
       "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
       "                T.reads(A[vi, vk], B[vk, vj])\n",
       "                T.writes(Y[vi, vj])\n",
       "                with T.init():\n",
       "                    Y[vi, vj] = T.float32(0)\n",
       "                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
       "        for i, j in T.grid(128, 128):\n",
       "            with T.block(\"C\"):\n",
       "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
       "                T.reads(Y[vi, vj])\n",
       "                T.writes(C[vi, vj])\n",
       "                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import IPython\n",
    "\n",
    "IPython.display.Code(MyModule.script(), language=\"python\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">])</span> <span class=\"o\">-&gt;</span> <span class=\"kc\">None</span><span class=\"p\">:</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">,</span> <span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j_0</span><span class=\"p\">,</span> <span class=\"n\">j_1</span><span class=\"p\">,</span> <span class=\"n\">k</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">32</span><span class=\"p\">,</span> <span class=\"mi\">4</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">j_1</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">vk</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">reduce</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">init</span><span class=\"p\">():</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SS&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)} \\PY{o}{\\PYZhy{}}\\PY{o}{\\PYZgt{}} \\PY{k+kc}{None}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j\\PYZus{}0}\\PY{p}{,} \\PY{n}{j\\PYZus{}1}\\PY{p}{,} \\PY{n}{k} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{32}\\PY{p}{,} \\PY{l+m+mi}{4}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{j\\PYZus{}1}\\PY{p}{)}\n",
       "                \\PY{n}{vk} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{reduce}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{k}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{init}\\PY{p}{(}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "                \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SS}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]) -> None:\n",
       "        # function attr dict\n",
       "        T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i, j_0, j_1, k in T.grid(128, 32, 4, 128):\n",
       "            with T.block(\"Y\"):\n",
       "                vi = T.axis.spatial(128, i)\n",
       "                vj = T.axis.spatial(128, j_0 * 4 + j_1)\n",
       "                vk = T.axis.reduce(128, k)\n",
       "                T.reads(A[vi, vk], B[vk, vj])\n",
       "                T.writes(Y[vi, vj])\n",
       "                with T.init():\n",
       "                    Y[vi, vj] = T.float32(0)\n",
       "                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
       "        for i, j in T.grid(128, 128):\n",
       "            with T.block(\"C\"):\n",
       "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
       "                T.reads(Y[vi, vj])\n",
       "                T.writes(C[vi, vj])\n",
       "                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 首先创建一个以给定的 MyModule 作为输入的 Schedule 辅助类。\n",
    "sch = tvm.tir.Schedule(MyModule)\n",
    "# 获得对块 Y 和相应循环的引用\n",
    "block_Y = sch.get_block(\"Y\", func_name=\"mm_relu\")\n",
    "i, j, k = sch.get_loops(block_Y)\n",
    "\n",
    "#将执行的第一个变换是将循环 j 分成两个循环，其中内部循环的长度为 4。\n",
    "j0, j1 = sch.split(j, factors=[None, 4])\n",
    "#查看存储在 sch.mod 中的变换结果。\n",
    "IPython.display.Code(sch.mod.script(), language=\"python\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">])</span> <span class=\"o\">-&gt;</span> <span class=\"kc\">None</span><span class=\"p\">:</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">,</span> <span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j_0</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">,</span> <span class=\"n\">j_1</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">32</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">j_1</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">vk</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">reduce</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">init</span><span class=\"p\">():</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SS&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)} \\PY{o}{\\PYZhy{}}\\PY{o}{\\PYZgt{}} \\PY{k+kc}{None}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j\\PYZus{}0}\\PY{p}{,} \\PY{n}{k}\\PY{p}{,} \\PY{n}{j\\PYZus{}1} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{32}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{j\\PYZus{}1}\\PY{p}{)}\n",
       "                \\PY{n}{vk} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{reduce}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{k}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{init}\\PY{p}{(}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "                \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SS}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]) -> None:\n",
       "        # function attr dict\n",
       "        T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i, j_0, k, j_1 in T.grid(128, 32, 128, 4):\n",
       "            with T.block(\"Y\"):\n",
       "                vi = T.axis.spatial(128, i)\n",
       "                vj = T.axis.spatial(128, j_0 * 4 + j_1)\n",
       "                vk = T.axis.reduce(128, k)\n",
       "                T.reads(A[vi, vk], B[vk, vj])\n",
       "                T.writes(Y[vi, vj])\n",
       "                with T.init():\n",
       "                    Y[vi, vj] = T.float32(0)\n",
       "                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
       "        for i, j in T.grid(128, 128):\n",
       "            with T.block(\"C\"):\n",
       "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
       "                T.reads(Y[vi, vj])\n",
       "                T.writes(C[vi, vj])\n",
       "                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 重新排序这两个循环。\n",
    "sch.reorder(j0, k, j1)\n",
    "IPython.display.Code(sch.mod.script(), language=\"python\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 另一个变体\n",
    "将块 C 移动到 Y 的内循环里\n",
    "\n",
    "将归约初始化和更新放在一个块体中。这种组合形式为循环变换带来了便利（因为初始化和更新的外循环 i、j 通常需要彼此保持同步）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">])</span> <span class=\"o\">-&gt;</span> <span class=\"kc\">None</span><span class=\"p\">:</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">,</span> <span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">32</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">for</span> <span class=\"n\">k</span><span class=\"p\">,</span> <span class=\"n\">j_1</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y&quot;</span><span class=\"p\">):</span>\n",
       "                    <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">j_1</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vk</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">reduce</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">init</span><span class=\"p\">():</span>\n",
       "                        <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span>\n",
       "            <span class=\"k\">for</span> <span class=\"n\">ax0</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">serial</span><span class=\"p\">(</span><span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                    <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">ax0</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)} \\PY{o}{\\PYZhy{}}\\PY{o}{\\PYZgt{}} \\PY{k+kc}{None}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{32}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{for} \\PY{n}{k}\\PY{p}{,} \\PY{n}{j\\PYZus{}1} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                    \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{j\\PYZus{}1}\\PY{p}{)}\n",
       "                    \\PY{n}{vk} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{reduce}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{k}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{init}\\PY{p}{(}\\PY{p}{)}\\PY{p}{:}\n",
       "                        \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\n",
       "            \\PY{k}{for} \\PY{n}{ax0} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{serial}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                    \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{ax0}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]) -> None:\n",
       "        # function attr dict\n",
       "        T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i, j_0 in T.grid(128, 32):\n",
       "            for k, j_1 in T.grid(128, 4):\n",
       "                with T.block(\"Y\"):\n",
       "                    vi = T.axis.spatial(128, i)\n",
       "                    vj = T.axis.spatial(128, j_0 * 4 + j_1)\n",
       "                    vk = T.axis.reduce(128, k)\n",
       "                    T.reads(A[vi, vk], B[vk, vj])\n",
       "                    T.writes(Y[vi, vj])\n",
       "                    with T.init():\n",
       "                        Y[vi, vj] = T.float32(0)\n",
       "                    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
       "            for ax0 in T.serial(4):\n",
       "                with T.block(\"C\"):\n",
       "                    vi = T.axis.spatial(128, i)\n",
       "                    vj = T.axis.spatial(128, j_0 * 4 + ax0)\n",
       "                    T.reads(Y[vi, vj])\n",
       "                    T.writes(C[vi, vj])\n",
       "                    C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block_C = sch.get_block(\"C\", \"mm_relu\")\n",
    "sch.reverse_compute_at(block_C, j0)\n",
    "IPython.display.Code(sch.mod.script(), language=\"python\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在循环变换之后，我们可以将 Y 元素的初始化与归约更新分开。我们可以通过 decompose_reduction 原语来做到这一点。（注意：这也是 TVM 在以后编译的时候隐式做的，所以这一步的主要目的是让它显式，看看最终效果）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">])</span> <span class=\"o\">-&gt;</span> <span class=\"kc\">None</span><span class=\"p\">:</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">,</span> <span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">32</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">for</span> <span class=\"n\">j_1_init</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">serial</span><span class=\"p\">(</span><span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y_init&quot;</span><span class=\"p\">):</span>\n",
       "                    <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">j_1_init</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">()</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "            <span class=\"k\">for</span> <span class=\"n\">k</span><span class=\"p\">,</span> <span class=\"n\">j_1</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y_update&quot;</span><span class=\"p\">):</span>\n",
       "                    <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">j_1</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vk</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">reduce</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vk</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">vk</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span>\n",
       "            <span class=\"k\">for</span> <span class=\"n\">ax0</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">serial</span><span class=\"p\">(</span><span class=\"mi\">4</span><span class=\"p\">):</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                    <span class=\"n\">vi</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">i</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">vj</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">spatial</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"n\">j_0</span> <span class=\"o\">*</span> <span class=\"mi\">4</span> <span class=\"o\">+</span> <span class=\"n\">ax0</span><span class=\"p\">)</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">])</span>\n",
       "                    <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">vi</span><span class=\"p\">,</span> <span class=\"n\">vj</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)} \\PY{o}{\\PYZhy{}}\\PY{o}{\\PYZgt{}} \\PY{k+kc}{None}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{32}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{for} \\PY{n}{j\\PYZus{}1\\PYZus{}init} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{serial}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y\\PYZus{}init}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                    \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{j\\PYZus{}1\\PYZus{}init}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "            \\PY{k}{for} \\PY{n}{k}\\PY{p}{,} \\PY{n}{j\\PYZus{}1} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y\\PYZus{}update}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                    \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{j\\PYZus{}1}\\PY{p}{)}\n",
       "                    \\PY{n}{vk} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{reduce}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{k}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vk}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{vk}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\n",
       "            \\PY{k}{for} \\PY{n}{ax0} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{serial}\\PY{p}{(}\\PY{l+m+mi}{4}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{vi} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{i}\\PY{p}{)}\n",
       "                    \\PY{n}{vj} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{spatial}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{n}{j\\PYZus{}0} \\PY{o}{*} \\PY{l+m+mi}{4} \\PY{o}{+} \\PY{n}{ax0}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{)}\n",
       "                    \\PY{n}{C}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{vi}\\PY{p}{,} \\PY{n}{vj}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]) -> None:\n",
       "        # function attr dict\n",
       "        T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i, j_0 in T.grid(128, 32):\n",
       "            for j_1_init in T.serial(4):\n",
       "                with T.block(\"Y_init\"):\n",
       "                    vi = T.axis.spatial(128, i)\n",
       "                    vj = T.axis.spatial(128, j_0 * 4 + j_1_init)\n",
       "                    T.reads()\n",
       "                    T.writes(Y[vi, vj])\n",
       "                    Y[vi, vj] = T.float32(0)\n",
       "            for k, j_1 in T.grid(128, 4):\n",
       "                with T.block(\"Y_update\"):\n",
       "                    vi = T.axis.spatial(128, i)\n",
       "                    vj = T.axis.spatial(128, j_0 * 4 + j_1)\n",
       "                    vk = T.axis.reduce(128, k)\n",
       "                    T.reads(Y[vi, vj], A[vi, vk], B[vk, vj])\n",
       "                    T.writes(Y[vi, vj])\n",
       "                    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]\n",
       "            for ax0 in T.serial(4):\n",
       "                with T.block(\"C\"):\n",
       "                    vi = T.axis.spatial(128, i)\n",
       "                    vj = T.axis.spatial(128, j_0 * 4 + ax0)\n",
       "                    T.reads(Y[vi, vj])\n",
       "                    T.writes(C[vi, vj])\n",
       "                    C[vi, vj] = T.max(Y[vi, vj], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sch.decompose_reduction(block_Y, k)\n",
    "IPython.display.Code(sch.mod.script(), language=\"python\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span id=\"构建和运行\"></span>\n",
    "\n",
    "# 构建和运行\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 原始\n",
    "rt_lib = tvm.build(MyModule, target=\"llvm\")\n",
    "a_nd = tvm.nd.array(a_np)\n",
    "b_nd = tvm.nd.array(b_np)\n",
    "c_nd = tvm.nd.empty((128, 128), dtype=\"float32\")\n",
    "func_mm_relu = rt_lib[\"mm_relu\"]\n",
    "func_mm_relu(a_nd, b_nd, c_nd)\n",
    "\n",
    "np.testing.assert_allclose(c_mm_relu, c_nd.numpy(), rtol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 变换后\n",
    "rt_lib_after = tvm.build(sch.mod, target=\"llvm\")\n",
    "rt_lib_after[\"mm_relu\"](a_nd, b_nd, c_nd)\n",
    "np.testing.assert_allclose(c_mm_relu, c_nd.numpy(), rtol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time cost of MyModule 0.00308494 sec\n",
      "Time cost of transformed sch.mod 0.000684898 sec\n"
     ]
    }
   ],
   "source": [
    "# 时间差\n",
    "f_timer_before = rt_lib.time_evaluator(\"mm_relu\", tvm.cpu())\n",
    "print(\"Time cost of MyModule %g sec\" % f_timer_before(a_nd, b_nd, c_nd).mean)\n",
    "f_timer_after = rt_lib_after.time_evaluator(\"mm_relu\", tvm.cpu())\n",
    "print(\"Time cost of transformed sch.mod %g sec\" % f_timer_after(a_nd, b_nd, c_nd).mean)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 分析影响性能的可能因素\n",
    "![CPU架构](img/cpu_arch.png)\n",
    "\n",
    "    - 访问 A 和 B 中的任何内存块的速度并不一致。现代 CPU 带有多级缓存，需要先将数据提取到缓存中，然后 CPU 才能访问它。\n",
    "    - 重要的是，访问已经在缓存中的数据要快得多。CPU 采用的一种策略是获取彼此更接近的数据。 当我们读取内存中的一个元素时，它会尝试将附近的元素（更正式的名称为“缓存行”）获取到缓存中。 因此，当你读取下一个元素时，它已经在缓存中。 因此，具有连续内存访问的代码通常比随机访问内存不同部分的代码更快。\n",
    "\n",
    "![可视化](img/tensor_func_loop_order.png)\n",
    "现在让我们看看上面的迭代可视化，分析一下是怎么回事。 在这个分析中，让我们关注最里面的两个循环：k 和 j1。高亮的地方显示了当我们针对 k 的一个特定实例迭代 j1 时迭代触及的 Y、A 和 B 中的相应区域。\n",
    "\n",
    "我们可以发现，j1 这一迭代产生了对 B 元素的连续访问。具体来说，它意味着在 j1=0 和 j1=1 时我们读取的值彼此相邻。这可以让我们拥有更好的缓存访问行为。此外，我们使 C 的计算更接近 Y，从而实现更好的缓存行为。\n",
    "\n",
    "我们当前的示例主要是为了证明不同的代码变体可以导致不同的性能。更多的变换步骤可以帮助我们获得更好的性能，我们将在以后的章节中介绍。本练习的主要目标是首先让我们获得程序变换工具，并首先体验通过变换可能实现的功能。\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<span id=\"创建TensorIR并与之交互的方法\"></span>\n",
    "\n",
    "# 创建TensorIR并与之交互的方法\n",
    "\n",
    "- 用TVMScript的方式\n",
    "- 用张量表达式TE生成TensorIR代码\n",
    "    - 帮助为给定的更高级别的输入生成TensorIR函数\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { line-height: 125%; }\n",
       "td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }\n",
       "td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       "span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }\n",
       ".output_html .hll { background-color: #ffffcc }\n",
       ".output_html { background: #f8f8f8; }\n",
       ".output_html .c { color: #3D7B7B; font-style: italic } /* Comment */\n",
       ".output_html .err { border: 1px solid #FF0000 } /* Error */\n",
       ".output_html .k { color: #008000; font-weight: bold } /* Keyword */\n",
       ".output_html .o { color: #666666 } /* Operator */\n",
       ".output_html .ch { color: #3D7B7B; font-style: italic } /* Comment.Hashbang */\n",
       ".output_html .cm { color: #3D7B7B; font-style: italic } /* Comment.Multiline */\n",
       ".output_html .cp { color: #9C6500 } /* Comment.Preproc */\n",
       ".output_html .cpf { color: #3D7B7B; font-style: italic } /* Comment.PreprocFile */\n",
       ".output_html .c1 { color: #3D7B7B; font-style: italic } /* Comment.Single */\n",
       ".output_html .cs { color: #3D7B7B; font-style: italic } /* Comment.Special */\n",
       ".output_html .gd { color: #A00000 } /* Generic.Deleted */\n",
       ".output_html .ge { font-style: italic } /* Generic.Emph */\n",
       ".output_html .gr { color: #E40000 } /* Generic.Error */\n",
       ".output_html .gh { color: #000080; font-weight: bold } /* Generic.Heading */\n",
       ".output_html .gi { color: #008400 } /* Generic.Inserted */\n",
       ".output_html .go { color: #717171 } /* Generic.Output */\n",
       ".output_html .gp { color: #000080; font-weight: bold } /* Generic.Prompt */\n",
       ".output_html .gs { font-weight: bold } /* Generic.Strong */\n",
       ".output_html .gu { color: #800080; font-weight: bold } /* Generic.Subheading */\n",
       ".output_html .gt { color: #0044DD } /* Generic.Traceback */\n",
       ".output_html .kc { color: #008000; font-weight: bold } /* Keyword.Constant */\n",
       ".output_html .kd { color: #008000; font-weight: bold } /* Keyword.Declaration */\n",
       ".output_html .kn { color: #008000; font-weight: bold } /* Keyword.Namespace */\n",
       ".output_html .kp { color: #008000 } /* Keyword.Pseudo */\n",
       ".output_html .kr { color: #008000; font-weight: bold } /* Keyword.Reserved */\n",
       ".output_html .kt { color: #B00040 } /* Keyword.Type */\n",
       ".output_html .m { color: #666666 } /* Literal.Number */\n",
       ".output_html .s { color: #BA2121 } /* Literal.String */\n",
       ".output_html .na { color: #687822 } /* Name.Attribute */\n",
       ".output_html .nb { color: #008000 } /* Name.Builtin */\n",
       ".output_html .nc { color: #0000FF; font-weight: bold } /* Name.Class */\n",
       ".output_html .no { color: #880000 } /* Name.Constant */\n",
       ".output_html .nd { color: #AA22FF } /* Name.Decorator */\n",
       ".output_html .ni { color: #717171; font-weight: bold } /* Name.Entity */\n",
       ".output_html .ne { color: #CB3F38; font-weight: bold } /* Name.Exception */\n",
       ".output_html .nf { color: #0000FF } /* Name.Function */\n",
       ".output_html .nl { color: #767600 } /* Name.Label */\n",
       ".output_html .nn { color: #0000FF; font-weight: bold } /* Name.Namespace */\n",
       ".output_html .nt { color: #008000; font-weight: bold } /* Name.Tag */\n",
       ".output_html .nv { color: #19177C } /* Name.Variable */\n",
       ".output_html .ow { color: #AA22FF; font-weight: bold } /* Operator.Word */\n",
       ".output_html .w { color: #bbbbbb } /* Text.Whitespace */\n",
       ".output_html .mb { color: #666666 } /* Literal.Number.Bin */\n",
       ".output_html .mf { color: #666666 } /* Literal.Number.Float */\n",
       ".output_html .mh { color: #666666 } /* Literal.Number.Hex */\n",
       ".output_html .mi { color: #666666 } /* Literal.Number.Integer */\n",
       ".output_html .mo { color: #666666 } /* Literal.Number.Oct */\n",
       ".output_html .sa { color: #BA2121 } /* Literal.String.Affix */\n",
       ".output_html .sb { color: #BA2121 } /* Literal.String.Backtick */\n",
       ".output_html .sc { color: #BA2121 } /* Literal.String.Char */\n",
       ".output_html .dl { color: #BA2121 } /* Literal.String.Delimiter */\n",
       ".output_html .sd { color: #BA2121; font-style: italic } /* Literal.String.Doc */\n",
       ".output_html .s2 { color: #BA2121 } /* Literal.String.Double */\n",
       ".output_html .se { color: #AA5D1F; font-weight: bold } /* Literal.String.Escape */\n",
       ".output_html .sh { color: #BA2121 } /* Literal.String.Heredoc */\n",
       ".output_html .si { color: #A45A77; font-weight: bold } /* Literal.String.Interpol */\n",
       ".output_html .sx { color: #008000 } /* Literal.String.Other */\n",
       ".output_html .sr { color: #A45A77 } /* Literal.String.Regex */\n",
       ".output_html .s1 { color: #BA2121 } /* Literal.String.Single */\n",
       ".output_html .ss { color: #19177C } /* Literal.String.Symbol */\n",
       ".output_html .bp { color: #008000 } /* Name.Builtin.Pseudo */\n",
       ".output_html .fm { color: #0000FF } /* Name.Function.Magic */\n",
       ".output_html .vc { color: #19177C } /* Name.Variable.Class */\n",
       ".output_html .vg { color: #19177C } /* Name.Variable.Global */\n",
       ".output_html .vi { color: #19177C } /* Name.Variable.Instance */\n",
       ".output_html .vm { color: #19177C } /* Name.Variable.Magic */\n",
       ".output_html .il { color: #666666 } /* Literal.Number.Integer.Long */</style><div class=\"highlight\"><pre><span></span><span class=\"c1\"># from tvm.script import tir as T</span>\n",
       "<span class=\"nd\">@tvm</span><span class=\"o\">.</span><span class=\"n\">script</span><span class=\"o\">.</span><span class=\"n\">ir_module</span>\n",
       "<span class=\"k\">class</span> <span class=\"nc\">Module</span><span class=\"p\">:</span>\n",
       "    <span class=\"nd\">@T</span><span class=\"o\">.</span><span class=\"n\">prim_func</span>\n",
       "    <span class=\"k\">def</span> <span class=\"nf\">mm_relu</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">],</span> <span class=\"n\">C</span><span class=\"p\">:</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">Buffer</span><span class=\"p\">[(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">),</span> <span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">])</span> <span class=\"o\">-&gt;</span> <span class=\"kc\">None</span><span class=\"p\">:</span>\n",
       "        <span class=\"c1\"># function attr dict</span>\n",
       "        <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">func_attr</span><span class=\"p\">({</span><span class=\"s2\">&quot;global_symbol&quot;</span><span class=\"p\">:</span> <span class=\"s2\">&quot;mm_relu&quot;</span><span class=\"p\">,</span> <span class=\"s2\">&quot;tir.noalias&quot;</span><span class=\"p\">:</span> <span class=\"kc\">True</span><span class=\"p\">})</span>\n",
       "        <span class=\"c1\"># body</span>\n",
       "        <span class=\"c1\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span class=\"n\">Y</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">alloc_buffer</span><span class=\"p\">([</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">],</span> <span class=\"n\">dtype</span><span class=\"o\">=</span><span class=\"s2\">&quot;float32&quot;</span><span class=\"p\">)</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i0</span><span class=\"p\">,</span> <span class=\"n\">i1</span><span class=\"p\">,</span> <span class=\"n\">i2</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;Y&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">,</span> <span class=\"n\">k</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SSR&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i0</span><span class=\"p\">,</span> <span class=\"n\">i1</span><span class=\"p\">,</span> <span class=\"n\">i2</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">],</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">k</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">])</span>\n",
       "                <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">init</span><span class=\"p\">():</span>\n",
       "                    <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">)</span>\n",
       "                <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">]</span> <span class=\"o\">+</span> <span class=\"n\">A</span><span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">k</span><span class=\"p\">]</span> <span class=\"o\">*</span> <span class=\"n\">B</span><span class=\"p\">[</span><span class=\"n\">k</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">]</span>\n",
       "        <span class=\"k\">for</span> <span class=\"n\">i0</span><span class=\"p\">,</span> <span class=\"n\">i1</span> <span class=\"ow\">in</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">grid</span><span class=\"p\">(</span><span class=\"mi\">128</span><span class=\"p\">,</span> <span class=\"mi\">128</span><span class=\"p\">):</span>\n",
       "            <span class=\"k\">with</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">block</span><span class=\"p\">(</span><span class=\"s2\">&quot;C&quot;</span><span class=\"p\">):</span>\n",
       "                <span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">axis</span><span class=\"o\">.</span><span class=\"n\">remap</span><span class=\"p\">(</span><span class=\"s2\">&quot;SS&quot;</span><span class=\"p\">,</span> <span class=\"p\">[</span><span class=\"n\">i0</span><span class=\"p\">,</span> <span class=\"n\">i1</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">reads</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">writes</span><span class=\"p\">(</span><span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">])</span>\n",
       "                <span class=\"n\">C</span><span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">]</span> <span class=\"o\">=</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">max</span><span class=\"p\">(</span><span class=\"n\">Y</span><span class=\"p\">[</span><span class=\"n\">i</span><span class=\"p\">,</span> <span class=\"n\">j</span><span class=\"p\">],</span> <span class=\"n\">T</span><span class=\"o\">.</span><span class=\"n\">float32</span><span class=\"p\">(</span><span class=\"mi\">0</span><span class=\"p\">))</span>\n",
       "    \n",
       "</pre></div>\n"
      ],
      "text/latex": [
       "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n",
       "\\PY{c+c1}{\\PYZsh{} from tvm.script import tir as T}\n",
       "\\PY{n+nd}{@tvm}\\PY{o}{.}\\PY{n}{script}\\PY{o}{.}\\PY{n}{ir\\PYZus{}module}\n",
       "\\PY{k}{class} \\PY{n+nc}{Module}\\PY{p}{:}\n",
       "    \\PY{n+nd}{@T}\\PY{o}{.}\\PY{n}{prim\\PYZus{}func}\n",
       "    \\PY{k}{def} \\PY{n+nf}{mm\\PYZus{}relu}\\PY{p}{(}\\PY{n}{A}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{,} \\PY{n}{C}\\PY{p}{:} \\PY{n}{T}\\PY{o}{.}\\PY{n}{Buffer}\\PY{p}{[}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{]}\\PY{p}{)} \\PY{o}{\\PYZhy{}}\\PY{o}{\\PYZgt{}} \\PY{k+kc}{None}\\PY{p}{:}\n",
       "        \\PY{c+c1}{\\PYZsh{} function attr dict}\n",
       "        \\PY{n}{T}\\PY{o}{.}\\PY{n}{func\\PYZus{}attr}\\PY{p}{(}\\PY{p}{\\PYZob{}}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{global\\PYZus{}symbol}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{mm\\PYZus{}relu}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{tir.noalias}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{:} \\PY{k+kc}{True}\\PY{p}{\\PYZcb{}}\\PY{p}{)}\n",
       "        \\PY{c+c1}{\\PYZsh{} body}\n",
       "        \\PY{c+c1}{\\PYZsh{} with T.block(\\PYZdq{}root\\PYZdq{})}\n",
       "        \\PY{n}{Y} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{alloc\\PYZus{}buffer}\\PY{p}{(}\\PY{p}{[}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{]}\\PY{p}{,} \\PY{n}{dtype}\\PY{o}{=}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{float32}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\n",
       "        \\PY{k}{for} \\PY{n}{i0}\\PY{p}{,} \\PY{n}{i1}\\PY{p}{,} \\PY{n}{i2} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{Y}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{,} \\PY{n}{k} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SSR}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i0}\\PY{p}{,} \\PY{n}{i1}\\PY{p}{,} \\PY{n}{i2}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{A}\\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{k}\\PY{p}{]}\\PY{p}{,} \\PY{n}{B}\\PY{p}{[}\\PY{n}{k}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{init}\\PY{p}{(}\\PY{p}{)}\\PY{p}{:}\n",
       "                    \\PY{n}{Y}\\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\n",
       "                \\PY{n}{Y}\\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]} \\PY{o}{=} \\PY{n}{Y}\\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]} \\PY{o}{+} \\PY{n}{A}\\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{k}\\PY{p}{]} \\PY{o}{*} \\PY{n}{B}\\PY{p}{[}\\PY{n}{k}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\n",
       "        \\PY{k}{for} \\PY{n}{i0}\\PY{p}{,} \\PY{n}{i1} \\PY{o+ow}{in} \\PY{n}{T}\\PY{o}{.}\\PY{n}{grid}\\PY{p}{(}\\PY{l+m+mi}{128}\\PY{p}{,} \\PY{l+m+mi}{128}\\PY{p}{)}\\PY{p}{:}\n",
       "            \\PY{k}{with} \\PY{n}{T}\\PY{o}{.}\\PY{n}{block}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{C}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{)}\\PY{p}{:}\n",
       "                \\PY{n}{i}\\PY{p}{,} \\PY{n}{j} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{axis}\\PY{o}{.}\\PY{n}{remap}\\PY{p}{(}\\PY{l+s+s2}{\\PYZdq{}}\\PY{l+s+s2}{SS}\\PY{l+s+s2}{\\PYZdq{}}\\PY{p}{,} \\PY{p}{[}\\PY{n}{i0}\\PY{p}{,} \\PY{n}{i1}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{reads}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{T}\\PY{o}{.}\\PY{n}{writes}\\PY{p}{(}\\PY{n}{C}\\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{)}\n",
       "                \\PY{n}{C}\\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]} \\PY{o}{=} \\PY{n}{T}\\PY{o}{.}\\PY{n}{max}\\PY{p}{(}\\PY{n}{Y}\\PY{p}{[}\\PY{n}{i}\\PY{p}{,} \\PY{n}{j}\\PY{p}{]}\\PY{p}{,} \\PY{n}{T}\\PY{o}{.}\\PY{n}{float32}\\PY{p}{(}\\PY{l+m+mi}{0}\\PY{p}{)}\\PY{p}{)}\n",
       "    \n",
       "\\end{Verbatim}\n"
      ],
      "text/plain": [
       "# from tvm.script import tir as T\n",
       "@tvm.script.ir_module\n",
       "class Module:\n",
       "    @T.prim_func\n",
       "    def mm_relu(A: T.Buffer[(128, 128), \"float32\"], B: T.Buffer[(128, 128), \"float32\"], C: T.Buffer[(128, 128), \"float32\"]) -> None:\n",
       "        # function attr dict\n",
       "        T.func_attr({\"global_symbol\": \"mm_relu\", \"tir.noalias\": True})\n",
       "        # body\n",
       "        # with T.block(\"root\")\n",
       "        Y = T.alloc_buffer([128, 128], dtype=\"float32\")\n",
       "        for i0, i1, i2 in T.grid(128, 128, 128):\n",
       "            with T.block(\"Y\"):\n",
       "                i, j, k = T.axis.remap(\"SSR\", [i0, i1, i2])\n",
       "                T.reads(A[i, k], B[k, j])\n",
       "                T.writes(Y[i, j])\n",
       "                with T.init():\n",
       "                    Y[i, j] = T.float32(0)\n",
       "                Y[i, j] = Y[i, j] + A[i, k] * B[k, j]\n",
       "        for i0, i1 in T.grid(128, 128):\n",
       "            with T.block(\"C\"):\n",
       "                i, j = T.axis.remap(\"SS\", [i0, i1])\n",
       "                T.reads(Y[i, j])\n",
       "                T.writes(C[i, j])\n",
       "                C[i, j] = T.max(Y[i, j], T.float32(0))\n",
       "    "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from tvm import te\n",
    "\n",
    "A = te.placeholder((128, 128), \"float32\", name=\"A\")\n",
    "B = te.placeholder((128, 128), \"float32\", name=\"B\")\n",
    "k = te.reduce_axis((0, 128), \"k\")\n",
    "Y = te.compute((128, 128), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name=\"Y\")\n",
    "C = te.compute((128, 128), lambda i, j: te.max(Y[i, j], 0), name=\"C\")\n",
    "\n",
    "# 由上述TE表达式转换为TensorIR\n",
    "te_func = te.create_prim_func([A, B, C]).with_attr({\"global_symbol\": \"mm_relu\"})\n",
    "MyModuleFromTE = tvm.IRModule({\"mm_relu\": te_func})\n",
    "IPython.display.Code(MyModuleFromTE.script(), language=\"python\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "MLC的开发过程\n",
    "![mlc_process](img/mlc_process.png)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('base': conda)",
   "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.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2f394aca7ca06fed1e6064aef884364492d7cdda3614a461e02e6407fc40ba69"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}