В прошлой публикации История одного провала я рассказал про свои попытки автоматизировать упрощение символьных выражений. Но почти совсем не коснулся вопроса – зачем мне это потребовалось, так что пришлось много объясняться по этому поводу в комментариях. В этой статье я расскажу про почти успешную часть того проекта – программу, которая должна была писать другие программы. За 10 лет до этого вашего ChatGPT.

Итак, коротко напомню, в чем состояла исходная проблема. Я – астрофизик. Когда меня спрашивают, в чем отличие астронома и астрофизика, я всегда говорю, что астроном наблюдает разные явления в телескоп, а астрофизик говорит ему, что он там должен видеть. Доля шутки, тут, конечно, есть, но, в целом, дела примерно так и обстоят. Астрофизики исследуют физические процессы в звездах, планетах, галактиках и даже Вселенных (по крайней мере в одной), для чего строят физические модели. Модели же нужны потому, что астрономы, по сути, мало что могут увидеть непосредственно. Например, мои любимые взаимодействующие двойные звезды настолько малы, по сравнению с расстояниями до них, что получить хоть сколько-то детальную картинку из наблюдений невозможно. Из наблюдений получаются только интегральные характеристики – кривые изменения яркости да изменения спектров (цвета, по сути). То есть, как если бы мы взяли изображение звезды и ужали его до одного пикселя. Вот из этого пикселя нам надо понять, что же там происходит. Мы, астрофизики, это умеем.

Для этого и нужны модели. Как по одному пикселю восстановить изображение? Нужно заранее знать, что там нарисовано! Мы, из общих физических соображений можем нарисовать картину течения, какой она должна быть в данной ситуации, после чего "ужав ее до одного пикселя" мы можем сравнить "синтетические" наблюдательные данные с реальными. Если совпало – мы молодцы, определили правильно физику процессов. Если не совпало – это редкостная удача, известные модели не работают, а значит мы стоим на пороге какого-то открытия! Вообще, есть еще один забавный способ получить двумерную и даже трехмерную картинку по наблюдениям "одного пикселя" – доплеровская томография, но про нее в другой раз.

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

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

Вычисления довольно тяжелые. Сотни часов счета на нескольких сотнях ядер – обычный рутинный расчет, не более. Так что код должен быть очень, очень хорошо оптимизирован. Выигрыш в несколько процентов скорости это сэкономленные дни счета! Ну и сам код довольно сложен. Во-первых, он параллельный, причем работает на системах без общей памяти, где обмен данными между вычислительными узлами задается явным образом, при помощи библиотеки MPI. Во-вторых, он включает в себя довольно развесистые формулы, используемые для расчета потоков, отдельную логику для обработки границ, обработку случаев, когда численная схема не смогла вычислить значения правильно, сохранение и восстановление данных, чтобы счет можно было прерывать и возобновлять и многое, многое другое. Написан на Фортране. Хотя некоторые мои коллеги предпочитают C и даже C++, неиссякаемый источник для взаимных шуток и подкалываний. В общем, внесение каких-то изменений в код, например, включение учета какой-нибудь дополнительной физики – задача нетривиальная и очень трудоемкая. Мне хотелось эту задачу упростить.

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

Первым делом мне нужно было получить выражения для пересчета физических величин. Если вы не читали предыдущую статью, то самое время прочитать ее сейчас, там я подробно описываю, какой эпический провал ожидал меня на этом этапе. Однако, хотя мне и не удалось получить оптимальные уравнения, они были достаточно хороши, чтобы продолжать. И я продолжил! У меня, в самом простом случае были трехмерные массивы с физическими величинами rho(:,:,:) – плотность, p(:,:,:) – давление и v(3,:,:,:) – скорость. В процессе пересчета с учетом потоков я складывал значения для следующего временного шага в массивы, соответственно, rho1(:,:,:), p1(:,:,:) и v1(3,:,:,:), которые потом копировал в исходные rho, p и v. После чего повторял всё заново. Условно, программа выглядела примерно так:

t=0
do while t<tmax
  dt=<тут вычисляем шаг по времени>
  do <цикл по всем ячейкам, индексы i,j,k>
    rho1(i,j,k)=<тут гигантское выражение от rho, p и v, используюся только соседние с i,j,k ячейки>
    p1(i,j,k)=<аналогично rho2>
    v1(:,i,j,k)=<и тут то же самое>
  end do
  <тут устанавливаем граничные условия и обмениваемся данными с соседними вычислительными ядрами>
  rho=rho1
  p=p1
  v=v1
  t=t+dt
