เอกสารนี้อธิบายการวิเคราะห์การจัดทำดัชนี HLO ซึ่งช่วยให้คุณคำนวณแผนที่การจัดทำดัชนีสำหรับการดำเนินการของ HLO ในเชิงสัญลักษณ์ได้ แผนที่การจัดทำดัชนีเป็นฟังก์ชันที่แมปดัชนีของ Tensor หนึ่งกับดัชนีของดัชนีอื่น เช่น ดัชนีผลลัพธ์คำสั่ง HLO กับดัชนีอินพุตคำสั่ง HLO หรือในทางกลับกัน
ตัวอย่าง
สำหรับการออกอากาศจาก tensor<20xf32> ถึง tensor<10x20x30xf32>
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
การแมปการจัดทำดัชนีจากเอาต์พุตถึงอินพุตคือ \((i, j, k) \mapsto (j)\) สำหรับ $i \in [0, 10]\(, \)j \in [0, 20]\( and \)k \in [0, 30]$
แรงจูงใจ
XLA GPU ใช้โซลูชันที่ออกแบบมาโดยเฉพาะเพื่อเหตุผลเกี่ยวกับการร่วมมือ การใช้งานตัวถูกดำเนินการ และรูปแบบการเรียงต่อกัน (ดูรายละเอียดเพิ่มเติมด้านล่าง) เป้าหมายของการวิเคราะห์การจัดทำดัชนีคือการจัดหาคอมโพเนนต์ที่นำมาใช้ใหม่ได้สำหรับกรณีการใช้งานดังกล่าว การวิเคราะห์การจัดทำดัชนีสร้างขึ้นจากโครงสร้างพื้นฐาน Affine Map ของ MLIR และเพิ่มอรรถศาสตร์ HLO
การประสาน
การให้เหตุผลเกี่ยวกับการรวมหน่วยความจำกลายเป็นเรื่องที่เป็นไปได้สำหรับกรณีที่ไม่สำคัญเมื่อเราทราบว่ามีการอ่านองค์ประกอบ/ส่วนของอินพุตเพื่อประมวลผลองค์ประกอบของเอาต์พุต
ตัวถูกดำเนินการประโยชน์
การใช้งานตัวถูกดำเนินการใน XLA จะระบุปริมาณอินพุตของคำแนะนำแต่ละรายการ ในกรณีที่ใช้เอาต์พุตเหล่านั้นอย่างเต็มรูปแบบ ในปัจจุบัน การใช้งานยังไม่คำนวณ สำหรับกรณีทั่วไปด้วย การวิเคราะห์การจัดทำดัชนีช่วยให้คำนวณการใช้งานได้อย่างแม่นยำ
การแปลงรหัสวิดีโอโดยใช้ไทล์
ชิ้นส่วน/ชิ้นส่วนคือชุดย่อยของ tensor ทรงสี่เหลี่ยมผืนผ้าที่มีพารามิเตอร์ด้วยออฟเซ็ต ขนาด และความก้าวหน้า การขยายตัวของชิ้นส่วนแผนที่เป็นวิธีการคำนวณพารามิเตอร์ของชิ้นส่วนของผู้ผลิต/ผู้บริโภคของการทดสอบโดยใช้พารามิเตอร์การเรียงตัวของการทดสอบเอง ตอนนี้เรามีไลบรารีสำหรับ Softmax และเครื่องหมายจุด การเผยแพร่ชิ้นส่วนแผนที่สามารถทำได้แบบทั่วไปมากขึ้นและทนทานมากขึ้นหากแสดงผ่านแผนที่การจัดทำดัชนี
ฟังก์ชันและโดเมน
แมปการจัดทำดัชนีคือฟังก์ชัน \(\boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\) ที่แมปดัชนีแบบหลาย \(\boldsymbol{d}\) ของ tensor \(A\) กับองค์ประกอบ/ช่วงของ tensor \(B\)พารามิเตอร์ \(\boldsymbol{s}\) หมายถึงช่วงดัชนีของมิติข้อมูลที่แสดงใน tensor \(B\)แต่ไม่ใช่ tensor \(A\)
ตัวอย่างเช่น หากเราลดลงจาก tensor<2x4x8x16xf32> เป็น tensor<4x8xf32> การแมปการจัดทำดัชนีจากเอาต์พุต 2 มิติเป็นอินพุต 4 มิติจะเป็น\((d_0, d_1) \mapsto (s_0, d_0, d_1, s_1)\)โดยที่ \(d_i\) คือพารามิเตอร์มิติข้อมูลที่สอดคล้องกับดัชนี Tensor เอาต์พุต พารามิเตอร์ \(s_j\)
จะเข้ารหัสหลายค่า กล่าวคือ เพื่อคำนวณ \((d_0, d_1)\) องค์ประกอบของผลลัพธ์ เราจำเป็นต้องมี \((s_0, d_0, d_1, s_1)\) องค์ประกอบของอินพุต โดยที่ \(s_0 \in [0, 2)\) และ
\(s_1 \in [0, 16)\)
การแมปนี้สร้างขึ้นจากแอตทริบิวต์ของคำสั่ง HLO หรืออาจเขียนการแมปของวิธีการที่ไม่ได้รวมเข้าด้วยกันเพื่อรับการจัดทำดัชนีสำหรับฟิวชัน นอกจากนี้ การแมปยังมีโดเมนซึ่งระบุองค์ประกอบของ tensor ที่การแมปที่มีอยู่
\[ \begin{eqnarray} \boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\; &s.t.& \\ \boldsymbol{lb}_d &\leq& \boldsymbol{d} \leq \boldsymbol{ub}_d \\ \boldsymbol{lb}_s &\leq& \boldsymbol{s} \leq \boldsymbol{ub}_s \\ \boldsymbol{lb}_g &\leq& \boldsymbol{g}(\boldsymbol{d}, \boldsymbol{s}) \leq \boldsymbol{ub}_g \end{eqnarray} \]
เนื่องจากเราต้องการลดการคำนวณใหม่ให้เหลือน้อยที่สุด เราจึงต้องมีไลบรารีสำหรับการคำนวณเชิงสัญลักษณ์ XLA อาศัย MLIR อยู่แล้ว เราจึงใช้ mlir::AffineMap แทนการเขียนไลบรารีเลขคณิตสัญลักษณ์
AffineMapทั่วไปจะมีหน้าตาดังนี้
(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)
AffineMap มีพารามิเตอร์ 2 ประเภทด้วยกัน ได้แก่ มิติข้อมูลและสัญลักษณ์ที่เรานำมาใช้ \(\boldsymbol d\) และ \(\boldsymbol s\) ตามลำดับ
AffineMap ไม่มีข้อมูลเมตาเกี่ยวกับช่วงของมิติข้อมูล เราจึงต้องจัดหาข้อมูลนี้ด้วยตนเอง
struct Range {
int64_t lower_bound;
int64_t upper_bound;
};
struct IndexingMap {
mlir::AffineMap affine_map;
std::vector<Range> dimension_ranges;
std::vector<Range> symbol_ranges;
llvm::DenseMap<mlir::AffineExpr, Range> expr_ranges;
};
dim_ranges จะเข้ารหัสข้อจำกัดกล่อง inclusive สำหรับพารามิเตอร์มิติข้อมูล \(\boldsymbol{d}\) ของแมปการจัดทำดัชนี ซึ่งมักจะสอดคล้องกับรูปร่างของ Tensor เอาต์พุตสำหรับการดำเนินการ เช่น สลับตำแหน่ง ลด องค์ประกอบ จุด แต่มีข้อยกเว้นบางอย่าง เช่น HloConcatenateInstruction
symbol_ranges เข้ารหัสค่าที่เป็นไปได้ที่ \(\boldsymbol {s}\) พารามิเตอร์ทำได้
ลองดูตัวอย่างเพื่อทำความเข้าใจความหมายของข้อความข้างต้น
แผนที่การจัดทำดัชนีสำหรับการดำเนินการที่ไม่ได้รวม
Elementwise
สำหรับ Ops แบบองค์ประกอบ แมปการจัดทำดัชนีคือข้อมูลประจำตัว
p0 = f32[10, 20] parameter(0)
p1 = f32[10, 20] parameter(1)
add = f32[10, 20] add(p0, p1)
เอาต์พุตของอินพุตแมป:
- เอาต์พุต -> Input_0: \((d_0, d_1) \mapsto (d_0, d_1)\) สำหรับ $\boldsymbol{d} \in [0,9] \times [0, 19]\(, i.e. \)\boldsymbol{d} \in {\rm Dom}(เอาต์พุต)$
- เอาต์พุต -> Input_1: \((d_0, d_1) \mapsto (d_0, d_1)\) สำหรับ $\boldsymbol{d} \in {\rm Dom} (เอาต์พุต)$
อินพุตที่จะส่งออกแมป
- Input_i -> เอาต์พุต: \((d_0, d_1) \mapsto (d_0, d_1)\) สำหรับ $\boldsymbol{d} \in {\rm Dom}(เอาต์พุต)$
ประกาศ
การออกอากาศหมายความว่ามิติข้อมูลบางส่วนจะถูกนำออกเมื่อเราจับคู่เอาต์พุตกับอินพุต และเมื่อเราจับคู่อินพุตกับเอาต์พุต
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
เอาต์พุตที่จะป้อนการแมป:
- เอาต์พุต -> อินพุต: \((d_0, d_1, d_2) \mapsto (d_1)\) สำหรับ $\boldsymbol{d} \in {\rm Dom}(เอาต์พุต)$
อินพุตไปยังเอาต์พุตแมป
- อินพุต -> เอาต์พุต: \((d_0) \mapsto (s_0, d_1, s_1)\) สำหรับ $\boldsymbol{d} \in {\rm Dom}(เอาต์พุต)\( and \)\boldsymbol{s} \in [0, 9] \times [0, 29]$
โปรดทราบว่าตอนนี้เรามี \(\boldsymbol s\) ทางด้านขวาสำหรับการแมปอินพุตกับเอาต์พุต สัญลักษณ์เหล่านั้นคือสัญลักษณ์ที่แสดงถึงช่วงของค่า ตัวอย่างเช่น ในกรณีนี้ ทุกองค์ประกอบของอินพุตที่มีดัชนี \(d_0\) จะแมปกับส่วนเอาต์พุตขนาด 10x1x30
ค่าคงที่และไอโอตา
เพื่อความสะดวก พารามิเตอร์เหล่านี้ไม่มีพารามิเตอร์อินพุต จึงไม่มีสิ่งใดที่จะต้องคำนวณการจัดทำดัชนี
สลับ
แมปการจัดทำดัชนีสําหรับการสับเปลี่ยนเป็นการเรียงสับเปลี่ยนมิติข้อมูลอินพุต/เอาต์พุต
p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}
เอาต์พุตที่จะป้อนการแมป:
- เอาต์พุต -> อินพุต: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_3, d_1, d_2)\) สำหรับ\(\boldsymbol{d} \in {\rm Dom}(output)\)
อินพุตไปยังเอาต์พุตแมป:
- อินพุต -> เอาต์พุต: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_2, d_3, d_1)\) สำหรับ\(\boldsymbol{d} \in {\rm Dom}(input)\)
รีเวิร์ส
การแมปการจัดทำดัชนีสำหรับแบบย้อนกลับจะเปลี่ยนมิติข้อมูลที่เปลี่ยนกลับเป็น $upper_bound(d_i) - d_i$:
p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}
เอาต์พุตที่จะป้อนการแมป:
- เอาต์พุต -> อินพุต: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(เอาต์พุต)$
อินพุตไปยังเอาต์พุตแมป:
- อินพุต -> เอาต์พุต: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(ข้อมูล)$
(Variadic)Reduce (ลด)
การลดตาม Variadic มีอินพุตหลายรายการและหลาย Init แล้ว แผนที่จากเอาต์พุตสู่อินพุตจะเพิ่มขนาดที่ลดลง ดังนั้นจึงเหมือนจะผกผันกับการออกอากาศ
p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
reduce = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
dimensions={0}, to_apply=min
เอาต์พุตของอินพุตแมป:
- เอาต์พุต -> Input_j: \((d_0) \mapsto (s_0, d_0)\) สำหรับ $\boldsymbol{d} \in {\rm Dom}(เอาต์พุต)\( and \)\boldsymbol{s} \in [0, 9]$
- เอาต์พุต -> init_j: \((d_0) \mapsto ()\) สำหรับ $\boldsymbol{d} \in {\rm Dom}(เอาต์พุต)$
อินพุตที่จะส่งออกแมป
- Input_i -> export_j: \((d_0, d_1) \mapsto (d_1)\) สำหรับ $\boldsymbol{d} \in {\rm Dom}(input)$
- init_i -> export_j: \(() \mapsto (s_0)\) สำหรับ \(\boldsymbol{s} \in [0, 9]\)
สำหรับ \(i, j = 0, \ldots, INPUT\\_COUNT\)
Slice
การจัดทำดัชนีจากเอาต์พุตถึงอินพุตสำหรับสไลซ์ผลลัพธ์ในแผนที่การจัดทำดัชนีแบบก้าว ซึ่งใช้ได้กับทุกองค์ประกอบของเอาต์พุต การแมปจากอินพุตถึงเอาต์พุตจํากัดไว้เฉพาะช่วงก้าวขององค์ประกอบในอินพุต
p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
slice={[5:10:1], [3:20:7], [0:50:2]}
เอาต์พุตที่จะป้อนการแมป:
- เอาต์พุต -> อินพุต: \((d_0, d_1, d_2) \mapsto (d_0 + 5, 7d_1 + 3, 2d_2)\) สำหรับ\(\boldsymbol{d} \in {\rm Dom}(output)\)
อินพุตไปยังเอาต์พุตแมป:
- อินพุต -> เอาต์พุต: \((d_0, d_1, d_2) \mapsto (d_0, d_1 / 7, d_2 / 2)\) สำหรับ \(\boldsymbol{d} \in [5, 9] \times [3, 19] \times [0, 49]\) ก้าวย่าง $[1, 7, 2]$
จะแจ้งภายหลัง: การจัดทำดัชนีอินพุตต่อเอาต์พุต
ปรับรูปร่าง
การเปลี่ยนรูปแบบมีหลากหลายรสชาติ
ยุบรูปร่าง
นี่คือการเปลี่ยนรูปร่าง "เชิงเส้น" จาก N-D เป็น 1D
p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)
เอาต์พุตที่จะป้อนการแมป:
- เอาต์พุต -> อินพุต: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) สำหรับ $\boldsymbol{d} \ใน {\rm Dom}(เอาต์พุต)$
อินพุตไปยังเอาต์พุตแมป:
- อินพุต -> เอาต์พุต: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) สำหรับ $\boldsymbol{d} \in {\rm Dom}(input)$
ขยายรูปร่าง
นี่คือตัวเลือก "รูปร่างยุบ" แบบผกผัน โดยเปลี่ยนรูปร่างอินพุต 1 มิติเป็นเอาต์พุต N-D
p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)
เอาต์พุตที่จะป้อนการแมป:
- เอาต์พุต -> อินพุต: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) สำหรับ $\boldsymbol{d} \in {\rm Dom}(เอาต์พุต)$
อินพุตไปยังเอาต์พุตแมป:
- อินพุต -> เอาต์พุต: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) สำหรับ $\boldsymbol{d} \in {\rm Dom}(input)$
การปรับรูปร่างทั่วไป
ต่อไปนี้เป็นการดำเนินการปรับรูปร่างที่ไม่จะแสดงเป็นรูปร่างขยายหรือยุบเดี่ยวไม่ได้ โฆษณาเหล่านี้สามารถแสดงเป็นองค์ประกอบหนึ่งของรูปร่างขยายหรือยุบ 2 รายการขึ้นไป
ตัวอย่างที่ 1: การแปลงเชิงเส้นเป็นเชิงเส้น
p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)
การปรับรูปร่างนี้สามารถแสดงเป็นองค์ประกอบของรูปร่างยุบของ tensor<4x8xf32> ถึง tensor<32xf32> แล้วขยายรูปร่างเป็น tensor<2x4x4xf32>
เอาต์พุตที่จะป้อนการแมป:
- เอาต์พุต -> อินพุต: $(d_0, d_1, d_2) \mapsto (2d_0 + (4d_1 + d_2) / 8, 4d_1 + d_2) \mod 8)$
เป็นเวลา \(\boldsymbol{d} \in {\rm Dom}(output)\)
อินพุตไปยังเอาต์พุตแมป:
- อินพุต -> เอาต์พุต: $(d_0, d_1) \mapsto ((8d_0 + d_1) / 16, ((8d_0 + d_1) \mod 16) / 4, d_1 \mod 4)$
สำหรับ \(\boldsymbol{d} \in {\rm Dom}(input)\)
ตัวอย่างที่ 2: รูปร่างย่อยที่ขยายและยุบ
p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)
การปรับรูปร่างนี้สามารถแสดงเป็นองค์ประกอบของการปรับรูปร่าง 2 แบบ รายการแรกจะยุบมิติข้อมูลด้านนอก tensor<4x8x12xf32> เป็น tensor<32x12xf32> และรายการที่ 2 ขยายมิติข้อมูลด้านในสุด tensor<32x12xf32> เป็น tensor<32x3x4xf32>
เอาต์พุตที่จะป้อนการแมป:
- เอาต์พุต -> อินพุต: \((d_0, d_1, d_2) \mapsto (d_0 / 8, d_0 \mod 8, 4d_1 + d_2)\) สำหรับ \(\boldsymbol{d} \in {\rm Dom}(output)\)
อินพุตไปยังเอาต์พุตแมป:
- อินพุต -> เอาต์พุต: \((d_0, d_1, d_2) \mapsto (8d_0 + d_1, d_2 / 4, d_2 \mod 4)\) สำหรับ \(\boldsymbol{d} \in {\rm Dom}(input)\)
บิตแคสต์
ออบเจ็กต์บิตแคสต์อาจแสดงเป็นลำดับของการสลับรูปร่างเปลี่ยนรูปร่างได้ ดังนั้น แผนที่การจัดทำดัชนีเป็นเพียงองค์ประกอบของแผนที่การจัดทำดัชนีสำหรับลำดับนี้
เชื่อมต่อ
มีการกำหนดการแมปเอาต์พุตไปยังอินพุตสำหรับ Concat สำหรับอินพุตทั้งหมด แต่จะมีโดเมนที่ไม่ทับซ้อนกัน กล่าวคือ ระบบจะใช้อินพุตเพียงครั้งละ 1 รายการเท่านั้น
p0 = f32[3,50] parameter(0)
p1 = f32[3,30] parameter(1)
concat = f32[3,80] concatenate(f32[3,50] p0, f32[3,30] p1),
dimensions={1}
เอาต์พุตที่จะป้อนการแมป:
- เอาต์พุต -> อินพุต 1:
\((d_0, d_1) \mapsto (d_0, d_1)\) สำหรับ \(\boldsymbol{d} \in [0, 2] \times [0, 49]\)
- เอาต์พุต -> อินพุต 2:
\((d_0, d_1) \mapsto (d_0, d_1 - 50)\) สำหรับ $\boldsymbol{d} \in [0, 2] \times [50, 79]$
อินพุตที่จะส่งออกแมป:
- อินพุต 1 -> เอาต์พุต: \((d_0, d_1) \mapsto (d_0, d_1)\) สำหรับ $\boldsymbol{d} \in {\rm Dom}(input_1)$
- อินพุต 2 -> เอาต์พุต: \((d_0, d_1) \mapsto (d_0, d_1 + 50)\) สำหรับ $\boldsymbol{d} \in {\rm Dom}(input_2)$
จุด (ใช้เอาต์พุตไปยังอินพุต
แผนที่การจัดทำดัชนีสำหรับจุดคล้ายคลึงกับแผนที่การลด
p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
dot = f32[4, 128, 64] dot(p0, p1),
lhs_batch_dims={0}, rhs_batch_dims={0},
lhs_contracting_dims={2}, rhs_contracting_dims={1}
เอาต์พุตไปยังอินพุตจะแมปดังนี้
- เอาต์พุต -> Input_1: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) สำหรับ\(\boldsymbol{d} \in {\rm Dom}(output)\) และ \(\boldsymbol{s} \in [0, 255]\)
- เอาต์พุต -> Input_2: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_2)\) สำหรับ\(\boldsymbol{d} \in {\rm Dom}(output)\) และ \(\boldsymbol{s} \in [0, 255]\)
อินพุตที่จะส่งออกแมป
- อินพุต_1 -> เอาต์พุต: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) สำหรับ\(\boldsymbol{d} \in {\rm Dom}(input_1)\) และ \(\boldsymbol{s} \in [0, 63]\)
- อินพุต_2 -> เอาต์พุต: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_1)\) สำหรับ\(\boldsymbol{d} \in {\rm Dom}(input_2)\) และ \(\boldsymbol{s} \in [0, 127]\)
ลดกรอบเวลา (จะแจ้งภายหลัง)
แผ่น (จะแจ้งภายหลัง)
การจัดทำดัชนีแผนที่สำหรับ Fusion
แผนที่การจัดทำดัชนีสำหรับการดำเนินการของฟิวชันคือองค์ประกอบของแผนที่การจัดทำดัชนีสำหรับการดำเนินการทั้งหมดในคลัสเตอร์ อาจมีการอ่านอินพุตบางรายการหลายครั้งด้วยรูปแบบการเข้าถึงที่แตกต่างกัน
อินพุตเดียว แมปการจัดทำดัชนีหลายรายการ
นี่คือตัวอย่างสำหรับ \(p_0 + p_0^T\)
f {
p0 = f32[1000, 1000] parameter(0)
transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0}
ROOT a0 = f32[1000, 1000] add(p0, transpose_p0)
}
แมปการจัดทำดัชนีเอาต์พุตสู่อินพุตสำหรับ p0 จะเป็น $(d_0, d_1) \mapsto (d_0,
d_1)\( and \)(d_0, d_1) \mapsto (d_1, d_0)$ ซึ่งหมายความว่าเราจะต้องอ่านพารามิเตอร์อินพุต 2 ครั้งเมื่อคำนวณองค์ประกอบเอาต์พุต 1 รายการ
ข้อมูล 1 รายการ แมปการจัดทำดัชนีที่ซ้ำกันออกแล้ว
มีบางกรณีที่แผนที่การจัดทำดัชนีจะเหมือนกัน แม้ว่าแผนที่การจัดทำดัชนีจะไม่ชัดเจนในทันทีก็ตาม
f {
p0 = f32[20, 10, 50] parameter(0)
lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2}
lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1)
lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1}
rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0}
rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1)
rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2}
ROOT add = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}
แมปการจัดทำดัชนีเอาต์พุตไปยังอินพุตสำหรับ p0 ในกรณีนี้มีเพียง $(d_0, d_1, d_2)
\mapsto (d_2, d_0, d_1)$
ซอฟต์แม็กซ์
การแมปการจัดทำดัชนีเอาต์พุตไปยังอินพุตสำหรับ parameter 0 สำหรับ Softmax มีดังนี้
- \((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\)
- \((d_0, d_1, d_2)[s_0] \mapsto (d_0, d_1, s_0)\)
\(\boldsymbol{d} \in {\rm Dom}(output)\) และ \(\boldsymbol{s} \in [0, 124]\) หมายถึงมิติข้อมูลด้านในสุดของอินพุต
ตัวย่อของแผนที่การจัดทำดัชนี
ค่าเริ่มต้นอย่างง่ายสำหรับอัปสตรีม mlir::AffineMap ไม่สามารถตั้งสมมติฐานเกี่ยวกับช่วงของมิติข้อมูล/สัญลักษณ์ ดังนั้นจึงอธิบายนิพจน์ด้วย mod และ div อย่างมีประสิทธิภาพไม่ได้
เราสามารถใช้ประโยชน์จากความรู้เกี่ยวกับขอบเขตล่างและบนของการแสดงออกย่อยในแผนที่ความสัมพันธ์เพื่อลดความซับซ้อนยิ่งขึ้น
เครื่องมือลดความซับซ้อนสามารถเขียนนิพจน์ต่อไปนี้ใหม่
- \((d_0, d_1) \mapsto (d_0 + d1 / 16, d1 \mod 16)\) สำหรับ $\boldsymbol{d} \in [0, 6] \time [0, 14]\( becomes \)(d_0, d_1) \mapsto (d_0, d_1)$
- $(d_0, d_1, d_2) \mapsto ((100d_0 + 10d_1 + d_2) /100, ((100d_0 + 10d_1 + d_2) \mod 100) / 10, d_2 \d_2, (100 วัน_0) \mod 100) / 10, d_2 \d_0, (100d_0)\( for \)\( becomes \)
- $(d_0, d_1, d_2) \mapsto ((16d_0 + 4d_1 + d_2) /8, (16d_0 + 4d_1 + d_2) \mod 8)\( for \)d_i \ใน [0, 9]\( becomes \)(d_0, ม็อด_2 + 4d_1)
- \((d_0, d_1) \mapsto (-(-11d_0 - d_1 + 109) / 11 + 9)\) สำหรับ $\boldsymbol{d} \in [0, 9] \times [0, 10]\( becomes \)(d_0, d_1) \mapsto (d_0)$
การแมปการจัดทำดัชนีที่ง่ายขึ้นช่วยให้เราเข้าใจว่าการเปลี่ยนรูปร่างบางส่วนใน HLO ยกเลิกกันเอง
p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)
หลังจากการเรียบเรียงแผนที่การจัดทำดัชนีและการทำให้แผนที่เรียบง่ายขึ้น เราจะได้รับ
\((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\).