การเปลี่ยนแปลงใน StableHLO

สถานะปัจจุบันของไดนามิกมีคำอธิบายอย่างเป็นทางการมากขึ้นใน Dynamism RFC หน้านี้จะให้ภาพรวมระดับสูงของ RFC และกล่าวถึง API และเครื่องมือที่สําคัญสําหรับการโต้ตอบกับโปรแกรมแบบไดนามิก

ภาพรวมของคำศัพท์และการสนับสนุนสำหรับ Dynamism

ก่อนอื่น เราขอพูดถึงคำศัพท์สั้นๆ ที่จะปรากฏในเอกสารนี้ รวมถึงการแนะนำสั้นๆ เกี่ยวกับการสนับสนุนใน StableHLO

มิติข้อมูลแบบไดนามิก

มิติข้อมูลแบบไดนามิกหมายถึงมิติข้อมูลที่ไม่ทราบขนาดมิติข้อมูล ใน StableHLO เราแสดงมิติข้อมูลแบบไดนามิกโดยใช้ ? นั่นคือ tensor<16x?xf32>

พลวัตที่มีขอบเขต

การเปลี่ยนแปลงแบบจํากัดหมายถึงมิติข้อมูลแบบไดนามิกที่มีค่าที่ทราบขีดจํากัดบน โดยทั่วไปวิธีนี้จะเป็นประโยชน์สำหรับ Padding ของ Tensor ระหว่างการดำเนินการ ใน StableHLO เราจะแสดงไดนามิกแบบจำกัดโดยใช้ #stablehlo.bounds เป็นการเข้ารหัสเทนเซอร์ เช่น เทนเซอร์อันดับ 2 ที่มีมิติข้อมูลแบบไดนามิก 1 รายการซึ่งถูกจำกัดที่ 16 และอีกมิติข้อมูลหนึ่งที่ไม่มีขีดจำกัดจะแสดงเป็น tensor<?x?xf32, #stablehlo.bounds<16, ?>>

StableHLO สามารถแสดงการเปลี่ยนแปลงแบบจำกัดได้ แต่เฟรมเวิร์กรองรับอย่างจำกัด โดยมาจาก TensorFlow และรองรับบางส่วนใน PyTorch/XLA

การเปลี่ยนแปลงแบบไม่จำกัด

ไดนามิกแบบไม่จำกัดตามชื่อที่บ่งบอกหมายถึงมิติข้อมูลแบบไดนามิกที่ไม่มีขอบเขตที่ทราบเกี่ยวกับขนาด ลักษณะการทำงานแบบไดนามิกประเภทนี้พบได้ทั่วไปใน StableHLO ซึ่งรองรับ JAX, PyTorch/XLA และ TF และมักใช้ในการส่งออกโมเดลที่มีขนาดกลุ่มหรือความยาวลำดับแบบไดนามิก

ใน StableHLO เราเพียงนำการเข้ารหัสขอบเขตสำหรับไดนามิกรูปแบบนี้ ซึ่งก็คือ tensor<?x?xf32>

หลายรูปแบบของรูปร่าง

รูปร่างหลายมิติคือคำที่เรารับช่วงมาจาก JAX

ผลที่ตามมาที่สำคัญ 2 ประการในการกำหนดลักษณะของพหุรูปแบบมีดังนี้

1. ความเป็นพลวัตทั้งหมดในโปรแกรมจะสืบย้อนกลับไปยังอาร์กิวเมนต์อินพุต
2. ความเป็นพลวัตทั้งหมดเกี่ยวข้องกับรูปร่างของ Tensor เท่านั้น ซึ่งไม่ได้ขึ้นอยู่กับข้อมูล

เมื่อทราบรูปแบบแบบคงที่ของโปรแกรมแล้ว เราจะนำโปรแกรมแบบไดนามิกมาปรับแต่งให้เป็นโปรแกรมแบบคงที่สำหรับการคอมไพล์ได้อย่างเต็มที่ (ดู"ช่วงการทำงานของคอมไพเลอร์สำหรับการปรับแต่งโปรแกรมแบบไดนามิก")

โดยทั่วไปแล้ว รูปแบบของตัวแปรจะใช้แบบไดนามิกแบบไม่จำกัด หากรูปแบบอาร์กิวเมนต์ที่ทราบจะนำไปสู่โปรแกรมแบบคงที่ทั้งหมด ก็ไม่จำเป็นต้องคาดเดาวิธีจำกัดค่า

การเปลี่ยนแปลงแบบไดนามิกโดยอิงตามข้อมูล

ไดนามิกที่ขึ้นอยู่กับข้อมูลหมายถึงขนาดมิติข้อมูลแบบไดนามิกที่เกี่ยวข้องกับข้อมูลภายในเทนเซอร์ ตัวอย่าง Canonical คือฟังก์ชัน nonzeros ซึ่งแสดงผลดัชนีองค์ประกอบทั้งหมดที่เป็น 0 ในค่า tensor ระบบจะไม่สามารถรู้รูปร่างได้หากไม่มีการประเมินข้อมูล แต่มักจะสามารถรวบรวมได้โดยใช้ไดนามิกแบบตายตัวและใช้หน่วยความจำเพิ่มเติมกับขนาด Tensor เอาต์พุตที่อาจเกิดขึ้น

การดำเนินการแบบไดนามิกจำนวนมากที่ขึ้นอยู่กับข้อมูลสามารถจำลองได้โดยใช้แบบไดนามิกแบบจำกัด ซึ่งจะมีการระบุขีดจำกัดบนของขนาดเทนเซอร์ และโดยทั่วไปฮาร์ดแวร์จะใช้การดำเนินการนี้ผ่านการเติมเทนเซอร์ ปัจจุบันมีการสนับสนุนบางอย่างสำหรับไดนามิกแบบขึ้นอยู่กับข้อมูลใน PyTorch/XLA และ TensorFlow แต่ JAX ยังไม่ติดตามการดำเนินการที่ทําให้เกิดไดนามิกแบบขึ้นอยู่กับข้อมูล

การส่งออกโปรแกรมด้วยมิติข้อมูลแบบไดนามิก

ดูบทแนะนำ StableHLO เพื่อดูข้อมูลเกี่ยวกับวิธีส่งออกโปรแกรมที่มีขนาดกลุ่มหรือความยาวลำดับแบบไดนามิก

• บทแนะนำ JAX > ส่งออกด้วยขนาดกลุ่มแบบไดนามิก
• บทแนะนำ PyTorch/XLA > ส่งออกด้วยขนาดกลุ่มแบบไดนามิก

คอมไพเลอร์ผ่านเพื่อปรับแต่งโปรแกรมแบบไดนามิก

นําไปป์ไลน์การส่งผ่านแบบไดนามิกออก

พาสที่มีประโยชน์ 2-3 รายการสำหรับการปรับแต่งรูปร่างมีดังนี้ ซึ่งรวมอยู่ในไปป์ไลน์พาส createStablehloRemoveDynamismPipeline อย่างสะดวก

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

บัตรแต่ละใบสําหรับการปรับแต่งความมีชีวิตชีวา

พาสที่มักจะมีประโยชน์ในการปรับแต่งรูปร่างมีดังนี้

• stablehlo-refine-arguments เพื่อแทนที่อาร์กิวเมนต์อินพุตด้วยประเภท Tensor ที่เป็นรูปธรรม
• stablehlo-refine-shapes เพื่อเผยแพร่อาร์กิวเมนต์อินพุตใหม่เพื่อจัดระเบียบข้อมูลตลอดทั้งโปรแกรม
• stablehlo-canonicalize-dynamism เพื่อแทนที่ออปชันแบบไดนามิกด้วยออปชันแบบคงที่

ดูข้อมูลและตัวอย่างล่าสุดได้ในเอกสารประกอบที่ลิงก์

ตัวอย่าง: ไดนามิกมีประโยชน์อย่างไรและฉันจะใช้ไดนามิกได้อย่างไร

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

โมเดล add_one แบบคงที่

เราจะใช้โมเดล add_one แบบง่ายต่อไปนี้เพื่อสาธิต

def add_one(x):
  return x + 1

เมื่อติดตามโดยใช้ tensor<4xf32> เราจะได้โปรแกรม StableHLO ต่อไปนี้

// File: add_one.mlir
func.func @add_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<4xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<4xf32>
  return %0 : tensor<4xf32>
}

รูปแบบนี้จะใช้ได้เฉพาะกับอาร์กิวเมนต์อินพุตที่มีรูปร่าง tensor<4xf32> หากเราเปลี่ยนขนาดการทดสอบเป็นกลุ่มหรือความยาวของลําดับ จะต้องติดตามซอร์สโค้ดอีกครั้งและลดระดับเป็น StableHLO อีกครั้ง และเราไม่สามารถรับประกันได้ว่าจะยังเข้าถึงซอร์สโค้ดได้อยู่

โมเดล add_one แบบไดนามิก

ซึ่งเป็นจุดที่มีอิทธิพลต่อรูปร่างของไดนามิกแบบโพลีมอร์ฟิก แต่ JAX และ PyTorch/XLA จะสร้างโมเดล add_one ที่มี IR ที่ถูกต้องแบบไดนามิกแทน ซึ่งจะออกอากาศค่าคงที่เพื่อให้ตรงกับรูปร่างอินพุตแบบไดนามิก ดังนี้

// File: add_one_dynamic.mlir
func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  %cst = stablehlo.constant dense<1.0> : tensor<f32>
  %0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor<?xf32>) -> tensor<i32>
  %1 = stablehlo.reshape %0 : (tensor<i32>) -> tensor<1xi32>
  %2 = stablehlo.dynamic_broadcast_in_dim %cst, %1, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %3 = stablehlo.add %arg0, %2 : tensor<?xf32>
  return %3 : tensor<?xf32>
}

การแสดงโมเดลนี้มีความยืดหยุ่นมากกว่ามาก และช่วยให้ระบุค่าต่างๆ เช่น ขนาดการทดสอบหรือความยาวของลําดับได้ โมเดลนี้สามารถนำไปใช้งานได้ในแพลตฟอร์มที่รองรับรูปร่างแบบไดนามิก (เช่น AI Edge) หรือจะปรับแต่งโดยใช้พาสแบบไดนามิกที่กล่าวถึงในเอกสารประกอบนี้ก็ได้

การปรับแต่งโมเดลแบบไดนามิก

ตัวอย่างเช่น การจัดลําดับบัตรต่อไปนี้สามารถปรับแต่งโปรแกรมนี้ได้อย่างเต็มที่

stablehlo-opt add_one_dynamic.mlir \
  --stablehlo-refine-arguments='types=tensor<16xf32>' \
  --stablehlo-refine-shapes \
  --stablehlo-canonicalize-dynamism

การเปลี่ยนแปลงโปรแกรมมีดังนี้

// After stablehlo-refine-arguments: Inputs updated, shapes not propagated
func.func public @main(%arg0: tensor<16xf32>) -> tensor<?xf32> {
  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {indices_of_shape_operands = dense<1> : tensor<1xi64>} : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
  %3 = stablehlo.dynamic_broadcast_in_dim %cst, %2, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %4 = stablehlo.add %0, %3 : tensor<?xf32>
  return %4 : tensor<?xf32>
}

// After stablehlo-refine-shapes: Shapes propagated, dynamic ops still exist
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// After stablehlo-canonicalize-dynamism: Dynamic ops replaced with static ops
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// (Bonus) Use ` --stablehlo-aggressive-simplification` pass to canonicalize the
// constant broadcast, leaving us with the original static program in this case.
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<16xf32>
  return %0 : tensor<16xf32>
}