end do

Слова <гигантское выражение> следует понимать буквально, там действительно сотни и тысячи членов. Почему в выражениях используются только соседние с i,j,k ячейки, понятно – мы же вычисляем потоки через границы и нам для каждого потока нужно знать только значения в соседней ячейке. На самом деле нужно по две соседние ячейки, так как все уважающие себя схемы имеют повышенный порядок точности, но для простоты будем считать, что у нас схема первого порядка.

Я, из любопытства, попробовал сгенерированный таким способом код скомпилировать, но не смог дождаться окончания компиляции. Нужно оптимизировать, подумал я. И оптимизировал!

Оптимизируем!

Первым делом, я сделал так, чтобы повторяющиеся части в выражениях вычислялись один раз и хранились во временных переменных. Таких переменных набралось несколько десятков и программа даже скомпилировалась за более-менее разумное время, это позволило мне убедиться, что код вообще работает. Далее я заметил, что часть подвыражений вообще не зависит от i,j,k, то есть является константами! Добавил автоматический вынос их за пределы цикла. Еще выяснилось, что некоторые подвыражения, связанные с геометрией сетки (сетка у меня была криволинейная, для вычисления потоков я просто перехордил в систему координат, связанную с центром грани, через которую считал поток) зависят только от i,j,k, но не зависят от значений других переменных. Я сделал автоматический вынос их в "константные массивы", вычисляемые один раз перед началом счета, но потом выкинул эту функциональность – она замедляла! В принципе, разумно – доступ к памяти стоит много дороже, чем эти вычисления. А вот дальше был нетривиальный момент.

Я заметил, что некоторые выражения повторяются, но со смещением на 1 по одному из индексов (индексу самого внутреннего цикла). В принципе, это ожидаемо – мы вычисляем поток через правую грань ячейки, потом сдвигаемся на 1 вправо и снова вычисляем то же значение для уже левой грани текущей ячейки. Почему бы не вычислить его один раз и не сохранить для следующей итерации цикла? Код, который умел находить и выносить такие выражения получился крайне замороченным, но, в итоге, мне удалось одной этой оптимизацией уменьшить количество вычислений более чем в два раза! Тут мне в голову пришла одна замечательная мысль.

Разделяй и властвуй!

У меня трехмерная сетка, разбитая на отдельные кусочки, по числу вычислительных ядер. То есть каждому ядру достается отдельная как бы сетка, меньшего размера, которую оно обрабатывает независимо, обмениваясь значениями граничных ячеек с соседями, чтобы решение сшивалось бесшовно. При этом вычислительные ядра могут физически находиться на разных компьютерах кластера, соединенных сетью (вычислительных узлах). Почему бы, подумал я, не использовать ядра в пределах одного узла более рационально. Для этого пилим сетку на более крупные куски, по числу вычислительных узлов, а не ядер. Внутри каждого узла делим трехмерную сетку на одномерные столбцы, каждое ядро узла берет себе пучок из этих столбцов и обрабатывает их, без всяких обменов, а потом узел уже обменивается готовыми данными с соседними узлами. Сказано – сделано! Внутри каждого узла я использовал OpenMP, чтобы загрузить ядра параллельной работой. Всё генерировалось автоматически и выглядело очень красиво в теории. Тестовый запуск же показал, что скорость работы уменьшилась и причем значительно. Вот это поворот...

На самом деле ларчик открывался просто. Рандомный доступ к памяти – это дорого, особенно, когда к одному куску памяти имеют доступ несколько ядер одновременно. В случае, когда все ядра были независимы и обменивались данными через MPI, конкурентного доступа к памяти почти не было. Каждое ядро эксклюзивно работало со сравнительно небольшим куском памяти, который хорошо помещался в кэш. Когда же я натравил все ядра на целый большой массив, то получил прелести конкуренции, которые съели весь профит. Решение проблемы, на самом деле, простое – сначала сложить все данные для каждого ядра в отдельный массив и пусть оно пишет тоже в свой личный массив, а когда все закончат – какое-нибудь одно ядро сложит все результаты в общий массив в памяти. Несмотря на то, что тут данные лишний раз копируются, отсутствие конкуренции и кэш позволяют получить хороший прирост производительности. В лохматые времена, когда ядро в процессоре было одно, этот метод тоже давал преимущество – при обходе трехмерного массива мы клали данные из каждой его строки в отдельный массив и потом копировали результаты обратно. За счет более эффективного использования кэша такой подход давал кратный (!) прирост производительности.

Поскольку я генерировал код автоматически, я мог точно определить, какие данные и в какой последовательности мне требовались при вычислениях. Так что я не делал отдельные кэш-массивы для всех величин, а делал один большой одномерный массив, куда клал данные из rho, p и v в нужном порядке. Данные в выходной массив тоже складывались в том порядке, как они вычислялись, так что и чтение и запись из/в память происходили строго последовательно. Сказать, что это дало прирост, значит ничего не сказать. Скорость увеличилась в разы!

И тут я вспомнил, что есть еще и GPU...

Вообще, это считается плохой практикой, смешивать при расчетах CPU и GPU, обычно используют что-то одно. Но я хотел выдоить все флопсы, какие были в системе. В принципе, идея была простая – при распределении работы часть столбцов выделить для GPU, пусть он тоже будет загружен обработкой. Код я генерирую, так что ничто не мешало мне написать генератор кода для OpenCL (на машине, где я всё это отлаживал, была видеокарта от AMD) и связать этот кусок с общей программой.

Тут мне сильно пригодилось то, что я складываю предварительно подготовленные данные в отдельный массив. Гораздо проще копировать в GPU один массив и потом обращаться к нему в коде ядра, чем передавать сразу несколько. Кроме того, тут открывается еще один путь для оптимизации, специфической для GPU. Дело в том, что типичный GPU состоит из сравнительно небольшого числа реальных процессорных ядер, но ядра эти суперскалярные, то есть команда сложения, например, может складывать не два числа одновременно, а, скажем, по 16 пар чисел за раз. Если вычислительные ядра "шагают в ногу", то есть выполняют строго одинаковый код, различающийся только входными данными, то производительность большинства математических операций сразу же возрастает в 16 раз! Число 16 не магическая константа, оно зависит от модели GPU, но обычно это 16. Я немного переписал код для создания входного массива так, чтобы он клал нужные данные по 16 штук за раз и загрузил GPU по-полной! К слову, CPU тоже может в векторные операции, но в те времена это было не сильно развито и я не стал такое реализовывать. Кстати, на моем рабочем iMac 2011 года GPU и CPU давали примерно одинаковую производительность, то есть, задействовав оба, я получал ускорение примерно вдвое.

Иногда они шагают не в ногу...

Численные схемы, которые я использую, содержат в себе условные операторы. То есть, в зависимости от некоторых соотношений, может подставляться одно или другое выражение при вычислениях. Понятно, что это напрочь ломает векторизацию на GPU, для CPU тоже ничего хорошего, так как при условном ветвлении, которое процессор не смог предсказать, он сбрасывает всякие заранее вычисленные данные и начинает вычислять их заново. Столкнувшись с этим, я изобрел (ну, может не я первый, не знаю) способ обойтись без ветвлений там, где без них никак не обойтись. Способ довольно наглый, но при этом красивый.

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

x = <большое выражение> + (<условие>?<вариант a>:<вариант b>) + <остальное выражение>

заменять на:

x = <большое выражение> + (C * <вариант a> + (1 - C) * <вариант b>) + <остальное выражение>

где C – некое число, принимающее значение 1.0 если <условие> было истинным и 0.0, если оно было ложным. Осталось придумать способ превращать логические выражения в арифметические, дающие 0.0 или 1.0 в качестве результата. Например:

a == b : C = 1-abs(sign(a-b))
a > b  : C = max(sign(a-b),0)
C1||C2 : max(C1+C2,1)
C1&&C2 : C1*C2
!C     : 1-C

...и так далее. Таким нехитрым способом я избавился от сбоев конвейера на CPU и от рассинхронизации ядер на GPU.

Но GPU слишком медленный для real-8!

Это правда. GPU сконструирован для работы с не очень точными данными, real-4 для него максимальная точность, на которой он может показать нормальную производительность. Но для моих моделей нужна была real-8! Нужно обезразмеривать, подумал я.

