Diseño para intérpretes

Modelo de datos

Los programas StableHLO son cálculos sobre tensores (arrays n-dimensionales), que, en el modelo actual, se implementan usando el clase Tensor. La clase de almacenamiento subyacente de un objeto Tensor detail::Buffer, almacena el mlir::ShapedType del tensor junto con un Un objeto mlir::HeapAsmResourceBlob que representa un BLOB mutable de tensor los datos dispuestos como un array de bytes contiguos en orden mayor a menor. Los objetos detail::Buffer se cuentan como referencia para simplificar la administración de la memoria.

Los elementos individuales de un tensor se representan con la clase Element, que usa una unión discriminada que contenga uno de estos valores: APInt, APFloat o pair<APFloat,APFloat> para el almacenamiento. El último se usa para almacenar elementos con tipos complejos.

Tensor tiene las siguientes APIs para interactuar con sus elementos individuales:

  • Element Tensor::get(llvm::ArrayRef<int64_t> index): Para extraer un elemento tensor individual en el índice multidimensional index como Element .
  • void Tensor::set(llvm::ArrayRef<int64_t> index, Element element);: Para actualizar un objeto Element element en un tensor multidimensional índice index.

Cómo funciona el intérprete

La función de entrada al intérprete es

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

que hace lo siguiente:

  1. Realiza un seguimiento de los argumentos de SSA de func y su tiempo de ejecución asociado Tensor. valores proporcionados en args, con un mapa de tabla de símbolos, M.
  2. Para cada operación de func, en orden SSACFG, haz lo siguiente:
    • Invoca a eval en la operación. Para cada operando SSA de la op, extrae su el valor del entorno de ejecución de M se proporcionará como un argumento en la invocación de eval.
    • Realiza un seguimiento de los resultados de SSA de la operación y el valor evaluado en M.

El eval al nivel de la operación mencionado en (2) es responsable de implementar las de ejecución de la op. A continuación, se incluye un ejemplo de stablehlo::AddOp. En el ejemplo, los elementos individuales de los tensores lhs y rhs se usan en pares se extraen como objetos Element que luego se agregan. El resultado de la suma, un objeto Element, se almacena en el tensor final 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;
}

En general, el diseño del intérprete está optimizado para facilitar la lectura de implementaciones de funciones eval para operaciones individuales y sirve como implementación de referencia para StableHLO. Por ejemplo, en lugar de definir eval como una función de plantilla y parametrizar con tipos de elementos encapsulamos detalles sobre cómo se manejan los diferentes tipos de elementos Element::operator+, etc., lo que simplifica la implementación de eval

Cómo usar el intérprete para el plegado constante

Podemos usar el mecanismo de intérprete para plegar las operaciones con operandos constantes de salida. En el siguiente fragmento de código, se demuestra una idea de la implementación para plegar stablehlo::AddOp con operandos de tipo de punto flotante:

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

En este momento, no estamos trabajando activamente para integrar al intérprete en plegado constante porque no planeamos implementar una carpeta para StableHLO. Sin embargo, en el futuro, planeamos aprovechar el intérprete para el plegado en MHLO. En ese momento, mejoraremos la ergonomía del fragmento de código de arriba (p.ej., podríamos tener una función auxiliar que incluya operandos de constantes Tensor y descomprime los resultados de Tensor en OpFoldResult).

Prueba el intérprete de StableHLO

El intérprete toma como entradas (A) un programa StableHLO y (B) valores de datos para datos de salida al programa y genera valores de datos de salida, que coinciden con los valores de datos esperados que proporcionó el usuario. Los valores de datos (B) son codificarse de forma fija en el propio programa mediante operaciones stablehlo.constant. El intérprete evalúa el programa de entrada. Los resultados de la operación sometida a prueba se verifica mediante verificaciones (p.ej., check.expect_eq, check.expect_almost_eq), como que se muestra a continuación. check.expect_eq y check.expect_eq_const comprueban la velocidad de bits igualdad para cualquier tipo admitido, y check.expect_almost_eq, y check.expect_almost_eq_const verifican si hay casi igualdad dentro de una tolerancia, se explica en la guía de pruebas (G6) para el punto flotante y los tipos complejos.

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

Una utilidad de prueba stablehlo-translate --interpret (código) es responsable de analizar el programa, interpretar cada función, incluida la las operaciones que constituyen la función. Tenemos un paquete de pruebas exclusivo, que consta de varias pruebas que ejecutan diversos comportamientos de tiempo de ejecución para cada Op. de StableHLO. Puedes encontrar las pruebas aquí.

Pautas para pruebas

(G1) ¿Necesitamos probar todos los tipos admitidos en cada operación?

