Quy trình phát triển XLA thường tập trung vào HLO IR, đại diện cho việc tính toán chức năng tách biệt do trình biên dịch cung cấp. XLA đi kèm với nhiều công cụ dòng lệnh (như 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ô giá đối với chu kỳ lặp lại compile->modify->run nhanh chóng, vì HLO vừa có thể trực quan hoá vừa có thể tấn công, đồng thời, việc thay đổi và chạy chế độ này thường là cách nhanh nhất để hiểu và khắc phục hiệu suất hoặc hành vi XLA.
Cách dễ nhất để lấy HLO cho chương trình đượ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
Tệp 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.
Đang chạy đoạn mã HLO: run_hlo_module
Công cụ run_hlo_module hoạt động trên HLO trước khi tối ưu hoá và theo mặc định sẽ 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
Giống như mọi công cụ khác, bạn có thể sử dụng --help để lấy danh sách đầy đủ các tuỳ chọn.
Chạy đoạn mã HLO có hỗ trợ SPMD: multihost_hlo_runner
Trình chạy HLO của nhiều máy chủ lưu trữ là một công cụ rất giống với một công cụ, nhưng xin lưu ý rằng công cụ này có hỗ trợ SPMD, bao gồm cả hoạt động giao tiếp giữa các máy chủ lưu trữ. Lệnh gọi thông thường sẽ có dạng như sau:
hlo_runner_main --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo
Chạy các bước/giai đoạn của quá trình biên dịch HLO: hlo-opt
Khi gỡ lỗi hoặc tìm hiểu hoạt động của trình biên dịch, bạn nên mở rộng tính năng này cho một phần cứng cụ thể tại một điểm cụ thể trong quy trình (có thể là HLO, HLO, TritonIR hoặc LLVM được tối ưu hoá) cho đầu vào HLO (Ổn định) 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ợ tuỳ thuộc vào nền tảng (ví dụ: PTX dành riêng cho NVIDIA) và bạn có thể xem bằng lệnh --list-stages:
$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx
Sau khi chọn một giai đoạn, người dùng có thể ghi kết quả chuyển đổi cho một nền tảng nhất định vào một luồng nhất định:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
Thao tác này sẽ in tệp kết xuất tới stdout (hoặc tới một tệp cụ thể nếu bạn chỉ định -o).
Sử dụng không theo thiết bị
Trong hầu hết quá trình biên dịch, bạn không cần quyền truy cập vào GPU. Chúng ta có thể nhận được thông số kỹ thuật GPU trên dòng lệnh, chẳng hạn như dữ liệu đầu ra PTX mà không cần quyền truy cập vào trình tăng tốc:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_80.txtpb input.hlo
Thông số kỹ thuật của các GPU phổ biến được gử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"
Quá trình biên dịch không có thiết bị có thể gặp sự cố nếu cần phải tự động điều chỉnh. May mắn là chúng tôi cũng có thể cung cấp các thành phần đó trên dòng lệnh:
hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=gpu_specs/a100_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: 2
results {
device: "sm_8.0 with 42331013120B RAM, 108 cores, 1410000KHz clock, 1215000KHz mem clock, 41943040B L2$"
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ể chuyển đổi tuần tự cơ sở dữ liệu tự động điều chỉnh bằng XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>
Chạy một gói trình biên dịch đơn
Các cờ từ XLA_FLAGS cũng được hỗ trợ, vì vậy, bạn có thể sử dụng công cụ này để kiểm thử việc chạy một lượt truyền:
hlo-opt --platform=CUDA --stage=hlo --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo