בנייה מהמקור

במאמר הזה מוסבר איך לבנות רכיבי XLA.

אם לא שיבטתם את מאגר XLA או לא התקנתם את Bazel, כדאי לעיין בקטעים הראשונים של מדריך הפיתוח של XLA.

Linux

הגדרה

הגדרות ה-build של XLA מוגדרות בקובץ .bazelrc שבתיקיית השורש של המאגר. אפשר להשתמש בסקריפט ./configure.py כדי לשנות הגדרות נפוצות.

אם צריך לשנות את ההגדרה, מריצים את הסקריפט ./configure.py מהספרייה הבסיסית של המאגר. בסקריפט הזה יש דגלים למיקום של יחסי תלות ב-XLA ואפשרויות נוספות להגדרת בנייה (דגלי קומפיילר, למשל). פרטים נוספים זמינים בקטע סשן לדוגמה.

תמיכה במעבד

מומלץ להשתמש בקובץ אימג' מתאים של Docker, כמו ml-build, שמשמש גם בתהליכי העבודה של CI ב-XLA ב-GitHub, כדי ליצור ולבדוק את XLA. תמונת ה-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

כדי ליצור את XLA עם תמיכה ב-CPU באמצעות קונטיינר Docker, מריצים את הפקודות הבאות:

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/...

אל תשתמשו בדגל --gpus all כשמפעילים את קונטיינר Docker.
מציינים את יכולות החישוב של CUDA באופן ידני, לדוגמה:

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

פרטים נוספים על תמונות Docker של TensorFlow GPU זמינים במסמך הזה.

אפשר גם ליצור יעדי XLA עם תמיכה ב-GPU בלי Docker. מגדירים ובונים יעדים באמצעות הפקודות הבאות:

./configure.py --backend=CUDA

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

פרטים נוספים על CUDA הרמטי

פיתוח XLA עם תמיכה ב-CUDA/cuDNN באמצעות קונטיינר CI/Release של JAX

‫XLA הוא קומפיילר שנמצא בשימוש פנימי ב-JAX. הפצת JAX מתבצעת באמצעות PyPI wheels. במסמכי התיעוד בנושא אינטגרציה רציפה של JAX מוסבר איך ליצור קובצי wheel של JAX באמצעות קונטיינר Docker‏ tensorflow/ml-build:latest.

אפשר להרחיב את ההוראות האלה גם כדי ליצור יעדי XLA בתוך מאגר JAX. כך אפשר לוודא שתצורת ה-build של יעדי XLA עקבית עם תצורת ה-build של JAX/XLA. זה יכול להיות שימושי אם רוצים לשחזר תוצאות של עומסי עבודה באמצעות כלי XLA שנוצרו במקור ב-JAX.

יצירת יעדי XLA בקונטיינר JAX CI

משכפלים את מאגר 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 עם הגדרות ה-build הנדרשות, כולל תמיכה ב-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