2_tensor_program_abstraction.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "2-tensor-program-abstraction.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "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",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Xe3vClsD9jlq",
        "outputId": "6d336487-f44b-45cc-ca2c-bae60a295cdb"
      },
      "source": [
        "!python3 -m  pip install mlc-ai-nightly -f https://mlc.ai/wheels"
      ],
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Looking in links: https://mlc.ai/wheels\n",
            "Requirement already satisfied: mlc-ai-nightly in /usr/local/lib/python3.7/dist-packages (0.9.dev1660+g1883c5b91)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (1.21.6)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (5.4.8)\n",
            "Requirement already satisfied: attrs in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (21.4.0)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (4.4.2)\n",
            "Requirement already satisfied: tornado in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (5.1.1)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (1.4.1)\n",
            "Requirement already satisfied: cloudpickle in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (1.3.0)\n",
            "Requirement already satisfied: synr==0.6.0 in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (0.6.0)\n"
          ]
        }
      ]
    },
    {
      "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",
      "metadata": {
        "id": "vvfOgcu-YdaB"
      },
      "source": [
        "import tvm\n",
        "from tvm.ir.module import IRModule\n",
        "from tvm.script import tir as T\n",
        "import numpy as np"
      ],
      "execution_count": 44,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qCViJNUNYfTW"
      },
      "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]"
      ],
      "execution_count": 47,
      "outputs": []
    },
    {
      "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",
      "source": [
        "type(MyModule)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QKsLAcDB8Npx",
        "outputId": "8534ef46-c656-4f36-961c-f6e59e04ad6d"
      },
      "execution_count": 48,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tvm.ir.module.IRModule"
            ]
          },
          "metadata": {},
          "execution_count": 48
        }
      ]
    },
    {
      "cell_type": "markdown",
      "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 script function get a string based representation of the IRModule. This function is quite useful for inspecting the module during each step of transformation."
      ],
      "metadata": {
        "id": "GpdPoa5q8Sj7"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VXy-4v3Czax9",
        "outputId": "c933d1e0-42d5-4df2-ad9a-6eb997deb10c"
      },
      "source": [
        "print(MyModule.script())"
      ],
      "execution_count": 50,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "@tvm.script.ir_module\n",
            "class Module:\n",
            "    @tir.prim_func\n",
            "    def main(A_1: tir.Buffer[128, \"float32\"], B_1: tir.Buffer[128, \"float32\"], C_1: tir.Buffer[128, \"float32\"]) -> None:\n",
            "        # function attr dict\n",
            "        tir.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
            "        # body\n",
            "        # with tir.block(\"root\")\n",
            "        for i in tir.serial(128):\n",
            "            with tir.block(\"C\"):\n",
            "                vi = tir.axis.spatial(128, i)\n",
            "                tir.reads(A_1[vi], B_1[vi])\n",
            "                tir.writes(C_1[vi])\n",
            "                C_1[vi] = A_1[vi] + B_1[vi]\n",
            "    \n"
          ]
        }
      ]
    },
    {
      "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",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oESoqN-xaTCf",
        "outputId": "58119fbd-b737-400d-bd94-776af0709501"
      },
      "source": [
        "rt_mod = tvm.build(MyModule, target=\"llvm\")  # The module for CPU backends.\n",
        "print(type(rt_mod))"
      ],
      "execution_count": 62,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<class 'tvm.driver.build_module.OperatorModule'>\n"
          ]
        }
      ]
    },
    {
      "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",
      "metadata": {
        "id": "5I3GqwnRz-Ne"
      },
      "source": [
        "func = rt_mod[\"main\"]"
      ],
      "execution_count": 88,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bngdW1eVl683",
        "outputId": "d5e0437c-bf16-4107-aa41-e11a4e1865ce"
      },
      "source": [
        "func"
      ],
      "execution_count": 89,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tvm.runtime.packed_func.PackedFunc at 0x7fed92a40fb0>"
            ]
          },
          "metadata": {},
          "execution_count": 89
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DKxo8uq_mNlp"
      },
      "source": [
        "a = tvm.nd.array(np.arange(128, dtype=\"float32\"))"
      ],
      "execution_count": 90,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1hAFAqv_mP8W"
      },
      "source": [
        "b = tvm.nd.array(np.ones(128, dtype=\"float32\")) "
      ],
      "execution_count": 91,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TseB1UBumivT"
      },
      "source": [
        "c = tvm.nd.empty((128,), dtype=\"float32\") "
      ],
      "execution_count": 92,
      "outputs": []
    },
    {
      "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",
      "metadata": {
        "id": "SMkcgO-L0Xr5"
      },
      "source": [
        "func(a, b, c)\n"
      ],
      "execution_count": 94,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AakkTpE50b6o",
        "outputId": "43971a60-2fbb-41bb-cc47-c496bfe5dda2"
      },
      "source": [
        "print(a)\n",
        "print(b)\n",
        "print(c)"
      ],
      "execution_count": 70,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "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"
          ]
        }
      ]
    },
    {
      "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",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xwyjwh51cjWI",
        "outputId": "d8a8687a-a722-4899-c181-9ca90c7d841e"
      },
      "source": [
        "sch = tvm.tir.Schedule(MyModule)\n",
        "print(type(sch))"
      ],
      "execution_count": 81,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<class 'tvm.tir.schedule.schedule.Schedule'>\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dw7Fgw8o8HPm"
      },
      "source": [
        "Let us first try to split the loops"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kNQf8D0ic4me",
        "outputId": "1cbfd7f9-a807-4571-b2e3-66ffe052037d"
      },
      "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",
        "print(sch.mod.script())"
      ],
      "execution_count": 82,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "@tvm.script.ir_module\n",
            "class Module:\n",
            "    @tir.prim_func\n",
            "    def main(A_1: tir.Buffer[128, \"float32\"], B_1: tir.Buffer[128, \"float32\"], C_1: tir.Buffer[128, \"float32\"]) -> None:\n",
            "        # function attr dict\n",
            "        tir.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
            "        # body\n",
            "        # with tir.block(\"root\")\n",
            "        for i_0, i_1, i_2 in tir.grid(8, 4, 4):\n",
            "            with tir.block(\"C\"):\n",
            "                vi = tir.axis.spatial(128, i_0 * 16 + i_1 * 4 + i_2)\n",
            "                tir.reads(A_1[vi], B_1[vi])\n",
            "                tir.writes(C_1[vi])\n",
            "                C_1[vi] = A_1[vi] + B_1[vi]\n",
            "    \n"
          ]
        }
      ]
    },
    {
      "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",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yJWBq7lRdDmn",
        "outputId": "1957cf67-f51c-4af1-c5ca-16c9718181a0"
      },
      "source": [
        "sch.reorder(i_0, i_2, i_1)\n",
        "print(sch.mod.script())"
      ],
      "execution_count": 83,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "@tvm.script.ir_module\n",
            "class Module:\n",
            "    @tir.prim_func\n",
            "    def main(A_1: tir.Buffer[128, \"float32\"], B_1: tir.Buffer[128, \"float32\"], C_1: tir.Buffer[128, \"float32\"]) -> None:\n",
            "        # function attr dict\n",
            "        tir.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
            "        # body\n",
            "        # with tir.block(\"root\")\n",
            "        for i_0, i_2, i_1 in tir.grid(8, 4, 4):\n",
            "            with tir.block(\"C\"):\n",
            "                vi = tir.axis.spatial(128, i_0 * 16 + i_1 * 4 + i_2)\n",
            "                tir.reads(A_1[vi], B_1[vi])\n",
            "                tir.writes(C_1[vi])\n",
            "                C_1[vi] = A_1[vi] + B_1[vi]\n",
            "    \n"
          ]
        }
      ]
    },
    {
      "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",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u95zQFuldHs_",
        "outputId": "b2205442-ac70-405c-8f42-ddb23e46e012"
      },
      "source": [
        "print(sch.mod.script())"
      ],
      "execution_count": 84,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "@tvm.script.ir_module\n",
            "class Module:\n",
            "    @tir.prim_func\n",
            "    def main(A_1: tir.Buffer[128, \"float32\"], B_1: tir.Buffer[128, \"float32\"], C_1: tir.Buffer[128, \"float32\"]) -> None:\n",
            "        # function attr dict\n",
            "        tir.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
            "        # body\n",
            "        # with tir.block(\"root\")\n",
            "        for i_0, i_2, i_1 in tir.grid(8, 4, 4):\n",
            "            with tir.block(\"C\"):\n",
            "                vi = tir.axis.spatial(128, i_0 * 16 + i_1 * 4 + i_2)\n",
            "                tir.reads(A_1[vi], B_1[vi])\n",
            "                tir.writes(C_1[vi])\n",
            "                C_1[vi] = A_1[vi] + B_1[vi]\n",
            "    \n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "sch.parallel(i_0)\n",
        "print(sch.mod.script())"
      ],
      "metadata": {
        "id": "m7NFPx9fy5Wy"
      },
      "execution_count": 85,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We can build and run the transformed program\n"
      ],
      "metadata": {
        "id": "OhGlqLTG_tNv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "transformed_mod = tvm.build(sch.mod, target=\"llvm\")  # The module for CPU backends.\n",
        "transformed_mod[\"main\"](a, b, c)"
      ],
      "metadata": {
        "id": "sCIYSDrI_wGq"
      },
      "execution_count": 99,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "wj_01P4mAfu2"
      }
    },
    {
      "cell_type": "code",
      "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",
        "print(ir_mod_from_te.script())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TZAPcqbGAesY",
        "outputId": "ac4d6407-8dea-4c75-d166-e071ffee8783"
      },
      "execution_count": 103,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "@tvm.script.ir_module\n",
            "class Module:\n",
            "    @tir.prim_func\n",
            "    def main(A_1: tir.Buffer[128, \"float32\"], B_1: tir.Buffer[128, \"float32\"], C_1: tir.Buffer[128, \"float32\"]) -> None:\n",
            "        # function attr dict\n",
            "        tir.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
            "        # body\n",
            "        # with tir.block(\"root\")\n",
            "        for i0 in tir.serial(128):\n",
            "            with tir.block(\"C\"):\n",
            "                i = tir.axis.spatial(128, i0)\n",
            "                tir.reads(A_1[i], B_1[i])\n",
            "                tir.writes(C_1[i])\n",
            "                C_1[i] = A_1[i] + B_1[i]\n",
            "    \n"
          ]
        }
      ]
    },
    {
      "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",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R9ExHE3BfYYv",
        "outputId": "0d82d527-8051-4bc3-c3a7-0cbb12d9bcce"
      },
      "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",
        "print(ir_module.script())\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)"
      ],
      "execution_count": 104,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "@tvm.script.ir_module\n",
            "class Module:\n",
            "    @tir.prim_func\n",
            "    def main(A_1: tir.Buffer[(1024, 1024), \"float32\"], B_1: tir.Buffer[(1024, 1024), \"float32\"], C_1: tir.Buffer[(1024, 1024), \"float32\"]) -> None:\n",
            "        # function attr dict\n",
            "        tir.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
            "        # body\n",
            "        # with tir.block(\"root\")\n",
            "        for i0, i1, i2 in tir.grid(1024, 1024, 1024):\n",
            "            with tir.block(\"C\"):\n",
            "                m, n, k = tir.axis.remap(\"SSR\", [i0, i1, i2])\n",
            "                tir.reads(A_1[m, k], B_1[k, n])\n",
            "                tir.writes(C_1[m, n])\n",
            "                with tir.init():\n",
            "                    C_1[m, n] = tir.float32(0)\n",
            "                C_1[m, n] = C_1[m, n] + A_1[m, k] * B_1[k, n]\n",
            "    \n",
            "Baseline: 3.692298\n"
          ]
        }
      ]
    },
    {
      "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",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W60q68KRgdNL",
        "outputId": "b49a101e-5148-4cf0-df88-0e112e741381"
      },
      "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",
        "print(sch.mod.script())\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)"
      ],
      "execution_count": 105,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<class 'tvm.tir.schedule.schedule.Schedule'>\n",
            "@tvm.script.ir_module\n",
            "class Module:\n",
            "    @tir.prim_func\n",
            "    def main(A_1: tir.Buffer[(1024, 1024), \"float32\"], B_1: tir.Buffer[(1024, 1024), \"float32\"], C_1: tir.Buffer[(1024, 1024), \"float32\"]) -> None:\n",
            "        # function attr dict\n",
            "        tir.func_attr({\"global_symbol\": \"main\", \"tir.noalias\": True})\n",
            "        # body\n",
            "        # with tir.block(\"root\")\n",
            "        for i0_0, i1_0, i2, i0_1, i1_1 in tir.grid(32, 32, 1024, 32, 32):\n",
            "            with tir.block(\"C\"):\n",
            "                m = tir.axis.spatial(1024, i0_0 * 32 + i0_1)\n",
            "                n = tir.axis.spatial(1024, i1_0 * 32 + i1_1)\n",
            "                k = tir.axis.reduce(1024, i2)\n",
            "                tir.reads(A_1[m, k], B_1[k, n])\n",
            "                tir.writes(C_1[m, n])\n",
            "                with tir.init():\n",
            "                    C_1[m, n] = tir.float32(0)\n",
            "                C_1[m, n] = C_1[m, n] + A_1[m, k] * B_1[k, n]\n",
            "    \n",
            "after transformation: 0.391784\n"
          ]
        }
      ]
    },
    {
      "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."
      ]
    }
  ]
}