Лекция 7. Использование AutoTVM для оптимизации матричного умножения¶

Цель -- продемонстрировать возможности использования AutoTVM для автоматической настройки параметров блочного матричного умножения $C = A * B$.

1. Подготовительный этап¶

1.1 Загрузка необходимых библиотек и определение параметров¶

  • Загрузка библиотек, необходимых для работы с AutoTVM и другими компонентами Apache TVM.
  • Определение размеров матриц и типа данных для хранения матриц.
  • Определение параметров целевого устройства для запуска.
In [1]:
import tvm
import tvm.testing
from tvm import te
import numpy
from tvm import autotvm

M = K = N = 1000

dtype = 'float32'
target = tvm.target.Target('llvm')
dev = tvm.cpu()

1.2 Реализация функций для измерения времени и проверки результатов¶

In [2]:
def timeit(a, b, c, func, dev):
    evaluator = func.time_evaluator(func.entry_name, dev, min_repeat_ms=1000, number=1)
    print(f"Время выполнения: {evaluator(a, b, c).mean:.4f}")

def check(func, M, K, N):
    # Генерация случайных данных для матриц
    a = tvm.nd.array(numpy.random.rand(M, K).astype(dtype), dev)
    b = tvm.nd.array(numpy.random.rand(K, N).astype(dtype), dev)

    c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)
    func(a, b, c)

    # Проверка результатов с помощью numpy
    answer = numpy.dot(a.numpy(), b.numpy())
    tvm.testing.assert_allclose(c.numpy(), answer, rtol=1e-5)
    return a, b, c

1.3 Функция для сборки и запуска оптимизированной функции умножения матриц¶

In [3]:
def build_and_run(s, args):
    A, B, C = args
    func = tvm.build(s, [A, B, C], target=target, name="mmult")

    # Вывод промежуточного представления TIR
    print(tvm.lower(s, [A, B, C], simple_mode=True))
    a, b, c = check(func, M, K, N)
    timeit(a, b, c, func, dev)

2. Наивный и блочный алгоритм матричного умножения¶

Вначале определим функцию gemm_compute, которая реализует матричное умножение согласно его определению, и создадим пустой план вычислений gemm_schedule.

In [4]:
def gemm_compute(M, K, N):
    k = te.reduce_axis((0, K), "k")
    A = te.placeholder((M, K), name="A")
    B = te.placeholder((K, N), name="B")
    C = te.compute((M, N), lambda m, n: te.sum(A[m, k] * B[k, n], axis=k), name="C")

    return A, B, C

def gemm_schedule(args):
    C = args[-1]
    s = te.create_schedule(C.op)
    return s
    
A, B, C = gemm_compute(M, K, N)
s = gemm_schedule((A, B, C))

build_and_run(s, [A, B, C])
# from tvm.script import ir as I
# from tvm.script import tir as T

@I.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer((1000, 1000), "float32"), B: T.Buffer((1000, 1000), "float32"), C: T.Buffer((1000, 1000), "float32")):
        T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)})
        for m, n in T.grid(1000, 1000):
            C_1 = T.Buffer((1000000,), data=C.data)
            C_1[m * 1000 + n] = T.float32(0)
            for k in range(1000):
                cse_var_2: T.int32 = m * 1000
                cse_var_1: T.int32 = cse_var_2 + n
                A_1 = T.Buffer((1000000,), data=A.data)
                B_1 = T.Buffer((1000000,), data=B.data)
                C_1[cse_var_1] = C_1[cse_var_1] + A_1[cse_var_2 + k] * B_1[k * 1000 + n]
Время выполнения: 0.9166

Теперь рассмотрим более эффективный блочный алгоритм. Для оптимизации вычислений разбьем матрицы на блоки и выполним умножение в этих блоках. Для блочного умножение будем использовать новое расписание gemm_block_schedule и старую функцию вычисления gemm_compute.

Скомпилируем и выполним полученные реализации.

In [5]:
def gemm_block_schedule(args, bn, kfactor):
    C = args[-1]

    s = te.create_schedule(C.op)
    mo, no, mi, ni = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
    ko, ki = s[C].split(s[C].op.reduce_axis[0], factor=kfactor)

    s[C].reorder(mo, no, ko, ki, mi, ni)
    return s


