mlc-ai-notebooks/2_tensor_program_abstraction.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Mpn1ti5Urdsv"
      },
      "source": [
        "# Tensor Program Abstraction in Action\n",
        "\n",
        "1. 用 TVM 的 2 种 DSL 写算子。\n",
        "    - TVMScript 的 TIR 写 矩阵加法\n",
        "    - Tensor Expression（简称 TE）写矩阵乘法\n",
        "2. 程序优化程序。tvm scheduler 把矩阵乘的性能提升 ~10 倍。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "vvfOgcu-YdaB"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tvm versin: 0.13.dev0\n"
          ]
        }
      ],
      "source": [
        "import tvm\n",
        "from tvm.ir.module import IRModule\n",
        "\n",
        "from tvm.script import tir as T\n",
        "from tvm import te\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "print('tvm versin: %s' % tvm.__version__)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "BBIuE2jc1DaU"
      },
      "source": [
        "## 1. 用 TVMScript 写矩阵加法\n",
        "\n",
        "### 1.1. 用 TVMScript 实现原始程序"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "qCViJNUNYfTW"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<class 'tvm.ir.module.IRModule'>\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(\n",
              "        A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "        B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "        C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "    ):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>, <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">128</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
              "                vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[vi], B[vi])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[vi])\n",
              "                C[vi] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[vi] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[vi]\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "@tvm.script.ir_module\n",
        "class MyModule:\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",
        "print(type(MyModule))\n",
        "MyModule.show()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "tOMSOuW6aIJg"
      },
      "source": [
        "### 1.2. Build and run\n",
        "\n",
        "1. build 成可执行 module，即，一种动态库。\n",
        "2. 用 tvm runtime 加载动态库并运行。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oESoqN-xaTCf",
        "outputId": "58119fbd-b737-400d-bd94-776af0709501"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "rt_mod type: <class 'tvm.driver.build_module.OperatorModule'>\n",
            "a:  [0. 1. 2. 3. 4. 5. 6. 7. 8. 9.]\n",
            "b:  [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]\n",
            "c:  [ 1.  2.  3.  4.  5.  6.  7.  8.  9. 10.]\n"
          ]
        }
      ],
      "source": [
        "rt_mod = tvm.build(MyModule, target=\"llvm\")  # The module for CPU backends.\n",
        "print('rt_mod type: %s' % type(rt_mod))\n",
        "\n",
        "# After build, mod contains a collection of runnable functions.\n",
        "# We can retrieve each function by its name.\n",
        "func = rt_mod[\"main\"]\n",
        "\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",
        "\n",
        "func(a, b, c)\n",
        "\n",
        "print('a: ', a.asnumpy()[:10])\n",
        "print('b: ', b.asnumpy()[:10])\n",
        "print('c: ', c.asnumpy()[:10])"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "i_MIDZCOcmwp"
      },
      "source": [
        "### 1.3. Transform the program\n",
        "\n",
        "使用程序 (tvm scheduler) 优化程序 (MyModule). schedule 是 TVM 提供的调度源语，通俗的理解，就是优化算法的集合。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kNQf8D0ic4me",
        "outputId": "1cbfd7f9-a807-4571-b2e3-66ffe052037d"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(\n",
              "        A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "        B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "        C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "    ):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>, <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i_0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>parallel(<span style=\"color: #008000\">8</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">for</span> i_2, i_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">4</span>):\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
              "                    vi <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">16</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_2)\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[vi], B[vi])\n",
              "                    T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[vi])\n",
              "                    C[vi] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[vi] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[vi]\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "sch = tvm.tir.Schedule(MyModule)\n",
        "\n",
        "# get block to transform(optimize)\n",
        "block_c = sch.get_block(\"C\")\n",
        "# get loop in the block to transform(optimize)\n",
        "(i,) = sch.get_loops(block_c)\n",
        "\n",
        "# 3 transformation steps: \n",
        "# 1. tile, 2. reorder, 3. parallel\n",
        "# note: split 128 to 8 * 4 * 4\n",
        "i_0, i_1, i_2 = sch.split(i, factors=[None, 4, 4])\n",
        "sch.reorder(i_0, i_2, i_1)\n",
        "sch.parallel(i_0)\n",
        "sch.mod.show()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "OhGlqLTG_tNv"
      },
      "source": [
        "Run 优化后的程序，并验证计算结果的正确性不变\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "sCIYSDrI_wGq"
      },
      "outputs": [],
      "source": [
        "c2 = tvm.nd.empty((128,), dtype=\"float32\")\n",
        "\n",
        "transformed_mod = tvm.build(sch.mod, target=\"llvm\")  # The module for CPU backends.\n",
        "transformed_mod[\"main\"](a, b, c2)\n",
        "\n",
        "# check results are the same as before\n",
        "np.testing.assert_allclose(c.asnumpy(), c2.asnumpy(), rtol=1e-5, atol=1e-5)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "wj_01P4mAfu2"
      },
      "source": [
        "## 2. TE (Tensor Expression) 写矩阵乘法\n",
        "\n",
        "### 2.1. 用 TE 实现原始程序\n",
        "\n",
        "与 TVMScript 的区别是：\n",
        "1. TE 抽象的层次更高，使用更简单。\n",
        "2. TVMScript 更底层，控制能力更强，但，tiling 等细节，也都需要开发者自己实现。"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "作为对比，先用 TE 实现矩阵加法。很明显，比 TVMScript 简单很多。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TZAPcqbGAesY",
        "outputId": "ac4d6407-8dea-4c75-d166-e071ffee8783"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(\n",
              "        A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "        B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "        C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">128</span>,), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "    ):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>, <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">128</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
              "                v_i <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">128</span>, i)\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_i], B[v_i])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_i])\n",
              "                C[v_i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[v_i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[v_i]\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# data. input & output 内存分配\n",
        "A = te.placeholder((128, ), name=\"A\")\n",
        "B = te.placeholder((128, ), name=\"B\")\n",
        "\n",
        "# 矩阵加法\n",
        "C = te.compute((128,), lambda i: A[i] + B[i], name=\"C\")\n",
        "\n",
        "# 根据 tvm 对设计，生成一个 prim func，以便用 TVM 做优化这个 func。\n",
        "func = te.create_prim_func([A, B, C])\n",
        "# function name 标记为 main。也是 TVM 的要求，必须有一个 main 函数作为 IRModule 的入口。\n",
        "func = func.with_attr(\"global_symbol\", \"main\")\n",
        "ir_mod_from_te = IRModule({\"main\": func})\n",
        "\n",
        "ir_mod_from_te.show()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "用 TE 实现矩阵乘法。并用 TVM build & run。-- 这是优化前的原始程序。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R9ExHE3BfYYv",
        "outputId": "0d82d527-8051-4bc3-c3a7-0cbb12d9bcce"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(\n",
              "        A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "        B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "        C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "    ):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>, <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> m, n, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
              "                v_m, v_n, v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [m, n, k])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_m, v_k], B[v_k, v_n])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_m, v_n])\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
              "                    C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>)\n",
              "                C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[v_m, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[v_k, v_n]\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Baseline: 2.032805\n"
          ]
        }
      ],
      "source": [
        "M = 1024\n",
        "K = 1024\n",
        "N = 1024\n",
        "\n",
        "# The default tensor type in tvm\n",
        "dtype = \"float32\"\n",
        "\n",
        "target = \"llvm\"\n",
        "dev = tvm.device(target, 0)\n",
        "\n",
        "# Algorithm\n",
        "k = te.reduce_axis((0, K), \"k\")\n",
        "A = te.placeholder((M, K), name=\"A\")\n",
        "B = te.placeholder((K, N), name=\"B\")\n",
        "C = te.compute((M, N), lambda m, n: te.sum(A[m, k] * B[k, n], axis=k), name=\"C\")\n",
        "\n",
        "# Default schedule\n",
        "func = te.create_prim_func([A, B, C])\n",
        "func = func.with_attr(\"global_symbol\", \"main\")\n",
        "ir_module = IRModule({\"main\": func})\n",
        "ir_module.show()\n",
        "\n",
        "# build and run\n",
        "func = tvm.build(ir_module, target=\"llvm\")  # The module for CPU backends.\n",
        "\n",
        "a = tvm.nd.array(np.random.rand(M, K).astype(dtype), dev)\n",
        "b = tvm.nd.array(np.random.rand(K, N).astype(dtype), dev)\n",
        "c = tvm.nd.array(np.zeros((M, N), dtype=dtype), dev)\n",
        "func(a, b, c)\n",
        "\n",
        "evaluator = func.time_evaluator(func.entry_name, dev, number=1)\n",
        "t_baseline = evaluator(a, b, c).mean\n",
        "print(\"Baseline: %f\" % t_baseline)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "swj-gMz-1vBE"
      },
      "source": [
        "用 TVM scheduler 优化程序。修改后的 loop，对 cache 更友好。实测，性能提升 10x。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W60q68KRgdNL",
        "outputId": "b49a101e-5148-4cf0-df88-0e112e741381"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(\n",
              "        A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "        B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "        C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((<span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">1024</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>),\n",
              "    ):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>, <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>bool(<span style=\"color: #008000; font-weight: bold\">True</span>)})\n",
              "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;):</span>\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> m_0, n_0, k, m_1, n_1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">1024</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;C&quot;</span>):\n",
              "                v_m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, m_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> m_1)\n",
              "                v_n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>spatial(<span style=\"color: #008000\">1024</span>, n_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> n_1)\n",
              "                v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>reduce(<span style=\"color: #008000\">1024</span>, k)\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[v_m, v_k], B[v_k, v_n])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[v_m, v_n])\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
              "                    C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>)\n",
              "                C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> C[v_m, v_n] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[v_m, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[v_k, v_n]\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "after transformation: 0.166623. baseline: 2.032805, improved: 12.20x\n"
          ]
        }
      ],
      "source": [
        "sch = tvm.tir.Schedule(ir_module)\n",
        "block_c = sch.get_block(\"C\")\n",
        "# Get loops surronding the block\n",
        "(y, x, k) = sch.get_loops(block_c)\n",
        "# step 1: tile (split)\n",
        "block_size = 32\n",
        "yo, yi = sch.split(y, [None, block_size])\n",
        "xo, xi = sch.split(x, [None, block_size])\n",
        "\n",
        "# step 2: reorder\n",
        "sch.reorder(yo, xo, k, yi, xi)\n",
        "sch.mod.show()\n",
        "\n",
        "# build and run\n",
        "func = tvm.build(sch.mod, target=\"llvm\")  # The module for CPU backends.\n",
        "\n",
        "c = tvm.nd.array(np.zeros((M, N), dtype=dtype), dev)\n",
        "func(a, b, c)\n",
        "\n",
        "evaluator = func.time_evaluator(func.entry_name, dev, number=1)\n",
        "t_new = evaluator(a, b, c).mean\n",
        "print(\"after transformation: %f. baseline: %f, improved: %.2fx\" % (\n",
        "    t_new, t_baseline, t_baseline/t_new))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "include_colab_link": true,
      "name": "2-tensor-program-abstraction.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3.8.10 64-bit",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.16"
    },
    "vscode": {
      "interpreter": {
        "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}