این صفحه به‌وسیله ‏Cloud Translation API‏ ترجمه شده است.
Switch to English

MLIR CodeGen برای XLA

XLA بر روی HloInstruction عمل می کند و بهینه سازی های زیادی را در این نمایش انجام می دهد ، بسیاری از این موارد را بین دستگاه های هدف به اشتراک می گذارد. همانطور که در برخی نقاط یک برنامه خطی محاسبه می شود و بافر حافظه به طور ثابت به هر مقدار اختصاص می یابد. کدنویسی خاص دستگاه با پیمایش این توالی و فراخوانی "ساطع کننده ها" برای ایجاد نمایشی مناسب برای دستگاه کار می کند (به عنوان مثال یک عملکرد LLVM برای محاسبات XLA بر روی پردازنده ، یا دنباله ای از "thunks" که عملیات GPU را کپسوله می کند و احتمالاً هنگام ایجاد PTX تولید می کند. هدف قرار دادن GPU).

به عنوان یک مرحله مرحله بندی ، در حال حاضر درست پس از اتمام مرحله اختصاصی بافر توسط XLA و رهگیری فرآیند رهگیری فرآیند رهگیری و انتشار ماژول MLIR در گویش lhlo . از آنجا ما بسته به دستگاه با استفاده از م componentsلفه های MLIR (Linalg ، affine و GPU GPL) کد نویسی را انجام می دهیم.

در زیر طرح رکورد برای مهاجرت تدریجی XLA / GPU با استفاده از lhlo به عنوان ورودی codegen آورده شده است.

وظایف

میزبان دستگاه
قالب ورودی HloInstruction * (وظیفه 1) HloInstruction * (وظیفه 1)
فرمت خروجی xla :: Thunk (وظیفه 2) LLVM IR (وظیفه 3)
  • وظیفه 1 هر دو قالب میزبان و ورودی دستگاه را از HloInstruction * به LHLO تغییر می دهد.
  • وظیفه 2 قالب خروجی میزبان را از Thunks به "مقداری صفحه فرود برای میزبان" تغییر می دهد (به زیر مراجعه کنید).
  • وظیفه 3 خروجی دستگاه را از LLVM IR به نوعی از MLIR منتقل می کند. این برای این پروژه اختیاری است و برای جزئیات به بخش "Migrating Device LLVM IR" مراجعه کنید.

این پروژه اولویت بندی داشتن مدلهای قابل اجرا از انتها به انتهای با انتشار دهنده های LHLO تا حد امکان را دارد. این بدان معنی است که لیست ترتیب زیر اهداف بر اساس اولویت است:

  • XLA / GPU را با انتشار دهنده های LHLO ، با Thunks و صادر کننده های موجود اصلاح نشده قابل اجرا کنید.
  • ارجاعات به HloInstruction * را در LHLO حذف کنید ، مورد به مورد:
    • یک امیتر قدیمی را به یک امیتر مبتنی بر MLIR تغییر دهید (به عنوان مثال Linalg) ، یا
    • فرستنده موجود را به صورت مکانیکی ترجمه کنید تا نمایانگر MLIR باشد (انتقال به استاندارد با GPU Dialect).

مهاجرت Thunks (وظیفه 2)

xla :: gpu :: Thunk یک ساختار داده است که:

  • می توان از میزبان فراخوانی کرد (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • داده های مختلفی را در زیر کلاس های خود حمل می کند.
  • با BufferAllocation تعامل دارد :: Slice و StreamExecutor.
  • هسته ها را راه اندازی می کند
  • تماس با همه کتابخانه های زمان اجرا.

هزینه آن شامل موارد زیر است:

  • نمایش داده های پیکربندی مخصوص op (به عنوان مثال پیکربندی های کانولوشن).
  • شکلهای op و عملوند را تغییر می دهد.
  • نمایندگی یک درخت تنه (در حالی که ، شرایط ، و غیره).

کار مهاجرت مستقل از مهاجرت LHLO / امیتر است. تحت منابع محدود ، اولویت آن پشت مهاجرت LHLO / انتشار دهنده است.

ما چندین گزینه در مورد نحوه پایین آوردن قسمت سمت میزبان از LHLO داریم:

  • TFRT
    • (نرم افزار) بسته بندی های عالی CUDA و HIP برای استفاده.
    • (نرم افزار) اجرای تماس های کتابخانه ای (cuDNN ، cuBLAS ، cuFFT و غیره) آسان است ، زیرا عملکردهای TFRT با کد ++ C تفسیر می شوند.
    • طرف میزبان در حال توسعه است و آزمایش نشده است.
  • کد CPU جنجالی
    • (حرفه ای) توانایی کمتری عالی چند حلقه و شرط ایجاد کنید و تمام شود.
    • (Con) GPUDialect هنوز زنجیره ها / جریان ها / عدم همزمان بودن / تخصیص دستگاه را مدل نمی کند.
    • (Con) پشتیبانی از زمان اجرا CUDA / HIP حداقل است (مسیر جعبه ابزار ، نسخه ، بارگذاری پویا و غیره).
  • زمان اجرا (تفسیر) XLA موجود

تصمیم گیری: TFRT را اتخاذ کنید ، اما از کد CPU متناوب در TFRT نیز پشتیبانی می کند.

مهاجرت دستگاه LLVM IR (وظیفه 3)

یک ساطع کننده اصلی با پر کردن عنصر به عنصر ، هدف را ایجاد می کند. هر عنصر خروجی به مجموعه ای از عناصر عملوند بستگی دارد. تمام عناصر با ترکیب بافر با شاخص های دینامیکی توصیف می شوند. توصیف تقریباً همه عملکردهای "ریاضی" کافی است ، اما به دلایل عملکردی فقط زیرمجموعه بزرگی از عملکردهای "ریاضی" مستقیماً در (Cpu | Gpu) ElementalIrEmitter اجرا می شود.

ElementalIrEmitter از این نظر منحصر به فرد است:

  • بخش بزرگی از کد بین XLA / GPU و CPU مشترک است.
  • این نمایانگر بخش بزرگی از عملیات است که در مدل ها دیده می شود ، از جمله همه اپ های عملیاتی که از نظر عناصر انتخاب شده اند.
  • اکثر ترکیبات فقط به ElementalIrEmitter بستگی دارد.
  • از لحاظ ساختاری ساده است ، زیرا یک DAG وابستگی به داده را بین عناصر op و عناصر عملوند توصیف می کند.
  • بیشتر قابل حمل و سطح بالایی است (مثلاً برخلاف GPU kReduce و GPU kCopy).
  • پشتیبانی از شکل پویا حداقل برای عملکردهای عنصر آسان است.

اکنون ، برای همه اپراتورها ، که از نظر ذاتی منتشر می شوند یا نه ، طعم های مختلفی از حالت نهایی هر عملکرد XLA وجود دارد:

  1. کد دستگاه به صورت LLVM IR باقی می ماند.
  2. بازتابنده انتشار دهنده قدیمی مانند LHLO -> MLIR LLVM گویش:
    • (هزینه) اگر بخواهیم در نهایت به استاندارد مهاجرت کنیم کار دور ریز خواهد بود.
    • (سود) آسان و مکانیکی است. در یک دوره کوتاه انجام می شود.
    • (سود) در مقایسه با (1) سود بیشتری ندارد.
  3. انتشار دهنده های قدیمی refactor مانند LHLO -> MLIR GPU + Standard + Loops:
    • (هزینه) بلند کردن آلاینده های موجود به استاندارد چالش هایی را به وجود می آورد. اشاره گرها و GEP ها باید به MemRefs و SubViews تبدیل شوند. اطمینان از کامل بودن amdgpu یکی دیگر از این موارد است.
    • (هزینه) XLA / GPU به شدت به فراداده LLVM متکی است:
      • range برای شاخص های بلوک / رشته.
      • align ، dereferenceable ، invariant.load ، alias.scope ، noalias برای بار / فروشگاه ها.
      • llvm.loop.unroll.disable ، llvm.loop.unroll.full ، llvm.loop.vectorize.enable برای حلقه های متوالی.
    • (سود) می تواند طولانی مدت باشد. قابل حمل تر
  4. انتشار دهنده های قدیمی را به LHLO تبدیل کنید -> Linalg ، و انتشار دهنده های جدید Linalg را بنویسید
    • (هزینه) این مورد به مورد است. در مقایسه با گزینه های قبلی ، اجرای جدیدی که با عملکرد XLA مطابقت داشته باشد باید معیار <-> بهینه سازی گردش کار را طی کند ، که برای برخی از گزینه ها می تواند هزینه قابل توجهی داشته باشد.
    • (منفعت) پشته متحد ؛ پشتیبانی جامعه؛ قابل حمل بودن پتانسیل های بهینه سازی بیشتر

نتیجه گیری:

  • دنبال (2) نروید. (1) یا (3) فقط بهتر از (2) هستند. (2) هزینه بیش از (1) دارد ، زیرا به بازسازی مکانیکی زیادی نیاز دارد. با استفاده از (1) هنوز می توانیم به هدف توانائی XLA برای برداشتن انتشار دهنده های MLIR برسیم. این کار با انجام LHLO -> LLVM IR -> پخش کننده های قدیمی دستگاه انجام می شود.
  • گزینه های ElementalIrEmitter برای (4) است ، اما نه به صورت افزایشی. هیچ راهی برای انجام آن توسط op با op وجود ندارد ، زیرا تمام عملیات op-منتشر شده به یک نمودار متصل می شوند. این کار همچنین می تواند به عنوان یک نقطه وحدت چندین نیروی در حال اجرا (xla / service / mlir_gpu ، هسته ساز ، Linalg) باشد.
  • همه گزینه های دیگر برای (1) است. به عنوان یک هدف کششی ، ممکن است به (3) یا (4) منتقل شوند.

اولویت بندی

در حالی که هر سه وظیفه ذکر شده در بالا قابل موازی سازی هستند ، اما تحت منابع محدود باید آنها را به صورت سریالی تنظیم کرد. اولویت بندی بر نتایج قابل مشاهده برای تکمیل هر کار متمرکز است.

اولویت بندی این است: Task1 (LHLO برای انتشار دهنده های قدیمی)> Task 2 (Thunks)> وظیفه 3 (MLIR emitters).

با پایان کار 1 ، کاربران XLA می توانند یک LHLO (به عنوان مثال ژنراتور هسته) تولید کرده و آنها را اجرا کنند. قالب تدوین MLIR قابل سریال نیست.

با پایان کار 2 ، LHLO به MLIR مناسب و سریالی کاهش می یابد. این امکان جمع آوری آفلاین را فراهم می کند.

با پایان کار 3 ، تمام انتشار دهنده های XLA در اجرای آن مبتنی بر MLIR هستند.

طراحی دقیق و با جزییات

مرحله 1: (وظیفه 1) LHLO را کامل کنید و Make Legacy Emitters Take LHLO کنید

این مرحله باعث می شود تا تمام انتشار دهنده های موجود XLA / GPU با عملکردهای MLIR تعامل داشته باشند. این مرحله بازسازی و NFC است.

این مرحله عمدتا مکانیکی است ، اما باید تفاوت های زیر را بین یک HloComputation بدون آزمایش و LHLO مشاهده کرد:

  • هر HloInstruction به عملوندهای خود دسترسی مستقیم دارد (DAG جریان داده). برعکس ، هر عملیات LHLO فقط به بافرهای عملوند خود دسترسی دارد (دو بخشی بین گزینه ها و بافرها). اپراتورهای LHLO برای دسترسی به اپراتورهای خود مجبور به گذر از زنجیره های استفاده از Def هستند.
  • انتشار دهنده های میراثی تودرتو از نظر تجربی تقریباً هرگز به مواد عملیاتی خود دسترسی ندارند. تنها استثنا kReduce است.
  • ساطع کنندگان میراث ثبت نشده فقط برای بدست آوردن برش ها به BufferAssign دسترسی پیدا می کنند نه برای دسترسی به ساختارهای داده aux مانند dataflow_analysis () یا alias_analysis (). llvm_ir نام مستعار_تحلیل () خود را بر اساس اطلاعات قطعه ایجاد می کند.

نتیجه گیری این است که LHLO باید بدون دردسر اساسی مناسب باشد.

مرحله 2: (اختیاری) پشتیبانی از پروفایل

این مرحله فقط درصورت نیاز است که برخی از منطق های XLA Thunk را کنار بگذاریم (مرحله بعدی را ببینید).

قبل از روشن کردن هر نوع ساطع کننده مبتنی بر MLIR ، به پروفایل برای ساطع کننده های MLIR نیاز داریم.

در حال حاضر XLA با تماس گرفتن با تایمر StreamExecutor ، پروفایل خود را انجام می دهد. تایمر زیر هود دو رویداد را قبل و بعد از راه اندازی هسته وارد می کند و زمان همگام سازی بین این دو رویداد را اندازه گیری می کند.

تقریباً سه روش برای پشتیبانی از پروفایل در MLIR وجود دارد:

  • پروفایلر را از انتها به انتهای دیگر اجرا کنید
  • با استفاده از یک پروفایل تزریقی ، برای هر op در LHLO یک نمایه نمایه اضافه کنید.

رویکرد "پایان به انتها" برای MLIR شفاف است ، اما همان مشکلی را دارد که باعث می شود XLA در وهله اول از آن استفاده نکند: تماس های کتابخانه ای جمع آوری شده توسط یک پروفایلر (nvprof / ...) نمی توانند به راحتی با HLO ارتباط برقرار کنند ops به عنوان مثال ، cuDNN چندین هسته را برای هر HLO راه اندازی می کند و نمی توان تشخیص داد که کدام هسته با کدام HLO مطابقت دارد.

رویکرد "پروفایل ساز تزریقی" به موارد زیر نیاز دارد:

  • LHLO برای گرفتن یک پروفایلر به عنوان یک پارامتر.
  • درج profile.start / profile.end قبل و بعد از هر عمل.
  • پاس از آن نمایه را کاهش می دهد. {شروع ، پایان} به اجرای ++ C.

مشخصات دقیق را نمی توان به راحتی برای برنامه های تولیدی MLIR انجام داد ، زیرا:

  • MLIR تایمر ندارد و همچنین به TFRT / StreamExecutor بستگی ندارد.
  • MLIR به راحتی توابع C را با پارامترهای پیچیده فراخوانی نمی کند.

مرحله 3: (وظیفه 2) مهاجرت به Thunks

به عنوان یک یادداشت ، تقریباً سه نوع نظر وجود دارد:

  • KernelThunk ، که یک هسته را راه اندازی می کند.
  • Thunks جریان کنترل ، که منطق جریان کنترل میزبان را دارد (شرطی می شود ، در حالی که ، برای ، توالی است) و هسته های بدن را راه اندازی می کند.
  • نظرات کتابخانه: cuDNN ، cuBLAS ، cuFFT ، NCCL و غیره

برنامه این است:

  • Thunks (غیر) را قابل سریال سازی کنید.
  • به بهبود TFRT در حالتی که بتواند از این معناشناسی پشتیبانی کند ، کمک کنید.
  • با بهبود وضعیت ، تک تک افراد را به تدریج مهاجرت کنید.

این موارد اقدام فقط جزئی سفارش داده می شوند. دستورالعمل اعدام واقعی / موازی کاری مهندسی قرار است به همین ترتیب ارزیابی شود.

مرحله 4: (وظیفه 3) ElementalIrEmitter منتقل شده

پس از آماده شدن پروفایل ، می توانیم تمام انتشار دهنده های مبتنی بر ElementalIrEmitter را در MLIR تکمیل و تنظیم کنیم. سپس با فرض اینکه همه این ساطع کننده های مبتنی بر MLIR از یک جریان واحد استفاده می کنند ، آنها را به طور پیش فرض روشن می کنیم.

توجه داشته باشید که انتقال ElementalIrEmitter از XLA / CPU نیز مفید است ، زیرا بخش بزرگی از کد را به اشتراک می گذارند.

با انجام همه معیارها و شکار عملکرد (TODO: تعریف برابری عملکرد) ، ما جدیدترین امیتر جدید مبتنی بر MLIR را روشن می کنیم و میراث ElementalIrEmitter را حذف می کنیم.

این مرحله همچنین انتقال های همجوشی (گزینه های تو در تو) را برای انتقال بعدی فراهم می کند.

مرحله 5: پشتیبانی از چند جریان یا رها کردن

تا زمانی که از MLIR پشتیبانی نکنیم ، نمی توانیم برخی از ساطع کننده ها را حذف کنیم یا اینکه ویژگی را رها کنیم. این مقدار نسبتاً زیاد کار در MLIR و اندکی سود برای XLA است. ما باید کاربران فعلی کاربران چند جریان XLA / GPU را بررسی کنیم و اگر منطقی است این ویژگی را حذف کنیم.

مرحله ششم: (وظیفه 3) عملیات دستگاه منتقل شده

این مرحله باعث انتقال همه عملکردهای غیرقابل توجیه می شود ، در این صورت می توانیم تمام انتشار دهنده های غیرقابل آزمایش را حذف کنیم.

این فراخوانی مجدد / refactor را برای kCopy و kReduce فراخوانی می کند. kReduce قبلاً به مقدار کافی کار شده است ، بنابراین مقدار واقعی کارهایی که باید انجام شود باید مشخص شود.