{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 编写定制 Pass\n", "\n", "\n", "**原作者**: [Jian Weng](https://were.github.io)\n", "\n", "TVM 是抽象出机器学习加速器异质性（heterogenity）的框架。有时用户可能希望定制一些分析和 IR 变换，使 TVM 适应他们自己的专用硬件。本教程帮助用户在 TVM 中编写定制的pass。\n", "\n", "## 前提条件\n", "\n", "在阅读本教程开始之前，假设读者已经很好地了解了以下主题：\n", "\n", "- 在 TVM 中编写算法并对其进行调度。否则，请参阅示例教程，如 {ref}`opt-gemm`。\n", "- HalideIR 的基本结构。否则，请参见 ``HalideIR/src/ir/IR.h`` 来了解 IR 节点定义了哪些属性。\n", "- 访问者设计模式（Visitor design pattern）。否则，请查看 [Python AST 模块](https://docs.python.org/3/library/ast.html)，查看 AST visitor 是如何实现的。\n", "- 如何将 Schedule 降格（lower）为 IRModule class 或 LLVM module。否则，请查看 ``python/tvm/build_module.py`` 以获得一些基础知识。" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import tvm\n", "from tvm import te" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "编写非常简单的向量加法，并使用默认的调度来构建它。然后，使用定制的 lower pass 直接操作 IR，而不是使用调度原语（primitives.）。" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\u001b[38;5;129m@tvm\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mscript\u001b[38;5;129;01m.\u001b[39;00mir_module\n", "\u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mModule\u001b[39;00m:\n", " \u001b[38;5;129m@T\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mprim_func\n", " \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmain\u001b[39m(a: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m128\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], b: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m128\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], c: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m128\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;129;01m-\u001b[39;00m\u001b[38;5;129;01m>\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n", " \u001b[38;5;30;03m# function attr dict\u001b[39;00m\n", " T\u001b[38;5;129;01m.\u001b[39;00mfunc_attr({\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfrom_legacy_te_schedule\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mglobal_symbol\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmain\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtir.noalias\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m})\n", " T\u001b[38;5;129;01m.\u001b[39;00mpreflattened_buffer(a, [\u001b[38;5;28m128\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, data\u001b[38;5;129;01m=\u001b[39;00ma\u001b[38;5;129;01m.\u001b[39;00mdata)\n", " T\u001b[38;5;129;01m.\u001b[39;00mpreflattened_buffer(b, [\u001b[38;5;28m128\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, data\u001b[38;5;129;01m=\u001b[39;00mb\u001b[38;5;129;01m.\u001b[39;00mdata)\n", " T\u001b[38;5;129;01m.\u001b[39;00mpreflattened_buffer(c, [\u001b[38;5;28m128\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, data\u001b[38;5;129;01m=\u001b[39;00mc\u001b[38;5;129;01m.\u001b[39;00mdata)\n", " \u001b[38;5;30;03m# body\u001b[39;00m\n", " \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mserial(\u001b[38;5;28m128\u001b[39m):\n", " c[i] \u001b[38;5;129;01m=\u001b[39;00m a[i] \u001b[38;5;129;01m+\u001b[39;00m b[i]\n", " \n", "\n" ] } ], "source": [ "n = tvm.tir.const(128, \"int32\")\n", "a = te.placeholder((n,), name=\"a\")\n", "b = te.placeholder((n,), name=\"b\")\n", "c = te.compute((n,), lambda i: a[i] + b[i], name=\"c\")\n", "\n", "sch = te.create_schedule(c.op)\n", "ir = tvm.lower(sch, [a, b, c])\n", "ir.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 编写 Pass\n", "\n", "本质上，“IR 变换 pass” 是将语句映射到新语句的函数。因此，下面定义了一个向量化函数，并逐步实现它。\n", "\n", "TVM 已经为用户提供了两个类来分析和变换 IR。" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### IR Visitor\n", "\n", "可以使用 ``tvm.tir.stmt_functor.post_order_visit(stmt, func)`` 从 Halide IR 收集信息。``func`` 是回调函数。该函数将在退出当前 IR 节点之前调用，即后序访问（post-order visit）。然后利用副作用来存储 IR 访问的结果，因为 ``func`` 的返回值会被忽略。\n", "\n", "```{note}\n", ":class: alert alert-info\n", "\n", "你必须使用一些数组来存储 IR 访问的结果。甚至值也是 single 变量。这主要是由于 Python-C 运行时中的约束。每次递归都会刷新变量值，但保留数组值。\n", "```" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "collapsed": false }, "outputs": [], "source": [ "loops = []\n", "\n", "def find_width8(op):\n", " \"\"\"找出所有范围能被 8 除的 'tir.For' 节点。\"\"\"\n", " if isinstance(op, tvm.tir.For):\n", " if isinstance(op.extent, tvm.tir.IntImm):\n", " if op.extent.value % 8 == 0:\n", " loops.append(op)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### IR 变换\n", "\n", "变换（transformation）接口与 visitor 接口略有不同。在 visitor 中只有 post-order 回调，但是 transformation visitor 同时支持 pre-order 和 post-order 回调。如果您想保留原始 IR 节点，只需返回 None。如果您想将当前节点更改为某个节点，请使用 TVM IR maker 接口来构建它并返回此值。\n", "\n", "```{note}\n", ":class: alert alert-info\n", "\n", "如果 pre-order 函数被调用并返回非 None 的值，则 post-order 函数将被跳过。\n", "```" ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def vectorize8(op):\n", " \"\"\"Split can vectorize the loops found in `find_width8`.\"\"\"\n", " if op in loops:\n", " extent = op.extent.value\n", " name = op.loop_var.name\n", " lo, li = te.var(name + \".outer\"), te.var(name + \".inner\")\n", " body = tvm.tir.stmt_functor.substitute(op.body, {op.loop_var: lo * 8 + li})\n", " body = tvm.tir.For(li, 0, 8, tvm.tir.ForKind.VECTORIZED, body)\n", " body = tvm.tir.For(lo, 0, extent // 8, tvm.tir.ForKind.SERIAL, body)\n", " return body\n", " return None\n", "\n", "\n", "@tvm.tir.transform.prim_func_pass(opt_level=0)\n", "def vectorize(f, mod, ctx):\n", " global loops\n", " tvm.tir.stmt_functor.post_order_visit(f.body, find_width8)\n", " if not loops:\n", " return f\n", " # 最后一个 list 参数表示要转换的节点类型。\n", " # 因此，在这种情况下，只有 `For` 节点会调用 `vectorize8`\n", " return f.with_body(tvm.tir.stmt_functor.ir_transform(f.body, None, vectorize8, [\"tir.For\"]))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Glue 到 lower pass\n", "\n", "到目前为止，已经完成了这个 IR 变换过程。接下来需要做的是将这个 pass 粘合到 TVM 的 lower pass 上。\n", "\n", "在本例中，通过向 ``tir.add_lower_pass`` 提供元组参数列表，将上面编写的 pass 注入到 TVM 标准 lower pass 中。\"Tuple\" 表明 lower 的不同阶段。在 TVM 中，有四个 lower 阶段，每个阶段(phase)完成后将调用用户自定义的阶段。\n", "\n", "```{note}\n", ":class: alert alert-info\n", "以下是每个阶段所做的基本变换：\n", "\n", "- 阶段 0：生成 raw IR 和循环级别（loop levels）。\n", "- 阶段 1：对 array storage 进行扁平化（flatten）。\n", "- 阶段 2：变换循环（transforms loops）：如 unroll、vectorization 和 thread-binding。\n", "- 阶段 3：做一些清理工作。\n", "```\n", "\n", "因此，将这个变换过程放置在阶段 1 之后是一个很好的地方。" ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\u001b[38;5;129m@tvm\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mscript\u001b[38;5;129;01m.\u001b[39;00mir_module\n", "\u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mModule\u001b[39;00m:\n", " \u001b[38;5;129m@T\u001b[39m\u001b[38;5;129;01m.\u001b[39;00mprim_func\n", " \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmain\u001b[39m(a: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m128\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], b: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m128\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m], c: T\u001b[38;5;129;01m.\u001b[39;00mBuffer[\u001b[38;5;28m128\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m]) \u001b[38;5;129;01m-\u001b[39;00m\u001b[38;5;129;01m>\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n", " \u001b[38;5;30;03m# function attr dict\u001b[39;00m\n", " T\u001b[38;5;129;01m.\u001b[39;00mfunc_attr({\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfrom_legacy_te_schedule\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mglobal_symbol\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmain\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtir.noalias\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;28;01mTrue\u001b[39;00m})\n", " T\u001b[38;5;129;01m.\u001b[39;00mpreflattened_buffer(a, [\u001b[38;5;28m128\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, data\u001b[38;5;129;01m=\u001b[39;00ma\u001b[38;5;129;01m.\u001b[39;00mdata)\n", " T\u001b[38;5;129;01m.\u001b[39;00mpreflattened_buffer(b, [\u001b[38;5;28m128\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, data\u001b[38;5;129;01m=\u001b[39;00mb\u001b[38;5;129;01m.\u001b[39;00mdata)\n", " T\u001b[38;5;129;01m.\u001b[39;00mpreflattened_buffer(c, [\u001b[38;5;28m128\u001b[39m], dtype\u001b[38;5;129;01m=\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mfloat32\u001b[39m\u001b[38;5;124m\"\u001b[39m, data\u001b[38;5;129;01m=\u001b[39;00mc\u001b[38;5;129;01m.\u001b[39;00mdata)\n", " \u001b[38;5;30;03m# body\u001b[39;00m\n", " \u001b[38;5;28;01mfor\u001b[39;00m i_outer \u001b[38;5;28;01min\u001b[39;00m T\u001b[38;5;129;01m.\u001b[39;00mserial(\u001b[38;5;28m16\u001b[39m):\n", " cse_var_1: T\u001b[38;5;129;01m.\u001b[39;00mint32 \u001b[38;5;129;01m=\u001b[39;00m i_outer \u001b[38;5;129;01m*\u001b[39;00m \u001b[38;5;28m8\u001b[39m\n", " c[cse_var_1:cse_var_1 \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m8\u001b[39m] \u001b[38;5;129;01m=\u001b[39;00m a[cse_var_1:cse_var_1 \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m8\u001b[39m] \u001b[38;5;129;01m+\u001b[39;00m b[cse_var_1:cse_var_1 \u001b[38;5;129;01m+\u001b[39;00m \u001b[38;5;28m8\u001b[39m]\n", " \n", "\n" ] } ], "source": [ "with tvm.transform.PassContext(config={\"tir.add_lower_pass\":\n", " [(1, vectorize)]}):\n", " tvm.lower(sch, [a, b, c]).show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 快速视图\n", "\n", "本教程提供了编写自定义 IR 变换 pass 的快速视图：\n", "\n", "- 使用 ``tvm.tir.stmt_functor.post_order_visit`` 收集每个 IR 节点的信息。\n", "- 使用 ``tvm.tir.stmt_functor.ir_transform`` 变换 IR 节点。\n", "- 包装上面的两个，写出 IR-transformation 函数。\n", "- 使用 ``tvm.transform.PassContext`` 将该函数放入 TVM lowering pass" ] } ], "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 }