تجزیه و تحلیل نمایه سازی

این سند تجزیه و تحلیل نمایه سازی HLO را توصیف می کند، که به شما امکان می دهد به طور نمادین نقشه های نمایه سازی را برای عملیات HLO محاسبه کنید. نقشه نمایه سازی تابعی است که شاخص های یک تانسور را به شاخص های تانسور دیگر نگاشت می کند، به عنوان مثال شاخص های خروجی دستورالعمل HLO به شاخص های ورودی دستورالعمل HLO یا بالعکس.

مثال

برای پخش از tensor<20xf32> به tensor<10x20x30xf32>

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

نقشه نمایه سازی از خروجی به ورودی (i, j, k) -> (j) برای i in [0, 10] ، j in [0, 20] و k in [0, 30] است.

انگیزه

XLA GPU از چندین راه حل سفارشی برای استدلال در مورد ادغام، استفاده از عملوند و طرح های کاشی کاری استفاده می کند (جزئیات بیشتر در زیر). هدف از تجزیه و تحلیل نمایه سازی، ارائه یک جزء قابل استفاده مجدد برای چنین موارد استفاده است. تجزیه و تحلیل نمایه سازی بر اساس زیرساخت نقشه Affine MLIR ساخته شده است و معنای HLO را اضافه می کند.

ادغام

استدلال در مورد ادغام حافظه برای موارد غیر پیش پا افتاده امکان پذیر می شود، زمانی که می دانیم چه عناصر / برش هایی از ورودی ها برای محاسبه یک عنصر از خروجی خوانده می شوند.

استفاده از عملوند

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

کاشی کاری

کاشی/برش زیرمجموعه‌ای فرامستطیلی شکل از یک تانسور است که با انحراف‌ها، اندازه‌ها و گام‌ها پارامتر می‌شود. انتشار کاشی راهی برای محاسبه پارامترهای کاشی تولید کننده/مصرف کننده عملیات با استفاده از پارامترهای کاشی کاری خود op است. در حال حاضر کتابخانه ای وجود دارد که این کار را برای softmax و dot انجام می دهد. اگر از طریق نقشه های نمایه سازی بیان شود، انتشار کاشی می تواند عمومی تر و قوی تر شود.

تابع و دامنه

نقشه نمایه سازی تابع f ( d , s ) است که d چند شاخص تانسور A را به عناصر/محدوده های تانسور B ترسیم می کند. پارامتر s به محدوده شاخص های ابعادی اشاره دارد که در تانسور B وجود دارد، اما در تانسور A وجود ندارد.

به عنوان مثال، اگر کاهشی از tensor<2x4x8x16xf32> به tensor<4x8xf32> داشته باشیم، نقشه نمایه سازی از خروجی 2 بعدی به ورودی 4 بعدی به صورت (d0, d1) -> (s0, d0, d1, s1) است. d_i پارامترهای ابعادی هستند که با شاخص های تانسور خروجی مطابقت دارند. پارامترهای s_j مقادیر متعددی را کد می کنند، یعنی برای محاسبه یک عنصر (d0, d1) خروجی، به عناصر ورودی (s0, d0, d1, s1) نیاز داریم که s0 in [0, 2) و s1 in [0, 16) .

این نگاشت را می توان از ویژگی های دستورالعمل های HLO ساخت یا نگاشت دستورالعمل های ترکیب نشده را می توان برای بدست آوردن نمایه سازی برای یک ادغام ایجاد کرد. نگاشت یک دامنه نیز دارد که مشخص می کند نگاشت برای چه عناصری از تانسور وجود دارد.

f ( d , s ) st

lb _d <= d <= ub _d

lb _s <= s <= ub _s

lb _g <= g <= ub _g

از آنجایی که می خواهیم محاسبه مجدد را به حداقل برسانیم، به یک کتابخانه برای محاسبات نمادین نیاز داریم. XLA قبلاً به MLIR بستگی دارد، بنابراین ما به جای نوشتن یک کتابخانه حسابی نمادین از mlir::AffineMap استفاده می کنیم.

یک AffineMap معمولی به نظر می رسد

(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)

AffineMap به راحتی دارای دو نوع پارامتر است: ابعاد و نمادهایی که می توانیم به ترتیب برای d و s استفاده کنیم. AffineMap حاوی هیچ ابرداده ای در مورد محدوده ابعاد نیست، بنابراین ما باید خودمان این داده ها را ارائه دهیم.

struct Range {
 int64_t lower_bound;
 int64_t upper_bound;
};

struct IndexingMap {
 mlir::AffineMap affine_map;
 std::vector<Range> dim_ranges;
 std::vector<Range> symbol_ranges;
 llvm::DenseMap<mlir::AffineExpr, Range> expr_ranges;
};

dim_ranges محدودیت‌های جعبه شامل پارامترهای بعد d نقشه نمایه‌سازی را رمزگذاری می‌کند، که معمولاً با شکل تانسور خروجی برای عملیات‌هایی مانند transpose، reduce، elementwise، dot منطبق است، اما استثناهایی مانند HloConcatenateInstruction وجود دارد.

symbol_ranges مقادیر ممکنی را که پارامترهای s می توانند دریافت کنند رمزگذاری می کنند.

بیایید به صورت مثال مطالعه کنیم تا بفهمیم همه موارد بالا واقعاً به چه معنا هستند.

نمایه سازی نقشه ها برای عملیات های Unfused

از نظر عنصری

برای عملیات عنصری، نقشه نمایه سازی یک هویت است.

  p0 = f32[10, 20] parameter(0)
  p1 = f32[10, 20] parameter(1)
  add = f32[10, 20] add(p0, p1)

خروجی نقشه های ورودی:

  • خروجی -> input_i:
(d0, d1) -> (d0, d1)
domain:
d0 in [0, 19]
d1 in [0, 19]

ورودی به نقشه های خروجی

  • input_i -> خروجی:
(d0, d1) -> (d0, d1)
domain:
d0 in [0, 19]
d1 in [0, 19]

پخش

پخش به این معنی است که وقتی خروجی را به ورودی نگاشت می کنیم برخی از ابعاد حذف می شود و وقتی ورودی به خروجی را نگاشت می کنیم اضافه می شود.

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

نقشه خروجی به ورودی:

(d0, d1, d2) -> (d1)
domain:
d0 in [0, 9]
d1 in [0, 19]
d2 in [0, 29]

نقشه ورودی به خروجی

(d0)[s0, s1] -> (s0, d0, s1)
domain:
d0 in [0, 19]
s0 in [0, 9]
s1 in [0, 29]

توجه داشته باشید که اکنون s در سمت راست برای نگاشت ورودی به خروجی داریم. اینها نمادهایی هستند که محدوده ای از مقادیر را نشان می دهند. به عنوان مثال، در این مورد خاص، هر عنصر ورودی با شاخص d0 به یک برش 10x1x30 از خروجی نگاشت می شود.

ثابت و آیوتا

به راحتی، آنها هیچ پارامتر ورودی ندارند، بنابراین چیزی برای محاسبه نمایه سازی وجود ندارد.

جابجا شود

نقشه نمایه سازی برای جابجایی، جایگشتی از ابعاد ورودی/خروجی است.

p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}

نقشه خروجی به ورودی:

(d0, d1, d2, d3) -> (d0, d3, d1, d2)
domain:
d0 in [0, 2]
d1 in [0, 5]
d2 in [0, 127]
d3 in [0, 12287]

نقشه ورودی به خروجی:

(d0, d1, d2, d3) -> (d0, d2, d3, d1)
domain:
d0 in [0, 2]
d1 in [0, 12287]
d2 in [0, 5]
d3 in [0, 127]

معکوس

نمایه سازی نقشه برای معکوس، ابعاد برگردانده شده را به upper_bound(d_i) - d_i :

p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}

نقشه خروجی به ورودی:

(d0, d1, d2, d3) -> (d0, -d1 + 16, -d2 + 8, d3)
domain:
d0 in [0, 0]
d1 in [0, 16]
d2 in [0, 8]
d3 in [0, 8]

نقشه ورودی به خروجی:

(d0, d1, d2, d3) -> (d0, -d1 + 16, -d2 + 8, d3)
domain:
d0 in [0, 0]
d1 in [0, 16]
d2 in [0, 8]
d3 in [0, 8]

(Variadic)کاهش دادن

کاهش متغیر دارای چندین ورودی و چندین ورودی است، نقشه از خروجی به ورودی ابعاد کاهش یافته را اضافه می کند. بنابراین، به نوعی مانند یک معکوس به یک پخش عمل می کند.

p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
reduce = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
  dimensions={0}, to_apply=min

خروجی نقشه های ورودی:

  • خروجی -> input_j:
(d0)[s0] -> (s0, d0)
domain:
d0 in [0, 9]
s0 in [0, 255]
  • خروجی -> init_j:
(d0) -> ()
domain:
d0 in [0, 9]

ورودی نقشه های خروجی:

  • input_i -> output_j:
(d0, d1) -> (d1)
domain:
d0 in [0, 255]
d1 in [0, 9]
  • init_i -> output_j:
()[s0] -> (s0)
domain:
s0 in [0, 9]

برای i، j = 0، ... INPUT_COUNT.

تکه

نمایه سازی از خروجی به ورودی برای برش منجر به یک نقشه نمایه سازی گام به گام می شود که برای هر عنصر خروجی معتبر است. نگاشت از ورودی به خروجی به محدوده ای از عناصر موجود در ورودی محدود می شود.

p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
  slice={[5:10:1], [3:20:7], [0:50:2]}

نقشه خروجی به ورودی:

(d0, d1, d2) -> (d0 + 5, d1 * 7 + 3, d2 * 2)
domain:
d0 in [0, 4]
d1 in [0, 2]
d2 in [0, 24]

نقشه ورودی به خروجی:

TBD : نمایه سازی ورودی به خروجی

تغییر شکل دهید

تغییر شکل ها طعم های مختلفی دارند.

شکل فرو ریختن

این یک تغییر شکل "خطی" از ND به 1D است.

p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)

نقشه خروجی به ورودی:

(d0) -> (d0 floordiv 8, d0 mod 8)
domain:
d0 in [0, 31]

نقشه ورودی به خروجی:

(d0) -> (d0 floordiv 8, d0 mod 8)
domain:
d0 in [0, 31]

شکل را گسترش دهید

این یک عملیات معکوس "شکل فروپاشی" است، یک ورودی 1 بعدی را به خروجی ND تغییر شکل می دهد.

p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)

نقشه خروجی به ورودی:

(d0, d1, d2) -> (d0 floordiv 8, d0 mod 8, d1 * 4 + d2)
domain:
d0 in [0, 31]
d1 in [0, 2]
d2 in [0, 3]

نقشه ورودی به خروجی:

(d0, d1, d2) -> (d0 * 8 + d1, d2 floordiv 4, d2 mod 4)
domain:
d0 in [0, 3]
d1 in [0, 7]
d2 in [0, 11]

تغییر شکل عمومی

اینها عملیات تغییر شکلی هستند که نمی‌توان آنها را به صورت یک شکل منبسط یا فروپاشی نشان داد. آنها را فقط می توان به عنوان ترکیبی از 2 یا بیشتر شکل گسترش یا فروپاشی نشان داد.

مثال 1: خطی سازی-خطی سازی.
p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)

این تغییر شکل را می توان به صورت ترکیبی از شکل فروپاشی tensor<4x8xf32> به tensor<32xf32> و سپس گسترش شکل به tensor<2x4x4xf32> نشان داد.

نقشه خروجی به ورودی:

(d0, d1, d2) -> (d0 * 2 + d1 floordiv 2, d2 + (d1 mod 2) * 4) 
domain:
d0 in [0, 1]
d1 in [0, 3]
d2 in [0, 3]

نقشه ورودی به خروجی:

(d0, d1) -> (d0 floordiv 2, d1 floordiv 4 + (d0 mod 2) * 2, d1 mod 4)
domain:
d0 in [0, 3]
d1 in [0, 7]
مثال 2: شکل های فرعی منبسط و جمع شده
p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)

این تغییر شکل را می توان به عنوان ترکیبی از دو تغییر شکل نشان داد. اولی بیرونی ترین tensor<4x8x12xf32> را به tensor<32x12xf32> می ریزد و دومی داخلی ترین tensor<32x12xf32> را به tensor<32x3x4xf32> گسترش می دهد.

نقشه خروجی به ورودی:

(d0, d1, d2) -> (d0 floordiv 8, d0 mod 8, d1 * 4 + d2)
domain:
d0 in [0, 31]
d1 in [0, 2]
d2 in [0, 3]

نقشه ورودی به خروجی:

(d0, d1, d2) -> (d0 * 8 + d1, d2 floordiv 4, d2 mod 4)
domain:
d0 in [0, 3]
d1 in [0, 7]
d2 in [0, 11]

بیت کست

یک عملیات بیت‌کست را می‌توان به صورت دنباله‌ای از transpose-reshape-transpose نشان داد. بنابراین، نقشه های نمایه سازی آن فقط ترکیبی از نقشه های نمایه سازی برای این دنباله است.

الحاق

نگاشت خروجی به ورودی برای concat برای همه ورودی ها تعریف شده است، اما با دامنه های غیر همپوشانی، یعنی تنها یکی از ورودی ها در یک زمان استفاده می شود.

p0 = f32[2, 5, 7] parameter(0)
p1 = f32[2, 11, 7] parameter(1)
p2 = f32[2, 17, 7] parameter(2)
ROOT concat = f32[2, 33, 7] concatenate(f32[2, 5, 7] p0, f32[2, 11, 7] p1, f32[2, 17, 7] p2), dimensions={1}

خروجی نقشه های ورودی:

  • خروجی -> ورودی 1:
(d0, d1, d2) -> (d0, d1, d2)
domain:
d0 in [0, 1]
d1 in [0, 4]
d2 in [0, 6]
  • خروجی -> ورودی 2:
(d0, d1, d2) -> (d0, d1 - 5, d2)
domain:
d0 in [0, 1]
d1 in [5, 15]
d2 in [0, 6]
  • خروجی -> ورودی 3:
(d0, d1, d2) -> (d0, d1 - 16, d2)
domain:
d0 in [0, 1]
d1 in [16, 32]
d2 in [0, 6]

ورودی های نقشه های خروجی:

  • ورودی 1 -> خروجی:
(d0, d1, d2) -> (d0, d1, d2)
domain:
d0 in [0, 1]
d1 in [0, 4]
d2 in [0, 6]
  • ورودی 2 -> خروجی:
(d0, d1, d2) -> (d0, d1 + 5, d2)
domain:
d0 in [0, 1]
d1 in [0, 10]
d2 in [0, 6]
  • ورودی 3 -> خروجی:
(d0, d1, d2) -> (d0, d1 + 16, d2)
domain:
d0 in [0, 1]
d1 in [0, 16]
d2 in [0, 6]

نقطه

نقشه های نمایه سازی برای نقطه بسیار شبیه نقشه های کاهش است.

p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
dot = f32[4, 128, 64] dot(p0, p1),
  lhs_batch_dims={0}, rhs_batch_dims={0},
  lhs_contracting_dims={2}, rhs_contracting_dims={1}

خروجی نقشه های ورودی:

  • خروجی -> ورودی_1:
(d0, d1, d2)[s0] -> (d0, d1, s0)
domain:
d0 in [0, 3]
d1 in [0, 127]
d2 in [0, 63]
s0 in [0, 255]
  • خروجی -> ورودی_2:
(d0, d1, d2)[s0] -> (d0, s0, d2)
domain:
d0 in [0, 3]
d1 in [0, 127]
d2 in [0, 63]
s0 in [0, 255]

ورودی های نقشه های خروجی:

  • input_1 -> خروجی:
(d0, d1, d2) -> (d0, d1, s0)
domain:
d0 in [0, 3]
d1 in [0, 127]
d2 in [0, 255]
s0 in [0, 63]
  • input_2 -> خروجی:
(d0, d1, d2) -> (d0, s_0, d1)
domain:
d0 in [0, 3]
d1 in [0, 255]
d2 in [0, 63]
s0 in [0, 127]

پد

نمایه سازی PadOp برعکس نمایه سازی SliceOp است.

p0 = f32[4, 4] parameter(0)
p1 = f32[] parameter(1)
pad = f32[12, 16] pad(p0, p1), padding=1_4_1x4_8_0

پیکربندی padding 1_4_1x4_8_0 نشان دهنده lowPad_highPad_interiorPad_dim_0 x lowPad_highPad_interiorPad_dim_1 است.

خروجی نقشه های ورودی:

  • خروجی -> ورودی:
(d0, d1) -> ((d0 - 1) floordiv 2, d1 - 4)
domain:
d0 in [1, 7]
d1 in [4, 7]
(d0 - 1) mod 2 in [0, 0]
  • خروجی -> init:
(d0, d1) -> ()
domain:
d0 in [0, 11]
d1 in [0, 15]

ReduceWindow

ReduceWindow در XLA نیز padding را انجام می دهد. بنابراین، نقشه های نمایه سازی را می توان به عنوان ترکیبی از نمایه سازی ReduceWindow محاسبه کرد که هیچ گونه padding و نمایه سازی PadOp را انجام نمی دهد.

c_inf = f32[] constant(-inf)
p0 = f32[1024, 514] parameter(0)
reduce-window = f32[1024, 3] reduce-window(p0, c_inf),
  window={size=1x512 pad=0_0x0_0}, to_apply=max

خروجی نقشه های ورودی:

  • خروجی -> ورودی:
(d0, d1)[s0] -> (d0, d1 + s0)
domain:
d0 in [0, 1023]
d1 in [0, 2]
s0 in [0, 511]
  • خروجی -> init:
(d0, d1) -> ()
domain:
d0 in [0, 1023]
d1 in [0, 2]

نمایه سازی نقشه ها برای فیوژن

نقشه نمایه سازی برای fusion op ترکیبی از نقشه های نمایه سازی برای هر عملیات در خوشه است. ممکن است برخی از ورودی ها چندین بار با الگوهای دسترسی متفاوت خوانده شوند.

یک ورودی، چندین نقشه نمایه سازی

در اینجا یک مثال برای p0 + transpose(p0) آورده شده است.

f {
  p0 = f32[1000, 1000] parameter(0)
  transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0}
  ROOT a0 = f32[1000, 1000] add(p0, transpose_p0)
}

نقشه های نمایه سازی خروجی به ورودی برای p0 (d0, d1) -> (d0, d1) و (d0, d1) -> (d1, d0) خواهد بود. به این معنی که برای محاسبه یک عنصر از خروجی ممکن است لازم باشد پارامتر ورودی را دو بار بخوانیم.

یک ورودی، نقشه نمایه سازی کپی شده

img

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

f {
  p0 = f32[20, 10, 50] parameter(0)
  lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2}
  lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1)
  lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1}
  rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0}
  rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1)
  rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2}
  ROOT add = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}

نقشه نمایه سازی خروجی به ورودی برای p0 در این مورد فقط (d0, d1, d2) -> (d2, d0, d1) است.

سافت مکس

img

نقشه های نمایه سازی خروجی به ورودی برای parameter 0 برای softmax:

(d0, d1, d2)[s0] -> (d0, d1, s0)
domain:
d0 in [0, 1]
d1 in [0, 64]
d2 in [0, 124]
s0 in [0, 124]

و

(d0, d1, d2) -> (d0, d1, d2)
domain:
d0 in [0, 1]
d1 in [0, 64]
d2 in [0, 124]

که در آن s_0 به درونی ترین بعد ورودی اشاره دارد.

نمایه سازی ساده کننده نقشه

ساده‌کننده پیش‌فرض mlir::AffineMap upstream نمی‌تواند هیچ فرضی در مورد محدوده ابعاد/نمادها داشته باشد. بنابراین، نمی تواند عبارات را با mod و div به طور موثر ساده کند.

ما می‌توانیم از دانش در مورد کران‌های پایین و بالای عبارات فرعی در نقشه‌های وابسته استفاده کنیم تا آنها را حتی بیشتر ساده کنیم.

ساده کننده می تواند عبارات زیر را بازنویسی کند.

  1. (d0, d1) -> (d0 + d1 floordiv 16, d1 mod 16) برای d در [0, 6] x [0, 14] می شود (d0, d1) -> (d0, d1)
  2. (d0, d1, d2) -> ((100d0 + 10d1 + d2) floorDiv 100, ((100d0 + 10d1 + d2) mod 100) floordiv 10, d2 mod 10) برای di in [0, 9] تبدیل به (d0, d1, d2) -> (d0, d1, d2) .
  3. (d0, d1, d2) -> ((16d0 + 4d1 + d2) floordiv 8, (16d0 + 4d1 + d2) mod 8) برای d_i in [0, 9] می شود (d0, d1, d2) -> (2d0 + (4d1 + d2) floordiv 8,(4d1 + d2) mod 8) .
  4. (d0, d1) -> (-(-11d0 - d1 + 109) floordiv 11 + 9) برای d در [0, 9] x [0, 10] می شود (d0, d1) -> (d0) .

ساده‌ساز نقشه نمایه‌سازی به ما امکان می‌دهد بفهمیم که برخی از شکل‌های زنجیره‌ای در HLO یکدیگر را خنثی می‌کنند.

p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)

پس از ترکیب نقشه های نمایه سازی و ساده سازی آنها به دست خواهیم آورد

(d0, d1, d2) -> (d0, d1, d2) .

ساده سازی نقشه نمایه سازی محدودیت ها را نیز ساده می کند.

  1. محدودیت های نوع lower_bound <= affine_expr (floordiv, +, -, *) constant <= upper_bound به صورت updated_lower_bound <= affine_expr <= updated_upped_bound بازنویسی می شوند.
  2. محدودیت هایی که همیشه ارضا می شوند، به عنوان مثال d0 + s0 in [0, 20] برای d0 in [0, 5] و s0 in [1, 3] حذف می شوند.
  3. عبارات Affine در محدودیت ها به عنوان نقشه affine نمایه سازی در بالا بهینه شده اند.