XLA araçlarını kullanma

XLA geliştirme iş akışı genellikle derleyiciye verilen izole işlevsel hesaplamayı temsil eden HLO IR çevresinde ortalanır. XLA, HLO'yu kullanan ve bu komutu çalıştıran veya bir ara derleme aşaması sağlayan birden fazla komut satırı aracı (aşağıda açıklanmıştır) ile birlikte gelir. HLO hem görselleştirilebilir hem de saldırıya açık olduğundan ve tekrar tekrar 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 paha biçilmez bir önem taşır.

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

optimizasyondan önceki tüm HLO dosyalarıyla birlikte başka birçok yararlı yapıyla birlikte belirtilen klasörde depolar.

HLO snippet'leri çalıştırılıyor: run_hlo_module

run_hlo_module aracı, optimizasyon öncesi HLO üzerinde çalışır ve varsayılan olarak derlemeyi, çalıştırmayı ve karşılaştırmayı referans çevirmen uygulamasıyla birlikte paket haline getirir. Örneğin, NVIDIA GPU'da bir giriş dosyasını computation.hlo çalıştırmak ve dosyanın doğruluğunu kontrol etmek için yapılan normal çağrı şu şekildedir:

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

Tüm araçlarda olduğu gibi, seçeneklerin tam listesini almak için --help kullanılabilir.

SPMD desteğiyle HLO snippet'lerini çalıştırma: multihost_hlo_runner

Birden fazla ana makine HLO çalıştırıcısı, ana makineler arası iletişim de dahil olmak üzere SPMD'yi desteklediğine dikkat etmekle buna çok benzer bir araçtır. Tipik bir çağrı şu şekilde görünür:

hlo_runner_main  --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo

HLO derlemesinin koşu kartları/aşamaları: hlo-opt

Derleyicinin çalışmalarında hata ayıklarken veya derleyicinin işleyişini öğrenirken, belirli bir (Kararlı) HLO girişi için ardışık düzendeki belirli bir noktada (HLO, optimize edilmiş HLO, TritonIR veya LLVM) belirli bir donanımın genişletmesini elde etmek genellikle faydalıdır.

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

$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx

Kullanıcı bir aşama seçtikten sonra belirli bir platforma ilişkin dönüşümün sonucunu belirli bir akışa yazabilir:

$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm

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

Cihazsız Kullanım

Derlemenin çoğu için GPU'ya erişim gerekli değildir ve komut satırında bir GPU spesifikasyonu belirterek GPU'ya erişim (ör. bir hızlandırıcıya erişmeden PTX çıkışı) alabiliriz:

$ hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_80.txtpb input.hlo

Popüler GPU'ların özellikleri derleyiciyle birlikte gönderilir ve sağlanan dosya, device_description.proto için dize serileştirmesidir:

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"

Otomatik ayarlama gerekiyorsa cihazsız derlemede sorunlarla karşılaşabilirsiniz. Neyse ki, komut satırında bunları da sağlayabiliriz:

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

Autotune dosyası, autotune_results.proto metninin serileştirilmesidir. Örnek:

version: 2
results {
  device: "sm_8.0 with 42331013120B RAM, 108 cores, 1410000KHz clock, 1215000KHz mem clock, 41943040B L2$"
  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_t=<myfile.pbtxt> kullanılarak serileştirilebilir

Tek Derleyici Kartı Çalıştırma

XLA_FLAGS işaretleri de desteklenir. Böylece araç, tek bir kartın koşusunu test etmek için kullanılabilir:

hlo-opt --platform=CUDA --stage=hlo --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo