Привет, Хабр! Я Кирилл Колодяжный, разработчик систем хранения данных в YADRO, ML-энтузиаст и автор книги "Hands-on Machine Learning with C++". В своих материалах я развеиваю миф о том, что машинное обучение — это сплошной Python. На самом деле под капотом моделей часто работает C++.
Этой теме я посвятил цикл статей: хочу рассказать, как привычные для «плюсовиков» инструменты используют для реализации ядра платформы машинного обучения. В первой части поговорим о стандартных библиотеках, идиомах программирования и алгоритмах управления памятью.
Где в ML пригодится C++
Когда речь заходит о машинном обучении, обычно представляют Python. Однако мне это кажется не совсем справедливым, поскольку в основе большинства систем лежит C++. Вот где явно применяют C++ в ML:
В инференсах моделей на edge-устройствах, например на телефонах или IoT. А еще в средах, где развертывание инфраструктуры Python нецелесообразно и требуется максимальная производительность. Например, в приложениях llama.cpp или gemma.cpp.
В реализации математических библиотек. Эти библиотеки лежат в основе всех популярных фреймворков, таких как TensorFlow и PyTorch. Например, библиотека линейной алгебры Eigen изначально использовалась в TensorFlow для математических вычислений.
Сегодня я сосредоточусь на «плюсовых» компонентах, которые выступают в роли связующих элементов для различных частей ML-платформ.
Об использовании С++ в ML я также рассказываю в рассылке для «плюсовиков». Вместе с коллегами из YADRO мы пишем о необычных задачах, с которыми столкнулись на работе, делимся примерами вопросов с собеседований и полезными материалами для изучения С++. Письма приходят раз в месяц, подписывайтесь →
Рассмотрим упрощенный процесс обучения нейронной сети. Каждый слой можно представить как математическую абстракцию: тензор, матрицу или их набор. Работа нейронной сети — это последовательность математических операций над тензорами, чаще всего это операции линейной алгебры, например матричное умножение.
Такие операции обычно реализуются в специализированных математических библиотеках. Далее я покажу, как реализовать один из таких ключевых компонентов — интерфейс для операций с линейной алгеброй — и какие архитектурные подходы и паттерны для этого наиболее применимы.
И еще один важный аспект, который рассмотрим в статье: математические библиотеки часто имеют несколько внутренних реализаций. Но при этом предоставляют единый интерфейс для пользователей, с которым мы взаимодействуем в среде машинного обучения.
В основе этой статьи — мои наработки по проекту с открытым кодом Adept. На данный момент это в большей степени образовательный проект, где можно познакомиться с примерами реализации основных компонентов ML-платформ: от реализации тензорной арифметики с помощью SIMD-интрисинков и вычислительных шейдеров до Python-интерфейса, который доходит до конечного пользователя, например эксперта по анализу данных или разработчика нейронных сетей. В этом материале буду показывать несколько упрощенные примеры, поэтому для полного погружения переходите в репозиторий.
Математическая библиотека: фасад, с которым сталкивается пользователь
Основной интерфейс математической библиотеки — это обычно абстракция тензора или матрицы. Какие операции нам от нее нужны?
Конструирование объекта тензора из значений
С помощью std::initializer_list можем проинициализировать тензор конкретными значениями. С каждым тензором также связаны его форма (количество измерений) и размер каждого измерения, тип данных и устройство, на котором мы хотим делать вычисления. У класса тензора обычно переопределяются математические операторы, с помощью которых может строиться граф вычислений, как в примере со сложением:
auto x = Tensor::from_values({1.f, 3.f}, Shape{1, 2}, device_t::CPU); auto y = ...; auto z = x + y;
Разные способы инициализации
Чаще всего мы работаем с тренировочными данными и получаем сырые указатели void*. При этом мы знаем все свойства тензора: размер, тип данных и целевое устройство. Поэтому интерфейс должен поддерживать разные способы инициализации. Так что в статье мы будем работать с формой тензора и перечислениями, которые указывают на устройство вычисления и тип данных.
class Shape { public: Shape() = default; Shape(const std::vector<index_t>& dims); Shape(std::initializer_list<index_t> dim_values); ... private: std::vector<index_t> dims_; };
enum class device_t : uint8_t { CPU = 1, GPU, EndOfDevices }; enum class dtype_t { Float32, Float64, Int32, Int8 };
Почему тип данных я задал через перечисление? Если бы мы работали только в рамках C++ API, то реализовали бы тензоры через темплейт, параметризовали статическим типом и наверняка выиграли бы в производительности. Например, библиотека Eigen так и делает.
Но популярные платформы машинного обучения, как правило, предоставляют внешний Python-интерфейс, что приводит к связке двух языков: C++ и Python. Например, динамическое управление типами в Python нужно «занести внутрь» платформы машинного обучения на С++. Поэтому далее в тексте расскажу, как работать с динамическим типом данных.
Применяем идиому pImpl в неклассическом виде
Как я уже упоминал, у класса Tensor могут быть разные внутренние реализации для CPU, GPU или других устройств. Один из способов гибкого переключения между ними — использование идиомы pImpl (Pointer to Implementation).
В этом случае публичный класс Tensor содержит лишь статический интерфейс и указатель на скрытую реализацию, а идиома pImpl позволяет:
подменить реализацию,
сократить время компиляции, вынося результат реализации в отдельные единицы трансляции,
четко отделить интерфейсы от реализации.
Интерфейс тензора может выглядеть так:
class Tensor { public: ... Tensor operator+=(const Tensor& other); Tensor operator-=(const Tensor& other); Tensor operator*=(const Tensor& other); ... Tensor mean() const; Tensor matmul(const Tensor& other) const; ... };
Здесь видим переопределенные математические операторы и методы для математических операций. А единственным полем в таком классе будет указатель на абстрактный класс DeviceTensor, который содержит много виртуальных функций:
class Tensor { public: ... private: std::shared_ptr<DeviceTensor> impl_; };
Использование идиомы сводится к простому перенаправлению вызовов. Каждый метод класса Tensor делегирует работу нужной нам реализации, которую мы указали при инициализации этого класса:
Tensor Tensor::sub_(const Tensor& other) { ... impl_->sub(*other.impl_); return {impl_}; }
Внутри конкретной реализации мы уже запрограммируем быстрый GPU-алгоритм с помощью CUDA или воспользуемся SIMD для быстрых вычислений на CPU:
void CPUTensorImpl::sub(const DeviceTensor& other) { const auto& other_tensor = static_cast<const CPUTensorImpl&>(other); ... FastSIMDSubtraction(*this, other_tensor); }
Скорее всего, нам придется приводить типы для конкретных реализаций. Но сделать это будет достаточно просто, потому что у тензора есть свойства, в которых указано и устройство, и тип данных, хранящихся внутри. Придется немного поработать за компилятор: не так эффективно, зато позволяет динамически управлять реализациями и типами.
Еще можно обратить внимание что в классе Tensor для хранения указателя на реализацию был использован умный указатель с подсчетом количества ссылок std::shared_ptr. Это сделано намеренно, для реализации механизма управления времени жизни объектов как в Python, т.е. при копировании объекта не происходит копирования памяти, это в своём роде применение паттерна проектирования Flyweight. Такой подход позволяет избежать многократного копирования больших объемов памяти при формировании математических выражений и сделать семантику C++ API максимально похожей на Python API. Копирование же данных будет происходить только при создании новых тензоров как результатов вычислений.
Отмечу, что это не классическое использование идиомы pImpl. Кроме разделения интерфейса и реализации, мы еще и добавляем работу с полиморфными реализациями.
Создаем реализацию для тензора
Один из самых очевидных подходов — использовать паттерн «Фабрика». Мы можем создать абстрактный интерфейс с реализациями для CPU, GPU и других устройств. Кроме этого, фабрика может отвечать за управление ресурсами, которые необходимы при создании тензоров.
В статье разберем пример, как фабрика управляет выделением памяти. Ее главная задача — конструировать тензор нужного типа и предоставлять ему корректно выделенную память. Основной интерфейс такой фабрики — набор методов для создания тензоров из различных источников, например from_blob() для работы с сырыми данными.
Одним из основных полей в моей фабрике будет словарь, где я буду хранить объекты типа MemoryPool, которые управляют памятью для тензоров. Эти MemoryPool будут инициализироваться с помощью специальных аллокаторов. Здесь, обратите внимание, я использовал flat_hash_map из сторонней библиотеки.
namespace adept::cpu { class TensorFactory { public: ... private: ska::flat_hash_map<dtype_t, std::shared_ptr<MemoryPool>> mem_pools_; };
Такой контейнер эффективен, когда операции чтения происходят часто, записи — редко, а элементов не очень много. Отмечу, что аналог такого словаря — std::flat_map — планируют включить в C++ 26.
Настраиваем аллокаторы и MemoryPool
Аллокаторы и MemoryPool нужны, потому что внутренние реализации тензора предназначены для разных платформ и процессоров, и часто для максимальной производительности требуется выделять память особым образом.
К тому же MemoryPool может работать как кеш для выделенных участков памяти. При обучении нейронной сети размер тензоров и матриц будет повторяться, а нам не нужно будет каждый раз запрашивать у системы новую память. Мы можем ею воспользоваться, отложить в кеш, а на следующем слое, когда понадобится такой же тензор, взять его снова. Например, такой подход используется в PyTorch, когда его собирают под мобильные платформы. Другим вариантом может быть использование арен, когда память резервируется большим блоком arena, а новые объекты размещаются последовательно. При этом освобождение происходит одномоментно при удалении всей арены.
Интерфейс аллокатора предельно простой — методы allocate и free. Внутри используем, например, aligned allocations, то есть выровненную в соответствии с некоторыми требованиями память. Здесь AllocateAlign — функция из библиотеки кросплатформенных SIMD-интринсиков Google Highway. Дальше я уже могу не задумываться, как именно был выделен буфер и как его следует освободить. Эта ответственность полностью делегирована аллокатору и объекту типа Buffer внутри MemoryPool, где интерфейсы для буфера могут быть определены так:
class Buffer { public: virtual ~Buffer() {} virtual size_t size() const = 0; virtual bool is_shared() const { return false; }; virtual bool is_aligned() const { return false; }; }; class CPUBuffer : public Buffer { public: virtual void* ptr() const = 0; }; typedef std::shared_ptr<Buffer> buffer_ptr_t;
А класс аллокатора определяется так:
template <typename T> requires std::integral<T> || std::floating_point<T> class SIMDAllocator : public Allocator { public: buffer_ptr_t allocate(size_t num_elements) { auto buffer = hwy::AllocateAligned<T>(elements); auto mem_size = to_mem_size(elements); return std::make_shared<AlignedBuffer<T>>(std::move(buffer), mem_size); } void free(buffer_ptr_t&& buffer); };
А реализация выровненного буфера для эффективного использования в SIMD операциях выглядит так:
template <typename T> class AlignedBuffer : public CPUBuffer { public: AlignedBuffer(hwy::AlignedFreeUniquePtr<T[]> data, size_t size) : data_(std::move(data)), size_(size) {} ~AlignedBuffer() = default; ... void* ptr() const override { return data_.get(); } size_t size() const override { return size_; } private: hwy::AlignedFreeUniquePtr<T[]> data_; size_t size_; };
Теперь самое интересное. Google Highway предоставляет шаблонные аллокаторы, специализированные для каждого типа данных, поэтому для корректной работы нужно создать отдельный пул памяти для каждого типа, который будет хранить нужный аллокатор.
Математические библиотеки обычно работают с небольшим и заранее известным набором типов — например, с несколькими типами чисел с плавающей запятой и целочисленными типами. Этот список обычно известен на этапе компиляции.
Основная задача — преобразовать этот список типов в набор соответствующих пулов памяти. Для этого можно применить подход, при котором для каждого типа из списка вызывается конструирующая функция mempool. На практике это можно реализовать с помощью лямбда-функции, параметризованной нужным типом. Такая лямбда выполняет две ключевые операции:
Получает enum-идентификатор типа, который будет использоваться как ключ для доступа к пулу в словаре.
Использует переданный тип для инициализации аллокатора.
Так во время инициализации фабрики для каждого поддерживаемого типа создается и регистрируется свой собственный, корректно настроенный пул памяти.
TensorFactory::TensorFactory() { auto make_mem_pool = [this](auto arg) { using T = decltype(arg); mem_pools_[to_dtype<T>()] = std::make_shared<MemoryPool>(std::make_shared<AlignedAllocator<T>>()); }; for_each_data_type<float32_t, float64_t, int32_t, int8_t>(make_mem_pool); }
Как преобразовать статический тип в динамический
Преобразовать статический тип в значение enumeration достаточно просто, если использовать выражение if-constexpr. Порядок действий:
Напишем такую цепочку
if-else, в которой используемtype-trait, который сравнивает типыstd::is_same.Сравним нужные нам типы с
enumeration.Вернем эти
enumeration.
template consteval dtype_t to_dtype() { if constexpr (std::is_same_v<T, float32_t>) { return dtype_t::Float32; } else if constexpr (std::is_same_v<T, float64_t>) { return dtype_t::Float64; } ... }
Поработать со списком типов можно с помощью Boost, у которого есть серьезные библиотеки для метапрограммирования. А можно сделать проще — использовать лямбду и функцию с переменным числом аргументов шаблона for_each_data_type.
... auto make_mem_pool = [this](auto arg) { using T = decltype(arg); ... } ... template <typename... DataTypes, typename F> void for_each_data_type(F&& lambda) { std::tuple<DataTypes...> t; std::apply([&lambda](auto&&... arg) { (lambda(arg), ...); }, t); }
Дальше через объект типа std::tuple из стандартной библиотеки, инициализированный с помощью template parameter pack, можно применить функцию std::apply для каждого элемента tuple. Конструирующая функция make_mem_pool принимает один аргумент, и с помощью decltype мы выведем тип и используем его для инициализации аллокатора. Такой подход имеет право на жизнь, но он не очень эффективный, так как мы создаем и копируем ненужные объекты только для получения типа.
Можно также использовать variadic template, с comma-fold expression, но взять параметризируемую шаблоном лямбду.
... auto make_mem_pool = []<typename T>() { ... } ... template <typename... DataTypes, typename F> void for_each_data_type(F&& func) { (void(func.template operator()<DataTypes>()),...); }
Тут мы не создаем промежуточных типов, однако синтаксис выходит достаточно специфический, на мой взгляд, и вот почему:
для вызова лямбды мы должны использовать синтаксис
.templateи оператор функции.запятая и три точки — это
comma-fold expression, которое вызовет нашу лямбду для каждого из типов.
Так мы создадим пулы памяти под каждый тип данных.
Реализацию алгоритма кеширующего MemoryPool мы пока рассматривать не будем, его базовая реализация может быть построена на основе LRU-кеша.
Создаем объект tensor
Рассмотрим случай, когда тензор создается из сырых данных.
std::shared_ptr<TensorImpl> TensorFactory::from_blob(const void* data, const TensorProperties& props) { ... auto numel = props.shape.numel(); auto buffer = mem_pools_[props.dtype]->allocate(numel); DISPATCH_TYPE(props.dtype, [&](){ hwy::CopyBytes(data, buffer->ptr(), numel * sizeof(scalar_t)); }); return std::make_shared<TensorImpl>(props, std::move(buffer)); }
Первый аргумент функции — это указатель на данные, а второй — структура TensorProperties, которая описывает метаданные: размерности тензора и тип хранимых данных. Из свойств можем получить количество элементов. По типу данных находим нужный memory pool, просим его создать буфер и конструируем тензор с нужными свойствами и буфером.
Однако тут возникает следующий момент: данные хоть и сырые, но тип динамический, заданный через enumeration. Пришел откуда-то извне, из Python прилетел. И с данными надо работать, используя конкретный тип. Откуда нам его получить? И как сделать так, чтобы мы могли данные эффективно переиспользовать? Я предлагаю подсмотренный в PyTorch подход — макросы.
Как использовать макросы с умом
DISPATCH_TYPE — это макрос. Он принимает на вход значение перечисления, определяющего тип данных, и лямбда-функцию, в которой будет доступен тип scalar_t, соответствующий конкретному типу данных (например, float, double, int). Да, макросы — это зло в С++, однако иногда, особенно в связки с шаблонами, они здорово облегчают жизнь, позволяют создать внутренний DSL (Domain Specific Language) и уменьшить количество кода.
Конкретно макрос DISPATCH_TYPE разворачивается в стандартный switch statement. Для каждого варианта в перечислении создается блок, внутри которого с помощью using определяется псевдоним scalar_t, соответствующий конкретному типу. Это позволяет писать обобщенную логику внутри лямбды, которая будет инстанцирована для каждого поддерживаемого типа данных. А после, в этом же скоупе, вызываем заданную лямбду.
Как можно сделать такой макрос? Сейчас научу плохому. Сначала описываю макрос, который создает одно выражение: принимает два аргумента, значение enumeration и конкретный тип данных.
#define DISPATCH_CASE(type_enum, type, ...) \ case type_enum: { \ using scalar_t = type; \ return __VA_ARGS__(); \ }
Макрос — просто строки, которые комбинируются в код, а позже будут скомпилированы. Макрос мы используем с переменным количеством аргументов, и применяем конструкцию __VA_ARGS__, чтобы обрабатывать список аргументов макроса.
Тут стоит обратить внимание, что после
__VA_ARGS__ идут две круглых скобки. Тут лямбда инстанцируется и сразу вызывается в этом месте.
Следующий макрос — это тот, который формирует нам полный switch statement.
#define DISPATCH_SWITCH(type_enum, ...) \ [&] { \ switch (type_enum) { \ __VA_ARGS__ \ default: \ THROW_ERROR(#type_enum " is not supported in current context"); \ } \ }()
В качестве параметров принимает значение enumeration для типа данных, а в качестве всех остальных параметров мы сюда передадим сгенерированные кейсы.
Объединить все вместе можно так:
#define DISPATCH_CASE_ALL_TYPES(...) \ DISPATCH_CASE(dtype_t::Float32, float32_t, __VA_ARGS__) \ DISPATCH_CASE(dtype_t::Float64, float64_t, __VA_ARGS__) \ DISPATCH_CASE(dtype_t::Int32, int32_t, __VA_ARGS__) #define DISPATCH_TYPE(TYPE, ...) \ DISPATCH_SWITCH(TYPE, DISPATCH_CASE_ALL_TYPES(__VA_ARGS__))
Сначала описываем все доступные нам case-случаи: DISPATCH_CASE_ALL_TYPES. Мы перечислили два floating point-типа и integer. В PyTorch таких наборов много, они отличаются, и DISPATCH_TYPE есть разные. Например, потому что в CUDA некоторые ядра работают только с одними типами данных. Некоторые ядра эффективно работают только с половинной точностью. А другое ядро может работать со всеми типами, но будет чуть менее производительным. И в конце это объединяется в макрос DISPATCH_TYPE, который внутри использует switch и все определенные case statements.
А для пользователя, мне кажется, это выглядит достаточно просто. Мы используем вызов лямбды с нужным типом. Последний берем из свойств тензора, определяем лямбду — она сразу вызовется в этом месте. В лямбде доступен тип scalar_t, который можем использовать для преобразования типов данных.
std::shared_ptr<DeviceTensor> TensorImpl::sum() const { auto new_tensor = TensorFactory::instance().empty(...); DISPATCH_TYPE(props_.dtype, [&]() { const auto* in_data = const_data_ptr<scalar_t>(); auto* out_data = new_tensor->mutable_data_ptr<scalar_t >(); out_data[0] = parallel_reduce<scalar_t>(0, props_.shape.numel(), [&](auto begin, auto end, scalar_t running_sum) { scalar_t partial_sum = 0; detail::vector_sum<scalar_t>(in_data + begin, partial_sum, end - begin); return running_sum + partial_sum; }, std::plus<scalar_t>()); }); return new_tensor; }
Это пример подсчета суммы элементов тензора уже с использованием параллелизации с помощью функции parallel_reduce, которая параметризируется конкретным типом. Это нужно для более эффективной работы с предзагрузкой данных в кеш процессора. Здесь я использовал библиотеку oneTBB. Функтор из стандартной библиотеки std::plus тоже специализирован типом scalar_t. Это значит, что такой тип можно легко использовать и писать алгоритмы в конкретных типах, один раз написав DISPATCH_TYPE.
В этой части мы рассмотрели, как может выглядеть интерфейс математической библиотеки, как задать реализации под конкретные устройства и как работать с динамическими типами, используя привычные нам подходы из C++. Отмечу, что эти методы можно использовать как для машинного обучения, так и для других проектов.
У статьи будет продолжение: поговорим о диспетчеризации вызовов свободных функций и автодифферинцированнии, которое используется в обучении нейронных сетей для реализации алгоритмов обратного распространения ошибки. Подписывайтесь на блог YADRO, чтобы не пропустить вторую часть.