assignment/assignment1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MLC Assignment 1: End-to-End Model Execution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 1: Model Preparation\n",
    "\n",
    "To get you familiar with the process of building and manipulating an end-to-end model using MLC, let's start from a simple image classification model.\n",
    "\n",
    "We first use the following commands to install necessary packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 -m pip install mlc-ai-nightly -f https://mlc.ai/wheels\n",
    "!python3 -m pip install torch torchvision torchaudio torchsummary --extra-index-url https://download.pytorch.org/whl/cpu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pickle as pkl\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torchvision\n",
    "import tvm\n",
    "import tvm.testing\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "from torch import nn\n",
    "from torchvision import transforms\n",
    "from tvm import topi, relax, te\n",
    "from tvm.script import tir as T\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is the model defined in PyTorch. It accepts a batch of images as input, and pass them through convolution layer, activation layer, pooling layer and fully-connected layers in order."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 4\n",
    "input_shape = (batch_size, 1, 28, 28)  # NCHW layout\n",
    "\n",
    "\n",
    "def pytorch_model():\n",
    "    list = []\n",
    "    list.append(nn.Conv2d(in_channels=1, out_channels=32, kernel_size=(3, 3), bias=True))\n",
    "    list.append(nn.ReLU())\n",
    "    list.append(nn.MaxPool2d(kernel_size=(2, 2)))\n",
    "    list.append(nn.Flatten())\n",
    "    list.append(nn.Linear(in_features=5408, out_features=100, bias=True))\n",
    "    list.append(nn.ReLU())\n",
    "    list.append(nn.Linear(in_features=100, out_features=10, bias=True))\n",
    "    list.append(nn.Softmax(dim=1))\n",
    "\n",
    "    model = nn.Sequential(*list).cpu()\n",
    "    name_map = {\n",
    "        \"0.weight\": \"conv2d_weight\",\n",
    "        \"0.bias\": \"conv2d_bias\",\n",
    "        \"4.weight\": \"linear0_weight\",\n",
    "        \"4.bias\": \"linear0_bias\",\n",
    "        \"6.weight\": \"linear1_weight\",\n",
    "        \"6.bias\": \"linear1_bias\",\n",
    "    }\n",
    "    for name, param in model.named_parameters():\n",
    "        param.data = torch.from_numpy(weight_map[name_map[name]]).cpu()\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We provide a pre-trained weight map for this model on the Fashion MNIST dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hide outputs\n",
    "!wget -nc https://github.com/mlc-ai/web-data/raw/main/models/fasionmnist_mlp_assignment_params.pkl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that its accuracy is about 84%."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the weight map from file.\n",
    "# The prediction accuracy of the weight map on test data is around 83.3%.\n",
    "weight_map = pkl.load(open(\"fasionmnist_mlp_assignment_params.pkl\", \"rb\"))\n",
    "class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',\n",
    "               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']\n",
    "\n",
    "\n",
    "def test(model, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        print_img = True\n",
    "        for data, label in test_loader:\n",
    "            data, label = data.cpu(), label.cpu()\n",
    "            output = model(data)\n",
    "            # sum up batch loss\n",
    "            test_loss += F.nll_loss(output, label, reduction=\"sum\").item()\n",
    "            # get the index of the max log-probability\n",
    "            pred = output.argmax(dim=1, keepdim=True)\n",
    "            if print_img:\n",
    "                imshow(data[0])\n",
    "                print(\"predict: {}, label: {}\".format(class_names[pred[0][0]], class_names[label[0]]))\n",
    "                print_img = False\n",
    "            correct += pred.eq(label.view_as(pred)).sum().item()\n",
    "\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "\n",
    "    print(\"\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n\".format(\n",
    "        test_loss, correct, len(test_loader.dataset),\n",
    "        100. * correct / len(test_loader.dataset)))\n",
    "\n",
    "\n",
    "def imshow(img):\n",
    "    img = img / 2 + 0.5\n",
    "    npimg = img.numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "test_data = torchvision.datasets.FashionMNIST(\n",
    "    \"./data\",\n",
    "    download=True,\n",
    "    train=False,\n",
    "    transform=transforms.Compose([transforms.ToTensor()])\n",
    ")\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    test_data, batch_size=batch_size, shuffle=False)\n",
    "test(pytorch_model(), test_loader)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2. Ingest Model From Pytorch\n",
    "\n",
    "To see the MLC abstraction of the end-to-end model, we need to ingest it from PyTorch and transform into TVMScript implementation. However, it is hard to manually do this. As you may have experienced in Exercise 1, writing a primitive tensor function for each model layer requires massive engineering efforts. Moreover, the manual writing process is error-prone - just imagine when you write dozens of lines of code while there exists some tiny bug in your implementation, finding the bug in could be annoying.\n",
    "\n",
    "Fortunately, in TVM there is a much simpler way of doing this. TVM provides a utility `relax.BlockBuilder` that can construct end-to-end models step by step in an IRModule that starts empty. (Recall that in Lecture 4 we introduced the dataflow block design of Relax, our MLC abstraction on computational graph level. And here the \"block\" in \"`BlockBuilder`\" stands for the dataflow blocks in Relax functions.)\n",
    "\n",
    "Specifically, in `BlockBuilder` we have an `emit_te` API, that helps convert a Tensor Expression operator description, which was introduced in Lecture 3, into a `call_tir` operation to the operator's corresponding TensorIR function (`call_tir` was introduced in Lecture 4 as well.) Compared with manually writing TensorIR functions, writing their Tensor Expression description can be done within only a few lines of code, which reduces the amount of efforts and is less likely for us to make mistakes.\n",
    "\n",
    "The signature of `emit_te` is `emit_te(func, *input)`, where `func` is a function that returns a Tensor Expression operator description, and `*input` is the inputs to `func`.\n",
    "\n",
    "Let's start with an introducing example. In the code block below, `relu` is a function that returns a Tensor Expression description of a ReLU operator. To construct a Relax function that executes a single ReLU operator, in function `emit_te_example` we first define a BlockBuilder instance `bb`. We also define a 2-dimensional 128x128 tensor variable `x`, which will serve as the input tensor of the ReLU operation (as well as the input of the Relax function).\n",
    "\n",
    "After that, we construct a Relax function `main` with `x` as input, using the `with bb.function(name, [*input])` API. Then we construct a dataflow block. Inside the dataflow block, we first have a `call_tir` to a TensorIR implementation of ReLU operator, through `emit_te`. The `emit_te` below generates a TensorIR function called \"`relu`\" in the IRModule, and add a `call_tir(relu, (x,), (128, 128), dtype=\"float32\")` operation in the dataflow block. And the `call_tir` is followed by a function return.\n",
    "\n",
    "After this construction, the BlockBuilder `bb` contains the constructed IRModule, which can be got by `bb.get()`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def relu(A):\n",
    "    B = te.compute(shape=(128, 128), fcompute=lambda i, j: te.max(A[i, j], 0), name=\"B\")\n",
    "    return B\n",
    "\n",
    "\n",
    "def emit_te_example():\n",
    "    bb = relax.BlockBuilder()\n",
    "    x = relax.Var(\"x\", (128, 128), relax.DynTensorType(2, \"float32\"))\n",
    "    with bb.function(\"main\", [x]):\n",
    "        with bb.dataflow():\n",
    "            lv0 = bb.emit_te(relu, x)\n",
    "            gv = bb.emit_output(lv0)\n",
    "        bb.emit_func_output(gv)\n",
    "    return bb.get()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function `emit_te_example` returns the constructed IRModule as output. To see what the BlockBuilder constructs, we print the IRModule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "\n",
    "mod = emit_te_example()\n",
    "IPython.display.Code(mod.script(), language=\"python\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the IRModule generated by the BlockBuilder does contain a TensorIR implementation of ReLU, and a Relax function which calls into the ReLU implementation via `call_tir`.\n",
    "\n",
    "Now it is your turn to use BlockBuilder and `emit_te` to create an IRModule equivalent to the PyTorch model defined above. You can write Tensor Expression descriptions for all the operators by yourself. Alternatively, TVM provides TOPI (short for \"TVM Operator Inventory\") library which wraps Tensor Expression descriptions for various operators. It is also encouraged if you can read the [documents](https://tvm.apache.org/docs/reference/api/python/topi.html) and find out a way to use them. The test function has been provided for you to check the correctness of your IRModule easily.\n",
    "\n",
    "Note that each Conv2d layer or linear layer in the model contains a bias add, which should be reflected in the IRModule you construct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model_via_emit_te():\n",
    "    bb = relax.BlockBuilder()\n",
    "    x = relax.Var(\"x\", input_shape, relax.DynTensorType(batch_size, \"float32\"))\n",
    "\n",
    "    conv2d_weight = relax.const(weight_map[\"conv2d_weight\"], \"float32\")\n",
    "    conv2d_bias = relax.const(weight_map[\"conv2d_bias\"].reshape(1, 32, 1, 1), \"float32\")\n",
    "    linear0_weight = relax.const(weight_map[\"linear0_weight\"], \"float32\")\n",
    "    linear0_bias = relax.const(weight_map[\"linear0_bias\"].reshape(1, 100), \"float32\")\n",
    "    linear1_weight = relax.const(weight_map[\"linear1_weight\"], \"float32\")\n",
    "    linear1_bias = relax.const(weight_map[\"linear1_bias\"].reshape(1, 10), \"float32\")\n",
    "\n",
    "    with bb.function(\"main\", [x]):\n",
    "        with bb.dataflow():\n",
    "           # TODO\n",
    "           ...\n",
    "        bb.emit_func_output(gv)\n",
    "\n",
    "    return bb.get()\n",
    "\n",
    "\n",
    "def build_mod(mod):\n",
    "    exec = relax.vm.build(mod, \"llvm\")\n",
    "    dev = tvm.cpu()\n",
    "    vm = relax.VirtualMachine(exec, dev)\n",
    "    return vm\n",
    "\n",
    "\n",
    "def check_equivalence(mod, torch_model, test_loader):\n",
    "    torch_model.eval()\n",
    "    with torch.no_grad():\n",
    "        rt_mod = build_mod(mod)\n",
    "        for data, label in test_loader:\n",
    "            data, label = data.cpu(), label.cpu()\n",
    "            output_from_pytorch = torch_model(data).numpy()\n",
    "            output_from_relax = rt_mod[\"main\"](tvm.nd.array(data, tvm.cpu())).numpy()\n",
    "            tvm.testing.assert_allclose(output_from_pytorch, output_from_relax, rtol=1e-4)\n",
    "\n",
    "\n",
    "test_data = torchvision.datasets.FashionMNIST(\n",
    "    \"./data\",\n",
    "    download=True,\n",
    "    train=False,\n",
    "    transform=transforms.Compose([transforms.ToTensor()])\n",
    ")\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "mod = create_model_via_emit_te()\n",
    "torch_model = pytorch_model()\n",
    "\n",
    "check_equivalence(mod, torch_model, test_loader)\n",
    "IPython.display.Code(mod.script(), language=\"python\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3. Use of Vendor Library\n",
    "\n",
    "As we have talked about in Lecture 4, we can integrate torch functions into an IRModule. The steps include registering an external runtime function and calling it inside the IRModule using `call_tir`.\n",
    "\n",
    "Here is an example of using torch matmul and torch add to implement a linear layer. You can also find this example in the Lecture 4 notes.\n",
    "\n",
    "```python\n",
    "@tvm.register_func(\"env.linear\", override=True)\n",
    "def torch_linear(x: tvm.nd.NDArray,\n",
    "                 w: tvm.nd.NDArray,\n",
    "                 b: tvm.nd.NDArray,\n",
    "                 out: tvm.nd.NDArray):\n",
    "    x_torch = torch.from_dlpack(x)\n",
    "    w_torch = torch.from_dlpack(w)\n",
    "    b_torch = torch.from_dlpack(b)\n",
    "    out_torch = torch.from_dlpack(out)\n",
    "    torch.mm(x_torch, w_torch.T, out=out_torch)\n",
    "    torch.add(out_torch, b_torch, out=out_torch)\n",
    "\n",
    "\n",
    "@tvm.script.ir_module\n",
    "class MyModuleWithExternCall:\n",
    "    @R.function\n",
    "    def main(x: Tensor((1, 784), \"float32\"),\n",
    "             w0: Tensor((128, 784), \"float32\"),\n",
    "             b0: Tensor((128,), \"float32\")):\n",
    "        # block 0\n",
    "        with R.dataflow():\n",
    "            lv0 = R.call_tir(\"env.linear\", (x, w0, b0), (1, 128), dtype=\"float32\")\n",
    "            ...\n",
    "        return ...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please register external functions for the convolution layer occurring in the IRModule you create in Section 2. You need to use NumPy or PyTorch as the function's implementation.\n",
    "\n",
    "You may use `BlockBuilder.emit` to directly add a `call_tir` operation to the end of the Relax function being constructed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def create_model_with_torch_func():\n",
    "    bb = relax.BlockBuilder()\n",
    "\n",
    "    x = relax.Var(\"x\", input_shape, relax.DynTensorType(4, \"float32\"))\n",
    "\n",
    "    conv2d_weight = relax.const(weight_map[\"conv2d_weight\"], \"float32\")\n",
    "    conv2d_bias = relax.const(weight_map[\"conv2d_bias\"].reshape(1, 32, 1, 1), \"float32\")\n",
    "    linear0_weight = relax.const(weight_map[\"linear0_weight\"], \"float32\")\n",
    "    linear0_bias = relax.const(weight_map[\"linear0_bias\"].reshape(1, 100), \"float32\")\n",
    "    linear1_weight = relax.const(weight_map[\"linear1_weight\"], \"float32\")\n",
    "    linear1_bias = relax.const(weight_map[\"linear1_bias\"].reshape(1, 10), \"float32\")\n",
    "\n",
    "    with bb.function(\"main\", [x]):\n",
    "        with bb.dataflow():\n",
    "            # TODO:\n",
    "            ...\n",
    "        bb.emit_func_output(gv)\n",
    "\n",
    "    return bb.get()\n",
    "\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False)\n",
    "mod = create_model_with_torch_func()\n",
    "check_equivalence(mod, torch_model, test_loader)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 4. Transformation in End-to–End Models\n",
    "\n",
    "In Exercise 1, we learned how to transform a single TensorIR Function. It is similar to do that in an end-to-end model.\n",
    "\n",
    "Compared with the batch matmul program, let's focus on a more challenging one: conv2d.\n",
    "\n",
    "To begin with, let's introduce some new primitives: \n",
    " - `compute_inline`: It inlines a block into another to reduce memory usage and memory access.\n",
    " - `fuse`: The opposite for `split`. Fuse multiple axes. Here `fuse` is used together with `parallel` / `vectorize` / `unroll` to further increase parallelism."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@T.prim_func\n",
    "def before_inline(a: T.handle, c: T.handle) -> None:\n",
    "    A = T.match_buffer(a, (128, 128))\n",
    "    B = T.alloc_buffer((128, 128))\n",
    "    C = T.match_buffer(c, (128, 128))\n",
    "    for i, j in T.grid(128, 128):\n",
    "        with T.block(\"B\"):\n",
    "            vi, vj = T.axis.remap(\"SS\", [i, j])\n",
    "            B[vi, vj] = A[vi, vj] * 2.0\n",
    "    for i, j in T.grid(128, 128):\n",
    "        with T.block(\"C\"):\n",
    "            vi, vj = T.axis.remap(\"SS\", [i, j])\n",
    "            C[vi, vj] = B[vi, vj] + 1.0\n",
    "\n",
    "\n",
    "sch = tvm.tir.Schedule(before_inline)\n",
    "sch.compute_inline(sch.get_block(\"B\"))\n",
    "IPython.display.Code(sch.mod[\"main\"].script(), language=\"python\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@T.prim_func\n",
    "def before_fuse(a: T.handle, b: T.handle) -> None:\n",
    "    A = T.match_buffer(a, (128, 128))\n",
    "    B = T.match_buffer(b, (128, 128))\n",
    "    for i, j in T.grid(128, 128):\n",
    "        with T.block(\"B\"):\n",
    "            vi, vj = T.axis.remap(\"SS\", [i, j])\n",
    "            B[vi, vj] = A[vi, vj] * 2.0\n",
    "\n",
    "\n",
    "sch = tvm.tir.Schedule(before_fuse)\n",
    "i, j = sch.get_loops(sch.get_block(\"B\"))\n",
    "sch.fuse(i, j)\n",
    "IPython.display.Code(sch.mod[\"main\"].script(), language=\"python\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we first create a schedule for the IRModule, and then transform the conv2d TensorIR function inside. Similar to Exercise 1, we provide you with a target function. But please note that, the target function does NOT serve as a \"standard transformation answer\" for several reasons:\n",
    " - it may not have the best performance on every hardware,\n",
    " - the original conv2d TensorIR implementation may vary, according to the Tensor Expression description you used in Section 2:\n",
    "   - if you described the conv2d computation along with the bias computation in Tensor Expression, then there should be a block which calculates the bias at the end of target TensorIR function,\n",
    "   - if you described conv2d and bias computation separately, or you used the conv2d provided by TOPI, then the target function should not have the bias block at the end. The original function of the target is generated by using TOPI conv2d.\n",
    "\n",
    "\n",
    "```python\n",
    "@T.prim_func\n",
    "def target_func(rxplaceholder: T.Buffer[(4, 1, 28, 28), \"float32\"], rxplaceholder_1: T.Buffer[(32, 1, 3, 3), \"float32\"], conv2d_nchw: T.Buffer[(4, 32, 26, 26), \"float32\"]) -> None:\n",
    "    T.func_attr({\"global_symbol\": \"conv2d\", \"tir.noalias\": True})\n",
    "    # body\n",
    "    # with T.block(\"root\")\n",
    "    for i0_0_i1_0_i2_0_i3_0_fused in T.parallel(2704):\n",
    "        for i0_1_i1_1_fused_init in T.unroll(8):\n",
    "            for i2_1_i3_1_fused_init in T.vectorized(4):\n",
    "                with T.block(\"conv2d_nchw_init\"):\n",
    "                    nn = T.axis.spatial(\n",
    "                        4, i0_0_i1_0_i2_0_i3_0_fused // 1352 * 2 + i0_1_i1_1_fused_init // 4)\n",
    "                    ff = T.axis.spatial(\n",
    "                        32, i0_0_i1_0_i2_0_i3_0_fused % 1352 // 169 * 4 + i0_1_i1_1_fused_init % 4)\n",
    "                    yy = T.axis.spatial(\n",
    "                        26, i0_0_i1_0_i2_0_i3_0_fused % 169 // 13 * 2 + i2_1_i3_1_fused_init // 2)\n",
    "                    xx = T.axis.spatial(\n",
    "                        26, i0_0_i1_0_i2_0_i3_0_fused % 13 * 2 + i2_1_i3_1_fused_init % 2)\n",
    "                    T.reads()\n",
    "                    T.writes(conv2d_nchw[nn, ff, yy, xx])\n",
    "                    conv2d_nchw[nn, ff, yy, xx] = T.float32(0)\n",
    "        for i4, i5, i6 in T.grid(1, 3, 3):\n",
    "            for i0_1_i1_1_fused in T.unroll(8):\n",
    "                for i2_1_i3_1_fused in T.vectorized(4):\n",
    "                    with T.block(\"conv2d_nchw_update\"):\n",
    "                        nn = T.axis.spatial(\n",
    "                            4, i0_0_i1_0_i2_0_i3_0_fused // 1352 * 2 + i0_1_i1_1_fused // 4)\n",
    "                        ff = T.axis.spatial(\n",
    "                            32, i0_0_i1_0_i2_0_i3_0_fused % 1352 // 169 * 4 + i0_1_i1_1_fused % 4)\n",
    "                        yy = T.axis.spatial(\n",
    "                            26, i0_0_i1_0_i2_0_i3_0_fused % 169 // 13 * 2 + i2_1_i3_1_fused // 2)\n",
    "                        xx = T.axis.spatial(\n",
    "                            26, i0_0_i1_0_i2_0_i3_0_fused % 13 * 2 + i2_1_i3_1_fused % 2)\n",
    "                        rc, ry, rx = T.axis.remap(\"RRR\", [i4, i5, i6])\n",
    "                        T.reads(conv2d_nchw[nn, ff, yy, xx], rxplaceholder[nn,\n",
    "                                rc, yy + ry, xx + rx], rxplaceholder_1[ff, rc, ry, rx])\n",
    "                        T.writes(conv2d_nchw[nn, ff, yy, xx])\n",
    "                        conv2d_nchw[nn, ff, yy, xx] = conv2d_nchw[nn, ff, yy, xx] + \\\n",
    "                            rxplaceholder[nn, rc, yy + ry, xx +\n",
    "                                          rx] * rxplaceholder_1[ff, rc, ry, rx]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unlike Exercise 1, this time the schedule is created for an IRModule, instead of a TensorIR function. Therefore, when using `sch.get_block`, a concrete function name should be provided, as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mod = create_model_via_emit_te()\n",
    "sch = tvm.tir.Schedule(mod)\n",
    "\n",
    "# Step 1. Get blocks\n",
    "# block = sch.get_block(name=\"your_block_name\", func_name=\"your_function_name\")\n",
    "\n",
    "# Step 2. Inline the padding block (if exists)\n",
    "\n",
    "# Step 3. Get loops\n",
    "\n",
    "# Step 4. Organize the loops\n",
    "\n",
    "# Step 5. decompose reduction\n",
    "\n",
    "# Step 6. fuse + vectorize / fuse + parallel / fuse + unroll\n",
    "\n",
    "IPython.display.Code(sch.mod.script(), language=\"python\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we can test the correctness of the transformed IRModule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=False)\n",
    "check_equivalence(sch.mod, torch_model, test_loader)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "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.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}