Формы и расположение

Структура XLA Op

Рассмотрим пример HLO:

add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)}
          add(exponential.183, broadcast.3115)

Он состоит из следующих компонентов:

  • Имя операции: add.936
    • Это уникальное название операции.
  • Форма: bf16[8,1,1280,16384]
    • Это выходная форма Op. Здесь dtype — bf316 , а форма — [8,1,1280,16384] .
  • Макет (с тайлингом): 3,2,0,1:T(8,128)(2,1)
    • Это описывает, как массив хранится в памяти. 3,2,0,1 обозначает порядок осей в памяти (т. е. по столбцам, по строкам и т. д.), а T(8,128)(2,1) обозначает используемые разбиение и заполнение.
    • Макет необязателен. Если он не указан, разбиение на плитки не производится, а размеры считаются упорядоченными от самых крупных к самым мелким.
  • Операция: add
    • Выполняемая операция. В данном случае это Add , что также упоминается в имени операции.
  • Аргументы: exponential.183 , broadcast.3115
    • Эта операция принимает два аргумента, указанные с их уникальными именами.

Давайте рассмотрим другой пример — слияние операций:

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)}
            fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32),
            kind=kCustom, calls=%all-reduce-scatter.3

Помимо ранее описанных компонентов, сюда входят:

  • Атрибуты: kind и calls
    • Они предоставляют дополнительную информацию о выполняемой операции, в данном случае: слиянии.
  • Ячейка памяти (идентификатор области памяти): S(1)
    • Это обозначает область памяти/местоположение, где хранится массив. S(1) здесь обозначает, что этот массив находится в VMEM (на TPU).
  • Подробности формы и макета для входного аргумента %fusion.32

В следующих разделах описываются фигуры, макет и идентификаторы пространства памяти . Подробнее о тайлинге можно узнать в разделе «Тайлинг» .

Формы

Прототип XLA ShapeProto ( xla_data.proto ) описывает количество измерений, размер и тип данных N-мерного массива (кратко — массива ).

Терминология, обозначения и условные обозначения

  • Истинное число измерений массива — это число измерений, размер которых больше 1.

  • Измерения нумеруются от 0 до N-1 для N -мерного массива. Размер измерения — неотрицательное целое число. В частности, допустим размер 0. Номера измерений заданы произвольно для удобства. Порядок этих номеров измерений не подразумевает определённого порядка (второстепенного/главного) в макете фигуры. Макет определяется прототипом LayoutProto .

  • По соглашению, измерения перечислены в порядке возрастания их номеров. Например, для трёхмерного массива размером [A x B x C] измерение 0 имеет размер A , измерение 1 — размер B , а измерение 2 — размер C

    Некоторые утилиты в XLA также поддерживают отрицательную индексацию, как в Python: измерение -1 — это последнее измерение (эквивалентно N-1 для N -мерного массива). Например, для трёхмерного массива, описанного выше, измерение -1 имеет размер C , измерение -2 — размер B и так далее.

  • Двух-, трёх- и четырёхмерные массивы часто имеют определённые буквы, соответствующие размерности. Например, для двумерного массива:

    • измерение 0: y
    • измерение 1: x

    Для трехмерного массива:

    • измерение 0: z
    • измерение 1: y
    • измерение 2: x

    Для 4D массива:

    • измерение 0: p
    • измерение 1: z
    • измерение 2: y
    • измерение 3: x

  • Функции в API XLA, принимающие измерения, делают это в порядке возрастания их номеров. Это соответствует порядку, используемому при передаче измерений в качестве initializer_list , например:

    ShapeUtil::MakeShape(F32, {A, B, C, D})

    создаст фигуру, массив размеров которой состоит из последовательности [A, B, C, D] .

Макет

Прототип LayoutProto описывает представление массива в памяти. Он включает следующие поля:

message LayoutProto {
  repeated int64 minor_to_major;
  int64 tail_padding_alignment_in_elements;
  ...
}

Упорядочивание измерений от малого к большему

Единственное обязательное поле — minor_to_major . Это поле описывает порядок измерений внутри фигуры от младшего к старшему. Значения в minor_to_major задают порядок измерений массива ( 0 до N-1 для N -мерного массива), где первое значение соответствует младшему измерению, а последнее — старшему. Второстепенное измерение — это измерение, которое изменяется наиболее быстро при пошаговом перемещении по элементам массива, расположенного в линейной памяти.

Например, рассмотрим следующий двумерный массив размером [2 x 3] :

a b c
d e f

Здесь измерение 0 имеет размер 2, а измерение 1 — размер 3. Если поле minor_to_major в макете равно [0, 1] , то измерение 0 является самым младшим измерением, а измерение 1 — самым старшим. Это соответствует следующему макету в линейной памяти:

a d b e c f

Этот порядок измерений от второстепенного к главному, от 0 до N-1 аналогичен порядку измерений по столбцам (для двумерных пространств). Если предположить монотонный порядок измерений, то в коде эту схему можно обозначить просто как «dim 0 — второстепенный».

С другой стороны, если поле minor_to_major в макете равно [1, 0] , то макет в линейной памяти будет следующим:

a b c d e f

Порядок измерений от младшего к старшему, от N-1 вплоть до 0 , для N -мерного массива аналогичен порядку измерений по строкам (для двумерных массивов). Если предположить монотонное упорядочение измерений, то в коде эту схему можно обозначить просто как «dim 0 is major».

Порядок по умолчанию от второстепенного к главному

Макет по умолчанию для вновь созданных фигур — «порядок измерений — от большего к меньшему» (т.е. [N-1, ..., 0] ).

Прокладка

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

Индексация в массивы

Класс IndexUtil в index_util.h предоставляет утилиты для преобразования многомерных индексов в линейные и наоборот с заданной формой и макетом. Многомерные индексы включают индекс int64 для каждого измерения. Линейные индексы представляют собой одно значение int64 , которое индексирует буфер, содержащий массив. См. shape_util.h и layout_util.h в том же каталоге, чтобы узнать об утилитах, упрощающих создание и обработку форм и макетов.

Идентификаторы пространства памяти

В HLO каждый массив может быть аннотирован идентификатором пространства памяти, записываемым как S(n).

  • S(0) (часто опускается) обозначает высокоскоростную память устройства (HBM).
  • S(1) представляет собой виртуальную память устройства (VMEM).
  • S(2) , S(3) и т. д. соответствуют дополнительным пространствам памяти, специфичным для устройства.
  • S(5) обозначает память хоста.