{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "(tutorial-tensor-expr-get-started)=\n",
        "# 使用张量表达式处理算子\n",
        "\n",
        "**Author**: [Tianqi Chen](https://tqchen.github.io)\n",
        "\n",
        "在本教程中,把注意力转向 TVM 如何使用张量表达式(Tensor Expression,简称 TE)定义张量计算并应用循环优化。TE 以纯函数式语言描述张量计算(即每个表达式都没有副作用)。从 TVM 的整体来看,Relay 将计算描述为一组算子,这些算子都可以表示为 TE 表达式,每个 TE 表达式都接受输入张量并生成输出张量。\n",
        "\n",
        "这是关于 TVM 中张量表达式语言的介绍性教程。TVM 使用领域专用张量表达式来进行有效的内核构建。通过两个使用张量表达式语言的例子,演示基本工作流程。第一个例子介绍了 TE 和用向量加法进行调度。第二个例子扩展了这些概念,用 TE 逐步优化矩阵乘法。这个矩阵乘法的例子将作为未来涵盖 TVM 更高级功能的教程的基础。\n",
        "\n",
        "## 例 1:为 CPU 编写和调度 TE 中的向量加法\n",
        "\n",
        "让我们看看 Python 中的例子,将实现向量加法的 TE,然后是针对 CPU 的调度。首先初始化 TVM 环境。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import warnings\n",
        "import numpy as np\n",
        "# import env\n",
        "\n",
        "import tvm\n",
        "from tvm import testing\n",
        "from tvm import te\n",
        "from tvm.target import Target\n",
        "\n",
        "warnings.filterwarnings('ignore')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "如果你能确定你所针对的 CPU 并指定它,你将获得更好的性能。如果你使用 LLVM,你可以从命令 ``llc --version`` 中得到这个信息,以获得 CPU 类型,你可以检查 ``/proc/cpuinfo``,了解你的处理器可能支持的额外扩展。例如,你可以使用 ``llvm -mcpu=skylake-avx512`` 来获取带有 AVX-512 指令的 CPU。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "tgt = Target(target=\"llvm\", host=\"llvm\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 描述张量计算\n",
        "\n",
        "描述矢量加法的计算。TVM 采用了张量语义,每个中间结果都表示为一个多维数组。用户需要描述生成张量的计算规则。首先定义符号变量 `n` 来表示形状。然后定义两个占位符张量 `A` 和 `B`,具有给定的形状 `(n,)`。然后用 ``compute`` 算子来描述结果张量 `C`。``compute`` 定义了计算,其输出符合指定的张量形状,计算将在张量中的每个位置进行,由 `lambda` 函数定义。注意,虽然 `n` 是变量,但它定义了 `A`、`B` 和 `C` 张量之间的一致形状。记住,在这个阶段没有实际的计算发生,因为只是声明了计算应该如何进行。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "n = te.var(\"n\")\n",
        "A = te.placeholder((n,), name=\"A\")\n",
        "B = te.placeholder((n,), name=\"B\")\n",
        "C = te.compute(A.shape, lambda i: A[i] + B[i], name=\"C\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "```{admonition} Lambda 函数\n",
        "`te.compute` 方法的第二个参数是执行计算的函数。在这个例子中,使用一个匿名函数,也被称为 `lambda` 函数,来定义计算,在本例中是对 `A` 和 `B` 的第 `i` 个元素进行加法。\n",
        "```\n",
        "\n",
        "### 为计算创建默认的调度\n",
        "\n",
        "虽然上面几行描述了计算规则,但可以用许多不同的方式计算 `C`,以适应不同的设备。对于有多个轴的张量,你可以选择先迭代哪个轴,或者计算可以分成不同的线程。TVM 要求用户提供调度,这是关于计算应该如何进行的描述。TE 中的调度操作可以改变循环顺序,在不同的线程中分割计算,并将数据块分组,以及其他操作。调度背后的重要概念是,它们只描述计算是如何进行的,所以同一个 TE 的不同调度会产生相同的结果。\n",
        "\n",
        "TVM 允许你创建一个自然的调度,通过以行为单位迭代的方式进行 `C` 运算。\n",
        "```c\n",
        "for (int i = 0; i < n; ++i) {\n",
        "  C[i] = A[i] + B[i];\n",
        "}\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [],
      "source": [
        "s = te.create_schedule(C.op)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 为计算创建默认调度\n",
        "\n",
        "有了 TE 表达式和调度,就可以为目标语言和架构(在这里是指 LLVM 和 CPU)生成可运行的代码。向 TVM 提供调度、调度中的 TE 表达式的列表、目标和主机,以及我们要产生的函数的名称。输出的结果是类型消除的(type-erased)函数,可以直接从 Python 中调用。\n",
        "\n",
        "在下面一行,使用 `tvm.build` 来创建函数。`build` 函数需要调度、所需的函数签名(包括输入和输出)以及要编译的目标语言。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fadd = tvm.build(s, [A, B, C], tgt, name=\"myadd\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "运行这个函数,并将其输出与 `numpy` 中的相同计算进行比较。编译后的 TVM 函数暴露了简洁的 C 语言 API,可以从任何语言调用。首先创建设备,也就是 TVM 可以编译调度的设备(本例中为 CPU)。在本例中,该设备是 LLVM CPU 目标。然后可以初始化设备中的张量,并执行自定义的加法运算。为了验证计算是否正确,我们可以将 `c` 张量的输出结果与 `numpy` 进行的相同计算进行比较。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "dev = tvm.device(tgt.kind.name, 0)\n",
        "\n",
        "n = 1024\n",
        "a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)\n",
        "b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)\n",
        "c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)\n",
        "fadd(a, b, c)\n",
        "testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "为了得到这个版本与 numpy 相比有多快的比较,创建辅助函数来运行 TVM 生成代码的配置文件。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Numpy running time: 0.000016\n",
            "naive: 0.000013\n"
          ]
        }
      ],
      "source": [
        "import timeit\n",
        "\n",
        "np_repeat = 100\n",
        "np_running_time = timeit.timeit(\n",
        "    setup=\"import numpy\\n\"\n",
        "    \"n = 32768\\n\"\n",
        "    'dtype = \"float32\"\\n'\n",
        "    \"a = numpy.random.rand(n, 1).astype(dtype)\\n\"\n",
        "    \"b = numpy.random.rand(n, 1).astype(dtype)\\n\",\n",
        "    stmt=\"answer = a + b\",\n",
        "    number=np_repeat,\n",
        ")\n",
        "print(\"Numpy running time: %f\" % (np_running_time / np_repeat))\n",
        "\n",
        "\n",
        "def evaluate_addition(func, target, optimization, log):\n",
        "    dev = tvm.device(target.kind.name, 0)\n",
        "    n = 32768\n",
        "    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)\n",
        "    b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)\n",
        "    c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)\n",
        "\n",
        "    evaluator = func.time_evaluator(func.entry_name, dev, number=10)\n",
        "    mean_time = evaluator(a, b, c).mean\n",
        "    print(\"%s: %f\" % (optimization, mean_time))\n",
        "\n",
        "    log.append((optimization, mean_time))\n",
        "\n",
        "\n",
        "log = [(\"numpy\", np_running_time / np_repeat)]\n",
        "evaluate_addition(fadd, tgt, \"naive\", log=log)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 更新调度以使用并行\n",
        "\n",
        "现在已经说明了 TE 的基本原理,更深入地了解调度的作用,以及如何使用它们来为不同的架构调度张量表达式。调度是一系列应用于表达式的步骤,以多种不同的方式对其进行转换。当调度应用于 TE 中的表达式时,输入和输出保持不变,但在编译时,表达式的实现可以改变。在默认的调度中,这个张量加法是并行运行的,但是很容易在所有的处理器线程中进行并行化。可以将并行调度的操作应用到计算中。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "s[C].parallel(C.op.axis[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "`tvm.lower` 命令将生成 TE 的中间表示(IR),以及相应的调度。通过在应用不同的调度操作时降低表达式,可以看到调度对计算的顺序的影响。使用旗标 `simple_mode=True` 来返回一个可读的 C 风格语句。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()\n",
            "  attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
            "  buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*n: int32)], [], type=\"auto\"),\n",
            "             B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*n)], [], type=\"auto\"),\n",
            "             C: Buffer(C_2: Pointer(float32), float32, [(stride_2: int32*n)], [], type=\"auto\")}\n",
            "  buffer_map = {A_1: A, B_1: B, C_1: C}\n",
            "  preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [n], [stride], type=\"auto\"), B_1: B_3: Buffer(B_2, float32, [n], [stride_1], type=\"auto\"), C_1: C_3: Buffer(C_2, float32, [n], [stride_2], type=\"auto\")} {\n",
            "  for (i: int32, 0, n) \"parallel\" {\n",
            "    C[(i*stride_2)] = (A[(i*stride)] + B[(i*stride_1)])\n",
            "  }\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print(tvm.lower(s, [A, B, C], simple_mode=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "现在,TVM 有可能在独立的线程上运行这些块。编译并运行这个应用了并行操作的新调度。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "parallel: 0.000022\n"
          ]
        }
      ],
      "source": [
        "fadd_parallel = tvm.build(s, [A, B, C], tgt, name=\"myadd_parallel\")\n",
        "fadd_parallel(a, b, c)\n",
        "\n",
        "tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())\n",
        "\n",
        "evaluate_addition(fadd_parallel, tgt, \"parallel\", log=log)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 更新调度以使用矢量化\n",
        "\n",
        "现代的 CPU 也有能力对浮点值进行 SIMD 操作,我们可以对我们的计算表达式应用另一个调度,以利用这一优势。实现这一点需要多个步骤:首先,我们必须使用分割调度原语将调度分割成内循环和外循环。内循环可以使用矢量化调度原语来使用 SIMD 指令,然后外循环可以使用并行调度原语来并行化。选择分割因子为你的 CPU 上的线程数。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "vector: 0.000017\n",
            "@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()\n",
            "  attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
            "  buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*n: int32)], [], type=\"auto\"),\n",
            "             B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*n)], [], type=\"auto\"),\n",
            "             C: Buffer(C_2: Pointer(float32), float32, [(stride_2: int32*n)], [], type=\"auto\")}\n",
            "  buffer_map = {A_1: A, B_1: B, C_1: C}\n",
            "  preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [n], [stride], type=\"auto\"), B_1: B_3: Buffer(B_2, float32, [n], [stride_1], type=\"auto\"), C_1: C_3: Buffer(C_2, float32, [n], [stride_2], type=\"auto\")} {\n",
            "  for (i.outer: int32, 0, floordiv((n + 3), 4)) \"parallel\" {\n",
            "    for (i.inner.s: int32, 0, 4) {\n",
            "      if @tir.likely((((i.outer*4) + i.inner.s) < n), dtype=bool) {\n",
            "        let cse_var_1: int32 = ((i.outer*4) + i.inner.s)\n",
            "        C[(cse_var_1*stride_2)] = (A[(cse_var_1*stride)] + B[(cse_var_1*stride_1)])\n",
            "      }\n",
            "    }\n",
            "  }\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# 重新创建调度,因为我们用并行操作修改了它\n",
        "n = te.var(\"n\")\n",
        "A = te.placeholder((n,), name=\"A\")\n",
        "B = te.placeholder((n,), name=\"B\")\n",
        "C = te.compute(A.shape, lambda i: A[i] + B[i], name=\"C\")\n",
        "\n",
        "s = te.create_schedule(C.op)\n",
        "\n",
        "# 这个 factor 应该被选择来匹配适合你的 CPU 的线程数。\n",
        "# 这将根据架构的不同而变化,\n",
        "# 但一个好的规则是将这个 factor 设置为等于可用的 CPU 内核数。\n",
        "factor = 4\n",
        "\n",
        "outer, inner = s[C].split(C.op.axis[0], factor=factor)\n",
        "s[C].parallel(outer)\n",
        "s[C].vectorize(inner)\n",
        "\n",
        "fadd_vector = tvm.build(s, [A, B, C], tgt, name=\"myadd_parallel\")\n",
        "\n",
        "evaluate_addition(fadd_vector, tgt, \"vector\", log=log)\n",
        "\n",
        "print(tvm.lower(s, [A, B, C], simple_mode=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 比较不同的调度\n",
        "\n",
        "我们现在可以比较不同的调度:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "            Operator\t              Timing\t         Performance\n",
            "               numpy\t1.5522249741479753e-05\t                 1.0\n",
            "               naive\t1.3208600000000001e-05\t  0.8509462365305823\n",
            "            parallel\t         2.17076e-05\t  1.3984828463358165\n",
            "              vector\t1.6507899999999998e-05\t  1.0634991882578924\n"
          ]
        }
      ],
      "source": [
        "baseline = log[0][1]\n",
        "print(\"%s\\t%s\\t%s\" % (\"Operator\".rjust(20), \"Timing\".rjust(20), \"Performance\".rjust(20)))\n",
        "for result in log:\n",
        "    print(\n",
        "        \"%s\\t%s\\t%s\"\n",
        "        % (result[0].rjust(20), str(result[1]).rjust(20), str(result[1] / baseline).rjust(20))\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "```{admonition} 代码特殊化\n",
        "正如你可能已经注意到的,`A`、`B` 和 `C` 的声明都采取了相同的形状参数 `n`。TVM 将利用这一点,只向内核传递一个形状参数,正如你在打印的设备代码中发现的那样。这是专业化的一种形式。\n",
        "\n",
        "在主机端,TVM 会自动生成检查代码,检查参数中的约束。所以如果你把不同形状的数组传入 `fadd`,就会出现错误。\n",
        "\n",
        "我们可以做更多的特殊化。例如,我们可以在计算声明中写 `n = tvm.runtime.convert(1024)`,而不是 `n = te.var(\"n\")` 。生成的函数将只接受长度为 1024 的向量。\n",
        "```\n",
        "\n",
        "我们已经定义、调度并编译了向量加法运算符,然后我们能够在 TVM 运行时上执行它。我们可以将运算符保存为一个库,然后我们可以在以后使用 TVM 运行时加载它。\n",
        "\n",
        "### 针对 GPU 的向量加法(可选)\n",
        "\n",
        "TVM 能够针对多种架构。在下一个例子中,将针对 GPU 的向量加法进行编译。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<class 'tvm.te.tensor.Tensor'>\n",
            "-----GPU code-----\n",
            "\n",
            "#ifdef _WIN32\n",
            "  using uint = unsigned int;\n",
            "  using uchar = unsigned char;\n",
            "  using ushort = unsigned short;\n",
            "  using int64_t = long long;\n",
            "  using uint64_t = unsigned long long;\n",
            "#else\n",
            "  #define uint unsigned int\n",
            "  #define uchar unsigned char\n",
            "  #define ushort unsigned short\n",
            "  #define int64_t long long\n",
            "  #define uint64_t unsigned long long\n",
            "#endif\n",
            "extern \"C\" __global__ void __launch_bounds__(64) myadd_kernel0(float* __restrict__ C, float* __restrict__ A, float* __restrict__ B, int n, int stride, int stride1, int stride2) {\n",
            "  if (((int)blockIdx.x) < (n >> 6)) {\n",
            "    C[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = (A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)] + B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride2)]);\n",
            "  } else {\n",
            "    if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {\n",
            "      C[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = (A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)] + B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride2)]);\n",
            "    }\n",
            "  }\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# 如果你想运行这段代码,改变 ``run_cuda = True``。\n",
        "# 注意,在默认情况下,这个例子不会在 docs CI 中运行。\n",
        "\n",
        "run_cuda = True\n",
        "if run_cuda:\n",
        "    # Change this target to the correct backend for you gpu. For example: cuda (NVIDIA GPUs),\n",
        "    # rocm (Radeon GPUS), OpenCL (opencl).\n",
        "    tgt_gpu = tvm.target.Target(target=\"cuda\", host=\"llvm\")\n",
        "\n",
        "    # Recreate the schedule\n",
        "    n = te.var(\"n\")\n",
        "    A = te.placeholder((n,), name=\"A\")\n",
        "    B = te.placeholder((n,), name=\"B\")\n",
        "    C = te.compute(A.shape, lambda i: A[i] + B[i], name=\"C\")\n",
        "    print(type(C))\n",
        "\n",
        "    s = te.create_schedule(C.op)\n",
        "\n",
        "    bx, tx = s[C].split(C.op.axis[0], factor=64)\n",
        "\n",
        "    ################################################################################\n",
        "    # Finally we must bind the iteration axis bx and tx to threads in the GPU\n",
        "    # compute grid. The naive schedule is not valid for GPUs, and these are\n",
        "    # specific constructs that allow us to generate code that runs on a GPU.\n",
        "\n",
        "    s[C].bind(bx, te.thread_axis(\"blockIdx.x\"))\n",
        "    s[C].bind(tx, te.thread_axis(\"threadIdx.x\"))\n",
        "\n",
        "    ######################################################################\n",
        "    # Compilation\n",
        "    # -----------\n",
        "    # After we have finished specifying the schedule, we can compile it\n",
        "    # into a TVM function. By default TVM compiles into a type-erased\n",
        "    # function that can be directly called from the python side.\n",
        "    #\n",
        "    # In the following line, we use tvm.build to create a function.\n",
        "    # The build function takes the schedule, the desired signature of the\n",
        "    # function (including the inputs and outputs) as well as target language\n",
        "    # we want to compile to.\n",
        "    #\n",
        "    # The result of compilation fadd is a GPU device function (if GPU is\n",
        "    # involved) as well as a host wrapper that calls into the GPU\n",
        "    # function. fadd is the generated host wrapper function, it contains\n",
        "    # a reference to the generated device function internally.\n",
        "\n",
        "    fadd = tvm.build(s, [A, B, C], target=tgt_gpu, name=\"myadd\")\n",
        "\n",
        "    ################################################################################\n",
        "    # The compiled TVM function exposes a concise C API that can be invoked from\n",
        "    # any language.\n",
        "    #\n",
        "    # We provide a minimal array API in python to aid quick testing and prototyping.\n",
        "    # The array API is based on the `DLPack <https://github.com/dmlc/dlpack>`_ standard.\n",
        "    #\n",
        "    # - We first create a GPU device.\n",
        "    # - Then tvm.nd.array copies the data to the GPU.\n",
        "    # - ``fadd`` runs the actual computation\n",
        "    # - ``numpy()`` copies the GPU array back to the CPU (so we can verify correctness).\n",
        "    #\n",
        "    # Note that copying the data to and from the memory on the GPU is a required step.\n",
        "\n",
        "    dev = tvm.device(tgt_gpu.kind.name, 0)\n",
        "\n",
        "    n = 1024\n",
        "    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)\n",
        "    b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)\n",
        "    c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)\n",
        "    fadd(a, b, c)\n",
        "    tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())\n",
        "\n",
        "    ################################################################################\n",
        "    # Inspect the Generated GPU Code\n",
        "    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n",
        "    # You can inspect the generated code in TVM. The result of tvm.build is a TVM\n",
        "    # Module. fadd is the host module that contains the host wrapper, it also\n",
        "    # contains a device module for the CUDA (GPU) function.\n",
        "    #\n",
        "    # The following code fetches the device module and prints the content code.\n",
        "\n",
        "    if (\n",
        "        tgt_gpu.kind.name == \"cuda\"\n",
        "        or tgt_gpu.kind.name == \"rocm\"\n",
        "        or tgt_gpu.kind.name.startswith(\"opencl\")\n",
        "    ):\n",
        "        dev_module = fadd.imported_modules[0]\n",
        "        print(\"-----GPU code-----\")\n",
        "        print(dev_module.get_source())\n",
        "    else:\n",
        "        print(fadd.get_source())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 保存和加载已编译的模块\n",
        "\n",
        "除了运行时编译,我们还可以将编译后的模块保存到文件中,以后再加载回来。\n",
        "\n",
        "下面的代码首先执行了以下步骤:\n",
        "\n",
        "- 它将编译后的主机模块保存到一个对象文件中。\n",
        "- 然后它将设备模块保存到 ptx 文件中。\n",
        "- `cc.create_shared` 调用编译器(gcc)来创建共享库"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "['myadd.so', 'myadd.o']\n"
          ]
        }
      ],
      "source": [
        "from tvm.contrib import cc\n",
        "from tvm.contrib import utils\n",
        "\n",
        "temp = utils.tempdir()\n",
        "fadd.save(temp.relpath(\"myadd.o\"))\n",
        "if tgt.kind.name == \"cuda\":\n",
        "    fadd.imported_modules[0].save(temp.relpath(\"myadd.ptx\"))\n",
        "if tgt.kind.name == \"rocm\":\n",
        "    fadd.imported_modules[0].save(temp.relpath(\"myadd.hsaco\"))\n",
        "if tgt.kind.name.startswith(\"opencl\"):\n",
        "    fadd.imported_modules[0].save(temp.relpath(\"myadd.cl\"))\n",
        "cc.create_shared(temp.relpath(\"myadd.so\"), [temp.relpath(\"myadd.o\")])\n",
        "print(temp.listdir())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "```{admonition} 模块存储格式\n",
        "CPU(主机)模块被直接保存为共享库(`.so`)。设备代码可以有多种自定义格式。在我们的例子中,设备代码被保存在 ptx 中,还有一个元数据 json 文件。它们可以通过导入分离加载和链接。\n",
        "```\n",
        "\n",
        "### 加载已编译的模块\n",
        "\n",
        "我们可以从文件系统中加载编译好的模块并运行代码。下面的代码分别加载主机和设备模块,并将它们链接在一起。我们可以验证新加载的功能是否工作。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fadd1 = tvm.runtime.load_module(temp.relpath(\"myadd.so\"))\n",
        "if tgt.kind.name == \"cuda\":\n",
        "    fadd1_dev = tvm.runtime.load_module(temp.relpath(\"myadd.ptx\"))\n",
        "    fadd1.import_module(fadd1_dev)\n",
        "\n",
        "if tgt.kind.name == \"rocm\":\n",
        "    fadd1_dev = tvm.runtime.load_module(temp.relpath(\"myadd.hsaco\"))\n",
        "    fadd1.import_module(fadd1_dev)\n",
        "\n",
        "if tgt.kind.name.startswith(\"opencl\"):\n",
        "    fadd1_dev = tvm.runtime.load_module(temp.relpath(\"myadd.cl\"))\n",
        "    fadd1.import_module(fadd1_dev)\n",
        "\n",
        "fadd1(a, b, c)\n",
        "tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 把所有东西都装进库\n",
        "\n",
        "在上面的例子中,分别存储了设备和主机代码。TVM 也支持将所有东西作为共享库导出。在底层,将设备模块打包成二进制的 blob,并将它们与主机代码连接在一起。目前支持打包 Metal、OpenCL 和 CUDA 模块。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fadd.export_library(temp.relpath(\"myadd_pack.so\"))\n",
        "fadd2 = tvm.runtime.load_module(temp.relpath(\"myadd_pack.so\"))\n",
        "fadd2(a, b, c)\n",
        "tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "```{admonition} 运行时 API 和线程安全\n",
        "TVM 的编译模块并不依赖于 TVM 编译器。相反,它们只依赖于一个最小的运行时库。TVM 运行库包装了设备驱动程序,并提供线程安全和设备无关的调用到编译的函数。\n",
        "\n",
        "这意味着你可以从任何线程、任何 GPU 上调用已编译的 TVM 函数,只要你已经为该 GPU 编译了代码。\n",
        "```\n",
        "\n",
        "## 生成 OpenCL 代码\n",
        "\n",
        "TVM 提供代码生成功能到多个后端。我们还可以生成 OpenCL 代码或 LLVM 代码,在 CPU 后端运行。\n",
        "\n",
        "下面的代码块生成 OpenCL 代码,在 OpenCL 设备上创建阵列,并验证代码的正确性。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "if tgt.kind.name.startswith(\"opencl\"):\n",
        "    fadd_cl = tvm.build(s, [A, B, C], tgt, name=\"myadd\")\n",
        "    print(\"------opencl code------\")\n",
        "    print(fadd_cl.imported_modules[0].get_source())\n",
        "    dev = tvm.cl(0)\n",
        "    n = 1024\n",
        "    a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)\n",
        "    b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)\n",
        "    c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)\n",
        "    fadd_cl(a, b, c)\n",
        "    tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "```{admonition} TE 调度原语\n",
        "TVM 包括一些不同的调度原语:\n",
        "\n",
        "- `split`:将一个指定的轴按定义的因子(factor)分成两个轴。\n",
        "- `tile`:将一个计算按定义的 factor 分成两个轴。\n",
        "- `fuse`:融合一个计算的两个连续轴。\n",
        "- `reorder`:可以将一个计算的轴重新排序到一个定义的顺序。\n",
        "- `bind`:可以将一个计算绑定到一个特定的线程,在GPU编程中很有用。\n",
        "- `compute_at`:默认情况下,TVM 会在函数的最外层计算张量,也就是默认的根。`compute_at` 指定一个张量应该在另一个运算符的第一个计算轴上计算。\n",
        "- `compute_inline`:当标记为内联时,一个计算将被展开,然后插入到需要张量的地址中。\n",
        "- `compute_root`:将一个计算移到函数的最外层,或根部。这意味着该阶段的计算将在进入下一阶段之前被完全计算。\n",
        "\n",
        "这些原语的完整描述可以在 [调度原语](schedule_primitives) 文档页中找到。\n",
        "```\n",
        "\n",
        "## 实例2:用 TE 手动优化矩阵乘法\n",
        "\n",
        "现在我们将考虑第二个更高级的例子,演示仅用 18 行 python 代码,TVM 如何将一个普通的矩阵乘法操作加快 18 倍。\n",
        "\n",
        "矩阵乘法是计算密集型运算。为了获得良好的 CPU 性能,有两个重要的优化措施:\n",
        "\n",
        "1. 提高内存访问的高速缓存命中率。复杂的数值计算和热点内存(hot-spot memory)访问都可以通过高缓存命中率(high cache hit rate)来加速。这就要求我们将原点内存(origin )访问模式转化为符合高速缓存策略的模式。\n",
        "2. SIMD(单指令多数据),也被称为矢量处理单元。在每个周期中,SIMD 可以处理一小批数据,而不是处理一个单一的值。这就要求我们将循环体中的数据访问模式转化为统一模式,以便 LLVM 后端可以将其降低到 SIMD。\n",
        "\n",
        "本教程中使用的技术是 [资源库](https://github.com/flame/how-to-optimize-gemm) 中提到的技巧的一个子集。其中一些已经被 TVM 抽象自动应用了,但由于 TVM 的限制,其中一些不能自动应用。\n",
        "\n",
        "### 准备和性能基线\n",
        "\n",
        "我们首先收集 `numpy` 实现矩阵乘法的性能数据。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Numpy running time: 0.023805\n"
          ]
        }
      ],
      "source": [
        "import tvm\n",
        "import tvm.testing\n",
        "from tvm import te\n",
        "import numpy\n",
        "\n",
        "# The size of the matrix\n",
        "# (M, K) x (K, N)\n",
        "# You are free to try out different shapes, sometimes TVM optimization outperforms numpy with MKL.\n",
        "M = 1024\n",
        "K = 1024\n",
        "N = 1024\n",
        "\n",
        "# The default tensor data type in tvm\n",
        "dtype = \"float32\"\n",
        "\n",
        "# You will want to adjust the target to match any CPU vector extensions you\n",
        "# might have. For example, if you're using using Intel AVX2 (Advanced Vector\n",
        "# Extensions) ISA for SIMD, you can get the best performance by changing the\n",
        "# following line to ``llvm -mcpu=core-avx2``, or specific type of CPU you use.\n",
        "# Recall that you're using llvm, you can get this information from the command\n",
        "# ``llc --version`` to get the CPU type, and you can check ``/proc/cpuinfo``\n",
        "# for additional extensions that your processor might support.\n",
        "\n",
        "target = tvm.target.Target(target=\"llvm\", host=\"llvm\")\n",
        "dev = tvm.device(target.kind.name, 0)\n",
        "\n",
        "# Random generated tensor for testing\n",
        "a = tvm.nd.array(numpy.random.rand(M, K).astype(dtype), dev)\n",
        "b = tvm.nd.array(numpy.random.rand(K, N).astype(dtype), dev)\n",
        "\n",
        "# Repeatedly perform a matrix multiplication to get a performance baseline\n",
        "# for the default numpy implementation\n",
        "np_repeat = 100\n",
        "np_running_time = timeit.timeit(\n",
        "    setup=\"import numpy\\n\"\n",
        "    \"M = \" + str(M) + \"\\n\"\n",
        "    \"K = \" + str(K) + \"\\n\"\n",
        "    \"N = \" + str(N) + \"\\n\"\n",
        "    'dtype = \"float32\"\\n'\n",
        "    \"a = numpy.random.rand(M, K).astype(dtype)\\n\"\n",
        "    \"b = numpy.random.rand(K, N).astype(dtype)\\n\",\n",
        "    stmt=\"answer = numpy.dot(a, b)\",\n",
        "    number=np_repeat,\n",
        ")\n",
        "print(\"Numpy running time: %f\" % (np_running_time / np_repeat))\n",
        "\n",
        "answer = numpy.dot(a.numpy(), b.numpy())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "现在用 TVM TE 编写基本的矩阵乘法,并验证它产生的结果与 `numpy` 的实现相同。我们还写了一个函数,它将帮助衡量调度优化的性能。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "none: 2.233156\n"
          ]
        }
      ],
      "source": [
        "# TVM Matrix Multiplication using TE\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 x, y: te.sum(A[x, k] * B[k, y], axis=k), name=\"C\")\n",
        "\n",
        "# Default schedule\n",
        "s = te.create_schedule(C.op)\n",
        "func = tvm.build(s, [A, B, C], target=target, name=\"mmult\")\n",
        "\n",
        "c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)\n",
        "func(a, b, c)\n",
        "tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)\n",
        "\n",
        "\n",
        "def evaluate_operation(s, vars, target, name, optimization, log):\n",
        "    func = tvm.build(s, [A, B, C], target=target, name=\"mmult\")\n",
        "    assert func\n",
        "\n",
        "    c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)\n",
        "    func(a, b, c)\n",
        "    tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)\n",
        "\n",
        "    evaluator = func.time_evaluator(func.entry_name, dev, number=10)\n",
        "    mean_time = evaluator(a, b, c).mean\n",
        "    print(\"%s: %f\" % (optimization, mean_time))\n",
        "    log.append((optimization, mean_time))\n",
        "\n",
        "\n",
        "log = []\n",
        "\n",
        "evaluate_operation(s, [A, B, C], target=target, name=\"mmult\", optimization=\"none\", log=log)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "让我们来看看使用 TVM 低级函数的运算器和默认调度的中间表示。请注意这个实现基本上是矩阵乘法的天真实现,在 A 和 B 矩阵的索引上使用三个嵌套循环。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()\n",
            "  attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
            "  buffers = {A: Buffer(A_2: Pointer(float32), float32, [1048576], []),\n",
            "             B: Buffer(B_2: Pointer(float32), float32, [1048576], []),\n",
            "             C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}\n",
            "  buffer_map = {A_1: A, B_1: B, C_1: C}\n",
            "  preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {\n",
            "  for (x: int32, 0, 1024) {\n",
            "    for (y: int32, 0, 1024) {\n",
            "      C[((x*1024) + y)] = 0f32\n",
            "      for (k: int32, 0, 1024) {\n",
            "        let cse_var_2: int32 = (x*1024)\n",
            "        let cse_var_1: int32 = (cse_var_2 + y)\n",
            "        C[cse_var_1] = (C[cse_var_1] + (A[(cse_var_2 + k)]*B[((k*1024) + y)]))\n",
            "      }\n",
            "    }\n",
            "  }\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print(tvm.lower(s, [A, B, C], simple_mode=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 优化1:阻塞\n",
        "\n",
        "提高缓冲区命中率的一个重要技巧是阻塞,在这个过程中,你的内存访问结构是在一个块的内部有一个小的邻域,具有很高的内存定位性。在本教程中,我们选择一个 32 的块因子。这将导致一个块充满 32 * 32 * sizeof(float) 的内存区域。这相当于一个 4KB 的缓存大小,而 L1 缓存的参考缓存大小为 32KB。\n",
        "\n",
        "我们首先为 ``C`` 操作创建一个默认的调度,然后用指定的块因子对其应用一个 `tile` 调度原语,调度原语返回所产生的循环顺序,从最外层到最内层,作为一个向量 `[x_outer, y_outer, x_inner, y_inner]`。然后我们得到操作输出的还原轴,并使用 4 的因子对其进行分割操作。这个因子并不直接影响我们现在正在进行的阻塞优化,但在以后我们应用矢量化时将会很有用。\n",
        "\n",
        "现在操作已经被阻塞了,我们可以重新调度计算的顺序,把减少操作放到计算的最外层循环中,帮助保证被阻止的数据仍然在缓存中。这样就完成了调度,我们可以建立并测试与原生的调度相比的性能。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "blocking: 0.249289\n"
          ]
        }
      ],
      "source": [
        "bn = 32\n",
        "\n",
        "# Blocking by loop tiling\n",
        "xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)\n",
        "(k,) = s[C].op.reduce_axis\n",
        "ko, ki = s[C].split(k, factor=4)\n",
        "\n",
        "# Hoist reduction domain outside the blocking loop\n",
        "s[C].reorder(xo, yo, ko, ki, xi, yi)\n",
        "\n",
        "evaluate_operation(s, [A, B, C], target=target, name=\"mmult\", optimization=\"blocking\", log=log)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "通过重新安排计算顺序以利用缓存,你应该看到计算的性能有了明显的改善。现在,打印内部表示,并将其与原始表示进行比较。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()\n",
            "  attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
            "  buffers = {A: Buffer(A_2: Pointer(float32), float32, [1048576], []),\n",
            "             B: Buffer(B_2: Pointer(float32), float32, [1048576], []),\n",
            "             C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}\n",
            "  buffer_map = {A_1: A, B_1: B, C_1: C}\n",
            "  preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {\n",
            "  for (x.outer: int32, 0, 32) {\n",
            "    for (y.outer: int32, 0, 32) {\n",
            "      for (x.inner.init: int32, 0, 32) {\n",
            "        for (y.inner.init: int32, 0, 32) {\n",
            "          C[((((x.outer*32768) + (x.inner.init*1024)) + (y.outer*32)) + y.inner.init)] = 0f32\n",
            "        }\n",
            "      }\n",
            "      for (k.outer: int32, 0, 256) {\n",
            "        for (k.inner: int32, 0, 4) {\n",
            "          for (x.inner: int32, 0, 32) {\n",
            "            for (y.inner: int32, 0, 32) {\n",
            "              let cse_var_3: int32 = (y.outer*32)\n",
            "              let cse_var_2: int32 = ((x.outer*32768) + (x.inner*1024))\n",
            "              let cse_var_1: int32 = ((cse_var_2 + cse_var_3) + y.inner)\n",
            "              C[cse_var_1] = (C[cse_var_1] + (A[((cse_var_2 + (k.outer*4)) + k.inner)]*B[((((k.outer*4096) + (k.inner*1024)) + cse_var_3) + y.inner)]))\n",
            "            }\n",
            "          }\n",
            "        }\n",
            "      }\n",
            "    }\n",
            "  }\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "print(tvm.lower(s, [A, B, C], simple_mode=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 优化 2: 矢量化\n",
        "\n",
        "另一个重要的优化技巧是矢量化。当内存访问模式是统一的,编译器可以检测到这种模式并将连续的内存传递给 SIMD 矢量处理器。在 TVM 中,我们可以使用 ``vectorize`` 接口来提示编译器这种模式,利用这一硬件特性。\n",
        "\n",
        "在本教程中,我们选择对内循环的行数据进行矢量化,因为在我们之前的优化中,它已经是缓存友好的。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "vectorization: 0.257297\n",
            "@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()\n",
            "  attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
            "  buffers = {A: Buffer(A_2: Pointer(float32), float32, [1048576], []),\n",
            "             B: Buffer(B_2: Pointer(float32), float32, [1048576], []),\n",
            "             C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}\n",
            "  buffer_map = {A_1: A, B_1: B, C_1: C}\n",
            "  preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {\n",
            "  for (x.outer: int32, 0, 32) {\n",
            "    for (y.outer: int32, 0, 32) {\n",
            "      for (x.inner.init: int32, 0, 32) {\n",
            "        C[ramp((((x.outer*32768) + (x.inner.init*1024)) + (y.outer*32)), 1, 32)] = broadcast(0f32, 32)\n",
            "      }\n",
            "      for (k.outer: int32, 0, 256) {\n",
            "        for (k.inner: int32, 0, 4) {\n",
            "          for (x.inner: int32, 0, 32) {\n",
            "            let cse_var_3: int32 = (y.outer*32)\n",
            "            let cse_var_2: int32 = ((x.outer*32768) + (x.inner*1024))\n",
            "            let cse_var_1: int32 = (cse_var_2 + cse_var_3)\n",
            "            C[ramp(cse_var_1, 1, 32)] = (C[ramp(cse_var_1, 1, 32)] + (broadcast(A[((cse_var_2 + (k.outer*4)) + k.inner)], 32)*B[ramp((((k.outer*4096) + (k.inner*1024)) + cse_var_3), 1, 32)]))\n",
            "          }\n",
            "        }\n",
            "      }\n",
            "    }\n",
            "  }\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Apply the vectorization optimization\n",
        "s[C].vectorize(yi)\n",
        "\n",
        "evaluate_operation(s, [A, B, C], target=target, name=\"mmult\", optimization=\"vectorization\", log=log)\n",
        "\n",
        "# The generalized IR after vectorization\n",
        "print(tvm.lower(s, [A, B, C], simple_mode=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 优化3:循环交换\n",
        "\n",
        "如果我们看一下上面的 IR,我们可以看到内循环的行数据被矢量化,B 被转化为 PackedB(这从内循环的 `(float32x32*)B2` 部分可以看出)。现在 PackedB 的遍历是顺序的。所以我们要看一下 A 的访问模式。在当前的计划中,A 是被逐列访问的,这对缓冲区不友好。如果我们改变 `ki` 和内轴 `xi` 的嵌套循环顺序,A 矩阵的访问模式将对缓存更友好。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "loop permutation: 0.130006\n",
            "@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()\n",
            "  attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
            "  buffers = {A: Buffer(A_2: Pointer(float32), float32, [1048576], []),\n",
            "             B: Buffer(B_2: Pointer(float32), float32, [1048576], []),\n",
            "             C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}\n",
            "  buffer_map = {A_1: A, B_1: B, C_1: C}\n",
            "  preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {\n",
            "  for (x.outer: int32, 0, 32) {\n",
            "    for (y.outer: int32, 0, 32) {\n",
            "      for (x.inner.init: int32, 0, 32) {\n",
            "        C[ramp((((x.outer*32768) + (x.inner.init*1024)) + (y.outer*32)), 1, 32)] = broadcast(0f32, 32)\n",
            "      }\n",
            "      for (k.outer: int32, 0, 256) {\n",
            "        for (x.inner: int32, 0, 32) {\n",
            "          for (k.inner: int32, 0, 4) {\n",
            "            let cse_var_3: int32 = (y.outer*32)\n",
            "            let cse_var_2: int32 = ((x.outer*32768) + (x.inner*1024))\n",
            "            let cse_var_1: int32 = (cse_var_2 + cse_var_3)\n",
            "            C[ramp(cse_var_1, 1, 32)] = (C[ramp(cse_var_1, 1, 32)] + (broadcast(A[((cse_var_2 + (k.outer*4)) + k.inner)], 32)*B[ramp((((k.outer*4096) + (k.inner*1024)) + cse_var_3), 1, 32)]))\n",
            "          }\n",
            "        }\n",
            "      }\n",
            "    }\n",
            "  }\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "s = te.create_schedule(C.op)\n",
        "xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)\n",
        "(k,) = s[C].op.reduce_axis\n",
        "ko, ki = s[C].split(k, factor=4)\n",
        "\n",
        "# re-ordering\n",
        "s[C].reorder(xo, yo, ko, xi, ki, yi)\n",
        "s[C].vectorize(yi)\n",
        "\n",
        "evaluate_operation(\n",
        "    s, [A, B, C], target=target, name=\"mmult\", optimization=\"loop permutation\", log=log\n",
        ")\n",
        "\n",
        "# Again, print the new generalized IR\n",
        "print(tvm.lower(s, [A, B, C], simple_mode=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 优化4:数组打包\n",
        "\n",
        "另一个重要的技巧是数组打包。这个技巧是对数组的存储维度进行重新排序,将某些维度上的连续访问模式转换为扁平化后的顺序模式。\n",
        "\n",
        "![](images/array-packing.png)\n",
        "\n",
        "正如上图所示,在阻塞计算后,我们可以观察到 B 的数组访问模式(扁平化后),它是有规律的,但是不连续的。我们期望经过一些转换后,我们可以得到一个连续的访问模式。通过将 `[16][16]` 数组重新排序为 `[16/4][16][4]` 数组,当从打包的数组中抓取相应的值时,B 的访问模式将是连续的。\n",
        "\n",
        "为了实现这一目标,将不得不从新的默认调度开始,考虑到 B 的新包装,值得花点时间来评论一下。TE 是编写优化运算符的强大而富有表现力的语言,但它往往需要对你所编写的底层算法、数据结构和硬件目标有一些了解。在本教程的后面,我们将讨论一些让 TVM 承担这一负担的选项。无论如何,让我们继续编写新的优化调度。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "array packing: 0.140257\n",
            "@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()\n",
            "  attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
            "  buffers = {A: Buffer(A_2: Pointer(float32), float32, [1048576], []),\n",
            "             B: Buffer(B_2: Pointer(float32), float32, [1048576], []),\n",
            "             C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}\n",
            "  buffer_map = {A_1: A, B_1: B, C_1: C}\n",
            "  preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {\n",
            "  allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope = global {\n",
            "    for (x: int32, 0, 32) \"parallel\" {\n",
            "      for (y: int32, 0, 1024) {\n",
            "        packedB_1: Buffer(packedB, float32x32, [32768], [])[((x*1024) + y)] = B[ramp(((y*1024) + (x*32)), 1, 32)]\n",
            "      }\n",
            "    }\n",
            "    for (x.outer: int32, 0, 32) {\n",
            "      for (y.outer: int32, 0, 32) {\n",
            "        for (x.inner.init: int32, 0, 32) {\n",
            "          C[ramp((((x.outer*32768) + (x.inner.init*1024)) + (y.outer*32)), 1, 32)] = broadcast(0f32, 32)\n",
            "        }\n",
            "        for (k.outer: int32, 0, 256) {\n",
            "          for (x.inner: int32, 0, 32) {\n",
            "            for (k.inner: int32, 0, 4) {\n",
            "              let cse_var_3: int32 = ((x.outer*32768) + (x.inner*1024))\n",
            "              let cse_var_2: int32 = (k.outer*4)\n",
            "              let cse_var_1: int32 = (cse_var_3 + (y.outer*32))\n",
            "              C[ramp(cse_var_1, 1, 32)] = (C[ramp(cse_var_1, 1, 32)] + (broadcast(A[((cse_var_3 + cse_var_2) + k.inner)], 32)*packedB_1[(((y.outer*1024) + cse_var_2) + k.inner)]))\n",
            "            }\n",
            "          }\n",
            "        }\n",
            "      }\n",
            "    }\n",
            "  }\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# 必须重写算法。\n",
        "packedB = te.compute((N / bn, K, bn), lambda x, y, z: B[y, x * bn + z], name=\"packedB\")\n",
        "C = te.compute(\n",
        "    (M, N),\n",
        "    lambda x, y: te.sum(A[x, k] * packedB[y // bn, k, tvm.tir.indexmod(y, bn)], axis=k),\n",
        "    name=\"C\",\n",
        ")\n",
        "\n",
        "s = te.create_schedule(C.op)\n",
        "\n",
        "xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)\n",
        "(k,) = s[C].op.reduce_axis\n",
        "ko, ki = s[C].split(k, factor=4)\n",
        "\n",
        "s[C].reorder(xo, yo, ko, xi, ki, yi)\n",
        "s[C].vectorize(yi)\n",
        "\n",
        "x, y, z = s[packedB].op.axis\n",
        "s[packedB].vectorize(z)\n",
        "s[packedB].parallel(x)\n",
        "\n",
        "evaluate_operation(s, [A, B, C], target=target, name=\"mmult\", optimization=\"array packing\", log=log)\n",
        "\n",
        "# Here is the generated IR after array packing.\n",
        "print(tvm.lower(s, [A, B, C], simple_mode=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 优化 5:通过缓存优化块的写入\n",
        "\n",
        "到目前为止,我们所有的优化都集中在有效地访问和计算 `A` 和 `B` 矩阵的数据以计算 `C` 矩阵上。在阻塞优化之后,运算器将逐块地将结果写入 `C`,而且访问模式不是顺序的。我们可以通过使用一个顺序缓存数组来解决这个问题,使用 `cache_write`、`compute_at` 和 `unroll` 的组合来保存块结果,并在所有块结果准备好后写入 `C`。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "block caching: 0.131626\n",
            "@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()\n",
            "  attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
            "  buffers = {A: Buffer(A_2: Pointer(float32), float32, [1048576], []),\n",
            "             B: Buffer(B_2: Pointer(float32), float32, [1048576], []),\n",
            "             C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}\n",
            "  buffer_map = {A_1: A, B_1: B, C_1: C}\n",
            "  preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {\n",
            "  allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope = global;\n",
            "  allocate(C.global: Pointer(global float32), float32, [1024]), storage_scope = global {\n",
            "    for (x: int32, 0, 32) \"parallel\" {\n",
            "      for (y: int32, 0, 1024) {\n",
            "        packedB_1: Buffer(packedB, float32x32, [32768], [])[((x*1024) + y)] = B[ramp(((y*1024) + (x*32)), 1, 32)]\n",
            "      }\n",
            "    }\n",
            "    for (x.outer: int32, 0, 32) {\n",
            "      for (y.outer: int32, 0, 32) {\n",
            "        for (x.c.init: int32, 0, 32) {\n",
            "          C.global_1: Buffer(C.global, float32, [1024], [])[ramp((x.c.init*32), 1, 32)] = broadcast(0f32, 32)\n",
            "        }\n",
            "        for (k.outer: int32, 0, 256) {\n",
            "          for (x.c: int32, 0, 32) {\n",
            "            let cse_var_4: int32 = (k.outer*4)\n",
            "            let cse_var_3: int32 = (x.c*32)\n",
            "            let cse_var_2: int32 = ((y.outer*1024) + cse_var_4)\n",
            "            let cse_var_1: int32 = (((x.outer*32768) + (x.c*1024)) + cse_var_4)\n",
            "             {\n",
            "              C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[cse_var_1], 32)*packedB_1[cse_var_2]))\n",
            "              C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 1)], 32)*packedB_1[(cse_var_2 + 1)]))\n",
            "              C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 2)], 32)*packedB_1[(cse_var_2 + 2)]))\n",
            "              C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 3)], 32)*packedB_1[(cse_var_2 + 3)]))\n",
            "            }\n",
            "          }\n",
            "        }\n",
            "        for (x.inner: int32, 0, 32) {\n",
            "          for (y.inner: int32, 0, 32) {\n",
            "            C[((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)) + y.inner)] = C.global_1[((x.inner*32) + y.inner)]\n",
            "          }\n",
            "        }\n",
            "      }\n",
            "    }\n",
            "  }\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "s = te.create_schedule(C.op)\n",
        "\n",
        "# Allocate write cache\n",
        "CC = s.cache_write(C, \"global\")\n",
        "\n",
        "xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)\n",
        "\n",
        "# Write cache is computed at yo\n",
        "s[CC].compute_at(s[C], yo)\n",
        "\n",
        "# New inner axes\n",
        "xc, yc = s[CC].op.axis\n",
        "\n",
        "(k,) = s[CC].op.reduce_axis\n",
        "ko, ki = s[CC].split(k, factor=4)\n",
        "s[CC].reorder(ko, xc, ki, yc)\n",
        "s[CC].unroll(ki)\n",
        "s[CC].vectorize(yc)\n",
        "\n",
        "x, y, z = s[packedB].op.axis\n",
        "s[packedB].vectorize(z)\n",
        "s[packedB].parallel(x)\n",
        "\n",
        "evaluate_operation(s, [A, B, C], target=target, name=\"mmult\", optimization=\"block caching\", log=log)\n",
        "\n",
        "# Here is the generated IR after write cache blocking.\n",
        "print(tvm.lower(s, [A, B, C], simple_mode=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 优化6:并行化\n",
        "\n",
        "到目前为止,我们的计算只被设计为使用单核。几乎所有的现代处理器都有多个内核,计算可以从并行运行的计算中获益。最后的优化是利用线程级并行化的优势。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "parallelization: 0.026403\n",
            "@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()\n",
            "  attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
            "  buffers = {A: Buffer(A_2: Pointer(float32), float32, [1048576], []),\n",
            "             B: Buffer(B_2: Pointer(float32), float32, [1048576], []),\n",
            "             C: Buffer(C_2: Pointer(float32), float32, [1048576], [])}\n",
            "  buffer_map = {A_1: A, B_1: B, C_1: C}\n",
            "  preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [1024, 1024], []), B_1: B_3: Buffer(B_2, float32, [1024, 1024], []), C_1: C_3: Buffer(C_2, float32, [1024, 1024], [])} {\n",
            "  allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope = global {\n",
            "    for (x: int32, 0, 32) \"parallel\" {\n",
            "      for (y: int32, 0, 1024) {\n",
            "        packedB_1: Buffer(packedB, float32x32, [32768], [])[((x*1024) + y)] = B[ramp(((y*1024) + (x*32)), 1, 32)]\n",
            "      }\n",
            "    }\n",
            "    for (x.outer: int32, 0, 32) \"parallel\" {\n",
            "      allocate(C.global: Pointer(global float32), float32, [1024]), storage_scope = global;\n",
            "      for (y.outer: int32, 0, 32) {\n",
            "        for (x.c.init: int32, 0, 32) {\n",
            "          C.global_1: Buffer(C.global, float32, [1024], [])[ramp((x.c.init*32), 1, 32)] = broadcast(0f32, 32)\n",
            "        }\n",
            "        for (k.outer: int32, 0, 256) {\n",
            "          for (x.c: int32, 0, 32) {\n",
            "            let cse_var_4: int32 = (k.outer*4)\n",
            "            let cse_var_3: int32 = (x.c*32)\n",
            "            let cse_var_2: int32 = ((y.outer*1024) + cse_var_4)\n",
            "            let cse_var_1: int32 = (((x.outer*32768) + (x.c*1024)) + cse_var_4)\n",
            "             {\n",
            "              C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[cse_var_1], 32)*packedB_1[cse_var_2]))\n",
            "              C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 1)], 32)*packedB_1[(cse_var_2 + 1)]))\n",
            "              C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 2)], 32)*packedB_1[(cse_var_2 + 2)]))\n",
            "              C.global_1[ramp(cse_var_3, 1, 32)] = (C.global_1[ramp(cse_var_3, 1, 32)] + (broadcast(A[(cse_var_1 + 3)], 32)*packedB_1[(cse_var_2 + 3)]))\n",
            "            }\n",
            "          }\n",
            "        }\n",
            "        for (x.inner: int32, 0, 32) {\n",
            "          for (y.inner: int32, 0, 32) {\n",
            "            C[((((x.outer*32768) + (x.inner*1024)) + (y.outer*32)) + y.inner)] = C.global_1[((x.inner*32) + y.inner)]\n",
            "          }\n",
            "        }\n",
            "      }\n",
            "    }\n",
            "  }\n",
            "}\n",
            "\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# parallel\n",
        "s[C].parallel(xo)\n",
        "\n",
        "x, y, z = s[packedB].op.axis\n",
        "s[packedB].vectorize(z)\n",
        "s[packedB].parallel(x)\n",
        "\n",
        "evaluate_operation(\n",
        "    s, [A, B, C], target=target, name=\"mmult\", optimization=\"parallelization\", log=log\n",
        ")\n",
        "\n",
        "# Here is the generated IR after parallelization.\n",
        "print(tvm.lower(s, [A, B, C], simple_mode=True))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 矩阵乘法实例总结\n",
        "\n",
        "在应用了上述仅有 18 行代码的简单优化后,我们生成的代码可以开始接近 `numpy` 与 Math Kernel Library(MKL)的性能。由于我们在工作中一直在记录性能,所以我们可以比较一下结果。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "            Operator\t              Timing\t         Performance\n",
            "                none\t         2.233155965\t                 1.0\n",
            "            blocking\t        0.2492886151\t 0.11163063351018566\n",
            "       vectorization\t         0.257297356\t  0.1152169217164373\n",
            "    loop permutation\t        0.1300062278\t0.058216367256731225\n",
            "       array packing\t         0.140256607\t 0.06280645382509144\n",
            "       block caching\t 0.13162604679999998\t 0.05894171695258194\n",
            "     parallelization\t         0.026402827\t0.011823100318029064\n"
          ]
        }
      ],
      "source": [
        "baseline = log[0][1]\n",
        "print(\"%s\\t%s\\t%s\" % (\"Operator\".rjust(20), \"Timing\".rjust(20), \"Performance\".rjust(20)))\n",
        "for result in log:\n",
        "    print(\n",
        "        \"%s\\t%s\\t%s\"\n",
        "        % (result[0].rjust(20), str(result[1]).rjust(20), str(result[1] / baseline).rjust(20))\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "请注意,网页上的输出反映了在一个非独家 Docker 容器上的运行时间,应该被认为是不可靠的。强烈建议你自己运行该教程,观察 TVM 取得的性能提升,并仔细研究每个例子,了解对矩阵乘法操作的迭代改进。\n",
        "\n",
        "## 最后说明和总结\n",
        "\n",
        "如前所述,如何使用 TE 和调度原语进行优化,可能需要对底层架构和算法有一些了解。然而,TE 的设计是作为更复杂的算法的基础,可以搜索潜在的优化。有了这篇关于 TE 的介绍中的知识,我们现在可以开始探索 TVM 如何将调度优化过程自动化。\n",
        "\n",
        "本教程提供了一个 TVM 张量表达(TE)工作流程的演练,使用了一个矢量添加和一个矩阵乘法的例子。一般的工作流程是:\n",
        "\n",
        "- 通过一系列的操作来描述你的计算。\n",
        "- 描述我们要如何计算使用调度原语。\n",
        "- 编译到我们想要的目标函数。\n",
        "- 可以选择保存该函数以便以后加载。\n",
        "\n",
        "即将出版的教程扩展了矩阵乘法的例子,并展示了你如何建立矩阵乘法和其他操作的通用模板,这些模板具有可调整的参数,允许你为特定平台自动优化计算。\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3.8.13 ('py38': conda)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.13"
    },
    "vscode": {
      "interpreter": {
        "hash": "28558e8daad512806f5c536a1a04c119185f99f65b79002708a12162d02a79c7"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}