{ "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# Compute and Reduce with Tuple Inputs\n**Author**: [Ziheng Jiang](https://github.com/ZihengJiang)\n\nOften we want to compute multiple outputs with the same shape within\na single loop or perform reduction that involves multiple values like\n:code:`argmax`. These problems can be addressed by tuple inputs.\n\nIn this tutorial, we will introduce the usage of tuple inputs in TVM.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from __future__ import absolute_import, print_function\n\n\nimport tvm\nfrom tvm import te\nimport numpy as np" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Describe Batchwise Computation\nFor operators which have the same shape, we can put them together as\nthe inputs of :any:`te.compute`, if we want them to be scheduled\ntogether in the next schedule procedure.\n\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "n = te.var(\"n\")\nm = te.var(\"m\")\nA0 = te.placeholder((m, n), name=\"A0\")\nA1 = te.placeholder((m, n), name=\"A1\")\nB0, B1 = te.compute((m, n), lambda i, j: (A0[i, j] + 2, A1[i, j] * 3), name=\"B\")\n\n# The generated IR code would be:\ns = te.create_schedule(B0.op)\nprint(tvm.lower(s, [A0, A1, B0, B1], simple_mode=True))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n## Describe Reduction with Collaborative Inputs\nSometimes, we require multiple inputs to express some reduction\noperators, and the inputs will collaborate together, e.g. :code:`argmax`.\nIn the reduction procedure, :code:`argmax` need to compare the value of\noperands, also need to keep the index of operand. It can be expressed\nwith :py:func:`te.comm_reducer` as below:\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# x and y are the operands of reduction, both of them is a tuple of index\n# and value.\ndef fcombine(x, y):\n lhs = tvm.tir.Select((x >= y), x, y)\n rhs = tvm.tir.Select((x >= y), x, y)\n return lhs, rhs\n\n\n# our identity element also need to be a tuple, so `fidentity` accepts\n# two types as inputs.\ndef fidentity(t0, t1):\n return tvm.tir.const(-1, t0), tvm.te.min_value(t1)\n\n\nargmax = te.comm_reducer(fcombine, fidentity, name=\"argmax\")\n\n# describe the reduction computation\nm = te.var(\"m\")\nn = te.var(\"n\")\nidx = te.placeholder((m, n), name=\"idx\", dtype=\"int32\")\nval = te.placeholder((m, n), name=\"val\", dtype=\"int32\")\nk = te.reduce_axis((0, n), \"k\")\nT0, T1 = te.compute((m,), lambda i: argmax((idx[i, k], val[i, k]), axis=k), name=\"T\")\n\n# the generated IR code would be:\ns = te.create_schedule(T0.op)\nprint(tvm.lower(s, [idx, val, T0, T1], simple_mode=True))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "

#### Note

For ones who are not familiar with reduction, please refer to\n `general-reduction`.

\n\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Schedule Operation with Tuple Inputs\nIt is worth mentioning that although you will get multiple outputs\nwith one batch operation, but they can only be scheduled together\nin terms of operation.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "n = te.var(\"n\")\nm = te.var(\"m\")\nA0 = te.placeholder((m, n), name=\"A0\")\nB0, B1 = te.compute((m, n), lambda i, j: (A0[i, j] + 2, A0[i, j] * 3), name=\"B\")\nA1 = te.placeholder((m, n), name=\"A1\")\nC = te.compute((m, n), lambda i, j: A1[i, j] + B0[i, j], name=\"C\")\n\ns = te.create_schedule(C.op)\ns[B0].compute_at(s[C], C.op.axis)\n# as you can see in the below generated IR code:\nprint(tvm.lower(s, [A0, A1, C], simple_mode=True))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Summary\nThis tutorial introduces the usage of tuple inputs operation.\n\n- Describe normal batchwise computation.\n- Describe reduction operation with tuple inputs.\n- Notice that you can only schedule computation in terms of operation instead of tensor.\n\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "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.16" } }, "nbformat": 4, "nbformat_minor": 0 }