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

Posted 2021-05-05

[TOC]

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

这篇教程是关于在TVM中如何执行`张量化`的介绍。

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

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

```python
from __future__ import absolute_import, print_function

import tvm
import numpy as np
```

## 定义矩阵乘法

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

```python
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))
```

输出：

```c
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。

```python
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))
```

输出：

```c
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`中完成。

```python
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自动绑定我们的推断值。

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

输出：

```c
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实现，仅用于演示。

```python
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执行之前进行。

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

输出：

```c
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`生成的版本进行比较，确保我们的实现是正确的。

```python
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轴分割如下，

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

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

```python
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计算进行张量化，构建并检查结果

```c
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](https://docs.tvm.ai/vta/index.html)。我们还演示了如何使用内联汇编导入，这有助于用户将汇编轻松导入到计算调度中。