{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "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[1] >= y[1]), x[0], y[0])\n    rhs = tvm.tir.Select((x[1] >= y[1]), x[1], y[1])\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": [
        "<div class=\"alert alert-info\"><h4>Note</h4><p>For ones who are not familiar with reduction, please refer to\n  `general-reduction`.</p></div>\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[0])\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.7.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}