Juru Bahasa StableHLO

Tujuan utama penafsir StableHLO adalah untuk menyediakan referensi pada semantik opset StableHLO sesuai dengan spesifikasi pendukung. Tujuan sekundernya adalah agar implementasinya mengikuti spesifikasi, mengutamakan keterbacaan daripada performa, untuk memberikan kejelasan tambahan ke semantik operasi yang paling terlibat seperti Convolution, Gather/Scatter, dan DotGeneral.

Saat ini, OpenXLA mendukung interpretasi 91 dari 96 spesifikasi Operasi StabilHLO. 3 operasi yang tersisa (FftOp, RngOp, RngBitGeneratorOp) telah semantik mereka didokumentasikan dalam spec.md, dan memiliki menyelesaikan penyelidikan awal tentang cara melangkah maju (lihat status.md untuk mengetahui daftar lengkap operasi dan status terbarunya). Laporan akhir akan diterapkan berdasarkan kebutuhan komunitas.

Cakupan

Kami mengategorikan opset StableHLO ke dalam 11 kategori yang terdiri dari 118 operasi di total (lihat Lampiran). Penerapan Referensi workstream mengatur tugas penerapan penerjemah untuk 100% operasi StableHLO seperti yang ditetapkan dalam spesifikasi StableHLO. Kami berencana untuk menyelesaikan semua atau hampir semua pekerjaan dalam alur kerja ini di StableHLO v1.0. Dari 96 operasi yang saat ini memiliki spesifikasi, kita dapat menafsirkan 91 operasi melalui OpenXLA (lihat Kasus Khusus untuk 5 kasus lainnya).

Spesifikasi

Persyaratan utama untuk penerjemah adalah memiliki korespondensi 1:1 dengan spesifikasi. Spesifikasi memungkinkan standarisasi penafsir di seluruh operasi serupa yang menjadikan implementasi penerjemah modular dan berkualitas tinggi.

Kasus Khusus

Lain-lain

Kategori ini memiliki operasi yang dapat diurai yang masa depannya tidak jelas saat ini. Ada ada tiga operasi khusus dalam kategori ini yang tidak didukung oleh penerjemah momen tersebut:

• FftOp
• RngOp
• RngBitGeneratorOp

FftOp dikategorikan sebagai Lain-Lain, tetapi tidak seperti operasi lain dalam kategori ini, operasi ini tidak memiliki pass peluas, dan mendukungnya di StableHLO adalah WIP.

RngOp dan RngBitGeneratorOp dapat diurai menjadi operasi MHLO, tetapi dekomposisi memperkenalkan XlaRngGetAndUpdateStateOp yang merupakan atau. Penafsiran pendukung kedua operasi ini adalah WIP.

Alat untuk mengonversi operasi yang tersisa dalam kategori ini ke operasi StableHLO yang penafsiran yang didukung berada di hlo_expand_main.cc.

Tidak di HLO

Terlepas dari operasi yang ditentukan, kategori ini terdiri dari 8 operasi yang tidak ditentukan (lihat Kategori Operasi SttableHLO) yang direncanakan pindah dari StableHLO. Sebagian besar operasi ini sudah memiliki kartu di mhlo ke mengubahnya ke operasi yang setara dengan StableHLO.

Alat untuk mengonversi operasi yang tersisa dalam kategori ini ke operasi Stabil yang setara yang didukung penafsir berada di mlir-hlo-opt.cc.

Kuantisasi

Dukungan penerjemah untuk operasi stablehlo.constant dengan jenis terkuantisasi adalah tidak didukung dan dilacak melalui #1691.

Petunjuk Penggunaan

Membangun Penerjemah Referensi

Penafsir dapat dibangun dan diuji melalui Bazel atau CMake (lebih disukai). Untuk lengkap lihat 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

Untuk menjalankan penafsir, kami memiliki alat terjemahan untuk menafsirkan program StableHLO yang ditulis dalam MLIR.

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

Dialek Penerjemah

Dialek Interpreter berisi berbagai operasi utilitas yang terkait dengan penerjemahan mesin. Secara khusus, interpreter.run_parallel (lihat InterpreterOps.td untuk semantik operasi dan contoh penggunaan) memungkinkan interpretasi operasi Distribusi, dan banyak lagi utilitas akan ditambahkan berdasarkan kebutuhan masyarakat.

Dialek Cek

Dialek Check digunakan untuk membandingkan nilai runtime penafsir dengan yang diharapkan masing-masing. Output program StableHLO dapat diuji melalui berbagai operasi pemeriksaan (lihat CheckOps.td untuk semantik operasi dan contoh penggunaan).

Program Ujian Menulis

Kita menggunakan alat lit LLVM untuk menjalankan dan membandingkan file yang dihasilkan untuk membedakan dengan output penafsir (lihat stablehlo/tests/interpret misalnya, pengujian).

Menguji AddOp (sampel dari 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
}

Pengujian operasi dalam kategori Distribusi mengharuskan untuk menjalankannya melalui Operasi utilitas interpreter.run_parallel.

Menguji AllReduceOp (sampel dari 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
  }
}

Men-debug StabilHLO

Dengan mengikuti langkah-langkah build StableHLO, biner StableHLO untuk alat di stablehlo/tools harus berada di /build/bin. Alat {i>debugging <i}yang umum seperti GDB dapat digunakan untuk menelusuri kode:

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

Lampiran

Konversikan Operasi Lain-Lain

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

Konversi Tidak Di Operasi 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>

Kategori Operasi StabilHLO