TVM Compiler中文教程：TVM使用张量化Tensorize利用硬件内联函数

ゞ浴缸里的玫瑰 2022-01-19 03:05 291阅读 0赞

### 文章目录 ###

*  TVM使用Tensorize利用硬件内联函数
 *   *  定义矩阵乘法
     *  调度矩阵乘法
     *  定义GEMV张量内联函数
     *  为张量化进行Reduce-update
     *  总结

# TVM使用Tensorize利用硬件内联函数 #

**张量化计算优化这篇教程没看太懂。**  
这篇教程是关于在TVM中如何执行`张量化`的介绍。

通过使用调度原语`tensorize`，人们可以用相应的内联函数替换计算单元，从而可以轻松利用手工制作的微内核函数，和扩展TVM来支持新的硬件架构。

本教程的目的是展示`tensorize`的功能和用法，而不是提供有效的解决方案。

from __future__ import absolute_import, print_function
    
    import tvm
    import numpy as np

## 定义矩阵乘法 ##

以矩阵乘法为例。 矩阵乘法首先将两个矩阵之间的相应元素相乘，然后在某个轴上累加。下面，在TVM中描述了`A*B^T`的计算。

N, M, L = 1024, 512, 64
    A = tvm.placeholder((N, L), name='A')
    B = tvm.placeholder((M, L), name='B')
    #设置reduce轴
    k = tvm.reduce_axis((0, L), name='k')
    C = tvm.compute((N, M), lambda i, j:
                    tvm.sum(A[i, k] * B[j, k], axis=k), name='C')
    s = tvm.create_schedule(C.op)
    print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C { 
      for (i, 0, 1024) { 
        for (j, 0, 512) { 
          C[((i*512) + j)] = 0.000000f
          for (k, 0, 64) { 
            C[((i*512) + j)] = (C[((i*512) + j)] + (A[((i*64) + k)]*B[((j*64) + k)]))
          }
        }
      }
    }

## 调度矩阵乘法 ##

现在，假设我们有一个加速器支持矩阵向量乘法（GEMV）作为硬件原语，它可以采用任意大小的reduce轴，但另一个轴需要不大于16.因此我们分解矩阵乘法循环使最里面的循环为(16x64)GEMV。

factor = 16
    # x,y对应N,M
    x, y = C.op.axis
    #z对应L,sum缩减
    z, = C.op.reduce_axis
    #分裂
    yo, yi = s[C].split(y, factor=factor)
    #换轴
    s[C].reorder(x, yo, yi, z)
    print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C { 
      for (i, 0, 1024) { 
        for (j.outer, 0, 32) { 
          for (j.inner, 0, 16) { 
            C[((((i*32) + j.outer)*16) + j.inner)] = 0.000000f
            for (k, 0, 64) { 
              C[((((i*32) + j.outer)*16) + j.inner)] = (C[((((i*32) + j.outer)*16) + j.inner)] + (A[((i*64) + k)]*B[((((j.outer*16) + j.inner)*64) + k)]))
            }
          }
        }
      }
    }

如上IP打印显示，内循环`j.inner`与k和起来形成一个GEMV计算。在最内层的两个循环内，索引i是固定的，对矩阵A的访问仅改变k，这使得A的访问模式成为“向量”。为了利用我们假设的硬件GEMV指令，我们可以对`j.inner`进行**张量化**。

## 定义GEMV张量内联函数 ##

在调度张量化之前，我们需要首先定义GEMV的内联函数。它包括两部分，第一部分是GEMV的计算定义，TVM使用它来匹配原始矩阵乘法调度中的计算模式；第二个是指定如何在设备上执行GEMV，这在下面的`intrin_func`中完成。

def intrin_gemv(m, l):
        a = tvm.placeholder((l,), name='a')
        b = tvm.placeholder((m,l), name='b')
        k = tvm.reduce_axis((0,l), name='k')
        c = tvm.compute((m,), lambda i: tvm.sum(a[k]*b[i,k],axis=k), name='c')
        #声明buffer
        Ab = tvm.decl_buffer(a.shape, a.dtype,
                            name='A',
                            offset_factor=1,
                            stride=[1])
        Bb = tvm.decl_buffer(b.shape, b.dtype,
                             name="B",
                             offset_factor=1,
                             strides=[tvm.var("s1"), 1])
        Cb = tvm.decl_buffer(c.shape, c.dtype,
                             name="C",
                             offset_factor=1,
                             strides=[1])
        #设备原生函数调用
        def intrin_func(ins,outs):
            ib = tvm.ir_builder.create()
            aa, bb = ins
            cc = outs[0]
            ib.emit(tvm.call_extern("int32","gemv_update",
                   cc.access_ptr('w'),
                   aa.access_ptr('r'),
                   bb.access_ptr('r'),
                   m,l,bb.stride[0]))
            return ib.get()
        #
        with tvm.build_config(offset_factor=1):
            return tvm.decl_tensor_intrin(c.op, intrin_func, binds={ a: Ab, b: Bb, c: Cb})

