Rust SIMD на GPU: core::simd теперь работает на видеокартах

Компания VectorWare, создающая первую GPU-native программную платформу, сообщила об успешном запуске портируемого SIMD-модуля Rust (core::simd) на GPU. Это шаг к тому, чтобы разработчики могли писать сложные высокопроизводительные приложения, использующие всю мощь GPU-железа, привычными средствами Rust.

Параллелизм ниже уровня потока

Ранее, когда в VectorWare перенесли потоки Rust на GPU, каждый std::thread отображался на GPU-warp. Это позволило запускать множество конкурентных потоков на GPU, но не задействовало параллельные линии выполнения внутри каждого потока/warp.

На CPU абстракцией параллелизма внутри потока служит SIMD. Одна инструкция оперирует сразу несколькими элементами данных, упакованными в векторный регистр: там, где скалярный код складывает два числа, SIMD-сложение берёт два вектора, скажем, из восьми значений f32, и разом получает восемь сумм. Этот параллелизм данных существует внутри одного потока, ниже уровня, на котором операционная система вообще что-либо планирует.

Портируемый SIMD в Rust

Исторически написание SIMD-кода на Rust означало обращение к архитектурно-специфичным векторным интринсикам из core::arch, таким как _mm256_add_ps на x86-64 или vaddq_f32 на Arm. Эти интринсики привязаны к конкретному набору инструкций, поэтому программа, работающая на нескольких архитектурах, требует отдельной реализации под каждую из них.

Портируемый SIMD в Rust добавляет уровень абстракции поверх этих интринсиков. Он предоставляет единый обобщённый тип Simd<T, N>, представляющий вектор из N элементов типа T. Арифметика, сравнения, редукции и перестановки элементов пишутся один раз для Simd, а компилятор сам преобразует их в те векторные инструкции, которые поддерживает целевой процессор.

В VectorWare пришли к выводу, что GPU — это просто ещё один вид векторного железа, на который может нацеливаться портируемый SIMD. Приятный бонус: портируемый SIMD находится в core, а не в std, и ему даже не требуется та std-поддержка, которую ранее перенесли на GPU.

SIMT — это и есть SIMD

GPU выполняют код в модели, которую NVIDIA называет SIMT — Single Instruction, Multiple Thread (одна инструкция, много потоков). Warp выдаёт одну инструкцию, и каждая из его 32 линий выполняет эту инструкцию над своими данными. Одна инструкция, работающая сразу с множеством элементов данных, — это буквально то, что означает SIMD, и адресация по отдельным линиям, которую добавляет SIMT, этого не меняет. Warp — это широкий векторный юнит, и портируемый SIMD-вектор отображается на него напрямую.

Например, Simd<i16, 32> отдаёт по одному элементу i16 каждой из 32 линий warp, а сложение двух таких векторов компилируется в единую warp-инструкцию, в которой каждая линия одновременно складывает свой элемент.

Такое отображение завершает иерархию параллелизма, выстроенную в предыдущих работах команды. На CPU поток содержит SIMD-линии, а на GPU std::thread соответствует warp'у, чьи аппаратные линии играют ту же роль. В обоих случаях этими линиями управляет core::simd.

Мировая премьера: `core::simd` на GPU

Как и в предыдущих публикациях, продемонстрировать это визуально сложно — код выглядит как обычный Rust. Те же типы core::simd, которые на ноутбуке компилируются в SIMD-инструкции x86-64, на GPU превращаются в warp-операции — без единого изменения в исходном коде.

Ниже определена небольшая портируемая SIMD-подпрограмма, вызываемая из main. Она задействует основные элементы модели: поэлементную арифметику, сравнение, дающее маску по линиям, select, управляемый этой маской, и горизонтальную редукцию по линиям.

#![feature(portable_simd)]
 
use core::simd::cmp::SimdPartialOrd;
use core::simd::num::SimdFloat;
use core::simd::{Select, Simd};
 
// Portable SIMD. This exact function also compiles and runs on the CPU,
// where it lowers to x86-64, Arm, or scalar code depending on the target.
fn relu_dot(a: Simd<f32, 32>, b: Simd<f32, 32>) -> f32 {
    // Elementwise multiply: 32 products computed at once.
    let products = a * b;
 
    // Per-lane comparison produces a mask, one boolean per lane.
    let positive = products.simd_gt(Simd::splat(0.0));
 
    // Keep the positive products, replace the rest with zero.
    let clamped = positive.select(products, Simd::splat(0.0));
 
    // Horizontal add across all lanes down to a single scalar.
    clamped.reduce_sum()
}
 
fn main() {
    // Two 32-wide vectors, built with ordinary Rust.
    let a = Simd::<f32, 32>::splat(2.0);
    let b = Simd::<f32, 32>::from_array(std::array::from_fn(|i| i as f32 - 16.0));
 
    // Elementwise ops, a comparison mask, a select, and a reduction:
    // all ordinary portable SIMD, all running on the GPU.
    let result = relu_dot(a, b);
 
    // Printed from the GPU using our std support.
    println!("relu_dot = {result}");
}

Точка входа — обычная fn main без каких-либо GPU-специфичных аннотаций. Тулчейн компилирует её в GPU-ядро, а результат выводится с устройства с помощью std-поддержки.

Программа, запущенная на GPU, выдаёт точно такой же результат, как и при запуске на CPU.

Реализация

Отображение опирается на одно наблюдение: warp — это векторный юнит, линии которого адресуются по отдельности. Как только Simd<T, N> раскладывается по линиям, у каждого семейства операций находится прямой аналог на уровне warp.

Поэлементные SIMD-операции — самый простой случай. Сложение, умножение, сравнение и другие поэлементные операторы приходят из обычных реализаций трейтов Rust для Simd, таких как Add. GPU выполняет их нативно.

SIMD-редукции, такие как reduce_sum и reduce_max, объединяют значения всех линий в один скаляр. Для этого используются warp shuffle инструкции GPU, обменивающие и объединяющие значения между линиями, так что итоговый скаляр получается одинаковым во всех линиях.

Кросс-линейные шаффлы SIMD, такие как simd_swizzle! и повороты, перемещают элементы между линиями. Поскольку SIMD-линия — это линия GPU-warp, эти операции отображаются на те же примитивы warp shuffle, благодаря которым GPU-линии так хорошо обмениваются данными.

Маски SIMD отображаются не менее прямолинейно. Mask<T, N> задаёт по одному предикату на каждую SIMD-линию. Mask::select выполняет выбор в каждой линии warp. Горизонтальные запросы к маске, такие как any и all, используют GPU-инструкции vote и ballot.

Скалярные значения в окружающем коде — например, счётчик цикла или константа — вычисляются одинаково каждой линией и потому просто реплицируются по всему warp, как в обычном CUDA. Это то же различие между «единообразным» и «варьирующимся», которое языки с явным параллелизмом данных вроде ISPC делают явным, только здесь оно вытекает прямо из типов самого Rust: обычный f32 — единообразен, Simd<f32, 32> — варьируется.

Работа с линиями

Единственное место, где абстракция и железо расходятся, — это количество линий. На CPU Simd<T, N> допускает любое N от 1 до 64, но у GPU фиксированная ширина: 32 линии у NVIDIA и 32 или 64 у AMD. Отображение один к одному работает только когда N совпадает с этой шириной. Меньшее N оставляет часть линий простаивать, а большее — заставляет некоторые или все линии обрабатывать более одного элемента.

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

