【MLC】 TensorIR 练习

前言

• 这两天重新看了一下天奇的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_func
      def 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.mod
      
  mod_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 + a
  
  print(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_func
    def 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})
      # TODO
      for 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 tvm
  
  N, 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_func
      def 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})
          # TODO
          mat_w = T.alloc_buffer((2, 9), "int64")
          # (8 - 3)/1 + 1 = 6
          mat_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_im2col
          for 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]
          
          # matmul
          for 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 IPython
  
  dtype = "float32"
  
  @tvm.script.ir_module
  class MyBmmRelu:
      @T.prim_func
      def 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 j
      block_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 cj1
      sch.compute_at(block_Y, cj0, preserve_unit_loops=False)
      sch.vectorize(cj1)
      
      # parallel n
      block_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_update
      sch.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 IPython
  
  dtype = "float32"
  
  @tvm.script.ir_module
  class MyBmmRelu:
      @T.prim_func
      def 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 j
      block_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 cj1
      sch.compute_at(block_Y, cj0, preserve_unit_loops=False)
      sch.vectorize(cj1)
      
      # parallel n
      block_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_update
      sch.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中的设计模式