Этот трюк обычно используется, если у вас при вычислениях используются очень большие или очень малые числа, на которых стандартная математическая библиотека дает слишком большую погрешность, при использовании типов менее точных, чем real-8. Нужно просто переписать все уравнения так, чтобы используемые числа были порядка единицы. Это несложно, просто все расстояния (которые у меня порядка нескольких радиусов Солнца) делим на радиус Солнца, а время делим на величину орбитального периода двойной звезды (у меня – порядка часов), аккуратно подставляем это в исходные уравнения и, после простых преобразований, у нас получаются точно такие же уравнения, но с другими константами. Особый шик – выбрать обезразмеривание так, чтобы и константы тоже стали равными 1. Поскольку у меня был, пусть и скверный, но рабочий блок символьной математики, я смог автоматизировать обезразмеривание. Пришлось прописать физическую размерность для всех констант и переменных, после чего я мог просто написать, что я хочу, чтобы орбитальный период был равен 1 и суммарная масса звезд тоже была равна 1, блок символьной математики решал простую систему уравнений и выводил новые константы, которые зашивались сразу в генерируемый код. Это позволило сильно повысить точность и без проблем обойтись real-4 при вычислениях на GPU.

Балансировка нагрузки

А вот тут меня снова ждал провал. Даже два провала. Как происходит обработка вычислительной сетки:

1. Ядро CPU делит трехмерную сетку на пучки столбцов, по числу решающих ядер (т.е. общему числу ядер всех CPU на узле + количеству GPU на нем же)

2. Для каждого пучка столбцов подготавливает "входной массив", в который нужные для расчетов данные записаны в нужном порядке. Причем для GPU они еще сгруппированы по размеру варпа

3. Отправляет данные на GPU и запускает их на счет

4. Все CPU-ядра берут свои массивы и обрабатывают их

5. Ядро CPU сгружает готовые данные с GPU

6. Ядро CPU переносит данные из результирующих массивов в исходные

Чтобы максимально избежать простоев железа необходимо запускать счет на каждом ядре и GPU сразу же, как только для него готовы входные данные. К моменту, когда ядро или GPU завершило обработку, для него должен уже быть готов новый входной массив, чтобы следующий цикл счета запускался без задержек. Ядро, которое занимается подготовкой и копированием данных тоже не должно простаивать, пока все заняты счетом, оно должно выделить для себя небольшой пучок столбцов и обработать их, но так, чтобы к моменту, когда кто-то другой закончит, оно уже было готово принять результаты. Как я ни бился, мне не удалось загрузить узел с 16 ядрами (2 процессора по 8 ядер) и двумя GPU на 100%. Всегда что-нибудь простаивало. Угадать, когда GPU закончит оказалось невозможно. Копирование данных может занимать сколько угодно времени и это непредсказуемо. Ну, а в редкие моменты, когда мне удавалось таки загрузить работой всех, узел банально зависал, подозреваю, от перегрева :)

Выводы из всей этой истории

Первое – генерация кода вполне рабочая практика. Фактически, я написал компилятор с некоторого DSL в код на Фротране и OpenCL. Сам "компилятор" был написан на Лисп-языке Scheme, а для Лиспов это типично – набором макросов превратить язык в черт знает что в некий псевдо-язык, приспособленный для решения некоей проблемы, что и было мною реализовано. Код получился довольно большим по моим меркам – порядка 20 тыс. строк (и генерировал код примерно на 15 тыс. строк, хех), при этом отдельные части оптимизатора были настолько сложными, что алгоритм его работы просто не помещался у меня в мозгу целиком. Я впервые столкнулся с тем, что программа была сложнее программиста, очень странное ощущение. Тем не менее, мне удалось ее написать и отладить, в чем большая заслуга функционального стиля программирования в целом и Лиспа, как языка, в частности. Когда мне нужно было исправить оптимизатор, я каждый раз заново изучал отдельный кусок кода, разбирался, как он работает, вносил нужные изменения и потом сразу же забывал его напрочь. В большинстве других языков это бы никогда не привело к работающему результату, но Лисп очень хорошо приспособлен для подобных экзерцисов, за что я его и полюбил в итоге.

Что мне не удалось. Во-первых, мой подход к символьной математике (см. предыдущую статью) оказался неверным и у меня просто не было тогда сил и времени для его переделки. Так что, несмотря на всю оптимизацию, формулы, код которых оптимизировался были изначально слишком раздутыми, что съедало весь профит. Во-вторых, подойдя к пределу производительности, я столкнулся с необходимостью четкой синхронизации работы всех вычислительных ядер и GPU, обеспечить которую я не смог. Как правило, у меня как минимум один GPU из двух, в среднем, простаивал. Побороть это, наверное, можно, но в тот момент у меня просто кончились силы и я забросил проект на годы. Может быть, я вернусь к нему через некоторое время. Как минимум, символьная математика у меня уже работает.