-chlo-legalize-to-stablehlo

Легализует поток операций CHLO в операции StableHLO и Shape.

-shape-legalize-to-stablehlo

Легализовать операции, связанные с фигурами, в StableHLO.

Экспериментальный проход, который легализует операции, связанные с фигурами, для операций StableHLO.

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

-stablehlo-canonicalize-dynamism

Канонизирует динамические операции StableHLO в статические операции.

Заменяет динамические операции StableHLO, такие как DynamicReshapeOp, на соответствующие статические аналоги, такие как DynamicReshapeOp на ReshapeOp или DynamicBroadcastInDim на BroadcastInDim если все динамические элементы этих операций фактически являются константами.

  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>

  ==>

  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>

-stablehlo-check-shape-assertions

_Проверьте операции утверждения stablehlo.custom_call @shape.

Проверка пользовательских вызовов shape_assertion.

Утверждения формы проверяют ограничения динамических измерений в StableHLO. Например, если фреймворку требуется обеспечить ограничение DimA < 2 , можно выдать следующий IR:

%dimA = <get_dimension_size or input arg> : tensor<i32>
%c2 = stablehlo.constant dense<2> : tensor<i32>
%is_lt = stablehlo.compare LT %dimA, %c2 : tensor<i1>
stablehlo.custom_call @shape_assertion(%is_lt) { error_message = "DimA must be less than 2" }

После прохода, если формы верны, stablehlo.custom_call будет удален.

Параметры 

-enable-shape-assertions : Whether shape assertions may generate errors.

-stablehlo-compatibility-expander

Расширитель совместимости для операций StableHLO.

Операции StableHLO обновляются или в последних версиях появляются новые операции. Этот дополнительный проход расширяет обратную совместимость со старыми версиями StableHLO, разбивая новые операции StableHLO на эквивалентные операции, поддерживаемые старыми версиями.

Почему это добровольный пропуск?

Иногда улучшения операций StableHLO используются для значительного упрощения обработки некоторых распространённых шаблонов в экосистеме OpenXLA. К ним относятся, например, TanOp, обладающий высокой поддержкой фреймворков и компиляторов, а также измерения пакетной обработки «сбор/рассеивание», которые могут быть представлены с помощью срезов, но значительно затрудняют шардинг. Для этой категории новых функций мы не предлагаем автоматическое понижение версии, поскольку это может привести к потере важной информации, используемой при последующих оптимизациях. Этот проход можно использовать для расширения этих операций на основе целевой версии, чтобы максимизировать совместимость за счёт потенциально менее оптимальной компиляции.

func.func @tan_op_non_complex(%arg0: tensor<4xf64>) -> tensor<4xf64> {
  %1 = stablehlo.tan %arg0 : tensor<4xf64>
  func.return %1 : tensor<4xf64>
}

==>

func.func @tan_op_non_complex(%arg0: tensor<4xf64>) -> tensor<4xf64> {
  %0 = stablehlo.sine %arg0 : tensor<4xf64>
  %1 = stablehlo.cosine %arg0 : tensor<4xf64>
  %2 = stablehlo.divide %0, %1 : tensor<4xf64>
  return %2 : tensor<4xf64>
}

Параметры 

-target : The target version. Must be a version of the form #.#.#.

-stablehlo-complex-math-expander

Расширитель для сложных математических операций StableHLO.

Комплексные математические операции StableHLO представляют собой разложения с использованием действительных математических операций StableHLO.

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

func.func @sqrt_op_complex(%arg0: tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>> {
  %1 = stablehlo.sqrt %arg0 : tensor<4xcomplex<f64>>
  func.return %1 : tensor<4xcomplex<f64>>
}

==>

func.func @sqrt_op_complex(%arg0: tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>> {
  TBD
  return %2 : tensor<4xcomplex<f64>>
}

-stablehlo-convert-to-signless

Проходим преобразование IR в целые числа без знака.

-stablehlo-legalize-composite-to-call

Заменяет составные операции вызовом их декомпозиции.

Заменяет составные операции вызовом их декомпозиции, например, как показано ниже:

stablehlo.composite "my_namespace.my_op" %arg0, %arg1 {
  decomposition = @bar,
  version = 1,
  composite_attributes = {
    "my_attribute": "my_value"
  }
}

Станет:

func.call @bar(%arg0, %arg1)

Подмножество композитов можно исключить из этого преобразования с помощью флага «исключение», например:

stablehlo-opt --stablehlo-legalize-composite-to-call=except='foo.baz,foo.qux'

Параметры 

-except : Names of composites that should not be replaced with calls.

-stablehlo-legalize-deprecated-ops

Легализовать устаревшие операции, превратив их в хорошо поддерживаемые.

В RFC по устареванию наборов операций StableHLO v1.0 (#2283) предлагается удалить несколько избыточных операций. Этот этап помогает оценить влияние удаления этих операций на различные конвейеры компиляции, легализуя их для поддерживаемых в долгосрочной перспективе аналогов.

Параметры 

-fail-on-unused : Fail on (mostly) unused ops that are deprecated without any fallback.

-stablehlo-legalize-qdq-to-quantized-op

Объединить шаблон (деквантование, операция с плавающей точкой и квантование) в квантованную операцию StableHLO

Объединить шаблон (деквантование, операция с плавающей точкой и квантование) в квантованную операцию StableHLO. Примечание: проход не удаляет никакие ранее существовавшие операции. Например, следующая программа 

func.func @add(%arg0: tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>) -> tensor<16x16xf32>
  %1 = stablehlo.abs %0 : tensor<16x16xf32>
  %2 = stablehlo.uniform_quantize %1 : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>
  func.return %2 : tensor<16x16x!quant.uniform<ui8:f32, 34.0:16>>
}

Станет: 

func.func @add(%arg0: tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>) -> tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>) -> tensor<16x16xf32>
  %1 = stablehlo.abs %0 : tensor<16x16xf32>
  %2 = stablehlo.abs %arg0 : tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
  %3 = stablehlo.uniform_quantize %1 : (tensor<16x16xf32>) -> tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
  return %2 : tensor<16x16x!quant.uniform<u8:f32, 3.400000e+01:16>>
}

-stablehlo-legalize-quant-to-math

Преобразование квантованных операций StableHLO в примитивные математические операции StableHLO.

Преобразовать программы StableHLO, использующие типы UniformQuantized, в семантически эквивалентные целочисленные математические операции. 

func.func @add(%arg0: tensor<!quant.uniform<i8:f32,1.0:0>>, %arg1: tensor<!quant.uniform<i8:f32,2.0:1>>) ->  tensor<!quant.uniform<i8:f32,3.0:2>> {
  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32,1.0:0>>, tensor<!quant.uniform<i8:f32,2.0:1>>) -> tensor<!quant.uniform<i8:f32,3.0:2>>
  func.return %0 : tensor<!quant.uniform<i8:f32,3.0:2>>
}

Станет: 

func.func @add(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
  %0 = stablehlo.convert %arg0 : (tensor<i8>) -> tensor<f32>
  %cst = stablehlo.constant dense<0.333333343> : tensor<f32>
  %1 = chlo.broadcast_multiply %0, %cst : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %cst_0 = stablehlo.constant dense<2.000000e+00> : tensor<f32>
  %2 = chlo.broadcast_add %1, %cst_0 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %3 = stablehlo.round_nearest_even %2 : tensor<f32>
  %4 = stablehlo.convert %3 : (tensor<f32>) -> tensor<i32>
  %5 = stablehlo.convert %arg1 : (tensor<i8>) -> tensor<f32>
  %cst_1 = stablehlo.constant dense<0.666666686> : tensor<f32>
  %6 = chlo.broadcast_multiply %5, %cst_1 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %cst_2 = stablehlo.constant dense<1.33333337> : tensor<f32>
  %7 = chlo.broadcast_add %6, %cst_2 : (tensor<f32>, tensor<f32>) -> tensor<f32>
  %8 = stablehlo.round_nearest_even %7 : tensor<f32>
  %9 = stablehlo.convert %8 : (tensor<f32>) -> tensor<i32>
  %c = stablehlo.constant dense<2> : tensor<i32>
  %10 = chlo.broadcast_add %4, %9 : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %11 = chlo.broadcast_subtract %10, %c : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %c_3 = stablehlo.constant dense<-128> : tensor<i32>
  %c_4 = stablehlo.constant dense<127> : tensor<i32>
  %12 = stablehlo.clamp %c_3, %11, %c_4 : tensor<i32>
  %13 = stablehlo.convert %12 : (tensor<i32>) -> tensor<i8>
  return %13 : tensor<i8>
}

-stablehlo-legalize-quantized-op-to-qdq

Разложить квантованную операцию StableHLO на шаблон (деквантование, операция с плавающей точкой и квантование).

Разложите квантованные программы StableHLO, используя операции равномерного квантования/деквантования. Например, следующая программа: 

func.func @add(%arg0: tensor<!quant.uniform<i8:f32,1.0:0>>, %arg1: tensor<!quant.uniform<i8:f32,2.0:1>>) ->  tensor<!quant.uniform<i8:f32,3.0:2>> {
  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32,1.0:0>>, tensor<!quant.uniform<i8:f32,2.0:1>>) -> tensor<!quant.uniform<i8:f32,3.0:2>>
  func.return %0 : tensor<!quant.uniform<i8:f32,3.0:2>>
}

