Subgroup против Warp - или как я провел вечер воскресенья

Всем привет! Хочу поделиться историей одного бага, который подарил мне незабываемый вечер. Как разработчик библиотеки машинного обучения, я иногда занимаюсь написанием и оптимизацией вычислительных ядер для GPU, и этот случай заставил меня по-новому взглянуть на subgroups в Vulkan. Это был не просто баг, а интереснейшая головоломка, разгадка которой расширила моё понимание низкоуровневой работы GPU.

Мой надёжный друг CUDA Warp
В мире CUDA есть понятие warp — группа из 32 потоков, которые выполняют инструкции в режиме lock-step. Это предсказуемо, как швейцарские часы. Большинство алгоритмов reduce и scan заточены под эту константу. Я знал это как свои пять пальцев и строил на этом фундаменте свои реализации таких алгоритмов для Vulkan.
С Vulkan всё казалось похожим. Subgroup — та же идея: lock-step потоки, обмен данными, барьеры. «Отлично! — подумал я. — Просто копируем логику». И это была моя роковая ошибка.

Призрак в системе: некорректные результаты и Device Lost
Запуски на интегрированной графике Intel преподнёсли сюрпризы. Мои тесты, отлично работавшие на NVIDIA, начали выдавать заметные расхождения в результатах. А иногда приложение просто падало с загадочным VK_ERROR_DEVICE_LOST, не оставляя никакой полезной отладочной информации.

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

Subgroup — это не константа, а переменная!
Прозрение наступило, когда я наткнулся на похожий баг в другом проекте. А потом в очередной раз перечитал спецификацию. Оказалось, размер subgroup в Vulkan не фиксирован. Более того, он может динамически меняться не только от GPU к GPU, но и в рамках одного конвейера!
Вот что я упустил из виду:
- Драйвер может запускать шейдер, допуская разный размер subgroup для одного и того же вычислительного ядра в зависимости от загрузки GPU и эвристик планировщика. Сейчас он выполняет его с subgroup в 32 потока, а в другой раз - в 16, чтобы лучше упаковать рабочие группы.
- Размер subgroup можно контролировать используя расширение VK_EXT_SUBGROUP_SIZE_CONTROL, которое вводит целый механизм для управления этим поведением. Структура vk::PhysicalDeviceSubgroupSizeControlProperties показывает поддерживаемый диапазон minSubgroupSize, maxSubgroupSize значений, а vk::PhysicalDeviceSubgroupSizeControlFeatures управляет в целом функциональностью в целом.
- Чтобы получить фиксированный размер, нужно явно задать requiredSubgroupSize при создании pipeline для шейдера и включить соответствующее расширение. Без этого вы играете в русскую рулетку с аппаратурой.

Почему всё сломалось?
Мои шейдеры, заточенные под 32 потока, на Intel попадали в подгруппы меньшего размера.
- Reduce ломался: алгоритмы, использующие жёстко закодированные смещения, начинали обращаться к несуществующим потокам, получая мусор и портя результаты вычислений.
- Синхронизация с использованием barrier() стала опасной: я использовал барьеры для синхронизации внутри subgroup, будучи уверены, что все 32 потока встретятся на этом барьере. Но в мире, где потоков всего 16, часть наших потоков (в логике, рассчитанной на 32) уже завершала работу и не доходила до барьера. Попытка синхронизироваться в таких условиях — это как undefined behavior. Драйвер Intel в таких случаях часто просто ронял GPU, вызывая ошибку VK_ERROR_DEVICE_LOST — которую непросто отлаживать, потому что она не оставляет следов, лишь ощущение полной потери контроля.