Công cụ XLA

Quy trình phát triển XLA thường tập trung vào IR HLO, đại diện cho hoạt động 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 một giai đoạn biên dịch trung gian. Việc sử dụng các công cụ như vậy là vô cùng quý giá đối với một chu kỳ lặp lại nhanh chóng compile->modify->run, vì HLO vừa có thể hình dung vừa có thể hack, đồng thời việc thay đổi và chạy nó theo cách 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

thao tác này sẽ 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 các 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ẽ kết hợp quá trình biên dịch, chạy và so sánh với việc triển khai trình thông dịch tham chiếu. Ví dụ: lệnh gọi thông thường để chạy một 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

run_hlo_module hỗ trợ lệnh gọi có nhiều mô-đun HLO. 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 các 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 ý rằng công cụ này hỗ trợ SPMD, bao gồm cả giao tiếp giữa các máy chủ. Hãy 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ệnh gọi với 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 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 (cho dù đó là HLO, HLO được tối ưu hoá, TritonIR hay 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 là dành riêng cho NVIDIA) và có thể được 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 lượt chuyển đổi cho một nền tảng nhất định vào một luồng dữ liệu nhất định:

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

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

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

Quá trình biên dịch không cần 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 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 phân phối 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ể xem thêm thông số kỹ thuật của GPU tại /xla/tools/hlo_opt/gpu_specs

Tự động dò

Đôi khi, bản tổng hợp có thể sử dụng tính năng tự động điều chỉnh dựa trên bản tổng hợp --stage. Để quá trình biên dịch không cần 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> (thu đượ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 autotune là quá trình chuyển đổi tuần tự văn bản của autotune_results.proto, với ví dụ có dạng 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 đường chuyền riêng lẻ mà không phụ thuộc vào 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 trên mô-đun hlo đầu vào và xác định chính xác 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

Dùng tuỳ chọn --list-passes để lấy chuỗi tên thẻ và vé.

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 đường chuyền cho lựa chọn --passes.

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

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

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

Đưa quy trình đăng ký thẻ và vé vào thông cáo báo chí(ví dụ) để tất cả người dùng hlo-opt đều có thể sử dụng thẻ và vé.
Tạo lại công cụ hlo-opt, xác thực việc đăng ký thẻ thành công bằng lựa chọn --list-passes rồi dùng lựa chọn --passes để chạy thẻ.
```
$ hlo-opt --passes=foo-pass input.hlo
```
Viết các bài 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 chi tiết.

Đo lường thời gian chạy của thẻ và vé

Đố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 bạn khó phát hiện được những điểm giảm hiệu suất nhỏ. Ngược lại, các lần chạy riêng lẻ bằng hlo-opt cho phép đo lường hiệu suất một cách chính xác và dễ dàng phát hiện ngay cả những mức tăng nhỏ về thời gian thực thi do các thay đổi về mã mới gây ra.

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 LLVM Module xuống PTX

Công cụ này sẽ chạy quy trình tối ưu hoá LLVMIR rồi 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>

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 các mô-đun HLO

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

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