编写可调模板并使用自动调谐器

编写可调模板并使用自动调谐器

这是TVM中自动调试模块的入门说明。

自动调试有两个步骤。第一步是定义搜索空间。第二步是运行搜索算法来探索这个空间。本文可以学习如何在TVM中执行这两个步骤。整个工作流程由矩阵乘法示例说明。

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

安装依赖

要在TVM中使用autotvm软件包,需要安装一些额外的依赖项。可以跳过此步骤(安装xgboost),因为它不需要XGBoost(如果使用python2,请将“ 3”更改为“ 2”):

pip3 install --user psutil xgboost

为了使TVM在调试中更快地运行,建议使用cython作为TVM的FFI。在TVM的根目录中,执行(如果使用python2,请将“ 3”更改为“ 2”):

pip3 install --user cython

sudo make cython3

现在返回python代码,导入包。

import logging

import sys

 

import numpy as np

import tvm

from tvm import te, testing

 

# the module is called `autotvm`

from tvm import autotvm

步骤1:定义搜索空间

在本节中,将确定性TVM调度代码重写为可调调度模板。可以将搜索空间定义的过程视为现有调度代码的参数化。

首先,这是如何在TVM中实现分块矩阵乘法。

# Matmul V0: Constant tiling factor

def matmul_v0(N, L, M, dtype):

    A = te.placeholder((N, L), name="A", dtype=dtype)

    B = te.placeholder((L, M), name="B", dtype=dtype)

 

    k = te.reduce_axis((0, L), name="k")

    C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")

    s = te.create_schedule(C.op)

 

    # schedule

    y, x = s[C].op.axis

    k = s[C].op.reduce_axis[0]

 

    yo, yi = s[C].split(y, 8)

    xo, xi = s[C].split(x, 8)

 

    s[C].reorder(yo, xo, k, yi, xi)

 

    return s, [A, B, C]

参数化调度

在先前的调度代码中,使用常数“ 8”作为切片因子。但是,可能不是最好的,因为最好的切片因子取决于实际的硬件环境和输入形态。

如果希望调度代码可在更广泛的输入形态和目标硬件中移植,则最好定义一组候选值并根据目标硬件上的测量结果选择最佳值。

在autotvm中,可以为此类值定义可调参数或“旋钮”。

# Matmul V1: List candidate values

@autotvm.template("tutorial/matmul_v1")  # 1. use a decorator

def matmul_v1(N, L, M, dtype):

    A = te.placeholder((N, L), name="A", dtype=dtype)

    B = te.placeholder((L, M), name="B", dtype=dtype)

 

    k = te.reduce_axis((0, L), name="k")

    C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")

    s = te.create_schedule(C.op)

 

    # schedule

    y, x = s[C].op.axis

    k = s[C].op.reduce_axis[0]

 

    # 2. get the config object

    cfg = autotvm.get_config()

 

    # 3. define search space

    cfg.define_knob("tile_y", [1, 2, 4, 8, 16])

    cfg.define_knob("tile_x", [1, 2, 4, 8, 16])

 

    # 4. schedule according to config

    yo, yi = s[C].split(y, cfg["tile_y"].val)

    xo, xi = s[C].split(x, cfg["tile_x"].val)

 

    s[C].reorder(yo, xo, k, yi, xi)

 

    return s, [A, B, C]

在这里,对先前的调度代码进行了四个修改,并获得了一个可调的“模板”。可以一一解释这些修改。

  • 使用decorator将此功能标记为简单模板。
  • 获取配置对象:可以将此cfg作为该函数的参数,但是可以通过其他方式获得。使用此参数,此功能不再是确定性的调度代码。相反,可以将不同的配置传递给此功能并获得不同的调度,因此此功能是一个“模板”。

为了使模板函数更紧凑,在单个函数中做了两件事。(1)定义搜索空间,(2)根据该空间中的实体进行调度。为了达到这个目的,将cfg其设为aConfigSpaceConfigEntityobject。

如果是ConfigSpace,将收集此功能中的所有可tune调旋钮并建立搜索空间。如果是ConfigEntity,它将忽略所有空间定义API(即cfg.define_XXXXX(...))。相反,存储所有可调旋钮的确定性值,根据这些值进行调度。

在自动调试期间,将首先使用一个ConfigSpace 对象调用此模板以构建搜索空间。然后,ConfigEntity 在构建空间中将此模板称为不同的模板,以获取不同的调度。最后,将测量由不同调度生成的代码,并选择最佳调度。

  • 定义两个可调旋钮。第一个tile_y具有5个可能的值。第二个tile_x具有相同的可能值列表。这两个旋钮是独立的,因此它们跨越的搜索空间的大小为5x5 = 25
  • 根据cfg中的确定性值进行调度 

使用更好的空间定义API

在上一个模板中,手动列出了旋钮的所有可能值。这是定义空间的最低级别的API。但是,还提供了另一组API,以使空间定义更轻松,更智能。建议使用这套高级API。

在以下示例中,用ConfigSpace.define_split定义拆分旋钮。它将列举所有可能的方式来分割轴并构造空间。

还有ConfigSpace.define_reorder用于重新排序的旋钮和 ConfigSpace.define_annotate用于展开,向量化,线程绑定的注释。当高级API无法满足要求时,始终可以回退而使用低级API。

@autotvm.template("tutorial/matmul")

def matmul(N, L, M, dtype):

    A = te.placeholder((N, L), name="A", dtype=dtype)

    B = te.placeholder((L, M), name="B", dtype=dtype)

 

    k = te.reduce_axis((0, L), name="k")

    C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")

    s = te.create_schedule(C.op)

 

    # schedule

    y, x = s[C].op.axis

    k = s[C].op.reduce_axis[0]

 

    ##### define space begin #####

    cfg = autotvm.get_config()

    cfg.define_split("tile_y", y, num_outputs=2)

    cfg.define_split("tile_x", x, num_outputs=2)

    ##### define space end #####

 

    # schedule according to config

    yo, yi = cfg["tile_y"].apply(s, C, y)

    xo, xi = cfg["tile_x"].apply(s, C, x)

 

    s[C].reorder(yo, xo, k, yi, xi)

 

    return s, [A, B, C]

注意

有关cfg.defile_split更多说明 

在此模板中,将枚举所有可以将y轴分解为长度为y的两个轴的可能组合。例如,如果y的长度为32,并且想使用32的因数将其分成两个轴,则(外轴的长度,内轴的长度)对有6个可能的值,即(32,1) ,(16、2),(8、4),(4、8),(2、16)或(1、32)。它们只有6个可能tile_y的值。cfg.define_split("tile_y", y, num_outputs=2)

