Công cụ XLA

Quy trình phát triển XLA thường tập trung vào HLO IR, đại diện cho quá trình tính toán chức năng riêng biệt được cung cấp cho trình biên dịch. XLA đi kèm với nhiều công cụ dòng lệnh (được mô tả bên dưới) sử dụng HLO và chạy HLO hoặc cung cấp giai đoạn biên dịch trung gian. Việc sử dụng các công cụ như vậy là vô giá đối với chu kỳ lặp lại nhanh chóng compile->modify->run, vì HLO có thể hình dung và có thể hack được, đồng thời việc thay đổi và chạy HLO một cách lặp đi lặp lại thường là cách nhanh nhất để hiểu và khắc phục hiệu suất hoặc hành vi của XLA.

Cách dễ nhất để lấy HLO cho một chương trình đang được biên dịch bằng XLA thường là sử dụng biến môi trường XLA_FLAGS:

$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point

biến này lưu trữ tất cả các tệp HLO trước khi tối ưu hoá trong thư mục được chỉ định, cùng với nhiều cấu phần phần mềm hữu ích khác.

[`run_hlo_module`] Chạy mô-đun HLO

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

Công cụ run_hlo_module hoạt động trên HLO trước khi tối ưu hoá và theo mặc định, công cụ này sẽ gói quá trình biên dịch, chạy và so sánh với quá trình triển khai trình thông dịch tham chiếu. Ví dụ: lệnh gọi thông thường để chạy tệp đầu vào computation.hlo trên GPU NVIDIA và kiểm tra tính chính xác của tệp đó là:

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

Chạy nhiều mô-đun HLO

Lệnh gọi có nhiều mô-đun HLO được hỗ trợ cho run_hlo_module. Cách chạy tất cả các mô-đun hlo từ một thư mục:

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

[`multihost_hlo_runner`] Chạy mô-đun HLO có hỗ trợ SPMD

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

Trình chạy HLO nhiều máy chủ là một công cụ rất tương tự, với lưu ý là công cụ này hỗ trợ SPMD, bao gồm cả giao tiếp giữa các máy chủ. Xem Trình chạy HLO nhiều máy chủ để biết thông tin chi tiết.

Chạy nhiều mô-đun HLO có hỗ trợ SPMD

Tương tự như run_hlo_module, multihost_hlo_runner cũng hỗ trợ lời gọi có nhiều mô-đun.

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

[`hlo-opt`] Biên dịch mô-đun HLO

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

Khi gỡ lỗi hoặc tìm hiểu cách hoạt động của trình biên dịch, bạn thường nên lấy bản mở rộng cho một phần cứng cụ thể tại một thời điểm cụ thể trong quy trình (có thể là HLO, HLO được tối ưu hoá, TritonIR hoặc LLVM), cho một đầu vào HLO hoặc StableHLO nhất định.

hlo-opt hỗ trợ nhiều giai đoạn đầu ra: có thể là PTX, HLO sau khi tối ưu hoá, LLVM IR trước khi tối ưu hoá hoặc TritonIR. Tập hợp chính xác các giai đoạn được hỗ trợ phụ thuộc vào nền tảng (ví dụ: PTX dành riêng cho NVIDIA) và có thể xem bằng lệnh --list-stages:

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

Sau khi chọn một giai đoạn, người dùng có thể ghi kết quả của quá trình chuyển đổi cho một nền tảng nhất định vào một luồng nhất định:

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

lệnh này sẽ in kết quả kết xuất vào stdout (hoặc vào một tệp nhất định nếu -o được chỉ định).

Biên dịch không có thiết bị cho GPU

Quá trình biên dịch không có thiết bị không cần quyền truy cập vào GPU. Quá trình biên dịch không có thiết bị cung cấp một cách để chỉ định thông số kỹ thuật GPU trên dòng lệnh (--xla_gpu_target_config_filename) cho các giai đoạn cần có quyền truy cập vào GPU, loại bỏ nhu cầu về thiết bị GPU.

Ví dụ: Đầu ra PTX mà không có quyền truy cập vào thiết bị GPU:

hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=/xla/tools/hlo_opt/gpu_specs/a100_pcie_80.txtpb input.hlo

Thông số kỹ thuật cho các GPU phổ biến được vận chuyển cùng với trình biên dịch và tệp được cung cấp là quá trình chuyển đổi tuần tự chuỗi của 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"

Bạn có thể tìm thấy thêm thông số kỹ thuật GPU tại /xla/tools/hlo_opt/gpu_specs

Tự động điều chỉnh

Đôi khi, quá trình biên dịch có thể liên quan đến việc tự động điều chỉnh dựa trên --stage biên dịch. Để quá trình biên dịch không có thiết bị hoạt động, người dùng cần
tắt tính năng tự động điều chỉnh bằng --xla_gpu_autotune_level=0
hoặc
tải kết quả tự động điều chỉnh đã có từ trước bằng --xla_gpu_load_autotune_results_from=<filename> (lấy được bằng --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

Tệp tự động điều chỉnh là quá trình chuyển đổi tuần tự văn bản của autotune_results.proto, với ví dụ như sau:

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
    }
  }
}

Bạn có thể tuần tự hoá cơ sở dữ liệu tự động điều chỉnh bằng cách sử dụng XLA_FLAGS=--xla_gpu_dump_autotune_results_to=<myfile.pbtxt>

[`hlo-opt`] Phát triển và gỡ lỗi HLO Pass

# 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>

Công cụ hlo-opt cho phép thực thi các lượt truyền riêng lẻ độc lập với các giai đoạn biên dịch nền tảng đã cho. Việc tách biệt này giúp nhanh chóng chạy các lượt truyền trên mô-đun hlo đầu vào và xác định nguyên nhân gốc rễ của lỗi.

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

Công cụ hlo-opt cũng hỗ trợ DebugOptions XLA_FLAGS.

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

Sử dụng tuỳ chọn --list-passes để lấy chuỗi tên lượt truyền.

hlo-opt --list-passes

Người dùng có thể tạo quy trình tuỳ chỉnh của riêng mình bằng cách chỉ định nhiều lượt truyền cho tuỳ chọn --passes.

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

Hỗ trợ phát triển HLO Pass mới

Trước tiên, hãy viết lượt truyền.
Đăng ký lượt truyền mới vào sổ đăng ký lượt truyền công cụ hlo-opt.
```
RegisterPass<FooPass>(FooPassInputOptions)
```
Dựa trên loại lượt truyền, hãy chọn một trong các vị trí sau để đăng ký:
opt_lib.cc Lượt truyền độc lập với phần cứng.
cpu_opt.cc Lượt truyền dành riêng cho CPU.
gpu_opt.cc Lượt truyền dành riêng cho GPU.
compiled_opt.cc Lượt truyền phổ biến cho CPU, GPU, XPU.
Đừng quên thêm phần phụ thuộc bản dựng.

Đưa quá trình đăng ký lượt truyền vào PR(ví dụ) để tất cả người dùng hlo-opt đều có thể sử dụng lượt truyền.
Xây dựng lại công cụ hlo-opt, xác thực quá trình đăng ký lượt truyền thành công bằng cách sử dụng tuỳ chọn --list-passes, sau đó sử dụng tuỳ chọn --passes để chạy lượt truyền.
```
$ hlo-opt --passes=foo-pass input.hlo
```
Viết mã kiểm thử đơn vị cho lượt truyền? Hãy tham khảo https://openxla.org/xla/test_hlo_passes để biết thêm thông tin.

Đo lường thời gian chạy lượt truyền

Đối với các mô hình lớn, quá trình biên dịch đầy đủ có thể mất vài phút, khiến việc phát hiện các hồi quy hiệu suất tinh tế trở nên khó khăn. Ngược lại, các lượt chạy lượt truyền riêng lẻ bằng cách sử dụng hlo-opt cho phép đo lường độ chính xác hiệu suất và dễ dàng phát hiện ngay cả khi thời gian thực thi tăng nhẹ do các thay đổi mã mới.

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

[`hlo-opt`] Chuyển đổi định dạng mô-đun HLO

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

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

Chuyển đổi `HLO Text` -> `HLO Proto`

hlo-opt --emit-proto input.hlo

Chuyển đổi `HLO Proto` hoặc `HLO Proto Binary` -> `HLO Text`

hlo-opt input.pbtxt or input.pb

[`ptx-opt`] Trình biên dịch mô-đun LLVM xuống PTX

Công cụ này sẽ chạy quy trình tối ưu hoá LLVMIR, sau đó gọi CompileToPtx.

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

Công cụ này cũng có thể kết xuất LLVMIR sau mỗi đường dẫn.

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

[`isolate_hlo`] Tách biệt các hướng dẫn HLO có vấn đề

Nếu bạn có một kết xuất HLO lớn và nghi ngờ một hướng dẫn hoặc phần cụ thể trong mô-đun HLO đang gây ra sự cố, bạn có thể sử dụng công cụ isolate_hlo.

Công cụ này trích xuất một hướng dẫn HLO duy nhất (và ngữ cảnh cần thiết của hướng dẫn đó) vào một mô-đun HLO mới, nhỏ hơn. Điều này cực kỳ hữu ích để tạo một trình tái tạo tối thiểu ở cấp trình biên dịch.

Tài liệu và nguồn: Công cụ isolate_hlo có trong kho lưu trữ OpenXLA. Xem thư mục xla/tools trong mã nguồn XLA.
Cách sử dụng: Xây dựng công cụ từ cây nguồn XLA. Công cụ này thường lấy một tệp mô-đun HLO đầu vào (văn bản hoặc proto), tên của hướng dẫn cần trích xuất và đường dẫn tệp đầu ra.
```
# Example usage after building XLA:
# ./build/tools/isolate_hlo --input=module.hlo --instruction_name=fusion.123 \
#   --output=isolated_fusion.123.hlo --input_format=txt --output_format=long_txt
```
Tham khảo thông báo trợ giúp của công cụ (--help) để biết các cờ và tuỳ chọn định dạng cụ thể.

Công cụ XLA Sử dụng bộ sưu tập để sắp xếp ngăn nắp các trang Lưu và phân loại nội dung dựa trên lựa chọn ưu tiên của bạn.

[run_hlo_module] Chạy mô-đun HLO

Chạy nhiều mô-đun HLO

[multihost_hlo_runner] Chạy mô-đun HLO có hỗ trợ SPMD

Chạy nhiều mô-đun HLO có hỗ trợ SPMD

[hlo-opt] Biên dịch mô-đun HLO