OpenBMB открыла ForgeStencil — агентную систему для автоматической оптимизации HPC-приложений

ForgeStencil нацелен на stencil-вычисления — повторяющиеся операции над точками регулярной сетки и их ближайшими соседями. Такие шаблоны лежат в основе моделирования жидкостей, электромагнитных полей, сейсмической разведки, медицинской реконструкции и многих CAE-задач. Они обычно ограничены не пиковой производительностью процессора, а пропускной способностью памяти и могут занимать основную часть времени выполнения приложения. Ручная настройка требует знания численных методов, архитектуры GPU и устройства конкретной кодовой базы, поэтому один HPC-инженер, как правило, глубоко оптимизирует не более 20 приложений в год.

Система разделяет работу между двумя агентами. Kernel Agent исследует пространство реализаций и синтезирует специализированные ядра под тип stencil-операции, форму данных и требуемую точность. App Agent занимается интеграцией: анализирует исходный код, находит горячие функции и возможности слияния операций, строит исходный GPU-бейзлайн, подставляет кандидатов и запускает штатные тесты приложения. После проверки корректности и замера производительности выбранная реализация встраивается обратно в проект — разработчики утверждают, что решения по оптимизации на всех этих этапах не требуют участия HPC-эксперта.

В тестах отдельных операторов ForgeStencil при вычислениях в FP32 показала геометрическое среднее ускорение в 2,35 раза относительно набора открытых решений, среди которых Halide, Devito, EBISU, DRStencil и FlashFFTStencil. Переход к чтению и записи в FP16 дал дополнительное ускорение в 1,95 раза, а на более сложных stencil-задачах с переменными коэффициентами преимущество над лучшей базовой реализацией составило в среднем 1,34 раза. Это результаты разработчиков на одинаковом оборудовании, поэтому их следует отличать от независимого бенчмарка и от ускорения приложения целиком.

Для прикладной проверки команда использовала собственные тесты и средства контроля реальных программ. Заявленное сквозное ускорение составило 3,86 раза для библиотеки многосеточных решателей Hypre, 5,78 раза для нейтронного mini-app Minisweep и 2,47 раза для электромагнитного симулятора gprMax/FDTD. Для обратной временной миграции результат достиг 1,81 раза, для расчета облигаций в QuantLib — 1,82 раза, для не cartesианской реконструкции MRI — 2,45 раза, а для цифрового томосинтеза молочной железы — 1,63 раза. Разброс показывает, что эффект зависит от структуры конкретного приложения и доли времени, приходящейся на оптимизируемое ядро.

По внутренней статистике ModelBest, более 100 приложений были оптимизированы за одну неделю, причем 42% набора относились непосредственно к программам промышленных производителей. Компания оценивает рост пропускной способности разработки в один-два порядка, а эффективности — примерно в 100 раз по сравнению с экспертным процессом, который занимает от нескольких дней до нескольких недель на проект. Эти оценки эквивалентны работе четырех-восьми инженеров в течение года и затратам в 3–10 млн юаней, однако пока они раскрыты самой командой ForgeStencil и нуждаются в воспроизведении на стороннем оборудовании и кодовых базах. Исходный код опубликован в репозитории OpenBMB на GitHub.