{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 2D 卷积优化\n",
"\n",
"**原作者**: [Thierry Moreau](https://homes.cs.washington.edu/~moreau/)\n",
"\n",
"本教程提供了关于如何使用 TVM 映射二维卷积工作负载有效的 VTA 设计的概述。建议先学习 {ref}`vta-mat-mult-opt` 教程。\n",
"\n",
"二维卷积在大多数计算机视觉深度神经网络中占主导地位。在本教程中，将演示 TVM 调度优化，将 NCHW 布局中的 2D 卷积算子映射到 VTA。还引入了延迟隐藏（latency hiding）的概念，它允许最大化 VTA 的计算和内存资源利用。\n",
"\n",
"## RPC 设置\n",
"\n",
"首先编程 Pynq 的 FPGA 并构建它的 RPC 运行时。"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"import os\n",
"import tvm\n",
"import tvm.testing\n",
"from tvm import te\n",
"import vta\n",
"import numpy as np\n",
"\n",
"from tvm import rpc\n",
"from tvm.contrib import utils\n",
"from vta.testing import simulator\n",
"\n",
"# Load VTA parameters from the 3rdparty/vta-hw/config/vta_config.json file\n",
"env = vta.get_env()\n",
"\n",
"# We read the Pynq RPC host IP address and port number from the OS environment\n",
"host = os.environ.get(\"VTA_RPC_HOST\", \"192.168.2.99\")\n",
"port = int(os.environ.get(\"VTA_RPC_PORT\", \"9091\"))\n",
"\n",
"# We configure both the bitstream and the runtime system on the Pynq\n",
"# to match the VTA configuration specified by the vta_config.json file.\n",
"if env.TARGET == \"pynq\":\n",
" # Make sure that TVM was compiled with RPC=1\n",
" assert tvm.runtime.enabled(\"rpc\")\n",
" remote = rpc.connect(host, port)\n",
"\n",
" # Reconfigure the JIT runtime\n",
" vta.reconfig_runtime(remote)\n",
"\n",
" # Program the FPGA with a pre-compiled VTA bitstream.\n",
" # You can program the FPGA with your own custom bitstream\n",
" # by passing the path to the bitstream file instead of None.\n",
" vta.program_fpga(remote, bitstream=None)\n",
"\n",
"# In simulation mode, host the RPC server locally.\n",
"elif env.TARGET in [\"sim\", \"tsim\"]:\n",
" remote = rpc.LocalSession()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 声明计算\n",
"\n",
"作为第一步，需要用 NCHW 格式描述 2D 卷积计算。\n",
"\n",
"通过 batch size、空间维度、输入通道、输出通道、核维度、核维度、填充维度和步长维度来定义二维卷积形状。\n",
"\n",
"选择 ResNet-18 架构的第 9 个卷积层的形状作为卷积 workload 参数。\n",
"\n",
"在 2D 卷积中添加了额外的算子，用于对输出进行移位和剪切，以模拟定点卷积之后的修正线性激活。将二维卷积层的 TVM 数据流图描述如下：\n",
"\n",
"```{image} images/conv2d_dataflow.png\n",
":align: center\n",
"```\n",
"\n",
"这个计算被故意设置得太大，以至于不能一次全部放入 VTA 的 on-chip buffers。因此，在调度阶段，将依靠计算分块策略将计算分解为可管理的块。\n",
"\n",
"````{admonition} 空间填充\n",
":class: alert alert-info\n",
"\n",
"注意，需要导入 TOPI 库来对输入特征映射张量应用空间填充（Spatial padding）。空间填充有助于在 2D 卷积环境中分块，因为如果卷积核窗口大小大于 1，那么任何给定层的输入特征映射的相同 `(x, y)` 空间位置将被读取多次。在 CPU 和 GPU 上，当并行工作时，提高内存访问效率的一种方法是空间打包（spatial packing），这需要重新布局数据。VTA load DMA 引擎可以自动插入填充，这样原始的输入特征映射就不必在内存中重新打包。\n",
"\n",
"当数据从 DRAM load 到 VTA 的 SRAM 时，下面展示了 VTA 对动态空间填充的影响，随后是 2D 跨步和填充内存读取。\n",
"\n",
"```{image} images/padding.png\n",
":align: center\n",
":width: 480px\n",
"```\n",
"````"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"from tvm import topi\n",
"\n",
"# 2D convolution layer dimensions taken from ResNet-18 architecture\n",
"# (9th convolutional layer)\n",
"batch_size = 1\n",
"height = 14\n",
"width = 14\n",
"in_channels = 256\n",
"out_channels = 256\n",
"kernel_h = 3\n",
"kernel_w = 3\n",
"pad_h = 1\n",
"pad_w = 1\n",
"stride_h = 1\n",
"stride_w = 1\n",
"assert batch_size % env.BATCH == 0\n",
"assert in_channels % env.BLOCK_IN == 0\n",
"assert out_channels % env.BLOCK_OUT == 0\n",
"\n",
"# Input feature map: (N, IC, H, W, n, ic)\n",
"data_shape = (\n",
" batch_size // env.BATCH,\n",
" in_channels // env.BLOCK_IN,\n",
" height,\n",
" width,\n",
" env.BATCH,\n",
" env.BLOCK_IN,\n",
")\n",
"# Kernel: (OC, IC, H, W, oc, ic)\n",
"kernel_shape = (\n",
" out_channels // env.BLOCK_OUT,\n",
" in_channels // env.BLOCK_IN,\n",
" kernel_h,\n",
" kernel_w,\n",
" env.BLOCK_OUT,\n",
" env.BLOCK_IN,\n",
")\n",
"# Derive output feature map dimensions\n",
"fout_height = (height + 2 * pad_h - kernel_h) // stride_h + 1\n",
"fout_width = (width + 2 * pad_w - kernel_w) // stride_w + 1\n",
"# Output feature map: (N, OC, H, W, n, oc)\n",
"output_shape = (\n",
" batch_size // env.BATCH,\n",
" out_channels // env.BLOCK_OUT,\n",
" fout_height,\n",
" fout_width,\n",
" env.BATCH,\n",
" env.BLOCK_OUT,\n",
")\n",
"\n",
"# Convolution reduction axes\n",
"dy = te.reduce_axis((0, kernel_h), name=\"dy\")\n",
"dx = te.reduce_axis((0, kernel_w), name=\"dx\")\n",
"ic = te.reduce_axis((0, in_channels // env.BLOCK_IN), name=\"ic\")\n",
"ic_tns = te.reduce_axis((0, env.BLOCK_IN), name=\"ic_tns\")\n",
"\n",
"# Input placeholder tensors\n",
"data = te.placeholder(data_shape, name=\"data\", dtype=env.inp_dtype)\n",
"kernel = te.placeholder(kernel_shape, name=\"kernel\", dtype=env.wgt_dtype)\n",
"\n",
"# Copy buffers:\n",
"# Apply spatial padding to input feature map\n",
"data_buf = topi.nn.pad(data, [0, 0, pad_h, pad_w, 0, 0], name=\"data_buf\")\n",
"kernel_buf = te.compute(kernel_shape, lambda *i: kernel(*i), \"kernel_buf\")\n",
"\n",
"# Declare 2D convolution\n",
"res_conv = te.compute(\n",
" output_shape,\n",
" lambda bo, co, i, j, bi, ci: te.sum(\n",
" data_buf[bo, ic, i * stride_h + dy, j * stride_w + dx, bi, ic_tns].astype(env.acc_dtype)\n",
" * kernel_buf[co, ic, dy, dx, ci, ic_tns].astype(env.acc_dtype),\n",
" axis=[ic, dy, dx, ic_tns],\n",
" ),\n",
" name=\"res_conv\",\n",
")\n",
"\n",
"# Add shift stage for fix-point normalization\n",
"res_shr = te.compute(output_shape, lambda *i: res_conv(*i) >> 8, name=\"res_shr\")\n",
"\n",
"# Apply clipping between (0, input max value)\n",
"inp_max = (1 << (env.INP_WIDTH - 1)) - 1\n",
"res_max = te.compute(output_shape, lambda *i: tvm.te.max(res_shr(*i), 0), \"res_max\")\n",
"res_min = te.compute(output_shape, lambda *i: tvm.te.min(res_max(*i), inp_max), \"res_min\")\n",
"\n",
"# Result Tensor\n",
"res = te.compute(output_shape, lambda *i: res_min(*i).astype(env.inp_dtype), name=\"res\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 调度计算\n",
"\n",
"将看到一组必要的调度变换，以有效的方式将 2D 卷积映射到 VTA。这些包括：\n",
"\n",
"- 分块计算\n",
"- 增加计算利用率（compute utilization）的虚拟线程（Virtual threading）\n",
"- Lowering 到 VTA 硬件 intrinsics"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"collapsed": false
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"@main = primfn(data_1: handle, kernel_1: handle, res_1: handle) -> ()\n",
" attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
" buffers = {data: Buffer(data_2: Pointer(int8), int8, [50176], []),\n",
" kernel: Buffer(kernel_2: Pointer(int8), int8, [589824], []),\n",
" res: Buffer(res_2: Pointer(int8), int8, [50176], [])}\n",
" buffer_map = {data_1: data, kernel_1: kernel, res_1: res}\n",
" preflattened_buffer_map = {data_1: data_3: Buffer(data_2, int8, [1, 16, 14, 14, 1, 16], []), kernel_1: kernel_3: Buffer(kernel_2, int8, [16, 16, 3, 3, 16, 16], []), res_1: res_3: Buffer(res_2, int8, [1, 16, 14, 14, 1, 16], [])} {\n",
" allocate(data_buf: Pointer(global int8), int8, [65536]), storage_scope = global;\n",
" allocate(kernel_buf: Pointer(global int8), int8, [589824]), storage_scope = global;\n",
" allocate(res_conv: Pointer(global int32), int32, [50176]), storage_scope = global {\n",
" for (i1: int32, 0, 16) {\n",
" for (i2: int32, 0, 16) {\n",
" for (i3: int32, 0, 16) {\n",
" for (i5: int32, 0, 16) {\n",
" let cse_var_1: int32 = (i3*16)\n",
" data_buf_1: Buffer(data_buf, int8, [65536], [])[((((i1*4096) + (i2*256)) + cse_var_1) + i5)] = @tir.if_then_else(((((1 <= i2) && (i2 < 15)) && (1 <= i3)) && (i3 < 15)), data[(((((i1*3136) + (i2*224)) + cse_var_1) + i5) - 240)], 0i8, dtype=int8)\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (i0: int32, 0, 16) {\n",
" for (i1_1: int32, 0, 16) {\n",
" for (i2_1: int32, 0, 3) {\n",
" for (i3_1: int32, 0, 3) {\n",
" for (i4: int32, 0, 16) {\n",
" for (i5_1: int32, 0, 16) {\n",
" let cse_var_2: int32 = ((((((i0*36864) + (i1_1*2304)) + (i2_1*768)) + (i3_1*256)) + (i4*16)) + i5_1)\n",
" kernel_buf_1: Buffer(kernel_buf, int8, [589824], [])[cse_var_2] = kernel[cse_var_2]\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (co: int32, 0, 16) {\n",
" for (i: int32, 0, 14) {\n",
" for (j: int32, 0, 14) {\n",
" for (ci: int32, 0, 16) {\n",
" res_conv_1: Buffer(res_conv, int32, [50176], [])[((((co*3136) + (i*224)) + (j*16)) + ci)] = 0\n",
" for (ic: int32, 0, 16) {\n",
" for (dy: int32, 0, 3) {\n",
" for (dx: int32, 0, 3) {\n",
" for (ic_tns: int32, 0, 16) {\n",
" let cse_var_4: int32 = (j*16)\n",
" let cse_var_3: int32 = ((((co*3136) + (i*224)) + cse_var_4) + ci)\n",
" res_conv_1[cse_var_3] = (res_conv_1[cse_var_3] + (cast(int32, data_buf_1[((((((ic*4096) + (i*256)) + (dy*256)) + cse_var_4) + (dx*16)) + ic_tns)])*cast(int32, kernel_buf_1[((((((co*36864) + (ic*2304)) + (dy*768)) + (dx*256)) + (ci*16)) + ic_tns)])))\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (i1_2: int32, 0, 16) {\n",
" for (i2_2: int32, 0, 14) {\n",
" for (i3_2: int32, 0, 14) {\n",
" for (i5_2: int32, 0, 16) {\n",
" let cse_var_5: int32 = ((((i1_2*3136) + (i2_2*224)) + (i3_2*16)) + i5_2)\n",
" res_conv_2: Buffer(res_conv, int32, [50176], [])[cse_var_5] = @tir.shift_right(res_conv_1[cse_var_5], 8, dtype=int32)\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (i1_3: int32, 0, 16) {\n",
" for (i2_3: int32, 0, 14) {\n",
" for (i3_3: int32, 0, 14) {\n",
" for (i5_3: int32, 0, 16) {\n",
" let cse_var_6: int32 = ((((i1_3*3136) + (i2_3*224)) + (i3_3*16)) + i5_3)\n",
" res_conv_3: Buffer(res_conv, int32, [50176], [])[cse_var_6] = max(res_conv_2[cse_var_6], 0)\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (i1_4: int32, 0, 16) {\n",
" for (i2_4: int32, 0, 14) {\n",
" for (i3_4: int32, 0, 14) {\n",
" for (i5_4: int32, 0, 16) {\n",
" let cse_var_7: int32 = ((((i1_4*3136) + (i2_4*224)) + (i3_4*16)) + i5_4)\n",
" res_conv_4: Buffer(res_conv, int32, [50176], [])[cse_var_7] = min(res_conv_3[cse_var_7], 127)\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (i1_5: int32, 0, 16) {\n",
" for (i2_5: int32, 0, 14) {\n",
" for (i3_5: int32, 0, 14) {\n",
" for (i5_5: int32, 0, 16) {\n",
" let cse_var_8: int32 = ((((i1_5*3136) + (i2_5*224)) + (i3_5*16)) + i5_5)\n",
" res[cse_var_8] = cast(int8, res_conv_4[cse_var_8])\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
"}\n",
"\n",
"\n"
]
}
],
"source": [
"# Create TVM schedule\n",
"s = te.create_schedule(res.op)\n",
"# Let's look at the default TVM schedule\n",
"print(tvm.lower(s, [data, kernel, res], simple_mode=True))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 分块计算\n",
"\n",
"默认情况下，2D 卷积太大，激活或卷积核权重无法同时适应 VTA 的 on-chip buffer。沿着输入通道、输出通道和高度空间维度应用分块。不沿宽度空间维度进行分块，因为它是 NCHW 布局中的最内层维度（因此，为了增加局部性，最好不要沿最内层维度进行分块）。"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"# Let's define tiling sizes\n",
"b_block = 1 // env.BATCH\n",
"oc_block = 128 // env.BLOCK_OUT\n",
"ic_block = 16 // env.BLOCK_IN\n",
"h_block = 7\n",
"w_block = 14\n",
"\n",
"# Tile the output tensor along the spatial and output channel dimensions\n",
"# (since by default we are doing single batch inference, the split along\n",
"# the batch dimension has no effect)\n",
"b, oc, y, x, b_tns, oc_tns = s[res].op.axis\n",
"b_out, b_inn = s[res].split(b, factor=b_block)\n",
"oc_out, oc_inn = s[res].split(oc, factor=oc_block)\n",
"y_out, y_inn = s[res].split(y, factor=h_block)\n",
"x_out, x_inn = s[res].split(x, factor=w_block)\n",
"s[res].reorder(b_out, oc_out, y_out, x_out, b_inn, oc_inn, y_inn, x_inn, b_tns, oc_tns)\n",
"\n",
"# Move intermediate computation into each output compute tile\n",
"s[res_conv].compute_at(s[res], x_out)\n",
"s[res_shr].compute_at(s[res], x_out)\n",
"s[res_max].compute_at(s[res], x_out)\n",
"s[res_min].compute_at(s[res], x_out)\n",
"\n",
"# Apply additional loop split along reduction axis (input channel)\n",
"b_inn, oc_inn, y_inn, x_inn, b_tns, oc_tns = s[res_conv].op.axis\n",
"ic_out, ic_inn = s[res_conv].split(ic, factor=ic_block)\n",
"\n",
"# Reorder axes.\n",
"# 1) Group the VTA tensor axes in the inner most position: b_tns, oc_tns, ic_tns\n",
"# to allow TVM to tensorize.\n",
"# 2) We move the ic_out axis all the way out of the convolution loop to block\n",
"# along the reduction axis.\n",
"# 3) Now we re-order the block axes: b_inn, oc_inn, y_inn, x_inn, ic_inn, dy, dx.\n",
"# VTA runtime/hardware requires us to write to a different output feature map\n",
"# location for every VTA tensor operation.\n",
"# This restriction requires us to order one of oc_inn, y_inn or x_inn right\n",
"# before b_tns, since they all affect output feature map indexing.\n",
"# Therefore, we choose to bring x_inn inside as shown below.\n",
"s[res_conv].reorder(ic_out, b_inn, oc_inn, y_inn, ic_inn, dy, dx, x_inn, b_tns, oc_tns, ic_tns)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 虚拟线程\n",
"\n",
"虚拟线程是一种在 VTA 硬件设计中增加任务级管道并行性的机制。换句话说，它通过隐藏内存访问延迟（hiding memory access latency）提高了计算资源的利用率。\n",
"\n",
"在下面的实现中，虚拟线程将工作分配给沿输出通道轴划分的两个线程。在下面的图中，展示了计算 2D 卷积时工作是如何分割的。\n",
"\n",
"```{image} images/virtual_threading.png\n",
":align: center\n",
":width: 480px\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"collapsed": false
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"@main = primfn(data_1: handle, kernel_1: handle, res_1: handle) -> ()\n",
" attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
" buffers = {data: Buffer(data_2: Pointer(int8), int8, [50176], []),\n",
" kernel: Buffer(kernel_2: Pointer(int8), int8, [589824], []),\n",
" res: Buffer(res_2: Pointer(int8), int8, [50176], [])}\n",
" buffer_map = {data_1: data, kernel_1: kernel, res_1: res}\n",
" preflattened_buffer_map = {data_1: data_3: Buffer(data_2, int8, [1, 16, 14, 14, 1, 16], []), kernel_1: kernel_3: Buffer(kernel_2, int8, [16, 16, 3, 3, 16, 16], []), res_1: res_3: Buffer(res_2, int8, [1, 16, 14, 14, 1, 16], [])} {\n",
" allocate(data_buf: Pointer(global int8), int8, [65536]), storage_scope = global;\n",
" allocate(kernel_buf: Pointer(global int8), int8, [589824]), storage_scope = global;\n",
" allocate(res_conv: Pointer(global int32), int32, [25088]), storage_scope = global {\n",
" for (i1: int32, 0, 16) {\n",
" for (i2: int32, 0, 16) {\n",
" for (i3: int32, 0, 16) {\n",
" for (i5: int32, 0, 16) {\n",
" let cse_var_1: int32 = (i3*16)\n",
" data_buf_1: Buffer(data_buf, int8, [65536], [])[((((i1*4096) + (i2*256)) + cse_var_1) + i5)] = @tir.if_then_else(((((1 <= i2) && (i2 < 15)) && (1 <= i3)) && (i3 < 15)), data[(((((i1*3136) + (i2*224)) + cse_var_1) + i5) - 240)], 0i8, dtype=int8)\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (i0: int32, 0, 16) {\n",
" for (i1_1: int32, 0, 16) {\n",
" for (i2_1: int32, 0, 3) {\n",
" for (i3_1: int32, 0, 3) {\n",
" for (i4: int32, 0, 16) {\n",
" for (i5_1: int32, 0, 16) {\n",
" let cse_var_2: int32 = ((((((i0*36864) + (i1_1*2304)) + (i2_1*768)) + (i3_1*256)) + (i4*16)) + i5_1)\n",
" kernel_buf_1: Buffer(kernel_buf, int8, [589824], [])[cse_var_2] = kernel[cse_var_2]\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (i2.outer: int32, 0, 2) {\n",
" for (co.init: int32, 0, 8) {\n",
" for (i.init: int32, 0, 7) {\n",
" for (j.init: int32, 0, 14) {\n",
" for (ci.init: int32, 0, 16) {\n",
" let cse_var_3: int32 = ((((co.init*1568) + (i.init*224)) + (j.init*16)) + ci.init)\n",
" {\n",
" res_conv_1: Buffer(res_conv, int32, [157351936], [])[cse_var_3] = 0\n",
" res_conv_1[(cse_var_3 + 12544)] = 0\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (ic.outer: int32, 0, 16) {\n",
" for (co: int32, 0, 8) {\n",
" for (i: int32, 0, 7) {\n",
" for (dy: int32, 0, 3) {\n",
" for (dx: int32, 0, 3) {\n",
" for (j: int32, 0, 14) {\n",
" for (ci: int32, 0, 16) {\n",
" for (ic_tns: int32, 0, 16) {\n",
" let cse_var_8: int32 = (j*16)\n",
" let cse_var_7: int32 = ((((co*1568) + (i*224)) + cse_var_8) + ci)\n",
" let cse_var_6: int32 = (cse_var_7 + 12544)\n",
" let cse_var_5: int32 = ((((((co*36864) + (ic.outer*2304)) + (dy*768)) + (dx*256)) + (ci*16)) + ic_tns)\n",
" let cse_var_4: int32 = (((((((ic.outer*4096) + (i2.outer*1792)) + (i*256)) + (dy*256)) + cse_var_8) + (dx*16)) + ic_tns)\n",
" {\n",
" res_conv_1[cse_var_7] = (res_conv_1[cse_var_7] + (cast(int32, data_buf_1[cse_var_4])*cast(int32, kernel_buf_1[cse_var_5])))\n",
" res_conv_1[cse_var_6] = (res_conv_1[cse_var_6] + (cast(int32, data_buf_1[cse_var_4])*cast(int32, kernel_buf_1[(cse_var_5 + 294912)])))\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (i1_2: int32, 0, 8) {\n",
" for (i2_2: int32, 0, 7) {\n",
" for (i3_2: int32, 0, 14) {\n",
" for (i5_2: int32, 0, 16) {\n",
" let cse_var_10: int32 = ((((i1_2*1568) + (i2_2*224)) + (i3_2*16)) + i5_2)\n",
" let cse_var_9: int32 = (cse_var_10 + 12544)\n",
" {\n",
" res_conv_2: Buffer(res_conv, int32, [157351936], [])[cse_var_10] = @tir.shift_right(res_conv_1[cse_var_10], 8, dtype=int32)\n",
" res_conv_2[cse_var_9] = @tir.shift_right(res_conv_1[cse_var_9], 8, dtype=int32)\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (i1_3: int32, 0, 8) {\n",
" for (i2_3: int32, 0, 7) {\n",
" for (i3_3: int32, 0, 14) {\n",
" for (i5_3: int32, 0, 16) {\n",
" let cse_var_12: int32 = ((((i1_3*1568) + (i2_3*224)) + (i3_3*16)) + i5_3)\n",
" let cse_var_11: int32 = (cse_var_12 + 12544)\n",
" {\n",
" res_conv_3: Buffer(res_conv, int32, [157351936], [])[cse_var_12] = max(res_conv_2[cse_var_12], 0)\n",
" res_conv_3[cse_var_11] = max(res_conv_2[cse_var_11], 0)\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (i1_4: int32, 0, 8) {\n",
" for (i2_4: int32, 0, 7) {\n",
" for (i3_4: int32, 0, 14) {\n",
" for (i5_4: int32, 0, 16) {\n",
" let cse_var_14: int32 = ((((i1_4*1568) + (i2_4*224)) + (i3_4*16)) + i5_4)\n",
" let cse_var_13: int32 = (cse_var_14 + 12544)\n",
" {\n",
" res_conv_4: Buffer(res_conv, int32, [157351936], [])[cse_var_14] = min(res_conv_3[cse_var_14], 127)\n",
" res_conv_4[cse_var_13] = min(res_conv_3[cse_var_13], 127)\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" for (i1.inner: int32, 0, 8) {\n",
" for (i2.inner: int32, 0, 7) {\n",
" for (i3.inner: int32, 0, 14) {\n",
" for (i5_5: int32, 0, 16) {\n",
" let cse_var_18: int32 = (i2.inner*224)\n",
" let cse_var_17: int32 = (i3.inner*16)\n",
" let cse_var_16: int32 = ((((i1.inner*1568) + cse_var_18) + cse_var_17) + i5_5)\n",
" let cse_var_15: int32 = (((((i1.inner*3136) + (i2.outer*1568)) + cse_var_18) + cse_var_17) + i5_5)\n",
" {\n",
" res[cse_var_15] = cast(int8, res_conv_4[cse_var_16])\n",
" res[(cse_var_15 + 25088)] = cast(int8, res_conv_4[(cse_var_16 + 12544)])\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
" }\n",
"}\n",
"\n",
"\n"
]
}
],
"source": [
"# VTA only supports 2 virtual threads\n",
"v_threads = 2\n",
"\n",
"# Perform virtual thread split along output channel outer axis\n",
"_, tx = s[res].split(oc_out, factor=v_threads)\n",
"s[res].reorder(tx, b_out)\n",
"s[res].bind(tx, te.thread_axis(\"cthread\"))\n",
"\n",
"# Let's look at the current TVM schedule after blocking and virtual threading\n",
"print(tvm.lower(s, [data, kernel, res], simple_mode=True))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Lowering Copies 到 DMA Transfers\n",
"\n",
"接下来，设置相应的 on-chip VTA SRAM buffers 的 buffers 作用域。将 load 循环移动到 2D 卷积计算循环，以 stage 内存加载，以便它们适合 on-chip SRAM buffers。最后，用 DMA 复制 pragma 注解了 load/store 循环外轴，以便在 VTA 上执行大容量内存传输。"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"collapsed": false
},
"outputs": [],
"source": [
"# Set scope of SRAM buffers\n",
"s[data_buf].set_scope(env.inp_scope)\n",
"s[kernel_buf].set_scope(env.wgt_scope)\n",
"s[res_conv].set_scope(env.acc_scope)\n",
"s[res_shr].set_scope(env.acc_scope)\n",
"s[res_min].set_scope(env.acc_scope)\n",
"s[res_max].set_scope(env.acc_scope)\n",
"\n",
"# Block data and kernel cache reads\n",
"s[data_buf].compute_at(s[res_conv], ic_out)\n",
"s[kernel_buf].compute_at(s[res_conv], ic_out)\n",
"\n",
"# Use DMA copy pragma on DRAM->SRAM operations\n",
"s[data_buf].pragma(s[data_buf].op.axis[0], env.dma_copy)\n",
"s[kernel_buf].pragma(s[kernel_buf].op.axis[0], env.dma_copy)\n",
"\n",
"# Use DMA copy pragma on SRAM->DRAM operation in each result block\n",
"# (this implies that these copies should be performed along b_inn,\n",
"# or result axis 4)\n",
"s[res].pragma(s[res].op.axis[4], env.dma_copy)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Lowering 计算到 VTA 计算 Intrinsics\n",
"\n",
"最后阶段是通过将二维卷积映射为张量 intrinsics，并将位移和剪切计算映射为向量 ALU，从而将计算循环 lower 到 VTA 硬件 intrinsics。"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"collapsed": false
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"@main = primfn(data_1: handle, kernel_1: handle, res_1: handle) -> ()\n",
" attr = {\"from_legacy_te_schedule\": True, \"global_symbol\": \"main\", \"tir.noalias\": True}\n",
" buffers = {data: Buffer(data_2: Pointer(int8), int8, [50176], []),\n",
" kernel: Buffer(kernel_2: Pointer(int8), int8, [589824], []),\n",
" res: Buffer(res_2: Pointer(int8), int8, [50176], [])}\n",
" buffer_map = {data_1: data, kernel_1: kernel, res_1: res}\n",
" preflattened_buffer_map = {data_1: data_3: Buffer(data_2, int8, [1, 16, 14, 14, 1, 16], []), kernel_1: kernel_3: Buffer(kernel_2, int8, [16, 16, 3, 3, 16, 16], []), res_1: res_3: Buffer(res_2, int8, [1, 16, 14, 14, 1, 16], [])} {\n",
" @tir.vta.coproc_dep_push(3, 2, dtype=int32)\n",
" @tir.vta.coproc_dep_push(3, 2, dtype=int32)\n",
" for (i2.outer: int32, 0, 2) {\n",
" for (cthread.s: int32, 0, 2) {\n",
" attr [IterVar(vta: int32, (nullptr), \"ThreadIndex\", \"vta\")] \"coproc_scope\" = 2 {\n",
" @tir.vta.coproc_dep_pop(3, 2, dtype=int32)\n",
" attr [IterVar(vta, (nullptr), \"ThreadIndex\", \"vta\")] \"coproc_uop_scope\" = \"VTAPushGEMMOp\" {\n",
" @tir.call_extern(\"VTAUopLoopBegin\", 8, 98, 0, 0, dtype=int32)\n",
" @tir.call_extern(\"VTAUopLoopBegin\", 7, 14, 0, 0, dtype=int32)\n",
" for (j.init: int32, 0, 14) {\n",
" @tir.vta.uop_push(0, 1, ((cthread.s*784) + j.init), 0, 0, 0, 0, 0, dtype=int32)\n",
" }\n",
" @tir.call_extern(\"VTAUopLoopEnd\", dtype=int32)\n",
" @tir.call_extern(\"VTAUopLoopEnd\", dtype=int32)\n",
" }\n",
" @tir.vta.coproc_dep_push(2, 1, dtype=int32)\n",
" }\n",
" }\n",
" for (ic.outer: int32, 0, 16) {\n",
" let cse_var_6: int32 = (i2.outer*7)\n",
" let cse_var_5: int32 = (ic.outer*9)\n",
" let cse_var_4: int32 = max((1 - cse_var_6), 0)\n",
" let cse_var_3: int32 = max((cse_var_6 - 6), 0)\n",
" let cse_var_2: int32 = ((9 - cse_var_4) - cse_var_3)\n",
" let cse_var_1: int32 = ((((ic.outer*196) + (i2.outer*98)) + (cse_var_4*14)) - 14)\n",
" {\n",
" attr [IterVar(vta, (nullptr), \"ThreadIndex\", \"vta\")] \"coproc_scope\" = 1 {\n",
" @tir.vta.coproc_dep_pop(2, 1, dtype=int32)\n",
" @tir.call_extern(\"VTALoadBuffer2D\", @tir.tvm_thread_context(@tir.vta.command_handle(, dtype=handle), dtype=handle), data_2, cse_var_1, 14, cse_var_2, 14, 1, cse_var_4, 1, cse_var_3, 0, 2, dtype=int32)\n",
" @tir.call_extern(\"VTALoadBuffer2D\", @tir.tvm_thread_context(@tir.vta.command_handle(, dtype=handle), dtype=handle), kernel_2, cse_var_5, 9, 8, 144, 0, 0, 0, 0, 0, 1, dtype=int32)\n",
" @tir.vta.coproc_dep_push(1, 2, dtype=int32)\n",
" }\n",
" attr [IterVar(vta, (nullptr), \"ThreadIndex\", \"vta\")] \"coproc_scope\" = 1 {\n",
" @tir.vta.coproc_dep_pop(2, 1, dtype=int32)\n",
" @tir.call_extern(\"VTALoadBuffer2D\", @tir.tvm_thread_context(@tir.vta.command_handle(, dtype=handle), dtype=handle), data_2, cse_var_1, 14, cse_var_2, 14, 1, cse_var_4, 1, cse_var_3, 144, 2, dtype=int32)\n",
" @tir.call_extern(\"VTALoadBuffer2D\", @tir.tvm_thread_context(@tir.vta.command_handle(, dtype=handle), dtype=handle), kernel_2, (cse_var_5 + 1152), 9, 8, 144, 0, 0, 0, 0, 72, 1, dtype=int32)\n",
" @tir.vta.coproc_dep_push(1, 2, dtype=int32)\n",
" }\n",
" for (cthread.s_1: int32, 0, 2) {\n",
" attr [IterVar(vta, (nullptr), \"ThreadIndex\", \"vta\")] \"coproc_scope\" = 2 {\n",
" @tir.vta.coproc_dep_pop(1, 2, dtype=int32)\n",
" attr [IterVar(vta, (nullptr), \"ThreadIndex\", \"vta\")] \"coproc_uop_scope\" = \"VTAPushGEMMOp\" {\n",
" @tir.call_extern(\"VTAUopLoopBegin\", 8, 98, 0, 9, dtype=int32)\n",
" @tir.call_extern(\"VTAUopLoopBegin\", 7, 14, 16, 0, dtype=int32)\n",
" for (dy: int32, 0, 3) {\n",
" for (dx: int32, 0, 3) {\n",
" for (j: int32, 0, 14) {\n",
" @tir.vta.uop_push(0, 0, ((cthread.s_1*784) + j), ((((cthread.s_1*144) + (dy*16)) + j) + dx), (((cthread.s_1*72) + (dy*3)) + dx), 0, 0, 0, dtype=int32)\n",
" }\n",
" }\n",
" }\n",
" @tir.call_extern(\"VTAUopLoopEnd\", dtype=int32)\n",
" @tir.call_extern(\"VTAUopLoopEnd\", dtype=int32)\n",
" }\n",
" @tir.vta.coproc_dep_push(2, 1, dtype=int32)\n",
" }\n",
" }\n",
" }\n",
" }\n",
" @tir.vta.coproc_dep_pop(2, 1, dtype=int32)\n",
" @tir.vta.coproc_dep_pop(2, 1, dtype=int32)\n",
" for (cthread.s_2: int32, 0, 2) {\n",
" let cse_var_7: int32 = (cthread.s_2*784)\n",
" attr [IterVar(vta, (nullptr), \"ThreadIndex\", \"vta\")] \"coproc_scope\" = 2 {\n",
" attr [IterVar(vta, (nullptr), \"ThreadIndex\", \"vta\")] \"coproc_uop_scope\" = \"VTAPushALUOp\" {\n",
" @tir.call_extern(\"VTAUopLoopBegin\", 784, 1, 1, 0, dtype=int32)\n",
" @tir.vta.uop_push(1, 0, cse_var_7, cse_var_7, 0, 3, 1, 8, dtype=int32)\n",
" @tir.call_extern(\"VTAUopLoopEnd\", dtype=int32)\n",
" }\n",
" attr [IterVar(vta, (nullptr), \"ThreadIndex\", \"vta\")] \"coproc_uop_scope\" = \"VTAPushALUOp\" {\n",
" @tir.call_extern(\"VTAUopLoopBegin\", 784, 1, 1, 0, dtype=int32)\n",
" @tir.vta.uop_push(1, 0, cse_var_7, cse_var_7, 0, 1, 1, 0, dtype=int32)\n",
" @tir.call_extern(\"VTAUopLoopEnd\", dtype=int32)\n",
" }\n",
" attr [IterVar(vta, (nullptr), \"ThreadIndex\", \"vta\")] \"coproc_uop_scope\" = \"VTAPushALUOp\" {\n",
" @tir.call_extern(\"VTAUopLoopBegin\", 784, 1, 1, 0, dtype=int32)\n",
" @tir.vta.uop_push(1, 0, cse_var_7, cse_var_7, 0, 0, 1, 127, dtype=int32)\n",
" @tir.call_extern(\"VTAUopLoopEnd\", dtype=int32)\n",
" }\n",
" @tir.vta.coproc_dep_push(2, 3, dtype=int32)\n",
" }\n",
" }\n",
" for (cthread.s_3: int32, 0, 2) {\n",
" attr [IterVar(vta, (nullptr), \"ThreadIndex\", \"vta\")] \"coproc_scope\" = 3 {\n",
" @tir.vta.coproc_dep_pop(2, 3, dtype=int32)\n",
" for (i1.inner: int32, 0, 8) {\n",
" for (i2.inner: int32, 0, 7) {\n",
" for (i3.inner: int32, 0, 14) {\n",
" let cse_var_8: int32 = (i2.inner*14)\n",
" @tir.call_extern(\"VTAStoreBuffer2D\", @tir.tvm_thread_context(@tir.vta.command_handle(, dtype=handle), dtype=handle), ((((cthread.s_3*784) + (i1.inner*98)) + cse_var_8) + i3.inner), 4, res_2, (((((cthread.s_3*1568) + (i1.inner*196)) + (i2.outer*98)) + cse_var_8) + i3.inner), 1, 1, 1, dtype=int32)\n",
" }\n",
" }\n",
" }\n",
" @tir.vta.coproc_dep_push(3, 2, dtype=int32)\n",
" }\n",
" }\n",
" }\n",
" @tir.vta.coproc_dep_pop(3, 2, dtype=int32)\n",
" @tir.vta.coproc_dep_pop(3, 2, dtype=int32)\n",
" @tir.vta.coproc_sync(, dtype=int32)\n",
"}\n",
"\n",
"\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"[16:02:31] /media/pc/data/4tb/lxw/books/tvm/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement required_alignment=256, provided_alignment=128\n"
]
}
],
"source": [
"# Apply tensorization over the batch tensor tile axis\n",
"s[res_conv].tensorize(b_tns, env.gemm)\n",
"\n",
"# Add an ALU pragma over the shift and clipping operations\n",
"s[res_shr].pragma(s[res_shr].op.axis[0], env.alu)\n",
"s[res_min].pragma(s[res_min].op.axis[0], env.alu)\n",
"s[res_max].pragma(s[res_max].op.axis[0], env.alu)\n",
"\n",
"# Let's look at the final lowered TVM schedule after lowering memory\n",
"# loads/stores down to DMA copy intrinsics, and the computation down to\n",
"# VTA compute intrinsics.\n",
"print(vta.lower(s, [data, kernel, res], simple_mode=True))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## TVM 计算和验证\n",
"\n",
"在指定调度之后，可以将其编译为 TVM 函数。保存模块，这样就可以通过 RPC 发送它。运行该函数并对 numpy 实现进行验证，以确保其正确性。"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"collapsed": false
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[16:02:32] /media/pc/data/4tb/lxw/books/tvm/src/tir/transforms/arg_binder.cc:95: Warning: Trying to bind buffer to another one with lower alignment requirement required_alignment=256, provided_alignment=128\n",
"/media/workspace/anaconda3/envs/mx/lib/python3.10/site-packages/tvm/driver/build_module.py:263: UserWarning: target_host parameter is going to be deprecated. Please pass in tvm.target.Target(target, host=target_host) instead.\n",
" warnings.warn(\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Execution statistics:\n",
"\tinp_load_nbytes : 114688\n",
"\twgt_load_nbytes : 1179648\n",
"\tacc_load_nbytes : 0\n",
"\tuop_load_nbytes : 1144\n",
"\tout_store_nbytes: 50176\n",
"\tgemm_counter : 451584\n",
"\talu_counter : 9408\n",
"Successful 2D convolution test!\n"
]
}
],
"source": [
"# This library facilitates 2D convolution testing\n",
"from tvm.topi.testing import conv2d_nchw_python\n",
"\n",
"# Compile the TVM module\n",
"with vta.build_config(disabled_pass={\"tir.CommonSubexprElimTIR\"}):\n",
" my_conv = vta.build(\n",
" s, [data, kernel, res], tvm.target.Target(\"ext_dev\", host=env.target_host), name=\"my_conv\"\n",
" )\n",
"temp = utils.tempdir()\n",
"my_conv.save(temp.relpath(\"conv2d.o\"))\n",
"remote.upload(temp.relpath(\"conv2d.o\"))\n",
"f = remote.load_module(\"conv2d.o\")\n",
"\n",
"# Get the remote device context\n",
"ctx = remote.ext_dev(0)\n",
"\n",
"# Initialize the data and kernel arrays randomly in the int range\n",
"# of (-128, 128] in NCHW layout\n",
"data_np = np.random.randint(-128, 128, size=(batch_size, in_channels, height, width)).astype(\n",
" data.dtype\n",
")\n",
"kernel_np = np.random.randint(\n",
" -128, 128, size=(out_channels, in_channels, kernel_h, kernel_w)\n",
").astype(kernel.dtype)\n",
"\n",
"# Apply packing to the data and kernel arrays from a 2D NCHW\n",
"# to a 4D NCHWnc packed layout\n",
"data_packed = data_np.reshape(\n",
" batch_size // env.BATCH, env.BATCH, in_channels // env.BLOCK_IN, env.BLOCK_IN, height, width\n",
").transpose((0, 2, 4, 5, 1, 3))\n",
"\n",
"kernel_packed = kernel_np.reshape(\n",
" out_channels // env.BLOCK_OUT,\n",
" env.BLOCK_OUT,\n",
" in_channels // env.BLOCK_IN,\n",
" env.BLOCK_IN,\n",
" kernel_h,\n",
" kernel_w,\n",
").transpose((0, 2, 4, 5, 1, 3))\n",
"\n",
"# Format the input/output arrays with tvm.nd.array to the DLPack standard\n",
"data_nd = tvm.nd.array(data_packed, ctx)\n",
"kernel_nd = tvm.nd.array(kernel_packed, ctx)\n",
"res_nd = tvm.nd.array(np.zeros(output_shape).astype(res.dtype), ctx)\n",
"\n",
"# Clear stats\n",
"if env.TARGET in [\"sim\", \"tsim\"]:\n",
" simulator.clear_stats()\n",
"\n",
"# Invoke the module to perform the computation\n",
"f(data_nd, kernel_nd, res_nd)\n",
"\n",
"# Verify against numpy implementation\n",
"res_ref = conv2d_nchw_python(\n",
" data_np.astype(env.acc_dtype),\n",
" kernel_np.astype(env.acc_dtype),\n",
" (stride_h, stride_w),\n",
" (pad_h, pad_w),\n",
").astype(env.acc_dtype)\n",
"res_ref = res_ref >> env.INP_WIDTH\n",
"res_ref = np.clip(res_ref, 0, inp_max)\n",
"res_ref = res_ref.astype(res.dtype)\n",
"res_ref = res_ref.reshape(\n",
" (\n",
" batch_size // env.BATCH,\n",
" env.BATCH,\n",
" out_channels // env.BLOCK_OUT,\n",
" env.BLOCK_OUT,\n",
" fout_height,\n",
" fout_width,\n",
" )\n",
").transpose((0, 2, 4, 5, 1, 3))\n",
"tvm.testing.assert_allclose(res_ref, res_nd.numpy())\n",
"\n",
"# Print stats\n",
"if env.TARGET in [\"sim\", \"tsim\"]:\n",
" sim_stats = simulator.stats()\n",
" print(\"Execution statistics:\")\n",
" for k, v in sim_stats.items():\n",
" print(\"\\t{:<16}: {:>16}\".format(k, v))\n",
"\n",
"print(\"Successful 2D convolution test!\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 小结\n",
"\n",
"本教程演示如何使用 TVM 调度原语 lower 硬件加速器 intrinsics 的 2D 卷积，利用特定于硬件的优化，比如使用带虚拟线程的隐藏延迟。"
]
}
],
"metadata": {
"interpreter": {
"hash": "ee40e4cbda3c4716866f133b45765e0887afdbc9aa3bd872ab229f889d521355"
},
"kernelspec": {
"display_name": "Python 3.10.4 ('mx')",
"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.10.4"
}
},
"nbformat": 4,
"nbformat_minor": 0
}