17 Jun 2025
Сначала я хотел написать класс для 2д случая, но в процессе улучшения кода вдруг осознал, что можно легко обобщить на 3д и 4д и получится простая, эффективная и универсальная библиотека.
Дизайн вдохновлён numpy и некоторыми возможностями третьей скалы.
Ниже расскажу, почему я выбрал те или иные решения.
Ключевая идея
Итак, сразу к делу, двухмерный массив будет выглядеть примерно так:
trait ArrayView2d[T]:
val data: Array[T]
def shape0: Int
def shape1: Int
def offset: Int
def stride0: Int
def stride1: Int
def getIndex(i0: Int, i1: Int): Int =
offset + i0 * stride0 + i1 * stride1
inline def apply(i0: Int, i1: Int): T =
data(getIndex(i0, i1))
inline def update(i0: Int, i1: Int, value: T): Unit =
data(getIndex(i0, i1)) = value
Т.е., я храню размеры по осям, положение “нулевого” элемента относительно начала массива и то, насколько меняется позиция в массиве при увеличении на того или иного индекса.
Эту идею я подсмотрел у numpy и это очень интересный способ делать самые разные “репрезентации” для одних и тех же данных.
- Взять часть массива: поправить
offset
,shape0
иshape1
- Транспонировать массив: поменять местами
stride0
иstride1
, а так жеshape0
иshape1
- Пропустить элементы с нечётными индексами - умножить
stride
на 2, поделить размер на 2. - Развернуть порядок на обратный - умножить
stride
на -1, поправитьoffset
. - Добавить третью ось с единичным размером - создать
view
для 3д, для третьей оси указатьshape2 = 1
- Сделать броадкаст оси с единичным размером до любого: занулить
stride
, поставить любойshape
.
С операцией reshape есть нюасны, но в некоторых случаях можно переиспользовать один и тот же массив данных во view с разным количеством измерений. Например, (48), (3, 16) и (3, 4, 4)
Т.е., на эту модель прекрасно ложится куча операций по манипулированию многомерными данными и при этом сами данные копировать или перекладывать почти не надо.
Почему это эффективно работает.
Если мы сделаем ArrayView[Int]
в скале, то в отличие от Kotlin и Java, компилятор скалы в data положит нормальный массив интов, а не массив объектов Integer. Этот замечательный факт сподвиг меня на написание библиотеки - можно использовать ArrayView2d[Int]
, ArrayView2d[Double]
, ArrayView2d[SomeType]
и всё это будет максимально эффективно храниться.
Мне очень хотелось сделать временные объекты для двухмерных индексов и вместо хардкода shape0 и shape1 хранить просто shape = Shape2(i0, i1)
, но в JVM всё ещё нет value типов и создание временных объектов на каждый доступ к массиву может плохо сказаться на производительности.
И ещё момент почему это плохо - у типичного процессора размер кеш-линии в 64 байта. Размер наших полей - указатель data на 4 или 8 байт и пять интов по 4 байта - в сумме 24 или 28 байт. Т.е., такой объект прекрасно помещается в одну линию кеша процессора, и доступ к полям объекта будет максимально быстрым. Если вдруг вместо этого я сделаю отдельный класс Shape2, то просто для получения размера надо будет перейти по ссылке shape и достать оттуда число - это определённо менее эффективно, чем просто вытащить поле-int.
Почему я использую одномерный массив data, а не массив массивов или что-то ещё - опять же, доступ к элементу внутри массива массивов потребует два прыжка в памяти - это медленнее. А ещё это будет много объектов в памяти вместо одного.
Финальный штрих - методы apply и update, которые используются для доступа к элементам, отмечены как инлайн.
Подвох в том, что если есть ArrayView[T]
и мы вызываем метод с возвращаемым типом T
, то этот тип будет объектом. И мы получим боксинг примитивных типов на ровном месте.
И тут нас спасает модификтора inline - для ArrayView[Double]
компилятор знает, что в data именно массив doulbe и в байткоде будет простой доступ к массиву типа такого:
val element: Doulbe = arrayView2d.data(arrayView.getPos(i0, i1))
Возможно, в каких-то случаях сработает escape analysis и инлайн на уровне JVM, а ещё java положит объект Shape рядом с ArrayView и оно будет работать быстро - но я решил вместо этого использовать максимально простой и близкий к железу дизайн.
Я замерил производительность в JMH - она близка к скорости кода, в котором вручную подставлены индексы. И быстрее, чем массив массивов. Считаю, что это успех!
transparent inline и context recievers
Если вы помните, в numpy метод arr[...]
очень гибкий и может принимать что угодно. Например arr[0, 1:-1, ::-1]
- здесь мы берём фиксированный первый индекс, а вторую и третью оси осталяем осями, просто по второй оси откидываем крайние значения, а по третьей оси меняем порядок на обратый. Получается, что размерность результата зависит от типов аргументов - это может быть либо число, либо интервал.
Причём есть хитрый момент - отрицательный индекс значит смещение “от конца” массива. Но если вы вдруг ошибётесь и в коде типа arr[i:j]
случайно окажутся отрицательные i
и j
- удачи в отладке.
В Scala можно сделать ещё круче c помощью transparent inline
:
transparent inline def view [T1 <: Int | Range,
T2 <: Int | Range](inline range0: AxisSize.Size ?=> T1,
inline range1: AxisSize.Size ?=> T2) = {
val t0 = AxisSize.withAxisSize(shape0, range0)
val t1 = AxisSize.withAxisSize(shape1, range1)
val start0 = ArrayViewUtil.getFirst(t0, shape0)
val start1 = ArrayViewUtil.getFirst(t1, shape1)
val offset = getIndex(start0, start1)
inline (t0, t1) match {
case (a: Int, b: Int) => ArrayView0dImpl(data, offset = offset)
case (a: Range, b: Int) => ArrayView1dImpl(data, shape0 = a.size, offset = offset, stride0 = stride0 * a.step)
case (a: Int, b: Range) => ArrayView1dImpl(data, shape0 = b.size, offset = offset, stride0 = stride1 * b.step)
case (a: Range, b: Range) => ArrayView2dImpl(data, shape0 = a.size, shape1 = b.size, offset = offset, stride0 = stride0 * a.step, stride1 = stride1 * b.step)
}
}
Так вот, каждый аргумент ещё имеет контекст, в котором передан size с размером вдоль оси.
Теперь можно написать так:
arrayView2d = arrayView3d.view(0, 1 until (size - 1), all.reversed)
Вообще мой подход не идеален, но это самое удобное из того что я смог придумать.
Для 3д и 4д всё аналогично, если не смотреть на комбинаторный взрыв до 16 вариантов для 4д. Но это будет во время компиляции, в байткоде будет только нужный код.
Итоги
Я не стал в библиотеку добавлять матричные операции типа арифметических, потому что это спецефический сценарий для массивов чисел.
Вместо этого я сфокусировался на том, чтобы получилась простая и эффективная библиотека для работы с многомерными данными любых типов.
Библиотека быстрая, простая, без зависимостей и написана на чистой скале (поддерживает scala js). Используйте на здоровье