소스에서 빌드

이 문서에서는 XLA 구성요소를 빌드하는 방법을 설명합니다.

XLA 저장소를 클론하지 않았거나 Bazel을 설치하지 않은 경우 XLA 개발자 가이드의 초기 섹션을 확인하세요.

Linux

구성

XLA 빌드는 저장소의 루트 디렉터리에 있는 .bazelrc 파일로 구성됩니다. ./configure.py 스크립트를 사용하여 일반적인 설정을 조정할 수 있습니다.

구성을 변경해야 하는 경우 저장소의 루트 디렉터리에서 ./configure.py 스크립트를 실행합니다. 이 스크립트에는 XLA 종속 항목 위치와 추가 빌드 구성 옵션 (예: 컴파일러 플래그)을 위한 플래그가 있습니다. 자세한 내용은 샘플 세션 섹션을 참고하세요.

CPU 지원

XLA를 빌드하고 테스트하려면 GitHub의 XLA CI 워크플로에서도 사용되는 ml-build와 같은 적절한 Docker 이미지를 사용하는 것이 좋습니다. ml-build 이미지에는 Clang 18이 사전 설치되어 있습니다.

docker run -itd --rm \
--name xla \
-w /xla \
-v $PWD:/xla \
us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest \
bash

Docker 컨테이너를 사용하여 다음 명령어를 실행하여 CPU 지원으로 XLA를 빌드할 수 있습니다.

docker exec xla ./configure.py --backend=CPU

docker exec xla bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

Docker를 사용하지 않고 CPU 지원으로 XLA 타겟을 빌드하려면 Clang을 설치해야 합니다. XLA는 현재 CI에서 Clang 18로 빌드되지만 이전 버전도 작동해야 합니다.

타겟을 구성하고 빌드하려면 다음 명령어를 실행합니다.

./configure.py --backend=CPU

bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

GPU 지원

위에서 언급한 것과 동일한 Docker 컨테이너를 사용하여 GPU 지원으로 XLA를 빌드하는 것이 좋습니다.

모든 GPU에 액세스할 수 있는 Docker 컨테이너를 시작하려면 다음 명령어를 실행합니다.

docker run -itd --rm \
  --gpus all \
  --name xla_gpu \
  -w /xla \
  -v $PWD:/xla \
  us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest \
  bash

GPU 지원으로 XLA를 빌드하려면 다음 명령어를 실행하세요.

docker exec xla_gpu ./configure.py --backend=CUDA

docker exec xla_gpu bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

Docker 컨테이너를 시작할 때는 --gpus all 플래그를 사용하지 마세요.
CUDA 컴퓨팅 기능을 수동으로 지정합니다. 예를 들면 다음과 같습니다.

docker exec xla_gpu ./configure.py --backend=CUDA \
  --cuda_compute_capabilities="9.0"

TensorFlow의 GPU Docker 이미지에 관한 자세한 내용은 이 문서를 참고하세요.

Docker 없이 GPU 지원으로 XLA 타겟을 빌드할 수도 있습니다. 다음 명령어를 사용하여 타겟을 구성하고 빌드합니다.

./configure.py --backend=CUDA

bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

Hermetic CUDA에 관한 자세한 내용은 이 문서를 참고하세요.

JAX CI/릴리스 컨테이너를 사용하여 CUDA/cuDNN 지원으로 XLA 빌드

XLA는 JAX에서 내부적으로 사용하는 컴파일러입니다. JAX는 PyPI 휠을 통해 배포됩니다. JAX 지속적 통합 문서에서는 tensorflow/ml-build:latest Docker 컨테이너를 사용하여 JAX 휠을 빌드하는 방법을 설명합니다.

이러한 안내를 확장하여 JAX 컨테이너 내에서 XLA 타겟을 빌드할 수도 있습니다. 이렇게 하면 XLA 타겟의 빌드 구성이 JAX/XLA 빌드 구성과 일치하므로 원래 JAX에서 생성된 XLA 도구를 사용하여 워크로드 결과를 재현하려는 경우에 유용할 수 있습니다.

JAX CI 컨테이너에서 XLA 타겟 빌드

JAX 저장소를 클론하고 'jax' 디렉터리로 이동합니다.

git clone <a href="https://github.com/jax-ml/jax.git">https://github.com/jax-ml/jax.git</a>

cd jax

다음을 실행하여 JAX CI/Release Docker 컨테이너를 시작합니다.

./ci/utilities/run_docker_container.sh

그러면 'jax'라는 Docker 컨테이너가 시작됩니다.

다음을 사용하여 컨테이너 내에서 jax-cuda-plugin 타겟을 빌드합니다.

docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin

이렇게 하면 CUDA/cuDNN 지원을 비롯한 필수 빌드 구성이 포함된 .jax_configure.bazelrc 파일이 생성됩니다.

컨테이너 내부의 대화형 셸에 액세스합니다.

docker exec -ti jax /bin/bash

이제 컨테이너 내의 /jax 디렉터리에 있습니다.

다음 명령어를 사용하여 XLA 타겟을 빌드합니다(예시 참고).

/usr/local/bin/bazel build \
  --config=cuda_libraries_from_stubs \
  --verbose_failures=true \
  @xla//xla/tools/multihost_hlo_runner:hlo_runner_main

선택적으로 HERMETIC 환경을 덮어쓸 수 있습니다(예:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"

필요한 경우 결과 아티팩트를 /jax/dist에 복사하여 호스트 OS에서 액세스합니다.

cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
  ./dist/

대화형 셸을 종료합니다.

exit