XLA 개발 워크플로는 일반적으로 컴파일러에 제공되는 격리된 함수 계산을 나타내는 HLO IR에 중점을 둡니다. XLA에는 HLO를 사용하고 실행하거나 중간 컴파일 단계를 제공하는 여러 명령줄 도구(아래에 설명됨)가 함께 제공됩니다. 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_runner_main --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo
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_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_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo
자동 조정 파일은 autotune_results.proto의 텍스트 직렬화이며, 예를 들면 다음과 같습니다.
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
}
}
}
자동 조정 데이터베이스는 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