GPU自动调度卷积层

GPU自动调度卷积层

本文对GPU使用自动调度程序。

与依靠手动模板定义搜索空间的基于模板的autotvm不同，自动调度程序不需要任何模板。用户只需要编写计算声明，无需任何调度命令或模板。自动调度程序可以自动生成一个较大的搜索空间，在该空间中找到良好的调度。

本文以卷积层为例。

本文无法在Windows或最新版本的macOS上运行。要使其运行，需要将本文的内容包装在一个if __name__ == "__main__":块中。

import os

import numpy as np

import tvm

from tvm import te, auto_scheduler, topi

from tvm.topi.testing import conv2d_nchw_python

定义计算

首先，定义卷积层的计算。该函数应返回输入/输出张量的列表。通过这些张量，自动调度器可以获得整个计算图。

@auto_scheduler.register_workload

def conv2d_layer(N, H, W, CO, CI, KH, KW, stride, padding):

data = te.placeholder((N, CI, H, W), name="data")

kernel = te.placeholder((CO, CI, KH, KW), name="kernel")

bias = te.placeholder((1, CO, 1, 1), name="bias")

conv = topi.nn.conv2d_nchw(data, kernel, stride, padding, dilation=1, out_dtype="float32")

out = topi.nn.relu(conv + bias)

return [data, kernel, bias, out]

创建搜索任务

然后，为resnet中的最后一个卷积层创建搜索任务。

target = tvm.target.Target("cuda")

# Use the last layer in ResNet-50

N, H, W, CO, CI, KH, KW, strides, padding = 1, 7, 7, 512, 512, 3, 3, (1, 1), (1, 1)

task = auto_scheduler.SearchTask(

func=conv2d_layer, args=(N, H, W, CO, CI, KH, KW, strides, padding), target=target

)

# Inspect the computational graph

print("Computational DAG:")

print(task.compute_dag)

输出：

Computational DAG:

data = PLACEHOLDER [1, 512, 7, 7]

pad_temp(i0, i1, i2, i3) = tir.if_then_else(((((i2 >= 1) && (i2 < 8)) && (i3 >= 1)) && (i3 < 8)), data[i0, i1, (i2 - 1), (i3 - 1)], 0f)

kernel = PLACEHOLDER [512, 512, 3, 3]

compute(nn, ff, yy, xx) += (pad_temp[nn, rc, (yy + ry), (xx + rx)]*kernel[ff, rc, ry, rx])

bias = PLACEHOLDER [1, 512, 1, 1]

T_add(ax0, ax1, ax2, ax3) = (compute[ax0, ax1, ax2, ax3] + bias[ax0, ax1, 0, 0])

compute(i0, i1, i2, i3) = max(T_add[i0, i1, i2, i3], 0f)

接下来，为自动调度程序设置参数。这些参数主要指定在搜索过程中如何进行测量。

measure_ctx启动不同的测量过程以提供隔离。保护主进程免受测量期间GPU崩溃的影响，避免其它运行时冲突。
min_repeat_ms定义每次测量中一次“重复”的最小持续时间。这样可以预热GPU，对于获得准确的测量结果是必不可少的。通常，建议值> = 300毫秒。
num_measure_trials是在搜索过程中可以使用的测量试验的数量。为了快速演示，在本文中仅进行了10次试用。在实践中，1000是使搜索收敛的一个好值。可以根据自己的时间预算进行更多试验。
此外，还用RecordToFile将测量记录转储到文件conv2d.json中。测量记录可用于最好地查询历史记录，恢复搜索以及以后进行更多分析。
有关更多参数auto_scheduler.TuningOptions，请参见auto_scheduler.LocalRPCMeasureContext。

log_file = "conv2d.json"

measure_ctx = auto_scheduler.LocalRPCMeasureContext(min_repeat_ms=300)

tune_option = auto_scheduler.TuningOptions(

num_measure_trials=10, # change this to 1000 to achieve the best performance

runner=measure_ctx.runner,

measure_callbacks=[auto_scheduler.RecordToFile(log_file)],

verbose=2,

)

输出：

Get devices for measurement successfully!

运行搜索

现在准备好所有输入。开始搜索，让自动调度程序发挥作用。经过一些测量试验之后，可以从日志文件中加载最佳调度并应用它。

# Run auto-tuning (search)

task.tune(tune_option)

# Apply the best schedule

sch, args = task.apply_best(log_file)

# Kill the measurement process

del measure_ctx

输出：

可以降低调度以在自动调度后查看IR。自动调度程序可以正确执行优化，包括多层平铺，协作提取，展开和算子融合。

print("Lowered TIR:")

print(tvm.lower(sch, args, simple_mode=True))

输出：

Lowered TIR:

primfn(data_1: handle, kernel_1: handle, bias_1: handle, compute_1: handle) -> ()

attr = {"global_symbol": "main", "tir.noalias": True}

buffers = {compute: Buffer(compute_2: Pointer(float32), float32, [1, 512, 7, 7], []),

kernel: Buffer(kernel_2: Pointer(float32), float32, [512, 512, 3, 3], []),

bias: Buffer(bias_2: Pointer(float32), float32, [1, 512, 1, 1], []),

data: Buffer(data_2: Pointer(float32), float32, [1, 512, 7, 7], [])}

buffer_map = {data_1: data, kernel_1: kernel, bias_1: bias, compute_1: compute} {

attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = 16;

attr [compute_3: Pointer(float32)] "storage_scope" = "local";

allocate(compute_3, float32, [14]);

attr [pad_temp.shared: Pointer(float32)] "storage_scope" = "shared";

allocate(pad_temp.shared, float32, [1296]);

attr [kernel.shared: Pointer(float32)] "storage_scope" = "shared";

allocate(kernel.shared, float32, [4608]);

attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 112 {

compute_3[0] = 0f32

compute_3[7] = 0f32

compute_3[1] = 0f32

compute_3[8] = 0f32

compute_3[2] = 0f32

compute_3[9] = 0f32

compute_3[3] = 0f32

compute_3[10] = 0f32

compute_3[4] = 0f32

compute_3[11] = 0f32

compute_3[5] = 0f32

compute_3[12] = 0f32

compute_3[6] = 0f32

compute_3[13] = 0f32

for (rc.outer.outer: int32, 0, 32) {

attr [IterVar(threadIdx.x_1: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 112;

pad_temp.shared[threadIdx.x_1] = @tir.if_then_else(((((9 <= floormod(threadIdx.x_1, 81)) && (floormod(threadIdx.x_1, 81) < 72)) && (1 <= floormod(threadIdx.x_1, 9))) && (floormod(threadIdx.x_1, 9) < 8)), (float32*)data_2[(((((rc.outer.outer*784) + (floordiv(threadIdx.x_1, 81)*49)) + (floordiv(floormod(threadIdx.x_1, 81), 9)*7)) + floormod(threadIdx.x_1, 9)) - 8)], 0f32, dtype=float32)