Привет, Хабр! Я Кирилл Колодяжный, разработчик систем хранения данных в YADRO и ML-энтузиаст. Продолжаю рассказывать о паттернах С++, которые легко адаптировать под задачи машинного обучения. В этой части поговорим о динамическом полиморфизме — технологии, которая помогает объединить интерфейс для запуска вычислений с разными условиями. Ссылку на первую часть найдете в конце статьи.

Что такое динамический полиморфизм

Когда разрабатывали высокопроизводительные математические библиотеки для машинного обучения (PyTorch или TensorFlow), возникла фундаментальная задача: как обеспечить единый интерфейс вычислений при наличии множества аппаратных платформ, вычислительных бэкендов и режимов работы?

Рассмотрим типичный сценарий:

Tensor loss(Tensor x, Tensor y) {
    Tensor z = mul(x, y);
    return mean(z);
}

Алгоритм здесь описан один раз и не должен меняться. Однако реализация операций mul и mean выбирается динамически в зависимости от:

  • Свойств аргументов — на каком устройстве (CPU/GPU) расположены данные тензоров.

  • Глобального контекста — какой бэкенд (SIMD, Vulkan, CUDA, OpenCL), какой режим работы inference или training сейчас активен.

  • Внешних условий — решений, которые пользователь принимает через Python API или UI.

Это и есть суть динамического полиморфизма для свободных функций: выбор реализации происходит не на этапе компиляции (как при использовании шаблонов) и не через виртуальные методы объектов во время выполнения, а на основе значений аргументов и контекста системы тоже во время выполнения. Ключевое отличие от классического ООП-полиморфизма — выбор реализации зависит не от типа объекта, а от значений его характеристик (значение поля device и текущего бэкенда системы).

Почему стандартные методы не подходят

Классический путь через виртуальные функции здесь применять нельзя: операции над тензорами чаще всего реализуются как свободные функции, а не методы класса. Это нужно, чтобы новые операции можно было добавлять без изменения класса. Наивный подход через if-else быстро становится неуправляемым:

Tensor mul(const Tensor& x, const Tensor& y) {
    if (x.device() == CPU && current_backend() == SIMD)
        return cpu::simd::mul(x, y);
    else if (x.device() == GPU && current_backend() == Vulkan)
        return gpu::vulkan::mul(x, y);
    else if (x.device() == GPU && current_backend() == CUDA)
        return gpu::cuda::mul(x, y);
    // ... и так для каждой операции и каждой комбинации устройств/бэкендов
}

Чем больше операций и платформ (CPU SIMD, GPU Vulkan/CUDA/OpenCL, NPU), тем сложнее расширять и поддерживать код.

Решение, примененное в PyTorch и реализованное в проекте Adept, использует централизованный механизм диспетчеризации.

Архитектура системы диспетчеризации

Система строится вокруг трех взаимодействующих компонентов, образующих иерархию от конкретного вызова до глобального управления:

  • Dispatcher — глобальный фасад, хранящий все таблицы функций.

  • FunctionTable — таблица реализаций для конкретной операции, например mul.

  • dispatch_key_t — ключ, кодирующий характеристики вызова (устройство + бэкенд + контекст).

Рассмотрим, как это работает, на примере вызова функции умножения тензоров mul — начнем с самого верхнего уровня:

Tensor x;
Tensor y;
Tensor z = mul(x, y);

Как устроена функция mul 

Используем интерфейс диспетчера и выполняем следующие действия: 

  1. Находим объект представления функции по ее названию. 

  2. Производим вызов функции тоже с использованием диспетчера.

Tensor mul(const Tensor& x, const Tensor& y) {
    static auto& func = Dispatcher::instance().find("mul");
    return Dispatcher::instance().call<Tensor, const Tensor&, const Tensor&>(
        func, x, y
    );
}

Здесь ключевой момент — явное указание типов шаблона при вызове call<...>(). Несмотря на то, что современные «плюсовые» компиляторы поддерживают автоматический вывод шаблонных параметров, полагаться на него опасно. Причина кроется в механизме стирания типов (type erasure) в основе хранения функций.

Почему необходим ручной контроль типов

Объект функции в данном случае представлен типом FunctionTable, где каждая реализация функции упаковывается в std::function, а затем — в std::any. Соответственно, при вызове требуется обратная операция — извлечение функции из std::any с точным совпадением сигнатуры:

