API ส่วนย่อยของ WMMA

หากต้องการเขียนโปรแกรมให้แกน Tensor โดยตรง ให้ใช้ WMMA ซึ่งเป็น API สำหรับการคูณและสะสมผลของเมทริกซ์ระดับวาร์ปในเนมสเปซ nvcuda::wmma 🧩

WMMA ทำงานในระดับวาร์ป โดยเธรดทั้ง 32 เธรดในวาร์ปจะทำงานร่วมกันบนไทล์เดียว คุณจึงคิดเป็นไทล์ ไม่ใช่เธรดเดี่ยว

แฟรกเมนต์คือชนิดข้อมูลของ WMMA ที่เก็บส่วนหนึ่งของไทล์เมทริกซ์ซึ่งเป็นของวาร์ปหนึ่งวาร์ป แต่ละเธรดจะถือครององค์ประกอบบางส่วนอย่างเงียบ ๆ

คุณประกาศแฟรกเมนต์ที่กำกับเป็น matrix_a, matrix_b หรือตัวสะสม ป้ายกำกับจะบอก WMMA ว่าไทล์นั้นถูกป้อนเข้าสู่การคูณและสะสมผลอย่างไร

แฟรกเมนต์ใช้ขนาดไทล์ที่กำหนดไว้ เช่น 16 คูณ 16 คูณ 16 คุณเลือกรูปร่างที่รองรับได้ เพราะฮาร์ดแวร์รับเฉพาะชุดค่าผสมเหล่านั้น

ขั้นแรกให้เรียก load_matrix_sync เพื่อดึงไทล์จากหน่วยความจำเข้าสู่แฟรกเมนต์ การโหลดนี้จะแบ่งข้อมูลไปทั่วทั้งวาร์ปให้คุณ

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

การเรียกหลักคือ mma_sync ซึ่งทำ D = A คูณ B บวก C บนแฟรกเมนต์ที่โหลดไว้ โดยใช้แกน Tensorโดยตรง

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

ส่วนต่อท้าย _sync เป็นเครื่องเตือนใจว่า การเรียก WMMA ทุกครั้งเป็นแบบทำงานพร้อมกันทั้งวาร์ปทั้ง 32 เลนต้องมาถึงพร้อมกัน มิฉะนั้นพฤติกรรมจะไม่เป็นที่กำหนด

การโหลดและจัดเก็บต้องใช้มิติหลัก ซึ่งคือระยะห่างระหว่างแถว รวมถึงเลย์เอาต์แบบเรียงตามแถวหรือคอลัมน์ เพื่ออ่านหน่วยความจำได้ถูกต้อง

คุณได้ทำความเข้าใจกระบวนการของ WMMA แล้ว ได้แก่ ประกาศแฟรกเมนต์ โหลดไทล์ ล้างตัวสะสม เรียก mma_sync แล้วจัดเก็บ ทุกขั้นตอนทำงานพร้อมกันทั้งวาร์ป 🙌