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

Структура операции XLA

Рассмотрим пример 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]
    • Это форма выходных данных операции. Здесь тип данных — bf16 , а форма — [8,1,1280,16384] .
  • Схема (с мозаичным расположением): 3,2,0,1:T(8,128)(2,1)
    • Это описывает способ хранения массива в памяти. 3,2,0,1 обозначает порядок расположения элементов в памяти (например, по столбцам, по строкам и т. д.), а T(8,128)(2,1) обозначает используемый порядок расположения элементов и отступов.
    • Расположение элементов необязательно. Если не указано, размещение элементов в виде плиток отсутствует, и предполагается, что размеры упорядочены от наиболее крупных к наиболее мелким.
  • Операция: 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

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

Формы

Прототип 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 -мерного массива). Например, для описанного выше 3-мерного массива измерение -1 имеет размер C , измерение -2 — размер B и так далее.

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

    • размерность 0: y
    • размер 1: x

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

    • размерность 0: z
    • размерность 1: y
    • размерность 2: x

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

    • размерность 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 , аналогичен столбцовому порядку (для двумерных объектов). Предполагая монотонный порядок размерностей, в коде можно просто обозначить такую ​​структуру как "размерность 0 — малая".

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

a b c d e f

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

Порядок от минорной к мажорной дисциплине по умолчанию

По умолчанию для вновь создаваемых фигур используется схема расположения «размеры от основных к второстепенным» (т.е. [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) обозначает память хоста.