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

این سند تجزیه و تحلیل نمایه سازی 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 ( x ) = f ( d , r , rt ) است که d چند شاخصی از یک تانسور A را به عناصر/محدوده های تانسور B ترسیم می کند. پارامتر r به محدوده شاخص‌های ابعادی اشاره دارد که در تانسور B وجود دارد، اما در تانسور A وجود ندارد. پارامتر rt به مقادیر زمان اجرا اشاره دارد، به عنوان مثال شاخص‌های یک عملیات جمع‌آوری.

برای مثال، اگر کاهشی از tensor<2x4x8x16xf32> به tensor<4x8xf32> داشته باشیم، نقشه نمایه سازی از خروجی 2 بعدی به ورودی 4 بعدی به صورت (d0, d1) -> (r0, d0, d1, r1) است. d_i متغیرهای بعدی هستند که با شاخص های تانسور خروجی مطابقت دارند. متغیرهای محدوده r_j چندین مقدار را رمزگذاری می کنند، یعنی برای محاسبه یک عنصر (d0, d1) خروجی، به عناصر ورودی (r0, d0, d1, r1) نیاز داریم که r0 in [0, 1] و r1 in [0, 15] .

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

f ( x ) خیابان

پوند <= g ( x ) <= ub

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

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

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

AffineMap دو نوع پارامتر دارد: ابعاد و نمادها . ابعاد مربوط به متغیرهای بعد d ، نمادها مربوط به متغیرهای محدوده r و متغیرهای RT rt هستند. AffineMap حاوی هیچ ابرداده ای در مورد محدوده ابعاد نیست، بنابراین ما باید خودمان این داده ها را ارائه دهیم.

struct Interval {
 int64_t lower;
 int64_t upper;
};

class IndexingMap {
   // Variable represents dimension, range or runtime variable.
  struct Variable {
    Interval bounds;
    // Name of the variable is used for nicer printing.
    std::string name = "";
  };

  mlir::AffineMap affine_map_;

  // DimVars represent dimensions of a tensor or of a GPU grid.
  std::vector<Variable> dim_vars_;

  // RangeVars represent ranges of values, e.g. to compute a single element of
  // the reduction's result we need a range of values from the input tensor.
  std::vector<Variable> range_vars_;

  // RTVars represent runtime values, e.g. a dynamic offset in
  // HLO dynamic-update-slice op.
  std::vector<Variable> rt_vars_;
  llvm::DenseMap<mlir::AffineExpr, Interval> constraints_;
};

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

range_vars_ مقادیر ممکنی را که پارامترهای r می توانند بگیرند رمزگذاری می کند.

rt_vars_ مقادیر قابل اجرا را در زمان اجرا رمزگذاری می کند. به عنوان مثال، افست برای یک دستورالعمل 1 بعدی HloDynamicSliceInstruction پویا است. RTVar مربوطه مقادیر قابل اجرا بین 0 و tensor_size - slice_size - 1 خواهد داشت.

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

نمایه سازی نقشه ها برای عملیات های 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, 9],
d1 in [0, 19]

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

  • input_i -> خروجی:
(d0, d1) -> (d0, d1),
domain:
d0 in [0, 9],
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 از خروجی نگاشت می شود.

ثابت و آیوتا

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

DynamicSlice

DynamicSlice درست مانند Slice است، اما افست ها پویا هستند.

src = s32[2,2,258] parameter(0)
of1 = s32[] parameter(1)
of2 = s32[] parameter(2)
of3 = s32[] parameter(3)
ds = dynamic-slice(s32[2,2,258] src, s32[] of1, s32[] of2, s32[] of3), dynamic_slice_sizes={1, 2, 32}

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

(d0, d1, d2){rt0, rt1, rt2} -> (d0 + rt0, d1 + rt1, d2 + rt2),
domain:
d0 in [0, 0],
d1 in [0, 1],
d2 in [0, 31],
rt0 in [0, 1],
rt1 in [0, 0],
rt2 in [0, 226]

توجه داشته باشید که اکنون rt را در سمت راست برای نگاشت ورودی به خروجی داریم. اینها نمادهایی هستند که مقادیر زمان اجرا را نشان می دهند. به عنوان مثال، در این مورد خاص برای هر عنصر خروجی با شاخص‌های d0, d1, d2 ما به انحراف‌های برش of1 ، of2 و of3 برای محاسبه شاخص ورودی دسترسی داریم. فواصل متغیرهای زمان اجرا با این فرض به دست می‌آیند که کل برش در محدوده باقی می‌ماند.

خروجی نقشه ورودی برای of1 ، of2 و of3 :

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

DynamicUpdateSlice

src = s32[20,30] parameter(0)
upd = s32[5,10] parameter(1)
of1 = s32[] parameter(2)
of2 = s32[] parameter(3)
dus = s32[20,30] dynamic-update-slice(
    s32[20,30] src, s32[5,10] upd, s32[] of1, s32[] of2)

خروجی نقشه ورودی برای src بی اهمیت است. می توان آن را با محدود کردن دامنه به شاخص های به روز نشده دقیق تر کرد، اما در حال حاضر نقشه های نمایه سازی از محدودیت های نابرابری پشتیبانی نمی کنند.

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

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

(d0, d1){rt0, rt1}  -> (d0 - rt0, d1 - rt1),
domain:
d0 in [0, 19],
d1 in [0, 29],
rt0 in [0, 15],
rt1 in [0, 20]

توجه داشته باشید که اکنون s در سمت راست برای نگاشت ورودی به خروجی داریم. اینها نمادهایی هستند که مقادیر زمان اجرا را نشان می دهند. به عنوان مثال، در این مورد خاص برای هر عنصر خروجی با شاخص‌های d0, d1 ما به انحراف‌های برش of1 و of2 برای محاسبه شاخص ورودی دسترسی داریم. فواصل متغیرهای زمان اجرا با این فرض به دست می‌آیند که کل برش در محدوده باقی می‌ماند.

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

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

جمع کنید

فقط جمع ساده شده پشتیبانی می شود. [gather_simplifier] را ببینید.( https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/gather_simplifier.h ).

operand = f32[33,76,70] parameter(0)
indices = s32[1806,2] parameter(1)
gather = f32[1806,7,8,4] gather(operand, indices),
  offset_dims={1,2,3},
  collapsed_slice_dims={},
  start_index_map={0,1},
  index_vector_dim=1,
  slice_sizes={7,8,4}

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

(d0, d1, d2, d3){rt0, rt1} -> (d1 + rt0, d2 + rt1, d3),
domain:
d0 in [0, 1805],
d1 in [0, 6],
d2 in [0, 7],
d3 in [0, 3],
rt0 in [0, 26],
rt1 in [0, 68]

توجه داشته باشید که اکنون s در سمت راست برای نگاشت ورودی به خروجی داریم. اینها نمادهایی هستند که مقادیر زمان اجرا را نشان می دهند. به عنوان مثال، در این مورد خاص برای هر عنصر خروجی با شاخص‌های d0, d1, d2, d3 عناصر (d0، 0) و (d0، 1) را از تانسور indices استخراج می‌کنیم.

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

(d0, d1, d2, d3)[s0] -> (d0, s0),
domain:
d0 in [0, 1805],
d1 in [0, 6],
d2 in [0, 7],
d3 in [0, 3],
s0 in [0, 1]

متغیر محدوده s0 نشان می دهد که برای محاسبه یک عنصر از خروجی به کل ردیف (d0, *) تانسور indices نیاز داریم.

جابجا شود

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

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=max

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

  • خروجی -> 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]

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

(d0, d1, d2) -> (d0 - 5, (d1 - 3) floordiv 7, d2 floordiv 2),
domain:
d0 in [5, 9],
d1 in [3, 17],
d2 in [0, 48],
(d1 - 3) mod 7 in [0, 0],
d2 mod 2 in [0, 0]

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

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

شکل فرو ریختن

این یک تغییر شکل "خطی" از 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, d1) -> (d0 * 8 + d1),
domain:
d0 in [0, 3],
d1 in [0, 7]

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

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

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

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

(d0, d1) -> (d0 * 8 + d1),
domain:
d0 in [0, 3],
d1 in [0, 7]

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

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

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

اینها عملیات تغییر شکلی هستند که نمی‌توان آنها را به صورت یک شکل منبسط یا فروپاشی نشان داد. آنها را فقط می توان به عنوان ترکیبی از 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)[s0] -> (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)[s0] -> (d0, s0, 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]

جایی که s0 به درونی ترین بعد ورودی اشاره دارد.

برای مثال‌های بیشتر به indexing_analysis_test.cc مراجعه کنید.

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

ساده‌کننده پیش‌فرض 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 نمایه سازی در بالا بهینه شده اند.

برای مثال‌های بیشتر به indexing_map_test.cc مراجعه کنید.

،

