مترجم StableHLO

هدف اصلی مفسر StableHLO ارائه یک پیاده‌سازی مرجع برای معانی opset StableHLO مطابق با مشخصات آن است. هدف ثانویه این است که پیاده‌سازی به دقت از مشخصات پیروی کند، خوانایی را بر عملکرد ترجیح دهد، تا وضوح بیشتری را برای معانی حتی پیچیده‌ترین عملیات مانند Convolution ، Gather / Scatter و DotGeneral فراهم کند.

در حال حاضر، OpenXLA از تفسیر ۹۱ مورد از ۹۶ عملیات StableHLO مشخص‌شده پشتیبانی می‌کند. ۲ عملیات باقی‌مانده ( RngOp و RngBitGeneratorOp ) دارای مستندات معنایی در spec.md هستند و تحقیقات اولیه در مورد نحوه‌ی پیشرفت آنها تکمیل شده است (برای مشاهده‌ی لیست کامل عملیات‌ها و آخرین وضعیت آنها به status.md مراجعه کنید). این پیشرفت‌های نهایی بر اساس نیاز جامعه پیاده‌سازی خواهند شد.

دامنه

ما مجموعه عملیات StableHLO را به 11 دسته شامل 118 عملیات در مجموع طبقه‌بندی کردیم ( به پیوست مراجعه کنید). جریان کاری پیاده‌سازی مرجع ، کار پیاده‌سازی یک مفسر برای 100٪ عملیات StableHLO را طبق مشخصات StableHLO تعریف شده، سازماندهی می‌کند. ما در حال برنامه‌ریزی برای تکمیل تمام یا تقریباً تمام کارها در این جریان کاری در StableHLO نسخه 1.0 هستیم. از 96 عملیاتی که در حال حاضر دارای مشخصات هستند، می‌توانیم 91 عملیات را از طریق OpenXLA تفسیر کنیم (برای 5 مورد باقی‌مانده به موارد خاص مراجعه کنید).

مشخصات

الزام اصلی برای مفسر، تطابق یک به یک با مشخصات است. این مشخصات امکان استانداردسازی مفسر را در میان عملیات مشابه فراهم می‌کند که منجر به پیاده‌سازی ماژولار و با کیفیت بالای مفسر می‌شود.

موارد خاص

متفرقه

این دسته دارای عملیات‌های تجزیه‌پذیر است که آینده‌ی آنها در حال حاضر مشخص نیست. سه عملیات مشخص‌شده در این دسته وجود دارد که مفسر در حال حاضر از آنها پشتیبانی نمی‌کند:

• RngOp
• RngBitGeneratorOp

RngOp و RngBitGeneratorOp می‌توانند به عملیات MHLO تجزیه شوند، اما این تجزیه یک XlaRngGetAndUpdateStateOp را معرفی می‌کند که یک عملیات خاص MHLO است. تفسیر پشتیبانی از این دو عملیات، یک WIP است.

ابزاری که عملیات‌های باقی‌مانده در این دسته را به عملیات‌های StableHLO که مفسر از آنها پشتیبانی می‌کند تبدیل می‌کند، در hlo_expand_main.cc قرار دارد.

در HLO نیست

جدا از عملیات‌های مشخص‌شده، این دسته شامل ۸ عملیات نامشخص است (به دسته‌های عملیات StableHLO مراجعه کنید) که قرار است از StableHLO خارج شوند. اکثر این عملیات‌ها دارای مجوزهای موجود در mhlo هستند تا بتوان آنها را به عملیات‌های معادل StableHLO تبدیل کرد.

ابزاری که عملیات‌های باقی‌مانده در این دسته را به عملیات‌های معادل StableHLO که مفسر از آنها پشتیبانی می‌کند تبدیل می‌کند، در mlir-hlo-opt.cc قرار دارد.

کوانتیزاسیون

پشتیبانی مفسر برای عملیات stablehlo.constant با نوع کوانتیزه پشتیبانی نمی‌شود و از طریق #1691 پیگیری می‌شود.

دستورالعمل استفاده

ساخت مفسر مرجع

این مفسر را می‌توان از طریق Bazel یا CMake (ترجیحاً) ساخت و آزمایش کرد. برای دستورالعمل‌های کامل، به README.md مراجعه کنید.

بازل:

bazel build //...

سی میک:

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

برای اجرای مفسر، ما یک ابزار ترجمه برای تفسیر برنامه‌های StableHLO نوشته شده در MLIR داریم.

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

گویش مترجم

گویش Interpreter شامل عملیات کاربردی مختلفی مربوط به مفسر است. به طور خاص، عملیات interpreter.run_parallel (برای معانی عملیات و مثال کاربرد به InterpreterOps.td مراجعه کنید) امکان تفسیر عملیات توزیع را فراهم می‌کند و قرار است بر اساس نیازهای جامعه، ابزارهای کاربردی بیشتری اضافه شوند.

گویش چک

گویش Check برای مقایسه مقادیر زمان اجرای مفسر با مقادیر مورد انتظار استفاده می‌شود. خروجی‌های برنامه StableHLO را می‌توان از طریق عملیات‌های بررسی مختلف آزمایش کرد (برای معانی عملیات و مثال‌های کاربرد به CheckOps.td مراجعه کنید).

نوشتن برنامه‌های تست

ما از ابزار lit در LLVM برای اجرا و مقایسه با فایل تولید شده برای مقایسه با خروجی مفسر استفاده می‌کنیم (برای مثال، تست‌ها را در stablehlo/tests/interpret ببینید).

تست AddOp (نمونه از 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
}

تست عملیات در دسته توزیع (Distribution) نیازمند اجرای آن از طریق ابزار interpreter.run_parallel op است.

تست AllReduceOp (نمونه از 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
  }
}

اشکال‌زدایی StableHLO

پس از مراحل ساخت StableHLO، فایل‌های باینری StableHLO برای ابزارهای موجود در stablehlo/tools باید در /build/bin قرار گیرند. ابزارهای اشکال‌زدایی رایج مانند GDB می‌توانند برای بررسی گام به گام کد استفاده شوند:

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

پیوست

تبدیل عملیات متفرقه

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

تبدیل در عملیات HLO نیست

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

دسته بندی های عملیات پایدار HLO