В Раст экосистеме уже есть несколько неплохих реализаций аллокаторов GPU памяти: gpu-allocator и Vulkano1 основанные на свободных списках, а так же gpu-alloc, основанный (как я понял), на гибриде свободных списков и buddy.
Перечисленные решения являются аллокаторами общего назначения, и имеют свои сильные и слабые стороны. В своем решении я решил пойти по пути избыточной внутренней фрагментации в угоду скорости выделения памяти. Кроме того, моё решение немного более дружелюбно в случае многопоточного доступа. У моего аллокатора память разного типа может выделяться параллельно в разных потоках.
Данное решение рассчитано на вариант использования, при котором пользователь часто выделяет и освобождает не слишком большое количество не слишком большой памяти переменного размера, а так же, возможно, выделяет куски большого, но константого размера, и использует их большую часть времени жизни программы.
Содержание
Память разделяется по типам, которые включают в себя собственно GPU индекс типа, а так же линейность, и доступ со стороны CPU. Если запрашивается host-visible память, она автоматически мапится, и остается замапленной все время жизни соответствтующего device memory объекта.
struct AllocInfo {
memory_type_index: u32, // GPU индекс типа памяти
linear: bool, // Флаг линейности
map: bool, // Флаг маппинга
}Поскольку память жестко разделяется на линейную и нелинейную, проблема image-buffer гранулярности отсутствует по построению.
Выделение осуществляется из каждой GPU кучи отдельно. Для каждой кучи есть общий кеш памяти, в котором кешируется device memory объекты, и два вида субаллокаторов: Slab и Buddy.
struct Allocator<D: AllocDevice> {
device: D, // Абстрактный GPU device, из которого изначально выделяется память
heaps: Vec<Heap<D>>, // Количество куч фиксировано (обычно их не более двух)
}
struct Heap<D: AllocDevice> {
index: usize, // Индекс кучи
pool: HeapPool<D>, // Основной кеш
buddies: Option<HeapBuddies<D>>, // Buddy-субаллокатор
slabs: Option<HeapSlabs<D>>, // Slab-субаллокатор
}При запросе на аллокацию, алогоритм сначала попытается выделить из Slabs. Если не получится, то из Buddy. Если и оттуда не получится, то из Heap Pool кеша.
По умолчанию память нигде не выделяется, выделение памяти из device происходит по запросу.
Под термином pool я имею ввиду простой кеш ранее выделенных device memory объектов. Из данного пула память может быть выдана как непосредственно пользователю аллокатора, так и отдельным субаллокаторам.
struct HeapPool<D: AllocDevice> {
cache: Table<MemoryInfo, VecDeque<DeviceMemory<D>>, RandomState>,
used: AtomicI64, // Трекает (не очень точно) размер выделенной из кучи памяти.
limit: u64, // Ограничение на выделение. Это число обычно равно 3/4 реальной памяти кучи.
}
// Ключ таблицы кешей
struct MemoryInfo {
alloc_size: u64, // Размер аллокации
alloc_info: AllocInfo, // Тип аллокации (см. выше)
}Здесь Table это специальный конкурентный HashMap, допускающий параллельный мутабельный доступ к записям из разных потоков. Я не буду вдаваться в детали реализации этой структуры, она похожа на DashMap, и подробности можно почитать в документации к DashMap.
DeviceMemory это Arc-подобная структура (без Weak части). Если в списке VecDeque<DeviceMemory<D>> счетчик равен 1, значит данный объект не используется, и Heap Pool его может выдать, либо вернуть обратно GPU устройству. При аллокации (или чистке), соответствующий список перебирается линейно.
В целом, предполагается, что доступ к HeapPool осуществляется не очень часто, и не очень интенсивно.
Slab-субаллокатор сделан очень просто. Он выделяет DeviceMemory размером в 64 мегабайта, и может разделить его на не более чем 64 куска по одному мегабайту. Занятость кусочков контролируется по принциау Slab-структуры.
struct HeapSlabs<D: AllocDevice> {
cache: Table<AllocInfo, Slab<D>, RandomState>,
limits: SlabLimits,
}
// Собственно Slab структура
struct Slab<D: AllocDevice> {
device_memory: DeviceMemory<D>, // Выделенная из кеша память
next_chunk: u64, // Смещение внутри device_memory; следующий кандидат на субаллокацию
chunk_size: u64, // 1 мегабайт
total_size: u64, // 64 мегабайта
entries: Vec<SlabEntry>, // Связный список занятых и незанятых, но уже субаллоцированных кусков
next_entry: usize, // Следующий незанятый субаллоцированный кусок
entries_occupied: usize, // Общее число занятый кусков
}
// Субаллоцированный кусок из device_memory
struct SlabEntry {
offset: MemorySize, // Смещение куска внутри device_memory
next: Option<usize>, // Ссылка на следующий элемент связного списка.
// Если None, то кусок занят (аллоцирован аллокатором);
// если Some, кусок доступен для аллокации.
}Выделение и возвращение памяти заданного типа и размера не более чем 1 мегабайт из Slab происходит очень быстро, за O(1).
Кроме того, если значение entries_occupied равно нулю, мы можем принять решение об удалении Slab целиком, неявно вернув память обратно в HeapPool.
Данный субаллокатор предназначен для очень быстрого и интенсивного выделения памяти переменного, но небольшого размера. Например, для выделения небольших host-visible SSBO буферов или небольших host-visible картинок для их последующего копирования в device-local память.
Buddy-субаллокатор выделяет память размером 256 мегабайт, и (по запросу) делит ее на субаллокации размером степени двойки: 128, 64, 32, 16...
При делении мы каждый раз получаем два куска (64Mb это два куска по 32Mb). Один из двух полученных кусков, если он подходящего размера, возвращается пользователю, в противном случае он делится дальше. Другой кусок кладется в список доступных кусков.
При запросе памяти заданного размера, она либо сразу берется из списка свободных кусков, либо субаллоцируется описанным выше делением свободных кусков большего размера.
При освобождении памяти, происходит процесс, обратный вышеописанному алгоритму: если в списке свободных кусков есть парный освобождаемому куску кусок (братский кусок), мы объединяем их в больший (coalescence): пару по 32Mb объединяем обратно в 64Mb. Затем полученный кусок мы так же пытаемся объединить с его парой из списка свободных кусков, либо, если это невозможно (парный кусок использует), кладет данный кусок в список свободных.
// Выделенный из DeviceMemory кусок
struct Chunk {
offset: u64, // Смещение памяти внутри DeviceMemory
place: ChunkPlace, // Отношение куска к его родительскому куску
left: Subdivision, // Используется ли левая половинка
right: Subdivision, // Используется ли правая половинка
}
enum ChunkPlace {
// Данный кусок является исходным куском, занимающим всю память (256 Mb).
// У него нет родителя
Top,
// Данный кусок является левой половинкой своего родителя.
// Число - индекс родителя в предыдущем ранге деления.
Left(usize),
// Данный кусок является правой половинкой своего родителя.
Right(usize),
}
enum Subdivision {
// Дочерний кусок не используется (нет разбиения).
Unused,
// Дочерний кусок аллоцирован.
Used,
// Дочерний кусок не аллоцирован, но доступен в списке кусков для аллокации
// Число - индекс в списке следующего ранга деления
Free(usize),
}Buddy алгоритм это классический алгоритм аллокации из 60-х годов, описанный в частности Дональдом Кнутом. В худшем случае он работает за O(log2(M)) шагов на аллокацию и деаллокацию, где M в данном случае это 256 мегабайт, то есть всего не более чем за 28 шагов.
Во многих случаях, тем не менее, алогоритм может отработать гораздо быстрее при условии, что мы не будет постоянно выделять куски небольшого размера, и сразу их возвращать обратно. Тем не менее, для небольших кусков (до 1Mb) у меня есть более простой Slab субаллокатор. Куски меньшего размера будут выделяться из Buddy только в случае, если Slab окажется переиспользуемым.
В дополнение к этому, из Buddy аллокатора я разрешаю выделять куски размера не больше 64 мегабайт, и размера, меньшего половины текущей памяти, выделенной из Buddy субаллокатора. Это сделано для того, чтобы память Buddy субаллокатора в среднем использовалась по назначению: выделение нескольких кусков памяти среднего размера от 1 до 64 мегабайт, когда Slab субаллокатора уже не хватает, но брать память из кучи еще рано.
Отдельно хотелось бы сказать пару слов об организации списков кусков Buddy-субаллокатора.
Эти списки представляют собой так же slab-подобные связные списки, из которых поиск и взятие свободного куска, и возвращение ранее аллоцированного осуществляется гарантированно за O(1).
const MAX_DEPTH: usize = 40; // 1Tb
struct Buddy<D: AllocDevice> {
device_memory: DeviceMemory<D>,
height: usize, // Максимальная глубина деления
used: u64, // Объем текущей памяти
limit: u64, // 256 магабайт
// Списки занятых и доступных кусков, разделенных по их размеру.
// Индекс в массиве это ранг размера кусока -- степень двойки от размера куска: 2^(self.height-index)
lists: [List; MAX_DEPTH],
}
struct List {
entries: Vec<ListEntry>,
next: usize,
// Следующий кусок, доступный для аллокации
// Если free == usize::MAX, то доступных кусков нет.
free: usize,
}
enum ListEntry {
// Вакантное место внутри List
Vacant {
next: usize,
},
// Место занято либо аллоцированным куском, либо куском, доступным для аллокации.
Occupied {
chunk: Chunk,
prev: usize,
next: usize,
},
}Фактически List представляет собой список, кодирующий две независимые системы связности внутри своей структуры:
- Основной слой связности занятых и вакантных слотов.
- Слой занятых слотов, которые представляют собой доступные для аллокации куски.
Любопытно, что в целом вся реализация Buddy-субаллокатора не потребовала прямой работы с Box памятью и raw-pointers не смотря на то, что фактически это рекурсивная структура данных. Фактически реализация сделана полностью алгоритмически, и с использованием средств стандартной библиотеки Rust, но при этом достаточно эффективно.
Повторюсь, что мой аллокатор в целом сначала пытается выделить память из Slab-субаллокатора. Если не получилось, то из Buddy-субаллокатора. И если и оттуда не получилось, то из HeapPool кеша.
На каждом этапе фактически достаточно проверить одно/два числа -- есть ли в соответствующем слое субалокации свободное место.
Если места нигде нет (включая HeapPool), система пытается очистить память в два этапа:
- Сначала удаляются неиспользуемые Slab-субаллокаторы, Buddy-субаллокаторы и слоты HeapPool кеша, относящиеся только к данному типу памяти.
- Если после этого аллоцировать так же не получилось, то чистятся все незанятые Slab, Buddy и HeapPool кеши, которые можно освободить.
Для целей оптимизации я решил не создавать более одного Slab и Buddy субаллокатора на один тип памяти. Деление и так уже достаточно гранулярно. И если пямяти не хватило, то не хватило. Во всяком случае подход с единственной Slab и Buddy парой на каждый тип памяти автоматически решает проблему фрагментации между разными субаллокаторами.
Внутри субаллокаторов я не забочусь о выравнивании. При запросе куска конкретного размера, каждый субаллокатор выделяет тот размер памяти и в том месте DeviceMemory, в котором сочтет нужным. Выравнивание смещения выданного куска осуществляется уже после субаллокации. Поэтому для субаллокатора всегда запрашивается память немного большего размера: size + align - 1. Это консервативный подход, но экономия лишенго байта на поиск лучшего кандидата внутри субаллокатора скорее-всего была бы излишней.
Проблема buffer-image granularity в моем дизайне так же отсутствует, поскольку, как говорилось ранее, субаллокаторы создаются на каждый тип памяти, включаящий в себя тип линейности. Фактически у меня может быть отдельный Buddy субаллокатор, например, на картинки, и отдельный Slab субаллокатор, например, на небольшие буферы.
Я пока не проводил серьезных сравнительных тестов с другими решениями. Наивные тесты на аллокации разных размеров показывают ожидаемые результат: аллокации происходят на порядок быстрее, чем просто выделение объектов из GPU устройства.
Использование субаллокаторов так же происходит достаточно консистетно, во всяком случае на тех тестовых данных, которые я обозначил в начале статьи: множество интенсивных аллокаций переменного, но относительно небольшого размера, и аллокации данных большого, но более-менее констатного размера.
// Подготовка тестовых данных.
// Я пробовал разные сиды.
// Результаты консистентно схожи, но я выбрал тот, на котором в случае с моим
// аллокатором они будут наихудшими (для справедливости).
let mut rng = StdRng::seed_from_u64(13123);
let mut sizes = Vec::new();
for _ in 0..100 {
let size: u64 = match rng.gen_range(0u8..100) {
0..70 => rng.gen_range(1..1 << 10), // 70% данных это аллокации до 1Kb
70..90 => rng.gen_range(1 << 10..1 << 20), // 20% данных от 1Kb до 1Mb
90.. => rng.gen_range(1 << 20..1 << 26), //10% данных от 1Mb до 64Mb
};
sizes.push(size);
}
// Сравнительный тест: аллокация непосредственно из GPU Device.
{
let mut acc = Vec::new();
let time = Instant::now();
// 1000 повторений теста
for x in 1..1000 {
for size in &sizes {
let mut size = *size;
if size <= 1 << 22 {
size += x;
}
let alloc = DeviceMemory::new(
&device,
DeviceMemoryInfo {
alloc_size: size,
alloc_info: AllocInfo {
memory_type_index: 0,
linear: false,
map: false,
},
},
)
.unwrap();
acc.push(alloc);
}
// После каждого прохода освобождаем 100 DeviceMemory объектов.
acc.clear();
}
info!("Time: {:?}", time.elapsed()); // 1.558503413s на 100,000 аллокаций.
}
// Тот же тест, но с использованием моего аллокатора
{
let mut acc = Vec::new();
let time = Instant::now();
for x in 1..1000 {
for size in &sizes {
let mut size = *size;
if size < 1 << 22 {
// Если размер меньше 4Mb, то чуть-чуть меняем его на каждой итерации.
// Так тест получается более честным: аллокатор будет каждый раз выделять объекты разного размера.
size += x;
}
let alloc = allocator
.allocate(
1,
DeviceMemoryInfo {
alloc_size: size,
alloc_info: AllocInfo {
memory_type_index: 0,
linear: false,
map: false,
},
},
)
.unwrap();
acc.push(alloc);
}
acc.clear();
}
info!("Time: {:?}", time.elapsed()); // 98.853935ms на 100,000 аллокаций.
// HeapPool использовался 5994 раза
// Buddy использовался 29970 раза
// Slab использовался 63936 раза
info!("Statistics: {:?}", allocator.stats());
}Итак, получилось, что на примерно 100 тысяч аллокаций с виделением объектов разного размера, но удалением всех аллокаций после каждой сотни, мой аллокатор примерно в 15 раз быстрее родного. Много это или мало я пока не готов сказать.
Результаты могут сильно варьироваться в зависимости от вариантов использования. В данном тесте мы освобождаем всю память целиком после каждой итерации, кроме того, мы используем только один тип памяти.
Если разнообразить использование типов памяти, и в течение итерации время от времени случайно освобождать часть аллокаций, мой аллокатор показывает примерно те же 100 миллисекунд на 100 тысяч итераций, а скорость нативного аллокатора падает до 8.5 секунд. То есть выходит, что мой аллокатор в 85 раз быстрее на таком тесте.
Очевидно, что можно найти случаи, в которых мой аллокатор наоборот справится хуже, или вообще не справится. К примеру, если не освобождать объекты вообще, мой аллокатор сможет выделить в целом меньше большой памяти из-за сильной внутренней фрагментации Slab и Buddy.
Настройки Slab и Buddy (например, ограничение в 1 мегабайт на Slab кусок) можно и нужно варьировать в заивисимости от результат профилировщика. Разные настройки могут давать существенно разные результаты в зависимости от сценариев использования.
Еще мне хотелось бы сказать несколько слов о FreeList стратегии субаллокации, которую я не использую в своём проекте.
Идея данной стратегии похожа на гибрид вышеописанных Buddy и Slab. По сути это такой своего рода Slab, в котором мы выделяем куски не фиксированного, а произвольного размера. А когда пользователь освобождает кусок, мы пытаемся его объединить с ранее освобожденными соседями в один большой свободный кусок для уменьшения фрагментации.
В целом это очень разумный подход для случая, когда мы заведомо не знаем варианты использования. Если пользователь будет выделять и освобождать куски непредсказуемо вариабельного размера, я бы предположил что в среднем данный алгоритм будет эвристически более надеждным решением, чем моё решение без профилирования. FreeList, как мне кажется, является разумным компромиссом проблем внутренней и внешней фрагментации. По этой причине скорее-всего и gpu-allocator, и Vulkano аллокатор по-умолчанию используют FreeList.
Главным недостатком FreeList является то, что при сильной фрагментации скорость нахождения наиболее подходящего свободного куска вероятнее всего снизится. В случае моего Slab, в котором фрагментация фиксирована и на всякий случай избыточна, поиск небольшого свободного куска всегда происходит за O(1). Buddy аллокатор так же лишен проблем внешней фрагментации: память дефрагментируется достаточно предсказуемо.
Данная характеристика важна именно в случае моего проекта. Мои варианты использования предполагают очень частое выделение и освобождение буферов памяти и картинок переменного, но как правило небольшого размера. А большие аллокации, если они вообще будут, скорее-всего будут использоваться почти все время жизни программы.
В большинстве руководств к Вулкану, и другим похожим GAPI, говорится о том, что встроенный аллокатор не предназначен для прямого использования, и пользователь GAPI должен строить свой собственный субаллокатор поверх встроенного, либо использовать существующие решения (например, VMA).
Собственно, я таким путём и пошел. Написал свой собственный аллокатор, специализированный под свои нужды. И, действительно, он вроде бы дает неплохой выигрыш в скорости по сравнению со встроенным.
Тем не менее, если посмотреть даже на мои тривиальные бенчмарки, 8.5 секунд на 100,000 аллокаций, честно говоря, это не то чтобы очень медленно. Во всяком случае такой результат я получил на своем очень стареньком Radeon Vega M, и я бы предположил, что на других десктопных платформах результат будет вряд ли хуже.
Поэтому предостережения статей по Вулкану, как мне кажется, во многих случаях избыточны для не AAA-проектов. Если в вашей игре будет задействовано небольшое число объектов памяти, и вы не планируете их пересоздавать слишком часто, я думаю, что можно обойтись вообще без специального субаллокатора, ну или в худшем случае можно использовать какую-то очень простую стратегию кеширования ранее выделенных объектов памяти. А преимуществом данного подхода будет отсутствие мучений с выравниваниями и прочими image-buffer granularity.
Footnotes
-
Стоит отметить, что Vulkano так же предлагает Buddy и Bump субаллокаторы в качестве альтернативы. ↩