مترجم StableHLO

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

در حال حاضر، OpenXLA از تفسیر 91 از 96 عملیات مشخص شده StableHLO پشتیبانی می کند. 3 عملیات باقیمانده ( FftOp ، RngOp ، RngBitGeneratorOp ) دارای معناشناسی خود در spec.md مستند شده اند، و تحقیقات اولیه در مورد چگونگی حرکت به جلو را تکمیل کرده اند (برای لیست کامل عملیات ها و آخرین وضعیت آن به status.md مراجعه کنید). این پیشرفت های نهایی بر اساس نیاز جامعه اجرا خواهد شد.

دامنه

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

مشخصات

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

موارد خاص

متفرقه

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

  • FftOp
  • RngOp
  • RngBitGeneratorOp

FftOp به عنوان متفرقه طبقه بندی می شود، اما بر خلاف سایر عملیات های این دسته، این عملیات دارای پاس توسعه دهنده نیست و پشتیبانی از آن در StableHLO یک WIP است.

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

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

در HLO نیست

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

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

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

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

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

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

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

بازل:

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

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

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

گویش مترجم

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

گویش چک

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

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

ما از ابزار روشن 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
}

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

تست 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 Ops 

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

دسته های عملیات StableHLO

دسته بندی ها یادگاری مجموع
119
کنترل جریان after_all، case، if، optimization_barrier، while 5
حرکت داده ها پخش_در_کم، الحاق، برش_دینامیک، تکه_به روز رسانی_دینامیک، جمع آوری، پد، تغییر شکل، معکوس، پراکنده، برش، مرتب سازی، جابجایی 12
توزیع all_gather, all_reduce, all_to_all, collective_permute, infeed, outfeed, partition_id, recv, small_scatter, replica_id, send 11
پویایی dynamic_broadcast_in_dim، dynamic_conv، dynamic_gather، dynamic_iota، dynamic_pad، dynamic_reshape، get_dimension_size، real_dynamic_slice، set_dimension_size 9
از نظر عنصری abs، افزودن، و، atan2، bitcast_convert، cbrt، سقف، گیره، مقایسه، مختلط، تبدیل، کسینوس، شمارش_صفرهای پیشرو، تقسیم، نمایی، نمایی_منهای_یک، طبقه، تصویر، محدود است، ورود به سیستم، log_plus_one، لجستیک، نقشه، حداکثر، ضرب، نفی، نه، یا، popcnt، توان، واقعی، کاهش_دقت، باقیمانده، دور_نزدیکترین_افز، گرد_نزدیکترین_زوج، rsqrt، انتخاب، shift_left، shift_right_hismetic، shift_right_logical، علامت، سینوس، sqrt، تفریق، tan، tanh، xor 48
توسعه پذیری custom_call، get_tuple_element، tuple 3
متفرقه batch_norm_grad, batch_norm_inference, batch_norm_training, cholesky, ثابت, fft, iota, rng, rng_bit_generator, triangular_solve 10
مدولار بودن تماس، تابع، ماژول، بازگشت 4
در HLO نیست پخش، ایجاد_توکن، جمع کپی متقابل، نقطه، einsum، torch_index_select، unary_einsum 8
کوانتیزاسیون یکنواخت_دکوانتیز کردن، یکنواخت_کوانتیزه کردن 2
کاهش پیچیدگی، نقطه_عمومی، کاهش، پنجره_کاهش، انتخاب_و_پراکندگی 5