这里`tvm.decl_tensor_intrin`声明怎么去执行`C.op`计算。我们的实现只需简单的输入和输出，将它们转换为指针并发出外部硬件内联函数调用。请注意，张量化需要用户指定`offset_factor`，利用此信息，TVM知道在原始数据结构的起始地址和传递给张量的偏移之间的**数据是否被对齐**，因此它有机会使用矢量化加载进行优化。为简化起见，我们将因子设置为1。

用于输入和输出的**缓冲区**也被声明，但这不是必需的，我们受益于缓冲区提供的额外信息。例如，我们将`bb.strides [0]`作为参数传递给外部函数`gemv_update`。现在`bb.strides [0] == l`，但稍后我们将看到它们如何与更复杂的调度有所不同。

注意，我们使用`tvm.var(“s1”)`作为B的第一个步幅维度。如果可以推断出步幅 - 在这种情况下，TVM知道张量B是紧凑的，因此步幅是`[L，1]` \- 这样的占位符可以让TVM自动绑定我们的推断值。

#gemv调用
    gemv = intrin_gemv(factor, L)
    #张量化
    s[C].tensorize(yi, gemv)
    print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C { 
      for (i, 0, 1024) { 
        for (j.outer, 0, 32) { 
          gemv_update(tvm_address_of(C[(((i*32) + j.outer)*16)]), tvm_address_of(A[(i*64)]), tvm_address_of(B[(j.outer*1024)]), 16, 64, 64)
        }
      }
    }

通过对`yi`进行张量化，最内部的两个循环现在被我们之前定义的内联函数所取代。为了构建和运行模块，让我们定义外部函数`gemv_update`，它是一个简单的GEMV实现，仅用于演示。

def gemv_impl():
        cc_code = """ extern "C" int gemv_update(float *cc, float *aa, float *bb, int m, int l, int stride) { for (int i = 0; i < m; ++i) { for (int j = 0; j < l; ++j) { cc[i] += aa[j] * bb[i * stride + j]; } } return 0; } """
        from tvm.contrib import util, clang
        temp = util.tempdir()
        ll_path = temp.relpath("temp.ll")
        # 从C源码创建LLVM IR
        ll_code = clang.create_llvm(cc_code, output=ll_path)
        return ll_code

现在我们利用编译属性`import_llvm`来导入llvm内联汇编。导入需要在张量化的GEMV执行之前进行。

s[C].pragma(x, "import_llvm", gemv_impl())
    print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C { 
      // attr [iter_var(i, )] pragma_import_llvm = "; ModuleID = '/tmp/tmps1q05tf9/input0.cc'\nsource_filename = \"/tmp/tmps1q05tf9/input0.cc\"\ntarget datalayout = \"e-m:e-i64:64-f80:128-n8:16:32:64-S128\"\ntarget triple = \"x86_64-pc-linux-gnu\"\n\n; Function Attrs: noinline nounwind optnone uwtable\ndefine i32 @gemv_update(float*, float*, float*, i32, i32, i32) #0 {\n %7 = alloca float*, align 8\n %8 = alloca float*, align 8\n %9 = alloca float*, align 8\n %10 = alloca i32, align 4\n %11 = alloca i32, align 4\n %12 = alloca i32, align 4\n %13 = alloca i32, align 4\n %14 = alloca i32, align 4\n store float* %0, float** %7, align 8\n store float* %1, float** %8, align 8\n store float* %2, float** %9, align 8\n store i32 %3, i32* %10, align 4\n store i32 %4, i32* %11, align 4\n store i32 %5, i32* %12, align 4\n store i32 0, i32* %13, align 4\n br label %15\n\n; <label>:15: ; preds = %50, %6\n %16 = load i32, i32* %13, align 4\n %17 = load i32, i32* %10, align 4\n %18 = icmp slt i32 %16, %17\n br i1 %18, label %19, label %53\n\n; <label>:19: ; preds = %15\n store i32 0, i32* %14, align 4\n br label %20\n\n; <label>:20: ; preds = %46, %19\n %21 = load i32, i32* %14, align 4\n %22 = load i32, i32* %11, align 4\n %23 = icmp slt i32 %21, %22\n br i1 %23, label %24, label %49\n\n; <label>:24: ; preds = %20\n %25 = load float*, float** %8, align 8\n %26 = load i32, i32* %14, align 4\n %27 = sext i32 %26 to i64\n %28 = getelementptr inbounds float, float* %25, i64 %27\n %29 = load float, float* %28, align 4\n %30 = load float*, float** %9, align 8\n %31 = load i32, i32* %13, align 4\n %32 = load i32, i32* %12, align 4\n %33 = mul nsw i32 %31, %32\n %34 = load i32, i32* %14, align 4\n %35 = add nsw i32 %33, %34\n %36 = sext i32 %35 to i64\n %37 = getelementptr inbounds float, float* %30, i64 %36\n %38 = load float, float* %37, align 4\n %39 = fmul float %29, %38\n %40 = load float*, float** %7, align 8\n %41 = load i32, i32* %13, align 4\n %42 = sext i32 %41 to i64\n %43 = getelementptr inbounds float, float* %40, i64 %42\n %44 = load float, float* %43, align 4\n %45 = fadd float %44, %39\n store float %45, float* %43, align 4\n br label %46\n\n; <label>:46: ; preds = %24\n %47 = load i32, i32* %14, align 4\n %48 = add nsw i32 %47, 1\n store i32 %48, i32* %14, align 4\n br label %20\n\n; <label>:49: ; preds = %20\n br label %50\n\n; <label>:50: ; preds = %49\n %51 = load i32, i32* %13, align 4\n %52 = add nsw i32 %51, 1\n store i32 %52, i32* %13, align 4\n br label %15\n\n; <label>:53: ; preds = %15\n ret i32 0\n}\n\nattributes #0 = { noinline nounwind optnone uwtable \"correctly-rounded-divide-sqrt-fp-math\"=\"false\" \"disable-tail-calls\"=\"false\" \"less-precise-fpmad\"=\"false\" \"no-frame-pointer-elim\"=\"true\" \"no-frame-pointer-elim-non-leaf\" \"no-infs-fp-math\"=\"false\" \"no-jump-tables\"=\"false\" \"no-nans-fp-math\"=\"false\" \"no-signed-zeros-fp-math\"=\"false\" \"no-trapping-math\"=\"false\" \"stack-protector-buffer-size\"=\"8\" \"target-cpu\"=\"x86-64\" \"target-features\"=\"+fxsr,+mmx,+sse,+sse2,+x87\" \"unsafe-fp-math\"=\"false\" \"use-soft-float\"=\"false\" }\n\n!llvm.module.flags = !{!0}\n!llvm.ident = !{!1}\n\n!0 = !{i32 1, !\"wchar_size\", i32 4}\n!1 = !{!\"clang version 6.0.1-svn334776-1~exp1~20190309042730.123 (branches/release_60)\"}\n"
      for (i, 0, 1024) { 
        for (j.outer, 0, 32) { 
          gemv_update(tvm_address_of(C[(((i*32) + j.outer)*16)]), tvm_address_of(A[(i*64)]), tvm_address_of(B[(j.outer*1024)]), 16, 64, 64)
        }
      }
    }

最后，我们将张量化版本与`numpy.dot`生成的版本进行比较，确保我们的实现是正确的。

func = tvm.build(s, [A, B, C], target="llvm", name="gemv")
    
    from topi.util import get_const_tuple
    dtype = A.dtype
    ctx = tvm.context("cpu", 0)
    a = np.random.uniform(size=get_const_tuple(A.shape)).astype(dtype)
    b = np.random.uniform(size=get_const_tuple(B.shape)).astype(dtype)
    c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=dtype), ctx)
    func(tvm.nd.array(a, ctx), tvm.nd.array(b, ctx), c)
    tvm.testing.assert_allclose(c.asnumpy(), np.dot(a, b.T), rtol=1e-3)

## 为张量化进行Reduce-update ##

到目前为止，您已经学会了张量化的基本概念，现在让我们向更复杂的情况迈进一步。

