Эта страница переведена с помощью Cloud Translation API.

Динамизм в StableHLO

Текущее состояние динамизма более формально изложено в Dynamism RFC . На этой странице будет представлен общий обзор RFC и обсуждены важные API и инструменты для взаимодействия с динамическими программами.

Терминология динамизма и обзор поддержки

Во-первых, рассмотрим несколько терминов, которые появятся в этом документе, а также краткое описание их поддержки в StableHLO:

Динамические размеры

Динамические размеры относятся к любому измерению, размер которого неизвестен. В StableHLO мы представляем динамические размеры, используя ? , т.е. tensor<16x?xf32> .

Ограниченный динамизм

Ограниченный динамизм относится к динамическому измерению, значение которого имеет известную верхнюю границу. Обычно это полезно для заполнения тензора во время выполнения. В StableHLO мы представляем ограниченный динамизм, используя #stablehlo.bounds в качестве тензорного кодирования, т.е. тензор ранга 2 с одним динамическим измерением, ограниченным значением 16, а другое без границы, может быть представлен как tensor<?x?xf32, #stablehlo.bounds<16, ?>> .

StableHLO способен отображать ограниченный динамизм, но имеет ограниченную поддержку платформы, начиная с TensorFlow и с некоторой поддержкой в PyTorch/XLA.

Безграничный динамизм

Неограниченный динамизм, как следует из названия, относится к динамическому измерению, размер которого не известен. Этот тип динамизма очень распространен в StableHLO с поддержкой JAX, PyTorch/XLA и TF, часто используемой для экспорта моделей с динамическим размером пакета или длиной последовательности.

В StableHLO мы просто исключаем кодирование границ для этой формы динамизма, т.е. tensor<?x?xf32> .

Полиморфизм формы

Полиморфизм формы — это термин, который мы унаследовали от JAX .

Есть два ключевых последствия полиморфизма формы:

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

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

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

Зависимый от данных динамизм

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

Многие динамические операции, зависящие от данных, можно смоделировать с использованием ограниченного динамизма, где указывается верхняя граница размера тензора, и аппаратное обеспечение обычно реализует это посредством заполнения тензора. Сегодня существует некоторая поддержка динамизма, зависящего от данных, в PyTorch/XLA и TensorFlow, но JAX в настоящее время не отслеживает операции, которые приводят к динамизму, зависящему от данных.

Экспорт программ с динамическими размерами

См. наши руководства по StableHLO для получения информации о том, как экспортировать программы с динамическими размерами пакетов или длиной последовательности:

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

Удаление динамического прохода конвейера

Есть несколько полезных проходов для уточнения фигур. Все они удобно объединены в конвейер проходов createStablehloRemoveDynamismPipeline :

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

Индивидуальные проходы для совершенствования динамизма

По отдельности проходы, которые могут оказаться полезными для уточнения формы:

stablehlo-refine-arguments для замены входных аргументов конкретными типами тензоров.
stablehlo-refine-shapes для распространения информации о форме нового входного аргумента по всей программе.
stablehlo-canonicalize-dynamism для замены динамических операций их статическими вариантами.

Актуальную информацию и примеры смотрите в связанной документации.

Пример: Чем полезен динамизм и как я могу его использовать?

Динамизм имеет множество применений, здесь мы в основном сосредоточимся на распространенном варианте использования полиморфизма формы — создании гибкого представления экспортированной модели, обычно используемого для представления динамического размера пакета или длины последовательности.

Статическая модель add_one

Чтобы продемонстрировать это, мы воспользуемся следующей простой моделью add_one :

def add_one(x):
  return x + 1

При трассировке с использованием tensor<4xf32> мы получим следующую программу StableHLO:

// File: add_one.mlir
func.func @add_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<4xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<4xf32>
  return %0 : tensor<4xf32>
}

Эта модель будет работать только для входных аргументов, имеющих форму tensor<4xf32> . Если бы мы когда-либо изменили размер пакета или длину последовательности, нам пришлось бы заново отследить исходный код и снова перейти к StableHLO, и нет никакой гарантии, что у нас все еще будет доступ к исходному коду!

Динамическая модель add_one

Именно здесь в игру вступает полиморфный динамизм формы. Вместо этого JAX и PyTorch/XLA могут создавать модель add_one с динамически допустимым IR, которая будет транслировать константу в соответствии с формой динамического ввода следующим образом:

// File: add_one_dynamic.mlir
func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  %cst = stablehlo.constant dense<1.0> : tensor<f32>
  %0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor<?xf32>) -> tensor<i32>
  %1 = stablehlo.reshape %0 : (tensor<i32>) -> tensor<1xi32>
  %2 = stablehlo.dynamic_broadcast_in_dim %cst, %1, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %3 = stablehlo.add %arg0, %2 : tensor<?xf32>
  return %3 : tensor<?xf32>
}

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

Уточнение динамической модели

Например, следующий порядок проходов может полностью усовершенствовать эту программу:

stablehlo-opt add_one_dynamic.mlir \
  --stablehlo-refine-arguments='types=tensor<16xf32>' \
  --stablehlo-refine-shapes \
  --stablehlo-canonicalize-dynamism

Поэтапно программа трансформируется следующим образом:

// After stablehlo-refine-arguments: Inputs updated, shapes not propagated
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) {indices_of_shape_operands = dense<1> : tensor<1xi64>} : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
  %3 = stablehlo.dynamic_broadcast_in_dim %cst, %2, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %4 = stablehlo.add %0, %3 : tensor<?xf32>
  return %4 : tensor<?xf32>
}

// After stablehlo-refine-shapes: Shapes propagated, dynamic ops still exist
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// After stablehlo-canonicalize-dynamism: Dynamic ops replaced with static ops
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// (Bonus) Use ` --stablehlo-aggressive-simplification` pass to canonicalize the
// constant broadcast, leaving us with the original static program in this case.
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<16xf32>
  return %0 : tensor<16xf32>
}