Станет: 

func.func @add(%arg0: tensor<!quant.uniform<i8:f32, 1.000000e+00>>, %arg1: tensor<!quant.uniform<i8:f32, 2.000000e+00:1>>) -> tensor<!quant.uniform<i8:f32, 3.000000e+00:2>> {
  %0 = stablehlo.uniform_dequantize %arg0 : (tensor<!quant.uniform<i8:f32, 1.000000e+00>>) -> tensor<f32>
  %1 = stablehlo.uniform_dequantize %arg1 : (tensor<!quant.uniform<i8:f32, 2.000000e+00:1>>) -> tensor<f32>
  %2 = stablehlo.add %0, %1 : tensor<f32>
  %3 = stablehlo.uniform_quantize %2 : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 3.000000e+00:2>>
  return %3 : tensor<!quant.uniform<i8:f32, 3.000000e+00:2>>
}

-stablehlo-legalize-to-vhlo

Легализовать StableHLO в VHLO.

Легализуйте StableHLO до последней версии операторов в VHLO. Эти операторы затем можно понизить до более старых версий VHLO для обеспечения прямой совместимости с помощью VhloToVersionPass . 

stablehlo.exponential %[[ARG0]] <{result_accuracy = DEFAULT}> : tensor<f32>
# ====>
"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.DEFAULT_v1}> : !vhlo.tensor_v1<!vhlo.f32_v1>

Подробную информацию об использовании VHLO для сохранения прямой и обратной совместимости см. на сайте vhlo.md > Диалект VHLO .

Параметры 

-allow-other-dialects : Allow serialization to use other (potentially unstable) dialects, inserts unrealized casts between dialects.

-stablehlo-refine-arguments

Уточняет формы аргументов основной функции.

Изменяет аргументы основной функции, используя сигнатуру типа входных данных. Оборачивает аргументы в custom_call @stablehlo.shape_refinement_operand_wrapper чтобы сохранить корректность IR до запуска уточнения формы. 

func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  ...
}

==>

func.func public @main(%arg0: tensor<16xf32>) -> tensor<?xf32> {
  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {...}
    : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
}

Параметр refinedTypesOption можно использовать для указания списка уточняемых типов. Его можно указать в MLIR с помощью --types='tensor<...>,tensor<...>' или передать методу pass create. Список уточняемых типов должен указывать тип каждого аргумента main метода.

Параметры 

-types : The new types to be used for the main function's arguments, specified as an MLIR TypeRange 'tensor<1x2xf32>, ...'

-stablehlo-refine-shapes

Уточняет формы в программе StableHLO.

Проходит через программу StableHLO, уточняющую формы в операциях.

Флагманский вариант использования этого прохода — специализация программ с динамическими формами на статические формы. Если программа StableHLO с динамическими формами имеет правильную структуру, то обновление типов её аргументов с динамических форм на статические и запуск этого прохода приведут к распространению статических форм по всей программе.

Этот проход удаляет custom_call @shape_refinement_operand_wrapper заменяя использование результата непосредственно операндом, и распространяет статические фигуры по всей программе. 

  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {...}
      : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  %1 = stablehlo.add %0, %0 : tensor<?xf32>

  ==>

  %1 = stablehlo.add %arg0, %arg0 : tensor<16xf32>

Модули, подходящие для уточнения формы, должны иметь следующие свойства:

  • Все динамические формы зависят только от входных форм (без зависимости формы от содержимого входного массива). Мы имеем в виду операции, которые транзитивно зависят только от входных форм (например, как задано функцией stablehlo.get_dimension_size ) или глобальных констант, таких как разрешённые значения символьных целых чисел (например, тензорные). : A = 5) как операции dimension . Все значения размерностей могут быть преобразованы в константы посредством межпроцедурного свёртывания констант.
  • Промежуточные функции могут принимать несколько аргументов-токенов (типа !stablehlo.token) в начале списка аргументов, за которыми следуют некоторые глобальные константные аргументы, которые являются постоянными целыми скалярами, такими как разрешенные значения символических целых чисел (т. е. тензорные : А = 5).
  • Некоторые промежуточные функции могут возвращать вычисления с глобальными константами, например, floordiv для значений symint. Эти функции возвращают только константные значения после уточнения. Эти функции являются встроенными.
  • Все вызовы одной функции разрешаются в одни и те же формы аргументов, и никакие рекурсивные/ко-рекурсивные вызовы функций не производятся.

-stablehlo-wrap-in-composite

Оборачивает несоставную операцию StableHLO в составную операцию.

Оборачивает операции StableHLO в операции stablehlo.composite .

Например, рассмотрим простую программу StableHLO: 

func.func @main(%arg0 : tensor<2xf32>, %arg1 : tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
  return %0 : tensor<2xf32>
}

Применение этого прохода для оболочки операций stablehlo.add приведет к следующей программе: 

func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.composite "stablehlo.add" %arg0, %arg1 {decomposition = @stablehlo.add.impl} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
  return %0 : tensor<2xf32>
}
func.func private @stablehlo.add.impl(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
  %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
  return %0 : tensor<2xf32>
}

Примечания:

  • Атрибут name сгенерированной операции stablehlo.composite всегда будет таким же, как имя исходной операции, которая была упакована (например, если вы оборачиваете операцию stablehlo.add , составная операция также будет называться "stablehlo.add" ).
  • Закрытая функция, инкапсулирующая исходную операцию (на которую ссылается атрибут decomposition операции stablehlo.composite ), будет названа с использованием шаблона <op_name>.impl[.N] , где <op_name> — имя исходной операции, а N — уникальный целочисленный идентификатор, сгенерированный для предотвращения конфликтов имен внутри модуля.

Этот пропуск можно использовать двумя способами:

Режим 1: использование командной строки

Этот режим предназначен для отладки или тестирования, поскольку обеспечивает минимальный контроль над атрибутами создаваемых операций stablehlo.composite . Он охватывает все экземпляры операций, заданные с помощью параметров op-names (списка имён операций, разделённых запятыми). Атрибуты вновь созданной операции stablehlo.composite будут такими же, как и атрибуты исходной операции.

Пример использования: 

stablehlo-opt input.mlir --stablehlo-wrap-in-composite=op-names='stablehlo.add,stablehlo.mul' -o output.mlir

Режим 2: Программное модульное обертывание с настраиваемой обработкой атрибутов

Этот режим расширяет программную обёртку на весь модуль, обеспечивая детальный контроль над обёртываемыми операциями и их атрибутами. Это достигается с помощью API createStablehloWrapInCompositePass , принимающего в качестве аргумента CompositeAttributeProviderMap .

CompositeAttributeProviderMap — это карта, которая определяет, какие операции следует рассматривать для обёртывания и как следует обрабатывать их атрибуты. Её семантика следующая:

  • Ключи (mlir::TypeID): TypeID операции MLIR. Если TypeID операции совпадает с ключом в карте, она становится кандидатом на перенос.
  • Значения (лямбда-функции): лямбда-функция типа std::function<std::optional<NamedAttrList>(Operation*)> . Эта функция применяется к каждой операции-кандидату.
    • Ввод: mlir::Operation* , который является экземпляром типа операции, соответствующего ключу TypeID .
    • Возвращаемое значение: std::optional<NamedAttrList> .
      • Если лямбда возвращает NamedAttrList (обёрнутый в std::optional ), операция оборачивается в операцию stablehlo::composite , а возвращаемые атрибуты используются для установки атрибутов композита.
      • Если лямбда-функция возвращает std::nullopt , операция не переносится. Это позволяет осуществлять выборочную пересылку на основе пользовательских критериев.

Пример (С++): 


stablehlo::CompositeAttributeProviderMap compositeAttributeProviderMap;

compositeAttributeProviderMap[mlir::TypeID::get<mlir::stablehlo::AddOp>()] =
  [](mlir::Operation* op) -> std::optional<mlir::NamedAttrList> {
  // Custom logic to determine if and how to wrap the operation.
  // Example: Only wrap if it's on a specific type.
  if (mlir::isa<mlir::Float32Type>(op->getOperand(0).getType())) {
    return mlir::NamedAttrList(op->getAttrs());
  }
  return std::nullopt; // Do not wrap.
};

pm.addPass(createStablehloWrapInCompositePass(compositeAttributeProviderMap, compositeVersion));
if (mlir::failed(pm.run(module))) {
  return;
}

Параметры 

-op-names : The names of the ops to wrap.
-version  : The version number of the composite op.

-vhlo-legalize-to-stablehlo

Легализовать VHLO в StableHLO.

-vhlo-to-version

Конвертируйте версии VHLO для совместимости.

Конвертирует версии VHLO между собой для повышения и понижения версии IR с целью сохранения прямой и обратной совместимости. 

"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = DEFAULT}>
# ==( -target=1.0.0 )==>
"vhlo.exponential_v1"(%[[ARG0]])
# ==( -target=1.9.0 )==>
"vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = DEFAULT}>

Подробную информацию об использовании VHLO для сохранения прямой и обратной совместимости см. на сайте vhlo.md > Диалект VHLO .

Параметры 

-target : The target version. Must be a version of the form #.#.# .