Podemos usar una combinación de las siguientes reglas para decidir:

  1. Cuando se implementa una op, si existe código en el eval correspondiente para manejar un tipo particular, es imperativo que se realicen pruebas para abordar ese tipo. Por ejemplo, para la op add, hay código exclusivo para manejar tipos de números enteros, booleanos, de punto flotante y complejos. Por lo tanto, necesitas una prueba para cada categoría de tipos.

  2. Si un conjunto de tipos se controla de manera uniforme en la función eval correspondiente, entonces una sola prueba para todos esos tipos debería ser suficiente. A modo de ejemplo, Para la op add, todas las variantes de tipos de números enteros (si4, u4, si8 y u8) etc.) se manejan de la misma manera con APIs de llvm::APInt, por lo que podemos omitir agregar pruebas para cada una de esas variantes y, en su lugar, agregar una prueba representativa. Para evitar ambigüedades en la selección del representante, debe seguir estos lineamientos:

    • Si todos los tipos, que se manejan de manera uniforme, tienen el mismo tipo primitivo. (es decir, si todos son números enteros, de punto flotante o de tipos complejos), entonces elige el que tenga el ancho de bits máximo.
    • Si todos los tipos, que se manejan de manera uniforme, tienen una combinación de tipos primitivos, entonces elige uno que tenga el siguiente tipo primitivo, en orden descendente de preferencia: número entero, punto flotante, booleano, complejo.

(G2) ¿Cómo decidimos la cantidad de pruebas necesarias para cubrir una operación el comportamiento de los usuarios?

El objetivo es abarcar por completo la lógica del intérprete para la operación. (es decir, todos los casos límite de la implementación) con una cantidad mínima de pruebas. Minimizar la cantidad de pruebas es importante para el mantenimiento. Cuantas menos pruebas que tenemos, más fácil será revisarlos y asegurarnos de que a cubrir toda la op. Como resultado, esperamos que la mayoría de los procesos ops terminarán con una sola prueba. Si, por un buen motivo, abarca la cobertura no es práctica, entonces está bien que se detenga en >= 90%. Esto se decidirá caso por caso durante la revisión de la solicitud de extracción.

(G3) ¿Quieres agregar pruebas para la infraestructura de intérpretes?

La infraestructura del intérprete es mayormente sencilla y se puede agregar a nuestra base de confianza. La única parte no trivial es cómo se empaquetan los distintos tipos en y se descomprimen del almacenamiento de intérpretes subyacente. Como se explicó en (G1), probarán solo los tipos de operaciones que se manejan de manera diferente. Con que es posible que el código de empaquetado/desempaquetado, que corresponde a los diferentes de números enteros o de punto flotante, es posible que no se cubran por completo y pruebas. Para garantizar una cobertura completa, podemos elegir una op como constant que admite todos los tipos de elementos StableHLO y escribe pruebas exhaustivas.

(G4) Si la implementación de una op depende de otras, ¿deberíamos escribir para las últimas pruebas?

No. Por ejemplo, la implementación de batch_norm_grad se puede basar en divide, subtract, multiply y otros. No debemos probar esta última opción ops mientras se prueba el primero.

(G5) ¿Debemos escribir pruebas para evaluar las características definidas por la implementación o no definidas comportamientos del modelo?

No debemos escribir pruebas que ejerzan las funciones definidas por la implementación o comportamientos indefinidos de la operación. Pruebas que ejecutan comportamientos definidos por la implementación demuestran un comportamiento local del intérprete que no debería generalizados. Las pruebas que llevan a cabo un comportamiento indefinido no contribuyen al la comprensión del comportamiento de la op.

(G6) Cuando se escriben pruebas para los tipos de punto flotante, ¿cuál es la precisión el resultado esperado se debe especificar en las verificaciones?

Para las operaciones primarias (suma, resta, multiplicación, división y cuadrada), se espera que una implementación siguiendo la especificación del IEEE proporcione una resultado redondeado dentro de 0.5 ULP del resultado matemáticamente exacto. Dicho esto, puedan imaginar con seguridad que el resultado esperado de estas operaciones la mayoría con 1 ULP de distancia. Sin embargo, es posible que esto no funcione para las funciones trascendentales. (sine, cosine, etc.) para las que se definidas por la implementación (razonamiento).

La implementación actual usa un modelo "único para todos" de tolerancia de 0.0001. El siguiente ejemplo demuestra la tolerancia anterior en acción.

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
}

Este es solo el primer paso para probar la exactitud numérica de las operaciones de StableHLO. Por el momento, esta es un área poco especificada de la especificación StableHLO, y hay trabajo continuo para resolverlo #1156 con base en nuestra experiencia con StableHLO en la práctica y en los comentarios de interesados. A medida que avance el proceso, actualizaremos la infraestructura según corresponda.

(G7) ¿Algo sobre el estilo de programación de las pruebas?

  1. Asegúrate de usar el nombre real de las entradas y salidas en lugar del valor predeterminado. a valores SSA (p. ej., %0, %1, etcétera)
  2. Asegúrate de que las pruebas usen un formato impreso, si existe.

(G8) ¿Deberíamos incluir el ejemplo ya proporcionado en la especificación? Sí (para que las pruebas sean completas).