TVM Reduction降低算力
這是有關(guān)如何降低算力TVM的介紹材料。像sum / max / min這樣的關(guān)聯(lián)約簡(jiǎn)運(yùn)算符是線性代數(shù)運(yùn)算的典型構(gòu)造塊。
本文將演示如何降低TVM算力。
from future import absolute_import, print_function

import tvm
import tvm.testing
from tvm import te
import numpy as np
描述行數(shù)
假設(shè)要計(jì)算行總數(shù)作為示例。用numpy語(yǔ)義可以寫(xiě)成B = numpy.sum(A, axis=1)
以下幾行描述了行求和算子。創(chuàng)建歸約公式，使用 te.reduce_axis來(lái)聲明歸約軸。te.reduce_axis降低算力的范圍。 te.sum接受要降低算力的表達(dá)式以及降低算力軸，并計(jì)算聲明范圍內(nèi)所有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”)
調(diào)度降低算力
有幾種調(diào)度降低算力的方法。在執(zhí)行任何操作之前，打印出默認(rèn)調(diào)度的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[(istride)] = 0f32
for (k: int32, 0, m) {
B_2[(istride)] = ((float32*)B_2[(istride)] + (float32)A_2[((istride_1) + (kstride_2))])
}
}
}
會(huì)發(fā)現(xiàn)IR代碼與C代碼非常相似。減速軸類(lèi)似于法線軸，可以拆分。
在下面的代碼中，將B的行軸和軸拆分為不同的因子。結(jié)果是嵌套歸約。
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.outer32) + i.inner) < n), dtype=bool) {
B_2[(((i.outer32) + i.inner)stride)] = 0f32
}
if @tir.likely((((i.outer32) + i.inner) < n), dtype=bool) {
for (k.outer: int32, 0, floordiv((m + 15), 16)) {
for (k.inner: int32, 0, 16) {
if @tir.likely((((k.outer16) + k.inner) < m), dtype=bool) {
B_2[(((i.outer32) + i.inner)stride)] = ((float32)B_2[(((i.outer32) + i.inner)stride)] + (float32)A_2[((((i.outer32) + i.inner)stride_1) + (((k.outer16) + k.inner)stride_2))])
}
}
}
}
}
}
}
要構(gòu)建GPU內(nèi)核，可以將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.x32) + threadIdx.x) < n), dtype=bool) {
B_2[(((blockIdx.x32) + threadIdx.x)stride)] = 0f32
}
for (k.outer: int32, 0, floordiv((m + 15), 16)) {
for (k.inner: int32, 0, 16) {
if @tir.likely((((blockIdx.x32) + threadIdx.x) < n), dtype=bool) {
if @tir.likely((((k.outer16) + k.inner) < m), dtype=bool) {
B_2[(((blockIdx.x32) + threadIdx.x)stride)] = ((float32)B_2[(((blockIdx.x32) + threadIdx.x)stride)] + (float32)A_2[((((blockIdx.x32) + threadIdx.x)stride_1) + (((k.outer16) + k.inner)stride_2))])
}
}
}
}
}
}
歸約分解和并行化
建立歸約的一個(gè)問(wèn)題是，不能簡(jiǎn)單地在歸約軸上并行化。需要對(duì)約簡(jiǎn)的算子進(jìn)行劃分，在對(duì)臨時(shí)數(shù)組進(jìn)行約簡(jiǎn)之前，將局部約簡(jiǎn)結(jié)果存儲(chǔ)在臨時(shí)數(shù)組中。
rfactor原語(yǔ)會(huì)重寫(xiě)計(jì)算。在下面的調(diào)度中，將B的結(jié)果寫(xiě)入臨時(shí)結(jié)果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, [(n16)]) {
for (k.inner: int32, 0, 16) {
for (i: int32, 0, n) {
B.rf[((k.innern) + i)] = 0f32
for (k.outer: int32, 0, floordiv((m + 15), 16)) {
if @tir.likely((((k.outer16) + k.inner) < m), dtype=bool) {
B.rf[((k.innern) + i)] = ((float32*)B.rf[((k.innern) + i)] + (float32)A_2[((istride_1) + (((k.outer16) + k.inner)stride_2))])
}
}
}
}
for (ax0: int32, 0, n) {
B_2[(ax0stride)] = 0f32
for (k.inner.v: int32, 0, 16) {
B_2[(ax0stride)] = ((float32)B_2[(ax0stride)] + (float32)B.rf[((k.inner.vn) + ax0)])
}
}
}
}
B的調(diào)度算子也將被重寫(xiě)為Bf縮減結(jié)果的第一軸上的和
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)]
降低算力跨線
現(xiàn)在，我們可以在分解后的軸上進(jìn)行并行化處理。在此，B的復(fù)位軸標(biāo)記為螺紋。TVM將算力減少軸標(biāo)記為線程，如果它是唯一的算力降低，則可以在設(shè)備中進(jìn)行交叉線程。
分解后的情況確實(shí)如此。也可以直接在還原軸上計(jì)算BF。最終生成的內(nèi)核將按blockIdx.x劃分行，按threadIdx.x劃分threadIdx.y列，最后對(duì)threadIdx.x進(jìn)行跨線程縮減
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))];
}
}
將結(jié)果內(nèi)核與numpy進(jìn)行比較，驗(yàn)證結(jié)果內(nèi)核的正確性。
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)
通過(guò)2D簡(jiǎn)化描述卷積
在TVM中，可以通過(guò)2D約簡(jiǎn)來(lái)描述卷積。這是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)]))
}
}
}
}
}
定義通用換向歸約運(yùn)算
除了內(nèi)置的如降低算力操作te.sum， tvm.te.min和tvm.te.max，還可以通過(guò)定義交換降低算力操作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”)
注意
執(zhí)行涉及多個(gè)值的歸約argmax，可以通過(guò)元組輸入來(lái)完成。有關(guān)更多詳細(xì)信息，請(qǐng)參見(jiàn)使用協(xié)作輸入來(lái)描述縮減。
總結(jié)
本文提供了降低算力調(diào)度的演練。
? 用reduce_axis描述歸約。
? 如果需要并行性，請(qǐng)使用rfactor分解軸。
? 定義新的歸約運(yùn)算 te.comm_reducer

總結(jié)

以上是生活随笔為你收集整理的TVM Reduction降低算力的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。