Sử dụng 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 phép tính 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 (mô tả bên dưới) sử dụng HLO và chạy 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 compile->modify->run nhanh, vì HLO vừa có thể hình dung vừa có thể xâm nhập, đồng thời việc thay đổi và chạy 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 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

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

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, công cụ này sẽ biên dịch, chạy và so sánh các gói 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 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ư tất cả các công cụ, bạn có thể sử dụng --help để lấy danh sách đầy đủ các tuỳ chọn.

Chạy các đoạn mã HLO có hỗ trợ SPMD: multihost_hlo_runner

Trình chạy HLO nhiều máy chủ là một công cụ rất giống, 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ủ. Hãy xem phần Trình chạy HLO nhiều máy chủ để biết thông tin chi tiết.

Phát lại nhiều HLO

Lệnh gọi với nhiều mô-đun được hỗ trợ cho cả run_hlo_modulehlo_runner_main, thường thuận tiện để phát lại tất cả mô-đun trong thư mục báo lỗi:

$ hlo_runner_main /dump/*before_optimizations*

Chạy các lượt/giai đoạn biên dịch HLO: hlo-opt

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 đ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 (Ổ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 các giai đoạn chính xác được hỗ trợ phụ thuộc vào nền tảng (ví dụ: PTX dành riêng cho NVIDIA) và có thể được 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

lệnh này sẽ in tệp báo lỗi vào stdout (hoặc vào một tệp nhất định nếu bạn chỉ định -o).

Sử dụng không cần thiết bị

Bạn không cần quyền truy cập vào GPU cho hầu hết các quá trình biên dịch. Bằng cách chỉ định thông số kỹ thuật GPU trên dòng lệnh, chúng ta có thể nhận được kết quả PTX mà không cầ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_pcie_80.txtpb input.hlo

Thông số kỹ thuật của 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"

Quá trình biên dịch không cần thiết bị có thể gặp vấn đề nếu cần tự động điều chỉnh. May mắn thay, chúng ta cũng có thể cung cấp các thông tin đó trên dòng lệnh:

$ 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ể chuyển đổi tuần tự 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_t=<myfile.pbtxt>

Chạy một lượt biên dịch

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ần:

$ hlo-opt --platform=CUDA --stage=hlo --passes=algebraic_simplifer input.hlo