Hardware accelerators, especially those designed for tensor processing, have become ubiquitous in today's computing landscape. However, even with significant efforts in building compilers, programming these tensor accelerators remains challenging, leaving much of their potential underutilized. Recently, large language models (LLMs), trained on large amounts of code, have shown significant promise in code generation and optimization tasks, but generating low-resource languages, such as specialized tensor accelerator code still poses a significant challenge. We tackle this challenge with Autocomp, an approach that empowers accelerator programmers to leverage domain knowledge and hardware feedback to optimize code via an automated LLM-driven search. We accomplish this by: 1) formulating each optimization pass as a structured two-phase prompt, divided into planning and code generation phases, 2) inserting domain knowledge during planning via a concise and adaptable optimization menu, and 3) integrating correctness and performance metrics from hardware as feedback at each search iteration. Across three distinct hardware platforms, we demonstrate that Autocomp-optimized code runs 5.6x faster than the vendor-provided library (Gemmini), outperforms expert-level hand-tuned code by 1.9x (AWS Trainium), and achieves 3.8x higher performance than a machine learning-based cost model for GPUs (NVIDIA L40S). Additionally, we demonstrate that optimization schedules generated from Autocomp can be reused across similar tensor operations, improving speedups by up to 24% under a fixed sample budget.
academic- معرّف الورقة: 2505.18574
- العنوان: Autocomp: A Powerful and Portable Code Optimizer for Tensor Accelerators
- المؤلفون: Charles Hong, Sahil Bhatia, Alvin Cheung, Yakun Sophia Shao (جامعة بيركلي)
- التصنيف: cs.PL cs.AI cs.AR cs.LG
- حالة النشر: نسخة أولية قيد المراجعة
- رابط الورقة: https://arxiv.org/abs/2505.18574
أصبحت معالجات الأجهزة المتخصصة، وخاصة تلك المصممة لمعالجة الموترات، منتشرة في بيئات الحوسبة الحديثة. ومع ذلك، رغم الجهود الكبيرة المبذولة في بناء المترجمات، لا تزال برمجة هذه المسرعات تشكل تحديًا كبيرًا، مما يترك إمكاناتها غير مستغلة إلى حد كبير. تقترح هذه الورقة Autocomp، وهي طريقة لتحسين الكود من خلال البحث الآلي المدفوع بنماذج اللغة الكبيرة (LLM)، مما يمكّن مبرمجي المسرعات من الاستفادة من المعرفة المتخصصة والتغذية الراجعة من الأجهزة. يتحقق هذا النهج من خلال ثلاث تقنيات رئيسية: 1) صياغة كل عملية تحسين كطلب منظم على مرحلتين، مقسم إلى مرحلة التخطيط وتوليد الكود؛ 2) إدراج المعرفة المتخصصة أثناء مرحلة التخطيط من خلال قائمة تحسين موجزة وقابلة للتكيف؛ 3) دمج مقاييس الصحة والأداء من الأجهزة كتغذية راجعة في كل تكرار بحث.
تشمل التحديات الرئيسية في برمجة مسرعات الموترات:
- تعقيد البرمجة: بخلاف برمجة المعالجات العامة، تتطلب مسرعات الموترات إدارة صريحة لحركة البيانات وتكوين الحالة وجدولة العمليات
- تكاليف تكييف المترجم: يتطلب تكييف المترجمات التقليدية للمنصات الجديدة عملاً هندسياً كبيراً، حيث تمثل تكاليف تطوير البرامج 40-50% من تكاليف تطوير الأجهزة الجديدة
- مشاكل جدولة التحسين: مشكلة الانفجار التوافقي في تحديد التحسينات المراد تطبيقها وترتيب تطبيقها
- تحديات اللغات منخفضة الموارد: مجموعات التعليمات المعمارية (ISA) واللغات المجال المتخصصة (DSL) للمسرعات المتخصصة ممثلة بشكل ناقص في مجموعات بيانات تدريب نماذج اللغة الكبيرة
- المترجمات التقليدية: XLA و TVM و Triton وغيرها تدعم عدداً قليلاً من الواجهات الخلفية، بشكل أساسي المعالجات والمعالجات الرسومية
- طرق DSL: توفر Halide و Exo بدائيات للتعبير عن الحسابات الموترية، لكن عبء التحسين لا يزال على عاتق المبرمج
- الطرق المدفوعة بالبيانات: تتطلب كميات كبيرة من بيانات الأداء للتدريب، وهي نادرة جداً بالنسبة لمسرعات الأجهزة المتخصصة
- تطبيق LLM المباشر: بالنسبة للغات المسرعات منخفضة الموارد، يكون توليد الكود بدون عينات غير موثوق به بشكل كبير
- أول طريقة لتحسين كود مسرعات الموترات منخفضة الموارد مدفوعة بـ LLM: تقترح إطار عمل Autocomp متخصص للمسرعات المتخصصة
- إطار عمل تحسين قابل للنقل بدرجة عالية: يمكن تكييفه مع منصات أجهزة جديدة بمجرد تعديل الطلبات، مما يقلل بشكل كبير من التكاليف الهندسية
- أداء متفوق: يتفوق بشكل ملحوظ على الطرق الموجودة على ثلاث منصات أجهزة مختلفة
- آلية إعادة استخدام الجدولة: يثبت أن جداول التحسين يمكن إعادة استخدامها بين العمليات الموترية المتشابهة، مما يحسن كفاءة العينة
الإدخال: كود مسرع موتر غير محسّن
الإخراج: كود معادل وظيفياً لكن محسّن الأداء
القيود: الحفاظ على التكافؤ الدلالي، التحقق من الصحة من خلال الأجهزة
يتضمن هيكل الطلب:
- وصف ISA للمسرع: دلالات التعليمات، معايير عنونة الذاكرة، وصف البنية الأجهزية
- الكود الحالي: الكود المراد تحسينه
- التغذية الراجعة للأداء: الكمون (عدد الدورات)، معدل استخدام الذاكرة وغيرها
- قائمة التحسين: خيارات تحسين عالية المستوى محددة مسبقاً (مثل تجانب الحلقات، إعادة الترتيب، الدمج، إلخ)
- معلومات تكرار البحث: رقم التكرار الحالي، لتوجيه اختيارات التحسين
يتضمن هيكل الطلب:
- وصف ISA للمسرع: نفس المرحلة 1
- الكود الحالي: نفس المرحلة 1
- الخطة المولدة: خطة التحسين المحددة من مخرجات المرحلة 1
- أمثلة التعلم السياقي: أمثلة كود للتحسينات المعقدة (مثل التجانب)
- تعليمات التنفيذ: تعليمات باللغة الطبيعية لتطبيق الخطة وإخراج الكود المحسّن
- بحث بحزم بعرض B=6، استكشاف متوازي لمسارات تحسين متعددة
- تصفية الصحة: التحقق من الكود المرشح من خلال مجموعة اختبارات وظيفية
- فحص الأداء: الاحتفاظ فقط بالمرشحين الذين يتفوقون على الأداء من العقدة الأب
- التحسين التكراري: عملية بحث بميزانية ثابتة T من التكرارات
- حذف قائمة التحسين: إزالة عشوائية لأجزاء من قائمة الخيارات أثناء كل تخطيط (احتمالية 70%)
- تكامل LLM: توزيع الطلبات على نماذج LLM متعددة لزيادة تنوع الاستجابات
- مقاييس الأداء الفورية (الكمون، معدل استخدام الذاكرة) توجه اختيار خطوات التحسين التالية
- محاكاة دقيقة على مستوى الدورة أو قياس الأداء على مستوى الشريحة
- تسجيل تسلسلات جدولة عالية الجودة
- إعادة استخدام الجداول المعروفة للعمليات الموترية المتشابهة (نفس النسبة أو الأبعاد المشتركة)
- بحث خفيف الوزن متبوع بمزيد من التحسين
- Gemmini: مولد مسرع مفتوح المصدر، يدعم مسرعات موترات بأسلوب الصفيف الانقباضي والمتجهات
- AWS Trainium: مسرع موتر عالي الأداء تجاري، يستخدم واجهة Neuron Kernel Interface (NKI)
- NVIDIA L40S GPU: معالج رسومات حديث لمراكز البيانات، يتضمن Tensor Core متخصص
- Gemmini: GEMM والالتفاف من ResNet-50، التنبؤ بالتحكم في نموذج TinyMPC
- Trainium: عمليات التعلم العميق على مستوى البرنامج التعليمي والمتقدم (RMSNorm و LayerNorm و GEMM و Mamba وغيرها)
- GPU: اختبارات KernelBench المستوى 1
- مكتبات البرامج المتقدمة: مكتبة Gemmini و PyTorch NeuronX و PyTorch
- كود منخفض المستوى غير محسّن: Exo غير محسّن وكود البرنامج التعليمي nki-samples
- كود محسّن يدويًا: تطبيقات محسّنة يدويًا على مستوى الخبراء
- نماذج تكلفة ML: TVM MetaSchedule (GPU)
- آلة الحالة الأجهزية: آلة الحالة المحدودة للأجهزة Gemmini (الحد الأعلى المرجعي)
- اختبار GEMM: تحسن 5.6× مقارنة بمكتبة Gemmini، يتفوق على كود التحسين اليدوي للخبراء بـ 1.4×
- اختبار الالتفاف: تحسن 2.6× مقارنة بمكتبة البرامج، يتفوق على التحسين اليدوي بـ 1.1×
- الجبر الخطي الدقيق: يتفوق على الكود غير المحسّن بـ 2.7×، وحتى يتفوق على تطبيق آلة الحالة المحدودة للأجهزة المحسّنة يدويًا بـ 1.6× (الانتشار الأمامي)
- أحمال العمل التعليمية: يتفوق على كود التحسين اليدوي بـ 1.36× (المتوسط الهندسي)، يتفوق على كود PyTorch NeuronX المترجم بـ 13.52×
- أحمال العمل المتقدمة: يتفوق على كود التحسين على مستوى الخبراء بـ 1.9× (المتوسط الهندسي)، تحسن يصل إلى 17.37× في الالتفاف العميق أحادي البعد
- اختبارات KernelBench: يتفوق على PyTorch بـ 2.05× (المتوسط الهندسي)، يتفوق على TVM MetaSchedule بـ 3.8×
- يتفوق على PyTorch في جميع الاختبارات المرجعية، بينما يتفوق TVM على PyTorch في اختبارين فقط
تتحقق الدراسات الاستئصالية التفصيلية من أهمية كل مكون:
- ISA للمسرع: إزالته تؤدي إلى انخفاض كبير في الأداء، لكن لا يزال يحقق تحسناً
- قائمة التحسين: ضرورية تماماً، إزالتها تؤدي إلى فشل كامل في الأداء
- حذف قائمة التحسين: يؤثر بشكل كبير على الأداء، مما يمنع انحياز النموذج نحو خيارات قائمة محدودة
- تكامل LLM: يوفر تنوعاً مهماً، أداء نموذج واحد أقل
- التغذية الراجعة لأداء الأجهزة: مفيدة لكن محدودة التأثير، لأن قائمة التحسين تتضمن بالفعل المقاييس ذات الصلة
- مع ميزانية 100 عينة: جدولة معاد استخدامها تحقق تسريع 4.6×، بدون إعادة استخدام فقط 3.7×
- مع ميزانية 200 عينة: جدولة معاد استخدامها تحقق تسريع 5.0×، بدون إعادة استخدام فقط 4.2×
- يثبت قابلية تعميم الجدولة، مما يقلل بشكل فعال من تكاليف البحث للاختبارات المرجعية المتشابهة
- نماذج الأداء: Timeloop و MAESTRO تستخدم نماذج معمارية أجهزة عالية المستوى
- الطرق الآلية: التعلم الآلي والبرمجة الخطية والتحسين الأسود الصندوق والتعلم المعزز
- القيود: التجريدات الموجودة تتجاهل التحسينات الخاصة بالتنفيذ ومستوى التعليمات
- نطاق التطبيق: البحث التطوري والاسترجاع المعزز والتحسين التكراري وما بعد التدريب
- التحسينات على مستوى النظام: CUDA وعوامل SIMD الجوهرية
- الفجوة البحثية: نقص العمل على تحسين كود LLM للأجهزة المتخصصة (غير CPU/GPU)
- فعالية التحسين المدفوع بـ LLM: يتفوق Autocomp بشكل ملحوظ على الطرق التقليدية على منصات أجهزة متعددة
- قابلية نقل عالية جداً: يمكن تكييفه مع أجهزة جديدة بمجرد تعديل الطلبات، بتكاليف هندسية منخفضة جداً
- قيمة إعادة استخدام الجدولة: جداول التحسين لها قابلية تعميم جيدة، مما يحسن بشكل ملحوظ من كفاءة العينة
- ضرورة التصميم على مرحلتين: فصل مرحلة التخطيط والتنفيذ يحسن معدل النجاح في مهام التحسين المعقدة
- أهمية المعرفة المتخصصة: المعرفة المتخصصة التي توفرها قائمة التحسين حاسمة للأداء
- قيمة التغذية الراجعة من الأجهزة: مقاييس الأداء الفورية توجه بشكل فعال اختيار اتجاه التحسين
- الاعتماد على قدرات LLM: يقتصر أداء الطريقة على قدرات توليد الكود والاستدلال للنموذج الأساسي
- تكاليف البحث: يتطلب استدعاءات LLM متعددة ومحاكاة أجهزة، بتكاليف حسابية أعلى
- الخصوصية المجالية: تتطلب قائمة التحسين تصميماً يدويًا لمنصات أجهزة مختلفة
- نطاق التقييم: يركز بشكل أساسي على أحمال عمل الحسابات الموترية، وتطبيقها على أنواع حسابات أخرى غير معروف
- توليد القائمة الآلي: البحث عن طرق لبناء قائمة التحسين تلقائياً
- هجرة الجدولة عبر المنصات: استكشاف نقل معرفة الجدولة بين منصات أجهزة مختلفة
- تحسين كفاءة التكلفة: تقليل عدد استدعاءات LLM ومحاكاة الأجهزة في عملية البحث
- تطبيقات أوسع: التوسع إلى مسرعات متخصصة أخرى بخلاف حسابات الموترات
- ابتكار قوي: أول تطبيق لـ LLM على تحسين كود مسرعات الموترات منخفضة الموارد، مع مسار تقني جديد
- قيمة عملية عالية: يحل مشاكل هندسية فعلية، مما يقلل بشكل كبير من تكاليف تطوير البرامج للأجهزة الجديدة
- تجارب شاملة: تقييم كامل على ثلاث منصات أجهزة مختلفة، مع نتائج مقنعة
- طريقة عامة: تصميم الإطار يتمتع بقابلية توسع وقابلية نقل جيدة
- أداء متفوقة: يتفوق بشكل ملحوظ على أفضل الطرق الموجودة في اختبارات مرجعية متعددة
- التكاليف الحسابية: يتطلب عدداً كبيراً من استدعاءات LLM ومحاكاة الأجهزة، قد يحد من التطبيق العملي
- الاعتماد على التصميم اليدوي: تتطلب قائمة التحسين معرفة خبراء يدوية، درجة الأتمتة محدودة
- قيود التقييم: يركز بشكل أساسي على أنواع معينة من حسابات الموترات، تحتاج التعميمية إلى التحقق
- نقص التحليل النظري: يفتقر إلى ضمانات نظرية حول تقارب الطريقة وأمثليتها
- القيمة الأكاديمية: يفتتح تطبيق LLM في مجال تحسين الترجمة للأجهزة المتخصصة، بأهمية أكاديمية كبيرة
- التأثير الصناعي: من المتوقع أن يقلل بشكل كبير من تكاليف تطوير مجموعة البرامج للأجهزة الجديدة، بقيمة صناعية كبيرة
- قابلية الاستنساخ: التزم المؤلفون بفتح المصدر والطلبات، مما يسهل الأبحاث اللاحقة
- الإلهام: يوفر مسار تقني جديد لتحسين الترجمة للأجهزة المتخصصة الأخرى
- تطوير نماذج أجهزة جديدة: توليد سريع لكود محسّن لمسرعات موترات مصممة حديثاً
- بناء مترجمات DSL: كمكمل أو بديل لطرق المترجمات التقليدية
- أدوات ضبط الأداء: مساعدة المطورين على تحسين كود المسرع الموجود
- البحث والتعليم: توفير أدوات آلية لبرمجة وتحسين المسرعات
تستشهد الورقة بعدد كبير من الأعمال ذات الصلة، تشمل بشكل أساسي:
- تصميم معالجات الأجهزة (Gemmini و TPU و Trainium وغيرها)
- المترجمات و DSL (XLA و TVM و Halide و Exo وغيرها)
- توليد كود LLM (CodeGen و Codex وغيرها)
- طرق التحسين الآلي (التعلم المعزز والخوارزميات التطورية وغيرها)
التقييم الإجمالي: هذه ورقة بحثية عالية الجودة تقدم مساهمات مهمة في مجال التقاطع الناشئ بين تطبيقات LLM وتحسين الترجمة للأجهزة المتخصصة. الطريقة مبتكرة، التقييم التجريبي شامل، والقيمة العملية كبيرة. على الرغم من وجود مجال للتحسين في التكاليس الحسابية ودرجة الأتمتة، فإن الورقة تفتتح اتجاهاً جديداً لتطور هذا المجال، بقيمة أكاديمية وصناعية مهمة.