Kernel คืออะไรกันแน่ ?

ตัวเลขสองตัวบน GPU

คุณรู้วิธีเขียน c = a + b ใน Python อยู่แล้ว คุณทำมันเป็นพันๆ ครั้งแล้ว ที่นี่เราจะพูดถึงมันในเทนเซอร์ของ PyTorch เทนเซอร์ก็คืออาร์เรย์ของตัวเลข การวางมันบน GPU หมายความว่าอาร์เรย์นั้นอยู่บนหน่วยความจำของ GPU แทนที่จะเป็นหน่วยความจำ CPU ทั่วไป เมื่อ a และ b เป็นเทนเซอร์สองตัวบน GPU บรรทัดเดียวนี้ก็ทำงานเสร็จเร็วพอจนคุณไม่ต้องคิดถึงมัน

ตอนนี้ย่อมันลง สมมติว่า a และ b เป็น float เดี่ยวๆ สองตัว ทั้งคู่อยู่บน GPU บรรทัดเดียวกัน c = a + b จริงๆ แล้วอะไรที่ทำงานบนชิป?

คำตอบคือ เคอร์เนล ในโลกนี้ เคอร์เนลคือโปรแกรมเล็กๆ ที่ GPU ใช้รันกับข้อมูลบางส่วน ไม่ใช่เคอร์เนล OS ที่คอมพิวเตอร์ของคุณบูต ไม่ใช่เคอร์เนลทางคณิตศาสตร์จากหนังสือพีชคณิตเชิงเส้น คำนี้ถูกใช้ซ้ำหลายครั้ง และนั่นไม่ใช่ความผิดของคุณ ในโลกของ GPU เคอร์เนลหมายถึง: ฟังก์ชันเล็กๆ ที่ GPU ถูกสั่งให้ไปรันทันที แบบขนาน บนข้อมูลที่คุณส่งให้

เมื่ออ่านจบ คุณจะสามารถดูโค้ด PyTorch สักส่วนแล้วนับได้ว่า GPU จะรันเคอร์เนลกี่ตัว มันฟังดูเหมือนเทคนิคเล็กๆ และมันก็ใช่ แต่ก็เป็นก้าวแรกที่จะขจัดความรู้สึกว่า "GPU เป็นกล่องดำ" ซึ่งเป็นความรู้สึกตอนที่โมเดลของคุณช้าและคุณไม่รู้ว่าทำไม ทุกคำถามที่คุณถามเกี่ยวกับประสิทธิภาพของ GPU จะกลับมาที่ "เคอร์เนลไหนทำงาน และพวกมันทำอะไรอยู่" ดังนั้นนี่คือจุดเริ่มต้น

เคอร์เนลแรกของคุณ

มาทำให้ a และ b ใหญ่ขึ้นอีกนิด คราวนี้เป็นเทนเซอร์ขนาด 8 ยังคงเป็น PyTorch หนึ่งบรรทัด: c = a + b

เมื่อคุณรันสิ่งนี้ CPU ของคุณ (เครื่องที่รัน Python จริงๆ) จะบอก GPU ว่า: เฮ้ ไปรันโปรแกรมนี้กับข้อมูลนี้ คำสั่งนั้นเรียกว่า การเปิดตัว สิ่งที่ถูกเปิดตัวคือเคอร์เนล: หนึ่งโปรแกรม พร้อมที่จะรัน การเปิดตัวตัวเองนั้นมีต้นทุนต่ำ ไม่กี่ไมโครวินาทีต่อครั้ง สิ่งที่ รอบๆ การเปิดตัว (ข้อมูลที่ไปยัง GPU ผลลัพธ์ที่กลับมา) คือต้นทุนที่แท้จริง และนั่นคือสิ่งที่เราจะนับไปเรื่อยๆ

ภายในเคอร์เนล งานจริงทำโดยผู้ทำงานจิ๋วที่เรียกว่า เธรด GPU มีเธรดนับพันให้ใช้ สำหรับการบวกขนาด 8 ของเรา 8 เธรดจะมารับงาน: เธรด 0 จัดการองค์ประกอบ 0, เธรด 1 จัดการองค์ประกอบ 1, ไปเรื่อยๆ จนถึงเธรด 7 แต่ละเธรดรันโปรแกรมเล็กๆ เดียวกัน: อ่านหนึ่งองค์ประกอบของ a อ่านองค์ประกอบที่ตรงกันของ b บวก它们 แล้วเขียนผลลัพธ์ไปยัง c

(ในทางปฏิบัติ GPU เปิดตัวเธรดเป็นกลุ่มขนาดคงที่ที่เรียกว่า warp ซึ่งมี 32 เธรดเสมอในการ์ด NVIDIA และปิดบังส่วนที่เกินเมื่ออาร์เรย์ของคุณหารไม่ลงตัว ปลอดภัยที่จะไม่สนใจตอนนี้)

ดังนั้นเรามี PyTorch หนึ่งบรรทัด หนึ่งการเปิดตัว หนึ่งเคอร์เนล 8 เธรดที่ทำงาน 8 การบวก ทีนี้มานับว่าอะไรข้ามชิปไปบ้าง ในการบวก แต่ละเธรดต้องการองค์ประกอบของ a และองค์ประกอบของ b นั่นคือการอ่าน a 8 ครั้งและการอ่าน b 8 ครั้ง จากนั้นแต่ละเธรดเขียนผลลัพธ์ไปยัง c นั่นคือการเขียน 8 ครั้ง

การอ่านและเขียนเหล่านั้นไปยังหน่วยความจำขนาดใหญ่ที่อยู่ติดกับชิป GPU ในการ์ดดาต้าเซ็นเตอร์ (A100, H100) หน่วยความจำนั้นเรียกว่า HBM (high-bandwidth memory) ในการ์ดผู้บริโภค (RTX 4090) และ T4 แบบ Colab ที่คนมักจะลองใช้จริง เรียกว่า GDDR ไม่ว่าจะแบบไหน มันคือหน่วยความจำที่เร็วและอยู่ติดกับชิป และเราจะเรียกมันว่าหน่วยความจำ GPU มันเร็ว แต่ไม่ฟรี และทุกการเดินทางไปยังมันมีต้นทุน

หนึ่งเคอร์เนล = หนึ่งการเปิดตัว = หนึ่งรอบผ่านข้อมูล ไม่ว่าเคอร์เนลจะทำอะไรภายในตัวมัน การอ่านและเขียนที่ขอบของมัน (การเดินทางออกไปยังหน่วยความจำ GPU เพื่อดึงข้อมูลเข้า การเดินทางกลับเพื่อเขียนผลลัพธ์) คือส่วนที่มีต้นทุน นั่นคือรูปแบบทั้งหมด

ไม่มีอะไรเปลี่ยนแปลงเมื่อเทนเซอร์ใหญ่ขึ้น PyTorch บรรทัดเดียวกัน หนึ่งเคอร์เนลเดียวกัน แค่มีเธรดมากขึ้น ถ้า a และ b มีองค์ประกอบละล้านตัว GPU จะเปิดตัวเคอร์เนลเดียวกันด้วยฝูงเธรดที่ใหญ่ขึ้น คณิตศาสตร์ปรับขนาด ไบต์ที่เคลื่อนย้ายปรับขนาด แต่แบบจำลองทางความคิดไม่เปลี่ยน หนึ่งบรรทัด หนึ่งเคอร์เนล

สิ่งที่เกิดขึ้นระหว่างสองโอเปอเรชัน

1c = (a + b).relu()

คุณรู้จัก Python ดีพอที่จะรู้ว่านี่คือสองโอเปอเรชัน: การบวกแล้วตามด้วย relu ต่อกัน ในอินเทอร์พรีเตอร์ นั่นคือการเรียกฟังก์ชันสองครั้ง บน GPU ในโหมด eager ของ PyTorch นั่นคือการเปิดตัวเคอร์เนลสองครั้ง: หนึ่งครั้งสำหรับการบวก หนึ่งครั้งสำหรับ relu ถึงตอนนี้ก็ยังไม่น่าแปลกใจ

สิ่งที่จริงๆ แล้วน่าสนใจคือสิ่งที่เกิดขึ้นระหว่างสองเคอร์เนลนั้น

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

นับการรับส่งหน่วยความจำสำหรับสองเคอร์เนลนั้น:

• การบวก: อ่าน a, อ่าน b, เขียน tmp สามการถ่ายโอนขนาดอาร์เรย์
• relu: อ่าน tmp, เขียน c อีกสองครั้ง

รวมห้าการถ่ายโอนขนาดอาร์เรย์ เปรียบเทียบกับการบวกขนาด 8 เพียงอย่างเดียวจากหัวข้อก่อนหน้าซึ่งมีสามครั้ง การเพิ่ม .relu() ในห่วงโซ่ไม่ได้ทำให้คุณเสียแค่การคำนวณของ relu มันทำให้คุณเสียการเดินทางรอบเต็มของอาร์เรย์ผ่านหน่วยความจำ GPU เพราะ tmp ต้องถูกเขียนออกมาเพื่อให้เคอร์เนลถัดไปอ่านกลับเข้าไป

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

ทำไม PyTorch ถึงทำแบบนี้? เพราะในโหมด eager เมื่อคุณเขียน a + b PyTorch จะรันมัน ทันที มันไม่รู้ว่าคุณกำลังจะเรียก .relu() ต่อไป แต่ละโอเปอเรชันถูกส่งทันทีที่บรรทัด Python ของมันทำงาน ไม่มีแผน ไม่มีการมองไปข้างหน้า ทุกโอเปอเรชันยืนด้วยตัวเอง สร้างอาร์เรย์จริง และส่งต่อให้สิ่งที่ตามมาผ่านหน่วยความจำ

นี่คือรูปแบบที่ต้องจำไว้ ทุกค่ากลางในโค้ด PyTorch ของคุณจะถูกเขียนออกไปทางกายภาพไปยังหน่วยความจำ GPU และอ่านกลับเข้ามาโดยโอเปอเรชันถัดไป ทุกอัน นั่นคือสิ่งที่ "จำนวนเคอร์เนล" วัดจริงๆ ทุกเคอร์เนลที่เพิ่มขึ้นคือการเดินทางรอบอีกครั้งที่ข้อมูลของคุณต้องทำผ่านหน่วยความจำ GPU

ฟิวชัน: สองโอเปอเรชัน หนึ่งเคอร์เนล

ลองนึกภาพเคอร์เนลหนึ่งตัวที่ทำทั้งหมดนี้ในครั้งเดียว: อ่านองค์ประกอบของ a อ่านองค์ประกอบของ b บวก它们 ใช้ relu กับผลลัพธ์ (ทั้งหมดภายในเคอร์เนล บนพื้นที่ชั่วคราวเล็กๆ ต่อเธรดที่ไม่เคยออกจากชิป) และจากนั้นก็เขียนค่าสุดท้ายไปยัง c ค่ากลาง (a + b) ยังคงมีอยู่ แต่เฉพาะภายในเคอร์เนล ในพื้นที่ส่วนตัวของแต่ละเธรด มันไม่เคยถูกเขียนออกไปยังหน่วยความจำ GPU tmp ในฐานะอาร์เรย์จริง ไม่มีอยู่เลย

นับการถ่ายโอนตอนนี้: การอ่าน a: 1 ครั้งต่อองค์ประกอบ การอ่าน b: 1 ครั้งต่อองค์ประกอบ การเขียน c: 1 ครั้งต่อองค์ประกอบ สามการถ่ายโอนขนาดอาร์เรย์ คณิตศาสตร์เดียวกันกับเวอร์ชันสองเคอร์เนล แต่ลดการเดินทางรอบลงสองครั้ง

ที่ขนาด 8 นี่เป็นข้อผิดพลาดการปัดเศษ ไม่มีใครสนใจ ที่ขนาด 1 ล้าน หรือ 100 ล้าน การเดินทางรอบพิเศษเหล่านั้นกลายเป็นส่วนใหญ่ของเวลารัน และนาฬิกาจับเวลาก็สะท้อนมัน ทำไม* การรับส่งหน่วยความจำถึงครอบงำแบบนั้นเป็นหัวข้อทั้งหมดของบทความที่ 2 ในซีรีส์นี้ ดังนั้นฉันจะปล่อย "ทำไม" ไว้ตรงนี้ ประเด็นตอนนี้คือ: คณิตศาสตร์เดียวกัน การเดินทางน้อยลง เร็วกว่าในทางปฏิบัติ

เทคนิคที่รวมโอเปอเรชันที่อาจเป็นเคอร์เนลแยกกันเป็นเคอร์เนลเดียวเพื่อให้ค่ากลางไม่ต้องไปเยือนหน่วยความจำ GPU มีชื่อเรียกว่า ฟิวชัน นั่นคือคำทั้งหมด นั่นคือแนวคิดทั้งหมด

ทีนี้ส่วนที่ยุ่งยาก การเขียนเคอร์เนลที่รวมกันนั้นด้วยมือดูง่ายสำหรับ add + relu สองโอเปอเรชัน หนึ่งบรรทัด "การคำนวณ" ตรงกลาง แต่โค้ด PyTorc h จริงมีโอเปอเรชันหลายสิบตัวต่อกัน แต่ละตัวมี shape และ dtype และกฎการกระจายเสียงของตัวเอง การเขียนเคอร์เนลแบบฟิวชันที่จัดการทั้งหมดนั้นถูกต้องเป็นงานวิศวกรรมจริงๆ โดยทั่วไปคุณจะไม่เขียนเคอร์เนล elementwise ทั่วไปเหล่านี้ด้วยมือ

ข่าวดี: PyTorch มีเครื่องมือที่ทำการเขียนใหม่ให้คุณ สำหรับกรณีแบบนี้โดยอัตโนมัติ มันเรียกว่า torch.compile

คุณเกือบจะเห็น torch.compile ถูกใช้แบบนี้:

1model = torch.compile(model)

หนึ่งบรรทัด มีคนบนอินเทอร์เน็ตบอกคุณว่ามันทำให้สิ่งต่างๆ เร็วขึ้น นี่คือสิ่งที่มันทำจริง ในภาษาอังกฤษธรรมดา: แทนที่จะรันโอเปอเรชันของคุณทีละตัวแบบโหมด eager torch.compile จะจับการทำงานของเทนเซอร์ที่ฟังก์ชันของคุณทำ หาโอกาสในการรวม它们 และสร้างโค้ดที่ปรับให้เหมาะสม การเรียกในภายหลังที่ตรงกับสมมติฐานเดียวกันสามารถใช้โค้ดนั้นซ้ำได้

ฟิวชันที่เราทำบนกระดาษข้างต้น (add และ relu แชร์เคอร์เนลเดียว tmp ไม่เคยแตะหน่วยความจำ) เป็นสิ่งที่ torch.compile จะทำกับโค้ดของคุณโดยอัตโนมัติ ตราบใดที่โอเปอเรชันนั้นเรียบง่ายพอ เมื่อคนพูดว่า torch.compile "ทำให้ PyTorch เร็วขึ้น" นี่คือส่วนใหญ่ของสิ่งที่พวกเขาหมายถึง

สำหรับกรณีที่ torch.compile ไม่สามารถฟิวชันเองได้ (โอเปอเรชันที่กำหนดเองที่มันจำไม่ได้ การลดรูปที่ผิดปกติ โครงร่างหน่วยความจำที่แปลก) ยังคงมีคนที่ต้องเขียนเคอร์เนลด้วยมือ นั่นคือสิ่งที่เครื่องมืออย่าง Triton และ CUDA มีไว้สำหรับ บทความแยกต่างหาก

เห็นด้วยตัวคุณเอง

ทุกอย่างข้างต้นเป็นการนับเคอร์เนลบนกระดาษ ถึงเวลานับมันบน GPU จริง ถ้าคุณมีเครื่องใดๆ ที่มี CUDA GPU อยู่ในมือ (เวิร์กสเตชัน โน๊ตบุ๊ค Colab อินสแตนซ์คลาวด์) คุณสามารถรันสิ่งนี้เองได้ในไม่กี่นาที

เครื่องมือคือ torch.profiler มันอยู่ใน PyTorc h อยู่แล้ว สิ่งที่มันทำคือบันทึกสิ่งที่ GPU ทำจริงขณะที่โค้ดของคุณทำงาน และส่งคืนตารางที่คุณสามารถอ่านได้

ขั้นตอนที่ 1: เวอร์ชัน eager

ห่อบรรทัดสองโอเปอเรชันในฟังก์ชันเพื่อให้เรามีสิ่งที่เรียก:

1import torch
2from torch.profiler import profile, ProfilerActivity
3
4def add_relu(a, b):
5    return (a + b).relu()
6
7a = torch.randn(1_000_000, device="cuda")
8b = torch.randn(1_000_000, device="cuda")
9
10with profile(activities=[ProfilerActivity.CUDA]) as prof:
11    add_relu(a, b)
12    torch.cuda.synchronize()
13
14print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

torch.cuda.synchronize() มีไว้เพื่อให้แน่ใจว่า GPU ทำงานเสร็จก่อนที่เราจะอ่านเวลาการทำงาน งาน GPU ทำงานแบบอะซิงโครนัส และถ้าไม่มีการซิงค์ คุณจะวัดค่าใช้จ่ายในการเปิดตัวแทนที่จะเป็นงานเคอร์เนลจริง

ขั้นตอนที่ 2: อ่านผลลัพธ์

ผลลัพธ์ profiler จริงของคุณจะมีแถวมากกว่าที่คุณคาดไว้ มีบรรทัดการจัดสรรหน่วยความจำและงานบัญชี PyTorc h หลายรายการปนอยู่ แถวที่เราสนใจคือ CUDA kernels ซึ่งเป็นฟังก์ชันจริงที่ GPU รัน มองหาแถวที่มี kernel ในชื่อ สองแถวนั้นจะมีลักษณะประมาณนี้:

1vectorized_elementwise_kernel<...BinaryFunctor_add...>    1    12us
2vectorized_elementwise_kernel<...threshold_kernel...>     1     9us

ชื่อเทมเพลตที่แน่นอนเปลี่ยนไปตามรุ่น PyTorch (relu มักปรากฏเป็น threshold เพราะนั่นคือโอเปอเรชันพื้นฐาน และ add บางครั้งเป็น CUDAFunctor_add) อย่าพยายามแยกวิเคราะห์ทั้งหมด แค่ นับแถว สองแถว สองเคอร์เนล หนึ่งสำหรับการบวก หนึ่งสำหรับ relu ตรงตามที่เราบอกว่าจะเกิดขึ้นในหัวข้อที่แล้ว

ขั้นตอนที่ 3: เวอร์ชันที่คอมไพล์แล้ว

เปลี่ยนหนึ่งบรรทัด ห่อฟังก์ชันใน torch.compile:

1compiled = torch.compile(add_relu)

ทีนี้มีข้อควรระวังที่ควรรู้ก่อนรัน การเรียก ครั้งแรก ไปยังฟังก์ชันที่ torch.compile จะช้า บางครั้งช้าอย่างน่าแปลก เพราะนั่นคือตอนที่ torch.compile ทำงาน: วิเคราะห์โค้ดของคุณ หาสิ่งที่จะฟิวชัน สร้างเคอร์เนลที่ฟิวชันแล้ว ถ้าคุณโปรไฟล์การเรียกครั้งแรก คุณกำลังวัดขั้นตอนการคอมไพล์ ไม่ใช่เคอร์เนล ดังนั้นรูปแบบคือ: เรียกมันครั้งหนึ่งเพื่ออุ่นเครื่อง ทิ้งผลลัพธ์ จากนั้น โปรไฟล์

1compiled(a, b)                # warm-up, throw away
2torch.cuda.synchronize()
3
4with profile(activities=[ProfilerActivity.CUDA]) as prof:
5    compiled(a, b)
6    torch.cuda.synchronize()
7
8print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

ขั้นตอนที่ 4: อ่านผลลัพธ์อีกครั้ง

คราวนี้:

1triton_poi_fused_add_relu_0    1    14us

หนึ่งแถว หนึ่งเคอร์เนล ชื่อยังบอกคุณว่ามันทำอะไร: fused add และ relu คณิตศาสตร์เดียวกันกับก่อนหน้านี้ หนึ่งการเปิดตัวแทนที่จะเป็นสอง

คุณเพิ่งทำสิ่งที่บทความนี้พูดถึง ในหนึ่งประโยค: คุณขอให้ PyTorc h รวมสองโอเปอเรชันเป็นหนึ่งเคอร์เนล ดู profiler และยืนยันว่าจำนวนลดลงจากสองเป็นหนึ่ง ฟิวชัน ในโลกจริง บนเครื่องของคุณ

ถ้าคุณอยากเห็นมันชัดเจนขึ้น ลองทำที่ขนาดเทนเซอร์ต่างๆ ที่ขนาด 100 ทั้งสองเวอร์ชันทำงานเร็วมากจนความแตกต่างหายไปในสัญญาณรบกวน ที่ขนาด 10 ล้านหรือ 100 ล้าน เวอร์ชันที่คอมไพล์แล้วเริ่มนำหน้าอย่างชัดเจน เพราะการเดินทางรอบที่เราตัดออกเป็นส่วนสำคัญของงานที่ขนาดนั้น

การนับเคอร์เนลไม่ใช่คำแนะนำเชิงนามธรรมอีกต่อไป คุณมีวิธีตรวจสอบแล้ว

สรุป

นี่คือทั้งหมดในครั้งเดียว

โค้ด PyTorch ของคุณ เมื่อมันทำงานบน GPU จะกลายเป็นลำดับของเคอร์เนล แต่ละเคอร์เนลคือหนึ่งการเปิดตัว หนึ่งรอบผ่านข้อมูลของคุณ หนึ่งการเดินทางรอบผ่านหน่วยความจำ GPU เพื่อดึงข้อมูลเข้าและเขียนผลลัพธ์ โอเปอเรชันธรรมดากลายเป็นหนึ่งเคอร์เนล ห่วงโซ่ของโอเปอเรชันกลายเป็นหนึ่งเคอร์เนล ต่อโอเปอเรชัน โดยค่าเริ่มต้น โดยมีค่ากลางเดินทางรอบผ่านหน่วยความจำระหว่างกัน torch.compile สามารถฟิวชันห่วงโซ่ธรรมดาให้คุณได้ ดังนั้นค่ากลางเหล่านั้นไม่เคยแตะหน่วยความจำ เคอร์เนลน้อยลงมักหมายถึงการรับส่งหน่วยความจำน้อยลง และการรับส่งหน่วยความจำน้อยลงมักหมายถึงเร็วกว่า