Стабильная спецификацияHLO

StableHLO — это набор операций для высокоуровневых операций (HLO) в моделях машинного обучения (ML). StableHLO служит связующим звеном между различными фреймворками и компиляторами машинного обучения: фреймворки машинного обучения, создающие программы StableHLO, совместимы с компиляторами машинного обучения, использующими программы StableHLO.

Наша цель — упростить и ускорить разработку машинного обучения, обеспечив большую совместимость между различными фреймворками машинного обучения (такими как TensorFlow, JAX и PyTorch) и компиляторами машинного обучения (такими как XLA и IREE). С этой целью данный документ представляет собой спецификацию языка программирования StableHLO.

Данная спецификация состоит из трёх основных разделов. В разделе «Программы» описывается структура программ StableHLO, состоящих из функций StableHLO, которые, в свою очередь, состоят из операций StableHLO. В рамках этой структуры раздел «Операции» определяет семантику отдельных операций. В разделе «Выполнение» описывается семантика всех этих операций, выполняемых совместно в рамках программы. Наконец, в разделе «Нотация» обсуждается нотация, используемая в спецификации.

Чтобы просмотреть спецификацию предыдущей версии StableHLO, откройте репозиторий интересующей вас версии с тегом . Например, StableHLO v0.19.0 Spec . Чтобы просмотреть изменения, произошедшие при каждом выпуске минорной версии StableHLO, обратитесь к журналу версий в VhloDialect.td .

Программы

Program ::= {Func}

Программы StableHLO состоят из произвольного количества функций StableHLO. Ниже приведён пример программы с функцией @main , которая имеет 3 входных данных ( %image , %weights и %bias ) и 1 выходной. Тело функции содержит 6 операций.

func.func @main(
  %image: tensor<28x28xf32>,
  %weights: tensor<784x10xf32>,
  %bias: tensor<1x10xf32>
) -> tensor<1x10xf32> {
  %0 = "stablehlo.reshape"(%image) : (tensor<28x28xf32>) -> tensor<1x784xf32>
  %1 = "stablehlo.dot"(%0, %weights) : (tensor<1x784xf32>, tensor<784x10xf32>) -> tensor<1x10xf32>
  %2 = "stablehlo.add"(%1, %bias) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  %3 = "stablehlo.constant"() {value = dense<0.0> : tensor<1x10xf32>} : () -> tensor<1x10xf32>
  %4 = "stablehlo.maximum"(%2, %3) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  "func.return"(%4): (tensor<1x10xf32>) -> ()
}

Функции

Func        ::= 'func' '.' 'func' FuncId FuncInputs FuncOutputs '{' FuncBody '}'
FuncInputs  ::= '(' [FuncInput {',' FuncInput}] `)`
FuncInput   ::= ValueId ':' ValueType
FuncOutputs ::= ['->' FuncOutput, {',' FuncOutput}]
FuncOutput  ::= ValueType
FuncBody    ::= {Op}

Функции StableHLO (также называемые именованными функциями ) имеют идентификатор, входы/выходы и тело. В будущем мы планируем ввести дополнительные метаданные для функций для достижения лучшей совместимости с HLO ( #425 , #626 , #740 , #744 ).

Идентификаторы

FuncId  ::= '@' letter {letter | digit}
ValueId ::= '%' digit {digit}
          | '%' letter {letter | digit}
letter  ::= 'a' | ... | 'z' | 'A' | ... | 'Z' | '_'
digit   ::= '0' | ... | '9'

Идентификаторы StableHLO похожи на идентификаторы во многих языках программирования, но имеют две особенности: 1) все идентификаторы имеют сигилы, которые различают различные виды идентификаторов, 2) идентификаторы значений могут быть полностью числовыми, что упрощает генерацию программ StableHLO.

Типы

Type         ::= ValueType | NonValueType
ValueType    ::= TensorType | QuantizedTensorType | TokenType | TupleType | BufferType
NonValueType ::= TensorElementType | QuantizedTensorElementType | FunctionType | StringType

Типы StableHLO подразделяются на типы значений (также называемые типами первого класса ), которые представляют значения StableHLO, и типы, не являющиеся значениями , которые описывают другие элементы программы. Типы StableHLO похожи на типы во многих языках программирования, но их главная особенность заключается в предметно-ориентированной природе StableHLO, что приводит к некоторым необычным результатам (например, скалярные типы не являются типами значений).

TensorType ::= 'tensor' '<' Shape TensorElementType '>'
Shape ::= {DimensionSize 'x'}
DimensionSize ::= digit {digit} | '?'

Типы тензоров представляют собой тензоры, то есть многомерные массивы. Они имеют форму и тип элемента , где форма представляет собой неотрицательные или неизвестные размеры измерений в порядке возрастания соответствующих измерений (которые также называются осями ), пронумерованных от 0 до R-1 . Количество измерений R называется рангом . Например, tensor<2x3xf32> — это тип тензора с формой 2x3 и типом элемента f32 . Он имеет два измерения (или, другими словами, две оси) — 0-е измерение и 1-е измерение, — размеры которых равны 2 и 3. Его ранг равен 2.

Формы могут быть частично или полностью неизвестны (динамическими), например, tensor<?x2xf64> частично неизвестен, а tensor<?x?xf64> полностью неизвестен. Динамические размеры обозначаются с помощью символа ? . Формы не могут быть неранжированы.

В будущем мы планируем изучить расширение типов тензоров за пределы размеров измерений и типов элементов, например, чтобы включить макеты ( #629 ) и разреженность ( #1078 ).

QuantizedTensorType ::= 'tensor' '<' Shape QuantizedTensorElementType '>'
QuantizedTensorElementType ::= '!quant.uniform' '<'
                  QuantizationStorageType
                  ['<' QuantizationStorageMin ':' QuantizationStorageMax '>']
                  ':' QuantizationExpressedType
                  [':' QuantizationDimension]
                  ',' QuantizationParameters '>'
QuantizationStorageType ::= IntegerType
QuantizationStorageMin ::= IntegerLiteral
QuantizationStorageMax ::= IntegerLiteral
QuantizationExpressedType ::= FloatType
QuantizationDimension ::= IntegerLiteral
QuantizationParameters ::= QuantizationParameter
                         | '{' QuantizationParameter {',' QuantizationParameter} '}'
QuantizationParameter ::= QuantizationScale [':' QuantizationZeroPoint]
QuantizationScale ::= FloatLiteral
QuantizationZeroPoint ::= IntegerLiteral

Типы квантованных элементов представляют собой целочисленные значения типа хранения в диапазоне от storage_min до storage_max (включительно), соответствующие значениям с плавающей точкой заданного типа . Для заданного целого значения i соответствующее значение с плавающей точкой f может быть вычислено как f = (i - zero_point) * scale , где scale и zero_point называются параметрами квантования . Параметры storage_min и storage_max необязательны в грамматике, но имеют значения по умолчанию min_value(storage_type) и max_value(storage_type) соответственно. Типы квантованных элементов имеют следующие ограничения:

• (C1) type(storage_min) = storage_type .
• (C2) type(storage_max) = storage_type .
• (C3) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type) .
• (C4) type(scales...) = expressed_type .
• (C5) 0 < scales .
• (C6) is_finite(scales...) .
• (C7) storage_min <= zero_points <= storage_max .
• (C8) type(zero_points...) = storage_type .
• (C9) size(scales) = size(zero_points) .
• (C10) Если is_empty(quantization_dimension) , то size(scales) = 1 .
• (C11) 0 <= quantization_dimension .

В настоящее время QuantizationScale — это константа с плавающей точкой, но существует большой интерес к целочисленным шкалам, представленным с помощью множителей и сдвигов. Мы планируем изучить это в ближайшем будущем ( #1404 ).

Продолжается обсуждение семантики QuantizationZeroPoint , включая тип, значения и возможность существования только одной или потенциально нескольких нулевых точек в квантованном тензорном типе. Исходя из результатов этого обсуждения, спецификация нулевых точек может измениться в будущем ( #1405 ).

Еще одно продолжающееся обсуждение касается семантики QuantizationStorageMin и QuantizationStorageMax для определения того, следует ли накладывать какие-либо ограничения на эти значения и на значения квантованных тензоров ( #1406 ).

Наконец, мы планируем изучить представление неизвестных масштабов и нулевых точек, аналогично тому, как мы планируем изучить представление неизвестных размеров измерений ( #1407 ).

Квантованные тензорные типы представляют собой тензоры с квантованными элементами. Эти тензоры в точности совпадают с обычными тензорами, за исключением того, что их элементы имеют квантованные типы элементов вместо обычных.

В квантованных тензорах квантование может быть потензорным (per-tensor) , то есть иметь одну scale и zero_point для всего тензора, или поосевым (per-axis) , то есть иметь несколько scales и zero_points , по одной паре на срез конкретного измерения quantization_dimension . Более формально, в тензоре t с квантованием по осям существует dim(t, quantization_dimension) срезов quantization_dimension : t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] т. д. Все элементы в i -м срезе используют scales[i] и zero_points[i] в ​​качестве своих параметров квантования. Типы квантованных тензоров имеют следующие ограничения:

• Для потензорного квантования:
  • Никаких дополнительных ограничений.
• Для квантования по осям:
  • (C12) quantization_dimension < rank(self) .
  • (C13) dim(self, quantization_dimension) = size(scales) .

TokenType ::= 'token'

Типы токенов представляют собой токены, то есть непрозрачные значения, создаваемые и потребляемые некоторыми операциями. Токены используются для задания порядка выполнения операций, как описано в разделе «Выполнение» .

TupleType ::= 'tuple' '<' TupleElementTypes '>'
TupleElementTypes ::= [ValueType {',' ValueType}]

Типы буферов представляют собой буферы. Например, в XLA буферы представляют собой многомерные массивы с согласованным хранилищем. Подобно тензорным типам , типы буферов имеют форму и тип элемента , где форма представляет неотрицательные или неизвестные размеры измерений в порядке возрастания соответствующих измерений (которые также называются осями ), пронумерованных от 0 до R-1 . Количество измерений R называется рангом . Например, memref<2x3xf32> — это тип буфера с формой 2x3 и типом элемента f32 . Он имеет два измерения (или, другими словами, две оси) — 0-е измерение и 1-е измерение, — размеры которых равны 2 и 3. Его ранг равен 2.

Буферы можно выделять с помощью custom_call функции CreateBuffer или Pin и освобождать с помощью custom_call функции Unpin . Только операторы custom_call могут читать и записывать содержимое буферов. Подробнее см. в описании custom_call .

Типы кортежей представляют кортежи, то есть разнородные списки. Кортежи — это устаревшая функция, существующая только для совместимости с HLO. В HLO кортежи используются для представления вариативных входных и выходных данных. В StableHLO вариативные входные и выходные данные поддерживаются изначально, и единственное применение кортежей в StableHLO — это полное представление HLO ABI, где, например, T , tuple<T> и tuple<tuple<T>> могут существенно различаться в зависимости от конкретной реализации. В будущем мы планируем внести изменения в HLO ABI, которые, возможно, позволят нам удалить типы кортежей из StableHLO ( #598 ).

TensorElementType ::= BooleanType | IntegerType | FloatType | ComplexType
BooleanType ::= 'i1'
IntegerType ::= SignedIntegerType | UnsignedIntegerType
SignedIntegerType ::= 'si2' | 'si4' | 'si8' | 'si16' | 'si32' | 'si64'
UnsignedIntegerType ::= 'ui2' | 'ui4' | 'ui8' | 'ui16' | 'ui32' | 'ui64'
FloatType ::= 'f4E2M1FN' | 'f6E2M3FN' | 'f6E3M2FN' | 'f8E3M4' | 'f8E4M3'
            | 'f8E4M3FN' | 'f8E4M3FNUZ' | 'f8E4M3B11FNUZ' | 'f8E5M2'
            | 'f8E5M2FNUZ' | 'f8E8M0FNU' | 'bf16' | 'f16' | 'f32' | 'f64'
TensorFloat32 ::= 'tf32'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

Типы элементов представляют элементы тензорных типов. В отличие от многих языков программирования, в StableHLO эти типы не являются типами первого класса. Это означает, что программы на StableHLO не могут напрямую представлять значения этих типов (в результате скалярные значения типа T идиоматично представлять с помощью 0-мерных тензорных значений типа tensor<T> ).

• Тип Boolean представляет логические значения true и false .
• Целочисленные типы могут быть знаковыми ( si ) или беззнаковыми ( ui ) и иметь одну из поддерживаемых разрядностей ( 2 , 4 , 8 , 16 , 32 или 64 ). Знаковые типы siN представляют целые значения от -2^(N-1) до 2^(N-1)-1 включительно, а беззнаковые типы uiN представляют целые значения от 0 до 2^N-1 включительно.
• Типы чисел с плавающей точкой могут быть одним из следующих:
  • f8E3M4 , f8E4M3 и f8E5M2 8-битные числа с плавающей запятой, соответствующие соглашениям IEEE-754.
  • Типы f8E4M3FN и f8E5M2 соответствуют кодировкам E4M3 и E5M2 формата FP8, описанного в документе Форматы FP8 для глубокого обучения .
  • Типы f8E4M3FNUZ и f8E5M2FNUZ , соответствующие кодировкам E4M3 и E5M2 форматов FP8, описанных в 8-битных числовых форматах для глубоких нейронных сетей .
  • Тип f8E4M3B11FNUZ , соответствующий кодировке E4M3 форматов FP8, описанных в разделе «Обучение и вывод гибридных чисел с 8-битной плавающей точкой (HFP8) для глубоких нейронных сетей» .
  • Тип bf16 , соответствующий формату bfloat16 , описанному в BFloat16: секрет высокой производительности на облачных TPU .
  • Типы f16 , f32 и f64 соответствуют форматам binary16 («половинная точность»), binary32 («одинарная точность») и binary64 («двойная точность»), описанным в стандарте IEEE 754 .
  • Тип tf32 соответствует формату TensorFloat32 и имеет ограниченную поддержку в StableHLO.
  • Типы MX (микромасштабирование) f4E2M1FN , f6E2M3FN , f6E3M2FN и f8E8M0FNU описаны в спецификации форматов микромасштабирования OCP .
• Комплексные типы представляют собой комплексные значения, имеющие действительную и мнимую части одного и того же типа элемента . Поддерживаемые комплексные типы: complex<f32> (обе части относятся к типу f32 ) и complex<f64> (обе части относятся к типу f64 ).

FunctionType ::= '(' InputTypes ')' '->' '(' OutputTypes ')'
InputTypes ::= [ValueType {',' ValueType}]
OutputTypes ::= [ValueType {',' ValueType}]

Типы функций представляют собой как именованные, так и анонимные функции. Они имеют входные типы (список типов слева от -> ) и выходные типы (список типов справа от -> ). Во многих языках программирования типы функций являются типами первого класса, но не в StableHLO.

StringType ::= 'string'

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

Операции

Операции StableHLO (также называемые ops ) представляют собой замкнутый набор высокоуровневых операций в моделях машинного обучения. Как уже говорилось, синтаксис StableHLO во многом вдохновлён MLIR, который не обязательно является самой эргономичной альтернативой, но, возможно, наилучшим образом соответствует цели StableHLO — повышению уровня взаимодействия между фреймворками машинного обучения и компиляторами машинного обучения.

Op            ::= [OpOutputs] OpName OpInputs ':' OpSignature
OpName        ::= '"' 'stablehlo' '.' OpMnemonic '"'
OpMnemonic    ::= 'abs' | 'add' | ...

Операции StableHLO (также называемые ops ) имеют имя, входы/выходы и сигнатуру. Имя состоит из префикса stablehlo. и мнемонического кода , который однозначно идентифицирует одну из поддерживаемых операций. Полный список всех поддерживаемых операций см. ниже.

OpInputs        ::= OpInputValues OpInputFuncs OpInputAttrs
OpInputValues   ::= '(' [OpInputValue {',' OpInputValue}] ')'
OpInputValue    ::= ValueId
OpInputFuncs    ::= ['(' OpInputFunc {',' OpInputFunc} ')']
OpInputAttrs    ::= ['{' OpInputAttr {',' OpInputAttr} '}']
OpOutputs       ::= [OpOutput {',' OpOutput} '=']
OpOutput        ::= ValueId

Операции потребляют входные данные и производят выходные данные . Входные данные подразделяются на входные значения (вычисляемые во время выполнения), входные функции (предоставляемые статически, поскольку в StableHLO функции не являются значениями первого класса) и входные атрибуты (также предоставляемые статически). Вид входных и выходных данных, потребляемых и создаваемых операцией, зависит от её мнемоники. Например, операция add потребляет 2 входных значения и производит 1 выходное значение. Для сравнения, операция select_and_scatter потребляет 3 входных значения, 2 входные функции и 3 входных атрибута.

OpInputFunc ::= '{' Unused FuncInputs ':' FuncBody '}'
Unused      ::= '^' digit {digit}
              | '^' letter {letter | digit}

Функции ввода (которые также называются анонимными функциями ) очень похожи на именованные функции, за исключением того, что: 1) у них нет идентификатора (отсюда и название «анонимные»), 2) они не объявляют типы вывода (типы вывода выводятся из return оператора внутри функции).

Синтаксис функций ввода включает в себя в настоящее время неиспользуемую часть (см. выше описание Unused продукции), которая добавлена ​​для совместимости с MLIR. В MLIR существует более общая концепция «регионов», которые могут содержать несколько «блоков» операций, соединённых вместе операциями перехода. Эти блоки имеют идентификаторы, соответствующие Unused продукции, что позволяет отличать их друг от друга. В StableHLO нет операций перехода, поэтому соответствующая часть синтаксиса MLIR не используется (но всё ещё присутствует).

OpInputAttr      ::= OpInputAttrName '=' OpInputAttrValue
OpInputAttrName  ::= letter {letter | digit}
OpInputAttrValue ::= Constant

Входные атрибуты имеют имя и значение, являющееся одной из поддерживаемых констант. Они являются основным способом задания статических метаданных для элементов программы. Например, операция concatenate использует dimension атрибута для указания измерения, по которому конкатенируются входные значения. Аналогично, операция slice использует несколько атрибутов, таких как start_indices и limit_indices для указания границ, используемых для среза входного значения.

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

• layout ( #629 ).
• mhlo.frontend_attributes ( #628 ).
• mhlo.sharding ( #619 ).
• output_operand_aliases ( #740 ).
• Метаданные местоположения ( #594 ).

OpSignature ::= '(' [ValueType {',' ValueType}] ')' '->' '(' [ValueType {',' ValueType}] ')'

Сигнатура операции состоит из типов всех входных значений (список типов слева от -> ) и типов всех выходных значений (список типов справа от -> ). Строго говоря, входные типы избыточны, как и выходные типы почти всегда избыточны (поскольку для большинства операций StableHLO выходные типы могут быть выведены из входных данных). Тем не менее, сигнатура операции намеренно включена в синтаксис StableHLO для совместимости с MLIR.

Ниже приведён пример операции с мнемоникой select_and_scatter . Она принимает 3 входных значения ( %operand , %source и %init_value ), 2 входные функции и 3 входных атрибута ( window_dimensions , window_strides и padding ). Обратите внимание, что сигнатура операции включает только типы входных значений (но не типы входных функций и атрибутов, которые предоставляются встроенными).

%result = "stablehlo.select_and_scatter"(%operand, %source, %init_value) ({
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "stablehlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GE>
    } : (tensor<i32>, tensor<i32>) -> tensor<i1>
    "stablehlo.return"(%0) : (tensor<i1>) -> ()
}, {
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    "stablehlo.return"(%0) : (tensor<i32>) -> ()
}) {
  window_dimensions = dense<[3, 1]> : tensor<2xi64>,
  window_strides = dense<[2, 1]> : tensor<2xi64>,
  padding = dense<[[0, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<4x2xi32>, tensor<2x2xi32>, tensor<i32>) -> tensor<4x2xi32>

Константы

Constant ::= BooleanConstant
           | IntegerConstant
           | FloatConstant
           | ComplexConstant
           | TensorConstant
           | QuantizedTensorConstant
           | StringConstant
           | EnumConstant

Константы StableHLO состоят из литерала и типа, которые вместе представляют значение StableHLO. Как правило, тип является частью синтаксиса константы, за исключением случаев, когда он однозначен (например, логическая константа однозначно имеет тип i1 , тогда как целочисленная константа может иметь несколько возможных типов).

BooleanConstant ::= BooleanLiteral
BooleanLiteral  ::= 'true' | 'false'

Булевы константы представляют собой логические значения true и false . Булевы константы имеют тип i1 .

IntegerConstant   ::= IntegerLiteral ':' IntegerType
IntegerLiteral    ::= ['-' | '+'] DecimalDigits
                    | ['-' | '+'] '0x' HexadecimalDigits
DecimalDigits     ::= decimalDigit {decimalDigit}
HexadecimalDigits ::= hexadecimalDigit {hexadecimalDigit}
decimalDigit      ::= '0' | ... | '9'
hexadecimalDigit  ::= decimalDigit | 'a' | ... | 'f' | 'A' | ... | 'F'

Целочисленные константы представляют целые значения в виде строк в десятичной или шестнадцатеричной системе счисления. Другие системы счисления, например, двоичная или восьмеричная, не поддерживаются. Целочисленные константы имеют следующие ограничения:

• (C1) is_wellformed(integer_literal, integer_type) .

FloatConstant  ::= FloatLiteral ':' FloatType
FloatLiteral   ::= SignPart IntegerPart FractionalPart ScientificPart
                 | '0x' [HexadecimalDigits]
SignPart       ::= ['-' | '+']
IntegerPart    ::= DecimalDigits
FractionalPart ::= ['.' [DecimalDigits]]
ScientificPart ::= [('e' | 'E') ['-' | '+'] DecimalDigits]

Константы с плавающей точкой представляют значения с плавающей точкой в ​​виде строк, использующих десятичную или экспоненциальную запись. Кроме того, шестнадцатеричная запись может использоваться для непосредственного указания базовых битов в формате с плавающей точкой соответствующего типа. Константы с плавающей точкой имеют следующие ограничения:

• (C1) Если используется не шестнадцатеричная запись, is_wellformed(float_literal, float_type) .
• (C2) Если используется шестнадцатеричная запись, size(hexadecimal_digits) = num_bits(float_type) / 4 .

ComplexConstant ::= ComplexLiteral ':' ComplexType
ComplexLiteral  ::= '(' RealPart ',' ImaginaryPart ')'
RealPart        ::= FloatLiteral
ImaginaryPart   ::= FloatLiteral

Комплексные константы представляют комплексные значения, используя списки, состоящие из действительной части (идет первой) и мнимой части (идет второй). Например, (1.0, 0.0) : complex<f32> представляет 1.0 + 0.0i , а (0.0, 1.0) : complex<f32> представляет 0.0 + 1.0i . Порядок, в котором эти части затем сохраняются в памяти, определяется реализацией. Комплексные константы имеют следующие ограничения:

• (C1) is_wellformed(real_part, complex_element_type(complex_type)) .
• (C2) is_wellformed(imaginary_part, complex_element_type(complex_type)) .

TensorConstant ::= TensorLiteral ':' TensorType
TensorLiteral  ::= 'dense' '<' (DenseLiteral | ElementLiteral) '>'
DenseLiteral   ::= DenseDimension | DenseElements
DenseDimension ::= '[' [DenseLiteral {',' DenseLiteral}] ']'
DenseElements  ::= [ElementLiteral {',' ElementLiteral}]
ElementLiteral ::= BooleanLiteral | IntegerLiteral | FloatLiteral | ComplexLiteral

Тензорные константы представляют значения тензора с помощью вложенных списков, заданных с помощью нотации NumPy. Например, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> представляет значение тензора со следующим отображением индексов в элементы: {0, 0} => 1 , {0, 1} => 2 , {0, 2} => 3 , {1, 0} => 4 , {1, 1} => 5 , {1, 2} => 6 . Порядок, в котором эти элементы затем сохраняются в памяти, определяется реализацией. Тензорные константы имеют следующие ограничения:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)) , где:
  • has_syntax(element_literal: Syntax, element_type: Type) = is_wellformed(element_literal, type) .
  • has_syntax(tensor_literal: List, element_type: Type) = has_syntax(tensor_literal..., element_type) .
• (C2) has_shape(tensor_literal, shape(tensor_type)) , где:
  • has_shape(element_literal: Syntax, []) = true .
  • has_shape(tensor_literal: List, shape: List) = size(tensor_literal) = shape[0] and has_shape(tensor_literal..., shape[1:]) .
  • в противном случае — false .

QuantizedTensorConstant ::= QuantizedTensorLiteral ':' QuantizedTensorType
QuantizedTensorLiteral  ::= 'dense' '<' (DenseLiteral | ElementLiteral) '>'

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

• (C1) has_syntax(quantized_tensor_literal, storage_type(quantized_tensor_type)) .
• (C2) has_shape(quantized_tensor_literal, shape(quantized_tensor_type)) .

StringConstant  ::= StringLiteral
StringLiteral   ::= '"' {stringCharacter | escapeSequence} '"'
stringCharacter ::= all ASCII characters except '\00', '\01', ... '\1f' and '"'
escapeSequence  ::= '\' ('"' | '\' | 'n' | 't' | (hexadecimalDigit hexadecimalDigit))

Строковые литералы состоят из байтов, заданных с помощью символов ASCII и экранированных последовательностей. Они не зависят от кодировки, поэтому интерпретация этих байтов определяется реализацией. Строковые литералы имеют тип string .

Операции

пресс

Семантика

Выполняет поэлементную операцию abs над тензором operand и возвращает result тензор. В зависимости от типа элемента выполняет следующие действия:

• Для знаковых целых чисел: целый модуль.
• Для чисел с плавающей точкой: abs из IEEE-754.
• Для комплексных чисел: комплексный модуль.
• Для квантованных типов: dequantize_op_quantize(abs, operand, type(result)) .

Входы

Выходы

Ограничения

• (C1) shape(result) = shape(operand) .
• (C2) baseline_element_type(result) определяется как:
  • complex_element_type(element_type(operand)) если is_complex(operand) .
  • baseline_element_type(operand) в противном случае.

Примеры

// %operand: [-2, 0, 2]
%result = "stablehlo.abs"(%operand) : (tensor<3xi32>) -> tensor<3xi32>
// %result: [2, 0, 2]

Больше примеров

добавлять

Семантика

Выполняет поэлементное сложение двух тензоров lhs и rhs и возвращает result тензор. В зависимости от типа элемента выполняет следующие действия:

• Для булевых значений: логическое ИЛИ.
• Для целых чисел: сложение целых чисел.
• Для поплавков: addition из IEEE-754.
• Для комплексных чисел: комплексное сложение.
• Для квантованных типов: dequantize_op_quantize(add, lhs, rhs, type(result)) .

Входы

Выходы

Ограничения

• Если операция использует неквантованные тензоры:
  • (C1) type(lhs) = type(rhs) = type(result) .
• Если операция использует квантованные тензоры:
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result) .
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result) .
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result) .
  • (C6) Если is_per_axis_quantized(lhs) , то quantization_dimension(lhs) = quantization_dimension(result) .
  • (C7) Если is_per_axis_quantized(rhs) , то quantization_dimension(rhs) = quantization_dimension(result) .

Примеры

// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.add"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[6, 8], [10, 12]]

Больше примеров

после всего

Семантика

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

Входы

Выходы

Примеры

// %input0: !stablehlo.token
// %input1: !stablehlo.token
%result = "stablehlo.after_all"(%input0, %input1) : (!stablehlo.token, !stablehlo.token) -> !stablehlo.token

Больше примеров

all_gather

Семантика

В каждой группе процессов в сетке процессов StableHLO объединяет значения тензоров operands из каждого процесса по all_gather_dim и выдает тензоры results .

Операция разбивает сетку процессов StableHLO на process_groups , которые определяются следующим образом:

• cross_replica(replica_groups) , если channel_id <= 0 and use_global_device_ids = false .
• cross_replica_and_partition(replica_groups) , если channel_id > 0 and use_global_device_ids = false .
• flattened_ids(replica_groups) , если channel_id > 0 and use_global_device_ids = true .

После этого в каждой process_group :

• operands...@receiver = [operand@sender for sender in process_group] для всех receiver в process_group .
• results...@process = concatenate(operands...@process, all_gather_dim) для всех process в process_group .

Входы

Выходы

Ограничения

• (C1) 0 <= all_gather_dim < rank(operands...) .
• (C2) is_unique(replica_groups) .
• (C3) size(replica_groups) определяется как:
  • num_replicas , если используется cross_replica .
  • num_replicas , если используется cross_replica_and_partition .
  • num_processes , если используется flattened_ids .
• (C4) 0 <= replica_groups < size(replica_groups) .
• (C5) Если use_global_device_ids = true , то channel_id > 0 .
• (C6) type(results...) = type(operands...) за исключением:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1) .

Примеры

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [[1, 2], [3, 4]]
// %operand0@(1, 0): [[5, 6], [7, 8]]
// %operand1@(0, 0): [[11, 12], [13, 14]]
// %operand1@(1, 0): [[15, 16], [17, 18]]
%result:2 = "stablehlo.all_gather"(%operand0, %operand1) {
  all_gather_dim = 1 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x4xi64>, tensor<2x4xi64>)
// %result0@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result0@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result1@(0, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]
// %result1@(1, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]

Больше примеров

all_reduce

Семантика

В каждой группе процессов в сетке процессов StableHLO применяет computation функции редукции к значениям тензоров operands из каждого процесса и выдает тензоры results .

Операция разбивает сетку процессов StableHLO на process_groups , которые определяются следующим образом:

• cross_replica(replica_groups) , если channel_id <= 0 and use_global_device_ids = false .
• cross_replica_and_partition(replica_groups) , если channel_id > 0 and use_global_device_ids = false .
• flattened_ids(replica_groups) , если channel_id > 0 and use_global_device_ids = true .

После этого в каждой process_group :

• results...@process[result_index] = exec(schedule) для некоторого schedule двоичного дерева, где:
  • exec(node) = computation(exec(node.left), exec(node.right)) .
  • exec(leaf) = leaf.value .
• schedule — это определяемое реализацией двоичное дерево, упорядоченный обход которого — to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])) .