این سند تجزیه و تحلیل نمایه سازی 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 ( x ) = f ( d , r , rt ) است که d چند شاخصی از یک تانسور A را به عناصر/محدوده های تانسور B ترسیم می کند. پارامتر r به محدوده شاخص‌های ابعادی اشاره دارد که در تانسور B وجود دارد، اما در تانسور A وجود ندارد. پارامتر rt به مقادیر زمان اجرا اشاره دارد، به عنوان مثال شاخص‌های یک عملیات جمع‌آوری.

برای مثال، اگر کاهشی از tensor<2x4x8x16xf32> به tensor<4x8xf32> داشته باشیم، نقشه نمایه سازی از خروجی 2 بعدی به ورودی 4 بعدی به صورت (d0, d1) -> (r0, d0, d1, r1) است. d_i متغیرهای بعدی هستند که با شاخص های تانسور خروجی مطابقت دارند. متغیرهای محدوده r_j چندین مقدار را رمزگذاری می کنند، یعنی برای محاسبه یک عنصر (d0, d1) خروجی، به عناصر ورودی (r0, d0, d1, r1) نیاز داریم که r0 in [0, 1] و r1 in [0, 15] .

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

f ( x ) خیابان

پوند <= g ( x ) <= ub

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

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

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

AffineMap دو نوع پارامتر دارد: ابعاد و نمادها . ابعاد مربوط به متغیرهای بعد d ، نمادها مربوط به متغیرهای محدوده r و متغیرهای RT rt هستند. AffineMap حاوی هیچ ابرداده ای در مورد محدوده ابعاد نیست، بنابراین ما باید خودمان این داده ها را ارائه دهیم.

struct Interval {
 int64_t lower;
 int64_t upper;
};

class IndexingMap {
   // Variable represents dimension, range or runtime variable.
  struct Variable {
    Interval bounds;
    // Name of the variable is used for nicer printing.
    std::string name = "";
  };

  mlir::AffineMap affine_map_;

  // DimVars represent dimensions of a tensor or of a GPU grid.
  std::vector<Variable> dim_vars_;

  // RangeVars represent ranges of values, e.g. to compute a single element of
  // the reduction's result we need a range of values from the input tensor.
  std::vector<Variable> range_vars_;

  // RTVars represent runtime values, e.g. a dynamic offset in
  // HLO dynamic-update-slice op.
  std::vector<Variable> rt_vars_;
  llvm::DenseMap<mlir::AffineExpr, Interval> constraints_;
};

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

range_vars_ مقادیر ممکنی را که پارامترهای r می توانند بگیرند رمزگذاری می کند.

rt_vars_ مقادیر قابل اجرا را در زمان اجرا رمزگذاری می کند. به عنوان مثال، افست برای یک دستورالعمل 1 بعدی HloDynamicSliceInstruction پویا است. RTVar مربوطه مقادیر قابل اجرا بین 0 و tensor_size - slice_size - 1 خواهد داشت.

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

نمایه سازی نقشه ها برای عملیات های 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, 9],
d1 in [0, 19]

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

  • input_i -> خروجی:
(d0, d1) -> (d0, d1),
domain:
d0 in [0, 9],
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 از خروجی نگاشت می شود.

ثابت و آیوتا

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

DynamicSlice

DynamicSlice درست مانند Slice است، اما افست ها پویا هستند.

src = s32[2,2,258] parameter(0)
of1 = s32[] parameter(1)
of2 = s32[] parameter(2)
of3 = s32[] parameter(3)
ds = dynamic-slice(s32[2,2,258] src, s32[] of1, s32[] of2, s32[] of3), dynamic_slice_sizes={1, 2, 32}

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

(d0, d1, d2){rt0, rt1, rt2} -> (d0 + rt0, d1 + rt1, d2 + rt2),
domain:
d0 in [0, 0],
d1 in [0, 1],
d2 in [0, 31],
rt0 in [0, 1],
rt1 in [0, 0],
rt2 in [0, 226]

توجه داشته باشید که اکنون rt را در سمت راست برای نگاشت ورودی به خروجی داریم. اینها نمادهایی هستند که مقادیر زمان اجرا را نشان می دهند. به عنوان مثال، در این مورد خاص برای هر عنصر خروجی با شاخص‌های d0, d1, d2 ما به انحراف‌های برش of1 ، of2 و of3 برای محاسبه شاخص ورودی دسترسی داریم. فواصل متغیرهای زمان اجرا با این فرض به دست می‌آیند که کل برش در محدوده باقی می‌ماند.

خروجی نقشه ورودی برای of1 ، of2 و of3 :

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

DynamicUpdateSlice

src = s32[20,30] parameter(0)
upd = s32[5,10] parameter(1)
of1 = s32[] parameter(2)
of2 = s32[] parameter(3)
dus = s32[20,30] dynamic-update-slice(
    s32[20,30] src, s32[5,10] upd, s32[] of1, s32[] of2)

خروجی نقشه ورودی برای src بی اهمیت است. می توان آن را با محدود کردن دامنه به شاخص های به روز نشده دقیق تر کرد، اما در حال حاضر نقشه های نمایه سازی از محدودیت های نابرابری پشتیبانی نمی کنند.

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

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

(d0, d1){rt0, rt1}  -> (d0 - rt0, d1 - rt1),
domain:
d0 in [0, 19],
d1 in [0, 29],
rt0 in [0, 15],
rt1 in [0, 20]

توجه داشته باشید که اکنون s در سمت راست برای نگاشت ورودی به خروجی داریم. اینها نمادهایی هستند که مقادیر زمان اجرا را نشان می دهند. به عنوان مثال، در این مورد خاص برای هر عنصر خروجی با شاخص‌های d0, d1 ما به انحراف‌های برش of1 و of2 برای محاسبه شاخص ورودی دسترسی داریم. فواصل متغیرهای زمان اجرا با این فرض به دست می‌آیند که کل برش در محدوده باقی می‌ماند.

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

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

جمع کنید

فقط جمع ساده شده پشتیبانی می شود. [gather_simplifier] را ببینید.( https://github.com/openxla/xla/blob/main/xla/hlo/transforms/simplifiers/gather_simplifier.h ).

operand = f32[33,76,70] parameter(0)
indices = s32[1806,2] parameter(1)
gather = f32[1806,7,8,4] gather(operand, indices),
  offset_dims={1,2,3},
  collapsed_slice_dims={},
  start_index_map={0,1},
  index_vector_dim=1,
  slice_sizes={7,8,4}

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

(d0, d1, d2, d3){rt0, rt1} -> (d1 + rt0, d2 + rt1, d3),
domain:
d0 in [0, 1805],
d1 in [0, 6],
d2 in [0, 7],
d3 in [0, 3],
rt0 in [0, 26],
rt1 in [0, 68]

توجه داشته باشید که اکنون s در سمت راست برای نگاشت ورودی به خروجی داریم. اینها نمادهایی هستند که مقادیر زمان اجرا را نشان می دهند. به عنوان مثال، در این مورد خاص برای هر عنصر خروجی با شاخص‌های d0, d1, d2, d3 عناصر (d0، 0) و (d0، 1) را از تانسور indices استخراج می‌کنیم.

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

(d0, d1, d2, d3)[s0] -> (d0, s0),
domain:
d0 in [0, 1805],
d1 in [0, 6],
d2 in [0, 7],
d3 in [0, 3],
s0 in [0, 1]

متغیر محدوده s0 نشان می دهد که برای محاسبه یک عنصر از خروجی به کل ردیف (d0, *) تانسور indices نیاز داریم.

جابجا شود

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

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=max

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

  • خروجی -> 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]

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

(d0, d1, d2) -> (d0 - 5, (d1 - 3) floordiv 7, d2 floordiv 2),
domain:
d0 in [5, 9],
d1 in [3, 17],
d2 in [0, 48],
(d1 - 3) mod 7 in [0, 0],
d2 mod 2 in [0, 0]

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

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

شکل فرو ریختن

این یک تغییر شکل "خطی" از 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, d1) -> (d0 * 8 + d1),
domain:
d0 in [0, 3],
d1 in [0, 7]

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

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

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

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

(d0, d1) -> (d0 * 8 + d1),
domain:
d0 in [0, 3],
d1 in [0, 7]

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

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

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

اینها عملیات تغییر شکلی هستند که نمی‌توان آنها را به صورت یک شکل منبسط یا فروپاشی نشان داد. آنها را فقط می توان به عنوان ترکیبی از 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)[s0] -> (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)[s0] -> (d0, s0, 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]

جایی که s0 به درونی ترین بعد ورودی اشاره دارد.

برای مثال‌های بیشتر به indexing_analysis_test.cc مراجعه کنید.

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

ساده‌کننده پیش‌فرض 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 نمایه سازی در بالا بهینه شده اند.

برای مثال‌های بیشتر به indexing_map_test.cc مراجعه کنید.