Миф о RAM — это верование о том, что память современного компьютера напоминает идеальную память с произвольным доступом. Кэш люди считают оптимизацией для малых данных: если они умещаются в L2, то будут обрабатываться быстрее; если нет, то тут уж ничего не поделаешь.

Вероятнее всего, что самым быстрым разбиения данных будет такой код (я использую в качестве псевдокода Python; можете представить, что я пишу это на вашем любимом низкоуровневом языке):

groups = [[] for _ in range(n_groups)]
for element in elements:
    groups[element.group].append(element)

Он и в самом деле линеен (то есть асимптотически оптимален), и мы всё равно должны выполнять доступ к произвольным индексам, так что кэш здесь нам ни в чём бы не помог.

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

Решение

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

elements.sort(key = lambda element: element.group)

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

elements.sort(key = lambda element: element.group)
groups = [
    group_elements
    for _, group_elements
    in itertools.groupby(elements, key = lambda element: element.group)
]

Существует множество алгоритмов сортировки, учитывающих кэш, но поскольку индексы — это просто целые числа, здесь больше всего подойдёт поразрядная сортировка (radix sort). Среди всех готовых реализаций мне больше всего подошла в Rust radsort.

Ускорение

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

Генераторы

API сортировки пытаются создать впечатление, что данные сортируются на месте, даже если это не так. Для этого нужно, чтобы отсортированные данные явным образом записывались в память в определённом формате. Но если нам нужно выполнять итерации для групп, этого позволяют избежать генераторы или обратные вызовы:

# Предполагаем использование 32-битных индексов
def radix_sort(elements, bits = 32):
    # Базовый случай -- сортировать нечего или все индексы одинаковы
    if len(elements) <= 1 or bits <= 0:
        yield from elements
        return

    # Разбиваем по старшему байту индекса, если ещё его не видели
    buckets = [[] for _ in range(256)]
    for element in elements:
        buckets[(element.index >> max(0, bits - 8)) & 0xff].append(element)

    # Рекурсивно сортируем бакеты
    for bucket in buckets:
        yield from radix_sort(bucket, bits - 8)

Мы можем даже избавиться от этапа groupby и выполнять yield отдельных групп:

# Базовый случай -- сортировать нечего или все индексы одинаковы
if bits <= 0:
    if elements:
        # Группа!
        yield elements
    return

Перераспределения

Ещё одна проблема этого кода заключается в том, что он постоянно выполняет перераспределения массивов bucket при append. Это приводит к тому, что memcpy вызывается чаще необходимого, и это плохо для кэша. Часто для решения этой проблемы размеры вычисляют заранее:

def get_bucket(element):
    return (element.index >> max(0, bits - 8)) & 0xff

sizes = Counter(map(get_bucket, elements))

# На самом деле, Python не может оставлять место для списков, но притворимся, что `reserve` всё равно это делает.
# В C++ это `std::vector::reserve`. В Rust это `Vec::with_capacity`.
buckets = [reserve(sizes[i]) for i in range(256)]
for element in elements:
    buckets[get_bucket(element)].append(element)

Однако это требует двух итераций, а нам бы идеале хотелось, чтобы код выполнялся за один проход. Если индекс случаен, то мы можем убить одним выстрелом двух зайцев: приблизительно оценить размер каждого бакета как len(elements) / 256 и зарезервировать пространство под них. Если наша приблизительная оценка будет заниженной, то возникнут остатки, которые можно сохранить в отдельном маленьком хранилище:

class Bucket:
    reserved: list
    leftovers: list

    def __init__(self, capacity):
        self.reserved = reserve(capacity) # псевдокод
        self.leftovers = []

    def append(self, element):
        if len(self.reserved) < self.reserved.capacity(): # псевдокод
            self.reserved.append(element)
        else:
            self.leftovers.append(element)

    def __len__(self):
        return len(self.reserved) + len(self.leftovers)

    def __iter__(self):
        yield from self.reserved
        yield from self.leftovers

Здесь вероятность распределения ведёт себя честно: при больших входных данных лишь крошечный процент элементов переполняется и попадает в leftovers, поэтому лишние затраты памяти относительно малы, перераспределения при отправке в leftovers быстры, а разбивка на бакеты (и итерации по бакетам) удобна для кэша.

Разбиение

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

Базовый случай

Переход к прямолинейному алгоритму при малом объёме входных данных повышает производительность, потому что в этом случае более явственно проявляется влияние кода O(n log n). Однако поскольку radix_sort рекурсивна, мы можем выполнять эту проверку на каждом уровне рекурсии, обеспечивая выигрыш даже при большом объёме данных:

# Базовый случай -- достаточно мал, чтобы использовать прямолинейный алгоритм
if len(elements) <= CUTOFF or bits <= 8:
    counts = [0] * 256
    for element in elements:
        counts[element.index & 0xff] += 1

    groups = [[] for _ in range(256)]
    for element in elements:
        groups[get_bucket(element)].append(element)

    for group in groups:
        if group:
            yield group
    return

Оптимальное значение CUTOFF сильно зависит от машины. Оно зависит от относительной скорости уровней кэша и RAM, а также от размеров кэша и типов данных. В случае 64-битных integer я встречала машины, для которых оптимальным значением было 50k, 200k и 1M. Наилучший способ определить его — бенчмарк в среде выполнения; это приемлемое решение для длительно работающего ПО, например, для баз данных.

Бенчмарк

Структура

Вот небольшой бенчмарк.

Входные данные — это массив случайных 64-битных integer. Мы хотим сгруппировать их по простому хэшу, основанному на умножении, и выполним простой анализ бакетов, допустим, вычислим сумму минимумов среди бакетов. (В реальности бакеты бы использовались ниже по конвейеру каким-то другим дружественным к кэшу алгоритмом.)

Мы сравним две реализации:

1. Прямолинейный алгоритм с оптимизированными распределениями.

2. Группирование на основании поразрядной сортировки, со всеми представленными выше оптимизациями и оптимальной отсечкой (cutoff).

Средний размер группы равен 10.

Код выложен на GitHub.

Результаты

Относительная эффективность оптимизированного алгоритма растёт с увеличением объёма данных. И прямолинейный, и оптимизированный алгоритмы в конечном итоге останавливаются на фиксированной скорости обработки. В зависимости от машины, в пределе улучшения могут составлять от 2,5 до 9 раз.

Результаты (A, Y, M — это разные устройства):

Заключение

Стоило ли оно того? Если вам абсолютно необходима производительность и сегментирование — важная часть вашего конвейера, то любыми способами постарайтесь это использовать. Например, я использую это для поиска хэша без коллизий для заданного набора данных. Но как и в случае с любой оптимизацией, вам нужно разобраться, стоит ли расплачиваться за неё повышением сложности кода.

По крайней мере, если вы работаете с big data, то про этот трюк стоит помнить.

Есть и ещё один вывод: все знают, что при работе с данными на диске не стоит просто отображать их в память и выполнять типичные алгоритмы работы с памятью. Это возможно, но производительность будет плохой. Урок заключается в том, что это применимо также к RAM и кэшу: если у вас больше, допустим, 32 МиБ данных, то нужно серьёзно задуматься над разбиением данных или над переходом к алгоритмам работы с внешней памятью.