เครื่องมือ XLA

โดยปกติแล้ว เวิร์กโฟลว์การพัฒนา XLA จะมุ่งเน้นไปที่ IR ของ HLO ซึ่งแสดงถึงการคำนวณฟังก์ชันที่แยกออกมาซึ่งส่งไปยังคอมไพเลอร์ XLA มาพร้อมกับเครื่องมือบรรทัดคำสั่งหลายรายการ (อธิบายไว้ด้านล่าง) ซึ่งใช้ HLO และเรียกใช้ หรือให้ ขั้นตอนการคอมไพล์ระดับกลาง การใช้เครื่องมือดังกล่าวมีประโยชน์อย่างมากสำหรับรอบการทำซ้ำที่รวดเร็ว compile->modify->run เนื่องจาก HLO สามารถแสดงภาพและ แก้ไขได้ และการเปลี่ยนแปลงและการเรียกใช้ซ้ำๆ มักเป็นวิธีที่เร็วที่สุดในการ ทำความเข้าใจและแก้ไขประสิทธิภาพหรือลักษณะการทำงานของ XLA

วิธีที่ง่ายที่สุดในการรับ HLO สำหรับโปรแกรมที่คอมไพล์ด้วย XLA มักจะเป็นการใช้ตัวแปรสภาพแวดล้อม 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 ก่อนการเพิ่มประสิทธิภาพ และโดยค่าเริ่มต้น จะรวมการคอมไพล์ การเรียกใช้ และการเปรียบเทียบกับการใช้งานอินเทอร์พรีเตอร์อ้างอิง ตัวอย่างเช่น การเรียกใช้ปกติเพื่อเรียกใช้ไฟล์อินพุต computation.hlo ใน NVIDIA GPU และเพื่อตรวจสอบความถูกต้องคือ

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

เรียกใช้โมดูล HLO หลายโมดูล

run_hlo_module รองรับการเรียกใช้ที่มีโมดูล HLO หลายโมดูล หากต้องการเรียกใช้ โมดูล HLO ทั้งหมดจากไดเรกทอรี ให้ทำดังนี้

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

[`multihost_hlo_runner`] เรียกใช้โมดูล HLO ด้วยการรองรับ SPMD

# 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

เรียกใช้โมดูล HLO หลายโมดูลด้วยการรองรับ SPMD

multihost_hlo_runner ยังรองรับการเรียกใช้ที่มีหลายโมดูลเช่นเดียวกับ run_hlo_module

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, HLO ที่เพิ่มประสิทธิภาพ, TritonIR หรือ LLVM) สำหรับอินพุต HLO หรือ StableHLO ที่กำหนดมักจะมีประโยชน์

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

ตัวอย่าง: เอาต์พุต PTX ที่ไม่มีสิทธิ์เข้าถึงอุปกรณ์ GPU

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 option เพื่อรับสตริงชื่อบัตร

hlo-opt --list-passes

ผู้ใช้สามารถสร้างไปป์ไลน์ที่กำหนดเองได้โดยการระบุการส่งผ่านมากกว่า 1 รายการไปยังตัวเลือก --passes

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

ช่วยในการพัฒนาบัตร HLO ใหม่

ก่อนอื่นให้เขียนบัตร
ลงทะเบียนบัตรใหม่กับhlo-optการลงทะเบียนบัตรเครื่องมือ
```
RegisterPass<FooPass>(FooPassInputOptions)
```
เลือกตำแหน่งใดตำแหน่งหนึ่งต่อไปนี้เพื่อลงทะเบียนบัตรตามประเภทบัตร
opt_lib.cc บัตรที่ใช้ได้กับฮาร์ดแวร์ทุกประเภท
cpu_opt.cc บัตรเฉพาะสำหรับ CPU
gpu_opt.cc บัตรเฉพาะ GPU
compiled_opt.cc ส่งไปยัง CPU, GPU, XPU
อย่าลืมเพิ่มการอ้างอิงบิลด์

รวมการลงทะเบียนบัตรเป็นส่วนหนึ่งของ PR(ตัวอย่าง) เพื่อให้ผู้ใช้ hlo-opt ทุกคนใช้บัตรได้
สร้างเครื่องมือ hlo-opt ใหม่ ตรวจสอบการลงทะเบียนบัตรที่สำเร็จโดยใช้ตัวเลือก --list-passes แล้วใช้ตัวเลือก --passes เพื่อเรียกใช้บัตร
```
$ hlo-opt --passes=foo-pass input.hlo
```
หากต้องการเขียน Unit Test สำหรับการส่งผ่าน โปรดดูรายละเอียดเพิ่มเติมที่ https://openxla.org/xla/test_hlo_passes

การวัดรันไทม์ของบัตร

สําหรับโมเดลขนาดใหญ่ การคอมไพล์แบบเต็มอาจใช้เวลาถึง 2-3 นาที ซึ่งทําให้ตรวจหาการถดถอยของประสิทธิภาพที่ละเอียดอ่อนได้ยาก ในทางตรงกันข้าม การเรียกใช้การทดสอบแต่ละรายการโดยใช้ 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 Module ลงไปจนถึง 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>

[run_hlo_module] เรียกใช้โมดูล HLO

เรียกใช้โมดูล HLO หลายโมดูล

[multihost_hlo_runner] เรียกใช้โมดูล HLO ด้วยการรองรับ SPMD

เรียกใช้โมดูล HLO หลายโมดูลด้วยการรองรับ SPMD

[hlo-opt] คอมไพล์โมดูล HLO