charpter7.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "81ee1b57",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.ir.module import IRModule\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T\n",
    "from tvm import relax, topi\n",
    "import numpy as np\n",
    "import tvm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f8db49d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.script.ir_module\n",
    "class MyModule:\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((3, 4), \"float32\"),\n",
    "             y: R.Tensor((3, 4), \"float32\")):\n",
    "        with R.dataflow():\n",
    "            lv0 = relax.op.multiply(x, y)\n",
    "            gv0 = relax.op.add(lv0, y)\n",
    "            R.output(gv0)\n",
    "        return gv0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "685d56b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[relax.expr.Var(0x56260ce6e0b0), relax.expr.Var(0x56260d2044a0)]\n",
      "[relax.expr.VarBinding(0x56260cea3710), relax.expr.VarBinding(0x56260d2541e0)]\n",
      "lv0\n",
      "R.multiply(x, y)\n"
     ]
    }
   ],
   "source": [
    "relax_func = MyModule[\"main\"]\n",
    "print(relax_func.params)\n",
    "\n",
    "func_body = relax_func.body\n",
    "dataflow_block = func_body.blocks[0]\n",
    "print(dataflow_block.bindings)\n",
    "\n",
    "binding = dataflow_block.bindings[0]\n",
    "print(binding.var)\n",
    "print(binding.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8225101b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "@R.function\n",
      "def main(x: R.Tensor((3, 4), dtype=\"float32\"), y: R.Tensor((3, 4), dtype=\"float32\")) -> R.Tensor((3, 4), dtype=\"float32\"):\n",
      "    # block 0\n",
      "    with R.dataflow():\n",
      "        gv0: R.Tensor((3, 4), dtype=\"float32\") = R.ewise_fma(x, y, y)\n",
      "        R.output(gv0)\n",
      "    return gv0\n",
      "    \n"
     ]
    }
   ],
   "source": [
    "@relax.expr_functor.mutator\n",
    "class EwiseFMARewriter(relax.PyExprMutator):\n",
    "    def visit_call_(self, call):\n",
    "        call = self.visit_expr_post_order(call)\n",
    "        add_op = tvm.ir.Op.get(\"relax.add\")\n",
    "        multiply_op = tvm.ir.Op.get(\"relax.multiply\")\n",
    "        ewise_fma_op = tvm.ir.Op.get(\"relax.ewise_fma\")\n",
    "    \n",
    "        if call.op != add_op:\n",
    "            return call\n",
    "        \n",
    "        value = self.lookup_binding(call.args[0])\n",
    "        if not isinstance(value, relax.Call) or value.op != multiply_op:\n",
    "            return call\n",
    "        \n",
    "        fma_call = relax.Call(ewise_fma_op, [value.args[0], value.args[1], call.args[1]], None, None)\n",
    "        return fma_call\n",
    "    \n",
    "updated_fn = EwiseFMARewriter().visit_expr(MyModule[\"main\"])\n",
    "updated_fn = relax.analysis.remove_all_unused(updated_fn)\n",
    "print(updated_fn.script())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ae7a9c6f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import pickle as pkl\n",
    "import torchvision\n",
    "import torch\n",
    "\n",
    "params = pkl.load(open(\"fasionmnist_mlp_params.pkl\", \"rb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7ec6b2cd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/workspace/relax/python/tvm/script/highlight.py:116: UserWarning: No module named 'black'\n",
      "To print formatted TVM script, please install the formatter 'Black':\n",
      "/opt/conda/bin/python3.8 -m pip install \"black==22.3.0\" --upgrade --user\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        <span style=\"color: #007979; font-style: italic\"># block 0</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(x, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">0</span>], out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;&quot;</span>)\n",
       "            lv1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">1</span>])\n",
       "            lv2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(lv1)\n",
       "            lv3: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(lv2, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">2</span>], out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;&quot;</span>)\n",
       "            lv4: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv3, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">3</span>])\n",
       "            gv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> lv4\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "        \n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def create_model():\n",
    "    bb = relax.BlockBuilder()\n",
    "    x = relax.Var(\"x\", relax.TensorStructInfo((1, 784), \"float32\"))\n",
    "    w0 = relax.const(params['w0'].T, \"float32\")\n",
    "    b0 = relax.const(params['b0'], \"float32\")\n",
    "    w1 = relax.const(params['w1'].T, \"float32\")\n",
    "    b1 = relax.const(params['b1'], \"float32\")\n",
    "    \n",
    "    with bb.function(\"main\", [x]):\n",
    "        with bb.dataflow():\n",
    "            lv0 = bb.emit(relax.op.matmul(x, w0))\n",
    "            lv1 = bb.emit(relax.op.add(lv0, b0))\n",
    "            lv2 = bb.emit(relax.op.nn.relu(lv1))\n",
    "            lv3 = bb.emit(relax.op.matmul(lv2, w1))\n",
    "            lv4 = bb.emit(relax.op.add(lv3, b1))\n",
    "            gv = bb.emit_output(lv4)\n",
    "        bb.emit_func_output(gv)\n",
    "        \n",
    "    return bb.get()\n",
    "\n",
    "MLPModel = create_model()\n",
    "MLPModel.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "88048a11",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/workspace/relax/python/tvm/script/highlight.py:116: UserWarning: No module named 'black'\n",
      "To print formatted TVM script, please install the formatter 'Black':\n",
      "/opt/conda/bin/python3.8 -m pip install \"black==22.3.0\" --upgrade --user\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">fused_dense_add0</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), w: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), b: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;Primitive&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># block 0</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(x, w, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;&quot;</span>)\n",
       "            lv1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv, b)\n",
       "            gv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> lv1\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "        \n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">fused_dense_add1</span>(x1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), w1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), b1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;Primitive&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># block 0</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>matmul(x1, w1, out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;&quot;</span>)\n",
       "            lv11: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>add(lv2, b1)\n",
       "            gv1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> lv11\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv1)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv1\n",
       "        \n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(x2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        <span style=\"color: #007979; font-style: italic\"># block 0</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv12: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> fused_dense_add0(x2, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">0</span>], metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">1</span>])\n",
       "            lv21: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(lv12)\n",
       "            lv4: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> fused_dense_add1(lv21, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">2</span>], metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">3</span>])\n",
       "            gv2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> lv4\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv2)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv2\n",
       "        \n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@relax.expr_functor.mutator\n",
    "class DenseAddFusor(relax.PyExprMutator):\n",
    "    def __init__(self, mod: IRModule) -> None:\n",
    "        super().__init__()\n",
    "        \n",
    "        self.mod = mod\n",
    "        self.add_op = tvm.ir.Op.get(\"relax.add\")\n",
    "        self.dense_op = tvm.ir.Op.get(\"relax.matmul\")\n",
    "        self.counter = 0\n",
    "        \n",
    "    def transform(self) -> IRModule:\n",
    "        for global_var, func in self.mod.functions.items():\n",
    "            if not isinstance(func, relax.Function):\n",
    "                continue\n",
    "            \n",
    "            if func.attrs and \"Primitive\" in func.attrs.keys() and func.attrs[\"Primitive\"] != 0:\n",
    "                continue\n",
    "            \n",
    "            updated_func = self.visit_expr(func)\n",
    "            updated_func = relax.analysis.remove_all_unused(updated_func)\n",
    "            self.builder_.update_func(global_var, updated_func)\n",
    "        return self.builder_.get()\n",
    "    \n",
    "    def visit_call_(self, call):\n",
    "        call = self.visit_expr_post_order(call)\n",
    "        \n",
    "        def match_call(node, op):\n",
    "            if not isinstance(node, relax.Call):\n",
    "                return False\n",
    "            else:\n",
    "                return node.op == op\n",
    "            \n",
    "        if not match_call(call, self.add_op):\n",
    "            return call\n",
    "        \n",
    "        value = self.lookup_binding(call.args[0])\n",
    "        if value is None:\n",
    "            return call\n",
    "        \n",
    "        if not match_call(value, self.dense_op):\n",
    "            return call\n",
    "        \n",
    "        x = value.args[0]\n",
    "        w = value.args[1]\n",
    "        b = call.args[1]\n",
    "        \n",
    "        param_x = relax.Var(\"x\", x.struct_info)\n",
    "        param_w = relax.Var(\"w\", w.struct_info)\n",
    "        param_b = relax.Var(\"b\", b.struct_info)\n",
    "        \n",
    "        bb = relax.BlockBuilder()\n",
    "        fn_name = f'fused_dense_add{self.counter}'\n",
    "        self.counter += 1\n",
    "        with bb.function(fn_name, [param_x, param_w, param_b]):\n",
    "            with bb.dataflow():\n",
    "                lv0 = bb.emit(relax.op.matmul(param_x, param_w))\n",
    "                lv1 = bb.emit(relax.op.add(lv0, param_b))\n",
    "                gv = bb.emit_output(lv1)\n",
    "            bb.emit_func_output(gv)\n",
    "            \n",
    "        fused_fn = bb.get()[fn_name].with_attr(\"Primitive\", 1)\n",
    "        global_var = self.builder_.add_func(fused_fn, fn_name)\n",
    "        \n",
    "        return relax.Call(global_var, [x, w, b], None, None)\n",
    "    \n",
    "@tvm.ir.transform.module_pass(opt_level=2, name=\"DenseAddFuse\")\n",
    "class FuseDenseAddPass:\n",
    "    def transform_module(self, mod, ctx):\n",
    "        return DenseAddFusor(mod).transform()\n",
    "    \n",
    "MLPFused = FuseDenseAddPass()(MLPModel)\n",
    "MLPFused.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "461eb43a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/workspace/relax/python/tvm/script/highlight.py:116: UserWarning: No module named 'black'\n",
      "To print formatted TVM script, please install the formatter 'Black':\n",
      "/opt/conda/bin/python3.8 -m pip install \"black==22.3.0\" --upgrade --user\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">matmul</span>(rxplaceholder: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], rxplaceholder_1: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], T_matmul: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>]):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># body</span>\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_matmul&quot;</span>):\n",
       "                v_ax0, v_ax1, v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [ax0, ax1, k])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(rxplaceholder[v_ax0, v_k], rxplaceholder_1[v_k, v_ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_matmul[v_ax0, v_ax1])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
       "                    T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>)\n",
       "                T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> rxplaceholder[v_ax0, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> rxplaceholder_1[v_k, v_ax1]\n",
       "    \n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">add</span>(rxplaceholder: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], rxplaceholder_1: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], T_add: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>]):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># body</span>\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(rxplaceholder[v_ax0, v_ax1], rxplaceholder_1[v_ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_add[v_ax0, v_ax1])\n",
       "                T_add[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> rxplaceholder[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> rxplaceholder_1[v_ax1]\n",
       "    \n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">fused_dense_add0</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), w: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">784</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), b: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;Primitive&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># block 0</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(matmul, (x, w), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(add, (lv, b), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            gv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> lv1\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "        \n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">fused_dense_add1</span>(x1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), w1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">128</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), b1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">10</span>,), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        R<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;Primitive&quot;</span>: <span style=\"color: #008000\">1</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># block 0</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(matmul1, (x1, w1), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv11 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(add1, (lv2, b1), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            gv1: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> lv11\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv1)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv1\n",
       "        \n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">add1</span>(rxplaceholder: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], rxplaceholder_1: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], T_add: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>]):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># body</span>\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(rxplaceholder[v_ax0, v_ax1], rxplaceholder_1[v_ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_add[v_ax0, v_ax1])\n",
       "                T_add[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> rxplaceholder[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> rxplaceholder_1[v_ax1]\n",
       "    \n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">relu</span>(rxplaceholder: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], compute: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>]):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># body</span>\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;compute&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(rxplaceholder[v_i0, v_i1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(compute[v_i0, v_i1])\n",
       "                compute[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(rxplaceholder[v_i0, v_i1], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>))\n",
       "    \n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">matmul1</span>(rxplaceholder: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], rxplaceholder_1: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], T_matmul: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>]):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># body</span>\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_matmul&quot;</span>):\n",
       "                v_ax0, v_ax1, v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [ax0, ax1, k])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(rxplaceholder[v_ax0, v_k], rxplaceholder_1[v_k, v_ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_matmul[v_ax0, v_ax1])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
       "                    T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>)\n",
       "                T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> rxplaceholder[v_ax0, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> rxplaceholder_1[v_k, v_ax1]\n",
       "    \n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(x2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        <span style=\"color: #007979; font-style: italic\"># block 0</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv12: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> fused_dense_add0(x2, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">0</span>], metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">1</span>])\n",
       "            lv21 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(relu, (lv12,), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv4: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> fused_dense_add1(lv21, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">2</span>], metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">3</span>])\n",
       "            gv2: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> lv4\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv2)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv2\n",
       "        \n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@relax.expr_functor.mutator\n",
    "class LowerToTensorIR(relax.PyExprMutator):\n",
    "    def __init__(self, mod: IRModule, op_map) -> None:\n",
    "        super().__init__()\n",
    "        self.mod = mod\n",
    "        self.op_map = {\n",
    "            tvm.ir.Op.get(k): v for k, v in op_map.items()\n",
    "        }\n",
    "        \n",
    "    def visit_call_(self, call):\n",
    "        call = self.visit_expr_post_order(call)\n",
    "        \n",
    "        if call.op in self.op_map:\n",
    "            return self.op_map[call.op](self.builder_, call)\n",
    "        \n",
    "        return call\n",
    "    \n",
    "    def transform(self) -> IRModule:\n",
    "        for global_var, func in self.mod.functions.items():\n",
    "            if not isinstance(func, relax.Function):\n",
    "                continue\n",
    "            updated_func = self.visit_expr(func)\n",
    "            self.builder_.update_func(global_var, updated_func)\n",
    "            \n",
    "        return self.builder_.get()\n",
    "    \n",
    "def map_dense(bb, call):\n",
    "    x, w = call.args\n",
    "    return bb.call_te(topi.matmul, x, w)\n",
    "\n",
    "def map_add(bb, call):\n",
    "    x, b = call.args\n",
    "    return bb.call_te(topi.add, x, b)\n",
    "\n",
    "def map_relu(bb, call):\n",
    "    x, = call.args\n",
    "    return bb.call_te(topi.nn.relu, x)\n",
    "\n",
    "op_map = {\n",
    "    \"relax.matmul\": map_dense,\n",
    "    \"relax.add\": map_add,\n",
    "    \"relax.nn.relu\": map_relu\n",
    "}\n",
    "\n",
    "@tvm.ir.transform.module_pass(opt_level=0, name=\"LowerToTensorIR\")\n",
    "class LowerToTensorIRPass:\n",
    "    def transform_module(self, mod, ctx):\n",
    "        return LowerToTensorIR(mod, op_map).transform()\n",
    "    \n",
    "MLPModelTIR = LowerToTensorIRPass()(MLPFused)\n",
    "MLPModelTIR.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e057cc34",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/workspace/relax/python/tvm/script/highlight.py:116: UserWarning: No module named 'black'\n",
      "To print formatted TVM script, please install the formatter 'Black':\n",
      "/opt/conda/bin/python3.8 -m pip install \"black==22.3.0\" --upgrade --user\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #AA22FF\">@tvm</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>script<span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">fused_dense_add0</span>(x: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], w: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], b: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], T_add: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>]):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># body</span>\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;)</span>\n",
       "        T_matmul <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer([T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)], dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">784</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_matmul&quot;</span>):\n",
       "                v_ax0, v_ax1, v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [ax0, ax1, k])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(x[v_ax0, v_k], w[v_k, v_ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_matmul[v_ax0, v_ax1])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
       "                    T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>)\n",
       "                T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> x[v_ax0, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> w[v_k, v_ax1]\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(T_matmul[v_ax0, v_ax1], b[v_ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_add[v_ax0, v_ax1])\n",
       "                T_add[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> b[v_ax1]\n",
       "    \n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">relu</span>(rxplaceholder: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], compute: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>]):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># body</span>\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;)</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> i0, i1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;compute&quot;</span>):\n",
       "                v_i0, v_i1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [i0, i1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(rxplaceholder[v_i0, v_i1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(compute[v_i0, v_i1])\n",
       "                compute[v_i0, v_i1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>max(rxplaceholder[v_i0, v_i1], T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>))\n",
       "    \n",
       "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">fused_dense_add1</span>(x: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], w: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], b: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), <span style=\"color: #BA2121\">&quot;float32&quot;</span>], T_add: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer[(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)), <span style=\"color: #BA2121\">&quot;float32&quot;</span>]):\n",
       "        <span style=\"color: #007979; font-style: italic\"># function attr dict</span>\n",
       "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>})\n",
       "        <span style=\"color: #007979; font-style: italic\"># body</span>\n",
       "        <span style=\"color: #007979; font-style: italic\"># with T.block(&quot;root&quot;)</span>\n",
       "        T_matmul <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>alloc_buffer([T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)], dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1, k <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">128</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_matmul&quot;</span>):\n",
       "                v_ax0, v_ax1, v_k <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SSR&quot;</span>, [ax0, ax1, k])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(x[v_ax0, v_k], w[v_k, v_ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_matmul[v_ax0, v_ax1])\n",
       "                <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>init():\n",
       "                    T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">0</span>)\n",
       "                T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> x[v_ax0, v_k] <span style=\"color: #AA22FF; font-weight: bold\">*</span> w[v_k, v_ax1]\n",
       "        <span style=\"color: #008000; font-weight: bold\">for</span> ax0, ax1 <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">1</span>), T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int64(<span style=\"color: #008000\">10</span>)):\n",
       "            <span style=\"color: #008000; font-weight: bold\">with</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>block(<span style=\"color: #BA2121\">&quot;T_add&quot;</span>):\n",
       "                v_ax0, v_ax1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>axis<span style=\"color: #AA22FF; font-weight: bold\">.</span>remap(<span style=\"color: #BA2121\">&quot;SS&quot;</span>, [ax0, ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>reads(T_matmul[v_ax0, v_ax1], b[v_ax1])\n",
       "                T<span style=\"color: #AA22FF; font-weight: bold\">.</span>writes(T_add[v_ax0, v_ax1])\n",
       "                T_add[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> T_matmul[v_ax0, v_ax1] <span style=\"color: #AA22FF; font-weight: bold\">+</span> b[v_ax1]\n",
       "    \n",
       "    <span style=\"color: #AA22FF\">@R</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(x: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">784</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        <span style=\"color: #007979; font-style: italic\"># block 0</span>\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>dataflow():\n",
       "            lv1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(fused_dense_add0, (x, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">0</span>], metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">1</span>]), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(relu, (lv1,), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">128</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            lv4 <span style=\"color: #AA22FF; font-weight: bold\">=</span> R<span style=\"color: #AA22FF; font-weight: bold\">.</span>call_tir(fused_dense_add1, (lv2, metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">2</span>], metadata[<span style=\"color: #BA2121\">&quot;relax.expr.Constant&quot;</span>][<span style=\"color: #008000\">3</span>]), out_sinfo<span style=\"color: #AA22FF; font-weight: bold\">=</span>R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>))\n",
       "            gv: R<span style=\"color: #AA22FF; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">10</span>), dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">=</span> lv4\n",
       "            R<span style=\"color: #AA22FF; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "        \n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "MLPModelFinal = relax.transform.FuseTIR()(MLPModelTIR)\n",
    "MLPModelFinal.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f19f25f4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAS4AAAD4CAYAAABSUAvFAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAAcTklEQVR4nO3df5RU9Znn8ffTTTctP0R+CKKgokJGohlMGNTRjGYwCSa7kpzNqsxJ1mTcgdkJ2biTyYnrzlHXOTtjZpM42R3HFSOJk010HZNMmB0S4hgnrsY4oBIVUGgRhZafggLyo7urnv2jCq3+cZ9bdFX3rQuf1zl1qKrn/vh2dfXDvd/73O/X3B0RkTxpyroBIiJHS4lLRHJHiUtEckeJS0RyR4lLRHJn2FDurNWGexsjh3KXIseVQ7xNpx+2Wrbx0Q+N9Dd2F6pa9unnDq9w93m17G8gakpcZjYP+CbQDHzL3W+Plm9jJBfa3Fp2KSKBp/yRmrexa3eBp1ZMqWrZlskvT6h5hwMw4MRlZs3AncCHgS3ASjNb5u5r69U4EcmCU/Bi1o0I1dLHNQdod/eN7t4JPADMr0+zRCQrDhTxqh5ZqeVU8TRgc8XrLcCFvRcys4XAQoA2RtSwOxEZKkUa+4hr0Dvn3X0JsATgRBun+4tEGpzjdDX4qWItiasDmFrxekr5PRHJMQcKGZ4GVqOWPq6VwHQzm2ZmrcC1wLL6NEtEsnTM9nG5e7eZLQZWUCqHWOrua+rWMnnHsDOmxgt0diWGurduq3NremoePy6M++SJibHu8SeE6x4a3xrGD46N/989PC65nKmpO1yVA5PjP8rp394VxgvrNoRxG5b8p+fdKY0bZA4UGnzUmJr6uNx9ObC8Tm0RkQbR2D1cQ1w5LyKNz/GG7+NS4hKRHtyhq7HzlhKXiPRmFKjpdsdBp8QlIj04UNQRl4jkjY64RCRXSgWoSlzHvqbmOF5MGdtozvlhuOPC0WH81EeSa4q2XX1WuG5TcgkYAMP3xhfGu0+Iv+Cte5PPOYa/FX8uI14/GMZHP7o5jBPctrLvd38jXHXvufHPXWyLa8zSeAOfiznQ5Y09xqgSl4j04BiFBh8cWYlLRPoouk4VRSRH1MclIjlkFNTHJSJ5UhoBVYlLRHLE3ej0lCvlGVPiqocaR4sctnlnGD/08bgcojhyeGJs7/md4bqTTn0zjI/5WHsYz9S508PwlitPToyN6oh/ZzMWrgzjNRczpJXIZKyoPi4RyZNS57xOFUUkV9Q5LyI5o855EcmlggpQRSRPHKPLGzs1NHbrRGTIqXNeRHLHMZ0qHhdqnMopbQqxaV/dG8Zt9KjE2Pgnzw7X3TVzQhgfYy+H8Vp+9u65Hwjjr8yPv54nbIuLJKfc/lRysMY6qrc+fVEYH/f0G2E8bfqyrKlzXkRyxR2VQ4hIvpQ653XLj4jkjDrnRSRXHNNAgiKSPzriEpFcKc2rqMQlIrmimaylCjYs/jUUDxyINxDEmwpxHVfbtH3xtlM0nzMtjK/7cnKd2Kj2lnDd6f/xlwNqUzUOfPLCML79mnhqtA+d9esw/ujGeKywadeG4UyVpidr7KuKNR0PmtkmM3vezFab2ap6NUpEsuNuFL2pqkc1zGyemb1kZu1mdmM/8dPN7FEze9bMnjOzj6Vtsx5HXB9y9+QZSUUkd+pVgGpmzcCdwIeBLcBKM1vm7msrFvtT4EF3v8vMZgLLgTOj7TZ2D5yIDLnSeFxW1aMKc4B2d9/o7p3AA8D8fnZ5Yvn5GOD1tI3WesTlwM/MzIG73X1J7wXMbCGwEKCNETXuTkQG31GNgDqhVzfRkl554DRgc8XrLUDvDsZbKeWRLwAjgSvSdlpr4rrU3TvMbCLwsJm96O6PVS5Q/iGWAJxo42qeY0BEBlepHKLqq4q73H12jbtcAHzH3b9uZhcD3zWz89yTZ6GpKXG5e0f53x1m9iNKh4WPxWuJSCOr872KHcDUitdTyu9Vuh6YB+DuT5pZGzAB2JG00QH3cZnZSDMbfeQ58BHghYFuT0QaR5Gmqh5VWAlMN7NpZtYKXAss67XMa8BcADM7F2gDwjn7ajnimgT8yMyObOf77v7TGrZ33PJibWfQUR3Yvo/vD9c9ZXQc3/9v43qnPTPiL++MRQOvxWo+8cQwvv6WmWF82qze/7G/67Lxcbse3xHXv63cdnoY79rdFsZpCo5o0sYKs+A0rg6dMaVhbepTgOru3Wa2GFgBNANL3X2Nmd0GrHL3ZcCXgHvM7D9R+gk+6x4P9DbgxOXuG4HfHOj6ItK46nmTtbsvp1TiUPnezRXP1wKXHM02VTkvIj2URodo7EopJS4R6aF0y48Sl4jkio64RCSHqqyKz4wSl4j0UM+rioNFiasBNLXGw7sUD8WXxztumJMYu+OCe8N1/2jFZ8P4mInxKcPUPw+mAAN2/uHFiTG/ck+47mVT2sP4ufxLGN9fGJ4Y+4dvfzBc9/D4uK7gzEtfC+MHT24N41u+klxmMuUv4lINaw5KKep0b4pOFUUkVzTmvIjkjgPdOuISkbzRqaKI5IvrVFFEcubIQIKNTIlLRPrQEZeI5MpRDiSYCSWuBlDs7Kpp/TFXbEuMPbI3Hvpl+nfjabiu/vbPwvinbnwljLd3PZm87ScWhev+6o54YM2T/u7ZMO6Hk3+2U6ht6rNrX4yHRX/Q4ra/yugB79u7u4PggDdbsQmju6jOeRHJGfVxiUi+uE4VRSRn1MclIrmkxCUiueIYBXXOi0jeqHNeRHLF1TkvQDwVFaROR1W87IIwfs3UnyfGvrspeawugLG/ei6MHygmj2kF8Hu/syCMd2/clBg7h7gOK81gTovuF8cTWD2zPx5DbX9n/LkdPDVlCrKMuRKXiOSLbrIWkRzSEZeI5Io7FIpKXCKSM7qqKCK54uhUUURyR53zIpJDPpi1JnWgxDUEap030brjb9EZrTsTYwcOx/P7jQ2j8MPXZ4Xx1qBOK2t+yazE2IZ/F/9Ornz/82G8fd/JYfxgV7z937/sF4mxX1x6Ubhu0+Orw3g9NPqpYuoNSWa21Mx2mNkLFe+NM7OHzWxD+d+077+I5ETpqmJTVY+sVLPn7wDzer13I/CIu08HHim/FpFjhHt1j6ykJi53fwzY3evt+cB95ef3AZ+ob7NEJEvuVtUjKwPt45rk7lvLz7cBk5IWNLOFwEKANkYMcHciMlScbJNSNWo+SXV3J7jf1d2XuPtsd5/dQnzjqYg0Bq/ykZWBJq7tZjYZoPzvjvo1SUQy5eBFq+pRDTObZ2YvmVm7mfXbH25mV5vZWjNbY2bfT9vmQBPXMuC68vPrgB8PcDsi0oDq1cdlZs3AncCVwExggZnN7LXMdOA/A5e4+3uBG9K2m9rHZWb3A5cDE8xsC3ALcDvwoJldD7wKXJ36ExzHiocP17S+PbE6jO8ujEqMda09Md72sPgrsPXJU8P4GbwaxiNp44y9clXctfDn//r+MP7BtscTYy91xZ/LX7/+u2F859vJnznAgcNxHddPX0+e73LrojghnJP8Y9VNHa8YzgHa3X0jgJk9QOni3tqKZf4AuNPd95T27alncKmJy92TRoqbm7auiOTPUd6rOMHMVlW8XuLuSypenwZsrni9Bbiw1zZmAJjZE0AzcKu7/zTaqSrnRaQnB6pPXLvcPZ62O90wYDqlM7spwGNmdr67v5m0QmNP5SEimahjAWoHMLXi9ZTye5W2AMvcvcvdXwHWU0pkiZS4RKSX6q4oVnlVcSUw3cymmVkrcC2li3uV/p7S0RZmNoHSqePGaKNKXCLSV50Kudy9G1gMrADWAQ+6+xozu83MriovtgJ4w8zWAo8CX3b3N6Ltqo9LRHry+o4O4e7LgeW93ru54rkDf1x+VEWJ6xhwYdumxFjnaZ3xyue/JwwX2uLVu+d+IIzv+sKBxJhzKFx32uhdYfwrP78mjJ9zf1dirOX1t8J1CxvCMxWG/WRMGO/qij+4js3jE2NXvG9tYgxgy3m/kRiz9ifCdaum8bhEJH8a+15FJS4R6auYdQNiSlwi0tPR1XFlQolLRPrQmPMikj9KXCKSOzpVFJG8MR1xyWB3GHzq7j9JjL3np3G90uZ5cT3S+RdtCOPPnnRmGLe3TkgOpvyv3tUVfz1POSMsrqZ9QXKt1IhXTwnXnfjMuDBeKKbUgXU3h/GTJu5LjHUciH8nb513UvJ+O+L9VsUNqhwkMCtKXCLSl464RCR3lLhEJHeUuEQkV1SAKiJ5pKuKIpI/Slwikjc64pJBN+UvfpkYS/v+Hf7yrDC++pWpYXxUezwN1wmXvZkY2/P8hHDdaXdsDuO7rjw7jPPb3YmhWz73vXDVMxfGY4Hdvf1DYfzllvhnO3F48lhkI4bFY6jtHJHc/+T1GtNYfVwikitVDsucJSUuEelLiUtE8sY0kKCI5I6OuEQkT8x1VVFE8khXFUUkd3TEJVjK/16WUnxTLAx417sWXRzGC/uSa50AZixaOeB9A7yxJ3n/n/vCz8N1H7/3vWF87HeeTIknx+5lWrju+ruvDuMffN9LYXxES1yL1RRkhlPa9obr7n18R2KsPeX3Wa1GP1VMLVczs6VmtsPMXqh471Yz6zCz1eXHxwa3mSIyZLx0VbGaR1aqqbP9DjCvn/fvcPdZ5cfyfuIiklde5SMjqYnL3R8Ddg9BW0SkUeQ9cQUWm9lz5VPJsUkLmdlCM1tlZqu6OFzD7kRkqBwpiUh7ZGWgiesu4GxgFrAV+HrSgu6+xN1nu/vsFoYPcHciIu8aUOJy9+3uXnD3InAPMKe+zRKRTB2Lp4pmNrni5SeBF5KWFZGcycFVxdQ6LjO7H7gcmGBmW4BbgMvNbBalnLsJWDR4TWwMNjz5NNe7Umpn0uqwaixS3n/1RYmxA5Pijddap5Vm/D3JtVZLL//tcN1/ePjOMH7DmfH6kaa2tniBlHkFf+vEV8P4Tw7FNWi7Do5MjL36ZmKXMQAT17+YGHOvUz9yg9dxpSYud1/Qz9v3DkJbRKQBGMdAAaqIHIfq2MdlZvPM7CUzazezG4Pl/o2ZuZnNTtumEpeI9FRlKUQ1R2Vm1gzcCVwJzAQWmNnMfpYbDXwReKqaJipxiUhfxSof6eYA7e6+0d07gQeA+f0s92fAV4HkwfgrKHGJSB9HccQ14UiBefmxsNemTgMqZz3ZUn7v3X2ZvR+Y6u7/WG37NDqEiPRVfef8LndP7ZNKYmZNwDeAzx7NekpcZVG5A4AfDi4zNzXXtvOUcom0oWm625Iv3Z9+W/LUZVk7+9Orw/i5HSPC+PpvxX8vI15uTYw9tOhr4bpPHdoYxv/rL68K462vx9O2Fc46mBw7EP9ZTgyjdVDf4tIOoHKOuynl944YDZwH/LOVhn86BVhmZle5+6qkjSpxiUgfdSyHWAlMN7NplBLWtcDvHQm6+1vAO5NQmtk/A38SJS1QH5eI9KdO5RDu3g0sBlYA64AH3X2Nmd1mZvFha0BHXCLSRz1v5ymP17e813s3Jyx7eTXbVOISkZ40k7WI5I1R8+2zg06JS0T60hGXiORNo99krcRVljo0TVSrVcP0YQAb/ueFYbz57Xj9s24MarXSpkbzDL+hKfuetuL6MP7B964P4ydfsD8xNv9XfxiuaxuSh50BsHFx73XXmDh++oQ3E2Mdz0xOjA0ZJS4RyRXPdpDAaihxiUhfOuISkbxRH5eI5I8Sl4jkjY64RCRfnGoHCcyMEpeI9JCHyTKUuI7wlP9iaqh32nZDPI1W655422fcnDzFVypLGQDEa6tBq8XmP40/l0/9ZvxzP7QyHo+rZXfy17vYEn/mhdO6wnhTa/y5+e7kscAAJpyQXGPW+WwDZI0GaEJEiUtE+rAsC5OroMQlIj1pdAgRySP1cYlI7uiWHxHJHx1xiUiuVDlLdZaUuESkLyWuXqLxoQbxEqwNi39U704Zjyvw2s1xPdJH5/9LGF/3gYHvO1WNY4XVav1dcxJj3/7I/wrXXbTq0/HGC/FYY10Tkj9Xa0npxOmK69+KB1P+dNri7Y8fnjzI2tsv7Y33He+5ZnkoQE2dnszMpprZo2a21szWmNkXy++PM7OHzWxD+d+xg99cERkKVvSqHlmpZl7FbuBL7j4TuAj4vJnNBG4EHnH36cAj5dciknfVzqmY4VFZauJy963u/kz5+T5KkzqeBswH7isvdh/wiUFqo4gMMStW98jKUfVxmdmZwAXAU8Akd99aDm0DJiWssxBYCNDGiAE3VESGUN77uI4ws1HAD4Ab3L1H76G7Jx44uvsSd5/t7rNbGF5TY0VkaJhX98hKVYnLzFooJa3vufsPy29vN7PJ5fhkYMfgNFFEhpRTusJfzSMjqaeKZmbAvcA6d/9GRWgZcB1we/nfH1e1x1p+2KiUImX4ltRyh5RpvLZ98eLE2KHTO8N1H970njB+xqN7wviLr50Sxs/9SkdirHvb9nDdWq2/57fC+P+em1zycN0PPh+uWxgZl3I0j4mHninsa0mMecr/2ZYybA2eMu1byuFIdzF5urumvQfCdYeia+lYuOXnEuAzwPNmtrr83k2UEtaDZnY98Cpw9aC0UESGVB7quFITl7s/Tuln6c/c+jZHRDKX8WlgNXTLj4j0kfsjLhE5DilxiUje6IhLRPLFgUJjZy4lLhHpQ0dclSweXsYLabUzwaeZMs1W8/hxYXzdV88K401tB5O3vT2+I6Br54lh/LW4zIumYfG3aMs1ZyfGTv3F+HDd4uq1YXzD/7gwjN8/92/C+GceXJwYK4yLa+taRsf1cV374inAGJ78nWhura1QyVOuujUPi7+PezpPSIwVXtsyoDbVVR2vKprZPOCbQDPwLXe/vVf8j4F/T2lAh53A77v7q9E2q77lR0SOH/W65cfMmoE7gSuBmcCC8ugylZ4FZrv7+4CHgL9M264Sl4j0VN9hbeYA7e6+0d07gQcojSzz7u7cH3X3I7cL/AqYkrZR9XGJSA8GWPWd8xPMbFXF6yXuvqTi9WnA5orXW4Co/+F64CdpO1XiEpE+jmIm613uPrsu+zT7NDAbuCxtWSUuEempvqObdgBTK15PKb/Xg5ldAfwX4DJ3P5y2UfVxiUgvVQ5pU91R2UpguplNM7NW4FpKI8u8w8wuAO4GrnL3qobH0hGXiPRRrzoud+82s8XACkrlEEvdfY2Z3QascvdlwH8HRgF/VxpFi9fc/apou0ObuLy2acBqsf6mGWG8ZVTKGEibRybGCqPjmp2mU+Mj3+amuKbId8Z1YqO2Jq/ftH13uO6mlKnVln48nkJswYr/EMYZG0wRFtRZAXQfTvl6NsV/XdEQa8XulPG0Usbbakqb3izF213B7zSjv5Ee6ljH5e7LgeW93ru54vkVR7tNHXGJSE9+VFcVM6HEJSJ9NXbeUuISkb6OohwiE0pcItKXEpeI5IozNDNy1ECJS0R6MFyniiKSQ8XGPuQa0sRVHDuSfR+9KDG+dV48T97o5wc+E3bTqW+H8WIheZ47gOLE5FqslIogLKXeqDVl7KbzL3gljL85M3lsp1dmxOOM/dE1/xjGP/fY58I4afVMwb0ZabVQbW3x96GlOWX8tsDwlrhWav+h+Lt26GA8Flha29863JYYGxOuOQR0qigieaRTRRHJHyUuEckXTQgrInmjWX5EJI/UxyUi+aPEJSK54kAx54nLzKYCfwtMovQjLXH3b5rZrcAfUJoHDeCm8rg7iZoPFRjz4luJ8a1zR4VtOfVfJU+15injJx0uxD/quLa4zuuqib9OjL3eOTZc9+UDJ4fxGSO3h/HhTXFN0FmtyYNG7jsjucYL4K82zA3jo8fG45RNGZP8+wTYG9QrpXn7cFwrdaizZcDbfnPb6DA+45ytYdxOqu0Pe1RLcl3gvpq2XA/HRud8N/Ald3/GzEYDT5vZw+XYHe7+tcFrnohkIu+Jy923AlvLz/eZ2TpKUw6JyLHIgUJjl84f1WQZZnYmcAHwVPmtxWb2nJktNbN+z5fMbKGZrTKzVZ3d8WmHiDQCBy9W98hI1YnLzEYBPwBucPe9wF3A2cAsSkdkX+9vPXdf4u6z3X1267ARtbdYRAZf/Wb5GRRVXVU0sxZKSet77v5DAHffXhG/B/i/g9JCERlaObiqmHrEZaX5gu4F1rn7Nyren1yx2CeBF+rfPBHJxDFwxHUJ8BngeTNbXX7vJmCBmc2ilJ83AYvSNuQHD1H89brE+IyFKetX0dgkI86ZFsb3T4qvN9w7KXl4mN3viYfE6Rodt/yJrveF8XFr476EMWveTIzZznh6spO7d4Vxa41LErq3xaUcI2v4cidPCJe95qlTwnj35i1hfP+wBi+hPAauKj5O/0NOhTVbIpJT7lAY+FhnQ6HB076IZCLvR1wichxS4hKRfPGGv6qoxCUiPTl4hsWl1VDiEpG+GvyWHyUuEenJXdOTNYpCezzFl7XH60c3K2V9I1Njf8WOTWl1Wmm8O54eLXPqnBeRvHEdcYlIvhwbAwmKyPEkBzdZK3GJSA8OeIPf8nNUAwmKyHHA6zuQoJnNM7OXzKzdzG7sJz7czP5POf5UecDSkBKXiPThRa/qkcbMmoE7gSuBmZRGlZnZa7HrgT3ufg5wB/DVtO0qcYlIX/U74poDtLv7RnfvBB4A5vdaZj5wX/n5Q8Dc8jiAiYa0j2sfe3b9kz9UOcfYBCAeECo7jdq2Rm0XqG0DVc+2nVHrBvaxZ8U/+UMTqly8zcxWVbxe4u5LKl6fBmyueL0FuLDXNt5Zxt27zewtYDzBZzKkicvde0wwaGar3H32ULahWo3atkZtF6htA9VobXP3eVm3IY1OFUVkMHUAUyteTym/1+8yZjYMGAO8EW1UiUtEBtNKYLqZTTOzVuBaYFmvZZYB15Wffwr4uXtcAZt1HdeS9EUy06hta9R2gdo2UI3ctpqU+6wWAyuAZmCpu68xs9uAVe6+jNJkPN81s3ZgN6XkFrKUxCYi0nB0qigiuaPEJSK5k0niSrsFIEtmtsnMnjez1b3qU7Joy1Iz22FmL1S8N87MHjazDeV/xzZQ2241s47yZ7fazD6WUdummtmjZrbWzNaY2RfL72f62QXtaojPLU+GvI+rfAvAeuDDlIrRVgIL3H3tkDYkgZltAma7e+bFimb2O8B+4G/d/bzye38J7Hb328tJf6y7f6VB2nYrsN/dvzbU7enVtsnAZHd/xsxGA08DnwA+S4afXdCuq2mAzy1PsjjiquYWAAHc/TFKV1kqVd4ecR+lL/6QS2hbQ3D3re7+TPn5PmAdpersTD+7oF1ylLJIXP3dAtBIvzwHfmZmT5vZwqwb049J7r61/HwbMCnLxvRjsZk9Vz6VzOQ0tlJ5pIELgKdooM+uV7ugwT63RqfO+b4udff3U7qb/fPlU6KGVC7Sa6R6lruAs4FZwFbg61k2xsxGAT8AbnD3vZWxLD+7ftrVUJ9bHmSRuKq5BSAz7t5R/ncH8CNKp7aNZHu5r+RIn8mOjNvzDnff7u4FL03Kdw8ZfnZm1kIpOXzP3X9Yfjvzz66/djXS55YXWSSuam4ByISZjSx3mmJmI4GPAC/Eaw25ytsjrgN+nGFbejiSFMo+SUafXXlIlHuBde7+jYpQpp9dUrsa5XPLk0wq58uXe/+Kd28B+G9D3oh+mNlZlI6yoHQ71PezbJuZ3Q9cTmnYk+3ALcDfAw8CpwOvAle7+5B3kie07XJKpzsObAIWVfQpDWXbLgX+H/A8787edhOl/qTMPrugXQtogM8tT3TLj4jkjjrnRSR3lLhEJHeUuEQkd5S4RCR3lLhEJHeUuEQkd5S4RCR3/j9URVm9+kTKzwAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Class:  Sandal\n"
     ]
    }
   ],
   "source": [
    "test_data = torchvision.datasets.FashionMNIST(root=\"data\", train=False, transform=torchvision.transforms.ToTensor())\n",
    "test_dataloader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=True)\n",
    "class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']\n",
    "img, lbl = next(iter(test_dataloader))\n",
    "img = img.reshape(1, 28, 28).numpy()\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(img[0])\n",
    "plt.colorbar()\n",
    "plt.grid(False)\n",
    "plt.show()\n",
    "\n",
    "print(\"Class: \", class_names[lbl[0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5b20b799",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PRED:  Sandal\n"
     ]
    }
   ],
   "source": [
    "ex = relax.vm.build(MLPModelFinal, target=\"llvm\")\n",
    "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
    "data_nd = tvm.nd.array(img.reshape(1, -1))\n",
    "\n",
    "nd_res = vm[\"main\"](data_nd)\n",
    "pred = np.argmax(nd_res.numpy(), axis=1)[0]\n",
    "print(\"PRED: \", class_names[pred])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a64d843c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}