Автор: Цянь Бао и Хайдун Лан

В предыдущем блоге мы узнали, что Taichi, язык высокопроизводительных вычислений, встроенный в Python, выходит за рамки инструмента разработки компьютерной графики и визуализаторов, но также удобен для числовых вычислений, включающих массовые операции с 2D- и 3D-массивами. Вычислительная гидродинамика (CFD) — типичный сценарий, в котором Тайчи может сыграть свою роль.

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

Случай начального уровня: сравнение производительности в операциях суммирования

Операция суммирования вычисляет сумму всех элементов в заданном массиве. Это обычная арифметическая операция в числовых вычислениях.

В этом разделе в качестве примера используется операция суммирования для сравнения вычислительной производительности Taichi, CUB, Thrust, CuPy и Numba.

Начнем с небольшой справочной информации. И CUB, и Thrust, разработанные Nvidia, предоставляют общие утилиты для параллельных вычислений. Несмотря на сходство в использовании, CUB был разработан специально для графических процессоров Nvidia и напоминает базовые инструкции CUDA, в то время как Thrust отличается более высоким уровнем абстракции для совместимости с несколькими платформами, поэтому в большинстве случаев производительность ниже, чем у CUB.

Общие решения для ускорения графического процессора, доступные пользователям Python, включают CuPy и Numba. Первый предоставляет интерфейс, аналогичный NumPy, позволяя пользователям вызывать CuPy так же, как они вызывают NumPy. Здесь мы выбираем наиболее производительный CUB в качестве бэкенда ускорения для реализации CuPy. Реализация Numba выполняет операцию сокращения в соответствии с инструкциями, предоставленными официальной документацией.

Мы также реализуем операцию суммирования в Taichi, и это довольно просто:

@ti.kernel
def reduce_sum_kernel():
    sum = 0.0
    for i in f:
        sum += f[i]

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

График показывает, что вычислительная производительность всех реализаций менее чем удовлетворительна при минимальном размере данных. Причин низкой производительности две: 1) маленькие массивы не могут полностью использовать возможности GPU; 2) вызов библиотек неизбежно влечет за собой накладные расходы в разной степени. По мере увеличения размера данных доля фиксированных накладных расходов снижается, а производительность значительно повышается. Примечательно, что реализации CUDA/CUB и CuPy обеспечивают впечатляющую производительность, превышая 90% пиковой пропускной способности оборудования. Производительность Taichi сравнима с высокооптимизированными версиями CUB и CuPy и значительно превосходит Thrust при всех объемах данных.

На самом деле компилятор Taichi опирается на сложную базовую инженерию для обеспечения высокой производительности. В вышеупомянутом фрагменте кода суммирование на самом деле является атомарной операцией, которую нельзя распараллелить, и поэтому ее эффективность ограничена. Наиболее часто используемый метод оптимизации параллельных вычислений при суммировании векторов — это редукция, которая является одним из обязательных навыков для изучения параллельных вычислений. Благодаря автоматической оптимизации сокращения, реализованной его компилятором, Taichi достигает производительности, сравнимой с CUB, реализованной вручную, и намного лучше, чем Numba. См. сравнительный отчет.

В случае простой арифметической операции Taichi может конкурировать с высокооптимизированными реализациями CUDA и CUB и легко превзойти Thrust и Numba благодаря интуитивно понятному коду. Но можем ли мы доверять Тайчи в реальных задачах численных вычислений, включающих более сложные формулы? Теперь давайте перейдем к более сложному случаю в CFD.

Сложный случай: сравнение производительности при расчете поля скорости

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

Написание этой формулы с помощью CUB неэффективно из-за отсутствия API для ее поддержки. Его можно написать в Thrust, который является библиотекой шаблонов, но программирование шаблонов превращает отладку в кошмар. Что касается CuPy, то он может реализовывать сложные операции только через код CUDA. Когда нет доступной библиотеки, нам нужен высокопроизводительный язык программирования для реализации этой формулы. Поэтому мы ограничиваем область сравнения производительности набором Numba, CUDA и Taichi.

Начнем с того, что и Taichi, и Numba — это языки программирования, встроенные в Python, что позволяет пользователям создавать алгоритмы, просто следуя синтаксису Python. Однако Numba может использовать только вычислительную мощность ЦП с кодом Python; чтобы получить доступ к графическому процессору, ожидается, что вы понимаете модель программирования CUDA и самостоятельно управляете планированием потоков. Напротив, Taichi не требует никакого предварительного опыта работы с CUDA. Пользователям нужно только указать серверную часть CUDA в вызове ti.init(), чтобы выгрузить код на GPU.

Что касается нативной реализации CUDA, то мы закончили писать ядра за полчаса, но потратили почти два часа на выравнивание значений. Но это не конец истории. Вызов кода CUDA из Python также влечет за собой дополнительный код инструментов для компиляции, инкапсуляции интерфейса, выделения и освобождения памяти, а также передачи данных CPU-GPU. Весь процесс по-прежнему накладывает значительную нагрузку на эффективность программирования, хотя для упрощения доступны различные библиотеки.

При всем удобстве программирования, которое предлагает Taichi, как он оценивается с точки зрения производительности? Мы запускаем каждую реализацию 1000 раз подряд и берем среднее время, затрачиваемое на вызов, в качестве показателя производительности. Результаты представлены ниже. Чем ниже столбец, тем выше производительность.

Судя по всему, три реализации с GPU-ускорением значительно превосходят версию Numba на CPU. Когда длина края полей скоростей достигает 2048 или больше, Taichi тратит только 1/3 времени Numba (версия CUDA) на завершение вычислений и работает даже быстрее, чем нативная реализация CUDA. Высокая производительность может быть связана с более быстрым доступом к памяти, обеспечиваемым компилятором Taichi, который автоматически использует аппаратные возможности. CUDA может достичь сравнимой производительности после оптимизации при условии, что пользователи хорошо разбираются в базовом оборудовании.

Исходный код различных реализаций доступен в этом репозитории.

Краткое содержание

В этом блоге мы исследуем производительность Taichi при суммировании векторов и операциях с жидкостным полем и сравниваем Taichi с другими распространенными решениями для ускорения. Наши выводы таковы:

  • В случае простых операций суммирования производительность Taichi находится на уровне основных высокооптимизированных библиотек с ускорением на GPU.
  • В случае сложных полевых операций Taichi превосходит Numba (версия CUDA) в 3-4 раза и немного превосходит реализованную вручную и семантически эквивалентную версию CUDA.
  • С точки зрения эффективности кодирования, Taichi обеспечивает гораздо более простой и простой способ кодирования, чем Numba (версия CUDA) и CUDA; он также позволяет плавно переключаться между бэкендами, избавляя от необходимости переписывать код.

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

Наш следующий блог продемонстрирует, как написать программу CFD, используя 99 строк кода Taichi. Оставайтесь в курсе!