Inferencia de tipos

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

Los lineamientos propuestos a continuación tienen como objetivo garantizar la implementación de verificadores de alta calidad y funciones de forma para las operaciones estables.

Propuesta

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

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

spec es la fuente de información para todos los verificadores y las funciones de forma de las operaciones estables. Es necesario revisar los verificadores y las funciones de forma existentes de cada operación para que se alineen por completo con la especificación. Ten en cuenta que el documento de especificación evoluciona constantemente. En los casos en que la especificación de una op no esté disponible, se debe usar la implementación de XLA como fuente de información, incluidas xla/service/shape_inference.cc y xla/service/hlo_verifier.cc. La implementación de XLA no cubre 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 operaciones con traits y tipos para cada operando, atributo o resultado, y realizarán las verificaciones. Por lo tanto, NO se necesita ningún código de verificación en los verificadores o las funciones de forma para las propiedades ya garantizadas por la ODS. Quita el código de verificación si está duplicado con ODS, ya que nunca se activarán.

¿Debemos agregar pruebas para las restricciones de la ODS? Consulta Cómo establecer lineamientos para pruebas.

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

Ambos:

  • verificadores: Se implementan mediante Op::verify().
  • funciones de forma: implementadas por elementos de InferTypeOpInterface, como Op::inferReturnTypes() o Op::inferReturnTypeComponents

puede tener un código de verificación para comprobar operandos/atributos/resultados. Una división inicial puede ser la siguiente: dejar que los verificadores revisen los operandos o 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 algunos problemas:

  • Las funciones build() generadas automáticamente pueden llamar a la función de forma sin llamar primero al verificador. Por lo tanto, las entradas relacionadas también deben verificarse en la función Shape.
  • Código duplicado: Por ejemplo, en los verificadores, procesamos 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 se deben calcular dos veces.
  • Carga de mantenimiento: las verificaciones de una op se contienen en dos métodos diferentes.

La solución es la siguiente:

  1. Para la mayoría de las operaciones sin regiones (como PadOp), coloca todo el código de verificación en las funciones de forma y descarta por completo los verificadores.

  2. Para operaciones con regiones (como ReduceOp/IfOp; aquí hay una lista completa): los compiladores generados automáticamente no toman regiones como parámetros, por lo que si estos compiladores involucran inferencia de tipos, 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 tipo (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) Establecer lineamientos para las pruebas

¿Es necesario agregar o mantener pruebas para las verificaciones que están cubiertas por la ODS?

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

Sin embargo, ten cuidado con las partes faltantes: por ejemplo, si la operación 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 la inferencia de tipo?

ops_stablehlo.mlir contiene los casos positivos de operaciones y (al menos) 1 prueba negativa para cada error de verificación. También puede verificar que el tipo de resultado inferido es 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 op por línea con hlo_test_infer.get_return_type_components"(%x):... y comprueba que el tipo inferido coincida exactamente como se espera. Una prueba positiva por operación en general.

Qué hacer

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

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

  2. Agrega una sola 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.