XLA 도구 사용

XLA 개발 워크플로는 일반적으로 HLO IR: 격리된 기능 공식을 사용할 수 있습니다. XLA는 여러 명령줄 도구와 함께 제공됩니다. (아래에 설명됨)은 HLO를 소비하고 이를 실행하거나 중간 컴파일 단계입니다 이러한 도구를 사용하는 것은 compile->modify->run 반복 주기(HLO는 시각화 가능하고) 반복적으로 변경하고 실행하는 것이 이해하는 데 도움이 될 수 있습니다

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

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

이 파일은 최적화 이전의 모든 HLO 파일을 지정된 폴더에 저장합니다. 여러 유용한 아티팩트가 있습니다.

실행 중인 HLO 스니펫: run_hlo_module

run_hlo_module 도구는 기본적으로 사전 최적화 HLO에서 작동합니다. 참조 인터프리터를 사용한 컴파일, 실행, 비교 번들 있습니다. 예를 들어 입력 파일을 실행하는 일반적인 호출은 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 --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo