Inferencia de tipos

En un principio, StableHLO se inició a partir del dialecto MHLO y heredó la implementación MHLO de la inferencia de tipo. El progreso de la implementación se registra en status.md.

Los lineamientos propuestos a continuación están destinados a garantizar la implementación de verificadores de alta calidad y funciones de forma para las operaciones de StableHLO.

Propuesta

Estas propuestas se aplican tanto a la revisión de implementaciones existentes como a la adquisición de operaciones nuevas hasta una cobertura integral.

(P1) Uso de la especificación StableHLO como fuente de confianza

spec es la fuente de información de todos los verificadores y funciones de forma de las operaciones StableHLO. Es necesario volver a revisar los verificadores y las funciones de forma existentes de cada operación para que estén completamente alineados con la especificación. Ten en cuenta que el documento de especificación evoluciona constantemente. En los casos en que la especificación para una op no esté disponible, se debe usar la implementación de XLA como fuente de confianza, lo que incluye xla/service/shape_inference.cc y xla/service/hlo_verifier.cc. La implementación de XLA no abarca el dinamismo ilimitado, por lo que, para este tipo de dinamismo, aplicaremos el sentido común hasta que esté disponible el RFC de dinamismo.

(P2) Aprovecha al máximo la ODS

Los archivos ODS (como StablehloOps.td) definen las operaciones con rasgos y tipos para cada operando, atributo o resultado, y realizarán las verificaciones. Por lo tanto, NO se necesita un código de verificación en los verificadores ni en las funciones de forma para las propiedades ya garantizadas por la ODS. Quita el código de verificación si está duplicado con el ODS, ya que nunca se activarán.

¿Es necesario agregar pruebas para las restricciones de la ODS? Consulta Establece lineamientos de prueba.

(P3) Mantener el código de verificación en los verificadores y las funciones de forma

Ambos:

  • verificadores: los implementa Op::verify()
  • Funciones de forma: Se implementan a través de InferTypeOpInterface, como Op::inferReturnTypes() o Op::inferReturnTypeComponents.

pueden tener un código de verificación para comprobar operandos/atributos/resultados. Una división inicial podría ser la siguiente: permitir que los verificadores verifiquen los operandos y los atributos, y luego permitir que las funciones de forma solo calculen tipos de resultados inferidos y verifiquen la compatibilidad con los tipos de resultados reales. Sin embargo, en realidad, esta división tiene los siguientes problemas:

  • Las funciones build() generadas automáticamente pueden llamar a la función de forma sin llamar primero al verificador. Así que las entradas relacionadas también deben verificarse en la función Shape.
  • Código duplicado: Por ejemplo, en los verificadores, realizamos un procesamiento en los operandos y, luego, verificamos algunos resultados intermedios. Luego, en las funciones de forma, estos resultados intermedios son útiles para inferir los resultados finales. Estos resultados intermedios deben calcularse dos veces.
  • Carga de mantenimiento: Las verificaciones de una op se encuentran en dos métodos diferentes.

La solución es la siguiente:

  1. Para la mayoría de las operaciones sin regiones (como PadOp), intenta colocar todo el código de verificación en las funciones de forma y descartar por completo los verificadores. En los casos en que esto no sea posible debido a la imposibilidad de inferir los tipos de datos que se muestran (como con ReshapeOp o BroadcastInDimOp), crea un verificador para que contenga la lógica de verificación necesaria. Por lo general, las operaciones inferibles, como AddOp, aún pueden necesitar un verificador para realizar verificaciones adicionales, ya que verifican restricciones de un tipo de datos que se muestra proporcionado, al que no se puede acceder en los métodos de inferencia de tipo/forma.

  2. Para operaciones con regiones (como ReduceOp/IfOp; hay una lista completa aquí): Los compiladores generados automáticamente no toman regiones como parámetros, por lo que si estos compiladores incluyen inferencia de tipo, se llamará a la función de forma con regiones vacías (consulta este ejemplo).

    1. Si las regiones no son necesarias para la inferencia de tipo (como ReduceOp), coloca la lógica de verificación relacionada con la región en verificadores, en lugar de las funciones de forma. Duplicar algún código si es inevitable.

    2. Si se necesitan las regiones para la inferencia de tipos (IfOp/CaseOp/MapOp), además, la función de forma debe verificar que las regiones no estén vacías de forma explícita, aunque la ODS ya pueda garantizar su existencia en la definición de Op.

(P4) Establecimiento de lineamientos para las pruebas

¿Necesitamos agregar o mantener pruebas para las verificaciones que están cubiertas por el ODS?

Nosotros no. Las pruebas deben enfocarse en los verificadores y las funciones de forma, mientras que los cambios en ODS necesitan una revisión de esta op.

Sin embargo, ten cuidado con las partes faltantes: por ejemplo, si la op contiene el trait SameOperandsAndResultShape, que solo verifica las formas, pero no el tipo de elemento, la verificación de los tipos de elementos de operandos o resultados aún necesita pruebas.

¿Dónde se colocan las pruebas de los verificadores y de la inferencia de tipo?

ops_stablehlo.mlir contiene los casos positivos de operaciones y (al menos) 1 prueba negativa por cada error de verificación. También puede verificar que el tipo de resultado inferido sea compatible con el tipo de resultado real (no es lo mismo).

infer_stablehlo.mlir verifica la existencia de la función de forma de una operación por línea con hlo_test_infer.get_return_type_components"(%x):... y comprueba que el tipo inferido coincida exactamente con lo esperado. Una prueba positiva por operación en general.

Qué hacer

Cuando implementes o vuelvas a revisar el verificador o la función de forma de una op, ten en cuenta lo siguiente:

  1. Coloca todos los casos positivos y negativos en ops_stablehlo.mlir.

  2. Agrega una única prueba positiva en infer_stablehlo.mlir para probar la interfaz.

  3. (Opcional) Si una op es complicada y puede contener muchas pruebas, considera agregar un archivo de prueba separado llamado verify_<op_name>.mlir o verify_<your_topic>.mlir dentro de la misma carpeta.