สร้างจากซอร์ส

เอกสารนี้อธิบายวิธีสร้างคอมโพเนนต์ XLA

หากไม่ได้โคลนที่เก็บ XLA หรือติดตั้ง Bazel โปรดดูส่วนเริ่มต้นของคู่มือนักพัฒนาซอฟต์แวร์ XLA

Linux

กำหนดค่า

ระบบจะกำหนดค่าบิลด์ XLA โดยไฟล์ .bazelrc ในรูทไดเรกทอรีของที่เก็บ คุณสามารถใช้สคริปต์ ./configure.py เพื่อปรับการตั้งค่าทั่วไปได้

หากต้องการเปลี่ยนการกำหนดค่า ให้เรียกใช้สคริปต์ ./configure.py จาก ไดเรกทอรีรากของที่เก็บ สคริปต์นี้มีแฟล็กสำหรับตำแหน่งของ XLA dependencies และตัวเลือกการกำหนดค่าบิลด์เพิ่มเติม (เช่น แฟล็กคอมไพเลอร์) ดูรายละเอียดได้ที่ส่วนเซสชันตัวอย่าง

การรองรับ CPU

เราขอแนะนำให้ใช้ Docker Image ที่เหมาะสม เช่น ml-build ซึ่งใช้ในเวิร์กโฟลว์ CI ของ XLA ใน GitHub ด้วย สำหรับการสร้างและทดสอบ<0x0AXLA อิมเมจ 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 คุณจะสร้าง XLA ที่รองรับ CPU ได้โดยเรียกใช้คำสั่งต่อไปนี้

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

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

หากต้องการสร้างเป้าหมาย XLA ที่รองรับ CPU โดยไม่ใช้ Docker คุณจะต้องติดตั้ง Clang ปัจจุบัน XLA สร้างขึ้นด้วย Clang 18 ใน CI แต่เวอร์ชันก่อนหน้าก็ควรใช้งานได้เช่นกัน

หากต้องการกำหนดค่าและสร้างเป้าหมาย ให้เรียกใช้คำสั่งต่อไปนี้

./configure.py --backend=CPU

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

การรองรับ GPU

เราขอแนะนำให้ใช้คอนเทนเนอร์ Docker เดียวกันกับที่กล่าวถึงข้างต้นเพื่อสร้าง XLA ที่รองรับ GPU

หากต้องการเริ่มคอนเทนเนอร์ Docker ที่มีสิทธิ์เข้าถึง GPU ทั้งหมด ให้เรียกใช้คำสั่งต่อไปนี้

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

หากต้องการสร้าง XLA ที่รองรับ GPU ให้เรียกใช้คำสั่งต่อไปนี้

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

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

อย่าใช้ Flag --gpus all เมื่อเริ่มต้นคอนเทนเนอร์ Docker
ระบุความสามารถในการประมวลผล CUDA ด้วยตนเอง เช่น

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

ดูรายละเอียดเพิ่มเติมเกี่ยวกับอิมเมจ Docker ของ GPU ของ TensorFlow ได้ในเอกสารนี้

คุณยังสร้างเป้าหมาย XLA ที่รองรับ GPU ได้โดยไม่ต้องใช้ Docker ด้วย กำหนดค่าและ สร้างเป้าหมายโดยใช้คำสั่งต่อไปนี้

./configure.py --backend=CUDA

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

ดูรายละเอียดเพิ่มเติมเกี่ยวกับ CUDA แบบปิดได้ในเอกสารนี้

สร้าง XLA ที่รองรับ CUDA/cuDNN โดยใช้คอนเทนเนอร์ CI/การเผยแพร่ของ JAX

XLA เป็นคอมไพเลอร์ที่ JAX ใช้ภายใน JAX มีการเผยแพร่ผ่าน Wheel ของ PyPI เอกสารประกอบการรวมอย่างต่อเนื่องของ JAX อธิบายวิธีสร้าง Wheel ของ JAX โดยใช้คอนเทนเนอร์ Docker tensorflow/ml-build:latest

เราสามารถขยายวิธีการเหล่านี้เพื่อสร้างเป้าหมาย XLA ภายในคอนเทนเนอร์ JAX ได้เช่นกัน ซึ่งจะช่วยให้มั่นใจได้ว่าการกำหนดค่าบิลด์ของเป้าหมาย XLA จะสอดคล้องกับการกำหนดค่าบิลด์ JAX/XLA ซึ่งอาจเป็นประโยชน์หากเราต้องการสร้างผลลัพธ์ของภาระงานซ้ำโดยใช้เครื่องมือ XLA ที่สร้างขึ้นใน JAX

สร้างเป้าหมาย XLA ในคอนเทนเนอร์ CI ของ JAX

โคลนที่เก็บ JAX แล้วไปที่ไดเรกทอรี "jax"

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

cd jax

เริ่มคอนเทนเนอร์ Docker ของ JAX CI/Release โดยการเรียกใช้คำสั่งต่อไปนี้

./ci/utilities/run_docker_container.sh

ซึ่งจะเริ่มต้นคอนเทนเนอร์ Docker ที่ชื่อ "jax"

สร้างเป้าหมาย jax-cuda-plugin ภายในคอนเทนเนอร์โดยใช้คำสั่งต่อไปนี้

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

ซึ่งจะสร้างไฟล์ .jax_configure.bazelrc ที่มีการกำหนดค่าบิลด์ที่จำเป็น รวมถึงการรองรับ CUDA/cuDNN

เข้าถึงเชลล์แบบอินเทอร์แอกทีฟภายในคอนเทนเนอร์

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 envs ก็ได้ เช่น

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"

คัดลอกอาร์ติแฟกต์ที่ได้ไปยัง /jax/dist เพื่อเข้าถึงจากระบบปฏิบัติการโฮสต์ หากจำเป็น

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

ออกจากเชลล์แบบอินเทอร์แอกทีฟ

exit