bn = 30
kfactor = 10

A, B, C = gemm_compute(M, K, N)
s = gemm_block_schedule((A, B, C), bn, kfactor)

build_and_run(s, [A, B, C])
# from tvm.script import ir as I
# from tvm.script import tir as T

@I.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer((1000, 1000), "float32"), B: T.Buffer((1000, 1000), "float32"), C: T.Buffer((1000, 1000), "float32")):
        T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)})
        for m_outer, n_outer in T.grid(34, 34):
            C_1 = T.Buffer((1000000,), data=C.data)
            for m_inner_init in range(30):
                if T.likely(m_outer * 3 + m_inner_init // 10 < 100):
                    for n_inner_init in range(30):
                        if T.likely(n_outer * 3 + n_inner_init // 10 < 100):
                            C_1[m_outer * 30000 + m_inner_init * 1000 + n_outer * 30 + n_inner_init] = T.float32(0)
            for k_outer, k_inner, m_inner in T.grid(100, 10, 30):
                if T.likely(m_outer * 3 + m_inner // 10 < 100):
                    for n_inner in range(30):
                        if T.likely(n_outer * 3 + n_inner // 10 < 100):
                            cse_var_3: T.int32 = n_outer * 30
                            cse_var_2: T.int32 = m_outer * 30000 + m_inner * 1000
                            cse_var_1: T.int32 = cse_var_2 + cse_var_3 + n_inner
                            A_1 = T.Buffer((1000000,), data=A.data)
                            B_1 = T.Buffer((1000000,), data=B.data)
                            C_1[cse_var_1] = C_1[cse_var_1] + A_1[cse_var_2 + k_outer * 10 + k_inner] * B_1[k_outer * 10000 + k_inner * 1000 + cse_var_3 + n_inner]
Время выполнения: 0.6032

3. Параметры оптимизации и общие рекомендации¶

При оптимизации важно уделить внимание параметрам методов автоматической настройки, поскольку они предоставляют множество возможностей, и постоянно обновляются и улучшаются.

3.1 Количество запусков для измерения времени¶

Чем больше раз будет выполнен запуск операторов или подграфа, тем более репрезентативная статистика будет собрана. Это позволит методам автоматической настройки принимать более качественные решения при дальнейшей оптимизации и выборе оптимального алгоритма. Количество итераций зависит от возможностей запуска (допустимого времени) и стабильности оборудования. Важно проверять стабильность запусков.

3.2 Параметры проведения измерений в TVM¶

TVM, как правило, имеет два параметра, отвечающих за количество запусков: number и repeat. Пример использования:

for r in range(repeat):
    time_start = now()
    for n in range(number):
        func_name()
    time_end = now()
    total_times.append ((time_end - time_start)/number)

3.3 Использование warm-up¶

TVM может использовать warm-up -- первый запуск, который не будет включен в статистику.

3.4 Надёжная оценка и количество испытаний¶

Для качественной оценки зачастую требуется много испытаний, иногда вплоть до десятков и сотен тысяч итераций. Важно учитывать, что "много" не всегда означает часы работы, в некоторых случаях это может быть и сутки. Для ускорения сходимости автоматической оптимизации можно экспериментировать с параметрами модели затрат и использовать раннюю остановку, если значения целевой функции не меняются или изменяются медленно.

5. Автоматическая настройка размера блока¶

Для автоматического выбора размера блоков (параметры bn и kfactor) воспользуемся интерфейсом AutoTVM.

При необходимости расширения существующего расписания требуется расширить стратегию для целевой платформы и добавить туда новые функции для вычислений и расписания. Однако, из-за того, что оптимизация матричного умножения выполняется в образовательных целях и, возможно, нет доступа к исходному коду TVM (например, TVM установлен как pip-пакет), то используем декоратор autotvm.template. Этот декоратор позволяет зарегистрировать функцию в AutoTVM. Затем функция вычислений и расписание, определённые в рамках зарегистрированной функции, будут использоваться во время тюнинга для поиска оптимальной конфигурации.

Функция autotvm_gemm будет выполнять роль интерфейса для реализации перебора для AutoTVM __gemm. Использованные функции AutoTVM:

  • autotvm.get_config() -- получение текущей конфигурации AutoTVM и добавление туда новых переменных для пространства поиска.
  • define_split -- функция, которая определяет пространство для перебора коэффициента разбиения цикла.
  • define_knob -- функция, которая позволяет определить набор параметров, перебираемых во время настройки.
  • define_annotate -- функция, которая позволяет обозначить циклы для последующего развертывания.
  • define_reorder -- функция, которая позволяет указать, что AutoTVM может переставлять вложенные циклы местами для поиска оптимальной конфигурации.
In [6]:
@autotvm.template("tutorial/matmul")
def autotvm_gemm(M, K, N):
    return __gemm(M, K, N)
In [7]:
def __gemm(M, K, N):
    A, B, C = gemm_compute(M, K, N)
    s = te.create_schedule(C.op)

    m, n = C.op.axis[0], C.op.axis[1]
    k = s[C].op.reduce_axis[0]

    cfg = autotvm.get_config()
    cfg.define_split("tile_m", m, num_outputs=2)
    cfg.define_split("tile_n", n, num_outputs=2)
    cfg.define_split("tile_k", k, num_outputs=2)

    mo, mi = cfg["tile_m"].apply(s, C, m)
    no, ni = cfg["tile_n"].apply(s, C, n)
    ko, ki = cfg["tile_k"].apply(s, C, k)

    s[C].reorder(mo, no, ko, ki, mi, ni)

    return s, [A, B, C]

Получим задачу, которая содержит информацию о пространстве поиска.

In [8]:
task = autotvm.task.create("tutorial/matmul", args=(N, M, K), target=target)
print(task.config_space)
ConfigSpace (len=4096, range_length=4096, space_map=
   0 tile_m: Split(policy=factors, product=1000, num_outputs=2) len=16
   1 tile_n: Split(policy=factors, product=1000, num_outputs=2) len=16
   2 tile_k: Split(policy=factors, product=1000, num_outputs=2) len=16
)

5.1 Эксперименты¶

Для проведения автоматической настройки параметров умножения матриц можно использовать различные методы настройки, а именно RandomTuner и XGBTuner.

Определим метод настройки и количество итераций для настройки, а затем запустим поиск лучшей конфигурации. На выходе будет представлена лучшая полученная конфигурация.

5.1.1 RandomTuner¶

In [9]:
log_file = "gemm_random.log"

tuner = autotvm.tuner.RandomTuner(task)
tuner.tune(
    n_trial=100,
    measure_option=autotvm.measure_option(builder="local", runner=autotvm.LocalRunner(repeat=1, number=3)),
    callbacks=[autotvm.callback.log_to_file(log_file)],
)
In [10]:
with autotvm.apply_history_best(log_file):
    with target:
        s, args = autotvm_gemm(N, M, K)
        func = tvm.build(s, args, target=target, name="mmult")

        a, b, c = check(func, M, K, N)
        timeit(a, b, c, func, dev)
        print(tvm.lower(s, args, simple_mode=True))
Время выполнения: 0.0691
# from tvm.script import ir as I
# from tvm.script import tir as T

@I.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer((1000, 1000), "float32"), B: T.Buffer((1000, 1000), "float32"), C: T.Buffer((1000, 1000), "float32")):
        T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)})
        for m_outer, n_outer in T.grid(50, 10):
            C_1 = T.Buffer((1000000,), data=C.data)
            for m_inner_init, n_inner_init in T.grid(20, 100):
                C_1[m_outer * 20000 + m_inner_init * 1000 + n_outer * 100 + n_inner_init] = T.float32(0)
            for k_outer, k_inner, m_inner, n_inner in T.grid(50, 20, 20, 100):
                cse_var_3: T.int32 = n_outer * 100
                cse_var_2: T.int32 = m_outer * 20000 + m_inner * 1000
                cse_var_1: T.int32 = cse_var_2 + cse_var_3 + n_inner
                A_1 = T.Buffer((1000000,), data=A.data)
                B_1 = T.Buffer((1000000,), data=B.data)
                C_1[cse_var_1] = C_1[cse_var_1] + A_1[cse_var_2 + k_outer * 20 + k_inner] * B_1[k_outer * 20000 + k_inner * 1000 + cse_var_3 + n_inner]
In [11]:
with autotvm.apply_history_best(log_file) as best:
    best_config = best.query(task.target, task.workload)
    print(f'пространство поиска: {best_config._entity_map}')    
    print(f'производительность: {best_config.cost:.4f}')
пространство поиска: OrderedDict([('tile_m', [-1, 20]), ('tile_n', [-1, 100]), ('tile_k', [-1, 20])])
производительность: 0.0700

5.2 XGBTuner¶

In [12]:
log_file = "gemm_xgb.log"

tuner = autotvm.tuner.XGBTuner(task, loss_type="reg")
tuner.tune(
    n_trial=100,
    measure_option=autotvm.measure_option(builder="local", runner=autotvm.LocalRunner(repeat=1, number=3)),
    callbacks=[autotvm.callback.log_to_file(log_file)],
)
In [13]:
with autotvm.apply_history_best(log_file):
    with target:
        s, args = autotvm_gemm(N, M, K)
        func = tvm.build(s, args, target=target, name="mmult")

        a, b, c = check(func, M, K, N)
        timeit(a, b, c, func, dev)
        print(tvm.lower(s, args, simple_mode=True))
Время выполнения: 0.0661
# from tvm.script import ir as I
# from tvm.script import tir as T

@I.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer((1000, 1000), "float32"), B: T.Buffer((1000, 1000), "float32"), C: T.Buffer((1000, 1000), "float32")):
        T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)})
        for m_outer, n_outer in T.grid(20, 10):
            C_1 = T.Buffer((1000000,), data=C.data)
            for m_inner_init, n_inner_init in T.grid(50, 100):
                C_1[m_outer * 50000 + m_inner_init * 1000 + n_outer * 100 + n_inner_init] = T.float32(0)
            for k_outer, k_inner, m_inner, n_inner in T.grid(50, 20, 50, 100):
                cse_var_3: T.int32 = n_outer * 100
                cse_var_2: T.int32 = m_outer * 50000 + m_inner * 1000
                cse_var_1: T.int32 = cse_var_2 + cse_var_3 + n_inner
                A_1 = T.Buffer((1000000,), data=A.data)
                B_1 = T.Buffer((1000000,), data=B.data)
                C_1[cse_var_1] = C_1[cse_var_1] + A_1[cse_var_2 + k_outer * 20 + k_inner] * B_1[k_outer * 20000 + k_inner * 1000 + cse_var_3 + n_inner]
In [14]:
with autotvm.apply_history_best(log_file) as best:
    best_config = best.query(task.target, task.workload)
    print(f'пространство поиска: {best_config._entity_map}')    
    print(f'производительность: {best_config.cost:.4f}')
пространство поиска: OrderedDict([('tile_m', [-1, 50]), ('tile_n', [-1, 100]), ('tile_k', [-1, 20])])
производительность: 0.0693

6. Расширение возможностей¶

Можно перебирать не только размер блоков, но и порядок циклов, а также указать, использовать ли параллелизм, разворачивать ли циклы. Попробуем!

In [15]:
def __gemm(M, K, N):
    A, B, C = gemm_compute(M, K, N)
    s = te.create_schedule(C.op)

    m, n = C.op.axis[0], C.op.axis[1]
    k = s[C].op.reduce_axis[0]

    cfg = autotvm.get_config()
    cfg.define_split("tile_m", m, num_outputs=2)
    cfg.define_split("tile_n", n, num_outputs=2)
    cfg.define_split("tile_k", k, num_outputs=2)

    mo, mi = cfg["tile_m"].apply(s, C, m)
    no, ni = cfg["tile_n"].apply(s, C, n)
    ko, ki = cfg["tile_k"].apply(s, C, k)

    s[C].reorder(mo, no, ko, ki, mi, ni)

    cfg.define_annotate("unroll", [ki, mi, ni], policy="try_unroll")
    cfg.define_reorder("reorder", [ki, mi, ni], policy="all")

    cfg["reorder"].apply(s, C, [ki, mi, ni])
    cfg["unroll"].apply(s, C, [ki, mi, ni])

    cfg.define_knob("parallel", [False, True])
    outer_axis = s[C].fuse(mo, no, ko)
    if cfg["parallel"].val:
        s[C].parallel(outer_axis)

    return s, [A, B, C]


task = autotvm.task.create("tutorial/matmul", args=(N, M, K), target=target)
print(task.config_space)
ConfigSpace (len=393216, range_length=393216, space_map=
   0 tile_m: Split(policy=factors, product=1000, num_outputs=2) len=16
   1 tile_n: Split(policy=factors, product=1000, num_outputs=2) len=16
   2 tile_k: Split(policy=factors, product=1000, num_outputs=2) len=16
   3 unroll: Annotate(policy=try_unroll) len=8
   4 reorder: Reorder(policy=all) len=6
   5 parallel: OtherOption([False, True]) len=2
)
In [16]:
log_file = 'gemm_extended.log'

tuner = autotvm.tuner.XGBTuner(task, loss_type="reg")
tuner.tune(
    n_trial=100,
    measure_option=autotvm.measure_option(builder="local", runner=autotvm.LocalRunner(repeat=1, number=3)),
    callbacks=[autotvm.callback.log_to_file(log_file)],
)
In [17]:
with autotvm.apply_history_best(log_file):
    with target:
        s, args = autotvm_gemm(N, M, K)
        func = tvm.build(s, args, target=target, name="mmult")

        a, b, c = check(func, M, K, N)
        timeit(a, b, c, func, dev)
        print(tvm.lower(s, args, simple_mode=True))
Время выполнения: 0.0195
# from tvm.script import ir as I
# from tvm.script import tir as T

@I.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer((1000, 1000), "float32"), B: T.Buffer((1000, 1000), "float32"), C: T.Buffer((1000, 1000), "float32")):
        T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)})
        C_1 = T.Buffer((1000000,), data=C.data)
        for m_outer_n_outer_fused_k_outer_fused_init in T.parallel(312500):
            for m_inner_init, n_inner_init in T.grid(8, 20):
                C_1[m_outer_n_outer_fused_k_outer_fused_init // 2500 * 8000 + m_inner_init * 1000 + m_outer_n_outer_fused_k_outer_fused_init % 2500 // 50 * 20 + n_inner_init] = T.float32(0)
        for m_outer_n_outer_fused_k_outer_fused in T.parallel(312500):
            for m_inner, n_inner, k_inner in T.grid(8, 20, 20):
                cse_var_4: T.int32 = m_outer_n_outer_fused_k_outer_fused % 2500 // 50 * 20
                cse_var_3: T.int32 = m_outer_n_outer_fused_k_outer_fused // 2500 * 8000 + m_inner * 1000
                cse_var_2: T.int32 = m_outer_n_outer_fused_k_outer_fused % 50
                cse_var_1: T.int32 = cse_var_3 + cse_var_4 + n_inner
                A_1 = T.Buffer((1000000,), data=A.data)
                B_1 = T.Buffer((1000000,), data=B.data)
                C_1[cse_var_1] = C_1[cse_var_1] + A_1[cse_var_3 + cse_var_2 * 20 + k_inner] * B_1[cse_var_2 * 20000 + k_inner * 1000 + cse_var_4 + n_inner]
In [18]:
with autotvm.apply_history_best(log_file) as best:
    best_config = best.query(task.target, task.workload)
    print(f'пространство поиска: {best_config._entity_map}')    
    print(f'производительность: {best_config.cost:.4f}')
пространство поиска: OrderedDict([('tile_m', [-1, 8]), ('tile_n', [-1, 20]), ('tile_k', [-1, 20]), ('unroll', ['none', 'none', 'none']), ('reorder', [1, 2, 0]), ('parallel', True)])
производительность: 0.0193