Дизайн переводчика

Модель данных

Программы StableHLO — это вычисления над тензорами (n-мерными массивами), которые в текущей модели реализованы с помощью класса Tensor . Базовый класс хранения для объекта Tensor , detail::Buffer , хранит mlir::ShapedType тензора вместе с объектом mlir::HeapAsmResourceBlob , представляющим изменяемый блок тензорных данных, расположенный в виде непрерывного массива байтов в порядке от старшего к младшему. заказ . detail::Buffer объекты подсчитываются ссылками для упрощения управления памятью.

Отдельные элементы тензора представлены с помощью класса Element , который использует дискриминируемое объединение, содержащее для хранения один из APInt , APFloat или pair<APFloat,APFloat> . Последний используется для хранения элементов сложных типов.

Tensor имеет следующие API для взаимодействия с отдельными элементами:

  • Element Tensor::get(llvm::ArrayRef<int64_t> index) : для извлечения отдельного тензорного элемента по многомерному индексному index как объекта Element .
  • void Tensor::set(llvm::ArrayRef<int64_t> index, Element element); : Чтобы обновить element объекта Element в тензор с многомерным индексом index .

Как работает переводчик

Функция входа в интерпретатор

SmallVector<Tensor> eval(func::FuncOp func, ArrayRef<Tensor> args);

который делает следующее:

  1. Отслеживает аргументы SSA функции func и связанные с ними значения Tensor времени выполнения, указанные в args , с использованием карты таблицы символов M.
  2. Для каждой операции внутри func в порядке SSACFG:
    • Вызывает eval в операции. Для каждого операнда SSA операции извлеките его значение времени выполнения из M, которое будет предоставлено в качестве аргумента для вызова eval .
    • Отслеживает результат(ы) операции SSA и оцененное значение в M.

eval уровня операции, упомянутая в (2), отвечает за реализацию семантики выполнения операции. Ниже приведен пример stablehlo::AddOp . В этом примере отдельные элементы тензоров lhs и rhs попарно извлекаются как объекты Element , которые затем добавляются. Результат сложения, объект Element , сохраняется в тензоре окончательного result .

Tensor eval(AddOp op, const Tensor &lhs, const Tensor &rhs) {
  Tensor result(op.getType());

  for (auto it = result.index_begin(); it != result.index_end(); ++it)
    result.set(*it, lhs.get(*it) + rhs.get(*it));

  return result;
}

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

Использование интерпретатора для константной свертки

Мы можем использовать механизм интерпретатора для свертывания операций с постоянными значениями операндов. Следующий фрагмент кода демонстрирует идею реализации свертывания stablehlo::AddOp с операндами с плавающей запятой:

OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
  auto attrs = adaptor.getOperands();
  DenseElementsAttr lhsData = dyn_cast<DenseElementsAttr>(attrs[0]);
  DenseElementsAttr rhsData = dyn_cast<DenseElementsAttr>(attrs[1]);
  if (!lhsData || !rhsData) return {};

  auto lhs = Tensor(lhsData);
  auto rhs = Tensor(rhsData);
  auto result = eval(*this, lhs, rhs);

  SmallVector<APFloat> values;
  for (auto i = 0; i < result.getNumElements(); ++i) {
    Element element = result.get(i);
    values.push_back(cast<FloatAttr>(element.getValue()).getValue());
  }

  return DenseElementsAttr::get(result.getType(), values);
}

На данный момент мы не работаем активно над интеграцией интерпретатора в константную свертку, поскольку не планируем реализовывать папку для StableHLO. Однако в будущем мы планируем использовать интерпретатор для свертывания констант в MHLO, после чего мы улучшим эргономику приведенного выше фрагмента кода (например, у нас может быть вспомогательная функция, которая упаковывает константные операнды в объекты Tensor и распаковывает их). Результаты Tensor в OpFoldResult ).

Тестирование интерпретатора StableHLO

Интерпретатор принимает в качестве входных данных (A) программу StableHLO и (B) значения данных, которые должны быть переданы в программу, и генерирует значения выходных данных, которые сопоставляются с ожидаемыми значениями данных, предоставленными пользователем. Значения данных (B) жестко закодированы в самой программе с помощью операций stablehlo.constant . Интерпретатор оценивает входную программу. Выходные данные тестируемой операции проверяются с помощью проверок (например check.expect_eq , check.expect_almost_eq ), как показано ниже. check.expect_eq и check.expect_eq_const проверяют побитовое равенство для любого поддерживаемого типа, а check.expect_almost_eq и check.expect_almost_eq_const проверяют почти равенство в пределах допуска, как описано в руководстве по тестированию (G6), для типов с плавающей запятой и сложных типов.

