【MLC】 TensorIR 练习

news/2024/10/18 9:24:47/

文章目录

    • 前言
    • TensorIR 练习
      • TensorIR: 张量程序抽象案例研究
      • 练习 1：广播加法
      • 练习 2：二维卷积
      • 练习 3：变换批量矩阵乘法程序
    • 总结

前言

  • 这两天重新看了一下天奇的mlc课程文档，把里边儿的TensorIR 练习写了一下，顺便推广一下相关资料
    • MLC-机器学习编译-B站，MLC-文档
    • bbuf老哥-B站
    • Archer哥的二三事
  • 有关于TVM其实我是跟着二三事入门的，然而当前TVM学习的难点不在于资料不够而是不够集成，初学者往往面出现对一堆资料无从下手的情况，此时需要有人指出一条学习路线，才能更好地利用把把散落在各地的资料解决学习途中的困惑，二三事就是这么个路子，从前端讲到后端，由浅入深，结合其他资料以及GPT(读代码很方便)，往往快速理解TVM各个阶段的操作与设计。
  • 话说回来MLC这门课程是讲AI编译原理，个人认为最大的优点是用python作为示例demo（文档比较完善，代码可以直接跑），加上天奇讲的很清楚，所以很适合作为入门。
  • bbuf是宝藏，dddd~

TensorIR 练习

  • 2.5. TensorIR 练习
  • 练习文档如上所示，在练习前建议复习一下2.4 TensorIR: 张量程序抽象案例研究，下面内容简单分为两部分，首先会简单总结2.4的例子，然后会贴上2.5的过程，结合实现回忆做一个记录。

TensorIR: 张量程序抽象案例研究

  • 2.4用的例子是一个matmul后边接一个relu这么个小计算图，首先会给出ir_module的代码，然后对其中用到的语法进行解释，这玩意主要是构造了一个计算的过程，表示C = Relu(A*B)。
    import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T@tvm.script.ir_module
class MyModule:@T.prim_funcdef mm_relu(A: T.Buffer((128, 128), "float32"),B: T.Buffer((128, 128), "float32"),C: T.Buffer((128, 128), "float32")):T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})Y = T.alloc_buffer((128, 128), dtype="float32")for i, j, k in T.grid(128, 128, 128):with T.block("Y"):vi = T.axis.spatial(128, i)vj = T.axis.spatial(128, j)vk = T.axis.reduce(128, k)with T.init():Y[vi, vj] = T.float32(0)Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]for i, j in T.grid(128, 128):with T.block("C"):vi = T.axis.spatial(128, i)vj = T.axis.spatial(128, j)C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
def transform(mod, jfactor):sch = tvm.tir.Schedule(mod)block_Y = sch.get_block("Y", func_name="mm_relu")i, j, k = sch.get_loops(block_Y)j0, j1 = sch.split(j, factors=[None, jfactor])sch.reorder(j0, k, j1)block_C = sch.get_block("C", "mm_relu")sch.reverse_compute_at(block_C, j0)return sch.modmod_transformed = transform(MyModule, jfactor=12)
rt_lib_transformed = tvm.build(mod_transformed, "llvm")
f_timer = rt_lib.time_evaluator("mm_relu", tvm.cpu())
f_timer_transformed = rt_lib_transformed.time_evaluator("mm_relu", tvm.cpu())
print("Time cost of mod %g sec" % f_timer(a_nd, b_nd, c_nd).mean)
print("Time cost of transformed mod_transformed %g sec" % f_timer_transformed(a_nd, b_nd, c_nd).mean)
# display the code below
# print(IPython.display.Code(mod_transformed.script(), language="python"))
  • 如何写一个T.prim_func(函数可用链接定位到源码)：
    • 参数里边儿定义了输入输出，类型是T.Buffer，可以用T.alloc_buffer申请
    • T.func_attr作用是PrimFuncFrame属性的赋值
    • T.grid，是一个构造多重循环(跟多用几个range一个效果)的语法糖
    • T.block，构造了一个BlockFrame，可以用get_block拿到，继续往下探索可以发现一个设计：
      class ConcreteScheduleNode : public ScheduleNode {friend class Schedule;friend class ScheduleCopier;......BlockRV GetBlock(const String& name, const Optional<String>& func_name) override;......};
      • 关于这个东西问了下GPT，得到的解答是：对外提供给用户的主要是通过Schedule类来操作和访问调度信息。Schedule类封装了ScheduleNode中的成员函数，并提供了更高级别的接口，使用户能够方便地进行调度和优化操作。
      • 类似的设计在TVM中很常见，挖个坑。
    • T.axis.remap也是个比较好用的语法糖，可以快速绑定块轴，轴类型有spatial，reduce代表空间轴和归约轴，理解为有在输出对应的轴是空间轴，没有对应的，会被归约减少的轴是归约轴
    • vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)
# 等价于
# SSR means the properties of each axes are "spatial", "spatial", "reduce"
# 计算块在 Y 的空间位置 (Y[vi, vj]) 处生成一个点值，该点值独立于 Y 中的其他位置（具有不同的vi, vj 值的位置）。我们可以称 vi、vj 为空间轴，因为它们直接对应于块写入的缓冲区空间区域的开始。 涉及归约的轴（vk）被命名为归约轴。
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
  • 至此，一个T.prim_func就可以写出来了，transform可以看这个：python/tvm/tir/schedule/schedule.py

练习 1：广播加法

  • 题目：请编写一个 TensorIR 函数，将两个数组以广播的方式相加。
  • 练习1就是练一个加法，这很简单(仅仅针对用例的输入来说)
  • 首先用numpy写个模拟版本，lnumpy_tadd，然而我在调用numpy本身方法的时候发现这个东西其实满足交换律，按道理也应该满足最后一维相等这类的条件，但这里不是很清楚TVM中如何判断这种情况(挖坑+1)，于是干脆在后边就先不考虑这个了，直接按照输入先写。
  • 然后是T.prim_func，用了T.grid和T.axis.remap这种语法糖（见上文）三两行就表示完了
    from tvm.script import tir as T
import tvm
import numpy as np
# 广播加法# init data
a = np.arange(16).reshape(-1, 4)
b = np.arange(4, 0, -1).reshape(4)
# numpy version
c_np = a + b
c_np2 = b + aprint(c_np)
print(c_np2)def lnumpy_tadd(A:np.ndarray, B:np.ndarray, C:np.ndarray):# 需要保证A是(x, y)或者(y) B是(1, y)或者(y)的shape才行# assert AB的合法性# 最后一维是相等的，有一个数必须是维度assert A.shape[-1] == B.shape[-1], "shape err!"assert A.shape == C.shape, "shape err!"B = B.reshape(B.shape[-1])for i in range(A.shape[0]):for j in range(B.shape[-1]):C[i][j] = A[i][j] + B[j]lnumpy_tadd_c = np.empty_like(a)
lnumpy_tadd(a, b, lnumpy_tadd_c)
print(lnumpy_tadd_c)
@tvm.script.ir_module
class MyAdd:@T.prim_funcdef add(A: T.Buffer((4, 4), "int64"),B: T.Buffer((4), "int64"),C: T.Buffer((4, 4), "int64")):T.func_attr({"global_symbol": "add", "tir.noalias": True})# TODOfor i, j in T.grid(4, 4):with T.block("C"):vi, vj = T.axis.remap("SS", [i,j])C[vi, vj] = A[vi, vj] + B[vj]rt_lib = tvm.build(MyAdd, target="llvm")
a_tvm = tvm.nd.array(a)
b_tvm = tvm.nd.array(b)
c_tvm = tvm.nd.array(np.empty((4, 4), dtype=np.int64))
rt_lib["add"](a_tvm, b_tvm, c_tvm)
np.testing.assert_allclose(c_tvm.numpy(), c_np, rtol=1e-5)

练习 2：二维卷积

  • 题目：
    在这里插入图片描述
  • 对于卷积计算想必都或多或少有些了解，常用的实现是im2col+gemm，其中gemm可以被tvm中的调度优化替代，因此主要实现im2col和col2im基本就好了（当然本题可以用原始定义写，我想作者原意可能也是用定义写，不过anyway，都一样）
  • 这份代码顺便复习了一下卷积的实现，缺点就是代码通用性不够，还是偷懒了
    # torch version
import torch
import numpy as np
from tvm.script import tir as T
import tvmN, CI, H, W, CO, K = 1, 1, 8, 8, 2, 3
OUT_H, OUT_W = H - K + 1, W - K + 1
data = np.arange(N*CI*H*W).reshape(N, CI, H, W)
weight = np.arange(CO*CI*K*K).reshape(CO, CI, K, K)
data_torch = torch.Tensor(data)
weight_torch = torch.Tensor(weight)
conv_torch = torch.nn.functional.conv2d(data_torch, weight_torch)
conv_torch = conv_torch.numpy().astype(np.int64)
print(conv_torch)print(data.shape)
print(weight.shape)
print(conv_torch.shape)@tvm.script.ir_module
class MyConv:@T.prim_funcdef conv(data:T.buffer((1, 1, 8, 8), "int64"),weight:T.buffer((2, 1, 3, 3), "int64"),output:T.buffer((1, 2, 6, 6), "int64")):T.func_attr({"global_symbol": "conv", "tir.noalias": True})# TODOmat_w = T.alloc_buffer((2, 9), "int64")# (8 - 3)/1 + 1 = 6mat_d = T.alloc_buffer((1, 6*6, 9), "int64")mat_o = T.alloc_buffer((1, 36, 2), "int64")# data_im2col:for i, sx, sy, c, w, h in T.grid(1, 6, 6, 1, 3, 3):with T.block("d_im2col"):vi, vx, vy, vc, vw, vh = T.axis.remap("SSSSSS", [i, sx, sy, c, w, h])with T.init():mat_d[vi, vx*6+vy, vc*3*3 + vw*3 + vh] = T.int64(0)mat_d[vi, vx*6+vy, vc*3*3 + vw*3 + vh] = data[vi, vc, vx+vw, vy+vh]# weight_im2colfor i, j, w, h in T.grid(2, 1, 3, 3):with T.block("w_im2col"):vi, vj, vw, vh = T.axis.remap("SSSS", [i, j, w, h])with T.init():mat_w[vi, vj*3*3 + w*3 + h] = T.int64(0)mat_w[vi, vj*3*3 + w*3 + h] = weight[vi, vj, vw, vh]# matmulfor b, h, i, j in T.grid(1, 36, 2, 9):with T.block("matmul"):vb, vh, vi, vj = T.axis.remap("SSSR", [b, h, i, j])with T.init():mat_o[vb, vh, vi] = T.int64(0)mat_o[vb, vh, vi] = mat_o[vb, vh, vi] + mat_d[vb, vh, vj] * mat_w[vi, vj]# col2img:for b, c, w, h in T.grid(1, 2, 6, 6):with T.block("col2img"):vb, vc, vw, vh = T.axis.remap("SSSS", [b, c, w, h])output[vb, vc, vw, vh] = mat_o[vb, vw*6+vh, vc]rt_lib = tvm.build(MyConv, target="llvm")
data_tvm = tvm.nd.array(data)
weight_tvm = tvm.nd.array(weight)
conv_tvm = tvm.nd.array(np.empty((N, CO, OUT_H, OUT_W), dtype=np.int64))
rt_lib["conv"](data_tvm, weight_tvm, conv_tvm)
np.testing.assert_allclose(conv_tvm.numpy(), conv_torch, rtol=1e-5)
print(conv_tvm)f_timer = rt_lib.time_evaluator("conv", tvm.cpu())
print("Time cost of conv %g sec" % f_timer(data_tvm, weight_tvm, conv_tvm).mean)

练习 3：变换批量矩阵乘法程序

  • 题目
    在这里插入图片描述

  • 我们回到mm_relu，这次需要实现的是bmm_relu，是多batch版的mm_relu，这很简单，多加一个轴就行

    import numpy as np
import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
import IPythondtype = "float32"@tvm.script.ir_module
class MyBmmRelu:@T.prim_funcdef bmm_relu(A: T.Buffer((16, 128, 128), dtype),B: T.Buffer((16, 128, 128), dtype),C: T.Buffer((16, 128, 128), dtype),):# 这里的 global_symbol 对应函数名，tir.noalias 是一个属性，表示所有的缓冲存储器不重叠。T.func_attr({"global_symbol": "bmm_relu", "tir.noalias": True})Y = T.alloc_buffer((16, 128, 128), dtype)for n, i, j, k in T.grid(16, 128, 128, 128):with T.block("Y"):# SSR means the properties of each axes are "spatial", "spatial", "reduce"vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k])with T.init():Y[vn, vi, vj] = T.float32(0)Y[vn, vi, vj] = Y[vn, vi, vj] + A[vn, vi, vk] * B[vn, vk, vj]for n, i, j in T.grid(16, 128, 128):with T.block("C"):vn, vi, vj = T.axis.remap("SSS", [n, i, j])C[vn, vi, vj] = T.max(Y[vn, vi, vj], T.float32(0))

  • 这份代码出来的程序是这个：
    在这里插入图片描述

  • 我们的目标程序是：
    在这里插入图片描述

  • 分析一下目标程序，原始的块轴是(16, 128, 128, 128)，我们需要将第一个batch的16抽出来合并并设置为T.parallel，第二维128的i抽出来合并，第三维128拆分成168，block(“Y”)中的最后一维需要拆分成324，剩下就是一些T.vectorized，T.unroll以及reorder等操作。各种转换方法的说明可以直接看源代码，注释写的很清楚了：python/tvm/tir/schedule/schedule.py

  • 于是我的思路是这样的

    • 先把两个block中的j轴拆成16*8，然后将拆出来的cj0和j0以外的循环合并成一个循环，并按照target设置vectorize和parallel，
    • 接下来就是拆分block Y，目的是吧init拆出来，这用到了sch.decompose_reduction，他会拆分出Y_init和Y_update
    • 最后对init部分设置T.vectorized，对update部分拆分32*4，reorder+unroll一下即可得到相似的结果：
    def transform(mod):sch = tvm.tir.Schedule(mod)# split jblock_Y = sch.get_block("Y", func_name="bmm_relu")n, i, j, k = sch.get_loops(block_Y)j0, j1 = sch.split(j, factors=[None, 8])block_C = sch.get_block("C", func_name="bmm_relu")cn, ci, cj = sch.get_loops(block_C)cj0, cj1 = sch.split(cj, factors=[None, 8])# compute_at cj0 && vectorize cj1sch.compute_at(block_Y, cj0, preserve_unit_loops=False)sch.vectorize(cj1)# parallel nblock_Y = sch.get_block("Y", func_name="bmm_relu")n, i, j0, j1, k = sch.get_loops(block_Y)sch.parallel(n)# # split Y_init Y_updatesch.decompose_reduction(block_Y, j1)block_Y_init = sch.get_block("Y_init", func_name="bmm_relu")n, i, j_0, ax0_init = sch.get_loops(block_Y_init)sch.vectorize(ax0_init)block_Y_update = sch.get_block("Y_update", func_name="bmm_relu")n, i, j0, jax0, jax1 = sch.get_loops(block_Y_update)ax1_0, ax1_1 = sch.split(jax1, factors=[None, 4])sch.reorder(ax1_0, ax1_1, jax0)sch.unroll(ax1_1)return sch.mod

  • 转换完之后的程序如下所示，可以看到已经十分接近了：
    在这里插入图片描述

  • 最后展示一下优化之后的时间测试，这里并没有尝试更多的优化方法，主要是想把原始程序变为目标程序：
    在这里插入图片描述

  • 贴一下这个作业的完整代码：

    import numpy as np
