{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "%%shell\n# Installs the latest dev build of TVM from PyPI. If you wish to build\n# from source, see https://tvm.apache.org/docs/install/from_source.html\npip install apache-tvm --pre" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Writing a Customized Pass\n**Author**: [Jian Weng](https://were.github.io)\n\nTVM is a framework that abstracts away the heterogenity of machine learning accelerators.\nSometimes users may want customize some analysis and IR transformations\nto adapt TVM to their own specialized hardware. This tutorial helps users write\na customized pass in TVM.\n\n## Prerequisites\n\nBefore reading this tutorial, we assume readers have already known these topics well:\n\n- Writing an algorithm in TVM and schedule it. Otherwise, see example tutorials like\n `opt-gemm`.\n- The basic structure of HalideIR. Otherwise, see ``HalideIR/src/ir/IR.h`` to learn what\n attributes of IR nodes are defined.\n- Visitor design pattern. Otherwise, check the\n [Python AST module](https://docs.python.org/3/library/ast.html) to see how an AST\n visitor is implemented.\n- How a Schedule is lowered to either an IRModule class or a LLVM module. Otherwise,\n take a look at ``python/tvm/build_module.py`` to get some basics.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import tvm\nfrom tvm import te\nimport numpy as np" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We first write a very simple vector add and build it with the default schedule. Then, we use\nour customized lowering pass to manipulate the IR directly instead of using schedule primitives.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "n = tvm.tir.const(128, \"int32\")\na = te.placeholder((n,), name=\"a\")\nb = te.placeholder((n,), name=\"b\")\nc = te.compute((n,), lambda i: a[i] + b[i], name=\"c\")\n\nsch = te.create_schedule(c.op)\nir = tvm.lower(sch, [a, b, c])\nprint(ir)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Writing a Pass\nEssentially, an \"IR transformation pass\" is a function which maps a statement to a new statement.\nThus, we define this vectorize function and implement it step by step.\n\n\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "TVM already provides two class for users to both analyze and transform IR.\n\n### IR Visitor\nWe can use ``tvm.tir.stmt_functor.post_order_visit(stmt, func)`` to gather information from the Halide IR.\n``func`` is a function callback. This function will be called before exiting the current IR node,\ni.e. post-order visit. Then we leverage side effects to store the result of IR visit, because the\nreturn value of ``func`` will be ignored.\n\n
You MUST use some array to store the result of IR visit. Even the value is a single variable.\n This is mainly due to the constraints in the Python-C runtime. The variable values will be\n refreshed every recursion but the array values will be preserved.
If the pre-order function is called and returns a value which is not None, the post-order\n function will be skipped.
Here are the essential transformations done by each phase:\n - Phase 0 generates the raw IR and loop levels.\n - Phase 1 flattens the array storage.\n - Phase 2 transforms loops, like unroll, vectorization and thread-binding.\n - Phase 3 does some cleanup work.