// CHECK-LABEL: Evaluated results of function: add_op_test_ui4
func.func @add_op_test_ui4() {
  %0 = stablehlo.constant dense<[0, 2]> : tensor<2xui4>
  %1 = stablehlo.constant dense<[15, 3]> : tensor<2xui4>
  %2 = stablehlo.add %0, %1 : tensor<2xui4>
  check.expect_eq_const %2, [15, 5] : tensor<2xui4>
  func.return
}

Тестовая утилита stablehlo-translate --interpret ( code ) отвечает за анализ программы, интерпретируя каждую функцию, включая операции, составляющие функцию. У нас есть специальный набор тестов, состоящий из нескольких тестов, проверяющих различное поведение во время выполнения для каждого StableHLO Op. Тесты можно найти здесь .

Рекомендации по тестированию

(G1) Нужно ли нам проверять все поддерживаемые типы для каждой операции?

Для принятия решения мы можем использовать комбинацию следующих правил:

  1. Если при реализации операции в соответствующей функции eval существует код для обработки определенного типа, обязательно должны быть тесты для покрытия этого типа. Например, для операции add существует эксклюзивный код для обработки целочисленных, логических типов, типов с плавающей запятой и комплексных типов, и, следовательно, нам нужен один тест для каждой категории типов.

  2. Если набор типов обрабатывается единообразно в соответствующей функции eval , то одного теста для всех этих типов должно быть достаточно. Например, для операции add все варианты целочисленных типов ( si4 , u4 , si8 , u8 и т. д.) обрабатываются одинаково с использованием API llvm::APInt , и, следовательно, мы можем пропустить добавление тестов для каждого из этих вариантов. и вместо этого добавьте один репрезентативный тест. Чтобы избежать двусмысленности при выборе представителя, следует руководствоваться следующими рекомендациями:

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

(G2) Как мы принимаем решение о количестве тестов, необходимых для проверки поведения операции?

Цель — всесторонне охватить логику интерпретатора операции (т.е. все крайние случаи реализации) с минимальным количеством тестов. Минимизация количества тестов важна для удобства сопровождения. Чем меньше у нас тестов, тем легче их просмотреть и убедиться, что они всесторонне охватывают операцию. В результате мы ожидаем, что большинство более простых операций будут иметь только один тест. Если по какой-то уважительной причине полный охват нецелесообразен, можно остановиться на уровне >= 90%. Это будет решаться в каждом конкретном случае во время рассмотрения запроса на включение.

(G3) Как насчет добавления тестов для инфраструктуры интерпретатора?

Инфраструктура переводчика в основном проста и может быть добавлена ​​в нашу базу доверия. Единственная нетривиальная часть — это то, как различные типы упаковываются и распаковываются из базового хранилища интерпретатора. Как обсуждалось в (G1), мы будем тестировать только те типы операций, которые обрабатываются по-разному. При этом не исключено, что код упаковки/распаковки, соответствующий различным вариантам целочисленных/плавающих типов, не будет полностью охвачен при тестировании. Чтобы обеспечить полный охват, мы можем выбрать constant типа op, которая поддерживает все типы элементов StableHLO, и написать исчерпывающие тесты.

(G4) Если реализация операции зависит от других операций, должны ли мы писать тесты для последних?

Нет. Например, реализация batch_norm_grad может быть основана на divide , subtract , multiply и других. Нам следует избегать тестирования последних операций при тестировании первых.

(G5) Должны ли мы писать тесты для проверки поведения, определяемого/неопределенного реализации?

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

(G6) При написании тестов для типов с плавающей запятой, с какой точностью необходимо указывать в проверках ожидаемый результат?

Ожидается, что для элементарных операций (сложение, вычитание, умножение, деление и возведение в квадрат) реализация, соответствующая спецификации IEEE, обеспечит округленный результат в пределах 0,5 ULP от математически точного результата. Тем не менее, мы можем с уверенностью предположить, что ожидаемый результат от этих операций будет отличаться не более чем на 1 ULP. Однако это может не работать для трансцендентных функций ( sine , cosine и т. д.), для которых гарантии точности определяются реализацией ( обоснование ).

В текущей реализации используется универсальное значение допуска 0,0001. Следующий пример демонстрирует вышеуказанный допуск в действии.

func.func @check_tolerance() {
  %0 = stablehlo.constant dense<0.2> : tensor<f32>

  // The following check succeeds as %0 is almost equal to the provided
  // constant modulo the tolerance, mentioned above.
  check.expect_almost_eq_const %0, dense<0.19999> : tensor<f32>

  // The following check fails as %0 is not bitwise equal to the provided
  // constant.
  check.expect_eq_const %0, dense<0.19999> : tensor<f32>

  func.return
}

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

(G7) Что-нибудь о стиле кодирования тестов?

  1. Обязательно используйте фактическое имя входов/выходов вместо значений SSA по умолчанию (например, %0, %1 и т. д.).
  2. Убедитесь, что тесты используют красивый печатный формат, если он существует.

(G8) Должны ли мы включить пример, уже представленный в спецификации? Да (для полноты тестирования).