XLA Araçları

XLA geliştirme iş akışı genellikle derleyiciye verilen izole işlevsel hesaplamayı temsil eden HLO ara dili etrafında şekillenir. XLA, HLO'yu kullanan ve onu çalıştıran ya da ara derleme aşaması sağlayan birden fazla komut satırı aracıyla (aşağıda açıklanmıştır) birlikte gelir. HLO hem görselleştirilebilir hem de değiştirilebilir olduğundan ve yinelemeli olarak değiştirilip çalıştırılması genellikle bir XLA performansını veya davranışını anlamanın ve düzeltmenin en hızlı yoludur. Bu nedenle, bu tür araçların kullanılması hızlı bir compile->modify->run yineleme döngüsü için çok değerlidir.

XLA ile derlenen bir program için HLO'yu almanın en kolay yolu genellikle XLA_FLAGS ortam değişkenini kullanmaktır:

$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point

Bu komut, belirtilen klasörde optimizasyon öncesindeki tüm HLO dosyalarını ve diğer birçok faydalı yapıyı depolar.

[run_hlo_module] HLO modüllerini çalıştırma

bazel run //xla/tools:run_hlo_module -- [flags] <filename>

Araç, run_hlo_module ön optimizasyonlu HLO üzerinde çalışır ve varsayılan olarak derleme, çalıştırma ve referans yorumlayıcı uygulamasıyla karşılaştırma işlemlerini birlikte yapar. Örneğin, bir giriş dosyasını NVIDIA GPU'da çalıştırmak ve doğruluğunu kontrol etmek için kullanılan normal çağırma işlemi şöyledir:computation.hlo

run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo

Birden Çok HLO Modülü Çalıştırma

run_hlo_module için birden fazla HLO modülüyle çağırma desteklenir. Bir dizindeki tüm hlo modüllerini çalıştırmak için:

bazel run //xla/tools:run_hlo_module -- [flags] /dump/*before_optimizations*

[multihost_hlo_runner] SPMD desteğiyle HLO modüllerini çalıştırma

# Note: Binary name is `hlo_runner_main`.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] <filename>

Çoklu ana makine HLO çalıştırıcısı, ana makineler arası iletişim de dahil olmak üzere SPMD'yi desteklemesi dışında çok benzer bir araçtır. Ayrıntılar için Çoklu Ana Makine HLO Çalıştırıcısı bölümüne bakın.

SPMD desteğiyle birden fazla HLO modülü çalıştırma

run_hlo_module'ya benzer şekilde, multihost_hlo_runner de birden fazla modülle çağırmayı destekler.

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] /dump/*before_optimizations*

[hlo-opt] HLO Modülünü Derle

bazel run //xla/tools:hlo-opt -- --platform=[gpu|cpu|...] [more flags] <filename>

Derleyiciyle ilgili hataları ayıklarken veya derleyicinin işleyişini anlamaya çalışırken, belirli bir HLO ya da StableHLO girişi için işlem hattının belirli bir noktasında (HLO, optimize edilmiş HLO, TritonIR veya LLVM) belirli bir donanımın genişletilmiş halini almak genellikle yararlıdır.

hlo-opt, PTX, optimizasyonlardan sonra HLO, optimizasyonlardan önce LLVM IR veya TritonIR gibi birden fazla çıkış aşamasını destekler. Desteklenen aşamaların tam kümesi platforma bağlıdır (ör. PTX, NVIDIA'ya özgüdür) ve --list-stages komutu kullanılarak görülebilir:

hlo-opt --platform=CUDA --list-stages
buffer-assignment
hlo
hlo-backend
html
llvm
llvm-after-optimizations
llvm-before-optimizations
ptx

Bir aşama seçtikten sonra kullanıcı, belirli bir platformdaki dönüşüm sonucunu belirli bir akışa yazabilir:

hlo-opt --platform=cpu --stage=hlo input.hlo

bu komut, dökümü stdout'a (veya -o belirtilmişse belirli bir dosyaya) yazdırır.

GPU için Cihazsız Derleme

Cihazsız derleme için GPU'ya erişim gerekmez. Cihazsız Derleme, GPU'ya erişimin gerekli olduğu aşamalarda komut satırında (--xla_gpu_target_config_filename) GPU spesifikasyonunu belirtmenin bir yolunu sunarak GPU cihazı ihtiyacını ortadan kaldırır.

Örnek: GPU cihazına erişim olmadan PTX çıkışı:

hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=/xla/tools/hlo_opt/gpu_specs/a100_pcie_80.txtpb input.hlo

Popüler GPU'ların özellikleri derleyiciyle birlikte gönderilir ve sağlanan dosya, device_description.proto öğesinin dize serileştirilmiş halidir:

gpu_device_info {
  cuda_compute_capability {
    major: 8
    minor: 0
  }
  threads_per_block_limit: 1024
  threads_per_warp: 32
  shared_memory_per_block: 127152
  shared_memory_per_core: 65536
  threads_per_core_limit: 2048
  core_count: 6192
  fpus_per_core: 64
  block_dim_limit_x: 2147483647
  block_dim_limit_y: 65535
  block_dim_limit_z: 65535
  memory_bandwidth: 2039000000000
  l2_cache_size: 4194304
  clock_rate_ghz: 1.1105
  device_memory_size: 79050250240
}
platform_name: "CUDA"

Daha fazla GPU spesifikasyonu için /xla/tools/hlo_opt/gpu_specs adresini ziyaret edin.

Otomatik ayarlama

Bazen derleme, derlemeye dayalı otomatik ayarlamayı içerebilir --stage. Cihazsız derlemenin çalışması için kullanıcının
otomatik ayarlamayı devre dışı bırakması --xla_gpu_autotune_level=0
veya
önceden var olan otomatik ayarlama sonuçlarını yüklemesi --xla_gpu_load_autotune_results_from=<filename> (--xla_gpu_dump_autotune_results_to=<filename> ile elde edilen) gerekir.

hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=gpu_specs/a100_pcie_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo

Otomatik ayarlama dosyası, autotune_results.proto öğesinin metin olarak serileştirilmiş halidir. Örneğin:

version: 3
results {
  device: "CUDA: 8.0, Cores: 108, GPU clock: 1.41 GHz, Memory bandwidth: 1555 GB/s, L2 cache: 40 MB"
  hlo: "{\n  tmp_0 = f16[1,16,17,3]{3,2,1,0} parameter(0)\n  tmp_1 = f16[16,51]{1,0} bitcast(f16[1,16,17,3]{3,2,1,0} tmp_0)\n  tmp_2 = s8[16,17,3]{2,1,0} parameter(1)\n  tmp_3 = s8[51,16]{0,1} bitcast(s8[16,17,3]{2,1,0} tmp_2)\n  tmp_4 = f16[51,16]{0,1} convert(s8[51,16]{0,1} tmp_3)\n  tmp_5 = f16[16,16]{1,0} dot(f16[16,51]{1,0} tmp_1, f16[51,16]{0,1} tmp_4), lhs_contracting_dims={1}, rhs_contracting_dims={0}\n  ROOT tmp_6 = f16[1,16,16]{2,1,0} bitcast(f16[16,16]{1,0} tmp_5)\n}"
  result {
    run_time {
      nanos: 31744
    }
    triton {
      block_m: 32
      block_n: 32
      block_k: 32
      split_k: 1
      num_stages: 1
      num_warps: 4
    }
  }
}

Otomatik ayarlama veritabanı, XLA_FLAGS=--xla_gpu_dump_autotune_results_to=<myfile.pbtxt> kullanılarak serileştirilebilir.

[hlo-opt] HLO Geçişi Geliştirme ve Hata Ayıklama

# If you are working with hardware independent passes from the
# `xla/hlo/transforms/` directory, prefer light-weight version
# of the `hlo-opt` tool with fewer dependencies:

bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>

# Otherwise, for hardware independent and CPU, GPU passes use
# the same binary from "Compile HLO Modules" section above:

bazel run //xla/tools:hlo-opt -- [flags] <filename>

hlo-opt aracı, belirli platform derleme aşamalarından bağımsız olarak tek bir geçişin yürütülmesine olanak tanır. Bu izolasyon, giriş hlo modülünde geçişleri hızlı bir şekilde çalıştırmaya ve hataların temel nedenini belirlemeye yardımcı olur.

hlo-opt --passes=schedule-aware-collective-cse input.hlo

hlo-opt aracı, DebugOptions XLA_FLAGS'ı da destekler.

hlo-opt --passes=schedule-aware-collective-cse
--xla_gpu_experimental_collective_cse_distance_threshold=20 input.hlo

Kart adı dizesini almak için--list-passes seçeneğini kullanın.

hlo-opt --list-passes

Kullanıcılar, --passes seçeneğine birden fazla geçiş belirterek kendi özel işlem hattını oluşturabilir.

hlo-opt --passes=pass1,pass2,pass3 input.hlo

Yeni HLO Pass Geliştirme Sürecine Yardımcı Olma

  1. Öncelikle şifrenizi yazın.
  2. Yeni kartı hlo-opt araç kartı kayıt defterine kaydedin.

    RegisterPass<FooPass>(FooPassInputOptions)
    

    Kart türüne bağlı olarak, kaydetmek için aşağıdaki konumlardan birini seçin:
    opt_lib.cc Donanımdan bağımsız kartlar.
    cpu_opt.cc CPU'ya özel geçişler.
    gpu_opt.cc GPU'ya özel geçişler.
    compiled_opt.cc CPU, GPU ve XPU'da ortak olan geçişler.
    Derleme bağımlılığı eklemeyi unutmayın.

    Tüm hlo-opt kullanıcılarının kullanabilmesi için kart kaydını PR'nize(örnek) dahil edin.

  3. hlo-opt aracını yeniden oluşturun, --list-passes seçeneğini kullanarak kart kaydının başarılı olduğunu doğrulayın ve ardından kartı çalıştırmak için --passes seçeneğini kullanın.

    $ hlo-opt --passes=foo-pass input.hlo
    
  4. Geçiş için birim testleri yazma hakkında daha fazla bilgi için https://openxla.org/xla/test_hlo_passes adresini ziyaret edin.

Çalışma Zamanı Ölçümüne Katılma

Büyük modellerde tam derleme çalıştırmaları birkaç dakika sürebilir. Bu durum, performanstaki küçük gerilemeleri tespit etmeyi zorlaştırır. Buna karşılık, hlo-opt kullanılarak yapılan bireysel geçişler, hassas performans ölçümüne ve yeni kod değişikliklerinin neden olduğu küçük artışların bile kolayca tespit edilmesine olanak tanır.

time hlo-opt --passes=reduce-window-rewriter,scatter_simplifier
--xla_reduce_window_rewrite_base_length=128 input.hlo

[hlo-opt] Convert HLO Module Formats

# Use the light weight version of the `hlo-opt` tool.

bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>

HLO Text -> HLO Proto birimini dönüştür

hlo-opt --emit-proto input.hlo

HLO Proto veya HLO Proto Binary -> HLO Text birimini dönüştürün.

hlo-opt input.pbtxt or input.pb

[ptx-opt] Compiler LLVM Module down to PTX

Araç, LLVMIR optimizasyon işlem hattını çalıştırır ve ardından CompileToPtx'i çağırır.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 <filename>

Araç, her yoldan sonra LLVMIR'yi de dökebilir.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 --xla_dump_to=<path> --xla_gpu_dump_llvmir <filename>