##01 Бенчмарк из четырех JIT-бэкендов

Связанный репозиторий GitHub: https://github.com/ssghost/JITS_tests

компиляция точно в срок (JIT) - это способ выполнения компьютерного кода, который включает компиляцию во время выполнения программы (во время выполнения), а не перед выполнением. Поскольку индустрия разработки фреймворков глубокого обучения шагает вперед в так называемые "римские времена" (2019-2020), большинство основных фреймворков включили свои собственные компиляции JIT в свои бэкенды в своих последних установочных пакетах. Это явление можно было бы описать как интенсивную конкуренцию различных JIT-компиляторов в области вооружения.

Так почему бы нам не сделать тест для измерения фактической эффективности этих JIT-компиляторов, чтобы нарисовать краткую иллюстрацию общего механизма в этих различных реализациях JIT?

Идея проста, у нас есть две метрики для измерения: время выполнения и точность результата (в данном случае мы используем близость к значению числа пи). Первое короче, второе выше, компилятор лучше, и наоборот.

Участники также показаны на изображении обложки, а именно: Numba, JAX, Tensorflow, Triton.

Функция, которая будет одинаково выполняться этими 4 JIT-бэкендами на процессорах моего ноутбука, - это аппроксимация Pi методом Монте-Карло, которая может приблизительно приблизиться к фактическому значению pi.

Пишем декораторы

Прежде чем мы начнем, необходимо хорошо спроектировать функции для измерения этих показателей. Python имеет встроенный мощный инструментарий под названием «Decorator», который является элегантным способом завершения наших процедур измерения. Здесь мы написали две функции декоратора, они имеют ту же структуру, что и ниже:

def benchmark_metric(func: Callable[..., Any]) -> Callable[..., Any]:
    @functools.wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        value = func(*args, **kwargs)
        metric = f(value)
        logging.info(f"Function {func.__name__}'s metric value is {metric:.2f}")
        return value

    return wrapper

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

После этого все, что вам нужно сделать, это просто поставить оператор «@» перед функциями вызова JIT.

Пишем вызывающие функции

Как мы можем видеть в следующих псевдокодах, интересно, что все 4 типа бэкендов JIT-компилятора также являются декораторами. Следовательно, каждая вызывающая функция будет иметь по крайней мере 3 слоя декораторов, 2 для бенчмарков, 1 для вызова JIT-бэкендов.

Во-первых, мы могли бы рассмотреть ситуацию, когда мы не используем серверные части JIT, просто реализуем аппроксимацию Pi методом Монте-Карло в простой функции Python, и псевдокод для этого алгоритма будет таким, как показано ниже:

def monte_carlo_pi(ln: int):
    acc = 0
    for _ in range(ln):
        x = random.random()
        y = random.random()
        if (x**2 + y**2) < 1.0:
            acc += 1
    return 4.0 * acc / ln

Очевидно, что это подход, основанный на цикле, который означает, что мы должны повторить эту процедуру много-много раз, чтобы получить оптимизированный результат, чем больше раз мы повторим, тем большую точность мы получим. Время повторения будет "Large Number" (ln). Построенная диаграмма, изображающая эту процедуру, выглядит примерно так:

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

Для Numba существенное отличие заключается в том, что нам больше не нужно вызывать Numpy или некоторые подобные Numpy пакеты для предварительной обработки (или, более конкретно, "сжатия") ваших данных в структурированное векторное пространство. Удивительно, но мы можем просто скопировать и вставить наш псевдокод под @ statements, и он будет работать правильно, как мы и ожидали. Вот почему я поставил Numba первым участником нашего конкурса.

Но с остальными тремя JIT-бэкендами все будет немного сложнее. Мы должны регулярно выполнять предварительную обработку наших данных. К счастью, JAX и Tensorflow имеют собственный встроенный модуль Numpy, мы можем просто импортировать и вызывать их для сжатия наших данных в массивы Numpy и выполнять все вычисления с этими массивами Numpy (это означает, что все другие типы данных не модифицируются для подачи в массивы Numpy). их бэкенды JIT).

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

Результат

Что касается описания результатов нашего конкурса JIT, Matplotlib должен быть отличным выбором. Мы построим 8 графиков, по 2 показателя для каждого участника, с массивом входных переменных - время повторения. Давайте перейдем непосредственно к проверке этих графиков.

Что ж, при первом взгляде на эти 4 пары графиков я мог бы сказать, что их формы казались почти идентичными. Только Numba получил тривиальное улучшение в метрике perf_time. Чтобы сэкономить время на чтение, здесь я, наконец, решил опубликовать результаты двух метрик Numba в качестве типичного примера:

Выводы

• Современные платформы DL создали почти идентичные эффективные серверные части JIT-компиляции с собственной структурой данных.
• Ключевым моментом повышения производительности бэкендов JIT, работающих в среде, использующей только CPU, является оптимизация процедур сжатия данных.
  
• Может быть наилучшая точка компромисса между perf_time и value_acc, которая составляет от 10^6 до 10^7 в соответствии с нашими диаграммами результатов, что может дать тривиальный намек для дальнейших исследований.