Входы

Выходы

Ограничения

• (C1) is_unique(replica_groups) .
• (C2) size(replica_groups) определяется как:
  • num_replicas , если используется cross_replica .
  • num_replicas , если используется cross_replica_and_partition .
  • num_processes , если используется flattened_ids .
• (C3) 0 <= replica_groups < size(replica_groups) .
• (C4) Если use_global_device_ids = true , то channel_id > 0 .
• (C5) computation имеет тип (tensor<E>, tensor<E>) -> (tensor<E>) где is_promotable(element_type(operand), E) .
• (C6) shape(results...) = shape(operands...) .
• (C7) element_type(results...) = E .

Примеры

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [1, 2, 3, 4]
// %operand0@(1, 0): [5, 6, 7, 8]
// %operand1@(0, 0): [9, 10, 11, 12]
// %operand1@(1, 0): [13, 14, 15, 16]
%result:2 = "stablehlo.all_reduce"(%operand0, %operand0) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>, tensor<4xi64>)
// %result0@(0, 0): [6, 8, 10, 12]
// %result0@(1, 0): [6, 8, 10, 12]
// %result1@(0, 0): [22, 24, 26, 28]
// %result1@(1, 0): [22, 24, 26, 28]

Больше примеров

все_ко_всем

Семантика

В каждой группе процессов в сетке процессов StableHLO значения тензоров operands разбиваются по split_dimension на части, распределяются по разным частям между процессами, объединяются разбросанные части по concat_dimension и формируются тензоры results . Эта операция разбивает сетку процессов StableHLO на process_groups , которые определяются следующим образом:

• cross_replica(replica_groups) если channel_id <= 0 .
• cross_partition(replica_groups) если channel_id > 0 .

После этого в каждой process_group :

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) для всех sender в process_group .
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] где receiver_index = process_group.index(receiver) .
• results...@process = concatenate(scattered_parts...@process, concat_dimension) .

Входы

Выходы

Ограничения

• (C1) 0 <= split_dimension < rank(operands...) .
• (C2) dim(operands..., split_dimension) % split_count = 0 .
• (C3) 0 <= concat_dimension < rank(operands...) .
• (C4) 0 < split_count .
• (C5) is_unique(replica_groups) .
• (C6) size(replica_groups) определяется как:
  • num_replicas , если используется cross_replica .
  • num_partitions , если используется cross_partition .
• (C7) 0 <= replica_groups < size(replica_groups) .
• (C8) dim(replica_groups, 1) = split_count .
• (C9) type(results...) = type(operands...) за исключением случаев, когда split_dimension != concat_dimension :
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count .
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count .

Примеры

// num_replicas: 2
// num_partitions: 1
// %operand1@(0, 0): [[1, 2, 3, 4],
//                    [5, 6, 7, 8]]
// %operand1@(1, 0): [[9, 10, 11, 12],
//                    [13, 14, 15, 16]]
// %operand2@(0, 0): [[17, 18, 19, 20],
//                    [21, 22, 23, 24]]
// %operand2@(1, 0): [[25, 26, 27, 28],
//                    [29, 30, 31, 32]]
%result:2 = "stablehlo.all_to_all"(%operand1, %operand2) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // channel_id = 0
} : (tensor<2x4xi64>, tensor<2x4xi64>) -> (tensor<4x2xi64>, tensor<4x2xi64>)
// %result#0@(0, 0): [[1, 2], [5, 6], [9, 10], [13, 14]]
// %result#0@(1, 0): [[3, 4], [7, 8], [11, 12], [15, 16]]
// %result#1@(0, 0): [[17, 18], [21, 22], [25, 26], [29, 30]]
// %result#1@(1, 0): [[19, 20], [23, 24], [27, 28], [31, 32]]

Больше примеров

и

Семантика

Выполняет поэлементную операцию «И» над двумя тензорами lhs и rhs и возвращает result тензор. В зависимости от типа элемента выполняет следующие действия:

• Для булевых значений: логическое И.
• Для целых чисел: побитовое И.

Входы

Выходы

Ограничения

• (C1) type(lhs) = type(rhs) = type(result) .

Примеры

// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.and"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[1, 2], [3, 0]]

Больше примеров

атан2

Семантика

Выполняет поэлементную операцию atan2 над lhs и rhs тензорами и возвращает result тензор. В зависимости от типа элемента выполняет следующие действия:

• Для чисел с плавающей точкой: atan2 из IEEE-754.
• Для комплексных чисел: комплексный atan2.
• Для квантованных типов: dequantize_op_quantize(atan2, lhs, rhs, type(result)) .

Входы

Выходы

Ограничения

• (C1) baseline_type(lhs) = baseline_type(rhs) = baseline_type(result) .

Примеры

// %lhs: [0.0, 1.0, -1.0]
// %rhs: [0.0, 0.0, 0.0]
%result = "stablehlo.atan2"(%lhs, %rhs) : (tensor<3xf64>, tensor<3xf64>) -> tensor<3xf64>
// %result: [0.0, 1.57079637, -1.57079637] // [0.0, pi/2, -pi/2]

Больше примеров

batch_norm_grad

Семантика

Вычисляет градиенты нескольких входных данных batch_norm_training , используя обратное распространение из grad_output , и создаёт тензоры grad_operand , grad_scale и grad_offset . Более формально, эту операцию можно выразить как разложение на существующие операции StableHLO с использованием синтаксиса Python следующим образом:

def compute_sum(operand, feature_index):
  (sum,) = reduce(
      inputs=[operand],
      init_values=[constant(0, element_type(operand))],
      dimensions=[i for i in range(rank(operand)) if i != feature_index],
      body=lambda x, y: add(x, y))
  return sum

def compute_mean(operand, feature_index):
  sum = compute_sum(operand, feature_index)
  divisor = constant(size(operand) / dim(operand, feature_index),
                     element_type(operand))
  divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
  return divide(sum, divisor_bcast)

def batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index):
  # Broadcast inputs to type(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance`
  # Intermediate values will be useful for computing gradients
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)

  # Use the implementation from batchnorm_expander.cc in XLA
  # Temporary variables have exactly the same names as in the C++ code
  elements_per_feature = broadcast_in_dim(
      constant(divide(size(operand), dim(operand, feature_index)),
               element_type(grad_output)),
      [], type(operand))
  i1 = multiply(grad_output, elements_per_feature)
  i2 = broadcast_in_dim(
      compute_sum(grad_output, feature_index), [feature_index], type(operand))
  i3 = broadcast_in_dim(
      compute_sum(multiply(grad_output, centered_operand), feature_index),
      [feature_index], type(operand))
  i4 = multiply(i3, centered_operand)
  i5 = divide(i4, add(variance_bcast, epsilon_bcast))
  i6 = subtract(subtract(i1, i2), i5)

  grad_operand =
      multiply(divide(divide(scale_bcast, stddev), elements_per_feature), i6)
  grad_scale =
      compute_sum(multiply(grad_output, normalized_operand), feature_index)
  grad_offset = compute_sum(grad_output, feature_index)

  return grad_operand, grad_scale, grad_offset

Для квантованных типов выполняет dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, mean, variance, grad_output: batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index), operand, scale, mean, variance, grad_output, type(grad_operand), type(grad_scale), type(feature_index)) .

Входы

Выходы

Ограничения

• (C1) 0 <= feature_index < rank(operand) .
• (C2) operand , scale , mean , variance , grad_output , grad_operand , grad_scale и grad_offset имеют одинаковый baseline_element_type .
• (C3) operand , grad_output и grad_operand имеют одинаковую форму.
• (C4) scale , mean , variance , grad_scale и grad_offset имеют одинаковую форму.
• (C5) size(scale) = dim(operand, feature_index) .

Примеры

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
// %grad_output: [
//                [[0.1, 0.1], [0.1, 0.1]],
//                [[0.1, 0.1], [0.1, 0.1]]
//               ]
%grad_operand, %grad_scale, %grad_offset =
"stablehlo.batch_norm_grad"(%operand, %scale, %mean, %variance, %grad_output) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>,
     tensor<2x2x2xf64>) -> (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %grad_operand: [
//                 [[0.0, 0.0], [0.0, 0.0]],
//                 [[0.0, 0.0], [0.0, 0.0]]
//                ]
// %grad_scale:  [0.0, 0.0]
// %grad_offset: [0.4, 0.4]

batch_norm_inference

Семантика

Нормализует тензор operand по всем измерениям, за исключением измерения feature_index , и создаёт result тензор. Более формально эту операцию можно выразить как разложение на существующие операции StableHLO с использованием синтаксиса Python следующим образом:

def batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index):
  # Broadcast inputs to shape(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  offset_bcast = broadcast_in_dim(offset, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance` instead of
  # computing them like `batch_norm_training` does.
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)
  return add(multiply(scale_bcast, normalized_operand), offset_bcast)

Для квантованных типов выполняет dequantize_op_quantize(lambda operand, scale, offset, mean, variance: batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index), operand, scale, offset, mean, variance, type(result)) .

Входы

Выходы

Ограничения

• (C1) 0 <= feature_index < rank(operand) .
• (C2) operand , scale , offset , mean , variance и result имеют один и тот же baseline_element_type .
• (C3) size(scale) = dim(operand, feature_index) .
• (C4) size(offset) = dim(operand, feature_index) .
• (C5) size(mean) = dim(operand, feature_index) .
• (C6) size(variance) = dim(operand, feature_index) .
• (C7) baseline_type(operand) = baseline_type(result) .

Примеры

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
%result = "stablehlo.batch_norm_inference"(%operand, %scale, %offset, %mean, %variance) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>) -> tensor<2x2x2xf64>
// %result: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]

пакетная_норма_обучения

Семантика

Вычисляет среднее значение и дисперсию по всем измерениям, кроме измерения feature_index , и нормализует тензор operand , создавая тензоры output , batch_mean и batch_var . Более формально эту операцию можно выразить как разложение на существующие операции StableHLO с использованием синтаксиса Python следующим образом:

def compute_mean(operand, feature_index):
  (sum,) = reduce(
      inputs=[operand],
      init_values=[constant(0, element_type(operand))],
      dimensions=[i for i in range(rank(operand)) if i != feature_index],
      body=lambda x, y: add(x, y))
  divisor = constant(size(operand) / dim(operand, feature_index),
                     element_type(operand))
  divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
  return divide(sum, divisor_bcast)

def compute_variance(operand, feature_index):
  mean = compute_mean(operand, feature_index)
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  centered_operand = subtract(operand, mean_bcast)
  return compute_mean(mul(centered_operand, centered_operand), feature_index)

def batch_norm_training(operand, scale, offset, epsilon, feature_index):
  mean = compute_mean(operand, feature_index)
  variance = compute_variance(operand, feature_index)
  return batch_norm_inference(operand, scale, offset, mean, variance, epsilon,
                              feature_index),
         mean, variance

Для квантованных типов выполняет dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, offset: batch_norm_training(operand, scale, offset, epsilon, feature_index), operand, scale, offset, type(output), type(batch_mean), type(batch_var)) .

Входы

Выходы

Ограничения

• (C1) 0 <= feature_index < rank(operand) .
• (C2) operand , scale , offset , batch_mean , batch_var и output имеют одинаковый baseline_element_type .
• (C3) size(scale) = dim(operand, feature_index) .
• (C4) size(offset) = dim(operand, feature_index) .
• (C5) size(batch_mean) = dim(operand, feature_index) .
• (C6) size(batch_var) = dim(operand, feature_index) .
• (C7) baseline_type(output) = baseline_type(operand) .

Примеры

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
%output, %batch_mean, %batch_var = "stablehlo.batch_norm_training"(%operand, %scale, %offset) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>) ->
    (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %output: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]
// %batch_mean: [2.0, 3.0]
// %batch_var: [1.0, 1.0]

bitcast_convert

Семантика

Выполняет операцию побитового преобразования тензора operand и создает тензор result , в котором биты всего тензора operand переинтерпретируются с использованием типа тензора result .

Более формально, учитывая E = element_type(operand) , E' = element_type(result) , и R = rank(operand) :

• Если num_bits(E') < num_bits(E) , bits(result[i0, ..., iR-1, :]) = bits(operand[i0, ..., iR-1]) .
• Если num_bits(E') > num_bits(E) , bits(result[i0, ..., iR-2]) = bits(operand[i0, ..., iR-2, :]) .
• If num_bits(E') = num_bits(E) , bits(result[i0, ..., iR-1]) = bits(operand[i0, ..., iR-1]) .

bits returns in-memory representation of a given value, and its behavior is implementation-defined because the exact representation of tensors is implementation-defined, and the exact representation of element types is implementation-defined as well.

Входы

Выходы

Ограничения

• (C1) Given E = is_quantized(operand) ? storage_type(operand) : element_type(operand) , E' = is_quantized(result) ? storage_type(result) : element_type(result) , and R = rank(operand) :
  • If num_bits(E') = num_bits(E) , shape(result) = shape(operand) .
  • If num_bits(E') < num_bits(E) :
  • rank(result) = R + 1 .
  • dim(result, i) = dim(operand, i) for all 0 <= i < R .
  • dim(result, R) * num_bits(E') = num_bits(E) .
  • If num_bits(E') > num_bits(E) :
  • rank(result) = R - 1 .
  • dim(result, i) = dim(operand, i) for all 0 <= i < R .
  • dim(operand, R - 1) * num_bits(E) = num_bits(E') .
• (C2) If is_complex(operand) or is_complex(result) , then is_complex(operand) and is_complex(result) .

Примеры

// %operand: 0x0123456789ABCDEF
%result = "stablehlo.bitcast_convert"(%operand) : (tensor<f64>) -> tensor<4xf16>
// %result: [0xCDEF, 0x89AB, 0x4567, 0x0123] // little-endian representation

Больше примеров

broadcast_in_dim

Семантика

Expands the dimensions and/or rank of an input tensor by duplicating the data in the operand tensor and produces a result tensor. More formally, result[result_index] = operand[operand_index] where for all d in axes(operand) :

• operand_index[d] = 0 if dim(operand, d) = 1 .
• operand_index[d] = result_index[broadcast_dimensions[d]] otherwise.

Входы

Выходы

Ограничения

• (C1) element_type(result) is given by:
  • element_type(operand) , if !is_per_axis_quantized(operand) .
  • element_type(operand) except that quantization_dimension(operand) , scales(operand) , and zero_points(operand) may differ from quantization_dimension(result) , scales(result) , and zero_points(result) resp., otherwise.
• (C2) size(broadcast_dimensions) = rank(operand) .
• (C3) 0 <= broadcast_dimensions < rank(result) .
• (C4) is_unique(broadcast_dimensions) .
• (C5) For all d in axes(operand) :
  • dim(operand, d) = 1 or
  • dim(operand, d) = dim(result, broadcast_dimensions[d]) .
• (C6) If is_per_axis_quantized(result) :
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)] .
  • If dim(operand, quantization_dimension(operand)) = 1 , then scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))) .

Примеры

// %operand: [
//            [1, 2, 3]
//           ]
%result = "stablehlo.broadcast_in_dim"(%operand) {
  broadcast_dimensions = array<i64: 2, 1>
} : (tensor<1x3xi32>) -> tensor<2x3x2xi32>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

Больше примеров

случай

Семантика

Produces the output from executing exactly one function from branches depending on the value of index . More formally, result = selected_branch() where:

• selected_branch = branches[index] if 0 <= index < size(branches) .
• selected_branch = branches[-1] otherwise.

Входы

Выходы

Ограничения

• (C1) 0 < size(branches) .
• (C2) input_types(branches...) = [] .
• (C3) same(output_types(branches...)) .
• (C4) type(results...) = output_types(branches[0]) .

Примеры

// %index: -1
// %result_branch0: [0, 0]
// %result_branch1: [1, 1]
%result0, %result1 = "stablehlo.case"(%index) ({
  "stablehlo.return"(%result_branch0, %result_branch0) : (tensor<2xi64>, tensor<2xi64>) -> ()
}, {
  "stablehlo.return"(%result_branch1, %result_branch1) : (tensor<2xi64>, tensor<2xi64>) -> ()
}) : (tensor<i32>) -> (tensor<2xi64>, tensor<2xi64>)
// %result0: [1, 1]
// %result1: [1, 1]

Больше примеров

cbrt

Семантика

Performs element-wise cubic root operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: rootn(x, 3) from IEEE-754.
• For complex numbers: complex cubic root.
• For quantized types: dequantize_op_quantize(cbrt, operand, type(result))

Входы

Выходы

Ограничения

• (C1) baseline_type(operand) = baseline_type(result) .

Примеры

// %operand: [0.0, 1.0, 8.0, 27.0]
%result = "stablehlo.cbrt"(%operand) : (tensor<4xf64>) -> tensor<4xf64>
// %result: [0.0, 1.0, 2.0, 3.0]

Больше примеров

потолок

Семантика

Performs element-wise ceil of operand tensor and produces a result tensor. Implements the roundToIntegralTowardPositive operation from the IEEE-754 specification. For quantized types, performs dequantize_op_quantize(ceil, operand, type(result)) .

Входы

Выходы

Ограничения

• (C1) baseline_type(operand) = baseline_type(result) .

Примеры

// %operand: [-0.8166, -0.2530, 0.2530, 0.8166, 2.0]
%result = "stablehlo.ceil"(%operand) : (tensor<5xf32>) -> tensor<5xf32>
// %result: [-0.0, -0.0, 1.0, 1.0, 2.0]

Больше примеров

cholesky

Семантика

Computes the Cholesky decomposition of a batch of matrices.

More formally, for all i in index_space(result) , result[i0, ..., iR-3, :, :] is a Cholesky decomposition of a[i0, ..., iR-3, :, :] , in the form of either of a lower-triangular (if lower is true ) or upper-triangular (if lower is false ) matrix. The output values in the opposite triangle, ie the strict upper triangle or strict lower triangle correspondingly, are implementation-defined.

If there exists i where the input matrix is not an Hermitian positive-definite matrix, then the behavior is undefined.

For quantized types, performs dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)) .

Входы

Выходы

Ограничения

• (C1) baseline_type(a) = baseline_type(result) .
• (C2) 2 <= rank(a) .
• (C3) dim(a, -2) = dim(a, -1) .

Примеры

// %a: [
//      [1.0, 2.0, 3.0],
//      [2.0, 20.0, 26.0],
//      [3.0, 26.0, 70.0]
//     ]
%result = "stablehlo.cholesky"(%a) {
  lower = true
} : (tensor<3x3xf32>) -> tensor<3x3xf64>
// %result: [
//           [1.0, 0.0, 0.0],
//           [2.0, 4.0, 0.0],
//           [3.0, 5.0, 6.0]
//          ]

зажим

Семантика

Clamps every element of the operand tensor between a minimum and maximum value and produces a result tensor. More formally, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element) , where min_element = rank(min) = 0 ? min[] : min[result_index] , max_element = rank(max) = 0 ? max[] : max[result_index] . For quantized types, performs dequantize_op_quantize(clamp, min, operand, max, type(result)) .

Imposing an ordering on complex numbers involves surprising semantics, so in the future we are planning to remove support for complex numbers for this operation ( #560 ).

Входы

Выходы

Ограничения

• (C1) rank(min) = 0 or shape(min) = shape(operand) .
• (C2) rank(max) = 0 or shape(max) = shape(operand) .
• (C3) baseline_element_type(min) = baseline_element_type(operand) = baseline_element_type(max) .
• (C4) baseline_type(operand) = baseline_type(result) .

Примеры

// %min: [5, 10, 15]
// %operand: [3, 13, 23]
// %max: [10, 15, 20]
%result = "stablehlo.clamp"(%min, %operand, %max) : (tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<3xi32>
// %result: [5, 13, 20]

Больше примеров

collective_broadcast

Семантика

Within each process group in the StableHLO process grid, send the value of the operand tensor from the source process to the target processes and produce a result tensor.

The operation splits the StableHLO process grid into process_groups which is defined as follows:

• cross_replica(replica_groups) if channel_id <= 0 .
• cross_partition(replica_groups) if channel_id > 0 .

Afterwards, result@process is given by:

• operand@process_groups[i, 0] if there exists an i such that the process is in process_groups[i] .
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) otherwise.

Входы

Выходы

Ограничения

• (C1) is_unique(replica_groups) .
• (C2) 0 <= replica_groups < N where N is defined as:
  • num_replicas if cross_replica is used.
  • num_partitions if cross_partition is used.
• (C3) type(result) = type(operand) .

Примеры

// num_replicas: 4
// num_partitions: 1
// %operand@(0, 0): [[1, 2]]
// %operand@(1, 0): [[3, 4]]
// %operand@(2, 0): [[5, 6]]
// %operand@(3, 0): [[7, 8]]
%result = "stablehlo.collective_broadcast"(%operand) {
  replica_groups = dense<[[2, 1]]> : tensor<1x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor1x2xi64>) -> tensor<1x2xi64>
// %result@(0, 0): [[0, 0]]
// %result@(1, 0): [[5, 6]]
// %result@(2, 0): [[5, 6]]
// %result@(3, 0): [[0, 0]]

collective_permute

Семантика

Within each process group in the StableHLO process grid, sends the value of the operand tensor from the source process to the target process and produces a result tensor.

The operation splits the StableHLO process grid into process_groups which is defined as follows:

• cross_replica(source_target_pairs) if channel_id <= 0 .
• cross_partition(source_target_pairs) if channel_id > 0 .

Afterwards, result@process is given by:

• operand@process_groups[i, 0] , if there exists an i such that process_groups[i, 1] = process .
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) otherwise.

Входы

Выходы

Ограничения

• (C1) dim(source_target_pairs, 1) = 2 .
• (C2) is_unique(source_target_pairs[:, 0]) .
• (C3) is_unique(source_target_pairs[:, 1]) .
• (C4) 0 <= source_target_pairs < N , where N is defined as:
  • num_replicas if cross_replica is used.
  • num_partitions if cross_partition is used.
• (C5) type(result) = type(operand) .

Примеры

// num_replicas: 3
// num_partitions: 1
// %operand@(0, 0): [[1, 2], [3, 4]]
// %operand@(1, 0): [[5, 6], [7, 8]]
// %operand@(2, 0): [[9, 10], [11, 12]]
%result = "stablehlo.collective_permute"(%operand) {
  source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<2x2xi64>) -> tensor<2x2xi64>
//
// %result@(0, 0): [[0, 0], [0, 0]]
// %result@(1, 0): [[1, 2], [3, 4]]
// %result@(2, 0): [[5, 6], [7, 8]]

Больше примеров

сравнивать

Семантика

Performs element-wise comparison of lhs and rhs tensors according to comparison_direction and compare_type , and produces a result tensor.

The values of comparison_direction and compare_type have the following semantics:

For boolean and integer element types:

• EQ : lhs = rhs .
• NE : lhs != rhs .
• GE : lhs >= rhs .
• GT : lhs > rhs .
• LE : lhs <= rhs .
• LT : lhs < rhs .

For floating-point element types with compare_type = FLOAT , the op implements the following IEEE-754 operations:

• EQ : compareQuietEqual .
• NE : compareQuietNotEqual .
• GE : compareQuietGreaterEqual .
• GT : compareQuietGreater .
• LE : compareQuietLessEqual .
• LT : compareQuietLess .

For floating-point element types with compare_type = TOTALORDER , the op uses the combination of totalOrder and compareQuietEqual operations from IEEE-754.

For complex element types, lexicographic comparison of (real, imag) pairs is performed using the provided comparison_direction and compare_type . Imposing an ordering on complex numbers involves surprising semantics, so in the future we are planning to remove support for complex numbers when comparison_direction is GE , GT , LE or LT ( #560 ).

For quantized types. performs dequantize_compare(lhs, rhs, comparison_direction) .

Входы

Выходы

Ограничения

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs) .
• (C2) shape(lhs) = shape(rhs) = shape(result) .
• (C3) compare_type is defined as:
  • SIGNED if is_signed_integer(element_type(lhs)) .
  • UNSIGNED if is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)) .
  • FLOAT or TOTALORDER if is_float(element_type(lhs)) .
  • FLOAT if is_complex(element_type(lhs)) .

Примеры

// %lhs: [1.0, 3.0]
// %rhs: [1.1, 2.9]
%result = "stablehlo.compare"(%lhs, %rhs) {
  comparison_direction = #stablehlo<comparison_direction LT>,
  compare_type = #stablehlo<comparison_type FLOAT>
} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>
// %result: [true, false]

Больше примеров

сложный

Семантика

Performs element-wise conversion to a complex value from a pair of real and imaginary values, lhs and rhs , and produces a result tensor.

Входы

Выходы

Ограничения

• (C1) type(lhs) = type(rhs) .
• (C2) shape(result) = shape(lhs) .
• (C3) element_type(result) has type complex<E> where E = element_type(lhs) .

Примеры

// %lhs: [1.0, 3.0]
// %rhs: [2.0, 4.0]
%result = "stablehlo.complex"(%lhs, %rhs) : (tensor<2xf64>, tensor<2xf64>) -> tensor<2xcomplex<f64>>
// %result: [(1.0, 2.0), (3.0, 4.0)]

Больше примеров

композитный

Семантика

Encapsulates an operation made up (composed) of other StableHLO operations, taking inputs and composite_attributes and producing results . The semantics of the op are implemented by the decomposition attribute. The composite op can be replaced with its decomposition without changing program semantics. In cases where inlining the decomposition does not provide the same op semantics, prefer using custom_call .

The version field (defaults to 0 ) is used to denote when a composite's semantics change.

Входы

Выходы

Ограничения

• (C1) is_namespaced_op_name(name)
• (C2) is_defined_in_parent_scope(decomposition)
• (C3) types(inputs...) == input_types(decomposition)
• (C4) types(results...) == output_types(decomposition)

Примеры

%results = "stablehlo.composite"(%input0, %input1) {
  name = "my_namespace.my_op",
  composite_attributes = {
    my_attribute = "my_value"
  },
  decomposition = @my_op,
  version = 1 : i32
} : (tensor<f32>, tensor<f32>) -> tensor<f32>

Больше примеров

конкатенировать

Семантика

Concatenates inputs along dimension dimension in the same order as the given arguments and produces a result tensor. More formally, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1] , where:

1. id = d0 + ... + dk-1 + kd .
2. d is equal to dimension , and d0 , ... are d th dimension sizes of inputs .

Входы

Выходы

Ограничения

• (C1) same(element_type(inputs...)) .
• (C2) same(shape(inputs...)) except for dim(inputs..., dimension) .
• (C3) 0 < size(inputs) .
• (C4) 0 <= dimension < rank(inputs[0]) .
• (C5) element_type(result) = element_type(inputs[0]) .
• (C6) shape(result) = shape(inputs[0]) except for:
  • dim(result, dimension) = dim(inputs[0], dimension) + ... .

Примеры

// %input0: [[1, 2], [3, 4], [5, 6]]
// %input1: [[7, 8]]
%result = "stablehlo.concatenate"(%input0, %input1) {
  dimension = 0 : i64
} : (tensor<3x2xi64>, tensor<1x2xi64>) -> tensor<4x2xi64>
// %result: [[1, 2], [3, 4], [5, 6], [7, 8]]

Больше примеров

постоянный

Семантика

Produces an output tensor from a constant value .

Входы

Выходы

Ограничения

• (C1) type(value) = type(output) .

Примеры

%output = "stablehlo.constant"() {
  value = dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>
} : () -> tensor<2x2xf32>
// %output: [[0.0, 1.0], [2.0, 3.0]]

Больше примеров

конвертировать

Семантика

Performs an element-wise conversion from one element type to another on operand tensor and produces a result tensor.

For boolean-to-any-supported-type conversions, the value false is converted to zero, and the value true is converted to one. For any-supported-type-to-boolean conversions, a zero value is converted to false , and non-zero values are converted to true . See below for how this work for complex types.

For conversions involving integer-to-integer , integer-to-floating-point or floating-point-to-floating-point , if the source value can be exactly represented in the destination type, the result value is that exact representation. Otherwise, the behavior is TBD ( #180 ).

For conversions involving floating-point-to-integer , the fractional part is truncated. If the truncated value cannot be represented in the destination type, the behavior is TBD ( #180 ).

Conversion involving complex-to-complex follow the same behavior of floating-point-to-floating-point conversions for converting real and imaginary parts.

For complex-to-any-other-type and any-other-type-to-complex conversions, the source imaginary value is ignored or the destination imaginary value is zeroed, respectively. The conversion of the real part follows the floating-point conversions.

In principle, this operation could express dequantization (conversion from quantized tensors to regular tensors), quantization (conversion from regular tensors to quantized tensors) and requantization (conversion between quantized tensors), but at the moment we have dedicated operations for that - uniform_dequantize for the first use case and uniform_quantize for the second and the third use cases. In the future, these two ops may be merged into convert ( #1576 ).

Входы

Выходы

Ограничения

• (C1) shape(operand) = shape(result) .

Примеры

// %operand: [-1, 0, 1]
%result = "stablehlo.convert"(%operand) : (tensor<3xi64>) -> tensor<3xcomplex<f64>>
// %result: [(-1.0, 0.0), (0.0, 0.0), (1.0, 0.0)]

Больше примеров

свертка

Семантика

Computes dot products between windows of lhs and slices of rhs and produces result . The following diagram shows how elements in result are computed from lhs and rhs using a concrete example.

More formally, consider the following reframing of the inputs in terms of lhs in order to be able to express windows of lhs :

• lhs_window_dimensions = lhs_shape(dim(lhs, input_batch_dimension), dim(rhs, kernel_spatial_dimensions), dim(lhs, input_feature_dimension)) .
• lhs_window_strides = lhs_shape(1, window_strides, 1) .
• lhs_padding = lhs_shape([0, 0], padding, [0, 0]) .
• lhs_base_dilations = lhs_shape(1, lhs_dilation, 1) .
• lhs_window_dilations = lhs_shape(1, rhs_dilation, 1) .

This reframing uses the following helper functions:

• lhs_shape(n, hw, c) = permute([n] + hw + [c], [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]) .
• result_shape(n1, hw, c1) = permute([n1] + hw + [c1], [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]) .
• permute([j0, j1, ..., jR-1], permutation) = [i0, i1, ..., iR-1] where j[d] = i[permutation[d]] .

If feature_group_count = 1 and batch_group_count = 1 , then for all output_spatial_index in index_space(dim(result, output_spatial_dimensions...)) , result[result_shape(:, output_spatial_index, :)] = dot_product where:

• padding_value = constant(0, element_type(lhs)) .
• padded_lhs = pad(lhs, padding_value, lhs_padding[:, 0], lhs_padding[:, 1], lhs_base_dilations - 1) .
• lhs_window_start = lhs_shape(0, output_spatial_index, 0) * lhs_window_strides .
• lhs_window = slice(padded_lhs, lhs_window_start, lhs_window_start + lhs_window_dimensions, lhs_window_dilations) .
• reversed_lhs_window = reverse(lhs_window, [input_spatial_dimensions[dim] for dim in range(size(window_reversal)) if window_reversal[dim] = true]) . This feature appears to be unused, so in the future we are planning to remove it ( #1181 ).
• dot_product = dot_general(reversed_lhs_window, rhs, lhs_batching_dimensions=[], lhs_contracting_dimensions=input_spatial_dimensions + [input_feature_dimension], rhs_batching_dimensions=[], rhs_contracting_dimensions=kernel_spatial_dimensions + [kernel_input_feature_dimension]) .

If feature_group_count > 1 :

• lhses = split(lhs, feature_group_count, input_feature_dimension) .
• rhses = split(rhs, feature_group_count, kernel_output_feature_dimension) .
• results... = convolution(lhses..., rhses..., ..., feature_group_count=1, ...) .
• result = concatenate(results, output_feature_dimension) .

If batch_group_count > 1 :

• lhses = split(lhs, batch_group_count, input_batch_dimension) .
• rhses = split(rhs, batch_group_count, kernel_output_feature_dimension) .
• results... = convolution(lhses..., rhses..., ..., batch_group_count=1, ...) .
• result = concatenate(results, output_feature_dimension) .

For quantized types, performs dequantize_op_quantize( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs, type(result)) .

For hybrid quantized types, performs hybrid_dequantize_then_op( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs) .

Входы

Выходы

Ограничения

• (C1) N = rank(lhs) = rank(rhs) .
• (C2) size(window_strides) = N - 2 .
• (C3) 0 < window_strides .
• (C4) shape(padding) = [N - 2, 2] .
• (C5) size(lhs_dilation) = N - 2 .
• (C6) 0 < lhs_dilation .
• (C7) size(rhs_dilation) = N - 2 .
• (C8) 0 < rhs_dilation .
• (C9) size(window_reversal) = N - 2 .
• (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0 .
• (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0 .
• (C12) size(input_spatial_dimensions) = N - 2 .
• (C13) Given input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension] :
  • is_unique(input_dimensions) .
  • 0 <= input_dimensions < N .
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count .
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0 .
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0 .
• (C17) size(kernel_spatial_dimensions) = N - 2 .
• (C18) Given kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension] :
  • is_unique(kernel_dimensions) .
  • 0 <= kernel_dimensions < N .
• (C19) size(output_spatial_dimensions) = N - 2 .
• (C20) Given output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension] :
  • is_unique(output_dimensions) .
  • 0 <= output_dimensions < N .
• (C21) 0 < feature_group_count .
• (C22) 0 < batch_group_count .
• (C23) feature_group_count = 1 or batch_group_count = 1 .
• (C24) size(precision_config) = 2 .
• (C25) dim(result, result_dim) is defined as:
  • dim(lhs, input_batch_dimension) / batch_group_count if result_dim = output_batch_dimension .
  • dim(rhs, kernel_output_feature_dimension) if result_dim = output_feature_dimension .
  • num_windows otherwise, where:
  • output_spatial_dimensions[spatial_dim] = result_dim .
  • lhs_dim = input_spatial_dimensions[spatial_dim] .
  • rhs_dim = kernel_spatial_dimensions[spatial_dim] .
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1 .
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1] .
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1 .
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim] .
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1 .
• (C26) rank(result) = N .
• If the operation uses non-quantized tensors:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result) .
• If the operation uses quantized tensors:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C29) If is_per_axis_quantized(rhs) , then quantization_dimension(rhs) = kernel_output_feature_dimension .
  • (C30) If is_per_axis_quantized(result) , then quantization_dimension(result) = output_feature_dimension .
  • If is_quantized(lhs) :
  • (C31) storage_type(lhs) = storage_type(rhs) .
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C33) If is_per_tensor_quantized(rhs) , then is_per_tensor_quantized(result) .
  • If !is_quantized(lhs) :
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result) .

Примеры

// %lhs: [[
//        [
//          [1], [2], [5], [6]
//        ],
//        [
//          [3], [4], [7], [8]
//        ],
//        [
//          [10], [11], [14], [15]
//        ],
//        [
//          [12], [13], [16], [17]
//        ]
//      ]]
//
// %rhs: [
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]]
//       ]
%result = "stablehlo.convolution"(%lhs, %rhs) {
  window_strides = array<i64: 4, 4>,
  padding = dense<0> : tensor<2x2xi64>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  // In the StableHLO dialect, dimension numbers are encoded via:
  // `[<input dimensions>]x[<kernel dimensions>]->[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" are spatial dimensions.
  dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
  batch_group_count = 1 : i64,
  feature_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

Больше примеров

косинус

Семантика

Performs element-wise cosine operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: cos from IEEE-754.
• For complex numbers: complex cosine.
• For quantized types: dequantize_op_quantize(cosine, operand, type(result)) .

Входы

Выходы

Ограничения

• (C1) baseline_type(operand) = baseline_type(result) .

Примеры

// %operand: [
//            [0.0, 1.57079632],       // [0, pi/2]
//            [3.14159265, 4.71238898] // [pi, 3pi/2]
//           ]
%result = "stablehlo.cosine"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[1.0, 0.0], [-1.0, 0.0]]

Больше примеров

count_leading_zeros

Семантика

Performs element-wise count of the number of leading zero bits in the operand tensor and produces a result tensor.

Входы

Выходы

Ограничения

• (C1) type(operand) = type(result) .

Примеры

// %operand: [[0, 1], [128, -1]]
%result = "stablehlo.count_leading_zeros"(%operand) : (tensor<2x2xi64>) -> tensor<2x2xi64>
// %result: [[64, 63], [56, 0]]

Больше примеров

custom_call

Семантика

Encapsulates an implementation-defined operation call_target_name that takes inputs and called_computations and produces results . has_side_effect , backend_config and api_version may be used to provide additional implementation-defined metadata.

At the moment, this operation contains a fairly disorganized collection of metadata which reflects organic evolution of its counterpart operation in the XLA compiler. In the future, we are planning to unify this metadata ( #741 ).

Входы

Выходы

(XLA GPU Support) Special custom_call targets

There are three special call_target_name related to buffer types: CreateBuffer creates an uninitialized buffer , Pin creates an initialized buffer and Unpin deallocates a buffer and returns the content of the buffer .

%uninitialized_buffer = "stablehlo.custom_call"() {
  call_target_name = "CreateBuffer",
  api_version = 4 : i32,
} : () -> memref<4xf64>

%initialized_buffer = "stablehlo.custom_call"(%init_value) {
  call_target_name = "Pin",
  api_version = 4 : i32,
} : (tensor<4xf64>) -> memref<4xf64>

%dealloc_buffer = "stablehlo.custom_call"(%initialized_buffer) {
  call_target_name = "Unpin",
  api_version = 4 : i32,
} : (memref<4xf64>) -> tensor<4xf64>

Псевдоним

Some custom_call ops may require a part in the outputs and a part in the operands to share the same memory. This can be expressed via output_operand_aliases . An alias pair representation consists a list of output tuple indices representing the output part, and an operand_index along with a list of operand tuple indices representing the operand part. The list of output or operand tuple indices is empty if the corresponding type is not a tuple type, and can be arbitrarily long for an arbitrarily nested tuple type. This is similar to the XLA alias representation .

The output part and the input part in an alias pair must have the same type. For custom_call ops that aren't call to CreateBuffer , Pin and Unpin , a buffer operand can appear in at most one pair of alias, and a buffer output must appear in one pair of alias.

Примеры

%results = "stablehlo.custom_call"(%input0) {
  call_target_name = "foo",
  has_side_effect = false,
  backend_config = {bar = 42 : i32},
  api_version = 4 : i32,
  called_computations = [@foo]
} : (tensor<f64>) -> tensor<f64>

%updated_buffer = "stablehlo.custom_call"(%buffer) {
  call_target_name = "Update",
  api_version = 4 : i32,
  output_operand_aliases = [
    #stablehlo.output_operand_alias<output_tuple_indices = [],
      operand_index = 0,
      operand_tuple_indices = []>]
} : (memref<4xf64>) -> memref<4xf64>

разделять

Семантика

Performs element-wise division of dividend lhs and divisor rhs tensors and produces a result tensor. Depending on the element type, does the following:

• For integers: integer division which produces the algebraic quotient with any fractional part discarded.
• For floats: division from IEEE-754.
• For complex numbers: complex division.
• For quantized types:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)) .

Входы

Выходы

Ограничения

• (C1) baseline_type(lhs) = baseline_type(rhs) = baseline_type(result) .

Примеры

// %lhs: [17.1, -17.1, 17.1, -17.1]
// %rhs: [3.0, 3.0, -3.0, -3.0]
%result = "stablehlo.divide"(%lhs, %rhs) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// %result: [5.66666651, -5.66666651, -5.66666651, 5.66666651]

Больше примеров

dot_general

Семантика

Computes dot products between slices of lhs and slices of rhs and produces a result tensor.

More formally, result[result_index] = dot_product , where:

• lhs_result_dimensions = [d for d in axes(lhs) and d not in lhs_batching_dimensions and d not in lhs_contracting_dimensions] .
• rhs_result_dimensions = [d for d in axes(rhs) and d not in rhs_batching_dimensions and d not in rhs_contracting_dimensions] .
• result_batching_index + result_lhs_index + result_rhs_index = result_index where size(result_batching_index) = size(lhs_batching_dimensions) , size(result_lhs_index) = size(lhs_result_dimensions) and size(result_rhs_index) = size(rhs_result_dimensions) .
• transposed_lhs = transpose(lhs, lhs_batching_dimensions + lhs_result_dimensions + lhs_contracting_dimensions) .
• transposed_lhs_slice = slice(transposed_lhs, result_batching_index + result_lhs_index + [:, ..., :]) .
• reshaped_lhs_slice = reshape(transposed_lhs_slice, dims(lhs, lhs_contracting_dimensions)) .
• transposed_rhs = transpose(rhs, rhs_batching_dimensions + rhs_result_dimensions + rhs_contracting_dimensions) .
• transposed_rhs_slice = slice(transposed_rhs, result_batching_index + result_rhs_index + [:, ..., :]) .
• reshaped_rhs_slice = reshape(transposed_rhs_slice, dims(rhs, rhs_contracting_dimensions)) .
• dot_product = reduce( inputs=[multiply(reshaped_lhs_slice, reshaped_rhs_slice)], init_values=[constant(0, element_type(result))], dimensions=range(size(lhs_contracting_dimensions)), body=lambda x, y: add(x, y)) .

For quantized types, performs dequantize_op_quantize( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs, type(result)) .

For hybrid quantized types, performs hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs) .

precision_config controls the tradeoff between speed and accuracy for computations on accelerator backends. This can be one of the following (at the moment, the semantics of these enum values is underspecified, but we are planning to address this in #755 ):

• DEFAULT : Fastest calculation, but least accurate approximation to the original number.
• HIGH : Slower calculation, but more accurate approximation to the original number.
• HIGHEST : Slowest calculation, but most accurate approximation to the original number.

A DotAlgorithm defines the main properties of the algorithm used to implement the dot operation, which also defines the precision. If the algorithm attribute fields are set, then the precision_config must be DEFAULT . DotAlgorithms do not have a default value, as the default parameters are implementation defined. As such, all dot algorithm fields may be set to None to specify an empty dot algorithm, which will instead use the precision_config value.

DotAlgorithm fields include:

• lhs_precision_type and rhs_precision_type , the precisions that the LHS and RHS of the operation are rounded to. Precision types are independent from the storage types of the inputs and the output.
• accumulation_type the precision used for accumulation.
• lhs_component_count , rhs_component_count , and num_primitive_operations apply when we are doing an algorithm which decomposes the LHS and/or RHS into multiple components and does multiple "primitive" dot operations on those values - usually to emulate a higher precision (eg Leveraging the bfloat16 Artificial Intelligence Datatype For Higher-Precision Computations : bf16_6x tf32_3x, etc). For algorithms with no decomposition, these values should be set to 1 .
• allow_imprecise_accumulation to specify if accumulation in lower precision is permitted for some steps (eg CUBLASLT_MATMUL_DESC_FAST_ACCUM ).

Example DotAlgorithm attributes:

// Inputs are casted to tf32, and then accumulated in f32:
{lhs_precision_type = tf32,
 rhs_precision_type = tf32,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = false}


// bf16_6x: each input is decomposed to 3 bf16 components, then 6 dot operations are done on those components, and the result is accumulated in f32.
{lhs_precision_type = bf16,
 rhs_precision_type = bf16,
 accumulation_type = f32,
 lhs_component_count = 3,
 rhs_component_count = 3,
 num_primitive_operations = 6,
 allow_imprecise_accumulation = false}


// Inputs are (casted to) f8e5m2, and we accumulate in f32, but for some steps we may accumulate in lower precision.
{lhs_precision_type = f8e5m2,
 rhs_precision_type = f8e5m2,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = true}

It is up to the implementations to decide which combinations are supported. In general, it is not guaranteed that each algorithm is supported on each accelerator type by the consumer of the StableHLO. If a given algorithm is not supported, an error should be raised as opposed to falling back to an alternative. StableHLO verification will provide best effort verification, preventing algorithms that are not known to be supported on any hardware.

See xla_data.proto > Algorithm for some supported algorithm values. Ticket #2483 captures the plan to create a centralized doc on supported algorithms by backend.

Входы

Выходы

Ограничения

• (C1) size(lhs_batching_dimensions) = size(rhs_batching_dimensions) .
• (C2) size(lhs_contracting_dimensions) = size(rhs_contracting_dimensions) .
• (C3) is_unique(lhs_batching_dimensions + lhs_contracting_dimensions) .
• (C4) is_unique(rhs_batching_dimensions + rhs_contracting_dimensions) .
• (C5) 0 <= lhs_batching_dimensions < rank(lhs) .
• (C6) 0 <= lhs_contracting_dimensions < rank(lhs) .
• (C7) 0 <= rhs_batching_dimensions < rank(rhs) .
• (C8) 0 <= rhs_contracting_dimensions < rank(rhs) .
• (C9) dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...) .
• (C10) dim(lhs, lhs_contracting_dimensions...) = dim(rhs, rhs_contracting_dimensions...) .
• (C11) size(precision_config) = 2 .
• (C12) shape(result) = dim(lhs, lhs_batching_dimensions) + dim(lhs, lhs_result_dimensions) + dim(rhs, rhs_result_dimensions) .
• If the operation uses non-quantized tensors:
  • (C13) element_type(lhs) = element_type(rhs) .
• If the operation uses quantized tensors:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C15) zero_points(rhs) = 0 .
  • (C16) If is_per_axis_quantized(rhs) , then quantization_dimension(rhs) not in rhs_contracting_dimensions .
  • If is_quantized(lhs) :
  • (C17) storage_type(lhs) = storage_type(rhs) .
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C19) If is_per_tensor_quantized(rhs) , then is_per_tensor_quantized(result) .
  • If !is_quantized(lhs) :
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result) .
• If !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation) :
  • (C21) precision_config... = DEFAULT .
  • (C22) 0 < lhs_component_count .
  • (C23) 0 < rhs_component_count .
  • (C24) 0 < num_primitive_operations .

Примеры

// %lhs: [
//        [[1, 2],
//         [3, 4]],
//        [[5, 6],
//         [7, 8]]
//       ]
// %rhs: [
//        [[1, 0],
//         [0, 1]],
//        [[1, 0],
//         [0, 1]]
//       ]
%result = "stablehlo.dot_general"(%lhs, %rhs) {
  dot_dimension_numbers = #stablehlo.dot<
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimensions = [1]
  >,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>],
  algorithm = #stablehlo.dot_algorithm<
    lhs_precision_type = tf32,
    rhs_precision_type = tf32,
    accumulation_type = f32,
    lhs_component_count = 1,
    rhs_component_count = 1,
    num_primitive_operations = 1,
    allow_imprecise_accumulation = false
  >
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

Больше примеров

dynamic_broadcast_in_dim

Семантика

This operation is functionally identical to broadcast_in_dim op, but the result shape is specified dynamically via output_dimensions .

The operation also accepts optional attributes known_expanding_dimensions , known_nonexpanding_dimensions to express static knowledge about the expanding behavior of dimensions. If not specified, all dimensions are assumed to be possibly expanding.

Входы

Выходы

Ограничения

• (C1) element_type(result) is given by:
  • element_type(operand) , if !is_per_axis_quantized(operand) .
  • element_type(operand) except that quantization_dimension(operand) , scales(operand) , and zero_points(operand) may differ from quantization_dimension(result) , scales(result) , and zero_points(result) resp., otherwise.
• (C2) size(broadcast_dimensions) = rank(operand) .
• (C3) 0 <= broadcast_dimensions < rank(result) .
• (C4) is_unique(broadcast_dimensions) .
• (C5) For all d in axes(operand) :
  • dim(operand, d) = 1 or
  • dim(operand, d) = dim(result, broadcast_dimensions[d]) .
• (C6) If is_per_axis_quantized(result) :
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)] .
  • If dim(operand, quantization_dimension(operand)) = 1 , then scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))) .
• (C7) size(output_dimensions) = rank(result) .
• (C8) is_unique(known_expanding_dimensions + known_nonexpanding_dimensions) .
• (C9) 0 <= known_expanding_dimensions < rank(operand) .
• (C10) 0 <= known_nonexpanding_dimensions < rank(operand) .

Примеры

// %operand: [
//            [1, 2, 3]
//           ]
%operand = stablehlo.constant dense<[[1, 2, 3]]> : tensor<1x3xi64>
%output_dimensions = stablehlo.constant dense<[2, 3, 2]> : tensor<3xi64>
%result = "stablehlo.dynamic_broadcast_in_dim"(%operand, %output_dimensions) {
  broadcast_dimensions = array<i64: 2, 1>,
  known_expanding_dimensions = array<i64: 0>,
  known_nonexpanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

Больше примеров

dynamic_conv

Семантика

This operation is functionally identical to convolution op, but the padding is specified dynamically via padding .

Входы