![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Mpn1ti5Urdsv"
},
"source": [
"# Ep9: Computational Graph Optimization"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "qXysoqn-vZuF"
},
"source": [
"## Install packages \n",
"\n",
"For this course, we will use some ongoing development in tvm, which is an open-source machine learning compilation framework. We provide the following command to install a packaged version for mlc course. "
]
},
{
"cell_type": "code",
"execution_count": 70,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Xe3vClsD9jlq",
"outputId": "24043ae0-92cf-4677-e274-8df5b2d8c3a6"
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
"Looking in links: https://mlc.ai/wheels\n",
"Requirement already satisfied: mlc-ai-nightly in /usr/local/lib/python3.7/dist-packages (0.9.dev2226+gf68b7661e)\n",
"Requirement already satisfied: attrs in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (22.1.0)\n",
"Requirement already satisfied: psutil in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (5.4.8)\n",
"Requirement already satisfied: synr==0.6.0 in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (0.6.0)\n",
"Requirement already satisfied: decorator in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (4.4.2)\n",
"Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (1.7.3)\n",
"Requirement already satisfied: tornado in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (5.1.1)\n",
"Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (1.21.6)\n",
"Requirement already satisfied: cloudpickle in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (1.5.0)\n"
]
}
],
"source": [
"!python3 -m pip install mlc-ai-nightly -f https://mlc.ai/wheels"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "i-14C4skxIrJ"
},
"source": [
"## Prelude\n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "CKH9qNeRrBwJ"
},
"source": [
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "HH7xhkFhrElt"
},
"source": [
"Most of the MLC process can be viewed as transformation among tensor functions. In the past chapters, we studied how to transform each primitive tensor functions individually. In this chapter, let us talk about high-level transformations among computational graphs.\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BBIuE2jc1DaU"
},
"source": [
"## Preparations\n",
"\n",
"To begin with, let us import the necessary dependencies.\n"
]
},
{
"cell_type": "code",
"execution_count": 71,
"metadata": {
"id": "BVp0fHyRkYj6"
},
"outputs": [],
"source": [
"# This is needed for deferring annotation parsing in TVMScript\n",
"from __future__ import annotations \n",
"\n",
"import tvm\n",
"from tvm.ir.module import IRModule\n",
"from tvm.script import tir as T, relax as R\n",
"from tvm import relax, topi\n",
"import numpy as np"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aBhwDhrlrSiK"
},
"source": [
"## Pattern Match and Rewriting\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "F8Zqd1_qrmXN"
},
"source": [
"To begin with, let us start with the following example."
]
},
{
"cell_type": "code",
"execution_count": 72,
"metadata": {
"id": "sR45I3mvm695"
},
"outputs": [],
"source": [
"@tvm.script.ir_module\n",
"class MyModule:\n",
" @R.function\n",
" def main(x: Tensor((3, 4), \"float32\"), y: Tensor((3, 4), \"float32\")):\n",
" with relax.dataflow():\n",
" lv0 = relax.multiply(x, y)\n",
" gv0 = relax.add(lv0, y)\n",
" relax.output(gv0)\n",
" return gv0"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ZsN9Q4acu-al"
},
"source": [
"`MyModule` contains a relax function with two high-level operators, relax.multiply and relax.add. Our goal is to find these two operators and replace it\n",
"with a call into `relax.ewise_fma` operator."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "SWGxwuDB0PMV"
},
"source": [
"Before we dive into how to do that exactly, let us first examine the data structure that makes up the MyModule. Each IRModule contains a collection of functions, and the function body is composed of a set of data structures called abstract syntax trees (AST)."
]
},
{
"cell_type": "code",
"execution_count": 73,
"metadata": {
"id": "8L9tLb4s0rIB"
},
"outputs": [],
"source": [
"relax_func = MyModule[\"main\"]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "TBUIvV6c4Gzw"
},
"source": [
"Each function is represented by a relax.Function node"
]
},
{
"cell_type": "code",
"execution_count": 74,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "SYM-u_IJ0vNy",
"outputId": "26709d64-71b9-4d45-86d2-2d110524e6e7"
},
"outputs": [
{
"data": {
"text/plain": [
"tvm.relax.expr.Function"
]
},
"execution_count": 74,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"type(relax_func)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "V4NAP23h4QM4"
},
"source": [
"The function contains a list of parameters"
]
},
{
"cell_type": "code",
"execution_count": 75,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "BeNcHrEd4Mw8",
"outputId": "74afca54-7afd-4632-ceb6-2b592766896f"
},
"outputs": [
{
"data": {
"text/plain": [
"[relax.expr.Var(0xa8f2d80), relax.expr.Var(0xa750480)]"
]
},
"execution_count": 75,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"relax_func.params"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "kQYT-1Ke4VQf"
},
"source": [
"The function contains a body fields that represents its return value and set of binding blocks in the function."
]
},
{
"cell_type": "code",
"execution_count": 76,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "0hhXxL_20x-N",
"outputId": "ae106676-8101-4efc-e31f-4dc4513cc524"
},
"outputs": [
{
"data": {
"text/plain": [
"tvm.relax.expr.SeqExpr"
]
},
"execution_count": 76,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"func_body = relax_func.body\n",
"type(func_body)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "J4dQoc525CxL"
},
"source": [
"The function body SeqExpr contains a sequence of (binding) blocks"
]
},
{
"cell_type": "code",
"execution_count": 77,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "3nM-F8Gh2PzX",
"outputId": "9df0aeb1-cc23-49b7-8b9c-fa8a81c6f5a3"
},
"outputs": [
{
"data": {
"text/plain": [
"[relax.expr.DataflowBlock(0x3444520)]"
]
},
"execution_count": 77,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"func_body.blocks"
]
},
{
"cell_type": "code",
"execution_count": 78,
"metadata": {
"id": "CV_byz4_2jmD"
},
"outputs": [],
"source": [
"dataflow_block = func_body.blocks[0]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "OsDRAIB_5Lig"
},
"source": [
"In our particular case, we have a single data flow block that contains\n",
"two bindings. Each binding corresponds to one of the following two lines\n",
"\n",
"```python\n",
"lv0 = relax.multiply(x, y)\n",
"gv0 = relax.add(lv0, y)\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Zw8Yw9SV2o7H"
},
"outputs": [],
"source": [
"dataflow_block.bindings"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "A8RZMKD22slK"
},
"outputs": [],
"source": [
"binding = dataflow_block.bindings[0]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yHKybEKE5hq3"
},
"source": [
"Each binding have a var field that corresponds to the left hand side of the binding (lv0, gv0)"
]
},
{
"cell_type": "code",
"execution_count": 79,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "nzzMhaSS29UR",
"outputId": "6e03bfef-c441-4f40-dfda-0b88bb69bfcb"
},
"outputs": [
{
"data": {
"text/plain": [
"relax.expr.DataflowVar(0x6c36300)"
]
},
"execution_count": 79,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"binding.var"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-cAOJAmF5o5Z"
},
"source": [
"And its value field corresponds to the right-hand side of the binding. Each value field corresponds to a `relax.Call` node representing a call into a primitive function"
]
},
{
"cell_type": "code",
"execution_count": 80,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "qfP5rEdD2_1c",
"outputId": "0151ffb1-c5a6-4a06-db9f-b89d263e0895"
},
"outputs": [
{
"data": {
"text/plain": [
"CallNode(Op(relax.multiply), [relax.expr.Var(0x3463080), relax.expr.Var(0x3463300)], (nullptr), [])"
]
},
"execution_count": 80,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"binding.value"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "1oV8O46kWsyG"
},
"source": [
""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "iogK4y4kWxm9"
},
"source": [
"The above figure summarizes the data structure involved in this particular function."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "_8Q1V8b16K3N"
},
"source": [
"One approach to rewrite the program would be to traverse MyModule's AST recursively and generate a transformed AST. We can certainly do that using the python API available. However, we can use extra tooling support to simplify the process. The following code block follows a design pattern called **visitor pattern** that allows us to visit each AST node and rewrite them to transformed versions."
]
},
{
"cell_type": "code",
"execution_count": 81,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 157
},
"id": "Rgz_VxzPsEx8",
"outputId": "434671e8-3b3e-4d45-dd44-d5c67badd0a8"
},
"outputs": [
{
"data": {
"text/html": [
"@R.function\n",
"def main(x: Tensor((3, 4), "float32"), y: Tensor((3, 4), "float32")) -> Tensor(None, "float32", ndim = 2):\n",
" # block 0\n",
" with R.dataflow():\n",
" lv0: Tensor((3, 4), "float32") = relax.multiply(x, y)\n",
" gv0: Tensor((3, 4), "float32") = relax.ewise_fma(x, y, y)\n",
" R.output(gv0)\n",
" return gv0\n",
"@R.function\n",
"def main(x: Tensor((3, 4), "float32"), y: Tensor((3, 4), "float32")) -> Tensor(None, "float32", ndim = 2):\n",
" # block 0\n",
" with R.dataflow():\n",
" gv0: Tensor((3, 4), "float32") = relax.ewise_fma(x, y, y)\n",
" R.output(gv0)\n",
" return gv0\n",
"@tvm.script.ir_module\n",
"class Module:\n",
" @R.function\n",
" def main(x: Tensor((1, 784), "float32")) -> Tensor(None, "float32", ndim = 2):\n",
" # block 0\n",
" with R.dataflow():\n",
" lv: Tensor((1, 128), "float32") = relax.nn.dense(x, meta[relay.Constant][0])\n",
" lv1: Tensor((1, 128), "float32") = relax.add(lv, meta[relay.Constant][1])\n",
" lv2: Tensor((1, 128), "float32") = relax.nn.relu(lv1)\n",
" lv3: Tensor((1, 10), "float32") = relax.nn.dense(lv2, meta[relay.Constant][2])\n",
" lv4: Tensor((1, 10), "float32") = relax.add(lv3, meta[relay.Constant][3])\n",
" gv: Tensor((1, 10), "float32") = lv4\n",
" R.output(gv)\n",
" return gv\n",
" \n",
"@tvm.script.ir_module\n",
"class Module:\n",
" @R.function\n",
" def fused_dense_add0(x: Tensor((1, 784), "float32"), w: Tensor((128, 784), "float32"), b: Tensor((128,), "float32")) -> Tensor(None, "float32", ndim = 2):\n",
" # block 0\n",
" with R.dataflow():\n",
" lv: Tensor((1, 128), "float32") = relax.nn.dense(x, w)\n",
" gv: Tensor((1, 128), "float32") = relax.add(lv, b)\n",
" R.output(gv)\n",
" return gv\n",
" \n",
" @R.function\n",
" def fused_dense_add1(x1: Tensor((1, 128), "float32"), w1: Tensor((10, 128), "float32"), b1: Tensor((10,), "float32")) -> Tensor(None, "float32", ndim = 2):\n",
" # block 0\n",
" with R.dataflow():\n",
" lv1: Tensor((1, 10), "float32") = relax.nn.dense(x1, w1)\n",
" gv1: Tensor((1, 10), "float32") = relax.add(lv1, b1)\n",
" R.output(gv1)\n",
" return gv1\n",
" \n",
" @R.function\n",
" def main(x2: Tensor((1, 784), "float32")) -> Tensor(None, "float32", ndim = 2):\n",
" # block 0\n",
" with R.dataflow():\n",
" lv11: Tensor((1, 128), "float32") = fused_dense_add0(x2, meta[relay.Constant][0], meta[relay.Constant][1])\n",
" lv2: Tensor((1, 128), "float32") = relax.nn.relu(lv11)\n",
" lv4: Tensor((1, 10), "float32") = fused_dense_add1(lv2, meta[relay.Constant][2], meta[relay.Constant][3])\n",
" gv2: Tensor((1, 10), "float32") = lv4\n",
" R.output(gv2)\n",
" return gv2\n",
" \n",
"@tvm.script.ir_module\n",
"class Module:\n",
" @R.function\n",
" def main(x: Tensor((1, 784), "float32")) -> Tensor(None, "float32", ndim = 2):\n",
" # block 0\n",
" with R.dataflow():\n",
" lv1: Tensor((1, 128), "float32") = fused_dense_add0(x, meta[relay.Constant][0], meta[relay.Constant][1])\n",
" lv2 = R.call_tir(relu, (lv1,), (1, 128), dtype="float32")\n",
" lv4: Tensor((1, 10), "float32") = fused_dense_add1(lv2, meta[relay.Constant][2], meta[relay.Constant][3])\n",
" gv: Tensor((1, 10), "float32") = lv4\n",
" R.output(gv)\n",
" return gv\n",
" \n",
" @T.prim_func\n",
" def dense(rxplaceholder: T.Buffer[(1, 784), "float32"], rxplaceholder_1: T.Buffer[(T.int64(128), T.int64(784)), "float32"], T_matmul_NT: T.Buffer[(1, T.int64(128)), "float32"]) -> None:\n",
" # function attr dict\n",
" T.func_attr({"global_symbol": "dense", "tir.noalias": True, "layout_free_buffers": [1]})\n",
" # body\n",
" # with T.block("root")\n",
" for i0, i1, i2 in T.grid(1, T.int64(128), 784):\n",
" with T.block("T_matmul_NT"):\n",
" i = T.axis.spatial(1, i0)\n",
" j = T.axis.spatial(T.int64(128), i1)\n",
" k = T.axis.reduce(784, i2)\n",
" T.reads(rxplaceholder[i, k], rxplaceholder_1[j, k])\n",
" T.writes(T_matmul_NT[i, j])\n",
" with T.init():\n",
" T_matmul_NT[i, j] = T.float32(0)\n",
" T_matmul_NT[i, j] = T_matmul_NT[i, j] + rxplaceholder[i, k] * rxplaceholder_1[j, k]\n",
" \n",
" @T.prim_func\n",
" def add1(rxplaceholder: T.Buffer[(1, T.int64(10)), "float32"], rxplaceholder_1: T.Buffer[T.int64(10), "float32"], T_add: T.Buffer[(1, T.int64(10)), "float32"]) -> None:\n",
" # function attr dict\n",
" T.func_attr({"global_symbol": "add1", "tir.noalias": True})\n",
" # body\n",
" # with T.block("root")\n",
" for i0, i1 in T.grid(1, T.int64(10)):\n",
" with T.block("T_add"):\n",
" ax0 = T.axis.spatial(1, i0)\n",
" ax1 = T.axis.spatial(T.int64(10), i1)\n",
" T.reads(rxplaceholder[ax0, ax1], rxplaceholder_1[ax1])\n",
" T.writes(T_add[ax0, ax1])\n",
" T_add[ax0, ax1] = rxplaceholder[ax0, ax1] + rxplaceholder_1[ax1]\n",
" \n",
" @T.prim_func\n",
" def add(rxplaceholder: T.Buffer[(1, T.int64(128)), "float32"], rxplaceholder_1: T.Buffer[T.int64(128), "float32"], T_add: T.Buffer[(1, T.int64(128)), "float32"]) -> None:\n",
" # function attr dict\n",
" T.func_attr({"global_symbol": "add", "tir.noalias": True})\n",
" # body\n",
" # with T.block("root")\n",
" for i0, i1 in T.grid(1, T.int64(128)):\n",
" with T.block("T_add"):\n",
" ax0 = T.axis.spatial(1, i0)\n",
" ax1 = T.axis.spatial(T.int64(128), i1)\n",
" T.reads(rxplaceholder[ax0, ax1], rxplaceholder_1[ax1])\n",
" T.writes(T_add[ax0, ax1])\n",
" T_add[ax0, ax1] = rxplaceholder[ax0, ax1] + rxplaceholder_1[ax1]\n",
" \n",
" @R.function\n",
" def fused_dense_add1(x1: Tensor((1, 128), "float32"), w: Tensor((10, 128), "float32"), b: Tensor((10,), "float32")) -> Tensor(None, "float32", ndim = 2):\n",
" # block 0\n",
" with R.dataflow():\n",
" lv = R.call_tir(dense1, (x1, w), (1, 10), dtype="float32")\n",
" gv1 = R.call_tir(add1, (lv, b), (1, 10), dtype="float32")\n",
" R.output(gv1)\n",
" return gv1\n",
" \n",
" @T.prim_func\n",
" def dense1(rxplaceholder: T.Buffer[(1, T.int64(128)), "float32"], rxplaceholder_1: T.Buffer[(T.int64(10), T.int64(128)), "float32"], T_matmul_NT: T.Buffer[(1, T.int64(10)), "float32"]) -> None:\n",
" # function attr dict\n",
" T.func_attr({"global_symbol": "dense1", "tir.noalias": True, "layout_free_buffers": [1]})\n",
" # body\n",
" # with T.block("root")\n",
" for i0, i1, i2 in T.grid(1, T.int64(10), T.int64(128)):\n",
" with T.block("T_matmul_NT"):\n",
" i = T.axis.spatial(1, i0)\n",
" j = T.axis.spatial(T.int64(10), i1)\n",
" k = T.axis.reduce(T.int64(128), i2)\n",
" T.reads(rxplaceholder[i, k], rxplaceholder_1[j, k])\n",
" T.writes(T_matmul_NT[i, j])\n",
" with T.init():\n",
" T_matmul_NT[i, j] = T.float32(0)\n",
" T_matmul_NT[i, j] = T_matmul_NT[i, j] + rxplaceholder[i, k] * rxplaceholder_1[j, k]\n",
" \n",
" @T.prim_func\n",
" def relu(rxplaceholder: T.Buffer[(1, T.int64(128)), "float32"], compute: T.Buffer[(1, T.int64(128)), "float32"]) -> None:\n",
" # function attr dict\n",
" T.func_attr({"global_symbol": "relu", "tir.noalias": True})\n",
" # body\n",
" # with T.block("root")\n",
" for i0, i1 in T.grid(1, T.int64(128)):\n",
" with T.block("compute"):\n",
" i0_1 = T.axis.spatial(1, i0)\n",
" i1_1 = T.axis.spatial(T.int64(128), i1)\n",
" T.reads(rxplaceholder[i0_1, i1_1])\n",
" T.writes(compute[i0_1, i1_1])\n",
" compute[i0_1, i1_1] = T.max(rxplaceholder[i0_1, i1_1], T.float32(0))\n",
" \n",
" @R.function\n",
" def fused_dense_add0(x2: Tensor((1, 784), "float32"), w1: Tensor((128, 784), "float32"), b1: Tensor((128,), "float32")) -> Tensor(None, "float32", ndim = 2):\n",
" # block 0\n",
" with R.dataflow():\n",
" lv3 = R.call_tir(dense, (x2, w1), (1, 128), dtype="float32")\n",
" gv2 = R.call_tir(add, (lv3, b1), (1, 128), dtype="float32")\n",
" R.output(gv2)\n",
" return gv2\n",
" \n",
"@tvm.script.ir_module\n",
"class Module:\n",
" @T.prim_func\n",
" def fused_dense_add0(x: T.Buffer[(1, 784), "float32"], w: T.Buffer[(T.int64(128), T.int64(784)), "float32"], b: T.Buffer[T.int64(128), "float32"], T_add: T.Buffer[(1, T.int64(128)), "float32"]) -> None:\n",
" # function attr dict\n",
" T.func_attr({"tir.noalias": True, "global_symbol": "fused_dense_add0"})\n",
" # body\n",
" # with T.block("root")\n",
" T_matmul_NT = T.alloc_buffer([1, T.int64(128)], dtype="float32")\n",
" for i0, i1, i2 in T.grid(1, T.int64(128), 784):\n",
" with T.block("T_matmul_NT"):\n",
" i = T.axis.spatial(1, i0)\n",
" j = T.axis.spatial(T.int64(128), i1)\n",
" k = T.axis.reduce(784, i2)\n",
" T.reads(x[i, k], w[j, k])\n",
" T.writes(T_matmul_NT[i, j])\n",
" with T.init():\n",
" T_matmul_NT[i, j] = T.float32(0)\n",
" T_matmul_NT[i, j] = T_matmul_NT[i, j] + x[i, k] * w[j, k]\n",
" for i0, i1 in T.grid(1, T.int64(128)):\n",
" with T.block("T_add"):\n",
" ax0 = T.axis.spatial(1, i0)\n",
" ax1 = T.axis.spatial(T.int64(128), i1)\n",
" T.reads(T_matmul_NT[ax0, ax1], b[ax1])\n",
" T.writes(T_add[ax0, ax1])\n",
" T_add[ax0, ax1] = T_matmul_NT[ax0, ax1] + b[ax1]\n",
" \n",
" @T.prim_func\n",
" def relu(rxplaceholder: T.Buffer[(1, T.int64(128)), "float32"], compute: T.Buffer[(1, T.int64(128)), "float32"]) -> None:\n",
" # function attr dict\n",
" T.func_attr({"global_symbol": "relu", "tir.noalias": True})\n",
" # body\n",
" # with T.block("root")\n",
" for i0, i1 in T.grid(1, T.int64(128)):\n",
" with T.block("compute"):\n",
" i0_1 = T.axis.spatial(1, i0)\n",
" i1_1 = T.axis.spatial(T.int64(128), i1)\n",
" T.reads(rxplaceholder[i0_1, i1_1])\n",
" T.writes(compute[i0_1, i1_1])\n",
" compute[i0_1, i1_1] = T.max(rxplaceholder[i0_1, i1_1], T.float32(0))\n",
" \n",
" @T.prim_func\n",
" def fused_dense_add1(x: T.Buffer[(1, T.int64(128)), "float32"], w: T.Buffer[(T.int64(10), T.int64(128)), "float32"], b: T.Buffer[T.int64(10), "float32"], T_add: T.Buffer[(1, T.int64(10)), "float32"]) -> None:\n",
" # function attr dict\n",
" T.func_attr({"tir.noalias": True, "global_symbol": "fused_dense_add1"})\n",
" # body\n",
" # with T.block("root")\n",
" T_matmul_NT = T.alloc_buffer([1, T.int64(10)], dtype="float32")\n",
" for i0, i1, i2 in T.grid(1, T.int64(10), T.int64(128)):\n",
" with T.block("T_matmul_NT"):\n",
" i = T.axis.spatial(1, i0)\n",
" j = T.axis.spatial(T.int64(10), i1)\n",
" k = T.axis.reduce(T.int64(128), i2)\n",
" T.reads(x[i, k], w[j, k])\n",
" T.writes(T_matmul_NT[i, j])\n",
" with T.init():\n",
" T_matmul_NT[i, j] = T.float32(0)\n",
" T_matmul_NT[i, j] = T_matmul_NT[i, j] + x[i, k] * w[j, k]\n",
" for i0, i1 in T.grid(1, T.int64(10)):\n",
" with T.block("T_add"):\n",
" ax0 = T.axis.spatial(1, i0)\n",
" ax1 = T.axis.spatial(T.int64(10), i1)\n",
" T.reads(T_matmul_NT[ax0, ax1], b[ax1])\n",
" T.writes(T_add[ax0, ax1])\n",
" T_add[ax0, ax1] = T_matmul_NT[ax0, ax1] + b[ax1]\n",
" \n",
" @R.function\n",
" def main(x: Tensor((1, 784), "float32")) -> Tensor(None, "float32", ndim = 2):\n",
" # block 0\n",
" with R.dataflow():\n",
" lv1 = R.call_tir(fused_dense_add0, (x, meta[relay.Constant][0], meta[relay.Constant][1]), (1, 128), dtype="float32")\n",
" lv2 = R.call_tir(relu, (lv1,), (1, 128), dtype="float32")\n",
" lv4 = R.call_tir(fused_dense_add1, (lv2, meta[relay.Constant][2], meta[relay.Constant][3]), (1, 10), dtype="float32")\n",
" gv: Tensor((1, 10), "float32") = lv4\n",
" R.output(gv)\n",
" return gv\n",
" \n",
"