В VectorWare для этой «машины» построили собственное промежуточное представление (IR). Вместо отдельной структуры данных оно закодировано прямо в системе типов Rust — через типы, дженерики, константные дженерики и ограничения трейтов. Программа состоит из типизированных операций: ballot, шаффлов, редукций, сканов, gather, scatter, атомарных операций и strip mining для векторов шире, чем warp. Операнды, форма выполнения и ёмкость также типизированы. Поскольку операции несут свою форму прямо в типах, многие некорректные программы попросту невозможно составить.

IR не требует интерпретатора на GPU: каждая операция преобразуется напрямую в соответствующие инструкции без каких-либо накладных расходов по сравнению с вручную написанным PTX. Те же самые типы позволяют запускать это представление и на CPU. Был построен эталонный интерпретатор, детерминированно выполняющий IR, — своего рода Miri для программирования на уровне warp-линий. Он используется для симуляции GPU-кода и для дифференциального тестирования.

Сейчас работа нацелена на NVIDIA, но ничего в подходе не привязано к CUDA специфически. AMD wavefront'ы и subgroups в Vulkan предоставляют похожие примитивы и семантику. Само IR — архитектурно-независимый Rust-код.

Преимущества

Один и тот же исходный код работает и на CPU, и на GPU. Код и библиотеки, уже использующие портируемый SIMD, становятся кандидатами на выполнение на GPU без переписывания.

Немодифицированный CPU-код может использовать параллелизм на уровне линий GPU. Код, написанный с учётом GPU, может пойти дальше и использовать интринсики core::arch, напрямую отображающиеся на PTX.

Simd<T, N> — это обычное владеющее значение. Заимствования, время жизни и проверка типов применяются к нему точно так же, как на CPU. В VectorWare не добавляли GPU-специфичный векторный тип или новый набор аннотаций — существующий портируемый SIMD Rust отображается на нативную модель исполнения GPU. Так GPU превращаются в обычную платформу для Rust.

Недостатки

Портируемый SIMD в Rust всё ещё нестабилен. Он требует nightly-канала и флага #![feature(portable_simd)], и его интерфейс может измениться до стабилизации.

Векторы уже, чем warp, оставляют часть линий простаивать, а векторы шире warp превращают каждую операцию в несколько инструкций. Абстракция бесплатна только тогда, когда ширина вектора совпадает с числом линий warp.

Не каждая кросс-линейная операция отображается на эффективную warp-инструкцию. Шаффлы, соответствующие поддерживаемым железом паттернам, дёшевы, а произвольные перестановки могут потребовать нескольких инструкций или обращения к разделяемой памяти. Горизонтальные операции вроде редукций и all/any также выступают точками синхронизации внутри warp, что ограничивает свободу планировщика в перекрытии работы.

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

Дальнейшие планы

После того как SIMD, потоки и async отображены на GPU, естественный следующий шаг — их совместное использование: потоки распределяют работу между warp'ами, core::simd распределяет данные по линиям внутри каждого warp, а async структурирует конкурентность между ними.

Есть интерес и к отображению матричного SIMD на тензорные ядра GPU, а также к автовекторизации обычных скалярных циклов Rust в операции Simd, чтобы код получал параллелизм на уровне warp, даже не будучи написанным напрямую под core::simd. Как участники команды разработки компилятора Rust, в VectorWare заинтересованы в том, чтобы выяснить, насколько многое из этого может делать сам компилятор.

Единое векторное представление для CPU и GPU было бы ценным, хотя не факт, что сегодняшние типы портируемого SIMD — подходящая основа для этого. В частности, они во многом существуют особняком в API core и std. Требуется дальнейшее исследование.

Фокусируется ли VectorWare только на Rust?

Скорость, с которой удаётся продвигаться в работе с GPU, — во многом заслуга мощи абстракций и экосистемы Rust.

Как компания, в VectorWare понимают, что не все используют Rust. Будущие продукты будут поддерживать несколько языков программирования и рантаймов. Тем не менее там считают, что именно Rust уникально хорошо подходит для создания высокопроизводительных, надёжных GPU-native приложений — и это то, что вызывает наибольший интерес у команды.