Design do intérprete

Modelo de dados

Programas StableHLO são cálculos sobre tensores (matrizes n), que, no modelo atual, são implementadas usando o classe Tensor. A classe de armazenamento subjacente de um objeto Tensor, detail::Buffer, armazena o mlir::ShapedType do tensor junto com uma Objeto mlir::HeapAsmResourceBlob que representa um blob mutável de tensor os dados são dispostos como uma matriz de bytes contígua ordem maior para menor. Os objetos detail::Buffer são contados por referência para simplificar o gerenciamento de memória.

Elementos individuais de um tensor são representados pela classe Element, que usa uma união discriminada contendo um de APInt, APFloat ou pair<APFloat,APFloat> para armazenamento. O último é usado para armazenar elementos com tipos complexos.

A Tensor tem as seguintes APIs para interagir com os elementos individuais:

  • Element Tensor::get(llvm::ArrayRef<int64_t> index): para extrair uma elemento do tensor individual no índice multidimensional index como Element objeto.
  • void Tensor::set(llvm::ArrayRef<int64_t> index, Element element);: Para atualizar um objeto Element element em um tensor multidimensional índice index.

Como o intérprete funciona

A função de entrada para o intérprete é

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

que:

  1. Monitora os argumentos SSA de func e o ambiente de execução associado Tensor fornecidos em args, usando um mapa de tabela de símbolos, M.
  2. Para cada operação em func, em ordem de SSACFG:
    • Invoca eval na operação. Para cada operando SSA da operação, extraia o que o valor do ambiente de execução de M seja fornecido como um argumento para a invocação eval.
    • Rastreia os resultados de SSA da operação e o valor avaliado em M.

O eval de nível operacional mencionado em (2) é responsável pela implementação da e a semântica de execução da operação. Veja a seguir um exemplo de stablehlo::AddOp. No exemplo, os elementos individuais dos tensores lhs e rhs são pares extraídos como objetos Element, que vão ser adicionados. O resultado da adição, um objeto Element, é armazenado no tensor result final.

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;
}

No geral, o design do intérprete é otimizado para facilitar a leitura implementações de funções eval para operações individuais porque tem como objetivo e serve como implementação de referência para o StableHLO. Por exemplo, em vez de definindo eval como uma função de modelo e parametrizando-a com tipos de elementos; encapsulamos detalhes sobre como os diferentes tipos de elementos são tratados Element::operator+ etc., simplificando a implementação de eval.

Como usar o intérprete para dobrar constante

Podemos usar o mecanismo intérprete para dobrar operações com operando constante e a distribuição dos valores dos dados. O snippet de código a seguir demonstra uma ideia da implementação para dobrar stablehlo::AddOp com operandos de tipo de ponto flutuante:

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);
}

No momento, não estamos trabalhando ativamente na integração do intérprete dobrável constante porque não estamos planejando implementar pastas para StableHLO. No entanto, no futuro, planejamos usar o intérprete para em MHLO, melhorando a ergonomia do snippet de código. acima (por exemplo, poderíamos ter uma função auxiliar que empacota operandos constantes em Tensor e descompacta resultados de Tensor em OpFoldResult).

Como testar o interpretador de StableHLO

O intérprete usa como entradas (A) um programa StableHLO e (B) valores de dados para alimenta o programa e gera valores de dados de saída, que são em relação aos valores dos dados esperados informados pelo usuário. Os valores de dados (B) são codificado no próprio programa usando operações stablehlo.constant. A intérprete avalia o programa de entrada. Saídas da operação em teste é conferido por verificações (por exemplo, check.expect_eq, check.expect_almost_eq), conforme mostradas abaixo. check.expect_eq e check.expect_eq_const verificam bit a bit a bit igualdade para qualquer tipo com suporte, e check.expect_almost_eq e check.expect_almost_eq_const verifica a quase igualdade dentro de uma tolerância, explicado nas diretrizes de teste (G6), para tipos complexos e de ponto flutuante.

// 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
}

Um utilitário de teste stablehlo-translate --interpret (código) é responsável por analisar o programa, interpretando cada função, incluindo a operações que constituem a função. Temos um pacote de testes dedicado, de vários testes que executam diversos comportamentos de tempo de execução para cada operação StableHLO. Os testes estão disponíveis aqui.

Orientações para teste

(G1) Precisamos testar todos os tipos com suporte em todas as operações?

