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 파일을 지정된 폴더에 다른 유용한 아티팩트와 함께 저장합니다.

[run_hlo_module] HLO 모듈 실행

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

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

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

여러 HLO 모듈 실행

여러 HLO 모듈을 사용한 호출은 run_hlo_module에서 지원됩니다. 디렉터리에서 모든 HLO 모듈을 실행하려면 다음을 실행하세요.

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

[multihost_hlo_runner] SPMD 지원으로 HLO 모듈 실행

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

다중 호스트 HLO 러너는 호스트 간 통신을 비롯한 SPMD를 지원한다는 점을 제외하고 매우 유사한 도구입니다. 자세한 내용은 Multi-Host HLO Runner를 참고하세요.

SPMD 지원으로 여러 HLO 모듈 실행

run_hlo_module와 마찬가지로 multihost_hlo_runner도 여러 모듈을 사용한 호출을 지원합니다.

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

[hlo-opt] HLO 모듈 컴파일

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

컴파일러의 작동 방식을 디버깅하거나 이해할 때 주어진 HLO 또는 StableHLO 입력에 대해 파이프라인의 특정 지점에서 특정 하드웨어의 확장을 가져오는 것이 유용한 경우가 많습니다 (HLO, 최적화된 HLO, TritonIR, LLVM 등).

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

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

단계를 선택한 후 사용자는 지정된 플랫폼의 변환 결과를 지정된 스트림에 쓸 수 있습니다.

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

이렇게 하면 덤프가 stdout에 출력됩니다 (또는 -o가 지정된 경우 지정된 파일에 출력됨).

GPU용 기기 없는 컴파일

기기 없는 컴파일에는 GPU에 대한 액세스가 필요하지 않습니다. 기기 없는 컴파일은 GPU 액세스가 필요한 단계에서 명령줄(--xla_gpu_target_config_filename)에 GPU 사양을 지정하는 방법을 제공하여 GPU 기기의 필요성을 없앱니다.

예: GPU 기기에 액세스하지 않고 PTX 출력

hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=/xla/tools/hlo_opt/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"

자세한 GPU 사양은 /xla/tools/hlo_opt/gpu_specs에서 확인할 수 있습니다.

자동 조정

컴파일에는 컴파일 --stage에 기반한 자동 조정이 포함될 수 있습니다. 기기 없는 컴파일이 작동하려면 사용자가 --xla_gpu_autotune_level=0
로 자동 조정 기능을 사용 중지하거나 --xla_gpu_load_autotune_results_from=<filename>기존 자동 조정 결과를 로드해야 합니다 (--xla_gpu_dump_autotune_results_to=<filename>로 획득).

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_to=<myfile.pbtxt>를 사용하여 직렬화할 수 있습니다.

[hlo-opt] HLO 패스 개발 및 디버깅

# 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 도구를 사용하면 지정된 플랫폼 컴파일 단계와 관계없이 개별 패스를 실행할 수 있습니다. 이렇게 격리하면 입력 hlo 모듈에서 패스를 빠르게 실행하고 실패의 근본 원인을 정확히 파악할 수 있습니다.

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

hlo-opt 도구는 DebugOptions XLA_FLAGS도 지원합니다.

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

--list-passes 옵션을 사용하여 패스 이름 문자열을 가져옵니다.

hlo-opt --list-passes

사용자는 --passes 옵션에 패스를 두 개 이상 지정하여 자체 맞춤 파이프라인을 만들 수 있습니다.

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

새 HLO 패스 개발 지원

  1. 먼저 패스를 작성합니다.
  2. 새 패스를 hlo-opt 도구 패스 레지스트리에 등록합니다.

    RegisterPass<FooPass>(FooPassInputOptions)
    

    등록할 위치는 패스 유형에 따라 다음 중 하나를 선택합니다.
    opt_lib.cc 하드웨어에 종속되지 않는 패스
    cpu_opt.cc CPU 관련 패스입니다.
    gpu_opt.cc GPU 관련 패스입니다.
    compiled_opt.cc CPU, GPU, XPU에 공통적으로 적용되는 패스입니다.
    빌드 종속 항목을 추가하는 것을 잊지 마세요.

    모든 hlo-opt 사용자가 패스를 사용할 수 있도록 패스 등록을 PR에 포함합니다().

  3. hlo-opt 도구를 다시 빌드하고 --list-passes 옵션을 사용하여 패스 등록이 성공했는지 확인한 다음 --passes 옵션을 사용하여 패스를 실행합니다.

    $ hlo-opt --passes=foo-pass input.hlo
    
  4. 패스에 대한 단위 테스트를 작성하려면 https://openxla.org/xla/test_hlo_passes를 참고하세요.

런타임 측정 통과

대규모 모델의 경우 전체 컴파일 실행에 몇 분이 걸릴 수 있으므로 미묘한 성능 회귀를 감지하기가 어렵습니다. 반면 hlo-opt를 사용하여 개별 패스를 실행하면 정확한 성능 측정이 가능하며 새로운 코드 변경으로 인해 발생하는 실행 시간의 미세한 증가도 쉽게 감지할 수 있습니다.

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

[hlo-opt] HLO 모듈 형식 변환

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

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

HLO Text -> HLO Proto 변환

hlo-opt --emit-proto input.hlo

HLO Proto 또는 HLO Proto Binary -> HLO Text 변환

hlo-opt input.pbtxt or input.pb

[ptx-opt] 컴파일러 LLVM 모듈에서 PTX로

이 도구는 LLVMIR 최적화 파이프라인을 실행한 후 CompileToPtx를 호출합니다.

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

이 도구는 모든 경로 후에 LLVMIR을 덤프할 수도 있습니다.

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