Intérprete de StableHLO

El objetivo principal del intérprete de StableHLO es proporcionar una referencia implementación de la semántica del opset StableHLO de acuerdo con su especificación. El objetivo secundario es que la implementación siga de cerca la especificación, lo que favorece la legibilidad sobre el rendimiento para brindar más claridad a la semántica, incluso de las operaciones más complejas, como Convolution Gather/Scatter y DotGeneral.

Por el momento, OpenXLA admite la interpretación de 91 de 96 Operaciones de StableHLO. Las 3 operaciones restantes (FftOp, RngOp, RngBitGeneratorOp) tienen su semántica documentada en spec.md y tienen hemos realizado investigaciones iniciales sobre cómo avanzar (consulta status.md para obtener una lista completa de las operaciones y su estado más reciente). Estos cambios finales se implementarán mejoras según sea necesario en la comunidad.

Alcance

Clasificamos el conjunto de operaciones StableHLO en 11 categorías con 118 operaciones en totales (consulta el Apéndice). Implementación de referencia de trabajo organiza el trabajo para implementar un intérprete para el 100% de las operaciones de StableHLO, según se define en la especificación de StableHLO. Somos planea completar la totalidad o casi todo el trabajo en este flujo de trabajo en StableHLO. Versión 1.0. De las 96 operaciones que tienen una especificación actualmente, podemos interpretar 91 operaciones OpenXLA (consulta los casos especiales para los 5 restantes).

Especificación

El requisito principal del intérprete es tener correspondencia 1:1 con el específica. La especificación permite la estandarización del intérprete en operaciones similares que conducir a una implementación modular y de alta calidad del intérprete.

Casos especiales

Varios

Esta categoría tiene operaciones descomponibles cuyo futuro no está claro en este momento. Hay hay tres operaciones específicas en esta categoría que el intérprete no admite en el momento:

• FftOp
• RngOp
• RngBitGeneratorOp

FftOp se clasifica como Varias, pero, a diferencia de otras operaciones de esta categoría, esta op no tiene un pase de expansión, y admitirlo en StableHLO es En construcción.

RngOp y RngBitGeneratorOp se pueden dividir en operaciones de MHLO, pero la descomposición introduce un XlaRngGetAndUpdateStateOp, que es un MHLO específico op. La interpretación de respaldo de estas dos operaciones es un trabajo en curso.

La herramienta para convertir las operaciones restantes de esta categoría en operaciones de StableHLO que que admite el intérprete reside en hlo_expand_main.cc.

No está en HLO

Además de las operaciones especificadas, esta categoría consta de 8 operaciones no especificadas (consulta categorías de operaciones de StableHLO), que están planificadas para y salí del StableHLO. La mayoría de estas operaciones tienen pases existentes en mhlo a y convertirlos en operaciones equivalentes a StableHLO.

La herramienta para convertir las operaciones restantes de esta categoría en operaciones de StableHLO equivalentes. que el intérprete admite reside en mlir-hlo-opt.cc.

Cuantización

La compatibilidad con el intérprete para la operación stablehlo.constant con tipo cuantizado es y se hace un seguimiento mediante #1691:

Instrucciones de uso

Creación del intérprete de referencia

El intérprete se puede compilar y probar a través de Bazel o CMake (opción preferida). Total consulta README.md.

Bazel:

bazel build //...

CMake:

mkdir -p build && cd build

cmake .. -GNinja \
  -DLLVM_ENABLE_LLD="$LLVM_ENABLE_LLD" \
  -DCMAKE_BUILD_TYPE=Release \
  -DLLVM_ENABLE_ASSERTIONS=On \
  -DMLIR_DIR=${PWD}/../llvm-build/lib/cmake/mlir

Para ejecutar el intérprete, tenemos una herramienta de traducción a fin de interpretar los programas de StableHLO. escritas en MLIR.

stablehlo-translate --interpret <path/to/program>

Dialecto de intérprete

El dialecto Interpreter contiene varias operaciones de utilidad relacionadas con la de un intérprete. Específicamente, el interpreter.run_parallel (consulta InterpreterOps.td (para semántica de op y uso de ejemplo), op permite interpretar las operaciones de distribución y mucho más de servicios públicos nuevos según las necesidades de la comunidad.

El dialecto de verificación

El dialecto Check se usa para comparar los valores del entorno de ejecución del intérprete con los esperados. de salida. Los resultados del programa StableHLO se pueden probar a través de varias operaciones de verificación (consulta CheckOps.td para la semántica de op y el uso de ejemplo).

Redacción de programas de prueba

Usamos la herramienta lit de LLVM para ejecutar y comparar con el archivo generado para diferenciar con el resultado del intérprete (consulta stablehlo/tests/interpret). como las pruebas).

Probando AddOp (muestra de interpret_add.mlir):

// RUN: stablehlo-translate --interpret %s

func.func @add_op_scalar() {
  %0 = stablehlo.constant dense<2> : tensor<i4>
  %1 = stablehlo.constant dense<3> : tensor<i4>
  %2 = stablehlo.add %0, %1 : tensor<i4>
  check.expect_eq_const %2, dense<5> : tensor<i4>
  func.return
}

Para probar las operaciones de la categoría Distribución, es necesario ejecutarla a través del Operación de utilidad interpreter.run_parallel.

Probando AllReduceOp (muestra de all_reduce.mlir):

// RUN: stablehlo-translate --interpret %s

module @cross_replica {
  func.func public @all_reduce(%operand : tensor<4xi64>) -> tensor<4xi64> {
    %result = "stablehlo.all_reduce"(%operand) ({
      ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
        %0 = stablehlo.add %arg0, %arg1 : tensor<i64>
        stablehlo.return %0 : tensor<i64>
    }) {
      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
      channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
    } : (tensor<4xi64>) -> tensor<4xi64>
    return %result : tensor<4xi64>
  }
  func.func public @main() {
    %inputs0 = stablehlo.constant dense<[1, 2, 3, 4]> : tensor<4xi64>
    %inputs1 = stablehlo.constant dense<[5, 6, 7, 8]> : tensor<4xi64>
    %results:2 = "interpreter.run_parallel"(%inputs0, %inputs1) {
      programs=[[@all_reduce], [@all_reduce]]
    } : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>, tensor<4xi64>)
    check.expect_eq_const %results#0, dense<[6, 8, 10, 12]> : tensor<4xi64>
    check.expect_eq_const %results#1, dense<[6, 8, 10, 12]> : tensor<4xi64>
    func.return
  }
}

Cómo depurar StableHLO

Siguiendo los pasos de compilación de StableHLO, los objetos binarios de StableHLO para herramientas en stablehlo/tools debe estar en /build/bin. Herramientas de depuración comunes, como GDB se puede utilizar para recorrer el código:

gdb --args ./build/bin/stablehlo-translate -allow-unregistered-dialect --interpret ./stablehlo/tests/interpret/<test>.mlir

Apéndice

Convertir otras operaciones

# batch_norm_grad
hlo-expand --batch_norm_grad_expander <path/to/hlo_module>

# batch_norm_inference
hlo-expand --batch_norm_inference_expander <path/to/hlo_module>

# batch_norm_training
hlo-expand --batch_norm_training_expander <path/to/hlo_module>

# cholesky
hlo-expand --cholesky_expander <path/to/hlo_module>

# constant
# Supported in StableHLO interpreter.

# fft
# TBD

# iota
# Supported in StableHLO interpreter.

# rng
# TBD

# rng_bit_generator
# TBD

# triangular_solve
hlo-expand --triangular_solve_expander <path/to/hlo_module>

Convertir en operaciones de HLO

# broadcast
mlir-hlo-opt -mhlo-legalize-broadcast-to-broadcast-in-dim <path/to/input>

# create_token
mlir-hlo-opt -mhlo-legalize-create-token-to-after-all <path/to/input>

# cross-replica-sum
mlir-hlo-opt -mhlo-legalize-cross-replica-sum-to-all-reduce <path/to/input>

# dot
mlir-hlo-opt -mhlo-legalize-dot-to-dot-general <path/to/input>

# einsum
mlir-hlo-opt -mhlo-legalize-einsum-to-dot-general <path/to/input>

# torch_index_select
mlir-hlo-opt -mhlo-legalize-torch-index-select-to-gather <path/to/input>

# unary_einsum
mlir-hlo-opt --canonicalize -mhlo-legalize-einsum-to-dot-general <path/to/input>

Categorías de operaciones de StableHLO