XLA 도구 사용

XLA 개발 워크플로는 일반적으로 컴파일러에 제공된 격리된 기능 계산을 나타내는 HLO IR을 중심으로 합니다. XLA에는 HLO를 사용하고 실행하거나 중간 컴파일 단계를 제공하는 여러 명령줄 도구(아래 설명)가 함께 제공됩니다. 이러한 도구는 빠른 compile->modify->run 반복 주기에 매우 유용합니다. HLO는 시각화 및 해킹이 가능하며 반복적으로 변경하고 실행하는 것이 XLA 성능이나 동작을 이해하고 수정하는 가장 빠른 방법인 경우가 많기 때문입니다.

XLA로 컴파일되는 프로그램의 HLO를 얻는 가장 쉬운 방법은 일반적으로 XLA_FLAGS 환경 변수를 사용하는 것입니다.

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

이는 지정된 폴더에 최적화 전 HLO 파일과 기타 여러 유용한 아티팩트를 모두 저장합니다.

HLO 스니펫 실행: run_hlo_module

run_hlo_module 도구는 사전 최적화 HLO에서 작동하며 기본적으로 컴파일, 실행, 참조 인터프리터 구현과의 비교를 번들로 묶습니다. 예를 들어 NVIDIA GPU에서 입력 파일 computation.hlo를 실행하고 정확성을 확인하는 일반적인 호출은 다음과 같습니다.

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

모든 도구와 마찬가지로 --help를 사용하여 전체 옵션 목록을 가져올 수 있습니다.

SPMD 지원으로 HLO 스니펫 실행: multihost_hlo_runner

다중 호스트 HLO 러너는 매우 유사한 도구이지만 교차 호스트 통신을 비롯한 SPMD를 지원한다는 점에 유의해야 합니다. 자세한 내용은 다중 호스트 HLO 런너를 참고하세요.

다중 HLO 재생

run_hlo_module 및 hlo_runner_main 모두에서 여러 모듈을 사용한 호출이 지원되므로 덤프 디렉터리의 모든 모듈을 재생하는 데 편리합니다.

$ hlo_runner_main /dump/*before_optimizations*

HLO 컴파일의 패스/단계 실행: hlo-opt

컴파일러의 작동을 디버그하거나 이해할 때는 지정된 (안정적인) HLO 입력에 대해 파이프라인의 특정 지점 (HLO, 최적화된 HLO, TritonIR 또는 LLVM)에서 특정 하드웨어의 확장을 가져오는 것이 종종 유용합니다.

hlo-opt는 PTX, 최적화 후 HLO, 최적화 전 LLVM IR 또는 TritonIR 등 여러 출력 단계를 지원합니다. 지원되는 정확한 단계 집합은 플랫폼에 따라 다르며 (예: PTX는 NVIDIA 전용) --list-stages 명령어를 사용하여 확인할 수 있습니다.

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

사용자가 스테이지를 선택한 후 특정 플랫폼의 변환 결과를 특정 스트림에 쓸 수 있습니다.

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

그러면 덤프가 stdout (또는 -o가 지정된 경우 지정된 파일)에 출력됩니다.

기기 없는 사용

대부분의 컴파일에는 GPU 액세스가 필요하지 않으며 명령줄에서 GPU 사양을 지정하면 가속기에 액세스하지 않고도 PTX 출력을 가져올 수 있습니다.

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

널리 사용되는 GPU의 사양은 컴파일러와 함께 제공되며 제공된 파일은 device_description.proto의 문자열 직렬화입니다.

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"

자동 튜닝이 필요한 경우 기기 없는 컴파일에 문제가 발생할 수 있습니다. 다행히 명령줄에서도 이러한 옵션을 제공할 수 있습니다.

$ 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

자동 조정 파일은 autotune_results.proto의 텍스트 직렬화이며 예시는 다음과 같습니다.

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

자동 조정 데이터베이스는 XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>를 사용하여 직렬화할 수 있습니다.

단일 컴파일러 패스 실행

XLA_FLAGS의 플래그도 지원되므로 이 도구를 사용하여 단일 패스 실행을 테스트할 수 있습니다.

$ hlo-opt --platform=CUDA --stage=hlo --passes=algebraic_simplifer input.hlo