Привет!

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



Итак, что же нам нужно?

1. N-мерные матрицы (тензоры)
2. Имплементация базового набора методов работы с тензором как со структурой данных
3. Имплементация базового набора математических функций (для матриц и векторов)
4. Типы Generic, то есть любые. И кастомные операторы

Хранение элементов, транспонирование, подтензор

Элементы будут храниться в одномерном массиве. Чтобы из набора индексов получить элемент, мы будем умножать индексы на определенные коэффициенты и складывать. А именно, положим, у нас есть тензор [3 x 4 x 5]. Тогда нам нужно сформировать массив из трех элементов — блоков (сам название придумал). Тогда последний элемент равен 1, предпоследний равен 5, а первый элемент — 5 * 4 = 20. То есть blocks = [20, 5, 1]. Например, при обращении по индексу [1, 0, 4] индекс в массиве будет выглядеть как 20 * 1 + 5 * 0 + 4 * 1 = 24. Пока все ясно

Транспонирование

… то есть изменение порядка осей. Тупой и простой способ — создать новый массив и загнать туда элементы в новом порядке. Но часто бывает удобно транспонировать, работать с нужным порядком осей, а затем менять порядок осей обратно. Конечно, в этом случае нельзя менять сам линейный массив (ЛМ), а при обращении к определенным индексам, мы будем просто подменять порядок.

Рассмотрим функцию:

private int GetFlattenedIndexSilent(int[] indices)
{
    var res = 0;
    for (int i = 0; i < indices.Length; i++)
        res += indices[i] * blocks[i];
    return res + LinOffset;
}

Как видим, если поменять blocks местами, то создастся видимость свапнутых осей. Поэтому так и запишем:

public void Transpose(int axis1, int axis2)
{
    (blocks[axis1], blocks[axis2]) = (blocks[axis2], blocks[axis1]);
    (Shape[axis1], Shape[axis2]) = (Shape[axis2], Shape[axis1]);
}

Просто меняем номера и длины осей местами.

Подтензор

Подтензор N-мерного тензора это M-мерный тензор (M < N), который является частью исходного. Например, нулевой элемент тензора формы [2 x 3 x 4] это тензор формы [3 x 4]. Его мы получим просто сдвигом.

Представим, что мы получаем подтензор по индексу n. Тогда его первый элемент это n * blocks[0] + 0 * blocks[1] + 0 * blocks[2] + .... То есть сдвиг равен n * blocks[0]. Соответственно, не копируя исходный тензор, мы запоминаем сдвиг, создаем новый тензор со ссылкой на наши данные, но уже со сдвигом. А еще нужно будет выкинуть элемент из blocks, а именно элемент blocks[0], ведь это первая ось, к ней обращений не будет.

Другие операции с композицией

Все остальные уже следуют из этих.

1. SetSubtensor форвардит элементы в нужный подтензор
2. Concat создает новый тензор, и туда форвардит элементы из двух (при этом длина первой оси — сумма длин осей двух тензоров)
3. Stack группирует некоторое число тензоров в один с дополнительной осью (например, stack([3 x 4], [3 x 4]) -> [2 x 3 x 4])
4. Slice возвращает Stack от определенных подтензоров

Все операции с композицией, которые я определил, можно найти здесь.

Математические операции

Тут уже все просто.

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

2) Операции над матрицами. Произведение, инверсия, и другие простые операции, мне кажется, объяснения не требуют. Хотя есть что рассказать о детерминанте.


3) Операции над веторами (dot и cross product).

Оптимизация

Темплейты?

В C# нет темплейтов. Поэтому приходится использовать костыли. Некоторые люди придумали динамическую компиляцию в делегат, например так у него реализована сумма двух типов.

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

var c = default(TWrapper).Addition(a, b);

Что инлайнится до вашего метода. Пример имплементации такой структуры.

Индексация

Далее, хотя кажется, что логично в индексаторе использовать params, то есть как-то так:

public T this[params int[] indices]

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

Исключения

Еще я загнал все исключения и проверки в блок #if ALLOW_EXCEPTIONS на случай, если точно нужно быстро, а проблем с индексами точно нет. Небольшой прирост по производительности есть.

На самом деле это не просто так микрооптимизация, которая много чего будет стоить в плане безопасности. Например, в ваш тензор идет запрос, в котором вы и так уже по своим соображениям сделали проверку на корректность данных. Тогда зачем вам еще одна проверка? А они не бесплатные, особенно, если мы экономим даже лишние арифметические операции с целыми числами.

Многопоточность

Спасибо Билли, это оказалось очень просто и реализуется через Parallel.For.

Хотя многопоточность это не панацея, и включать ее надо аккуратно. Я провел бенчмарк для поточечного сложения тензоров на i7-7700HQ:



Где Y-ось показывает время (в микросекундах) на выполнение одной операции с двумя целочисленными тензорами определенного размера (размер по оси X).

То есть есть определенный порог, начиная с которого многопоток имеет смысл. Чтобы не надо было думать, сделал флажок Threading.Auto и тупо захардкодил константы, начиная с объема равного которым можно включать многопоток (есть более умный автоматический метод?).

При этом библиотека все равно никак не быстрее игровых матриц, или тем более тех, что подсчитываются на CUDA. А зачем? Те уже написаны, а у нас главное — кастомный тип.

Вот так вот

Вот такая коротенькая заметка, спасибо за прочтение. Ссылка на гитхаб проекта здесь. А главный его пользователь — библиотека символической алгебры AngouriMath.