在调度期间,cfg["tile_y"]是一个SplitEntity对象。将外轴和内轴的长度存储在cfg['tile_y'].size (具有两个元素的元组)中。在此模板中,使用yo, yi = cfg['tile_y'].apply(s, C, y)来应用它。实际上,这等于 yo, yi = s[C].split(y, cfg["tile_y"].size[1])或 yo, yi = s[C].split(y, nparts=cfg['tile_y"].size[0])

使用cfg.apply API的优点是,使多层拆分(当num_outputs> = 3时)更加容易。

步骤2:搜寻空间

在第1步中,通过将旧的调度代码扩展到模板中来构建搜索空间。下一步是选择一个tune调谐器,并在这个空间中进行探索。

TVM中的自动tune调谐器

tune调谐器的工作可以通过以下伪代码来描述

ct = 0

while ct < max_number_of_trials:

    propose a batch of configs

    measure this batch of configs on real hardware and get results

    ct += batch_size

当提议下一批配置时,调谐器可以采取不同的策略。在autotvm中为四个调谐器提供了不同的策略。

  • RandomTuner:以随机顺序枚举空间
  • GridSearchTuner:按网格搜索顺序枚举空间
  • GATuner:使用遗传算法搜索空间
  • XGBTuner:使用基于模型的方法。训练XGBoost模型以预测降低的IR的速度,并根据预测选择下一批。

可以根据空间大小,时间预算和其他因素选择调谐器。例如,如果空间很小(小于1000),那么使用gridsearch调谐器或随机调谐器就足够了。如果空间为10 ^ 9的水平(这是CUDA GPU上conv2d运算符的空间大小),则XGBoostTuner可以更有效地进行探索并找到更好的配置。

开始调试

在这里,继续矩阵乘法示例。首先,应该创建一个调试任务。还可以检查初始化的搜索空间。在这种情况下,对于512x512方阵乘法,空间大小为10x10 = 100

N, L, M = 512, 512, 512

task = autotvm.task.create("tutorial/matmul", args=(N, L, M, "float32"), target="llvm")

print(task.config_space)

出:

ConfigSpace (len=100, space_map=

   0 tile_y: Split(policy=factors, product=512, num_outputs=2) len=10

   1 tile_x: Split(policy=factors, product=512, num_outputs=2) len=10

)

然后,需要定义如何测量生成的代码并选择一个调谐器。由于的空间很小,所以可以使用随机调谐器。

在本文中,仅进行10个试验进行演示。实际上,可以根据自己的时间预算进行更多试验。会将调试结果记录到日志文件中。此文件可用于稍后获得最佳配置。

# logging config (for printing tuning log to the screen)

logging.getLogger("autotvm").setLevel(logging.DEBUG)

logging.getLogger("autotvm").addHandler(logging.StreamHandler(sys.stdout))

 

# There are two steps for measuring a config: build and run.

# By default, we use all CPU cores to compile program. Then measure them sequentially.

# We measure 5 times and take average to reduce variance.

measure_option = autotvm.measure_option(builder="local", runner=autotvm.LocalRunner(number=5))

 

# Begin tuning with RandomTuner, log records to file `matmul.log`

# You can use alternatives like XGBTuner.

tuner = autotvm.tuner.RandomTuner(task)

tuner.tune(

    n_trial=10,

    measure_option=measure_option,

    callbacks=[autotvm.callback.log_to_file("matmul.log")],

)

出:

Get devices for measurement successfully!

No: 1   GFLOPS: 0.52/0.52       result: MeasureResult(costs=(0.5179643672,), error_no=0, all_cost=8.699557542800903, timestamp=1607225778.9184623)      [('tile_y', [-1, 64]), ('tile_x', [-1, 1])],None,6

No: 2   GFLOPS: 2.05/2.05       result: MeasureResult(costs=(0.1307110214,), error_no=0, all_cost=2.452157735824585, timestamp=1607225781.4836178)      [('tile_y', [-1, 512]), ('tile_x', [-1, 8])],None,39

No: 3   GFLOPS: 2.77/2.77       result: MeasureResult(costs=(0.0968108324,), error_no=0, all_cost=2.015434741973877, timestamp=1607225783.5040994)      [('tile_y', [-1, 2]), ('tile_x', [-1, 8])],None,31

No: 4   GFLOPS: 7.71/7.71       result: MeasureResult(costs=(0.0348177938,), error_no=0, all_cost=0.9887301921844482, timestamp=1607225784.5313203)     [('tile_y', [-1, 1]), ('tile_x', [-1, 32])],None,50

No: 5   GFLOPS: 13.46/13.46     result: MeasureResult(costs=(0.0199451586,), error_no=0, all_cost=0.7833263874053955, timestamp=1607225785.3334467)     [('tile_y', [-1, 256]), ('tile_x', [-1, 64])],None,68

No: 6   GFLOPS: 11.91/13.46     result: MeasureResult(costs=(0.0225446656,), error_no=0, all_cost=0.7622959613800049, timestamp=1607225786.1802726)     [('tile_y', [-1, 256]), ('tile_x', [-1, 512])],None,98

No: 7   GFLOPS: 0.92/13.46      result: MeasureResult(costs=(0.2913359364,), error_no=0, all_cost=5.074311971664429, timestamp=1607225791.3119547)      [('tile_y', [-1, 128]), ('tile_x', [-1, 2])],None,17

No: 8   GFLOPS: 2.37/13.46      result: MeasureResult(costs=(0.1133100596,), error_no=0, all_cost=2.2167930603027344, timestamp=1607225793.595454)      [('tile_y', [-1, 8]), ('tile_x', [-1, 4])],None,23

No: 9   GFLOPS: 11.52/13.46     result: MeasureResult(costs=(0.0233022846,), error_no=0, all_cost=0.7279143333435059, timestamp=1607225795.1428313)     [('tile_y', [-1, 256]), ('tile_x', [-1, 32])],None,58

No: 10  GFLOPS: 14.67/14.67     result: MeasureResult(costs=(0.0182990712,), error_no=0, all_cost=0.7626948356628418, timestamp=1607225795.9127738)     [('tile_y', [-1, 64]), ('tile_x', [-1, 128])],None,76

最后,最好从缓存文件中应用历史记录,并检查其正确性。可以matmul直接在 autotvm.apply_history_best上下文中调用该函数。当调用此函数时,它将使用其参数查询调度上下文,并使用相同的参数获得最佳配置。

# apply history best from log file

with autotvm.apply_history_best("matmul.log"):

    with tvm.target.Target("llvm"):

        s, arg_bufs = matmul(N, L, M, "float32")

        func = tvm.build(s, arg_bufs)

 

# check correctness

a_np = np.random.uniform(size=(N, L)).astype(np.float32)

b_np = np.random.uniform(size=(L, M)).astype(np.float32)

c_np = a_np.dot(b_np)

 

c_tvm = tvm.nd.empty(c_np.shape)

func(tvm.nd.array(a_np), tvm.nd.array(b_np), c_tvm)

 

tvm.testing.assert_allclose(c_np, c_tvm.asnumpy(), rtol=1e-2)

 

上一篇:.webp是什么文件?怎么打开这种文件


下一篇:得物技术在搭建会场下的页面性能优化