2_tensor_program_abstraction.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mlc-ai/notebooks/blob/main/2_tensor_program_abstraction.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mpn1ti5Urdsv"
      },
      "source": [
        "# Tensor Program Abstraction in Action\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qXysoqn-vZuF"
      },
      "source": [
        "## Install packages \n",
        "\n",
        "For the purpose of this course, we will use some on-going development in tvm, which is an open source machine learning compilation framework. We provide the following command to install a packaged version for mlc course."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Xe3vClsD9jlq",
        "outputId": "6d336487-f44b-45cc-ca2c-bae60a295cdb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Looking in links: https://mlc.ai/wheels\n",
            "Requirement already satisfied: mlc-ai-nightly in /usr/local/lib/python3.8/dist-packages (0.9.dev2227+g06c0ef3be)\n",
            "Requirement already satisfied: attrs in /usr/local/lib/python3.8/dist-packages (from mlc-ai-nightly) (22.1.0)\n",
            "Requirement already satisfied: synr==0.6.0 in /usr/local/lib/python3.8/dist-packages (from mlc-ai-nightly) (0.6.0)\n",
            "Requirement already satisfied: decorator in /root/.local/lib/python3.8/site-packages (from mlc-ai-nightly) (5.1.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.8/dist-packages (from mlc-ai-nightly) (1.23.3)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.8/dist-packages (from mlc-ai-nightly) (5.9.2)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.8/dist-packages (from mlc-ai-nightly) (1.9.1)\n",
            "Requirement already satisfied: tornado in /usr/local/lib/python3.8/dist-packages (from mlc-ai-nightly) (6.2)\n",
            "Requirement already satisfied: cloudpickle in /usr/local/lib/python3.8/dist-packages (from mlc-ai-nightly) (2.2.0)\n"
          ]
        }
      ],
      "source": [
        "!python3 -m  pip install mlc-ai-nightly -f https://mlc.ai/wheels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BBIuE2jc1DaU"
      },
      "source": [
        "## Constructing Tensor Program\n",
        "\n",
        "Let us begin by constructing a tensor program that performs addition among two vectors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "vvfOgcu-YdaB"
      },
      "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": "code",
      "execution_count": 3,
      "metadata": {
        "id": "qCViJNUNYfTW"
      },
      "outputs": [],
      "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]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4PJd0Pw8zVQD"
      },
      "source": [
        "TVMScript is a way for us to express tensor program in python ast. Note that this code do not actually correspond to a python program, but a tensor program  that can be used in MLC process. The language is designed to align with python syntax with additional structures to facilitate analysis and transformation. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QKsLAcDB8Npx",
        "outputId": "8534ef46-c656-4f36-961c-f6e59e04ad6d"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "tvm.ir.module.IRModule"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "type(MyModule)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GpdPoa5q8Sj7"
      },
      "source": [
        "MyModule is an instance of an **IRModule** data structure, which is used to hold a collection of tensor functions. \n",
        "\n",
        "We can use the `show()` function to get a highlighted string based representation of the IRModule. This function is quite useful for inspecting the module during each step of transformation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VXy-4v3Czax9",
        "outputId": "c933d1e0-42d5-4df2-ad9a-6eb997deb10c"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #1E90FF; 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: #1E90FF\">main</span>(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>], 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>], 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>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
              "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\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>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
              "        <span style=\"color: #007979; font-style: italic\"># body</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> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>serial(<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",
              "    \n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "MyModule.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tOMSOuW6aIJg"
      },
      "source": [
        "### Build and run\n",
        "\n",
        "Any any time point, we can turn an IRModule to runnable functions by calling a build function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oESoqN-xaTCf",
        "outputId": "58119fbd-b737-400d-bd94-776af0709501"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<class 'tvm.driver.build_module.OperatorModule'>\n"
          ]
        }
      ],
      "source": [
        "rt_mod = tvm.build(MyModule, target=\"llvm\")  # The module for CPU backends.\n",
        "print(type(rt_mod))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y2ZfGrH1z6SV"
      },
      "source": [
        "After build, mod contains a collection of runnable functions. We can retrieve each function by its name."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "5I3GqwnRz-Ne"
      },
      "outputs": [],
      "source": [
        "func = rt_mod[\"main\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bngdW1eVl683",
        "outputId": "d5e0437c-bf16-4107-aa41-e11a4e1865ce"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<tvm.runtime.packed_func.PackedFunc at 0x7fdf00501e40>"
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "func"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "DKxo8uq_mNlp"
      },
      "outputs": [],
      "source": [
        "a = tvm.nd.array(np.arange(128, dtype=\"float32\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "1hAFAqv_mP8W"
      },
      "outputs": [],
      "source": [
        "b = tvm.nd.array(np.ones(128, dtype=\"float32\")) "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "TseB1UBumivT"
      },
      "outputs": [],
      "source": [
        "c = tvm.nd.empty((128,), dtype=\"float32\") "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p68xZ0_P0MPw"
      },
      "source": [
        "To invoke the function, we can create three NDArrays in the tvm runtime, and then invoke the generated function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "SMkcgO-L0Xr5"
      },
      "outputs": [],
      "source": [
        "func(a, b, c)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AakkTpE50b6o",
        "outputId": "43971a60-2fbb-41bb-cc47-c496bfe5dda2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[  0.   1.   2.   3.   4.   5.   6.   7.   8.   9.  10.  11.  12.  13.\n",
            "  14.  15.  16.  17.  18.  19.  20.  21.  22.  23.  24.  25.  26.  27.\n",
            "  28.  29.  30.  31.  32.  33.  34.  35.  36.  37.  38.  39.  40.  41.\n",
            "  42.  43.  44.  45.  46.  47.  48.  49.  50.  51.  52.  53.  54.  55.\n",
            "  56.  57.  58.  59.  60.  61.  62.  63.  64.  65.  66.  67.  68.  69.\n",
            "  70.  71.  72.  73.  74.  75.  76.  77.  78.  79.  80.  81.  82.  83.\n",
            "  84.  85.  86.  87.  88.  89.  90.  91.  92.  93.  94.  95.  96.  97.\n",
            "  98.  99. 100. 101. 102. 103. 104. 105. 106. 107. 108. 109. 110. 111.\n",
            " 112. 113. 114. 115. 116. 117. 118. 119. 120. 121. 122. 123. 124. 125.\n",
            " 126. 127.]\n",
            "[1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.\n",
            " 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.\n",
            " 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.\n",
            " 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.\n",
            " 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.\n",
            " 1. 1. 1. 1. 1. 1. 1. 1.]\n",
            "[  1.   2.   3.   4.   5.   6.   7.   8.   9.  10.  11.  12.  13.  14.\n",
            "  15.  16.  17.  18.  19.  20.  21.  22.  23.  24.  25.  26.  27.  28.\n",
            "  29.  30.  31.  32.  33.  34.  35.  36.  37.  38.  39.  40.  41.  42.\n",
            "  43.  44.  45.  46.  47.  48.  49.  50.  51.  52.  53.  54.  55.  56.\n",
            "  57.  58.  59.  60.  61.  62.  63.  64.  65.  66.  67.  68.  69.  70.\n",
            "  71.  72.  73.  74.  75.  76.  77.  78.  79.  80.  81.  82.  83.  84.\n",
            "  85.  86.  87.  88.  89.  90.  91.  92.  93.  94.  95.  96.  97.  98.\n",
            "  99. 100. 101. 102. 103. 104. 105. 106. 107. 108. 109. 110. 111. 112.\n",
            " 113. 114. 115. 116. 117. 118. 119. 120. 121. 122. 123. 124. 125. 126.\n",
            " 127. 128.]\n"
          ]
        }
      ],
      "source": [
        "print(a)\n",
        "print(b)\n",
        "print(c)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i_MIDZCOcmwp"
      },
      "source": [
        "## Transform the Tensor Program\n",
        "\n",
        "Now let us start to transform the Tensor Program. A tensor prigram can be transformed using an auxiliary data structure called schedule.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xwyjwh51cjWI",
        "outputId": "d8a8687a-a722-4899-c181-9ca90c7d841e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<class 'tvm.tir.schedule.schedule.Schedule'>\n"
          ]
        }
      ],
      "source": [
        "sch = tvm.tir.Schedule(MyModule)\n",
        "print(type(sch))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dw7Fgw8o8HPm"
      },
      "source": [
        "Let us first try to split the loops"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "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: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #1E90FF; 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: #1E90FF\">main</span>(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>], 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>], 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>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
              "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\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>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
              "        <span style=\"color: #007979; font-style: italic\"># body</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, i_1, i_2 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">8</span>, <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",
              "    \n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# 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",
        "sch.mod.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nzrbvqBSdC-D"
      },
      "source": [
        "We can also reorder the loops. Now we move loop i_2 to outside of i_1.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yJWBq7lRdDmn",
        "outputId": "1957cf67-f51c-4af1-c5ca-16c9718181a0"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #1E90FF; 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: #1E90FF\">main</span>(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>], 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>], 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>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
              "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\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>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
              "        <span style=\"color: #007979; font-style: italic\"># body</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, 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\">8</span>, <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",
              "    \n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "sch.reorder(i_0, i_2, i_1)\n",
        "sch.mod.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UmUr6b_L07-8"
      },
      "source": [
        "Finally, we can add hints to the program generator that we want to vectorize the inner most loop."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u95zQFuldHs_",
        "outputId": "b2205442-ac70-405c-8f42-ddb23e46e012"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #1E90FF; 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: #1E90FF\">main</span>(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>], 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>], 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>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
              "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\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>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
              "        <span style=\"color: #007979; font-style: italic\"># body</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, 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\">8</span>, <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",
              "    \n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "sch.mod.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "m7NFPx9fy5Wy"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #1E90FF; 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: #1E90FF\">main</span>(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>], 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>], 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>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
              "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\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>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
              "        <span style=\"color: #007979; font-style: italic\"># body</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",
              "    \n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "sch.parallel(i_0)\n",
        "sch.mod.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OhGlqLTG_tNv"
      },
      "source": [
        "We can build and run the transformed program\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "sCIYSDrI_wGq"
      },
      "outputs": [],
      "source": [
        "transformed_mod = tvm.build(sch.mod, target=\"llvm\")  # The module for CPU backends.\n",
        "transformed_mod[\"main\"](a, b, c)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wj_01P4mAfu2"
      },
      "source": [
        "## Constructing Tensor Program using Tensor Expression\n",
        "\n",
        "In the previous example, we directly use TVMScript to construct the tensor program. In practice, it is usually helpful to construct these functions pragmatically from existing definitions. Tensor expression is an API that helps us to build some of the expression-like array computations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "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: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #1E90FF; 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: #1E90FF\">main</span>(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>], 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>], 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>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
              "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\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>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
              "        <span style=\"color: #007979; font-style: italic\"># body</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> i0 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>serial(<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",
              "                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>, i0)\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[i], B[i])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[i])\n",
              "                C[i] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A[i] <span style=\"color: #AA22FF; font-weight: bold\">+</span> B[i]\n",
              "    \n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "# namespace for tensor expression utility\n",
        "from tvm import te\n",
        "\n",
        "# declare the computation using the expression API\n",
        "A = te.placeholder((128, ), name=\"A\")\n",
        "B = te.placeholder((128, ), name=\"B\")\n",
        "C = te.compute((128,), lambda i: A[i] + B[i], name=\"C\")\n",
        "\n",
        "# create a function with the specified list of arguments. \n",
        "func = te.create_prim_func([A, B, C])\n",
        "# mark that the function name is main\n",
        "func = func.with_attr(\"global_symbol\", \"main\")\n",
        "ir_mod_from_te = IRModule({\"main\": func})\n",
        "\n",
        "ir_mod_from_te.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GEqpO14Lf0Lq"
      },
      "source": [
        "## Transforming a matrix multiplication program\n",
        "\n",
        "In the above example, we showed how to transform an vector add. Now let us try to apply that to a slightly more complicated program(matrix multiplication). Let us first try to build the initial code using the tensor expression API.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "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: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #1E90FF; 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: #1E90FF\">main</span>(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>], 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>], 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>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
              "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\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>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
              "        <span style=\"color: #007979; font-style: italic\"># body</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> i0, i1, i2 <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",
              "                m, n, 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>, [i0, i1, i2])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[m, k], B[k, n])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[m, n])\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
              "                    C[m, 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[m, n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> C[m, n] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[m, k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[k, n]\n",
              "    \n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Baseline: 4.970728\n"
          ]
        }
      ],
      "source": [
        "from tvm import te\n",
        "\n",
        "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",
        "\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",
        "print(\"Baseline: %f\" % evaluator(a, b, c).mean)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "swj-gMz-1vBE"
      },
      "source": [
        "We can transform the loop access pattern to make it more cache friendly. Let us use the following schedule."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W60q68KRgdNL",
        "outputId": "b49a101e-5148-4cf0-df88-0e112e741381"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<class 'tvm.tir.schedule.schedule.Schedule'>\n"
          ]
        },
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #1E90FF; 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: #1E90FF\">main</span>(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>], 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>], 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>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> <span style=\"color: #008000; font-weight: bold\">None</span>:\n",
              "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\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>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
              "        <span style=\"color: #007979; font-style: italic\"># body</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> i0_0, i1_0, i2, i0_1, i1_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",
              "                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>, i0_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i0_1)\n",
              "                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>, i1_0 <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i1_1)\n",
              "                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>, i2)\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(A[m, k], B[k, n])\n",
              "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(C[m, n])\n",
              "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
              "                    C[m, 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[m, n] <span style=\"color: #AA22FF; font-weight: bold\">=</span> C[m, n] <span style=\"color: #AA22FF; font-weight: bold\">+</span> A[m, k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B[k, n]\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.101731\n"
          ]
        }
      ],
      "source": [
        "sch = tvm.tir.Schedule(ir_module)\n",
        "print(type(sch))\n",
        "block_c = sch.get_block(\"C\")\n",
        "# Get loops surronding the block\n",
        "(y, x, k) = sch.get_loops(block_c)\n",
        "block_size = 32\n",
        "yo, yi = sch.split(y, [None, block_size])\n",
        "xo, xi = sch.split(x, [None, block_size])\n",
        "\n",
        "sch.reorder(yo, xo, k, yi, xi)\n",
        "sch.mod.show()\n",
        "\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",
        "print(\"after transformation: %f\" % evaluator(a, b, c).mean)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h1RQGOBjn4w_"
      },
      "source": [
        "Try to change the value of bn to see what performance you can get. In pratice, we will leverage an automated system to search over a set of possible transfromations to find an optimal one."
      ]
    }
  ],
  "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.10"
    },
    "vscode": {
      "interpreter": {
        "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}