Практическая работа №3. Реализация и оптимизация тензорного выражения на примере умножения матриц¶

1. Цели и задачи работы¶

Цель работы – изучить программный интерфейс для реализации и оптимизации слоев нейронной сети в Apache TVM на примере запуска на процессорах архитектуры x86-64, а также рассмотреть возможность использования оптимизированных библиотек для матричных операций.

Достижение указанной цели предполагает решение следующих задач:

  1. Сборка и установка необходимых инструментов и библиотек. Сборка LLVM, OpenBLAS для линковки и вызова высокопроизводительного матричного умножения. Сборка Apache TVM с LLVM и OpenBLAS.

  2. Реализация "наивного" умножения матриц по определению с помощью Apache TVM (тензорное выражение с пустым планом вычислений). Разработка вспомогательной функции (обертки) для компиляции тензорного выражения, а также обеспечение измерения времени выполнения реализации и проверки корректности посредством сравнения с библиотекой NumPy.

  3. Рассмотрение нескольких этапов оптимизации матричного умножения с помощью реализации новых планов вычислений.

    1. Перестановка циклов для устранения непоследовательного доступа к памяти. Вызов функции для компиляции, обеспечение измерения времени выполнения реализации и проверки корректности посредством сравнения с библиотекой NumPy.
    2. Реализация блочного алгоритма для улучшения повторного использования данных. Вызов функции для компиляции, обеспечение измерения времени выполнения реализации и проверки корректности посредством сравнения с библиотекой NumPy.
  4. Замена тензорного выражения на вызов функции матричного умножения из высокопроизводительной библиотеки OpenBLAS. Вызов функции для компиляции, обеспечение измерения времени выполнения реализации и проверки корректности посредством сравнения с библиотекой NumPy.

  5. Сравнение производительности разработанных реализаций.

2. Постановка задачи умножения матриц¶

Пусть даны две прямоугольные матрицы $A$ и $B$ с размерами $ m\times k $ и $ k \times n $ соответственно. Результатом умножения матриц $A$ и $B$ является матрица $C$ размером $ m \times n $. Элемент $c_{ij}$ матрицы $C$ вычисляется по формуле:

$$ c_{ij} = \sum_{k=1}^{n} a_{ik} \cdot b_{kj}$$

3. Сборка Apache TVM¶

3.1. Сборка LLVM¶

В данной практической работе не используется кросс-компиляция моделей или слоев. Компиляция происходит на устройстве с архитектурой x86-64. Рекомендуется использовать версию 15 <= LLVM <= 17.

3.1.1. Установка с помощью менеджера пакетов¶

Установка с помощью менеджера пакетов предполагает выполнение следующей команды:

sudo apt install clang-17 llvm-17*

3.1.2. Сборка LLVM из исходных кодов (для версии llvmorg-17.0.6)¶

Для сборка LLVM из исходных кодов требуется загрузить необходимую версию LLVM из репозитория GitHub. В данной работе используется версия llvmorg-17.0.6, далее, используя утилиту CMake, сгенерировать make-файлы и выполнить сборку. Ниже приведена соответствующая последовательность команд.

git clone https://github.com/llvm/llvm-project.git -b llvmorg-17.0.6
cd llvm-project

mkdir _build
cd _build

cmake -DCMAKE_BUILD_TYPE="Release" \
      -DLLVM_ENABLE_PROJECTS=clang \
      -DBUILD_SHARED_LIBS=True \
      -DLLVM_USE_SPLIT_DWARF=True \
      -DCMAKE_INSTALL_PREFIX="../../_install" ../llvm

make

Примечание: в случае сборки LLVM из исходных кодов перед сборкой Apache TVM необходимо указать путь к LLVM в переменной окружения PATH и создать переменную окруженияLLVM_CONFIG. Ниже показан пример.

PATH="~/llvm-project/_build/bin:$PATH"
export LLVM_CONFIG=~/llvm-project/_build/bin/llvm-config

3.2. Установка OpenBLAS¶

Далее необходимо установить OpenBLAS, используя менеджер пакетов.

sudo apt-get install libopenblas-dev

3.3. Настройка окружения Python¶

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

python3 -m venv ~/tvm_cpu/
source ~/tvm_cpu/bin/activate

