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 и это очень интересный способ делать самые разные “репрезентации” для одних и тех же данных.

  1. Взять часть массива: поправить offset, shape0 и shape1
  2. Транспонировать массив: поменять местами stride0 и stride1, а так же shape0 и shape1
  3. Пропустить элементы с нечётными индексами - умножить stride на 2, поделить размер на 2.
  4. Развернуть порядок на обратный - умножить stride на -1, поправить offset.
  5. Добавить третью ось с единичным размером - создать view для 3д, для третьей оси указать shape2 = 1
  6. Сделать броадкаст оси с единичным размером до любого: занулить 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). Используйте на здоровье