Podemos usar uma combinação das seguintes regras para decidir:

  1. Ao implementar uma operação, se houver um código no eval correspondente para processar um tipo específico, é essencial que os testes sejam para cobrir esse tipo. Por exemplo, para a operação add, há um código exclusivo para lidar com tipos inteiros, booleanos, flutuantes e complexos. Portanto, precisa de um teste para cada categoria de tipos.

  2. Se um conjunto de tipos for processado de maneira uniforme na função eval correspondente, então um único teste para todos esses tipos será suficiente. Por exemplo, para a operação add, todas as variantes de tipos inteiros (si4, u4, si8, u8) e assim por diante) são tratadas da mesma forma usando APIs llvm::APInt e, portanto, podemos pular adicionar testes para cada uma dessas variantes e, em vez disso, adicionar uma teste representativo. Para evitar ambiguidade na seleção do representante, deve seguir as seguintes diretrizes:

    • Se todos os tipos, tratados de maneira uniforme, tiverem o mesmo tipo primitivo (ou seja, se todos forem tipos inteiros, de ponto flutuante ou complexos), escolha aquele com a largura máxima de bits.
    • Se todos os tipos, tratados de maneira uniforme, tiverem uma mistura de tipos primitivos, escolha aquela com o seguinte tipo primitivo, em ordem decrescente de preferência: inteiro, ponto flutuante, booleano, complexo.

(G2) Como decidimos o número de testes necessários para cobrir uma operação e comportamento?

O objetivo é abordar de forma abrangente a lógica do intérprete para a operação (ou seja, todos os casos isolados da implementação) com um número mínimo de testes. Minimizar o número de testes é importante para manter a capacidade de manutenção. Quanto menos testes nós temos, mais fácil será revisá-los e certificar-se de que eles abordaremos a operação de forma abrangente. Como resultado, esperamos que a maior parte das ops terão apenas um teste. Se, por algum bom motivo, forem abrangentes for impraticável, convém parar em >= 90%. Isto será decidido caso a caso durante a análise da solicitação de envio.

(G3) E quanto à adição de testes para a infraestrutura do intérprete?

A infraestrutura do intérprete é basicamente simples e pode ser adicionada nossa base de confiança. A única parte não trivial é como vários tipos são empacotados em e descompactados do armazenamento do intérprete. Como discutimos em (G1), testaremos apenas os tipos de operações que são tratados de forma diferente. Com que o código de empacotamento/descompactação, correspondente a diferentes variantes de tipos inteiros/de ponto flutuante, podem não ser totalmente cobertas testes. Para garantir a cobertura completa, podemos escolher uma operação como constant que oferece suporte a todos os tipos de elementos StableHLO e cria testes completos.

(G4) Se a implementação de uma operação depender de outras operações, será que testes para o último caso?

Não. Por exemplo, a implementação de batch_norm_grad pode se basear no divide, subtract, multiply e outros. Devemos evitar testar a última opção ops ao testar o primeiro.

(G5) Devemos criar testes para exercer a função definida pela implementação / indefinida o comportamento deles?

Não devemos criar testes que tenham o comportamento definido ou comportamentos indefinidos da operação. Testa o exercício de comportamentos definidos pela implementação demonstrar um comportamento local do intérprete que não deve ser ser generalizada. Testes que exercem comportamento indefinido não contribuem para a compreensão do comportamento das operações.

(G6) Ao escrever testes para tipos de ponto flutuante, com que precisão o o resultado esperado precisa ser especificado nas verificações?

Para operações elementares (adição, subtração, multiplicação, divisão e quadrado), uma implementação que segue a especificação IEEE deve fornecer uma resultado arredondado em 0,5 ULP do resultado matematicamente exato. Dito isso, podemos imaginar com segurança que o resultado esperado dessas operações seja com o máximo de 1 ULP de diferença. No entanto, isso pode não funcionar para funções transcendentais (sine, cosine etc.) para os quais as garantias de precisão são definido pela implementação (justificativa).

A implementação atual usa uma abordagem com uma tolerância de 0,0001. O exemplo a seguir demonstra a tolerância acima em ação.

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
}

Essa é apenas a primeira etapa no teste da acurácia numérica das operações StableHLO. No momento, esta é uma área subespecificada da especificação StableHLO, e há trabalho contínuo para descobrir isso #1156 com base na nossa experiência com o StableHLO na prática e no feedback de com as partes interessadas. Vamos atualizar a infraestrutura à medida que o processo continuar de maneira adequada.

(G7) Algo sobre o estilo de programação dos testes?

  1. Use o nome real das entradas/saídas em vez de padronizar. para valores SSA (por exemplo, %0, %1 etc.)
  2. Certifique-se de que os testes usem um formato bem impresso, se houver.

(G8) Devemos incluir o exemplo já fornecido na especificação? Sim (para concluir o teste).