สถานะปัจจุบันของความสามารถในการเปลี่ยนแปลงระบุไว้ในRFC ความสามารถในการเปลี่ยนแปลงอย่างเป็นทางการมากขึ้น หน้านี้จะให้ภาพรวมระดับสูงของ RFC และอธิบาย API และเครื่องมือที่สำคัญสำหรับการโต้ตอบกับโปรแกรมแบบไดนามิก
คำศัพท์เกี่ยวกับไดนามิกและภาพรวมการสนับสนุน
ก่อนอื่น เราจะอธิบายคำศัพท์บางคำที่จะปรากฏในเอกสารนี้ รวมถึงการแนะนำสั้นๆ เกี่ยวกับการรองรับคำศัพท์เหล่านั้นใน StableHLO
มิติข้อมูลแบบไดนามิก
มิติข้อมูลแบบไดนามิกหมายถึงมิติข้อมูลใดก็ตามที่ไม่ทราบขนาดของมิติข้อมูล
ใน StableHLO เราจะแสดงมิติข้อมูลแบบไดนามิกโดยใช้ ? นั่นคือ tensor<16x?xf32>
ความสามารถในการเปลี่ยนแปลงที่จำกัด
ความผันแปรที่จำกัดหมายถึงมิติข้อมูลแบบไดนามิกที่มีค่าขอบเขตบนที่ทราบ โดยทั่วไปแล้ว การดำเนินการนี้มีประโยชน์สำหรับการเพิ่มระยะขอบให้กับเทนเซอร์ในระหว่างการดำเนินการ
ใน 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>
การพ้องรูปของรูปร่าง
Shape Polymorphism เป็นคำที่เราได้รับมาจาก JAX
การเปลี่ยนแปลงรูปร่างมีผลกระทบที่สำคัญ 2 ประการ ดังนี้
- ความสามารถในการเปลี่ยนแปลงทั้งหมดในโปรแกรมจะย้อนกลับไปที่อาร์กิวเมนต์อินพุต
- ความสามารถในการเปลี่ยนแปลงทั้งหมดเกี่ยวข้องกับรูปร่างของเทนเซอร์เท่านั้น กล่าวคือ ไม่ขึ้นอยู่กับข้อมูล
เมื่อทราบรูปร่างแบบคงที่ของโปรแกรมแล้ว เราจะใช้กฎ 2 ข้อนี้เพื่อ นำโปรแกรมแบบไดนามิกมาปรับแต่งให้เป็นโปรแกรมแบบคงที่สำหรับการ คอมไพล์ได้อย่างเต็มที่ (ดู"การส่งผ่านคอมไพเลอร์สำหรับการปรับแต่งโปรแกรมแบบไดนามิก")
โดยทั่วไปแล้ว การเปลี่ยนรูปแบบรูปร่างจะใช้ความสามารถในการเปลี่ยนแปลงแบบไม่จำกัด หากทราบรูปร่างอาร์กิวเมนต์ จะทำให้โปรแกรมเป็นแบบคงที่ทั้งหมดได้ จึงไม่จำเป็นต้องคาดเดาว่าจะจำกัดค่าอย่างไร
ความสามารถในการเปลี่ยนแปลงตามข้อมูล
ความไดนามิกที่ขึ้นอยู่กับข้อมูลหมายถึงขนาดมิติข้อมูลแบบไดนามิกที่เกี่ยวข้องกับข้อมูลภายในเทนเซอร์ ตัวอย่าง Canonical คือฟังก์ชัน nonzeros ซึ่ง
ส่งคืนดัชนีขององค์ประกอบทั้งหมดที่เป็น 0 ในค่าเทนเซอร์ ไม่สามารถทราบรูปร่างได้โดยไม่ต้องประเมินข้อมูล แต่โดยทั่วไปจะคอมไพล์ได้โดยใช้
bounded dynamism ซึ่งใช้หน่วยความจำเพิ่มเติมกับขนาดของเทนเซอร์เอาต์พุตที่อาจเกิดขึ้น
การดำเนินการแบบไดนามิกที่ขึ้นอยู่กับข้อมูลหลายอย่างสามารถจำลองได้โดยใช้ไดนามิกแบบมีขอบเขต ซึ่งจะมีการระบุขอบเขตบนของขนาดเทนเซอร์ และโดยทั่วไปฮาร์ดแวร์จะใช้การดำเนินการนี้ผ่านการเพิ่มพื้นที่ว่างในเทนเซอร์ ปัจจุบัน PyTorch/XLA และ TensorFlow รองรับความไดนามิกที่ขึ้นอยู่กับข้อมูลอยู่บ้าง แต่ JAX ยังไม่ติดตามการดำเนินการที่ทำให้เกิดความไดนามิกที่ขึ้นอยู่กับข้อมูล
การส่งออกโปรแกรมที่มีมิติข้อมูลแบบไดนามิก
ดูข้อมูลเกี่ยวกับวิธีส่งออกโปรแกรมที่มีขนาดกลุ่มแบบไดนามิกหรือความยาวลำดับได้ในบทแนะนำ StableHLO
การส่งผ่านคอมไพเลอร์สำหรับการปรับแต่งโปรแกรมแบบไดนามิก
นำไปป์ไลน์การส่งผ่านความไดนามิกออก
มีพาสที่มีประโยชน์ 2-3 รายการสำหรับการปรับแต่งรูปร่าง ซึ่งทั้งหมดรวมอยู่ในไปป์ไลน์พาส createStablehloRemoveDynamismPipeline อย่างสะดวก
void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
TypeRange refinedTypes);
การส่งบัตรแต่ละใบเพื่อปรับแต่งความไดนามิก
โดยพาสที่มักจะมีประโยชน์ในการปรับแต่งรูปร่างมีดังนี้
stablehlo-refine-argumentsเพื่อแทนที่อาร์กิวเมนต์อินพุต ด้วยประเภทเทนเซอร์ที่เฉพาะเจาะจงstablehlo-refine-shapesเพื่อเผยแพร่ข้อมูลอาร์กิวเมนต์อินพุตใหม่ shape ทั่วทั้งโปรแกรมstablehlo-canonicalize-dynamismเพื่อแทนที่การดำเนินการแบบไดนามิก ด้วยการดำเนินการแบบคงที่stablehlo-check-shape-assertionsเพื่อตรวจสอบและ นำการเรียกที่กำหนดเองของการยืนยันรูปร่างออก
ดูข้อมูลและตัวอย่างล่าสุดได้ในเอกสารประกอบที่ลิงก์
ตัวอย่าง: ความสามารถในการเปลี่ยนแปลงมีประโยชน์อย่างไร และฉันจะใช้ได้อย่างไร
ความสามารถในการเปลี่ยนแปลงมีประโยชน์มากมาย ในที่นี้เราจะมุ่งเน้นที่กรณีการใช้งานทั่วไปสำหรับ การเปลี่ยนรูปทรง - การสร้างการแสดงโมเดลที่ส่งออกที่ยืดหยุ่น ซึ่งโดยทั่วไปใช้เพื่อแสดงขนาดกลุ่มแบบไดนามิกหรือความยาวลำดับ
โมเดล 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>
}