TVM Reduction降低算力

技术文章 3个月前 完美者
1,378 0

标签:详细   使用   where   c代码   直接   https   tle   lob   modules   

TVM Reduction降低算力

这是有关如何降低算力TVM的介绍材料。像sum / max / min这样的关联约简运算符是线性代数运算的典型构造块。

本文将演示如何降低TVM算力。

from __future__ import absolute_import, print_function

 

import tvm

import tvm.testing

from tvm import te

import numpy as np

描述行数

假设要计算行总数作为示例。用numpy语义可以写成B = numpy.sum(A, axis=1)

以下几行描述了行求和算子。创建归约公式,使用 te.reduce_axis来声明归约轴。te.reduce_axis降低算力的范围。 te.sum接受要降低算力的表达式以及降低算力轴,并计算声明范围内所有k的值之和。

等效的C代码如下:

for (int i = 0; i < n; ++i) {

  B[i] = 0;

  for (int k = 0; k < m; ++k) {

    B[i] = B[i] + A[i][k];

  }

}

n = te.var("n")

m = te.var("m")

A = te.placeholder((n, m), name="A")

k = te.reduce_axis((0, m), "k")

B = te.compute((n,), lambda i: te.sum(A[i, k], axis=k), name="B")

调度降低算力

有几种调度降低算力的方法。在执行任何操作之前,打印出默认调度的IR代码。

s = te.create_schedule(B.op)

print(tvm.lower(s, [A, B], simple_mode=True))

输出:

primfn(A_1: handle, B_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {B: Buffer(B_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
             A: Buffer(A_2: Pointer(float32), float32, [n, m: int32], [stride_1: int32, stride_2: int32], type="auto")}
  buffer_map = {A_1: A, B_1: B} {
  for (i: int32, 0, n) {
    B_2[(i*stride)] = 0f32
    for (k: int32, 0, m) {
      B_2[(i*stride)] = ((float32*)B_2[(i*stride)] + (float32*)A_2[((i*stride_1) + (k*stride_2))])
    }
  }
}

会发现IR代码与C代码非常相似。减速轴类似于法线轴,可以拆分。

在下面的代码中,将B的行轴和轴拆分为不同的因子。结果是嵌套归约。

ko, ki = s[B].split(B.op.reduce_axis[0], factor=16)
xo, xi = s[B].split(B.op.axis[0], factor=32)
print(tvm.lower(s, [A, B], simple_mode=True))

输出:

primfn(A_1: handle, B_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {B: Buffer(B_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
             A: Buffer(A_2: Pointer(float32), float32, [n, m: int32], [stride_1: int32, stride_2: int32], type="auto")}
  buffer_map = {A_1: A, B_1: B} {
  for (i.outer: int32, 0, floordiv((n + 31), 32)) {
    for (i.inner: int32, 0, 32) {
      if @tir.likely((((i.outer*32) + i.inner) < n), dtype=bool) {
        B_2[(((i.outer*32) + i.inner)*stride)] = 0f32
      }
      if @tir.likely((((i.outer*32) + i.inner) < n), dtype=bool) {
        for (k.outer: int32, 0, floordiv((m + 15), 16)) {
          for (k.inner: int32, 0, 16) {
            if @tir.likely((((k.outer*16) + k.inner) < m), dtype=bool) {
              B_2[(((i.outer*32) + i.inner)*stride)] = ((float32*)B_2[(((i.outer*32) + i.inner)*stride)] + (float32*)A_2[((((i.outer*32) + i.inner)*stride_1) + (((k.outer*16) + k.inner)*stride_2))])
            }
          }
        }
      }
    }
  }
}

要构建GPU内核,可以将B的行绑定到GPU线程。

s[B].bind(xo, te.thread_axis("blockIdx.x"))
s[B].bind(xi, te.thread_axis("threadIdx.x"))
print(tvm.lower(s, [A, B], simple_mode=True))

输出:

primfn(A_1: handle, B_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {B: Buffer(B_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
             A: Buffer(A_2: Pointer(float32), float32, [n, m: int32], [stride_1: int32, stride_2: int32], type="auto")}
  buffer_map = {A_1: A, B_1: B} {
  attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = floordiv((n + 31), 32);
  attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 32 {
    if @tir.likely((((blockIdx.x*32) + threadIdx.x) < n), dtype=bool) {
      B_2[(((blockIdx.x*32) + threadIdx.x)*stride)] = 0f32
    }
    for (k.outer: int32, 0, floordiv((m + 15), 16)) {
      for (k.inner: int32, 0, 16) {
        if @tir.likely((((blockIdx.x*32) + threadIdx.x) < n), dtype=bool) {
          if @tir.likely((((k.outer*16) + k.inner) < m), dtype=bool) {
            B_2[(((blockIdx.x*32) + threadIdx.x)*stride)] = ((float32*)B_2[(((blockIdx.x*32) + threadIdx.x)*stride)] + (float32*)A_2[((((blockIdx.x*32) + threadIdx.x)*stride_1) + (((k.outer*16) + k.inner)*stride_2))])
          }
        }
      }
    }
  }
}

归约分解和并行化

建立归约的一个问题是,不能简单地在归约轴上并行化。需要对约简的算子进行划分,在对临时数组进行约简之前,将局部约简结果存储在临时数组中。

rfactor原语会重写计算。在下面的调度中,将B的结果写入临时结果B.rf。分解后的尺寸成为B.rf的第一尺寸。

s = te.create_schedule(B.op)
ko, ki = s[B].split(B.op.reduce_axis[0], factor=16)
BF = s.rfactor(B, ki)
print(tvm.lower(s, [A, B], simple_mode=True))

输出:

primfn(A_1: handle, B_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {B: Buffer(B_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
             A: Buffer(A_2: Pointer(float32), float32, [n, m: int32], [stride_1: int32, stride_2: int32], type="auto")}
  buffer_map = {A_1: A, B_1: B} {
  attr [B.rf: Pointer(float32)] "storage_scope" = "global";
  allocate(B.rf, float32, [(n*16)]) {
    for (k.inner: int32, 0, 16) {
      for (i: int32, 0, n) {
        B.rf[((k.inner*n) + i)] = 0f32
        for (k.outer: int32, 0, floordiv((m + 15), 16)) {
          if @tir.likely((((k.outer*16) + k.inner) < m), dtype=bool) {
            B.rf[((k.inner*n) + i)] = ((float32*)B.rf[((k.inner*n) + i)] + (float32*)A_2[((i*stride_1) + (((k.outer*16) + k.inner)*stride_2))])
          }
        }
      }
    }
    for (ax0: int32, 0, n) {
      B_2[(ax0*stride)] = 0f32
      for (k.inner.v: int32, 0, 16) {
        B_2[(ax0*stride)] = ((float32*)B_2[(ax0*stride)] + (float32*)B.rf[((k.inner.v*n) + ax0)])
      }
    }
  }
}

B的调度算子也将被重写为Bf缩减结果的第一轴上的和

print(s[B].op.body)

输出:

[reduce(combiner=comm_reducer(result=[(x + y)], lhs=[x], rhs=[y], identity_element=[0f]), source=[B.rf[k.inner.v, ax0]], init=[], axis=[iter_var(k.inner.v, range(min=0, ext=16))], where=(bool)1, value_index=0)]

降低算力跨线

现在,我们可以在分解后的轴上进行并行化处理。在此,B的复位轴标记为螺纹。TVM将算力减少轴标记为线程,如果它是唯一的算力降低,则可以在设备中进行交叉线程。

分解后的情况确实如此。也可以直接在还原轴上计算BF。最终生成的内核将按blockIdx.x划分行,按threadIdx.x划分threadIdx.y列,最后对threadIdx.x进行跨线程缩减

xo, xi = s[B].split(s[B].op.axis[0], factor=32)
s[B].bind(xo, te.thread_axis("blockIdx.x"))
s[B].bind(xi, te.thread_axis("threadIdx.y"))
tx = te.thread_axis("threadIdx.x")
s[B].bind(s[B].op.reduce_axis[0], tx)
s[BF].compute_at(s[B], s[B].op.reduce_axis[0])
s[B].set_store_predicate(tx.var.equal(0))
fcuda = tvm.build(s, [A, B], "cuda")
print(fcuda.imported_modules[0].get_source())

输出:

extern "C" __global__ void default_function_kernel0(float* __restrict__ A, float* __restrict__ B, int m, int n, int stride, int stride1, int stride2) {
  float B_rf[1];
  __shared__ float red_buf0[512];
  B_rf[(0)] = 0.000000e+00f;
  for (int k_outer = 0; k_outer < (m >> 4); ++k_outer) {
    if (((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) < n) {
      B_rf[(0)] = (B_rf[(0)] + A[(((((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) * stride) + (((k_outer * 16) + ((int)threadIdx.x)) * stride1)))]);
    }
  }
  for (int k_outer1 = 0; k_outer1 < (((m & 15) + 15) >> 4); ++k_outer1) {
    if (((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) < n) {
      if (((((m >> 4) * 16) + (k_outer1 * 16)) + ((int)threadIdx.x)) < m) {
        B_rf[(0)] = (B_rf[(0)] + A[(((((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) * stride) + (((((m >> 4) * 16) + (k_outer1 * 16)) + ((int)threadIdx.x)) * stride1)))]);
      }
    }
  }
  __syncthreads();
  ((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] = B_rf[(0)];
  __syncthreads();
  if (((int)threadIdx.x) < 8) {
    ((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] = (((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] + ((volatile float*)red_buf0)[((((((int)threadIdx.y) * 16) + ((int)threadIdx.x)) + 8))]);
    ((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] = (((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] + ((volatile float*)red_buf0)[((((((int)threadIdx.y) * 16) + ((int)threadIdx.x)) + 4))]);
    ((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] = (((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] + ((volatile float*)red_buf0)[((((((int)threadIdx.y) * 16) + ((int)threadIdx.x)) + 2))]);
    ((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] = (((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] + ((volatile float*)red_buf0)[((((((int)threadIdx.y) * 16) + ((int)threadIdx.x)) + 1))]);
  }
  __syncthreads();
  if (((int)threadIdx.x) == 0) {
    B[((((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) * stride2))] = ((volatile float*)red_buf0)[((((int)threadIdx.y) * 16))];
  }
}

将结果内核与numpy进行比较,验证结果内核的正确性。

nn = 128
ctx = tvm.gpu(0)
a = tvm.nd.array(np.random.uniform(size=(nn, nn)).astype(A.dtype), ctx)
b = tvm.nd.array(np.zeros(nn, dtype=B.dtype), ctx)
fcuda(a, b)
tvm.testing.assert_allclose(b.asnumpy(), np.sum(a.asnumpy(), axis=1), rtol=1e-4)

通过2D简化描述卷积

在TVM中,可以通过2D约简来描述卷积。这是2D卷积的示例,滤波器大小= [3,3],步幅= [1,1]。

n = te.var("n")
Input = te.placeholder((n, n), name="Input")
Filter = te.placeholder((3, 3), name="Filter")
di = te.reduce_axis((0, 3), name="di")
dj = te.reduce_axis((0, 3), name="dj")
Output = te.compute(
    (n - 2, n - 2),
    lambda i, j: te.sum(Input[i + di, j + dj] * Filter[di, dj], axis=[di, dj]),
    name="Output",
)
s = te.create_schedule(Output.op)
print(tvm.lower(s, [Input, Filter, Output], simple_mode=True))

出:

primfn(Input_1: handle, Filter_1: handle, Output_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {Output: Buffer(Output_2: Pointer(float32), float32, [(n: int32 - 2), (n - 2)], []),
             Filter: Buffer(Filter_2: Pointer(float32), float32, [3, 3], []),
             Input: Buffer(Input_2: Pointer(float32), float32, [n, n], [stride: int32, stride_1: int32], type="auto")}
  buffer_map = {Input_1: Input, Filter_1: Filter, Output_1: Output} {
  for (i: int32, 0, (n - 2)) {
    for (j: int32, 0, (n - 2)) {
      Output_2[((i*(n - 2)) + j)] = 0f32
      for (di: int32, 0, 3) {
        for (dj: int32, 0, 3) {
          Output_2[((i*(n - 2)) + j)] = ((float32*)Output_2[((i*(n - 2)) + j)] + ((float32*)Input_2[(((i + di)*stride) + ((j + dj)*stride_1))]*(float32*)Filter_2[((di*3) + dj)]))
        }
      }
    }
  }
}

定义通用换向归约运算

除了内置的如降低算力操作te.sum, tvm.te.mintvm.te.max,还可以通过定义交换降低算力操作te.comm_reducer

n = te.var("n")
m = te.var("m")
product = te.comm_reducer(lambda x, y: x * y, lambda t: tvm.tir.const(1, dtype=t), name="product")
A = te.placeholder((n, m), name="A")
k = te.reduce_axis((0, m), name="k")
B = te.compute((n,), lambda i: product(A[i, k], axis=k), name="B")

注意

执行涉及多个值的归约argmax,可以通过元组输入来完成。有关更多详细信息,请参见使用协作输入来描述缩减。

总结

本文提供了降低算力调度的演练。

  • 用reduce_axis描述归约。
  • 如果需要并行性,请使用rfactor分解轴。
  • 定义新的归约运算 te.comm_reducer

 

TVM Reduction降低算力

标签:详细   使用   where   c代码   直接   https   tle   lob   modules   

原文地址:https://www.cnblogs.com/wujianming-110117/p/14176685.html

版权声明:完美者 发表于 2020-12-29 11:09:27。
转载请注明:TVM Reduction降低算力 | 完美导航

暂无评论

暂无评论...