pip install scipy numpy matplotlib pandas
pip install cloudpickle traitlets typing-extensions psutil pybind11 decorator attrs 
pip install notebook

3.4. Сборка Apache TVM¶

Для сборки Apache TVM не обязательно использовать созданную виртуальную среду для Python.

git clone --recursive https://github.com/apache/tvm
cd tvm

mkdir build
cd build

cmake -DUSE_LLVM=ON -DUSE_BLAS=openblas ..
make

3.5. Активация окружения для практической работы¶

Для активации виртуальной среды с целью решения задач практической работы необходимо выполнить следующие команды:

source ~/tvm_cpu/bin/activate
export PYTHONPATH=<PATH TO TVM>/python:${PYTHONPATH}

4. Программная реализация¶

Для использования функционала Apache TVM и других вспомогательных библиотек импортируем необходимые пакеты.

Также определим переменную, содержащую используемый тип данных для элементов тензоров - float32, и установим в качестве целевого устройства для запуска CPU.

In [1]:
import numpy as np
import tvm
from tvm import te

import matplotlib.pyplot as plt

M = K = N = 640

dtype = 'float32'
dev = tvm.cpu()

4.1. Строка компиляции¶

На данном этапе определим строку компиляции target.

Для определения архитектуры устройства необходимо создать обьект строки компиляции по умолчанию для LLVM - tvm.target.Target('llvm'). Далее с помощью атрибута mtriple выбрать строку компиляции:

  • Если атрибут mtriple отсутствует или содержит подстроку x86_64, используется стандартная строка компиляции llvm.
  • Если mtriple содержит подстроку riscv64, используется строка компиляции llvm -jit=orcjit -mtriple=riscv64-unknown-linux-gnu -mcpu=generic-rv64 -mabi=lp64d -mattr=+64bit,+m,+a,+f,+d.
  • В противном случае генерируется исключение.

Примечание: если TVM устанавливался через PyPI, то mtriple пустой.

In [2]:
print(f"mtriple устройства {tvm.target.Target('llvm').attrs.get('mtriple')}")

if tvm.target.Target('llvm').attrs.get('mtriple') is None or 'x86_64' in tvm.target.Target('llvm').attrs.get('mtriple'):
    target = tvm.target.Target('llvm')
elif 'riscv64' in tvm.target.Target('llvm').attrs.get('mtriple'):
    target = tvm.target.Target(
        'llvm -jit=orcjit -mtriple=riscv64-unknown-linux-gnu '
        '-mcpu=generic-rv64 -mabi=lp64d -mattr=+64bit,+m,+a,+f,+d'
    )
else:
    raise ValueError("Unsupported architecture")

print(f'{target = }')
mtriple устройства x86_64-pc-linux-gnu
target = llvm -keys=cpu -mtriple=x86_64-pc-linux-gnu

4.2. Вспомогательные функции¶

Выполним реализацию следующих вспомогательных функций.

  1. timeit - функция для измерения времени выполнения умножения матриц с использованием API TVM. Для замера времени используется функция func.time_evaluator, которое позволяет запустить несколько раз тензорное выражение и получить различную статистику. В данном случае настройки запуска следующие:

    1. Тензорное выражение запускается такое количество раз, чтобы суммарное время работы составляло 1 секунду (min_repeat_ms=1000). Это удобно, поскольку времена между различными версиями могут значительно различаться. Также, вместо данного параметра можно указать фиксированное число запусков через параметр number.
    2. Такая процедура повторяется пять раз (repeat=5). Для оценки времени используется среднее значение по всем процедурам.
  2. check - функция для проверки корректности умножения матриц. В функции генерируются случайные данные. После этого выполняется умножение как с помощью реализованной версии, так и с использованием функций библиотеки NumPy. Корректность определяется посредством сравнения результатов умножения.

  3. build_and_run - функция, которая реализует следующую последовательность действий:

    1. Компиляция тензорного выражения.
    2. Запуск операции умножения матриц.
    3. Измерение времени работы.
    4. Проверка корректности результата.

Примечание: псевдокод работы методов оценки времени работы в Apache TVM может быть записан следующим образом.

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) 
In [3]:
def timeit(a, b, c, func, dev):
    evaluator = func.time_evaluator(func.entry_name, dev, min_repeat_ms=1000, repeat=5)
    result = evaluator(a, b, c)
    print(f"\n\n\t Время выполнения: {result.mean:.4f} с")
    return result


def check(func, M, K, N):
    a = tvm.nd.array(np.random.rand(M, K).astype(dtype), dev)
    b = tvm.nd.array(np.random.rand(K, N).astype(dtype), dev)

    c = tvm.nd.array(np.zeros((M, N), dtype=dtype), dev)
    # вызов тензорного выражения
    func(a, b, c)

    expected = np.dot(a.numpy(), b.numpy())
    np.testing.assert_allclose(c.numpy(), expected, rtol=1e-5)
    return a, b, c


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

    # tvm.lower возвращает TIR(TensorIR) тензорного выражения. 
    print(tvm.lower(s, [A, B, C], simple_mode=True))
    a, b, c = check(func, M, K, N)
    return timeit(a, b, c, func, dev)

5. Программные реализации¶

5.1. "Наивное" матричное умножение¶

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

  1. Для умножения матриц определяются размеры двух входных матриц, а также цикл редукции. С помощью функции te.compute определяется выражение $A[m, k] * B[k, n]$ для вычисления выходной матрицы. Реализация тензорного выражения далее не изменяется, за исключением использования версии из OpenBLAS.

  2. Для реализации плана вычислений необходимо описать входные данные с помощью метода te.placeholder и при необходимости циклы редукции с помощью метода te.reduce_axis. После этого в функции te.compute необходимо указать размер выходного тензора и выражение для его вычисления. Можно использовать несколько последовательных или параллельных функций te.compute. Такая реализация не содержит оптимизаций. Для наивного матричного умножения достаточно пустого плана вычислений, который генерируется по умолчанию с помощью te.create_schedule(C.op). Сгенерированный TIR будет в точности соответствовать тензорному выражению в gemm_compute.

In [4]:
def gemm_compute(
    M: int, K: int, N: int
) -> tuple[tvm.te.tensor.Tensor, tvm.te.tensor.Tensor, tvm.te.tensor.Tensor]:
    """    
    Параметры:
        M: Число строк в матрице A.
        K: Число столбцов в матрице A и строк в матрице B.
        N: Число столбцов в матрице B.

    Возвращаемое значение:
        Три тензора: два входных тензора и результирующий тензор
    """
    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(С: tvm.te.tensor.Tensor) -> tvm.te.schedule.Schedule:
    """
    Параметры:
        C: Тензор результирующей матрицы.
        
    Возвращаемое значение:
        План вычислений.
    """
    s = te.create_schedule(C.op)
    return s

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

In [5]:
A, B, C = gemm_compute(M, K, N)
s = gemm_schedule(C)

naive_result = 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((640, 640), "float32"), B: T.Buffer((640, 640), "float32"), C: T.Buffer((640, 640), "float32")):
        T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)})
        for m, n in T.grid(640, 640):
            C_1 = T.Buffer((409600,), data=C.data)
            C_1[m * 640 + n] = T.float32(0.0)
            for k in range(640):
                cse_var_2: T.int32 = m * 640
                cse_var_1: T.int32 = cse_var_2 + n
                A_1 = T.Buffer((409600,), data=A.data)
                B_1 = T.Buffer((409600,), data=B.data)
                C_1[cse_var_1] = C_1[cse_var_1] + A_1[cse_var_2 + k] * B_1[k * 640 + n]


	 Время выполнения: 0.2298 с

5.2. Устранение непоследовательного доступа к памяти¶

Следующая задача практической работы - реализация плана вычислений с последовательным доступом к памяти.

Чтобы повысить эффективность алгоритма умножения матриц, можно изменить порядок циклов. Такая перестановка приведет к значительному ускорению выполнения за счет оптимизации доступа к памяти. В исходной реализации доступ к памяти для матрицы $B$ был регулярным, но непоследовательным. При итерации по переменной $k$ шаг доступа к элементам матрицы $B$ составлял $K$ элементов. Перестановка циклов в порядке i,k,j обеспечивает последовательный доступ к памяти, что существенно улучшает производительность за счет более эффективного использования кэш-памяти процессора.

