Цель -- продемонстрировать возможности использования AutoTVM для автоматической настройки параметров блочного матричного умножения $C = A * B$.
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()
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
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)
Вначале определим функцию gemm_compute, которая реализует матричное умножение
согласно его определению, и создадим пустой план вычислений gemm_schedule.
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.
Скомпилируем и выполним полученные реализации.
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
При оптимизации важно уделить внимание параметрам методов автоматической настройки, поскольку они предоставляют множество возможностей, и постоянно обновляются и улучшаются.
Чем больше раз будет выполнен запуск операторов или подграфа, тем более репрезентативная статистика будет собрана. Это позволит методам автоматической настройки принимать более качественные решения при дальнейшей оптимизации и выборе оптимального алгоритма. Количество итераций зависит от возможностей запуска (допустимого времени) и стабильности оборудования. Важно проверять стабильность запусков.
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)TVM может использовать warm-up -- первый запуск, который не будет включен в статистику.
Для качественной оценки зачастую требуется много испытаний, иногда вплоть до десятков и сотен тысяч итераций. Важно учитывать, что "много" не всегда означает часы работы, в некоторых случаях это может быть и сутки. Для ускорения сходимости автоматической оптимизации можно экспериментировать с параметрами модели затрат и использовать раннюю остановку, если значения целевой функции не меняются или изменяются медленно.
Для автоматического выбора размера блоков (параметры bn и kfactor) воспользуемся
интерфейсом AutoTVM.
При необходимости расширения существующего расписания требуется расширить стратегию для целевой платформы и добавить туда новые функции для вычислений и расписания. Однако, из-за того, что оптимизация матричного умножения выполняется в образовательных целях и, возможно, нет доступа к исходному коду TVM (например, TVM установлен как pip-пакет), то используем декоратор autotvm.template. Этот декоратор позволяет зарегистрировать функцию в AutoTVM. Затем функция вычислений и расписание, определённые в рамках зарегистрированной функции, будут использоваться во время тюнинга для поиска оптимальной конфигурации.
Функция autotvm_gemm будет выполнять роль интерфейса для реализации перебора для AutoTVM
__gemm. Использованные функции AutoTVM:
@autotvm.template("tutorial/matmul")
def autotvm_gemm(M, K, N):
return __gemm(M, K, N)
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]
Получим задачу, которая содержит информацию о пространстве поиска.
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 )
Для проведения автоматической настройки параметров умножения матриц можно использовать
различные методы настройки, а именно RandomTuner и XGBTuner.
Определим метод настройки и количество итераций для настройки, а затем запустим поиск лучшей конфигурации. На выходе будет представлена лучшая полученная конфигурация.
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)],
)
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]
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
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)],
)
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]
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
Можно перебирать не только размер блоков, но и порядок циклов, а также указать, использовать ли параллелизм, разворачивать ли циклы. Попробуем!
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 )
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)],
)
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]
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