import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
import IPythondtype = "float32"@tvm.script.ir_module
class MyBmmRelu:@T.prim_funcdef bmm_relu(A: T.Buffer((16, 128, 128), dtype),B: T.Buffer((16, 128, 128), dtype),C: T.Buffer((16, 128, 128), dtype),):# 这里的 global_symbol 对应函数名，tir.noalias 是一个属性，表示所有的缓冲存储器不重叠。T.func_attr({"global_symbol": "bmm_relu", "tir.noalias": True})Y = T.alloc_buffer((16, 128, 128), dtype)for n, i, j, k in T.grid(16, 128, 128, 128):with T.block("Y"):# SSR means the properties of each axes are "spatial", "spatial", "reduce"vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k])with T.init():Y[vn, vi, vj] = T.float32(0)Y[vn, vi, vj] = Y[vn, vi, vj] + A[vn, vi, vk] * B[vn, vk, vj]for n, i, j in T.grid(16, 128, 128):with T.block("C"):vn, vi, vj = T.axis.remap("SSS", [n, i, j])C[vn, vi, vj] = T.max(Y[vn, vi, vj], T.float32(0))def transform(mod):sch = tvm.tir.Schedule(mod)# split jblock_Y = sch.get_block("Y", func_name="bmm_relu")n, i, j, k = sch.get_loops(block_Y)j0, j1 = sch.split(j, factors=[None, 8])block_C = sch.get_block("C", func_name="bmm_relu")cn, ci, cj = sch.get_loops(block_C)cj0, cj1 = sch.split(cj, factors=[None, 8])# compute_at j0, vectorize cj1sch.compute_at(block_Y, cj0, preserve_unit_loops=False)sch.vectorize(cj1)# parallel nblock_Y = sch.get_block("Y", func_name="bmm_relu")n, i, j0, j1, k = sch.get_loops(block_Y)sch.parallel(n)# # split Y_init Y_updatesch.decompose_reduction(block_Y, j1)block_Y_init = sch.get_block("Y_init", func_name="bmm_relu")n, i, j_0, ax0_init = sch.get_loops(block_Y_init)sch.vectorize(ax0_init)block_Y_update = sch.get_block("Y_update", func_name="bmm_relu")n, i, j0, jax0, jax1 = sch.get_loops(block_Y_update)ax1_0, ax1_1 = sch.split(jax1, factors=[None, 4])sch.reorder(ax1_0, ax1_1, jax0)sch.unroll(ax1_1)return sch.mod# organize the loops
sch = tvm.tir.Schedule(MyBmmRelu)
print(IPython.display.Code(sch.mod.script(), language="python"))
mod = transform(MyBmmRelu)
sch = tvm.tir.Schedule(mod)
print(IPython.display.Code(sch.mod.script(), language="python"))# test data
a_tvm = tvm.nd.array(np.random.rand(16, 128, 128).astype("float32"))
b_tvm = tvm.nd.array(np.random.rand(16, 128, 128).astype("float32"))
c_tvm = tvm.nd.array(np.random.rand(16, 128, 128).astype("float32"))# runtime
before_rt_lib = tvm.build(MyBmmRelu, target="llvm")
after_rt_lib = tvm.build(sch.mod, target="llvm")# after_rt_lib["bmm_relu"](a_tvm, b_tvm, c_tvm)# time_evaluator
before_timer = before_rt_lib.time_evaluator("bmm_relu", tvm.cpu())
print("Before transformation:")
print(before_timer(a_tvm, b_tvm, c_tvm))f_timer = after_rt_lib.time_evaluator("bmm_relu", tvm.cpu())
print("After transformation:")
print(f_timer(a_tvm, b_tvm, c_tvm))

    总结

    • 练习了一下tir怎么手动构造，以及sch的transform，我个人认为这个玩意实际业务上不会用到，但是得会看（不然你报错都不知道怎么看），所以这个基础练习还是有点必要看看。
    • 接下来的计划：
      • 1.总结TVM的总体链路
      • 2.对比TVM与竞品的区别、优劣势
      • 3.分析TVM C++与Python相互调用的机制
      • 4.尝试记录Pass的构造过程，如何新增pass
      • 5.尝试记录runtime codegen的过程
      • 6.分析TVM中的设计模式

http://www.ppmy.cn/news/101852.html

相关文章

中低压母线室弧光保护装置在水电站的应用

中低压母线室弧光保护装置在水电站的应用

摘要&#xff1a;本文介绍了电弧光保护在水电站的配置及应用&#xff0c;提供给相关人员参考。 关键词&#xff1a;水电站&#xff1b;开关柜&#xff1b;电弧光 0前言 电弧光是由于发生相间短路或接地短路时空气电离而形成的。在我国电力系统中开关柜内部电弧光故障时有发生…
阅读更多...
华为OD机试真题 Java 实现【基站维修工程师】【2023Q1 200分】，附详细解题思路

华为OD机试真题 Java 实现【基站维修工程师】【2023Q1 200分】，附详细解题思路

一、题目描述 小王是一名基站维护工程师，负责某区域的基站维护。 某地方有n个基站（1<n<10），已知各基站之间的距离s（0<s<500），并且基站x到基站y的距离，与基站y到基站x的距离并不一定会相同。 小王从基站1出发，途径每个基站1次，然后返回基站1，需要请你…
阅读更多...
【SonarQube】下载、安装、配置、使用介绍

【SonarQube】下载、安装、配置、使用介绍

文章目录 SonarQube安装运行使用root启动问题处理修改文件数限制JDK版本问题创建Project创建token扫描代码数据持久化在线文档 SonarQube安装 官网下载地址: http://www.sonarqube.org/downloads/9.9.1.69595下载地址: https://binaries.sonarsource.com/Distribution/sonarqu…
阅读更多...
【大数据处理与可视化】六、数据可视化

【大数据处理与可视化】六、数据可视化

【大数据处理与可视化】六、数据可视化 实验目的实验内容实验步骤一、案例——画图分析某年旅游景点数据1、河北省总面积和游客量位居前三的景点2、河北省旅游量的占比哪个最多&#xff0c;哪个最少。 实验小结 实验目的 1.能够详述常见图表的类型和特点。 2.能够熟练运用Matp…
阅读更多...
携手高通，移远通信以全栈式车载产品实力重新定义汽车

携手高通，移远通信以全栈式车载产品实力重新定义汽车

5月25日至26日&#xff0c;2023高通汽车技术与合作峰会在苏州成功举办。 移远通信作为高通重要的长期战略合作伙伴&#xff0c;受邀参加此次峰会&#xff0c;并现场展示了支持5G/4G、C-V2X、算力、UWB、Wi-Fi、高精定位等领先技术的多款车规级模组、天线等产品&#xff0c;更有…
阅读更多...
cesium EntityCollection和EntityView的关系

cesium EntityCollection和EntityView的关系

Cesium中的EntityCollection和EntityView是两个紧密相关的概念&#xff0c;它们之间存在着重要的关系。 EntityCollection是Cesium中用于管理Entity对象的容器&#xff0c;它提供了一系列方法和属性&#xff0c;用于添加、移除、更新和查询Entity对象。当您向EntityCollection…
阅读更多...
连续通过式清洗机 头部带磁选的皮带机 热油烟气风机轮毂 单主梁门式起重机 小H型钢轧机 环链斗式提升机 U型螺旋输送机…CAD

连续通过式清洗机 头部带磁选的皮带机 热油烟气风机轮毂 单主梁门式起重机 小H型钢轧机 环链斗式提升机 U型螺旋输送机…CAD

连续通过式清洗机头部带磁选的皮带机热油烟气风机轮毂链板机&#xff08;图纸资料&#xff09;滚轮MDG30-5t-17m-a5mdg型单主梁门式起重机链板机&#xff08;图纸资料&#xff09;电机及减速机操作平台总图拐弯皮带机螺旋机A1链板机&#xff08;图纸资料&#xff09;DT01P传动装…
阅读更多...
跳跃游戏（贪心算法）

跳跃游戏（贪心算法）

题目&#xff1a; 给定一个非负整数数组 nums &#xff0c;你最初位于数组的 第一个下标 。 数组中的每个元素代表你在该位置可以跳跃的最大长度。 判断你是否能够到达最后一个下标。 输入&#xff1a;nums [2,3,1,1,4] 输出&#xff1a;true 解释&#xff1a;可以先跳 1 步…
阅读更多...
最新文章

Copyright @ 2022~2023