auto func = std::any_cast<std::function<Result(Args...)>>(entry);

Если типы Result и Args... не со��падут с оригинальной сигнатурой, зарегистрированной при добавлении функции, будет выброшено исключение std::bad_any_cast.

Автоматический вывод типов может привести к неочевидным расхождениям:

Переданный аргумент

Автоматически выведенный тип

Ожидаемый тип регистрации

const Tensor& x

Tensor (из-за value category)

const Tensor&

Tensor&& tmp

Tensor& (из-за ссылочного коллапса)

Tensor&&

Подобные расхождения особенно вероятны при работе с: 

  • универсальными ссылками (T&& в шаблонах),

  • константными квалификаторами (const),

  • временными объектами (rvalue).

Реализация класса FunctionTable

FunctionTable — это аналог таблицы виртуальных функций (vtable) из ООП. Но в отличие от классического ООП, где выбор реализации определяется типом объекта, здесь ключом диспетчеризации становятся характеристики аргументов и контекста выполнения. Таблица представляет собой отображение «ключ → реализация», где ключ кодирует набор параметров, влияющих на выбор конкретной версии функции.

Название: wikilink - описание: 1000

В этом примере ключ включает только тип вычислительного устройства. Однако архитектура поддерживает многомерные ключи и позволяет учитывать дополнительные измерения. К ключу можно добавить точность вычислений (FP32/FP16), требования к автодифференцированию, специализированные оптимизации для тензоров определенной формы и другие характеристики, при этом не нарушить единую архитектуру диспетчеризации.

Основная задача класса — хранение функций с различными сигнатурами в единой таблице. Решение достигается через механизм стирания типов (type erasure) с использованием комбинации std::function и std::any. Альтернативное решение  — объявление FunctionTable как шаблонного класса с параметрами сигнатуры:

template<typename Result, typename... Args>
class FunctionTable { /* ... */ };

Однако такой подход приводит к ряду проблем:

  • Взрывное увеличение количества инстанциаций шаблонов — для каждой уникальной сигнатуры создается отдельная таблица.

  • Сложность централ��зованного управления — диспетчеру пришлось бы хранить гетерогенную коллекцию таблиц разного типа, что вновь потребовало бы стирания типов на более высоком уровне.

  • Усложнение интерфейса регистрации — каждая реализация должна была бы явно указывать полную сигнатуру при добавлении в таблицу.

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

PyTorch использует похожий подход, но с важным отличием: вместо стандартных std::any и std::function применяются кастомные аналоги, разработанные специально для фреймворка.

Итак, класс FunctionTable может быть определен следующим образом:

struct FunctionTable final {
  FunctionTable(const std::string& name) : name_(name) {}

  template <typename Func>
  void add_func(const dispatch_key_t& key, const Func& func);

  template <typename Result, typename... Args>
  Result call(const dispatch_key_t& key, Args... args);
  
  std::string name_;
  std::array<std::any, detail::num_runtime_items> table_;
};

У него простой интерфейс. Нам достаточно двух методов:

  • Добавить функцию в таблицу, которая соответствует какому-то ключу с набором характеристик.

  • Произвести вызов. Тут нам нужен будет тип результата и тип аргумента

Метод для добавления реализуется следующим образом:

template <typename Func>
void add_func(const dispatch_key_t& key, const Func& func) {
  auto& entry = table_[key.index()];
  …
  entry = std::function(func);
}

Сначала мы находим соответствующую ключу запись в таблице. Потом переданная через аргумент функция упаковывается в std::function, а затем — в std::any с помощью неявного преобразования при сохранении в entry. Отмечу, что используется один шаблонный параметр. И тип функции выводится автоматически компилятором, что достаточно удобно.

Метод для вызова реализуется так:

template <typename Result, typename... Args>
Result call(const dispatch_key_t& key, Args... args) const {
  auto& entry = table_[key.index()];
  if (entry.has_value()) {
    auto func = std::any_cast<std::function<Result(Args...)>>(entry);
    return func(std::forward<Args>(args)...);
  }
  …
}

При вызове нам надо явно указать как тип возвращаемого результата, так и тип аргументов функций в шаблонных параметрах. Далее находим запись в таблице и, воспользовавшись этими типами для правильного задания сигнатуры функции, распаковываем ее из std::any. Затем просто вызываем функцию и передаем туда значение аргументов с использованием std::forward. Если же параметры вывелись неправильно или мы ошиблись в указателе, мы получим std::bad_any_cast. 

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

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

enum class device_t : uint8_t { CPU = 1, GPU, EndOfDevices };
enum class backend_t : uint8_t { CPUNaive = 1, SIMD, EndOfBackends };

Это перечисления, их может быть больше, но сейчас будем рассматривать только два: для идентификации устройства и бэкенда для него. Добавим дополнительное поле EndOf* к перечислению и сделаем его типизированным — тогда сможем использовать это значение для определения количества элементов:

constexpr uint8_t num_devices = static_cast<uint8_t>(device_t::EndOfDevices) - 1;
constexpr uint8_t num_backends = static_cast<uint8_t>(backend_t::EndOfBackends) - 1;

Надо будет просто перемножить и получить максимальное значение размера таблицы, в нашем случае — массива функций:

constexpr size_t num_runtime_items = num_devices * num_backends;
...
std::array<std::any, detail::num_runtime_items> table_;

Кодирование ключа как линейного индекса

Следующий важный элемент работы таблицы функций — представление ключа диспетчеризации dispatch_key_t. Это многомерная координата, каждый компонент которой кодирует одну характеристику вычислительного контекста для функции. В приведенной выше реализации используются два измерения:

  • устройство (device_t) — CPU или GPU,

  • бэкенд (backend_t) — SIMD, Vulkan, CUDA.

Каждая пара значений однозначно определяет точку в двумерном пространстве диспетчеризации. Для доступа в одномерном массиве координаты преобразуются в линейный индекс:

index = y × width + x

Тут y — индекс по измерению «устройство», x — индекс по измерению «бэкенд», width — количество возможных значений бэкенда.

Реализация в коде выходит так:

class dispatch_key_t {
public:
    dispatch_key_t(device_t device, backend_t backend) {
        const auto device_idx = static_cast<uint8_t>(device) - 1;
        const auto backend_idx = static_cast<uint8_t>(backend) - 1;
        
        // Линейная индексация: index = device_idx * num_backends + backend_idx
        index_ = device_idx * detail::num_backends + backend_idx;
    }
    
    size_t index() const { return index_; }
    
private:
    size_t index_{0};
};

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

Как устроен алгоритм диспетчеризации

Процесс диспетчеризации — это конвейер, преобразующий высокоуровневый вызов функции в выполнение конкретной платформенно-зависимой реализации. При вызове функции происходят следующие шаги:

  1. В диспетчере идет поиск нужной функции.

  2. Диспетчер производит анализ аргументов.

  3. Диспетчер формирует ключ диспетчеризации.

  4. Делегирует вызов объекту FunctionTable для вызова конкретной реализации.

Такой многоуровневый подход обеспечивает четкое разделение ответственности: интерфейсная функция остается простой и неизменной, в то время как механизмы выбора реализации помещены в диспетчер. Это позволяет расширять систему новыми бэкендами и устройствами без модификации кода и сохранять при этом предсказуемую и эффективную маршрутизацию вызовов.

Глобальный диспетчер: фасад системы динамической маршрутизации

Класс Dispatcher служит центральным фасадом всей системы диспетчеризации и выполняет следующие обязанности:

  • Хранит таблицы функций — управляет коллекцией FunctionTable, по одной на каждую операцию (mul, add, mean и т.д.).

  • Анализирует контекст вызова — извлекает характеристики из аргументов и глобального состояния системы

  • Маршрутизирует вызовы — формирует ключ и делегирует в соответствующую таблицу

Интерфейс диспетчера минималистичен и интуитивно понятен:

class Dispatcher final {
public:
	// Регистрация реализации для операции
	template<typename Func>
	void add_func(const std::string& name,
              	const dispatch_key_t& key,
              	const Func& func);
   
	// Поиск таблицы функций по имени операции
	FunctionTable& find(const std::string& name);
   
	// Вызов функции с автоматическим анализом контекста
	template<typename Result, typename... Args>
	Result call(FunctionTable& func, Args... args);
   
private:
	ska::flat_hash_map<std::string, FunctionTable> lookup_table_;
};

При регистрации новой реализации через add_func диспетчер автоматически создает таблицу для операции, если она еще не существует. Использование функции find мы уже рассматривали выше в реализации глобальной функции mul.

Хранение таблиц функций организовано через ska::flat_hash_map — высокопроизводительную реализацию хэш-таблицы на основе плоского массива. Такой выбор обусловлен практическими соображениями: даже в крупных математических библиотеках количество операций редко превышает несколько сотен. Это делает плоскую структуру данных более эффективной за счет лучшей локальности кеша и отсутствия указателей.

Самый интересный аспект работы диспетчера — метод call, где происходит преобразование семантики вызова в ключ диспетчеризации:

template<typename Result, typename... Args>
Result Dispatcher::call(FunctionTable& func, Args... args) {
    // 1. Анализ аргументов: извлечение общего устройства
    auto device = get_device(args...);
    
    // 2. Учёт глобального контекста: текущий бэкенд для устройства
    dispatch_key_t key{device, get_current_backend(device)};
    
    // 3. Делегирование в таблицу функций
    return func.call<Result, Args...>(key, std::forward<Args>(args)...);
}

В этом упрощенном примере метод get_device проходится по всем аргументам, которые переданы в функцию, анализирует их тип, берет из них значения и выдает идентификатор общего для всех аргументов устройства. Далее используем функцию get_current_backend для полученного устройства. На ее примере мы видим, как применяются внешние характеристики и параметры системы. Например, для конкретного вычислительного устройства (GPU) может быть глобально установлен бэкенд OpenCL или CUDA. После формируется двухэлементный ключ для выборки таблицы функций. И в этой уже таблице происходит вызов функций, который мы разобрали выше.

Анализ аргументов функции

Рассмотрим функцию, которая принимает на вход два тензора. Конечно, у аргументов, могут быть и другие типы, но суть в том, что мы знаем, какие типы аргументов мы хотим анализировать, и знаем, что у них есть определенные свойства (поля). В данном случае нас будет интересовать поле device. Это идентификатор устройства: на нем расположена память тензора, и тут же будут производиться вычисления. Мы должны проанализировать значения этого свойства для всех переданных аргументов, сравнить их и, если они не совпадают, сгенерировать ошибку или же сформировать ключ для последующего вызова функции.

Для реализации функции get_device используется обобщенный анализ аргументов с помощью паттерна визитор и CRTP (Curiously Recurring Template Pattern):

template<typename... Args>
device_t get_device(const Args&... args) {
    return DeviceExtractor().apply(args...).device;
}

Рассмотрим сначала базовый класс ArgsVisitor от которого наследуется DeviceExtractor — он предоставляет рекурсивный механизм обхода параметров:

template<typename F>
struct ArgsVisitor {
    // Терминальный случай: аргументы закончились
    template<typename... Args>
    F& apply() { return self(); }
    
    // Рекурсивный случай: обработка "головы" и переход к "хвосту"
    template<typename T, typename... Rest>
    F& apply(T&& arg, Rest&&... rest) {
        self()(std::forward<T>(arg));          // обработка текущего аргумента
        return apply(std::forward<Rest>(rest)...); // рекурсия
    }
    
private:
    // CRTP: получение ссылки на конкретный анализатор (наследника)
    F& self() { return *static_cast<F*>(this); }
};

В примере выше видим набор методов apply для итерации по аргументам функций и метод self, который реализуется в контексте паттерна CRTP. 

CRTP — это идиома шаблонного программирования в C++, при которой класс наследуется от шаблонного базового класса, параметризованного самим производным классом. Это позволяет базовому классу получать доступ к интерфейсу производного класса без виртуальных функций. 

С помощью этого подхода можно вызвать конкретную реализацию какого-то анализатора.

Mетод self использ��ется внутри функции apply, чтобы передать значение аргумента для его анализа в конкретной реализации. Механизм работает по принципу «голова-хвост»: 

  • Первый аргумент (голова) передается в operator() наследника для анализа.

  • Оставшиеся аргументы (хвост) обрабатываются рекурсивным вызовом apply. 

  • Рекурсия завершается при достижении терминального случая — вызова apply без параметров.