Это изменение не требует модификации тензорного выражения, необходимо скорректировать план вычислений. Для перестановки циклов можно использовать метод s[C].reorder, в который необходимо передать новый порядок циклов. Результирующая матрица содержит информацию о циклах, в частности, C.op.axis[0] позволяет получить объект первого цикла.

In [6]:
def gemm_schedule(С: tvm.te.tensor.Tensor) -> tvm.te.schedule.Schedule:
    """
    Параметры:
        C: Тензор результирующей матрицы.
        
    Возвращаемое значение:
        План вычислений.
    """
    s = te.create_schedule(C.op)

    s[C].reorder(C.op.axis[0], s[C].op.reduce_axis[0], C.op.axis[1])

    return s

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

In [7]:
A, B, C = gemm_compute(M, K, N)
s = gemm_schedule(C)
reorder_result = 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((640, 640), "float32"), B: T.Buffer((640, 640), "float32"), C: T.Buffer((640, 640), "float32")):
        T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)})
        for m in range(640):
            C_1 = T.Buffer((409600,), data=C.data)
            for n_init in range(640):
                C_1[m * 640 + n_init] = T.float32(0.0)
            for k, n in T.grid(640, 640):
                cse_var_2: T.int32 = m * 640
                cse_var_1: T.int32 = cse_var_2 + n
                A_1 = T.Buffer((409600,), data=A.data)
                B_1 = T.Buffer((409600,), data=B.data)
                C_1[cse_var_1] = C_1[cse_var_1] + A_1[cse_var_2 + k] * B_1[k * 640 + n]


	 Время выполнения: 0.0282 с

5.3. Блочный алгоритм умножения матриц¶

Следующая задача работы предполагает реализацию плана вычислений для блочного алгоритма умножения матриц. Реализуем блочный алгоритм, чтобы повысить эффективность использования кэш-памяти. Этот метод разбивает большие матрицы на более мелкие блоки, которые лучше соответствуют размерам кэша. Блочный алгоритм позволяет каждому блоку данных, загруженному в кэш, использоваться многократно до его вытеснения. Это приводит к улучшению повторного использования данных и уменьшению времени выполнения алгоритма.

Для этого необходимо изменить план, добавив разделение на блоки s[C].tile, разделение цикла редукции на два можно выполнить с помощью метода s[C].split, а переупорядочение циклов - по аналогии с предыдущим разделом.

In [8]:
def gemm_schedule(С: tvm.te.tensor.Tensor, bn: int = 8, kfactor: int = 8) -> tvm.te.schedule.Schedule:
    """
    Параметры:
        C: Тензор результирующей матрицы.
        bn: Фактор разбиения для осей матрицы.
        kfactor: Фактор разбиения для оси редукции.
        
    Возвращаемое значение:
        План вычислений.
    """
    
    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

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

In [9]:
bn = 32
kfactor = 32

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

block_result = 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((640, 640), "float32"), B: T.Buffer((640, 640), "float32"), C: T.Buffer((640, 640), "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, 20):
            C_1 = T.Buffer((409600,), data=C.data)
            for m_inner_init, n_inner_init in T.grid(32, 32):
                C_1[m_outer * 20480 + m_inner_init * 640 + n_outer * 32 + n_inner_init] = T.float32(0.0)
            for k_outer, k_inner, m_inner, n_inner in T.grid(20, 32, 32, 32):
                cse_var_3: T.int32 = n_outer * 32
                cse_var_2: T.int32 = m_outer * 20480 + m_inner * 640
                cse_var_1: T.int32 = cse_var_2 + cse_var_3 + n_inner
                A_1 = T.Buffer((409600,), data=A.data)
                B_1 = T.Buffer((409600,), data=B.data)
                C_1[cse_var_1] = C_1[cse_var_1] + A_1[cse_var_2 + k_outer * 32 + k_inner] * B_1[k_outer * 20480 + k_inner * 640 + cse_var_3 + n_inner]


	 Время выполнения: 0.0199 с

6. Повышение производительности с помощью вызова высокопроизводительной реализации из OpenBLAS¶

Для дальнейшей оптимизации можно пробовать разные существующие подходы: подбирать размеры блоков, добавлять многопоточность, новые уровни блоков и т.д. Другим вариантом оптимизации является подключение готовых высокопроизводительных реализаций. Apache TVM поддерживает использование OpenBLAS, cuBLAS, MKL, DNNL и других пакетов. В результате, можно как вызывать конкретную математическую операцию, так и встроить вызовы математических библиотек в нейронную сеть.

Модуль tvm.contrib предоставляет интерфейс для вызова различных математических библиотек. Для вызова реализации OpenBLAS на языке C используется CBLAS.

In [10]:
def gemm_compute(
    M: int, K: int, N: int
) -> tuple[tvm.te.tensor.Tensor, tvm.te.tensor.Tensor, tvm.te.tensor.Tensor]:
    """    
    Параметры:
        M: Число строк в матрице A.
        K: Число столбцов в матрице A и строк в матрице B.
        N: Число столбцов в матрице B.

    Возвращаемое значение:
        Три тензора: два входных тензора и результирующий тензор
    """
    A = te.placeholder((M, K), name="A")
    B = te.placeholder((K, N), name="B")

    C = tvm.contrib.cblas.matmul(A, B)
    return A, B, C

def gemm_schedule(С):
    """
    Параметры:
        C: Тензор результирующей матрицы.
        
    Возвращаемое значение:
        План вычислений.
    """
    s = te.create_schedule(C.op)
    return s

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

In [11]:
A, B, C = gemm_compute(M, K, N)
s = gemm_schedule(C)

openblas_result = 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((640, 640), "float32"), B: T.Buffer((640, 640), "float32"), C: T.Buffer((640, 640), "float32")):
        T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)})
        T.attr(0, "extern_scope", 0)
        T.call_packed("tvm.contrib.cblas.matmul", T.tvm_stack_make_array(A.data, T.tvm_stack_make_shape(640, 640), 0, 2, T.float32(0.0), 0), T.tvm_stack_make_array(B.data, T.tvm_stack_make_shape(640, 640), 0, 2, T.float32(0.0), 0), T.tvm_stack_make_array(C.data, T.tvm_stack_make_shape(640, 640), 0, 2, T.float32(0.0), 0), T.bool(False), T.bool(False))


	 Время выполнения: 0.0010 с

7. Сравнение производительности разработанных реализаций¶

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

In [12]:
print(naive_result)
print(reorder_result)
print(block_result)
print(openblas_result)
Execution time summary:
 mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  
  229.8411     229.4699     232.1100     228.8213      1.1892                  
Execution time summary:
 mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  
  28.2098      28.1506      28.5591      28.0885       0.1765                  
Execution time summary:
 mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  
  19.9282      20.0377      20.5360      19.2591       0.4726                  
Execution time summary:
 mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  
   1.0267       1.0395       1.1441       0.8829       0.0837                  
In [13]:
plt.rcParams['font.size'] = 20

fig, ax = plt.subplots(figsize=(15, 5))

name = ['Наивная\nреализация', 'Переупорядочивание\nциклов', 'Блочный\nалгоритм', 'OpenBLAS']
times = [naive_result.mean, reorder_result.mean, block_result.mean, openblas_result.mean]
bar_labels = ['red', 'blue', '_red', 'orange']
bar_colors = ['tab:blue', 'tab:red', 'tab:green', 'tab:orange']

bars = ax.bar(name, times, label=name, color=bar_colors)

for bar, n, time in zip(bars, name, times):
    h = bar.get_height()
    if n != 'OpenBLAS': h = h / 6 
    ax.text(
        bar.get_x() + bar.get_width() / 2,
        h,
        f'{round(time, 4)} с',
        ha='center',
        va='bottom',
        fontsize=20,
    )

ax.xaxis.label.set_size(40)
ax.set_title('Среднее время\nвыполнения (с)', fontsize=18)
plt.yscale('log')
plt.grid()
No description has been provided for this image

Вывод: из графиков видно, что две последовательные оптимизации работы с памятью значительно уменьшают время выполнения алгоритма умножения матриц. Лучшую производительность демонстрирует реализация из OpenBLAS. Поэтому, если есть возможность использования высокопроизводительной библиотеки, стоит рассмотреть ее использование.

Примечание: при использовании реализаций из внешних библиотек TVM не может оптимизировать сами реализации, а также не может применять некоторые оптимизации на уровне графа. Это может замедлить суммарное время выполнения графа вычислений нейронной сети.

In [ ]: