اشکال و چیدمان

ساختار یک عمل جراحی XLA

یک مثال HLO را در نظر بگیرید:

add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)}
          add(exponential.183, broadcast.3115)

این شامل اجزای زیر است:

  • نام اپراتور: add.936
    • این نام منحصر به فرد برای این عملیات است.
  • شکل: bf16[8,1,1280,16384]
    • این شکل خروجی Op است. در اینجا dtype برابر با bf316 و شکل آن [8,1,1280,16384] است.
  • طرح‌بندی (با کاشی‌کاری): 3,2,0,1:T(8,128)(2,1)
    • این نحوه ذخیره آرایه در حافظه را توصیف می‌کند. 3,2,0,1 ترتیب محورها در حافظه (یعنی ستون اصلی، ردیف اصلی و غیره) را نشان می‌دهد و T(8,128)(2,1) نشان دهنده کاشی‌کاری و لایه‌گذاری مورد استفاده است.
    • طرح‌بندی اختیاری است. اگر مشخص نشود، کاشی‌کاری انجام نمی‌شود و فرض می‌شود ابعاد از بزرگ به کوچک مرتب شده‌اند.
  • عملیات: add
    • عملیاتی که انجام می‌شود. در اینجا، Add است که در نام Op نیز ذکر شده است.
  • آرگومان‌ها: exponential.183 ، broadcast.3115
    • این عملیات دو آرگومان می‌گیرد که با نام‌های منحصر به فردشان مشخص شده‌اند.

بیایید مثال دیگری، یک عملیات فیوژن، را در نظر بگیریم:

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)}
            fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32),
            kind=kCustom, calls=%all-reduce-scatter.3

علاوه بر اجزای شرح داده شده قبلی، این شامل موارد زیر است:

  • ویژگی‌ها: kind و calls
    • این‌ها اطلاعات بیشتری در مورد عملیاتی که انجام می‌شود، در این مورد: فیوژن، ارائه می‌دهند.
  • مکان حافظه (شناسه فضای حافظه): S(1)
    • این نشان دهنده فضای حافظه/مکان ذخیره آرایه است. S(1) در اینجا نشان دهنده این آرایه است که در VMEM (روی یک TPU) قرار دارد.
  • جزئیات شکل و طرح‌بندی برای آرگومان ورودی %fusion.32

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

اشکال

پروتو XLA ShapeProto ( xla_data.proto ) تعداد ابعاد، اندازه و نوع داده یک آرایه N بعدی (به طور خلاصه آرایه ) را توصیف می‌کند.

اصطلاحات، نمادگذاری و قراردادها

  • تعداد واقعی ابعاد یک آرایه، تعداد ابعادی است که اندازه آنها بزرگتر از ۱ است.

  • ابعاد برای یک آرایه N بعدی از 0 تا N-1 شماره گذاری می‌شوند. اندازه یک بعد یک عدد صحیح غیر منفی است. به طور خاص، اندازه 0 معتبر است. شماره‌های بعد برای راحتی، برچسب‌های دلخواه هستند. ترتیب این شماره‌های بعد به معنای ترتیب جزئی/کلی خاصی در طرح شکل نیست. طرح توسط Proto LayoutProto تعیین می‌شود.

  • طبق قرارداد، ابعاد به ترتیب افزایش شماره بعد فهرست می‌شوند. برای مثال، برای یک آرایه سه‌بعدی با اندازه [A x B x C] ، بعد ۰ اندازه A ، بعد ۱ اندازه B و بعد ۲ اندازه C را دارد.

    برخی از ابزارهای XLA نیز از اندیس‌گذاری منفی مشابه پایتون پشتیبانی می‌کنند: بُعد -۱ آخرین بُعد است (معادل N-1 برای یک آرایه N بعدی). برای مثال، برای آرایه سه‌بعدی که در بالا توضیح داده شد، بُعد -۱ اندازه C ، بُعد -۲ اندازه B و غیره دارد.

  • آرایه‌های دو، سه و چهار بعدی اغلب حروف خاصی دارند که به ابعاد آنها مربوط می‌شود. برای مثال، برای یک آرایه دوبعدی:

    • بُعد ۰: y
    • بُعد ۱: x

    برای یک آرایه سه بعدی:

    • بُعد ۰: z
    • بُعد ۱: y
    • بُعد ۲: x

    برای یک آرایه چهاربعدی:

    • بعد ۰: p
    • بُعد ۱: z
    • بُعد ۲: y
    • بُعد ۳: x
  • توابعی در XLA API که ابعاد می‌گیرند، این کار را به ترتیب افزایشی شماره بعد انجام می‌دهند. این با ترتیبی که هنگام ارسال ابعاد به عنوان یک initializer_list استفاده می‌شود، مطابقت دارد؛ مثلاً

    ShapeUtil::MakeShape(F32, {A, B, C, D})

    شکلی ایجاد می‌کند که آرایه اندازه ابعاد آن شامل دنباله [A, B, C, D] است.

طرح بندی

پروتوتایپی LayoutProto نحوه نمایش یک آرایه در حافظه را شرح می‌دهد. این پروتوتایپ شامل فیلدهای زیر است:

message LayoutProto {
  repeated int64 minor_to_major;
  int64 tail_padding_alignment_in_elements;
  ...
}

مرتب‌سازی ابعاد کوچک به بزرگ

تنها فیلد مورد نیاز، minor_to_major است. این فیلد ترتیب ابعاد درون یک شکل را از کوچک به بزرگ توصیف می‌کند. مقادیر موجود در minor_to_major ترتیبی از ابعاد آرایه هستند ( 0 تا N-1 برای یک آرایه N بعدی) که اولین مقدار، کوچکترین بُعد تا آخرین مقدار که بزرگترین بُعد است، می‌باشد. کوچکترین بُعد، بُعدی است که هنگام حرکت در عناصر آرایه که در حافظه خطی قرار گرفته‌اند، با بیشترین سرعت تغییر می‌کند.

برای مثال، آرایه دوبعدی زیر با اندازه [2 x 3] را در نظر بگیرید:

a b c
d e f

در اینجا بُعد 0 اندازه ۲ و بُعد 1 اندازه ۳ دارد. اگر فیلد minor_to_major در طرح‌بندی [0, 1] باشد، بُعد 0 جزئی‌ترین بُعد و بُعد 1 اصلی‌ترین بُعد است. این مربوط به طرح‌بندی زیر در حافظه خطی است:

a d b e c f

این ترتیب ابعاد از کوچک به بزرگ از 0 تا N-1 مشابه ستون-بزرگ (برای ابعاد دوبعدی) است. با فرض ترتیب یکنواخت ابعاد، روش دیگری که می‌توانیم در کد به این طرح اشاره کنیم این است که «کم‌رنگی ۰ کوچک است».

از طرف دیگر، اگر فیلد minor_to_major در طرح‌بندی [1, 0] باشد، طرح‌بندی در حافظه خطی به صورت زیر خواهد بود:

a b c d e f

ترتیب ابعاد کوچک به بزرگ از N-1 تا 0 برای یک آرایه N بعدی مشابه row-major (برای آرایه‌های دو بعدی) است. با فرض ترتیب یکنواخت ابعاد، روش دیگری که می‌توانیم در کد به این طرح اشاره کنیم این است که "0 کوچک، بزرگ است".

ترتیب پیش‌فرض از جزئی به کلی

طرح پیش‌فرض برای شکل‌های تازه ایجاد شده «ترتیب ابعاد از بزرگ به کوچک است» (یعنی [N-1, ..., 0] ).

بالشتک

فیلد tail_padding_alignment_in_elements ترازبندی آرایه کاشی‌کاری شده را بر اساس تعداد عناصر تعریف می‌کند. پس از اعمال کاشی‌کاری، عناصر پدگذاری شده در انتهای طرح‌بندی اضافه می‌شوند تا تعداد کل عناصر مضربی از این مقدار شود.

اندیس گذاری در آرایه ها

کلاس IndexUtil در index_util.h ابزارهایی برای تبدیل بین شاخص‌های چندبعدی و شاخص‌های خطی با توجه به شکل و طرح‌بندی ارائه می‌دهد. شاخص‌های چندبعدی شامل یک شاخص int64 برای هر بعد هستند. شاخص‌های خطی یک مقدار int64 واحد هستند که در بافر نگهدارنده آرایه فهرست می‌شوند. برای ابزارهایی که ایجاد و دستکاری شکل‌ها و طرح‌بندی‌ها را ساده می‌کنند، به shape_util.h و layout_util.h در همان دایرکتوری مراجعه کنید.

شناسه‌های فضای حافظه

در HLO، هر آرایه ممکن است با یک شناسه فضای حافظه، که به صورت S(n) نوشته می‌شود، حاشیه‌نویسی شود.

  • S(0) (که اغلب حذف می‌شود) نشان‌دهنده حافظه با پهنای باند بالا (HBM) در دستگاه است.
  • S(1) نشان دهنده حافظه مجازی روی دستگاه (VMEM) است.
  • S(2) ، S(3) و غیره، متناظر با فضاهای حافظه اضافی مختص دستگاه هستند.
  • S(5) حافظه میزبان را نشان می‌دهد.