'sdy' گویش

گویش شاردی (SDY).

گویش Shardy (SDY) یک نمایش تقسیم بندی تانسور مبتنی بر محور و اجزای API اضافی را برای اتصال شاردنگ ها به تانسورها تعریف می کند.

گزارش نسخه: 0.0.1: محورهای کاهش نیافته را به TensorShardingAttr اضافه کنید.

عملیات

sdy.all_gather (sdy::AllGatherOp)

یک ارتباط همه جمعی را در امتداد محورها انجام می دهد

نحو:

operation ::= `sdy.all_gather` $gathering_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

تکه های یک تانسور را در امتداد محورهای مشخص شده در gathering_axes جمع می کند.

gathering_axes فهرستی از فهرست‌های محورها است. فهرست بیرونی بیش از ابعاد تانسور است. هر فهرست داخلی محورهایی را مشخص می کند که در امتداد آنها جمع آوری جداگانه در بعد مربوطه باید انجام شود. برای به دست آوردن تقسیم بندی نتیجه ( out_sharding ) روی تقسیم بندی عملوند ( tensor ) اعمال خواهد شد.

توجه داشته باشید که out_sharding برای تعیین تقسیم بندی نتیجه استفاده نمی شود. در عوض، تقسیم‌بندی نتیجه با تقسیم عملوند و gathering_axes تعیین می‌شود، و out_sharding باید با این تقسیم‌بندی استنباط‌شده مطابقت داشته باشد.

مثال:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_gather [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\]> : tensor<8x8x8xf32>

محدودیت ها:

• باید محدودیت های فهرست شده در Sdy_CollectiveOpInterface را برآورده کند.
• عناصر در gathering_axes باید محدودیت های فهرست شده در AxisRefListAttr را برآورده کنند.
• اعمال gathering_axes به عملوند به اشتراک گذاری out_sharding می رسد.

صفات: SameOperandsAndResultType

رابط ها: InferTypeOpInterface ، Sdy_CollectiveOpInterface ، SymbolUserOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.all_reduce (sdy::AllReduceOp)

ارتباط همه جانبه را در امتداد محورها انجام دهید

نحو:

operation ::= `sdy.all_reduce` $reduction_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

تکه های یک تانسور را در امتداد محورهای مشخص شده در reduction_axes کاهش می دهد. ترتیب reduction_axes برای نتیجه مهم نیست، اما می تواند ترتیب گروه های مشابه مربوطه را تحت تاثیر قرار دهد.

محدودیت ها:

• باید محدودیت های فهرست شده در Sdy_CollectiveOpInterface را برآورده کند.
• reduction_axes باید محدودیت های فهرست شده در AxisRefListAttr را برآورده کند.
• عملوند sharding و out_sharding باید دارای ابعاد مشابه باشند.
• reduction_axes نباید با محورهای تقسیم‌بندی بعد عملوند و محورهای تکرار شده همپوشانی داشته باشد (می‌تواند با محورهای کاهش‌یافته همپوشانی داشته باشد).
• reduction_axes نباید با محورهای تقلیل نشده out_sharding همپوشانی داشته باشد. به عبارت دیگر، out_sharding باید در امتداد reduction_axes (به طور ضمنی یا صریح) تکرار شود.

صفات: SameOperandsAndResultType

رابط ها: CollectiveOpInterface ، InferTypeOpInterface ، SymbolUserOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.all_slice (sdy::AllSliceOp)

عملیات برش پویا را در امتداد محورها انجام می دهد

نحو:

operation ::= `sdy.all_slice` $slicing_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

تکه های یک تانسور را در امتداد محورهای مشخص شده در slicing_axes برش می دهد. یک دوگانگی جبری بین sdy.all_slice و sdy.all_gather وجود دارد.

slicing_axes لیستی از لیست‌های محورها است. فهرست بیرونی بیش از ابعاد تانسور است. هر فهرست داخلی محورهایی را مشخص می کند که در امتداد آنها یک برش باید در بعد مربوطه انجام شود. برای به دست آوردن تقسیم بندی نتیجه ( out_sharding ) روی تقسیم بندی عملوند ( tensor ) اعمال خواهد شد.

توجه داشته باشید که out_sharding برای تعیین تقسیم بندی نتیجه استفاده نمی شود. در عوض، تقسیم‌بندی نتیجه توسط تقسیم‌بندی عملوند و slicing_axes تعیین می‌شود، و out_sharding باید با این تقسیم‌بندی استنباط‌شده مطابقت داشته باشد.

مثال:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}, {}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_slice [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a", "b", "c"}, {}, {"d"}\]> : tensor<8x8x8xf32>

محدودیت ها:

• باید محدودیت های فهرست شده در Sdy_CollectiveOpInterface را برآورده کند.
• عناصر موجود در slicing_axes باید محدودیت های فهرست شده در AxisRefListAttr را برآورده کنند.
• اعمال slicing_axes به عملوند sharding out_sharding می شود.

صفات: SameOperandsAndResultType

رابط ها: CollectiveOpInterface ، InferTypeOpInterface ، SymbolUserOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.all_to_all (sdy::AllToAllOp)

ارتباط همه جانبه را در امتداد محورها انجام می دهد

نحو:

operation ::= `sdy.all_to_all` $params $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

برای هر تاپل (محور، src_dim، tgt_dim) در لیست پارامترها، این عملیات تکه‌های یک تانسور را در امتداد بعد tgt_dim و محورهای مشخص‌شده در axes برش می‌دهد، آن تکه‌ها را در امتداد محورها پراکنده می‌کند و آنها را در امتداد بعد src_dim به هم متصل می‌کند.

این عملیات اساساً ترکیبی از یک all-gather در امتداد src_dim و axes و به دنبال آن یک قطعه تمام در امتداد tgt_dim و axes است، یعنی پسوندی از بعد تقسیم محور src_dim در تانسور ورودی به بعد تقسیم محور tgt_dim در خروجی اضافه می‌شود.

all-to-all برای به اشتراک گذاری عملوند ( tensor ) برای به دست آوردن تقسیم بندی نتیجه ( out_sharding ) اعمال می شود.

توجه داشته باشید که out_sharding برای تعیین تقسیم بندی نتیجه استفاده نمی شود. در عوض، تقسیم‌بندی نتیجه با تقسیم‌بندی عملوند، src_dim ، tgt_dim ، و axes تعیین می‌شود، و out_sharding باید با این تقسیم‌بندی استنباط‌شده مطابقت داشته باشد.

مثال:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b"}, {"c"}, {}, {}\]>]>} : tensor<8x8x4x4x32>
%2 = sdy.all_to_all [{"b"}: 0->2, {"c"}: 1->3] %1 out_sharding=<@mesh, [{"a"}, {}, {"b"}, {"c"}\]> : tensor<8x8x4x4x32>

محدودیت ها:

• باید محدودیت های فهرست شده در Sdy_CollectiveOpInterface را برآورده کند.
• لیست پارامترها نباید خالی باشد.
• برای هر پارامتر در params :
  • عناصر موجود در axes باید محدودیت های AxisRefAttr را برآورده کنند.
  • src_dim و tgt_dim باید ابعاد معتبر (غیر منفی و کمتر از رتبه تانسور) باشند.
  • هر src_dim یا tgt_dim باید در تمام پارامترها منحصر به فرد باشد.
  • src_dim باید در همه پارامترها به ترتیب صعودی مرتب شود.
• جابجایی axes از src_dim به tgt_dim در عملوند sharding out_sharding می شود.

صفات: SameOperandsAndResultType

رابط ها: InferTypeOpInterface ، Sdy_CollectiveOpInterface ، SymbolUserOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.collective_permute (sdy::CollectivePermuteOp)

یک ارتباط جمعی-جایگزین برای جایگزینی محورها انجام می دهد

نحو:

operation ::= `sdy.collective_permute` $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

تکه‌ای از تانسور ورودی را از هر دستگاه به دستگاه دیگر می‌فرستد تا محورهایی را که تانسور را خرد می‌کنند، مجدداً مرتب یا جایگزین کند.

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

این برای مرتب کردن مجدد محورها در یک بعد یا در ابعاد مختلف و تعویض محورهای خرد شده با محورهای تکراری مفید است.

در مثال زیر، اندازه تانسور خرد شده tensor<1x4x2xf32> است، و توسط جایگشت جمعی حفظ می‌شود.

مثال:

sdy.mesh @mesh = <["a"=2, "b"=2, "c"=4, "d"=2, "e"=2, "f"=2]>
%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "c"}, {"f"}, {"d", "e"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.collective_permute %1 out_sharding=<@mesh, [{"c":(1)2, "b", "f"}, {"a"}, {"e", "d"}\]> : tensor<8x8x8xf32>

محدودیت ها:

• باید محدودیت های فهرست شده در Sdy_CollectiveOpInterface را برآورده کند.
• اگر اشتراک گذاری ورودی و خروجی مش های متفاوتی داشته باشند، آن مش ها باید دقیقاً محورهای یکسان و ترتیب متفاوتی از شناسه های دستگاه داشته باشند.
• برای هر بعد، حاصل ضرب اندازه‌های محور تقسیم در out_sharding باید با تقسیم‌بندی بعد عملوند مربوطه مطابقت داشته باشد.

صفات: SameOperandsAndResultType

رابط ها: CollectiveOpInterface ، InferTypeOpInterface ، SymbolUserOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.constant (sdy::ConstantOp)

عملیات ثابت

یک تانسور output از یک value ثابت تولید می کند.

ببینید: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

مثال:

%output = sdy.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>

صفات: AlwaysSpeculatableImplTrait

رابط ها: ConditionallySpeculatable ، InferTypeOpInterface ، NoMemoryEffect (MemoryEffectOpInterface)

جلوه ها: MemoryEffects::Effect{}

ویژگی ها:

نتایج:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

عملیات لبه جریان داده

نحو:

operation ::= `sdy.data_flow_edge` $input (`sharding````=``` $sharding^)? attr-dict `:` type($result)

لبه جریان داده برخی از عملیات X، پلی بین مجموعه ای از منابع (هر کدام یک عملوند از X یا یک عملوند پایان دهنده بلوک X است) و مجموعه ای از اهداف (هر کدام نتیجه ای از X یا آرگومان بلوکی X هستند) را تعریف می کند، به طوری که همه منابع و اهداف باید به طور یکسان تقسیم شوند.

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

به عنوان مثال:

  y_0, ..., y_n = while (x_0, ..., x_n)
                  ((pred_arg_0,... , pred_arg_n) { ... })
                  ((body_arg_0,..., body_arg_n) {
                    ...
                    return return_value_0, ..., return_value_n
                  })

در حالی که op دارای n لبه جریان داده است، iمین لبه جریان داده بین منابع x_i ، return_value_i و اهداف y_i ، pred_arg_i ، body_arg_i است.

یک sdy.data_flow_edge صاحب یک یال را به عنوان ورودی می‌گیرد (می‌تواند هر یک از اهداف باشد، اما ترجیحاً یک نتیجه عملیاتی به جای آرگومان بلوک)، که نباید کاربرد دیگری داشته باشد. این عملیات خالص نیست زیرا می‌تواند ورودی‌هایی را دریافت کند که در ابتدا هیچ استفاده‌ای نداشته است.

sdy.data_flow_edge همچنین دارای یک اشتراک گذاری اختیاری برای همه اهداف لبه است، و این تقسیم بندی باید به جای اشتراک گذاری اهداف (در صورت امکان پیوست) در طول انتشار به روز شود. این زمانی مفید است که یک عملیات دارای لبه‌های زیادی باشد، زیرا در موارد زیر بسیار کارآمدتر است:

• در هر لبه به طور جداگانه انتشار دهید.
• اشتراک گذاری هر لبه را به‌جای همه هدف‌ها به طور همزمان به‌روزرسانی کنید (مثلاً یک عملیات دارای یک TensorShardingPerValueAttr منفرد غیرقابل تغییر برای تقسیم‌بندی نتایج است).
• هنگامی که اشتراک گذاری منبع تغییر کرد، هر لبه را جداگانه به لیست کاری اضافه کنید.

Propagation تقسیم‌بندی‌ها را بین همه منابع و اهداف یک sdy.data_flow_edge منتشر می‌کند، به‌گونه‌ای که گویی یک عملیات معمولی با منابع به‌عنوان عملوند و اهداف به‌عنوان نتایج، و هویت sdy.op_sharding_rule است. این بدان معناست که انتشار رو به جلو از منبع به هدف و انتشار به عقب از هدف به منبع است.

ما اجازه نمی دهیم ورودی یک sdy.data_flow_edge توسط یک عملیات SdyDialect تعریف شود، بنابراین می توانیم فرض کنیم که توسط یک op که دارای ویژگی sdy.sharding ثبت نشده است، تعریف شده است.

صفات: SameOperandsAndResultType

رابط ها: InferTypeOpInterface ، SymbolUserOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.manual_computation (sdy::ManualComputationOp)

عملکرد موازی چند دستگاه با مجموعه های دستی

نحو:

operation ::= `sdy.manual_computation` `(`operands`)`
              `in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)
              `out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)
              `manual_axes````=```$manual_axes
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:`
              functional-type(operands, results)

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

بدنه بر روی manual_axes محلی است. انتشار از طریق بدنه روی هر محور آزاد - آنهایی که در لیست manual_axes نیستند - رخ می دهد.

توجه داشته باشید که انتظار می‌رود هر تانسور رتبه‌بندی نشده دارای یک تقسیم‌بندی با رتبه 0 باشد، یعنی کاملاً تکرار شود.

محدودیت ها:

• عناصر موجود در in_shardings و out_shardings باید محدودیت های فهرست شده در TensorShardingAttr را برآورده کنند.
• تعداد ورودی/خروجی تانسور سراسری و محلی منطقه عملیاتی باید مطابقت داشته باشد.
• محورهای دستی باید قبل از هر محور آزاد در هر شاردینگ کم نور قرار گیرند.
• محورهای دستی نمی توانند بالشتک را معرفی کنند. یعنی اندازه ابعاد باید بر اندازه محورهای دستی مربوطه قابل تقسیم باشد.
• شکل‌های جهانی و محلی آرگومان‌ها/نتایج مناطق عملیاتی باید مطابقت داشته باشند.

صفات: IsolatedFromAbove ، RecursiveMemoryEffects ، SingleBlockImplicitTerminator<ReturnOp> ، SingleBlock

رابط ها: ShardableDataFlowOpInterface ، SymbolUserOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.mesh (sdy::MeshOp)

مش به نام

نحو:

operation ::= `sdy.mesh` $sym_name `=` $mesh attr-dict

مش با نام جدیدی را تعریف می کند. همه مش های موجود در یک ماژول باید تعداد دستگاه های یکسانی داشته باشند (به جز مش های دارای یک device_id). مش یک عملیات Symbol است که در SymbolTable ماژول ظاهر می شود و می توان به name آن اشاره کرد.

صفات: HasParent<ModuleOp>

رابط ها: Symbol

ویژگی ها:

sdy.named_computation (sdy::NamedComputationOp)

نام عملیات محاسباتی

نحو:

operation ::= `sdy.named_computation` `<`$name`>` `` `(` $operands `)`
              (`in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)^)?
              (`out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)^)?
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:` functional-type($operands, results)

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

این را می توان برای انتشار از طریق دستورالعمل های فراخوانی به عملکردهای دیگر استفاده کرد. همه کاربران Shardy باید یک مجوز import/export بنویسند که عملیات فراخوانی آنها را به sdy.named_computation ops تبدیل کند و بدنه تابع فراخوانده شده را در بدنه named_computation کپی/کپی کند.

نوع آرگومان های هر بلوک و مقادیر برگشتی در منطقه باید با نوع عملوندها و نوع نتایج عملیات یکسان باشد.

مثال:

%1 = sdy.named_computation<"foo">(%0) (%arg1: tensor<16x32xf32>) {
  sdy.return %arg1 : tensor<16x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

صفات: IsolatedFromAbove ، RecursiveMemoryEffects ، RecursivelySpeculatableImplTrait ، SingleBlockImplicitTerminator<ReturnOp> ، SingleBlock

رابط ها: ConditionallySpeculatable ، InferTypeOpInterface ، ShardableDataFlowOpInterface ، SymbolUserOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

عملیات مانع انتشار

نحو:

operation ::= `sdy.propagation_barrier` $input `allowed_direction````=```$allowed_direction attr-dict `:` type($input)

این اپ مانند یک عملیات هویتی عمل می کند و همان مقداری را که به عنوان ورودی گرفته بود، خروجی می دهد. اما از نظر انتشار، این فقط اجازه می دهد تا انتشار از طریق آن در جهت خاصی جریان یابد.

این مانع از انتشار خرده‌ها بین استفاده از نتیجه عملیات مانع و عملوند آن می‌شود.

• FORWARD به این معنی است که تقسیم‌بندی‌ها فقط می‌توانند از عملوند به نتیجه جریان پیدا کنند.
• BACKWARD به این معنی است که تکه تکه ها فقط می توانند از نتیجه به عملوند جریان داشته باشند.
• NONE به این معنی است که هیچ اشتراک گذاری نمی تواند از طریق این عملیات منتشر شود.
• نمی توان BOTH مشخص کرد، زیرا این عملیات اضافی است.

صفات: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

رابط ها: ConditionallySpeculatable ، InferTypeOpInterface ، NoMemoryEffect (MemoryEffectOpInterface)

جلوه ها: MemoryEffects::Effect{}

ویژگی ها:

عملگرها:

نتایج:

sdy.reduce_scatter (sdy::ReduceScatterOp)

ارتباط کاهش پراکندگی را در امتداد محورها انجام می دهد

نحو:

operation ::= `sdy.reduce_scatter` $reduce_scatter_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

تکه های یک تانسور را در امتداد محورهای مشخص شده در reduce_scatter_axes کاهش می دهد و سپس نتیجه را در امتداد همان محورها پراکنده می کند. این عملیات اساساً ترکیبی از یک sdy.all_reduce به دنبال یک sdy.all_slice در امتداد همان reduce_scatter_axes است.

محدودیت ها:

• باید محدودیت های فهرست شده در Sdy_CollectiveOpInterface را برآورده کند.
• عناصر موجود در reduce_scatter_axes باید محدودیت های فهرست شده در AxisRefListAttr را برآورده کنند.
• با استفاده از reduce_scatter_axes به عملوند، به اشتراک گذاری out_sharding تبدیل می شود.

صفات: SameOperandsAndResultType

رابط ها: CollectiveOpInterface ، InferTypeOpInterface ، SymbolUserOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.reshard (sdy::ReshardOp)

یک تانسور را مجدداً به یک تقسیم بندی متفاوت تبدیل می کند

نحو:

operation ::= `sdy.reshard` $input $sharding attr-dict `:` type($result)

تانسور ورودی را با تقسیم بندی مشخص شده مجدداً شارد می کند، که با اشتراک گذاری موجود تانسور ورودی متفاوت است.

هم ShardingConstraintOp و هم ReshardOp یک شاردینگ را به یک تانسور متصل می کنند. طول عمر آنها عبارت است از:

1. قبل از انتشار به اشتراک گذاری، ShardingConstraintOp توسط کاربران اضافه می شود.
2. انتشار Sharding، ShardingConstraintOp را مصرف می کند. هیچ ShardingConstraintOp در نتایج انتشار شاردینگ وجود ندارد. در عوض، ReshardOp ممکن است در صورت نیاز اضافه شود.
3. یک پارتیشن‌کننده یک ReshardOp را به یک op جمعی (یا یک عملیات هویتی) تبدیل می‌کند. در نتایج پارتیشن‌کننده نباید ReshardOp وجود داشته باشد.

// TODO(b/331680067). یک الگوی متعارف اضافه کنید تا عملیات اضافی // reshard حذف شود.

صفات: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

رابط ها: ConditionallySpeculatable ، InferTypeOpInterface ، NoMemoryEffect (MemoryEffectOpInterface) ، SymbolUserOpInterface

جلوه ها: MemoryEffects::Effect{}

ویژگی ها:

عملگرها:

نتایج:

sdy.return (sdy::ReturnOp)

عملیات sdy.return مناطق متصل به عملیات مبتنی بر منطقه sdy و هر عملیات مبتنی بر منطقه Shardy دیگر را خاتمه می دهد. متغیر است: فهرستی از مقادیر را به عنوان آرگومان می‌گیرد که انواع آن‌ها می‌تواند هر (اما از یک نوع، به عنوان مثال AnyTensor ) باشد و بنابراین می‌تواند در سطوح مختلف پشته IR Shardy دوباره استفاده شود.

نحو:

operation ::= `sdy.return` attr-dict ($results^ `:` type($results))?

ویژگی ها: AlwaysSpeculatableImplTrait , Terminator

رابط ها: ConditionallySpeculatable ، NoMemoryEffect (MemoryEffectOpInterface)

جلوه ها: MemoryEffects::Effect{}

عملگرها:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

یک تانسور را به اشتراک گذاری مشخص شده محدود می کند

نحو:

operation ::= `sdy.sharding_constraint` $input $sharding attr-dict `:` type($result)

یک تقسیم بندی را به یک تانسور میانی (مثلاً نتیجه یک ماتمول) متصل می کند تا نشان دهد که این تانسور یا زیر مجموعه ای از کاربردهای آن باید به این ترتیب تقسیم شود.

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

این عملیات می تواند:

• بدون استفاده (آویزان) - به این معنی که شاردینگ متصل به این صورت است که خود تانسور ورودی چگونه باید خرد شود.
• دارای کاربردها باشد - به این معنی که تقسیم بندی متصل به این صورت است که استفاده از عملیات محدودیت تقسیم بندی باید چگونه تقسیم شود، در حالی که سایر کاربردهای تانسور ورودی ممکن است تقسیم بندی متفاوتی داشته باشند (اگر تانسور ورودی کاربرد دیگری نداشته باشد، رفتار مشابه حالت بدون استفاده است).

صفات: SameOperandsAndResultType

رابط ها: InferTypeOpInterface ، SymbolUserOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.sharding_group (sdy::ShardingGroupOp)

تانسورها را در گروه محدود می کند تا به اشتراک گذاری یکسانی داشته باشند.

نحو:

operation ::= `sdy.sharding_group` $input `group_id````=```$group_id attr-dict `:` type($input)

این عملیات یک رابط برای تخصیص تانسورها به گروه‌های شاردینگ (گروه‌هایی از تانسورهایی که مجبور می‌شوند دارای تقسیم‌بندی‌های یکسان باشند) ارائه می‌کند. در حین انتشار، به محض اینکه یکی از عناصر گروه خرد شد، همه اعضای دیگر دقیقاً به همان روش خرد می شوند. این عملیات شناسه گروه آرگومان را می گیرد و هیچ نتیجه ای را برمی گرداند، اما در عوض نمایش گروه اشتراک گذاری داخلی را تغییر می دهد تا تانسور ورودی را به گروه با شناسه داده شده اضافه کند.

رابط ها: InferTypeOpInterface

ویژگی ها:

عملگرها:

صفات

AllToAllParamAttr

پارامتر همه به همه

نحو:

#sdy.all_to_all_param<
  ::llvm::ArrayRef<AxisRefAttr>,   # axes
  int64_t,   # src_dim
  int64_t   # tgt_dim
>

یک تاپل حاوی محورها و ابعاد منبع/هدف برای اجرای همه به همه.

پارامترها:

AllToAllParamListAttr

فهرست پارامترهای همه به همه

نحو:

#sdy.all_to_all_param_list<
  ::llvm::ArrayRef<AllToAllParamAttr>   # value
>

پارامترها:

AxisRefAttr

ارجاع به یک محور کامل یا یک محور فرعی تقسیم شده

نحو:

#sdy.axis_ref<
  ::llvm::StringRef,   # name
  SubAxisInfoAttr   # sub_axis_info
>

محدودیت ها:

• name باید در MeshAttr محدود شده باشد.
• اگر sub_axis_info وجود داشته باشد، باید محدودیت های SubAxisInfoAttr را برآورده کند.

پارامترها:

AxisRefListAttr

فهرست ارجاعات محوری

نحو:

#sdy.axis_ref_list<
  ::llvm::ArrayRef<AxisRefAttr>   # value
>

محدودیت ها:

• عناصر در value باید محدودیت های AxisRefAttr را برآورده کنند.
• هیچ محور تکراری یا محور فرعی وجود ندارد که با یکدیگر همپوشانی داشته باشند.
• هیچ دو محور همجوار، زیرمحورهای متوالی از همان محور کامل نیستند، یعنی می‌توانند در یک محور فرعی یا محور کامل ادغام شوند.

پارامترها:

AxisToPropagationDetailsAttr

جزئیات جریان لبه انتشار برای یک محور و منبع خاص.

نحو:

#sdy.axis_to_propagation_details<
  ::mlir::sdy::AxisRefAttr,   # axis_name
  ::mlir::sdy::EdgeValueRefAttr,   # source
  ::llvm::ArrayRef<EdgeValueRefAttr>   # targets
>

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

پارامترها:

DimMappingAttr

فهرست شاخص های عاملی برای یک بعد

یک لیست خالی نشان می دهد که این یک نگاشت تهی است (این با * تجزیه/چاپ شده است)، یعنی بعد به هیچ عاملی نگاشت نشده است.

محدودیت ها:

• حداقل یک شاخص عامل وجود دارد.
• شاخص های عامل باید در محدوده [0, $factor_sizes ) باشد.
• اگر چندین فاکتور وجود داشته باشد، هیچ یک از آنها نمی توانند سایز 1 داشته باشند.
• بدون شاخص های عامل تکراری.

پارامترها:

DimensionShardingAttr

تقسیم بندی ابعاد

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

محدودیت ها:

• عناصر موجود در axes باید محدودیت های فهرست شده در AxisRefListAttr را برآورده کنند.
• اگر تقسیم ابعاد دارای اولویت است:
  • اولویت بزرگتر یا مساوی 0 است.
  • بعد اگر بسته باشد حداقل یک محور دارد.

پارامترها:

EdgeValueRefAttr

ارجاع به یک شاخص خاص از لبه مقدار از type .

نحو:

#sdy.edge_value_ref<
  ::mlir::sdy::EdgeNodeType,   # type
  int64_t   # index
>

پارامترها:

ListOfAxisRefListsAttr

لیست لیست های مرجع

نحو:

#sdy.list_of_axis_ref_lists<
  ::llvm::ArrayRef<AxisRefListAttr>   # value
>

پارامترها:

ManualAxesAttr

فهرستی از محورهایی که ManualComputationOp روی آنها به صورت دستی است

نحو:

#sdy.manual_axes<
  ::llvm::ArrayRef<StringAttr>   # value
>

پارامترها:

MeshAttr

مش محورها و لیستی از دستگاه ها

نحو:

#sdy.mesh<
  ::llvm::ArrayRef<MeshAxisAttr>,   # axes
  ::llvm::ArrayRef<int64_t>   # device_ids
>

مش فهرستی از محورها و فهرست اختیاری شناسه‌های دستگاه است که سفارش دستگاه را مشخص می‌کند.

اگر فهرست محورها خالی باشد، مش دارای یک محور ضمنی بدون نام به اندازه 1 است. در این حالت، اگر فهرست شناسه دستگاه ارائه نشده باشد، فهرست شناسه دستگاه ضمنی [0] است. اگر لیست شناسه دستگاه ارائه شده باشد، باید دارای یک عدد صحیح با هر مقدار غیر منفی باشد. ما این مورد را حداکثر تقسیم بندی می نامیم.

برای همه موارد غیر حداکثر اشتراک‌گذاری، اگر فهرست شناسه دستگاه مشخص شده باشد، حاصل ضرب اندازه‌های محور باید با تعداد دستگاه‌ها مطابقت داشته باشد. اگر لیست شناسه دستگاه مشخص نشده باشد، لیست شناسه دستگاه ضمنی iota (محصول(محور)) است. برای سادگی، ما همچنین تعیین لیست شناسه دستگاه که همان iota(product(axes)) است را ممنوع می کنیم. در این مورد، لیست شناسه دستگاه نباید مشخص شود.

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

• یک مش خالی نشان دهنده یک شبکه نگهدارنده است که می تواند در حین انتشار جایگزین شود: <[]>
• یک مش با یک محور بدون نام و یک شناسه دستگاه صریح، که معمولاً برای نمایش حداکثر اشتراک‌گذاری استفاده می‌شود: <[], device_ids=[3]>
• مش با دو محور و شناسه‌های دستگاه ضمنی iota(6): <["a"=2، "b"=3]>
• مش با دو محور و شناسه‌های صریح دستگاه که سفارش دستگاه را مشخص می‌کند: <["a"=3، "b"=2]، device_ids=[0، 2، 4، 1، 3، 5]>

محدودیت ها:

• عناصر موجود در axes نباید نام تکراری داشته باشند.
• اگر device_ids مشخص شده باشد:
  • حاصل ضرب اندازه های محور باید با تعداد دستگاه ها مطابقت داشته باشد.
  • همه عناصر آن باید غیر منفی باشند.
  • device_ids نباید برابر با iota(product(axis_sizes)) باشد.
  • device_ids مرتب شده باید iota(product(axis_sizes)) باشد.

پارامترها:

MeshAxisAttr

نام محور در یک مش

نحو:

#sdy.mesh_axis<
  ::llvm::StringRef,   # name
  int64_t   # size
>

پارامترها:

OpShardingRuleAttr

نحوه پارتیشن بندی یک عملیات را مشخص می کند.

نحو:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  ::llvm::ArrayRef<int64_t>,   # reduction_factors
  ::llvm::ArrayRef<int64_t>,   # need_replication_factors
  ::llvm::ArrayRef<int64_t>,   # permutation_factors
  ::llvm::ArrayRef<int64_t>,   # blocked_propagation_factors
  bool   # is_custom_rule
>

یک قانون تقسیم بندی مشخص می کند که چگونه یک عملیات را می توان بر اساس ویژگی های مختلف در عملیات پارتیشن بندی کرد - هر ویژگی، شکل عملوندها، شکل نتایج و غیره. به عنوان مثال:

%0 = stablehlo.add %arg0, %arg1 {
    sdy.sharding_rule = #sdy.op_sharding_rule<
        ([i, j],[i, j])->([i, j])
        {i=8, j=8}>
} : tensor<8x8xf32>

%1 = stablehlo.dot_general %arg2, %arg3, contracting_dims = [1] x [0] {
  sdy.sharding_rule = #sdy.op_sharding_rule<
      ([i, k],[k, j])->([i, j])
      {i=8, j=16, k=8}>
}: (tensor<8x8xf32>, tensor<8x16xf32>) -> tensor<8x16xf32>

توجه داشته باشید که ما فاکتورهایی را با اندازه 1 مجاز می‌کنیم، حتی اگر نمی‌توان آنها را تقسیم کرد، این عمدتاً برای کامل بودن است زیرا بسیاری از عملیات‌ها مانند عملیات نقطه‌ای دارای ابعاد اندازه یک هستند که با عملوندها و نتایج مطابقت دارند.

انواع عوامل:

• reduction_factors شامل شاخص های عواملی است که نیاز به کاهش دارند، مانند ابعاد انقباض در یک عمل نقطه. این عوامل می توانند در عملوندها باشند اما در نتایج نه.
• need_replication_factors شامل شاخص های عواملی است که نیاز به تکرار کامل دارند، مانند بعد مرتب شده در یک عملیات مرتب سازی.
• permutation_factors شامل شاخص های فاکتورهایی است که در صورت تکه تکه شدن، به جایگشت جمعی نیاز دارند، مانند ابعاد padding در یک عملیات پد.
• همه عوامل دیگر به عنوان فاکتورهای عبوری در نظر گرفته می شوند، به عنوان مثال، فاکتورهایی که اگر به طور یکسان در همه تانسورهایی که به آنها نگاشت می شوند، به هیچ ارتباطی نیازی نداشته باشند.

blocked_propagation_factors شامل فاکتورهایی است که در آنها به اشتراک گذاری اجازه انتشار داده نمی شود. این متعامد به انواع عامل است. یعنی یک عامل انتشار مسدود شده می تواند هر یک از انواع فاکتورها باشد.

is_custom_rule توضیح می دهد که آیا این یک قانون تعریف شده توسط یک کاربر است یا خیر. کاربران می توانند قوانین اشتراک گذاری را برای تماس های سفارشی خود تعریف کنند یا قوانین اشتراک گذاری از پیش تعریف شده را برای عملیات استاندارد بازنویسی کنند. یک قانون سفارشی همیشه حفظ می شود/هرگز حذف نمی شود.

محدودیت ها:

• تعداد عملوند/نتایج نگاشت باید با تعداد عملوند/نتایج عملیات مطابقت داشته باشد.
• حداقل یک نقشه برداری وجود دارد (نمی توان برای عملیات بدون عملوند/نتیجه قاعده ای داشت).
• رتبه هر TensorMappingAttr با رتبه نوع تانسور مربوطه مطابقت دارد.
• برای هر گروه از عوامل ( reduction_factors ، need_replication_factors ، permutation_factors ):
  • عناصر باید در محدوده [0, $factor_sizes ] باشند.
  • هیچ شاخص عامل تکراری در هر گروه و در بین گروه ها وجود ندارد.

پارامترها:

PropagationEdgesAttr

فراداده لبه انتشار برای همه مراحل انتشار.

نحو:

#sdy.propagation_edges<
  ::llvm::ArrayRef<PropagationOneStepAttr>   # value
>

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

پارامترها:

PropagationOneStepAttr

فراداده انتشار هر مرحله

نحو:

#sdy.propagation_one_step<
  int64_t,   # step_index
  ::llvm::ArrayRef<AxisToPropagationDetailsAttr>   # axis_entries
>

جزئیات انتشار برای همه محورها برای یک مرحله انتشار.

پارامترها:

SubAxisInfoAttr

اطلاعاتی در مورد نحوه استخراج این محور فرعی از محور کامل

نحو:

#sdy.sub_axis_info<
  int64_t,   # pre_size
  int64_t   # size
>

هنگام تقسیم یک محور کامل به n محور فرعی، محور به [k_1,...,k_n] تغییر شکل می‌دهد و محور فرعی i را می‌توان با حاصلضرب تمام اندازه‌های محور در سمت چپ آن m=prod(k_1,...,k_(i-1)) (با نام پیش‌اندازه) و اندازه k_i بیان کرد. بنابراین، ویژگی sub-axis-info آن دو عدد را نگه می‌دارد و به صورت زیر نشان داده می‌شود: (m)k برای پیش اندازه m و اندازه k.

محدودیت ها:

• pre-size حداقل 1 است.
• size بزرگتر از 1 است.
• pre-size باید اندازه محور کامل را تقسیم کند، یعنی هم pre-size و هم size ، اندازه محور کامل را تقسیم می کنند و محور فرعی فراتر از محور کامل نمی رود.
• اندازه محور فرعی با اندازه محور کامل مربوطه برابر نیست، در این صورت باید از محور کامل به جای آن استفاده شود.

پارامترها:

TensorMappingAttr

نگاشت عاملی برای هر بعد یک تانسور.

نحو:

#sdy.tensor_mapping<
  ::llvm::ArrayRef<DimMappingAttr>   # dim_mappings
>

محدودیت ها:

• عناصر موجود در dim_mappings باید محدودیت های موجود در DimMappingAttr را برآورده کنند.
• بدون شاخص فاکتورهای تکراری در ابعاد.

پارامترها:

TensorShardingAttr

تقسیم تانسور

نحو:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>,   # replicated_axes
  ::llvm::ArrayRef<AxisRefAttr>   # unreduced_axes
>

یک تقسیم بندی تانسور به یک مش خاص محدود شده است و فقط می تواند به نام محورها از آن مش ارجاع دهد. تقسیم‌بندی ابعاد برای هر بعد تانسور به ما می‌گوید که در امتداد کدام محورها (یا محورهای فرعی) از اصلی به فرعی تقسیم می‌شود. تمام محورهای دیگری که یک بعد را تکه تکه نمی کنند، به طور ضمنی یا صریح (اگر در لیست محورهای تکرار شده ظاهر شوند) تکرار می شوند.

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

شبکه ای که این تقسیم بندی به آن محدود می شود را می توان با نام نماد، اشاره به نماد MeshOp مربوطه یا یک MeshAttr خطی ​​مشخص کرد.

شاردینگ می تواند محورهای کاهش نیافته (مشخص شده توسط unreduced_axes ) داشته باشد، به این معنی که تانسور در امتداد این محورها کاهش نیافته است. به عنوان مثال، اگر بعد انقباض یک ماتمول در امتداد محور x در هر دو lhs و rhs تقسیم شود، نتیجه در امتداد x کاهش نمی یابد. اعمال یک کاهش کامل بر روی تانسور در امتداد محورهای کاهش نیافته باعث می شود تانسور در امتداد آن محورها تکرار شود. با این حال، یک تانسور با محورهای کاهش نیافته لزومی ندارد که فوراً تماماً کاهش یابد، می‌تواند در صورت انتقال به عملیات خطی مانند stablehlo.add (تا زمانی که هم lhs و هم rhs کاهش نیافته باشند) کاهش نیافته و پس از آن کاملاً کاهش می‌یابد. ما نوع کاهش را مجموع فرض می‌کنیم، کاهش‌های دیگری ممکن است در آینده پشتیبانی شوند.

محدودیت ها:

• عناصر موجود در dim_shardings باید محدودیت های فهرست شده در DimensionShardingAttr را برآورده کنند.
• عناصر موجود در replicated_axes باید محدودیت های فهرست شده در AxisRefListAttr را برآورده کنند.
• عناصر موجود در unreduced_axes باید محدودیت های فهرست شده در AxisRefListAttr را برآورده کنند.
• اگر نوع تانسور مربوطه یک ShapedType نیست، تقسیم بندی باید دارای رتبه 0 و بدون محورهای تکراری باشد.
• اگر یک ShapedType است، پس:
  • تانسور باید دارای رتبه باشد.
  • تعداد تقسیم‌بندی‌های ابعادی برابر با رتبه تانسور است.
  • ابعاد اندازه 0 تکه تکه نشده است.
• هیچ محور تکراری یا محور فرعی وجود ندارد که در بین dim_shardings ، replicated_axes و unreduced_axes با یکدیگر همپوشانی داشته باشند.
• آیتم‌های موجود در replicated_axes و unreduced_axes به صورت wrt mesh_or_ref مرتب می‌شوند (به AxisRefAttr::getMeshComparator مراجعه کنید).

پارامترها:

TensorShardingPerValueAttr

تقسیم بندی تانسور در هر عملوند/نتیجه یک عملیات

نحو:

#sdy.sharding_per_value<
  ::llvm::ArrayRef<TensorShardingAttr>   # shardings
>

فهرستی از TensorShardingAttr s، یکی برای هر عملوند/نتیجه یک عملیات.

محدودیت ها:

• عناصر موجود در shardings باید محدودیت های TensorShardingAttr را برآورده کنند.

پارامترها:

Enums

EdgeNodeType

نوع گره لبه enum

موارد:

جهت انتشار

جهت انتشار enum

موارد: