'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 فهرستی از فهرست‌های محورها است. فهرست بیرونی روی ابعاد تانسور قرار دارد. هر فهرست درونی، محورهایی را مشخص می‌کند که در امتداد آنها یک جمع‌آوری جداگانه باید روی بُعد مربوطه انجام شود. این فهرست روی خرد کردن عملوند ( tensor ) اعمال می‌شود تا خرد کردن نتیجه ( out_sharding ) به دست آید.

توجه داشته باشید که 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 روی عملوند sharding باعث می‌شود 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 را برآورده کند.
• reduction_axes باید نسبت به مش مرتب شوند.
• عملوندهای sharding و out_sharding باید 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 فهرستی از فهرست‌های محورها است. فهرست بیرونی روی ابعاد تانسور قرار دارد. هر فهرست درونی، محورهایی را مشخص می‌کند که یک برش باید در امتداد آنها روی بُعد مربوطه انجام شود. این فهرست روی خرد کردن عملوند ( tensor ) اعمال می‌شود تا خرد کردن نتیجه ( out_sharding ) به دست آید.

توجه داشته باشید که 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)

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

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

همه به همه برای تقسیم‌بندی عملوند ( 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 تغییرناپذیر برای خرد کردن نتایج است).
• وقتی شاردینگ یک منبع تغییر کرد، هر یال را جداگانه به لیست کاری اضافه کنید.

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

ما اجازه نمی‌دهیم ورودی یک sdy.data_flow_edge توسط یک عملیات SdyDialect تعریف شود، بنابراین می‌توانیم فرض کنیم که توسط عملیاتی تعریف شده است که دارای ویژگی 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 نیستند - رخ خواهد داد.

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

محدودیت‌ها:

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

ویژگی‌ها: 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 تبدیل کند و بدنه تابع فراخوانی شده را در بدنه named_computation کپی/تکثیر کند.

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

مثال:

%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 روی عملوند sharding باعث می‌شود out_sharding .

صفات: SameOperandsAndResultType

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

ویژگی‌ها:

عملوندها:

نتایج:

sdy.reshard (sdy::ReshardOp)

یک تانسور را به یک شاردینگ متفاوت Reshard می‌کند

نحو:

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

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

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

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

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

صفات: AlwaysSpeculatableImplTrait ، SameOperandsAndResultType

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

اثرات: MemoryEffects::Effect{}

ویژگی‌ها:

عملوندها:

نتایج:

sdy.return (sdy::ReturnOp)

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

نحو:

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)

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

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

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

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

صفات: 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
>

پارامترها:

محورRefAttr

ارجاع به یک محور کامل یا یک زیرمحور مجزا

نحو:

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

محدودیت‌ها:

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

پارامترها:

محورهاRefListAttr

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

نحو:

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

محدودیت‌ها:

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

پارامترها:

جزئیات محور به انتشارAttr

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

نحو:

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

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

پارامترها:

DimMappingAttr

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

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

محدودیت‌ها:

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

پارامترها:

DimensionShardingAttr

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

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

محدودیت‌ها:

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

پارامترها:

EdgeValueRefAttr

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

نحو:

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

پارامترها:

فهرست محورهاRefلیست‌هاAttr

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

نحو:

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

پارامترها:

محورهای دستیAttr

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

نحو:

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

پارامترها:

مش اتر

شبکه‌ای از محورها و فهرستی از دستگاه‌ها

نحو:

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

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

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

• اگر device_ids ارائه نشده باشد، یک مش خالی است.
• اگر device_ids ارائه شود، باید یک عدد صحیح غیر منفی باشد، که ما آن را مش با حداکثر شاردینگ می‌نامیم.

اگر فهرست محورها ارائه شود

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

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

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

محدودیت‌ها:

• عناصر موجود در device_ids باید غیر منفی باشند.
• اگر axes خالی باشد، اندازه device_ids می‌تواند ۰ (مش خالی) یا ۱ (مش با حداکثر شاردینگ) باشد.
• اگر axes خالی نباشند،
  • عناصر موجود در 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>

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

انواع فاکتور:

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

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

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

محدودیت‌ها:

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

پارامترها:

لبه‌های انتشارAttr

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

نحو:

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

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

پارامترها:

تکثیر تک مرحله‌ای

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

نحو:

#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)) (معروف به pre-size) و اندازه k_i. بنابراین، ویژگی sub-axis-info این دو عدد را در خود نگه می‌دارد و به صورت زیر نشان داده می‌شود: (m)k برای pre-size m و اندازه k.

محدودیت‌ها:

• pre-size حداقل ۱ است.
• size بزرگتر از ۱ است.
• 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 مشخص می‌شود)، به این معنی که تانسور در امتداد این محورها کاهش نیافته است. برای مثال، اگر بُعد انقباضی یک matmul در امتداد محور x در هر دو lhs و rhs تقسیم‌بندی شود، نتیجه در امتداد x کاهش نیافته است. اعمال یک کاهش کامل روی تانسور در امتداد محورهای کاهش نیافته باعث می‌شود تانسور در امتداد آن محورها تکرار شود. با این حال، یک تانسور با محورهای کاهش نیافته لازم نیست بلافاصله کاهش یافته باشد، می‌تواند هنگام انتقال به عملیات خطی مانند stablehlo.add (تا زمانی که هر دو lhs و rhs کاهش نیافته باشند) و پس از آن کاهش یافته، کاهش نیافته باقی بماند. ما فرض می‌کنیم نوع کاهش sum است، کاهش‌های دیگر ممکن است در آینده پشتیبانی شوند.

محدودیت‌ها:

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

پارامترها:

TensorShardingPerValueAttr

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

نحو:

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

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

محدودیت‌ها:

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

پارامترها:

انوم‌ها

نوع گره لبه

نوع گره لبه‌ای enum

موارد:

جهت انتشار

جهت انتشار شمارشی

موارد: