ล่าม StableHLO

เป้าหมายหลักของล่าม StableHLO คือการให้ข้อมูลอ้างอิง การนำความหมายของแนวคิด StableHLO มาปรับใช้ เป้าหมายรองคือให้การปรับใช้อย่างใกล้ชิด ข้อมูลจำเพาะ ซึ่งให้ความอ่านง่ายมากกว่าประสิทธิภาพ เพื่อให้ความชัดเจนยิ่งขึ้น ต่อความหมายของการดำเนินการที่เกี่ยวข้องที่สุดอย่าง Convolution Gather/Scatter และ DotGeneral

ในขณะนี้ OpenXLA รองรับการตีความ 91 จาก 96 คะแนน การดำเนินการ HLO ที่เสถียร การดำเนินการ 3 รายการที่เหลือ (FftOp, RngOp, RngBitGeneratorOp) มี อรรถศาสตร์ของตนถูกบันทึกไว้ใน spec.md และมี ได้ดำเนินการตรวจสอบเบื้องต้นเกี่ยวกับวิธีการดำเนินการต่อแล้ว (ดู status.md เพื่อดูรายการการดำเนินการทั้งหมดและสถานะล่าสุด) ขั้นสุดท้าย การเพิ่มประสิทธิภาพจะได้รับการนำไปใช้กับชุมชนที่จำเป็น

ขอบเขต

เราได้จัดหมวดหมู่เครื่องมือ StableHLO เป็น 11 หมวดหมู่ ซึ่งประกอบด้วยการดำเนินการ 118 รายการใน ทั้งหมด (ดูภาคผนวก) การใช้งานข้อมูลอ้างอิง สตรีมงานจะจัดระเบียบงานด้านการใช้งานล่าม สำหรับการดำเนินการ StableHLO 100% ตามที่ระบุไว้ในข้อกำหนด StableHLO เรา วางแผนที่จะทำงานทั้งหมดหรือเกือบทั้งหมดให้เสร็จในสตรีมงานนี้ใน StableHLO v1.0 จากการดำเนินการ 96 รายการที่มีข้อมูลจำเพาะในปัจจุบัน เราตีความการดำเนินการ 91 ครั้งผ่าน OpenXLA (ดูกรณีพิเศษสำหรับอีก 5 กรณีที่เหลือ)

ข้อมูลจำเพาะ

ข้อกำหนดหลักสำหรับล่ามคือการโต้ตอบแบบ 1:1 กับ ข้อกำหนด ข้อกำหนดนี้จะช่วยให้กำหนดมาตรฐานล่ามให้กับการดำเนินการที่คล้ายกันซึ่ง ซึ่งจะนำไปสู่การใช้งานล่ามคุณภาพสูงแบบแยกส่วน

คดีพิเศษ

เบ็ดเตล็ด

หมวดหมู่นี้มีปฏิบัติการที่ย่อยได้ซึ่งมีอนาคตที่ไม่ชัดเจนในขณะนี้ มี เป็นการดำเนินการ 3 รายการในหมวดหมู่นี้ซึ่งล่ามไม่รองรับ ช่วงเวลา:

• FftOp
• RngOp
• RngBitGeneratorOp

FftOp จัดอยู่ในหมวดหมู่เบ็ดเตล็ด แต่แตกต่างจากการดำเนินการอื่นๆ ในหมวดหมู่นี้ กระบวนการนี้ไม่มี Extender Pass และการสนับสนุนนี้ใน StableHLO คือ อยู่ระหว่างดำเนินการ

สามารถแยก RngOp และ RngBitGeneratorOp เป็น MHLO Ops ได้ แต่ decomposition ทำให้เกิด XlaRngGetAndUpdateStateOp ซึ่งเป็น MHLO ที่เฉพาะเจาะจง การดำเนินการ การสนับสนุนการตีความทั้งสองสิ่งนี้เป็นการดำเนินการทางกฎหมาย

เครื่องมือในการแปลงการดำเนินการที่เหลือในหมวดหมู่นี้เป็น Ops ของ StableHLO อินเตอร์พรีเตอร์รองรับอยู่ใน hlo_expand_main.cc

ไม่ได้อยู่ใน HLO

นอกเหนือจากการดำเนินการที่กำหนดแล้ว หมวดหมู่นี้ยังประกอบด้วยการดำเนินการที่ไม่ได้ระบุ 8 รายการ (โปรดดู StableHLO Ops category) ซึ่งมีแผนจะ ย้ายออกจาก 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 สำหรับความหมายความเห็นและตัวอย่าง)

การเขียนโปรแกรมทดสอบ

เราใช้เครื่องมือไฟของ 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

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

หมวดหมู่การดำเนินการของ StableHLO