เขียนโค้ดคูณเมทริกซ์ ARM ด้วย NEON Assembly บน M2
เจาะลึกการเขียน ARM Assembly และ NEON instructions สำหรับคูณเมทริกซ์บนชิป Apple M2 ดันประสิทธิภาพแตะหลัก Gflops
ภาพประกอบจากคลังภาพสต็อก ไม่ใช่ภาพจากเหตุการณ์จริง
- นักพัฒนาทดลองเขียน ARM assembly เพื่อคูณเมทริกซ์บน MacBook Air ชิป M2
- ทดสอบวิธี Naive ได้ 1.98 Gflops ก่อนขยับมาใช้ SIMD ดันขึ้นเป็น 7.86 Gflops
- การใช้เทคนิค Tiling และขยายขนาด Tile ช่วยเพิ่มประสิทธิภาพคำนวณอย่างก้าวกระโดด
- เตรียมต่อยอดสู่การคำนวณ FP16 และคำสั่ง SME หรือ SVE2 ในบทความถัดไป
การเขียนเคอร์เนลสำหรับโครงข่ายประสาทเทียมมักเป็นเรื่องที่นักพัฒนากลับมาหยิบจับอยู่เสมอ แม้ปัจจุบันเครื่องมือบนฝั่ง GPU อย่าง CUDA จะมีความก้าวหน้าไปมาก แต่อุปกรณ์สถาปัตยกรรม ARM ในระบบคลาวด์และคอมพิวเตอร์ทั่วไปก็มีจำนวนเพิ่มขึ้นอย่างต่อเนื่อง ทำให้เกิดคำถามว่าเราจะสามารถเขียนโค้ดคูณเมทริกซ์โดยใช้ภาษาแอสเซมบลีของ ARM ให้ทำงานด้วยความเร็วที่น่าพอใจได้อย่างไร โดยเริ่มต้นทดลองเขียนด้วย NEON assembly บนเครื่อง MacBook Air ที่ใช้โปรเซสเซอร์ M2 ซึ่งมีขีดจำกัดทางทฤษฎีของการคูณเมทริกซ์แบบ f32 แบบเธรดเดี่ยวอยู่ที่ 112 GFlops
การทดสอบประสิทธิภาพครั้งนี้อ้างอิงจากการกวาดขนาดเมทริกซ์ (shape sweep) ตั้งแต่ 64x64 จนถึง 1024x1024 โดยรันจำนวน 20 ครั้งต่อหนึ่งสเต็ป แล้วนำค่าเฉลี่ยของแต่ละเคอร์เนลมาคำนวณผลลัพธ์ เริ่มต้นจากการเขียนฟังก์ชันแบบ Naive ด้วยแอสเซมบลีพื้นฐาน ซึ่งผลลัพธ์ที่ได้คือ 1.98 Gflops หรือคิดเป็นราว 1.8% ของขีดจำกัดทางทฤษฎีเท่านั้น สะท้อนให้เห็นว่าคอมไพเลอร์ภาษา C ยังคงมีความฉลาดในการจัดการโค้ดได้ดีกว่าในขั้นต้น
ภาพประกอบจากคลังภาพสต็อก ไม่ใช่ภาพจากเหตุการณ์จริง
เมื่อลองปรับมาใช้งานคำสั่ง SIMD โดยการโหลดค่า 4 ค่าจากเมทริกซ์ B เข้าสู่รีจิสเตอร์ 128-bit เดียว และกระจายค่า (broadcast) จากเมทริกซ์ A หนึ่งค่าลงในรีจิสเตอร์ 128-bit เพื่อใช้คำสั่ง fmla ในการคำนวณ 4 ค่าพร้อมกัน ส่งผลให้ประสิทธิภาพกระโดดขึ้นมาอยู่ที่ 7.86 Gflops หรือดีกว่าเดิมประมาณ 4 เท่าตัว แม้จะมีข้อจำกัดว่าขนาดของ N จะต้องหารด้วย 4 ลงตัวก็ตาม
เพื่อเพิ่มความเข้มข้นในการคำนวณต่อการโหลดข้อมูลแต่ละครั้ง จึงได้นำเทคนิคการทำ Tiling มาใช้ ซึ่งเป็นเทคนิคคลาสสิกที่นิยมใช้ในเคอร์เนลของ GPU โดยใช้รีจิสเตอร์ 128-bit จำนวน 32 ตัวสำหรับคำนวณเลขทศนิยม ทำให้สามารถสร้าง Output Tile ขนาด 8x8 โดยใช้ 16 รีจิสเตอร์สำหรับผลลัพธ์ และอีก 16 รีจิสเตอร์สำหรับอินพุต ส่งผลให้มีการโหลดข้อมูล 16 ครั้งและการดำเนินการ fmla ถึง 64 ครั้งต่อการวนซ้ำในมิติ K แต่ละรอบ นอกจากนี้ยังมีการแทรกคำสั่งโหลดข้อมูลบางส่วนของ B สลับกับการคำนวณ fmla เพื่อทำหน้าที่คล้ายระบบ Double buffering แบบง่าย ซึ่งช่วยเพิ่มประสิทธิภาพขึ้นมาได้อีกราว 5 Gflops
การเพิ่มประสิทธิภาพการคูณเมทริกซ์ (Matrix Multiplication) บนซีพียูระดับล่างอย่าง ARM ด้วย NEON ถือเป็นโจทย์คลาสสิกที่ช่วยให้เข้าใจการทำงานของฮาร์ดแวร์ในระดับฮาร์ดแวร์เวกเตอร์ (Vector Registers) และการบริหารจัดการแคชหน่วยความจำ การใช้เทคนิค Tiling ช่วยลดคอขวดของการดึงข้อมูลจากแรมหลักมาไว้ในรีจิสเตอร์ ซึ่งเป็นหัวใจสำคัญในการรีดประสิทธิภาพสูงสุดของชิปตระกูล Apple Silicon ก่อนจะขยับไปสู่ชุดคำสั่งขั้นสูงอย่าง SVE2 หรือ SME ในอนาคต
ท้ายที่สุด การขยับไปสู่ความแม่นยำแบบ FP16 ด้วย Tile ขนาด 8x16 ช่วยให้สามารถเพิ่มจำนวนค่าที่คำนวณเป็นสองเท่าโดยใช้จำนวนรีจิสเตอร์เท่าเดิม ปิดท้ายการทดลองที่สนุกสนานนี้ด้วยการเตรียมลุยต่อกับการคูณเมทริกซ์เลขจำนวนเต็ม (Integer matmul) และการทดลองใช้ SME หรือ SVE2 ในโอกาสต่อไป
ที่มา: Dev.to
พบข้อมูลผิดพลาดในบทความนี้? แจ้งปัญหาบทความนี้
ความคิดเห็น
แสดงความคิดเห็น