У метода apply два аргумента: первый — это как бы голова списка, а второй — это список остальных аргументов различных типов. Используя template parameter pack expansion при вызове apply, мы раскладываем аргумент Rest&&... rest на отдельные элементы: голова и хвост. Подход достаточно обыденный для функционального программирования. 

Конкретная реализация: DeviceExtractor

Конкретный анализатор наследуется от ArgsVisitor, указывая себя в качестве параметра шаблона для корректной работы CRTP:

struct DeviceExtractor : ArgsVisitor<DeviceExtractor> {
    device_t device{device_t::CPU}; // значение по умолчанию
    
    // Обработка тензора: извлечение устройства из поля
    void operator()(const Tensor& tensor) {
        if (tensor.defined())
            device = tensor.device();
    }
    
    // Обработка свойств тензора (альтернативный интерфейс)
    void operator()(const TensorProperties& props) {
        device = props.device;
    }
    
    // Игнорирование нерелевантных типов (скаляры, числа и т.д.)
    template<typename T>
    void operator()(const T&) { /* nop */ }
};

Основную функциональность мы реализуем при переопределении оператора функции с аргументами тех типов, которые нас интересуют. У тензора есть свойство device, которое мы можем проанализировать. У объекта типа TensorProperties мы также можем получить доступ к идентификатору устройства. Используя пустое шаблонное переопределение оператора, мы можем проигнорировать неинтересные нам типы аргументов. Такой подход обеспечивает селективный анализ: только аргументы типов Tensor и TensorProperties влияют на результат, остальные автоматически игнорируются. При вызове get_device(x, 3.14, y, "label") будут проанализированы исключительно x и y, а числовые и строковые аргументы пропущены.

DeviceExtractor должен также проверять согласованность устройств всех тензоров. Если один аргумент расположен на CPU, а другой — на GPU, система либо сгенерирует ошибку несовместимости, либо автоматически инициирует перенос данных в соответствии с политикой библиотеки. Так ведет себя и PyTorch, когда выполняет операции над тензорами на разных устройствах.

Регистрация реализаций

Последний элемент — это наполнение объекта Dispatcher конкретными реализациями функций. Один из вариантов реализации может быть основан на атрибутах конструкторов в компиляторах GCC или сlang:

__attribute__((__constructor__(INIT_ORDER))) 
void cpu_init() {
    Dispatcher::instance().add_func(
        "mul", 
        {device_t::CPU, backend_t::SIMD}, 
        cpu::mul
    );
}

__attribute__((__constructor__(INIT_ORDER))) 
void vulkan_init() {
    Dispatcher::instance().add_func(
        "mul", 
        {device_t::GPU, backend_t::Vulkan}, 
        gpu::vulkan::mul
    );
}

Преимущества такого подхода:

  • Модульность: каждый бэкенд регистрирует свои реализации в своих .cpp-файлах без изменения общего кода.

  • Ленивая инициализация: регистрация происходит автоматически при загрузке библиотеки

  • Порядок инициализации: атрибут INIT_ORDER позволяет управлять порядком вызова функций с атрибутом конструктора, то есть порядком загрузки модулей библиотеки.

Однако такой подход использует не стандартные возможности языка С++, а специфичные расширения компилятора.

Преимущества и недостатки динамической диспетчеризации

Исследования команды разработчиков PyTorch показывают реальную цену динамической диспетчеризации:

Сценарий

Накладные расходы

Причина

Легкие операции (малые тензоры, скаляры)

До 10%

Время диспетчеризации сравнимо с временем самой операции

Тяжелые операции (матричные умножения, свёртки)

<0.1%

Время диспетчеризации пренебрежимо мало на фоне вычислений

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

Предложенный подход обладает рядом преимуществ:

  • Расширяемость — новые реализации добавляются без изменения существующего кода.

  • Модульность — каждая реализация живет в своем модуле.

  • Гибкость ключей — легко расширить ключ новыми измерениями (алгоритм, точность).
    Динамическая регистрация — возможность добавлять реализации во время выполнения.
    Интеграция с генераторами кода — автоматическая регистрация сгенерированных ядер,

Однако у решения есть и ограничения:

  • Накладные расходы ~10% на легких операциях.

  • Сложность отладки из-за стирания типов.

  • Требование явного указания типов при вызове.

Что еще изучить по теме: