Intérprete de StableHLO

O principal objetivo do intérprete do StableHLO é fornecer uma referência implementação à semântica da opset StableHLO de acordo com especificação. A meta secundária é que a implementação acompanhe de perto da especificação, favorecendo a legibilidade em vez do desempenho, para oferecer mais clareza à semântica até mesmo das operações mais envolvidas, como Convolution, Gather/Scatter e DotGeneral.

No momento, o OpenXLA suporta a interpretação de 91 de 96 especificações Operações de StableHLO. As três operações restantes (FftOp, RngOp, RngBitGeneratorOp) têm a semântica é documentada em spec.md e têm concluído as investigações iniciais sobre como avançar (consulte status.md para uma lista completa das operações e o status mais recente delas). Os resultados as melhorias serão implementadas conforme necessário.

Escopo

Categorizamos a opset StableHLO em 11 categorias, que consistem em 118 operações em total (consulte o Apêndice). Implementação de referência fluxo de trabalho organiza o trabalho de implementação de um intérprete para 100% das operações StableHLO, conforme definido na especificação StableHLO. Estamos planeja concluir todo ou quase todo o trabalho neste fluxo de trabalho no StableHLO v1.0 Das 96 operações que têm uma especificação atualmente, podemos interpretar 91 operações por meio OpenXLA (consulte Casos especiais para saber os cinco restantes).

Especificação

O principal requisito para o intérprete é ter uma correspondência individual com o especificação. A especificação permite padronizar o intérprete em operações semelhantes que levam à implementação modular e de alta qualidade do intérprete.

Casos especiais

Diversos

Esta categoria tem operações decombináveis com um futuro incerto no momento. Há três operações específicas nesta categoria que o intérprete não oferece suporte momento:

• FftOp
• RngOp
• RngBitGeneratorOp

FftOp está classificado como "Diversos", mas, ao contrário de outras operações nessa categoria, essa operação não tem passagem de expansão, e o suporte a ela no StableHLO é Em andamento.

RngOp e RngBitGeneratorOp podem ser decompostos em operações MHLO, mas a decomposição introduz uma XlaRngGetAndUpdateStateOp, que é uma classe MHLO específica operação A interpretação de suporte dessas duas operações está em andamento.

A ferramenta para converter as operações restantes nessa categoria em operações StableHLO que o interpretador reside em hlo_expand_main.cc.

Não está em HLO

Além das operações especificadas, essa categoria consiste em 8 operações não especificadas (consulte categorias de operações do StableHLO), que serão movido para fora de StableHLO. A maioria dessas operações tem cartões mhlo para e convertê-los em operações equivalentes a StableHLO.

Ferramenta para converter as operações restantes nessa categoria em operações StableHLO equivalentes. que o intérprete suporta fica em mlir-hlo-opt.cc.

Quantização.

O suporte a intérpretes para a operação stablehlo.constant com tipo quantizado é incompatível e rastreado por 1691 (link em inglês).

Instruções de uso

Como criar o intérprete de referência

O intérprete pode ser criado e testado usando o Bazel ou o CMake (preferencial). Para total consulte README.md.

Bazel:

bazel build //...

O 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 executar o intérprete, temos uma ferramenta de tradução que interpreta programas StableHLO escritos em MLIR.

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

O dialeto do intérprete

O dialeto Interpreter contém várias operações utilitárias relacionadas à intérprete. Especificamente, o interpreter.run_parallel (consulte InterpreterOps.td para semântica de operações e uso de exemplo), o op permite a interpretação de operações de distribuição e muito mais os concessionárias de serviços públicos planejam ser adicionados com base nas necessidades da comunidade.

O dialeto de verificação

O dialeto Check é usado para comparar os valores de ambiente de execução do interpretador com o esperado e a distribuição dos valores dos dados. As saídas do programa StableHLO podem ser testadas com várias operações de verificação (consulte CheckOps.td para semântica de operações e uso de exemplo).

Como criar programas de teste

Usamos a ferramenta lit do LLVM para executar e comparar com o arquivo gerado para diferenciá-lo da saída do intérprete. Consulte stablehlo/tests/interpret como testes).

Testando AddOp (amostra 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
}

As operações de teste na categoria "Distribuição" exigem a execução delas por meio do Operação utilitária interpreter.run_parallel

Testando AllReduceOp (amostra 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
  }
}

Como depurar StableHLO

Seguindo as etapas de build do StableHLO, os binários do StableHLO para ferramentas stablehlo/tools precisa residir em /build/bin. Ferramentas de depuração comuns, como O GDB pode ser usado para percorrer o código:

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

Apêndice

Converter operações diversas

# 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>

Converter que não está em operações 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>

Categorias de operações do StableHLO