StableHLO Çevirmeni

StableHLO çevirmeninin ana hedefi, söz konusu kullanıcıya mantığına göre StableHLO işletim sisteminin anlamlarına bakın. İkincil hedef, uygulamanın izleme ve daha anlaşılır hale getirmek için performansa kıyasla okunabilirliği hatta Convolution gibi en karmaşık işlemlerin bile anlamında Gather/Scatter ve DotGeneral.

Şu anda OpenXLA, uzman 96 testten 91'inin StableHLO operasyonları Geri kalan 3 işlem (FftOp, RngOp, RngBitGeneratorOp) anlamları spec.md ve nasıl ilerleyeceğimize dair ilk araştırmaları tamamladık (bkz. status.md (operasyonların tam listesi ve en güncel durumlarını görebilirsiniz). Bu son gerekli geliştirmeler topluluk düzeyinde uygulanacaktır.

Kapsam

StableHLO işletim sistemini 118 operasyondan oluşan 11 kategoride toplam (Ek'e bakın). Referans Uygulama Çevirmen uygulama ile ilgili çalışmaları düzenleyen iş akışı StableHLO spesifikasyonunda tanımlandığı gibi StableHLO işlemlerinin% 100'ü için. Biz Bu iş akışındaki çalışmaların tamamını veya neredeyse tamamını StableHLO’da tamamlamayı planlamak sürüm 1.0 Şu anda spesifikasyona sahip olan 96 operasyondan 91 işlemi, OpenXLA (kalan 5 özel durum için Özel Durumlar'a bakın).

Spesifikasyon

Çevirmenin temel şartı, uzmanla spesifikasyonları. Bu spesifikasyon, çevirmenin farklı işlemlerde standartlaştırılmasına olanak tanır. çevirmenin modüler ve yüksek kaliteli uygulanmasını sağlar.

Özel Durumlar

Çeşitli

Bu kategoride, şu anda geleceği belirsiz olan ayrıştırılabilir operasyonlar var. Orada çevirmenin, bu kategoride desteklemediği ve denetlenmeyen üç işlem an:

  • FftOp
  • RngOp
  • RngBitGeneratorOp

FftOp Çeşitli olarak sınıflandırılır ancak bu kategorideki diğer işlemlerin aksine bu operasyon biriminin genişletme kartı yoktur ve bunu StableHLO'da desteklemek HAYIR.

RngOp ve RngBitGeneratorOp MHLO işlemlerine ayrıştırılabilir ancak ayrışma, MHLO'ya özgü bir XlaRngGetAndUpdateStateOp sağlar. f. Bu iki işlemin yorumlanmasını desteklemek bir WIP'dir.

Bu kategoride kalan işlemleri StableHLO operasyonlarına dönüştüren, çevirmen, hlo_expand_main.cc içinde bulunur.

HLO'da değil

Bu kategori, denetlenmiş operasyonların dışında 8 benzersiz işlemden oluşur ( StableHLO Ops Kategorileri) StableHLO'dan taşındı. Bu operasyonların çoğunun şurada mevcut kartları var: mhlo'dan bunları StableHLO eşdeğeri işlemlere dönüştürür.

Bu kategoride kalan işlemleri eşdeğer StableHLO operasyonlarına dönüştüren araç mlir-hlo-opt.cc dosyasında bulunur.

Nicelendirme

Sayısal türlerde stablehlo.constant işlemi için çevirmen desteği desteklenmemekte ve #1691.

Kullanım Talimatları

Referans Çevirmeni Oluşturma

Çevirmen, Bazel veya CMake ile geliştirilip test edilebilir (tercih edilir). Tam fiyat için README.md dosyasına göz atın.

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

Çevirmeni çalıştırmak için StableHLO programlarını MLIR dilinde yazılır.

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

Çevirmen Lehçesi

Interpreter diyalektinde, kullanabilirsiniz. Daha ayrıntılı olarak belirtmek gerekirse, interpreter.run_parallel (bkz. InterpreterOps.td işlem anlamı ve örnek kullanımı için), işlem, Dağıtım işlemlerinin ve diğer işlemlerin yorumlanmasına olanak tanır. toplumun ihtiyaçlarına göre kamu hizmetleri eklenmeyi planlamaktadır.

Kontrol Lehçesi

Çevirmen çalışma zamanı değerlerini beklenenlerle karşılaştırmak için Check diyalektini kullanılır değerler. StableHLO programı çıkışları, çeşitli kontrol işlemleriyle test edilebilir (bkz. CheckOps.td (işlem anlamları ve örnek kullanımı).

Test Programları Yazma

LLVM'nin ışık aracını kullanarak oluşturulan dosyayla karşılaştırarak, yorumlayıcının çıkışıyla karşılaştırabilirsiniz. (bkz. stablehlo/tests/interpret .

AddOp test ediliyor (örnek: 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
}

Dağıtım kategorisindeki işlemleri test etmek için, interpreter.run_parallel hizmet operasyonu

AllReduceOp test ediliyor (örnek: 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
  }
}

StableHLO'da hata ayıklama

StableHLO derleme adımlarının ardından, stablehlo/tools, /build/bin bölgesinde bulunmalıdır. Sık kullanılan hata ayıklama araçları: Kodda gezinmek için GDB kullanılabilir:

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

Ek

Çeşitli İşlemleri Dönüştür

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

HLO'da Olmayan İşlemleri Dönüştür

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

StableHLO Operasyon Kategorileri

Kategoriler Hafıza güçlendiriciler Toplam
119
Kontrol Akışı after_all, case, if, optimizasyon_barrier, ise 5
Veri Aktarımı yayın_in_dim, bitiş, dinamik_dilim, dinamik_güncelleme_dizisi, toplama, dolgu, yeniden şekillendirme, ters, dağılım, dilim, sıralama, ters çevirme 12
Dağıtım all_topher, all_reduce, all_to_all, toplu_permute, feed içi, outfeed, bölüm_kimliği, recv, azaltacak_dağılım, replika_kimliği, gönder 11
Dinamik dinamik_yayındaki_dim_ayar, dinamik_dönş, dinamik_toplama, dinamik_iota, dinamik_pad, dinamik_yeniden şekillendirme, al_boyut_boyutu, gerçek_dinamik_bölüm, ayarla_boyut_boyutu 9
Elementwise abs, add, and, atan2, bitcast_convert, cbrt, ceil, clamp, compare, complex, convert, cosine, count_leading_zeros, divide, exponential, exponential_minus_one, floor, imag, is_finite, log, log_plus_one, logistic, map, maximum, minimum, multiply, negate, not, or, popcnt, power, real, reduce_precision, remainder, round_nearest_afz, round_nearest_even, rsqrt, select, shift_left, shift_right_arithmetic, shift_right_logical, sign, sine, sqrt, subtract, tan, tanh, xor 48
Genişletilebilirlik custom_call, get_tuple_element, demet 3
Çeşitli toplu_norm_grad, toplu_norm_inference, toplu_norm_eğitimi, cholesky, sabit, fft, iota, rng, rng_bit_generator, üçgen_çözüm 10
Modularity ara, fonksiyon, modül, dönüş 4
HLO'da değil yayın, create_token, cross-replica-sum, nokta, einsum, torch_index_select, unary_einsum 8
Nicelendirme uniform_dequantize, uniform_quantize 2
Düşüş konvolüsyon, nokta_genel, azalt, pencere_azaltma, seç_ve_dağılım 5