假设我们的加速器只能将向量乘以一个方阵，其中向量大小不需要大于16.鉴于这样的硬件约束，现在我们需要将reduce轴分割如下，

zo, zi = s[C].split(z, factor=factor)
    s[C].reorder(x, yo, zo, yi, zi)

但是，由于张量化内联计算现在只覆盖reduce轴的一部分，而不是使用一个“body”函数，TVM将在reduce循环之前需要调用一个`reduce_reset`函数，而`reduce_update`函数则定义“更新”计算调度。

def gemv_impl():
        cc_code = """ extern "C" int gemv_update(float *cc, float *aa, float *bb, int m, int l, int stride) { for (int i = 0; i < m; ++i) { for (int j = 0; j < l; ++j) { cc[i] += aa[j] * bb[i * stride + j]; } } return 0; } extern "C" int gemv_reset(float *cc, int m) { for (int i = 0; i < m; ++i) { cc[i] = 0.0; } return 0; } """
        from tvm.contrib import util, clang
        temp = util.tempdir()
        ll_path = temp.relpath("temp.ll")
        # Create LLVM ir from c source code
        ll_code = clang.create_llvm(cc_code, output=ll_path)
        return ll_code
    def intrin_gemv(m, l):
        a = tvm.placeholder((l,), name='a')
        b = tvm.placeholder((m, l), name='b')
        k = tvm.reduce_axis((0, l), name='k')
        c = tvm.compute((m,), lambda i:
        tvm.sum(a[k] * b[i, k], axis=k), name='c')
        Ab = tvm.decl_buffer(a.shape, a.dtype,
                             name="A",
                             offset_factor=1,
                             strides=[1])
        Bb = tvm.decl_buffer(b.shape, b.dtype,
                             name="B",
                             offset_factor=1,
                             strides=[tvm.var("s1"), 1])
        Cb = tvm.decl_buffer(c.shape, c.dtype,
                             name="C",
                             offset_factor=1,
                             strides=[1])
        def intrin_func(ins, outs):
            aa, bb = ins
            cc = outs[0]
            def _body():
                ib = tvm.ir_builder.create()
                ib.emit(tvm.call_extern("int32", "gemv_update",
                                        cc.access_ptr("w"),
                                        aa.access_ptr("r"),
                                        bb.access_ptr("r"),
                                        m, l, bb.strides[0]))
                return ib.get()
            
            def _reduce_reset():
                ib = tvm.ir_builder.create()
                ib.emit(tvm.call_extern("int32", "gemv_reset", cc.access_ptr("w"), m))
                return ib.get()
            def _reduce_update():
                return _body()
            #返回三个函数
            return _body(), _reduce_reset(), _reduce_update()
            with tvm.build_config(offset_factor=1):
            return tvm.decl_tensor_intrin(c.op, intrin_func, binds={ a: Ab, b: Bb, c: Cb})

请注意，现在`intrin_func`函数返回一个三元组:( body，reduce\_reset，reduce\_update）。如果张量化计算包含所有reduce轴，则将调用函数body()，否则将使用reduce\_reset()和reduce\_update()。在我们的示例中，body()和reduce\_update()共享相同的实现，而在其他情况下，硬件可能对这两个函数有不同的指令。此外，我们可以看到，由于平铺，`bb.strides [0]`与`l`是不同的。

对方形GEMV计算进行张量化，构建并检查结果

gemv = intrin_gemv(factor, factor)
    s[C].tensorize(yi, gemv)
    s[C].pragma(yo, "import_llvm", gemv_impl())
    
    func = tvm.build(s, [A, B, C], target="llvm", name="gemv")
    a = np.random.uniform(size=get_const_tuple(A.shape)).astype(dtype)
    b = np.random.uniform(size=get_const_tuple(B.shape)).astype(dtype)
    c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=dtype), ctx)
    func(tvm.nd.array(a, ctx), tvm.nd.array(b, ctx), c)
    tvm.testing.assert_allclose(c.asnumpy(), np.dot(a, b.T), rtol=1e-3)

## 总结 ##

本教程演示了TVM中使用张量化内联计算。张量化`Tensorize`为用户提供了一种通过微内核获得完全优化调度的方法。例如，Intel CPU上的INT8量化可以使用张量来直接调用AVX指令。它还使TVM能够编译到ASIC-详情查看[VTA][]。我们还演示了如何使用内联汇编导入，这有助于用户将汇编轻松导入到计算调度中。

[VTA]: https://docs.tvm.ai/vta/index.html