Phiên dịch viên StableHLO

Mục tiêu chính của trình phiên dịch StableHLO là cung cấp tài liệu tham khảo triển khai ngữ nghĩa của bộ mã hoá ổn định HLO theo đặc điểm kỹ thuật. Mục tiêu phụ là để việc triển khai bám sát chặt chẽ thông số kỹ thuật, ưu tiên tính dễ đọc hơn hiệu suất, để làm rõ hơn vào ngữ nghĩa của ngay cả những thao tác liên quan nhất như Convolution, Gather/Scatter và DotGeneral.

Hiện tại, OpenXLA hỗ trợ giải thích 91 trên 96 được suy đoán Hoạt động ổn địnhHLO. 3 hoạt động còn lại (FftOp, RngOp, RngBitGeneratorOp) có ngữ nghĩa của chúng được ghi lại bằng spec.md và có hoàn thành cuộc điều tra ban đầu về cách thức tiếp tục (xem status.md để xem danh sách đầy đủ các hoạt động và trạng thái mới nhất của hoạt động đó). Các vòng chung kết này các điểm cải tiến sẽ được triển khai trên cơ sở cộng đồng khi cần.

Phạm vi

Chúng tôi đã phân loại hoạt động ổn định HLO thành 11 danh mục bao gồm 118 hoạt động trong tổng số (xem Phụ lục). Cách triển khai tệp đối chiếu luồng công việc sắp xếp công việc triển khai phiên dịch cho 100% hoạt động ổn định HLO như xác định trong thông số kỹ thuật StableHLO. Chúng tôi lên kế hoạch hoàn thành tất cả hoặc gần như toàn bộ công việc trong luồng công việc này trong StableHLO phiên bản 1.0. Trong số 96 hoạt động hiện có thông số kỹ thuật, chúng ta có thể diễn giải 91 hoạt động thông qua OpenXLA (xem Các trường hợp đặc biệt để biết 5 trường hợp còn lại).

Thông số kỹ thuật

Yêu cầu chính đối với phiên dịch viên là phải có khả năng trao đổi 1:1 với quy cách Thông số kỹ thuật cho phép chuẩn hoá trình phiên dịch trong các hoạt động tương tự dẫn đến việc triển khai trình phiên dịch theo mô-đun ở chất lượng cao.

Các trường hợp đặc biệt

Khác

Danh mục này có các hoạt động phân tách có tương lai không rõ ràng tại thời điểm này. Có là ba hoạt động được chỉ định trong danh mục này mà trình thông dịch không hỗ trợ khoảnh khắc:

• FftOp
• RngOp
• RngBitGeneratorOp

FftOp được phân loại là Hoạt động khác, nhưng không giống như các hoạt động khác trong danh mục này, hoạt động này không có đường truyền mở rộng và việc hỗ trợ tính năng này trong StableHLO là WIP.

RngOp và RngBitGeneratorOp có thể được phân rã thành các hoạt động MHLO, nhưng quá trình phân tách giới thiệu một XlaRngGetAndUpdateStateOp, một thuộc tính cụ thể của MHLO tối đa Việc hỗ trợ diễn giải hai hoạt động này là WIP.

Công cụ để chuyển đổi các hoạt động còn lại trong danh mục này sang hoạt động ổn định hỗ trợ trình phiên dịch nằm trong hlo_expand_main.cc.

Không có trong HLO

Ngoài các hoạt động được chỉ định, danh mục này bao gồm 8 hoạt động chưa được chỉ định (xem StableHLO Ops) (Danh mục hoạt động ổn định) theo kế hoạch đã chuyển ra khỏi Ổn định HLO. Hầu hết các hoạt động này đều có thẻ và vé trong mhlo sang chuyển đổi chúng thành hoạt động tương đương StableHLO.

Công cụ để chuyển đổi các hoạt động còn lại trong danh mục này thành các hoạt động ổn định tương đương mà trình thông dịch hỗ trợ nằm trong mlir-hlo-opt.cc.

Lượng tử hoá

Hỗ trợ phiên dịch cho thao tác stablehlo.constant với loại lượng tử hoá là không được hỗ trợ và theo dõi qua #1691.

Hướng dẫn sử dụng

Xây dựng người phiên dịch tài liệu tham khảo

Bạn có thể tạo và kiểm thử trình thông dịch thông qua Bazel hoặc CMake (ưu tiên). Dùng cho toàn bộ Vui lòng xem hướng dẫn tại README.md.

Bazel:

bazel build //...

CMake:

mkdir -p build && cd build

cmake .. -GNinja \
  -DLLVM_ENABLE_LLD="$LLVM_ENABLE_LLD" \
  -DCMAKE_BUILD_TYPE=Release \
  -DLLVM_ENABLE_ASSERTIONS=On \
  -DMLIR_DIR=${PWD}/../llvm-build/lib/cmake/mlir

Để chạy trình phiên dịch, chúng tôi có một công cụ dịch để dịch các chương trình StableHLO được viết bằng MLIR.

stablehlo-translate --interpret <path/to/program>

Phương ngữ của phiên dịch

Phương ngữ Interpreter chứa nhiều hoạt động tiện ích liên quan đến phiên dịch. Cụ thể, interpreter.run_parallel (xem InterpreterOps.td cho ngữ nghĩa hoạt động và ví dụ về cách sử dụng) op cho phép diễn giải Hoạt động phân phối, v.v. phần mềm tiện ích được lập kế hoạch bổ sung dựa trên nhu cầu của cộng đồng.

Phương ngữ kiểm tra

Phương ngữ Check được dùng để so sánh các giá trị thời gian chạy của phiên dịch với các giá trị dự kiến giá trị. Bạn có thể kiểm tra dữ liệu đầu ra của chương trình ổn định thông qua nhiều hoạt động kiểm tra (xem CheckOps.td cho ngữ nghĩa vận hành và ví dụ về cách sử dụng).

Viết chương trình kiểm thử

Chúng tôi sử dụng công cụ lit của LLVM để chạy và so sánh với tệp được tạo để so sánh với kết quả của trình phiên dịch (xem stablehlo/tests/interpret ví dụ: kiểm thử).

Đang kiểm tra AddOp (mẫu từ interpret_add.mlir):

// RUN: stablehlo-translate --interpret %s

func.func @add_op_scalar() {
  %0 = stablehlo.constant dense<2> : tensor<i4>
  %1 = stablehlo.constant dense<3> : tensor<i4>
  %2 = stablehlo.add %0, %1 : tensor<i4>
  check.expect_eq_const %2, dense<5> : tensor<i4>
  func.return
}

Hoạt động kiểm thử trong danh mục Phân phối yêu cầu phải chạy hoạt động đó qua Hoạt động dịch vụ tiện ích interpreter.run_parallel

Đang kiểm tra AllReduceOp (mẫu từ all_reduce.mlir):

// RUN: stablehlo-translate --interpret %s

module @cross_replica {
  func.func public @all_reduce(%operand : tensor<4xi64>) -> tensor<4xi64> {
    %result = "stablehlo.all_reduce"(%operand) ({
      ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
        %0 = stablehlo.add %arg0, %arg1 : tensor<i64>
        stablehlo.return %0 : tensor<i64>
    }) {
      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
      channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
    } : (tensor<4xi64>) -> tensor<4xi64>
    return %result : tensor<4xi64>
  }
  func.func public @main() {
    %inputs0 = stablehlo.constant dense<[1, 2, 3, 4]> : tensor<4xi64>
    %inputs1 = stablehlo.constant dense<[5, 6, 7, 8]> : tensor<4xi64>
    %results:2 = "interpreter.run_parallel"(%inputs0, %inputs1) {
      programs=[[@all_reduce], [@all_reduce]]
    } : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>, tensor<4xi64>)
    check.expect_eq_const %results#0, dense<[6, 8, 10, 12]> : tensor<4xi64>
    check.expect_eq_const %results#1, dense<[6, 8, 10, 12]> : tensor<4xi64>
    func.return
  }
}

Gỡ lỗi ổn địnhHLO

Làm theo các bước tạo bản dựng StableHLO, tệp nhị phân StableHLO cho các công cụ trong stablehlo/tools phải thuộc /build/bin. Những công cụ gỡ lỗi phổ biến như Bạn có thể sử dụng GDB để duyệt qua mã:

gdb --args ./build/bin/stablehlo-translate -allow-unregistered-dialect --interpret ./stablehlo/tests/interpret/<test>.mlir

Phụ lục

Chuyển đổi các hoạt động khác

# batch_norm_grad
hlo-expand --batch_norm_grad_expander <path/to/hlo_module>

# batch_norm_inference
hlo-expand --batch_norm_inference_expander <path/to/hlo_module>

# batch_norm_training
hlo-expand --batch_norm_training_expander <path/to/hlo_module>

# cholesky
hlo-expand --cholesky_expander <path/to/hlo_module>

# constant
# Supported in StableHLO interpreter.

# fft
# TBD

# iota
# Supported in StableHLO interpreter.

# rng
# TBD

# rng_bit_generator
# TBD

# triangular_solve
hlo-expand --triangular_solve_expander <path/to/hlo_module>

Chuyển đổi không thành hoạt động HLO

# broadcast
mlir-hlo-opt -mhlo-legalize-broadcast-to-broadcast-in-dim <path/to/input>

# create_token
mlir-hlo-opt -mhlo-legalize-create-token-to-after-all <path/to/input>

# cross-replica-sum
mlir-hlo-opt -mhlo-legalize-cross-replica-sum-to-all-reduce <path/to/input>

# dot
mlir-hlo-opt -mhlo-legalize-dot-to-dot-general <path/to/input>

# einsum
mlir-hlo-opt -mhlo-legalize-einsum-to-dot-general <path/to/input>

# torch_index_select
mlir-hlo-opt -mhlo-legalize-torch-index-select-to-gather <path/to/input>

# unary_einsum
mlir-hlo-opt --canonicalize -mhlo-legalize-einsum-to-dot-general <path/to/input>

Danh mục hoạt động ổn địnhHLO