tvm-custom-datatypes-notebook-example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Q4aC_aNcUP75"
   },
   "source": [
    "# Bring Your Own Datatypes to TVM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "seyxX9HZUP76"
   },
   "source": [
    "In this tutorial, we will show you how you can use your own custom datatypes in TVM, utilizing TVM's Bring Your Own Datatypes framework.\n",
    "Note that the Bring Your Own Datatypes framework currently only handles **software emulated versions of datatypes** right now, which is what we'll be discussing in this tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Datatype Libraries\n",
    "\n",
    "The central idea of the Bring Your Own Datatypes framework is to make datatypes in TVM more abstract, so that users can bring their own software-emulated datatype implementations.\n",
    "\n",
    "In the wild, these datatype implementations usually appear as libraries. For example:\n",
    "- [libposit](https://github.com/cjdelisle/libposit), a posit library\n",
    "- [Stillwater Universal](https://github.com/stillwater-sc/universal), a library with posits, fixed-point numbers, and other types\n",
    "- [SoftFloat](https://github.com/ucb-bar/berkeley-softfloat-3), Berkeley's software implementation of IEEE 754\n",
    "\n",
    "The Bring Your Own Datatypes framework allows libraries such as these to be easily plugged in to TVM!\n",
    "\n",
    "In this section, we will explore our example library, [libposit](https://github.com/cjdelisle/libposit).\n",
    "**Posits** are a datatype developed to compete with IEEE 754 floating point numbers.\n",
    "We won't go into much detail about the datatype itself.\n",
    "If you'd like to learn more, read through John Gustafson's [Beating Floating Point at its Own Game](https://posithub.org/docs/BeatingFloatingPoint.pdf).\n",
    "\n",
    "First, let's clone and build the library. Note that we're using my branch of libposit which includes some fixes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! git clone https://github.com/uwsampl/libposit\n",
    "! cd libposit && git checkout 7c1788f291c1b5f74ded9acf4ffae7911c2df28c && autoreconf -f -i && ./configure && make"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The library contains operations over posits.\n",
    "Let's see what we can do with 16-bit posits (which are much like 16-bit floats):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! nm libposit/libposit.a | grep posit16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see many operations that we might expect for datatypes: `posit16_abs` for absolute value, `posit16_add` for addition, `posit16_cmp` for comparsion, etc.\n",
    "\n",
    "Now, let's imagine we build a library which uses libposit. This is the library we want to use from TVM. Right now, we will just include a dummy function in the library, and then build it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "libposit_wrapper_source = \"\"\"\n",
    "#include \"posit.h\"\n",
    "\n",
    "void test() {\n",
    "  posit16_fromf(1.0);\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "! echo '{libposit_wrapper_source}' > libposit-wrapper.cc\n",
    "! g++ -Ilibposit/generated -lmpfr -lgmp --std=c++14 -shared -o libposit-wrapper.so -fPIC libposit-wrapper.cc libposit/libposit.a\n",
    "! ls -alh libposit-wrapper.so"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*TODO: The Bring Your Own Datatypes framework is easiest to use when the datatype library is a shared library which can be loaded straight from Python with `CDLL`. libposit, on the other hand, generates a static library, and so the above code is actually a hack to turn libposit into a shared library. By writing a small wrapper library and using a libposit function, the libposit functions get linked into our shared library and can be loaded at runtime. It would be nice to have a cleaner example that doesn't require such a hack, or figure out a way around this hack.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load the library into the process space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ctypes import CDLL, RTLD_GLOBAL\n",
    "\n",
    "lib = CDLL('/root/libposit-wrapper.so', RTLD_GLOBAL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Great! Now we have a library which uses libposit loaded into our process space. We will add to this library as needed and reload it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0fxZ29ZrUP78"
   },
   "source": [
    "## A Simple TVM Program\n",
    "We'll begin by writing a simple program in TVM; afterwards, we will re-write it to use custom datatypes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ghinqYT5UP79"
   },
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm import relay\n",
    "\n",
    "# Our basic program: Z = X + Y\n",
    "x = relay.var('x', shape=(3, ), dtype='float32')\n",
    "y = relay.var('y', shape=(3, ), dtype='float32')\n",
    "z = x + y\n",
    "program = relay.Function([x,y], z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rOyHNW9jUP8E"
   },
   "source": [
    "Now, we create random inputs to feed into this program using numpy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EIJNyNqJUP8F"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "np.random.seed(23) # for reproducibility\n",
    "\n",
    "# Create random input arrays on the above context\n",
    "x_input = np.random.rand(3).astype('float32')\n",
    "y_input = np.random.rand(3).astype('float32')\n",
    "print(\"x: {}\".format(x_input))\n",
    "print(\"y: {}\".format(y_input))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "fm8ueDQlUP8O"
   },
   "source": [
    "Finally, we're ready to run the program:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ld39h2rvUP8P"
   },
   "outputs": [],
   "source": [
    "ex = relay.create_executor(\"graph\")\n",
    "\n",
    "compiled = ex.evaluate(program)\n",
    "\n",
    "z_output = compiled(x_input, y_input)\n",
    "print(\"z: {}\".format(z_output))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7XOfXquqUP8T"
   },
   "source": [
    "## Adding Custom Datatypes\n",
    "\n",
    "Now, we will do the same, but we will use a custom datatype for our intermediate computation.\n",
    "\n",
    "We use the same input variables `x` and `y` as above, but before adding `x + y`, we first cast both `x` and `y` to a custom datatype via the `relay.cast(...)` call.\n",
    "\n",
    "Note how we specify the custom datatype: we indicate it using the special `custom[...]` syntax. Additionally, note the \"16\" after the datatype: this is the bitwidth of the custom datatype. This tells TVM that each instance of `posit` is 16 bits wide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pHhq0Tk9UP8U"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    x_posit = relay.cast(x, dtype='custom[posit]16')\n",
    "    y_posit = relay.cast(y, dtype='custom[posit]16')\n",
    "    z_posit = x_posit + y_posit\n",
    "    z = relay.cast(z_posit, dtype='float32')\n",
    "except tvm.TVMError as e:\n",
    "    # Print last line of error\n",
    "    print(str(e).split('\\n')[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RJqUIBorUP8X"
   },
   "source": [
    "Trying to generate this program throws an error from TVM. Unsurprisingly, TVM does not know how to handle any custom datatype out of the box. We first have to register the custom type with TVM, giving it a name and a type code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kJKVCD7OUP8Y"
   },
   "outputs": [],
   "source": [
    "tvm.datatype.register('posit', 150)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HJruYl15UP8b"
   },
   "source": [
    "Note that the type code, 150, is currently chosen manually by the programmer. See `TVMTypeCode::kCustomBegin` in [include/tvm/runtime/c_runtime_api.h](https://github.com/dmlc/tvm/blob/master/include/tvm/runtime/c_runtime_api.h)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QI96dWiIUP8c"
   },
   "source": [
    "Now we can generate our program again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3wZWHxd7UP8f"
   },
   "outputs": [],
   "source": [
    "x_posit = relay.cast(x, dtype='custom[posit]16')\n",
    "y_posit = relay.cast(y, dtype='custom[posit]16')\n",
    "z_posit = x_posit + y_posit\n",
    "z = relay.cast(z_posit, dtype='float32')\n",
    "program = relay.Function([x, y], z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have a Relay program that uses posits!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(program)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "YSq-jAPRUP8h"
   },
   "source": [
    "Now that we can express our program without errors, let's try running it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3WNSbuvkUP8i"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    compiled = ex.evaluate(program)\n",
    "    z_output_posit = compiled(x_input, y_input)\n",
    "    print(\"z: {}\".format(z_output_posit))\n",
    "    \n",
    "except tvm.TVMError as e:\n",
    "    # Print last line of error\n",
    "    print(str(e).split('\\n')[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rVsPSXRQUP8l"
   },
   "source": [
    "Now, trying to compile this program throws an error.\n",
    "Let's dissect this error.\n",
    "\n",
    "The error is occurring during the process of lowering the custom datatype code to code that TVM can compile and run. TVM is telling us that it cannot find a _lowering function_ for the `Cast` operation, when casting from source type 2 (`float`, in TVM), to destination type 150 (our custom datatype). When lowering custom datatypes, if TVM encounters an operation over a custom datatype, it looks for a user-registered *lowering function*, which tells it how to lower the operation to an operation over datatypes it understands. We have not told TVM how to lower `Cast` operations for our custom datatypes; thus, the source of this error."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jtbjLKUIUP8m"
   },
   "source": [
    "To fix this error, we simply need to specify a lowering function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ETiJMdqTUP8n"
   },
   "outputs": [],
   "source": [
    "tvm.datatype.register_op(tvm.datatype.create_lower_func(\"posit16_fromf\"),\n",
    "                         \"Cast\", \"llvm\", \"posit\", \"float\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "84bDvz9fUP8q"
   },
   "source": [
    "The `register_op(...)` call takes a lowering function, and a number of parameters which specify exactly the operation which should be lowered with the provided lowering function. In this case, the arguments we pass specify that this lowering function is for lowering a `Cast` from `float` to `posit` for target `\"llvm\"`.\n",
    "\n",
    "The lowering function passed into this call is very general: it should take an operation of the specified type (in this case, `Cast`) and return another operation which only uses datatypes which TVM understands.\n",
    "\n",
    "In the general case, we expect users to implement operations over their custom datatypes using calls to an external library. In our example, our `posit` library implements a `Cast` from `float` to 16-bit `posit` in the function `posit16_fromf`. To provide for the general case, we have made a helper function, `create_lower_func(...)`, which does just this: given a function name, it replaces the given operation with a `Call` to the function name provided. It additionally removes usages of the custom datatype by storing the custom datatype in an opaque `uint` of the appropriate width; in our case, a `uint16_t`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aaFggLxIUP8q"
   },
   "source": [
    "We can now re-try running the program:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2vUIx6jeUP8r"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    compiled = ex.evaluate(program)\n",
    "    z_output_posit = compiled(x_input, y_input)\n",
    "    print(\"z: {}\".format(z_output_posit))\n",
    "    \n",
    "except tvm.TVMError as e:\n",
    "    # Print last line of error\n",
    "    print(str(e).split('\\n')[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "stqafBUoUP8u"
   },
   "source": [
    "This new error tells us that the `Add` lowering function is not found, which is good news, as it's no longer complaining about the `Cast`! We know what to do from here: we just need to register the lowering functions for the other operations in our program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YKXYgL1WUP8u"
   },
   "outputs": [],
   "source": [
    "tvm.datatype.register_op(tvm.datatype.create_lower_func(\"posit16_add\"),\n",
    "                         \"Add\", \"llvm\", \"posit\")\n",
    "tvm.datatype.register_op(tvm.datatype.create_lower_func(\"posit16_tof\"),\n",
    "                         \"Cast\", \"llvm\", \"float\", \"posit\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pQxwXrvLUP8x"
   },
   "source": [
    "Now, we can run our program without errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2RFjZ0vAUP8z"
   },
   "outputs": [],
   "source": [
    "compiled = ex.evaluate(program)\n",
    "z_output_posit = compiled(x_input, y_input)\n",
    "print(\"z: {}\".format(z_output_posit))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3_FpH-N6UP89"
   },
   "outputs": [],
   "source": [
    "print(\"x:\\t\\t{}\".format(x_input))\n",
    "print(\"y:\\t\\t{}\".format(y_input))\n",
    "print(\"z (float32):\\t{}\".format(z_output))\n",
    "print(\"z (posit16):\\t{}\".format(z_output_posit))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nqrIBhGwUP9A"
   },
   "source": [
    "Perhaps as expected, the `posit16` results are very close to the `float` results, but with some loss in precision!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_ACtzZp6UP9A"
   },
   "source": [
    "## Running Models With Custom Datatypes\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VnV17w8tUP9B"
   },
   "source": [
    "\n",
    "We will first choose the model which we would like to run with posits. In this case we use [Mobilenet](https://arxiv.org/abs/1704.04861). We choose Mobilenet due to its small size. In this alpha state of the Bring Your Own Datatypes framework, we have not implemented any software optimizations for running software emulations of custom datatypes; the result is poor performance due to many calls into our datatype emulation library.\n",
    "\n",
    "Relay has packaged up many models within its [python/tvm/relay/testing](https://github.com/dmlc/tvm/tree/master/python/tvm/relay/testing) directory. We will go ahead and grab Mobilenet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "qocTy2P_UP9C"
   },
   "outputs": [],
   "source": [
    "from tvm.relay.testing.mobilenet import get_workload as get_mobilenet\n",
    "\n",
    "module, params = get_mobilenet(image_shape=(3, 32, 32), num_classes=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "a4m2cZ_sUP9E"
   },
   "source": [
    "It's easy to execute MobileNet with native TVM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VbJ1QKrYUP9F"
   },
   "outputs": [],
   "source": [
    "ex = tvm.relay.create_executor(\"graph\", mod=module)\n",
    "input = np.random.rand(1, 3, 32, 32).astype(\"float32\")\n",
    "result = ex.evaluate()(input, **params)\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "v9WLoUbyUP9H"
   },
   "source": [
    "Now, we would like to change the model to use posits internally. To do so, we need to convert the network. To do this, we first define a function which will help us convert tensors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "agghCRhZUP9L"
   },
   "outputs": [],
   "source": [
    "def convert_ndarray(dst_dtype, array):\n",
    "    \"\"\"Converts an NDArray into the specified datatype\"\"\"\n",
    "    ex = relay.create_executor('graph')\n",
    "    x = relay.var('x', shape=array.shape, dtype=str(array.dtype))\n",
    "    cast = relay.Function([x], x.astype(dst_dtype))\n",
    "    return ex.evaluate(cast)(array)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "40DeWdNMUP9O"
   },
   "source": [
    "Now, to actually convert the entire network, we have written [a pass in Relay](https://github.com/gussmith23/tvm/blob/ea174c01c54a2529e19ca71e125f5884e728da6e/python/tvm/relay/frontend/change_datatype.py#L21) which simply converts all nodes within the model to use the new datatype."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EsTvKV10UP9P"
   },
   "outputs": [],
   "source": [
    "from tvm.relay.frontend.change_datatype import ChangeDatatype\n",
    "\n",
    "src_dtype = 'float32'\n",
    "dst_dtype = 'custom[posit]16'\n",
    "\n",
    "# Currently, custom datatypes only work if you run simplify_inference beforehand\n",
    "module = tvm.relay.transform.SimplifyInference()(module)\n",
    "\n",
    "# Run type inference before changing datatype\n",
    "module = tvm.relay.transform.InferType()(module)\n",
    "\n",
    "# Change datatype from float to posit and re-infer types\n",
    "cdtype = ChangeDatatype(src_dtype, dst_dtype)\n",
    "expr = cdtype.visit(module['main'])\n",
    "module = tvm.relay.transform.InferType()(tvm.relay.Module.from_expr(expr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also convert the parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params = dict(\n",
    "  (p, convert_ndarray(dst_dtype, params[p])) for p in params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rgQmep2ZUP9Y"
   },
   "source": [
    "We also need to convert our input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BU5S8jwXUP9Z"
   },
   "outputs": [],
   "source": [
    "input = convert_ndarray(dst_dtype, input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ilSoA5S3UP9b"
   },
   "source": [
    "Finally, we can try to run the converted model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Zj2qLnDcUP9c"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    # Vectorization is not implemented with custom datatypes.\n",
    "    with tvm.build_config(disable_vectorize=True):\n",
    "        result_posit = ex.evaluate(expr)(input, **params)\n",
    "\n",
    "except tvm.TVMError as e:\n",
    "    print(str(e).split('\\n')[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6ORDcLtJUP9R"
   },
   "source": [
    "When we attempt to run the model, we get a familiar error telling us that more funcions need to be registerd for posits. \n",
    "\n",
    "Because this is a neural network, many more operations are required.\n",
    "Here, we register all the needed functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "F4O4AKeHUP9T"
   },
   "outputs": [],
   "source": [
    "tvm.datatype.register_op(\n",
    "    tvm.datatype.create_lower_func(\"posit16_fromf\"), \"FloatImm\",\n",
    "    \"llvm\", \"posit\")\n",
    "tvm.datatype.register_op(tvm.datatype.lower_ite,\n",
    "                         \"Call\",\n",
    "                         \"llvm\",\n",
    "                         \"posit\",\n",
    "                         intrinsic_name=\"tvm_if_then_else\")\n",
    "tvm.datatype.register_op(tvm.datatype.create_lower_func(\"posit16_mul\"),\n",
    "                         \"Mul\", \"llvm\", \"posit\")\n",
    "tvm.datatype.register_op(tvm.datatype.create_lower_func(\"posit16_div\"),\n",
    "                         \"Div\", \"llvm\", \"posit\")\n",
    "tvm.datatype.register_op(tvm.datatype.create_lower_func(\"posit16_sqrt\"),\n",
    "                         \"Call\",\n",
    "                         \"llvm\",\n",
    "                         \"posit\",\n",
    "                         intrinsic_name=\"sqrt\")\n",
    "tvm.datatype.register_op(tvm.datatype.create_lower_func(\"posit16_sub\"),\n",
    "                         \"Sub\", \"llvm\", \"posit\")\n",
    "tvm.datatype.register_min_func(lambda _: -268435456, \"posit\")\n",
    "tvm.datatype.register_op(tvm.datatype.create_lower_func(\"posit16_exp\"),\n",
    "                         \"Call\",\n",
    "                         \"llvm\",\n",
    "                         \"posit\",\n",
    "                         intrinsic_name=\"exp\")\n",
    "\n",
    "libposit_wrapper_source = \"\"\"\n",
    "#include \"posit.h\"\n",
    "\n",
    "extern \"C\" posit16_t posit16_max(posit16_t a, posit16_t b) {\n",
    "  return posit16_cmp(a, b) == 1 ? a : b;\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "! echo '{libposit_wrapper_source}' > libposit-wrapper.cc\n",
    "! g++ -Ilibposit/generated -lmpfr -lgmp --std=c++14 -shared -o libposit-wrapper-2.so -fPIC libposit-wrapper.cc libposit/libposit.a\n",
    "\n",
    "lib = CDLL('/root/libposit-wrapper-2.so', RTLD_GLOBAL)\n",
    "\n",
    "tvm.datatype.register_op(tvm.datatype.create_lower_func(\"posit16_max\"),\n",
    "                         \"Max\", \"llvm\", \"posit\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that, to implement the `Max` operator, we needed to rewrite our wrapper library with a new function, `posit16_max`.\n",
    "This is the only function we needed to implement by hand! All other functions we took straight from the posit library.\n",
    "\n",
    "\n",
    "Now we can finally run the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Vectorization is not implemented with custom datatypes.\n",
    "with tvm.build_config(disable_vectorize=True):\n",
    "    result_posit = ex.evaluate(expr)(input, **params)\n",
    "    result_posit = convert_ndarray(src_dtype, result_posit)\n",
    "    print(result_posit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, note that the output using 16-bit posits is understandably different from that of 32-bit floats, but is still within a sane distance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(result.asnumpy(), result_posit.asnumpy(), rtol=1e-6, atol=1e-5)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "08_TVM_Tutorial_BringYourOwnDatatypes.ipynb",
   "provenance": []
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}