'sdy' گویش

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

The Shardy (SDY) dialect defines an axis-based tensor sharding representation and additional API components to attach shardings to tensors.

گزارش نسخه: 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 به دست آید.

Traits: SameOperandsAndResultType

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

ویژگی‌ها:

عملوندها:

نتایج:

sdy.all_reduce (sdy::AllReduceOp)

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

نحو:

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

Reduces chunks of a tensor along axes specified in reduction_axes . The order of reduction_axes is not important for the result, but can affect the order of the corresponding replica groups.

محدودیت‌ها:

• باید محدودیت‌های ذکر شده در Sdy_CollectiveOpInterface را برآورده کند.
• reduction_axes باید محدودیت‌های ذکر شده در AxisRefListAttr را برآورده کند.
• reduction_axes باید نسبت به مش مرتب شوند.
• The operand sharding and out_sharding must have equivalent dimension shardings.
• 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 ) به دست آید.

Note that out_sharding is not used to determine the sharding of the result. Instead, the sharding of the result is determined by the sharding of the operand and the slicing_axes , and out_sharding must match this inferred 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 .

Traits: 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 :
  • Elements in axes must satisfy the constraints of AxisRefAttr .
  • src_dim and tgt_dim must be valid dimensions (non-negative and less than rank of tensor).
  • هر 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 باید با حاصلضرب اندازه‌های عملوند مربوطه در شاردینگ بُعد مطابقت داشته باشد.

Traits: SameOperandsAndResultType

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

ویژگی‌ها:

عملوندها:

نتایج:

sdy.constant (sdy::ConstantOp)

عملیات ثابت

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

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

Example:

%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 قرار دارد.

An sdy.data_flow_edge takes as input the owner of an edge (can be any of the targets, but preferably an op result rather than a block argument), which shouldn't have any other uses. This op isn't pure because it can take an input that originally didn't have any uses.

The sdy.data_flow_edge also holds an optional sharding for all targets of the edge, and that sharding should be updated instead of the targets' sharding (if can be attached) during propagation. This is useful when an op has many edges, as it's much more efficient to:

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

Propagation will propagate shardings between all sources and targets of a sdy.data_flow_edge as if it was a regular op with the sources as operands and targets as results, and an identity sdy.op_sharding_rule . That means that forward propagation is from sources to targets and backwards propagation is from targets to sources.

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

صفات: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , SymbolUserOpInterface

Attributes:

عملوندها:

نتایج:

sdy.func_data_flow_edge (sdy::FuncDataFlowEdgeOp)

تابع ورودی/خروجی جریان داده، لبه عملیات.

نحو:

operation ::= `sdy.func_data_flow_edge` $operand attr-dict `:` type($result)

یک لبه جریان داده op اما برای آرگومان‌های func یا نتایج فراخوانی. وقتی عملوند آن یک BlockArgument باشد؛ پلی از آرگومان فراخواننده callOp به کاربران آرگومان func است. برای هر آرگومان func یک لبه جریان داده func وجود دارد. وقتی عملوند آن یک OpResult باشد؛ پلی از مقدار برگشتی funcOp فراخوانی شده به کاربران نتیجه فراخوانی است. برای هر نتیجه فراخوانی یک لبه جریان داده func وجود دارد.

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

Interfaces: 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)

Groups a computation, ie a block of operations, and gives it a name. Propagation will flow in/out of the region as if everything was inlined.

This can be used to handle propagating through call instructions to other functions. Any users of Shardy should write an import/export pass that converts their call ops to sdy.named_computation ops, duplicating/copying the body of the called function into the body of the 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

Interfaces: ConditionallySpeculatable , InferTypeOpInterface , ShardableDataFlowOpInterface , SymbolUserOpInterface

ویژگی‌ها:

عملوندها:

نتایج:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

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

نحو:

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

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

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

• FORWARD به این معنی است که شاردینگ‌ها فقط می‌توانند از عملوند به نتیجه جریان پیدا کنند.
• BACKWARD means shardings can only flow from the result to the operand.
• NONE به این معنی است که هیچ شاردینگی نمی‌تواند از طریق این عملیات منتشر شود.
• نمی‌توان BOTH مشخص کرد، زیرا این گزینه اضافی خواهد بود.

صفات: AlwaysSpeculatableImplTrait ، SameOperandsAndResultType

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

Effects: MemoryEffects::Effect{}

ویژگی‌ها:

عملوندها:

نتایج:

sdy.reduce_scatter (sdy::ReduceScatterOp)

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

نحو:

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

Reduces chunks of a tensor along axes specified in reduce_scatter_axes , and then scatters the result along the same axes. This operation is essentially a combination of an sdy.all_reduce followed by an sdy.all_slice along the same reduce_scatter_axes .

محدودیت‌ها:

• باید محدودیت‌های ذکر شده در Sdy_CollectiveOpInterface را برآورده کند.
• عناصر موجود در reduce_scatter_axes باید محدودیت‌های ذکر شده در AxisRefListAttr را برآورده کنند.
• Applying reduce_scatter_axes to the operand sharding gets out_sharding .

صفات: SameOperandsAndResultType

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

ویژگی‌ها:

عملوندها:

نتایج:

sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

محورهای تکثیر شده به صورت ضمنی یا صریح را به محورهای کاهش نیافته منتقل کنید.

نحو:

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

axes باید به صورت ضمنی یا صریح در عملوند تکرار شوند. این عملیات باعث می‌شود که آنها در نتیجه کاهش نیابند. رابطه زیر را داریم:

تمام-کاهش(تکثیر شده به کاهش نیافته(x، محورها)، محورها) = x

مثال:

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

محدودیت‌ها:

• باید محدودیت‌های ذکر شده در Sdy_CollectiveOpInterface را برآورده کند.
• axes باید محدودیت‌های ذکر شده در AxisRefListAttr را برآورده کنند.
• axes باید نسبت به مش مرتب شوند.
• axes خالی نیستند.
• شاردینگ ورودی و خروجی باید شاردینگ‌هایی با ابعاد یکسان داشته باشند.
• axes باید به طور ضمنی یا صریح در تقسیم‌بندی عملوند تکرار شوند.
• محورهای کاهش نیافته + محورهای کاهش نیافته = محورهای کاهش نیافته خروجی.

صفات: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , Sdy_CollectiveOpInterface , 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 در نتایج پارتیشن‌بند وجود داشته باشد.

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

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

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

Operands:

sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

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

نحو:

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

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

all-gather(x, axiss) = all-reduce(sharded-to-unreduced(x, axiss), axiss)، که در آن all-gather، sharded-to-unreduced، و all-reduce روی محورهای یکسان اعمال می‌شوند.

مثال:

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

محدودیت‌ها:

• باید محدودیت‌های ذکر شده در Sdy_CollectiveOpInterface را برآورده کند.
• عناصر موجود در axes باید محدودیت‌های ذکر شده در AxisRefListAttr را برآورده کنند.
• اعمال axes به عملوند sharding باعث می‌شود out_sharding .

صفات: SameOperandsAndResultType

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

ویژگی‌ها:

عملوندها:

نتایج:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

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

نحو:

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

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

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

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

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

صفات: SameOperandsAndResultType

رابط‌ها: InferTypeOpInterface ، SymbolUserOpInterface

Attributes:

عملوندها:

نتایج:

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
>

Maps a source value reference to a list of target value references along a particular axis.

پارامترها:

DimMappingAttr

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

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

محدودیت‌ها:

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

پارامترها:

DimensionShardingAttr

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

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

محدودیت‌ها:

• عناصر موجود در axes باید محدودیت‌های ذکر شده در AxisRefListAttr را برآورده کنند.
• اگر تقسیم‌بندی ابعاد اولویت داشته باشد:
  • The priority is greater than or equal to 0.
  • اگر بُعد بسته باشد، حداقل یک محور دارد.

پارامترها:

EdgeValueRefAttr

Reference to a particular index of a value edge of type type .

نحو:

#sdy.edge_value_ref<
  `operand` | `result`,   # 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)) باشد را مجاز نمی‌دانیم؛ در این حالت، فهرست شناسه دستگاه نباید مشخص شود.
• حتی اگر اندازه کل محورها ۱ باشد، این یک مش با حداکثر شاردینگ نیست.

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

• یک مش خالی نشان‌دهنده یک مش حفره‌دار است که می‌تواند در طول انتشار جایگزین شود: <[]>
• A mesh without axes list and a single non-negative device ID, which is a maximal-sharding mesh: <[], 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>

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

Factor types:

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

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

is_custom_rule describes whether this is a rule defined by a user. Users can define sharding rules for their custom calls or overwrite the pre-defined sharding rules for the standard operations. A custom rule is always preserved/never removed.

محدودیت‌ها:

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

پارامترها:

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

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

نحو:

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

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

پارامترها:

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

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

نحو:

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

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

Parameters:

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 must divide the size of the full axis, ie, both pre-size and size divide the size of the full axis, and the sub-axis doesn't go beyond the full axis.
• اندازه زیرمحور با اندازه محور کامل مربوطه برابر نیست، در این صورت باید از محور کامل استفاده شود.

Parameters:

TensorMappingAttr

Factor mappings for each dimension of a tensor.

نحو:

#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
  `sum` | `max` | `min`   # reduction_op
>

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

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

مش‌بندی که این شاردینگ به آن متصل است، می‌تواند یا با نام یک نماد که به نماد MeshOp مربوطه اشاره می‌کند، یا با یک MeshAttr درون‌خطی مشخص شود.

A sharding can have unreduced axes (specified by unreduced_axes ), meaning the tensor is unreduced along these axes. For example, if the contracting dimension of a matmul is sharded along axis x in both the lhs and rhs, the result is unreduced along x . Applying an all-reduce on the tensor along the unreduced axes will make the tensor replicated along those axes. However, a tensor with unreduced axes doesn't have to be all-reduced immediately, it can remain unreduced when passed to linear operations like stablehlo.add (as long as both lhs and rhs are unreduced) and all-reduced afterwards. We assume the reduction type is sum, other reductions may be supported in the future.

Constraints:

• Elements in dim_shardings must satisfy the constraints listed in DimensionShardingAttr .
• Elements in replicated_axes must satisfy the constraints listed in AxisRefListAttr .
• Elements in unreduced_axes must satisfy the constraints listed in AxisRefListAttr .
• If the corresponding tensor type isn't a ShapedType , the sharding must have rank 0 and no replicated axes.
• If it is a ShapedType , then:
  • The tensor should have a rank.
  • The number of dimension shardings is equal to the rank of the tensor.
  • Dimensions of size 0 aren't sharded.
• There are no duplicate axis-refs or sub-axes that overlap with one another across dim_shardings , replicated_axes , and unreduced_axes .
• Items in replicated_axes and unreduced_axes are ordered wrt mesh_or_ref (see AxisRefAttr::getMeshComparator ).

Parameters:

TensorShardingPerValueAttr

Tensor sharding per operand/result of an op

نحو:

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

A list of TensorShardingAttr s, one for each operand/result of an op.

Constraints:

• Elements in shardings must satisfy the constraints of TensorShardingAttr .

Parameters:

انوم‌ها

EdgeNodeType

Edge node type enum

Cases:

PropagationDirection

Propagation direction enum

Cases:

ReductionOp

Reduction op enum

Cases:

The Shardy (SDY) dialect

The Shardy (SDY) dialect defines an axis-based tensor sharding representation and additional API components to attach shardings to tensors.

Version log: 0.0.1: Add unreduced axes to TensorShardingAttr.

عملیات

sdy.all_gather (sdy::AllGatherOp)

Performs an all-gather communication along axes

نحو:

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

Gathers chunks of a tensor along axes specified in gathering_axes .

The gathering_axes is a list of lists of axes. The outer list is over the dimensions of the tensor. Each inner list specifies the axes along which a separate gather should be performed on the respective dimension. It will be applied to the sharding of the operand ( tensor ) to obtain the sharding of the result ( out_sharding ).

Note that out_sharding is not used to determine the sharding of the result. Instead, the sharding of the result is determined by the sharding of the operand and the gathering_axes , and out_sharding must match this inferred sharding.

Example:

%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>

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• Elements in gathering_axes must satisfy the constraints listed in AxisRefListAttr .
• Applying gathering_axes to the operand sharding gets out_sharding .

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.all_reduce (sdy::AllReduceOp)

Perform an all-reduce comunication along axes

نحو:

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

Reduces chunks of a tensor along axes specified in reduction_axes . The order of reduction_axes is not important for the result, but can affect the order of the corresponding replica groups.

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• reduction_axes must satisfy the constraints listed in AxisRefListAttr .
• reduction_axes must be sorted wrt the mesh.
• The operand sharding and out_sharding must have equivalent dimension shardings.
• reduction_axes must not overlap with the operand dimension sharding and replicated axes (it can overlap with unreduced axes).
• reduction_axes must not overlap with the unreduced axes of out_sharding . In other words, out_sharding must be be replicated along reduction_axes (implicitly or explicitly).

Traits: SameOperandsAndResultType

Interfaces: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.all_slice (sdy::AllSliceOp)

Performs a dynamic-slice operation along axes

نحو:

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

Slices chunks of a tensor along axes specified in slicing_axes . There is an algebric duality between sdy.all_slice and sdy.all_gather .

The slicing_axes is a list of lists of axes. The outer list is over the dimensions of the tensor. Each inner list specifies the axes along which a slice should be performed on the respective dimension. It will be applied to the sharding of the operand ( tensor ) to obtain the sharding of the result ( out_sharding ).

Note that out_sharding is not used to determine the sharding of the result. Instead, the sharding of the result is determined by the sharding of the operand and the slicing_axes , and out_sharding must match this inferred sharding.

Example:

%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>

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• Elements in slicing_axes must satisfy the constraints listed in AxisRefListAttr .
• Applying slicing_axes to the operand sharding gets out_sharding .

Traits: SameOperandsAndResultType

Interfaces: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.all_to_all (sdy::AllToAllOp)

Performs an all-to-all communication along axes

نحو:

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

For each (axes, src_dim, tgt_dim) tuple in the parameter list, this operation slices chunks of a tensor along dimension tgt_dim and axes specified in axes , scatteres those chunks along the axes, and concatenates them along dimension src_dim .

This operation is essentially a combination of an all-gather along src_dim and axes , followed by an all-slice along tgt_dim and axes , ie, a suffix of the axes sharding dimension src_dim on the input tensor is appended to the axes sharding dimension tgt_dim on the output tensor.

The all-to-all will be applied to the sharding of the operand ( tensor ) to obtain the sharding of the result ( out_sharding ).

Note that out_sharding is not used to determine the sharding of the result. Instead, the sharding of the result is determined by the sharding of the operand, src_dim , tgt_dim , and axes , and out_sharding must match this inferred sharding.

Example:

%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>

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• The parameter list must not be empty.
• For each parameter in params :
  • Elements in axes must satisfy the constraints of AxisRefAttr .
  • src_dim and tgt_dim must be valid dimensions (non-negative and less than rank of tensor).
  • Any src_dim or tgt_dim must be unique across all parameters.
  • src_dim must be sorted in ascending order across all parameters.
• Moving axes from src_dim to tgt_dim in the operand sharding gets out_sharding .

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.collective_permute (sdy::CollectivePermuteOp)

Performs a collective-permute communication to replace axes

نحو:

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

Sends a chunk of the input tensor from each device to another to reorder/replace the axes that shard the tensor.

A collective permute can transform the input sharding such that each dimension must be as sharded as it was before, ie, it must be sharded along axes whose product of sizes matches that of the axes that previously sharded the tensor.

This is useful for reordering axes in a single dimension or across different dimensions, and swapping sharded axes with replicated ones.

In the below example, the sharded tensor size is tensor<1x4x2xf32> , and that is preserved by the collective permute.

Example:

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>

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• If input and output sharding have different meshes, then those meshes must have exactly the same axes and different order of device ids.
• For each dimension, the product of sharding axis sizes in out_sharding must match that of the corresponding operand dimension sharding.

Traits: SameOperandsAndResultType

Interfaces: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.constant (sdy::ConstantOp)

Constant operation

Produces an output tensor from a constant value .

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

Example:

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

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable , InferTypeOpInterface , NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

نتایج:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Data flow edge op.

نحو:

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

A data flow edge of some op X defines a bridge between a set of sources (each is either an operand of X or an operand of X's block terminator) and a set of targets (each is either a result of X or a block argument of X), such that all sources and targets should be sharded in the same way.

An op can have multiple data flow edges that are orthogonal to one another.

For example:

  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
                  })

This while op has n data flow edges, the i-th data flow edges is between sources x_i , return_value_i and targets y_i , pred_arg_i , body_arg_i .

An sdy.data_flow_edge takes as input the owner of an edge (can be any of the targets, but preferably an op result rather than a block argument), which shouldn't have any other uses. This op isn't pure because it can take an input that originally didn't have any uses.

The sdy.data_flow_edge also holds an optional sharding for all targets of the edge, and that sharding should be updated instead of the targets' sharding (if can be attached) during propagation. This is useful when an op has many edges, as it's much more efficient to:

• propagate through each edge separately.
• update the sharding of each edge separately instead of all targets at once (eg an op has a single immutable TensorShardingPerValueAttr for result shardings).
• add each edge to the worklist separately when the sharding of a source has changed.

Propagation will propagate shardings between all sources and targets of a sdy.data_flow_edge as if it was a regular op with the sources as operands and targets as results, and an identity sdy.op_sharding_rule . That means that forward propagation is from sources to targets and backwards propagation is from targets to sources.

We don't allow the input of a sdy.data_flow_edge to be defined by an SdyDialect op, so we can assume that it's defined by an op that has unregistered sdy.sharding attribute.

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.func_data_flow_edge (sdy::FuncDataFlowEdgeOp)

Func input/output data flow edge op.

نحو:

operation ::= `sdy.func_data_flow_edge` $operand attr-dict `:` type($result)

A data flow edge op but for func arguments or call results. When its operand is a BlockArgument; it is a bridge from the caller callOp's argument to the users of the func argument. There is one func data flow edge for each func argument. When its operand is an OpResult; it is a bridge from the called funcOp's return value to the users of the call result. There is one func data flow edge for each call result.

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , SymbolUserOpInterface

Operands:

نتایج:

sdy.manual_computation (sdy::ManualComputationOp)

Multi-device parallelism operation with manual collectives

نحو:

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)

Jump into a region written in terms of per-device local code with explicit collectives, where logical shapes match local per-device physical buffer shapes and collectives correspond exactly to physical cross-device communication.

The body is local wrt the manual_axes. Propagation will occur through the body on any free axes - those not in the manual_axes list.

Note that any unranked tensors are expected to have a sharding with rank 0, ie fully replicated.

Constraints:

• Elements in in_shardings and out_shardings must satisfy the constraints listed in TensorShardingAttr .
• The number of global and local tensor inputs/outputs of the op region must match.
• The manual axes must come before any free axes in each dim sharding.
• The manual axes cannot introduce padding. Namely, the dimension size must be divisible by the corresponding manual axes size.
• The global and local shapes of the op regions arguments/results must match.

Traits: IsolatedFromAbove , RecursiveMemoryEffects , SingleBlockImplicitTerminator<ReturnOp> , SingleBlock

Interfaces: ShardableDataFlowOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.mesh (sdy::MeshOp)

Named mesh

نحو:

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

Defines a new named mesh. All meshes in a module must have the same number of devices (except for meshes with a single device_id). The mesh is a Symbol operation that appears in the module's SymbolTable and can be referenced by its name .

Traits: HasParent<ModuleOp>

Interfaces: Symbol

Attributes:

sdy.named_computation (sdy::NamedComputationOp)

Named computation operation

نحو:

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)

Groups a computation, ie a block of operations, and gives it a name. Propagation will flow in/out of the region as if everything was inlined.

This can be used to handle propagating through call instructions to other functions. Any users of Shardy should write an import/export pass that converts their call ops to sdy.named_computation ops, duplicating/copying the body of the called function into the body of the named_computation .

The type of each block arguments and returned values in the region must be the same as the type of the operands and results type of the op.

Example:

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

Traits: IsolatedFromAbove , RecursiveMemoryEffects , RecursivelySpeculatableImplTrait , SingleBlockImplicitTerminator<ReturnOp> , SingleBlock

Interfaces: ConditionallySpeculatable , InferTypeOpInterface , ShardableDataFlowOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Propagation barrier operation

نحو:

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

This op operates like an identity op, outputting the same value it took as input. But in terms of propagation, this will only allow propagation to flow through it in a certain direction.

This prevents shardings from being propagated between the uses of the result of the barrier op and its operand.

• FORWARD means shardings can only flow from the operand to the result.
• BACKWARD means shardings can only flow from the result to the operand.
• NONE means no sharding can propagate through this op.
• Cannot specify BOTH , as this op would be redundant.

Traits: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

Interfaces: ConditionallySpeculatable , InferTypeOpInterface , NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

نتایج:

sdy.reduce_scatter (sdy::ReduceScatterOp)

Performs a reduce-scatter communication along axes

نحو:

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

Reduces chunks of a tensor along axes specified in reduce_scatter_axes , and then scatters the result along the same axes. This operation is essentially a combination of an sdy.all_reduce followed by an sdy.all_slice along the same reduce_scatter_axes .

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• Elements in reduce_scatter_axes must satisfy the constraints listed in AxisRefListAttr .
• Applying reduce_scatter_axes to the operand sharding gets out_sharding .

Traits: SameOperandsAndResultType

Interfaces: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

Move implicitly or explicitly replicated axes to unreduced axes.

نحو:

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

The axes should be implicitly or explicitly replicated in the operand. This operation makes them unreduced in the result. We have the following relationship:

all-reduce(replicated-to-unreduced(x, axes), axes) = x

Example:

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

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• axes must satisfy the constraints listed in AxisRefListAttr .
• axes must be sorted wrt the mesh.
• axes are not empty.
• The input and output sharding must have the same dimension shardings.
• axes must be implicitly or explicitly replicated in the operand sharding.
• inUnreducedAxes + axes = outUnreducedAxes.

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.reshard (sdy::ReshardOp)

Reshards a tensor to a different sharding

نحو:

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

Reshards the input tensor with the specified sharding, which is different from the input tensor's existing sharding.

Both ShardingConstraintOp and ReshardOp attach a sharding to a tensor. Their lifespan is:

1. Before sharding propagation, ShardingConstraintOp is added by users.
2. Sharding propagation consumes ShardingConstraintOp. There is no ShardingConstraintOp in the results of sharding propagation. Instead, ReshardOp may be added if needed.
3. A partitioner converts a ReshardOp into a collective op (or an identity op). There should be no ReshardOp in the results of the partitioner.

Traits: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

Interfaces: ConditionallySpeculatable , InferTypeOpInterface , NoMemoryEffect (MemoryEffectOpInterface) , SymbolUserOpInterface

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

نتایج:

sdy.return (sdy::ReturnOp)

The sdy.return operation terminates the regions attached to sdy region-based ops and any other Shardy region-based ops. It is variadic: it takes as arguments a list of values whose types can be any (but of the same kind, eg AnyTensor ) and therefore can be reused at various levels of the Shardy IR stack.

نحو:

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

Traits: AlwaysSpeculatableImplTrait , ReturnLike , Terminator

Interfaces: ConditionallySpeculatable , NoMemoryEffect (MemoryEffectOpInterface) , RegionBranchTerminatorOpInterface

Effects: MemoryEffects::Effect{}

Operands:

sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

Move some sharded axes of the operand to unreduced axes of the result.

نحو:

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

The axes should be used to shard the operand. This operation makes them unreduced in the result. We have the following relationship:

all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes), where all-gather, sharded-to-unreduced, all-reduce are applied on the same axes.

Example:

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

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• Elements in axes must satisfy the constraints listed in AxisRefListAttr .
• Applying axes to the operand sharding gets out_sharding .

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Constrains a tensor to the specified sharding

نحو:

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

Attaches a sharding to an intermediate tensor (eg the result of a matmul) to indicate that this is how that tensor, or a subset of its uses, should be sharded.

If the sharding has open dimensions and unconstraint axes, it means the tensor can be further sharded along the open dimensions.

This op can either:

• Have no uses (dangling) - which means the attached sharding is how the input tensor itself should be sharded.
• Have uses - which means the attached sharding is how the uses of the sharding constraint op should be sharded, while other uses of the input tensor might have a different sharding (if the input tensor has no other uses then the behavior is the same as the no uses case).

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.sharding_group (sdy::ShardingGroupOp)

Constrains tensors in the group to have the same sharding.

نحو:

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

This op provides an interface to assign tensors to sharding groups ( groups of tensors that will be enforced to have identical shardings). During propagation, as soon as one group element is sharded, all other members will be sharded in exactly the same way. This operation takes the argument group ID and returns no result, but instead modifies the internal sharding group representation to add the input tensor to the group with the given ID.

Interfaces: InferTypeOpInterface

Attributes:

Operands:

ویژگی‌ها

AllToAllParamAttr

All-to-all parameter

نحو:

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

A tuple containing the axes and source/target dimensions to perform all-to-all on.

Parameters:

AllToAllParamListAttr

List of all-to-all parameters

نحو:

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

Parameters:

AxisRefAttr

Reference to either a full axis or a split sub-axis

نحو:

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

Constraints:

• name must be present in the bound MeshAttr .
• If sub_axis_info is present, it must satisfy the constraints of SubAxisInfoAttr .

Parameters:

AxisRefListAttr

List of axis refs

نحو:

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

Constraints:

• Elements in value must satisfy the constraints of AxisRefAttr .
• There are no duplicate axis-refs or sub-axes that overlap with one another.
• No two adjacent axis-refs are consecutive sub-axes of that same full axis, ie, they can be merged into one sub-axis or the full axis.

Parameters:

AxisToPropagationDetailsAttr

Propagation edge flow details for a specific axis and source.

نحو:

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

Maps a source value reference to a list of target value references along a particular axis.

Parameters:

DimMappingAttr

List of factor indices for a dimension

An empty list indicates that this is a null mapping (this is parsed/printed with * ), ie the dimension isn't mapped to any factors.

Constraints:

• There is at least one factor index.
• Factor indices must be in range [0, $factor_sizes ).
• If there are multiple factors, none of them can have size 1.
• No duplicate factor indices.

Parameters:

DimensionShardingAttr

Dimension sharding

List of axis names to shard a tensor dimension on from major to minor, a boolean indicating whether the dimension can be further sharded, and an optional integer denoting the priority of this dimension sharding, which will respected during sharding propagation. Priorities originate from user sharding annotations and a lower value denotes a higher priority. The highest priority is assumed when the priority is missing in the annotation.

Constraints:

• Elements in axes must satisfy the constraints listed in AxisRefListAttr .
• If a dimension sharding has a priority:
  • The priority is greater than or equal to 0.
  • The dimension has at least one axis if it is closed.

Parameters:

EdgeValueRefAttr

Reference to a particular index of a value edge of type type .

نحو:

#sdy.edge_value_ref<
  `operand` | `result`,   # type
  int64_t   # index
>

Parameters:

ListOfAxisRefListsAttr

List of axis ref lists

نحو:

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

Parameters:

ManualAxesAttr

A list of axes that a ManualComputationOp is manual on

نحو:

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

Parameters:

MeshAttr

Mesh of axes and a list of devices

نحو:

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

A mesh is a list of axes and an optional list of device IDs specifying the device ordering.

If the list of axes is empty

• If the device_ids is not provided, it is an empty mesh.
• If the device_ids is provided, it must be a single non-negative integer, we call it a maximal-sharding mesh .

If the list of axes is provided

• If a device ID list is specified, the product of the axis sizes should match the number of devices.
• If a device ID list is not specified, the implicit device ID list is iota(product(axes)). For simplicity, we also disallow specifying a device ID list that is the same as iota(product(axes)); in this case, a device ID list shouldn't be specified.
• It is not a maximal-sharding mesh even if the total size of axes is 1.

Here are some examples of meshes:

• An empty mesh represents a placeholder mesh that can be replaced during propagation: <[]>
• A mesh without axes list and a single non-negative device ID, which is a maximal-sharding mesh: <[], device_ids=[3]>
• A mesh with two axes and implicit device IDs iota(6): <["a"=2, "b"=3]>
• A mesh with two axes and explicit device IDs specifying the device ordering: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Constraints:

• Elements in device_ids should be non-negative.
• If axes is empty, the size of device_ids can be 0 (empty mesh) or 1 (maximal-sharding mesh).
• If axes is not empty,
  • Elements in axes must not have duplicate names.
  • If device_ids is specified, the original device_ids is not iota(product(axis_sizes)) and the sorted device_ids is iota(product(axis_sizes)) .

Parameters:

MeshAxisAttr

Named axis in a mesh

نحو:

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

Parameters:

OpShardingRuleAttr

Specifies how an operation can be partitioned.

نحو:

#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
>

A sharding rule specifies how an operation can be partitioned according to various properties on the op - any attributes, the shape of operands, the shape of the results, etc. For example:

%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>

Note that we allow factors with size 1 even though they cannot be sharded, this is mainly for completeness as many ops such as pointwise ops have size one dimensions that correspond across operands and results.

Factor types:

• reduction_factors contains the indices of factors requiring reduction, such as the contracting dimensions in a dot operation. These factors can be in operands but not in results.
• need_replication_factors contains the indices of factors requiring full replication, such as the sorted dimension in a sort operation.
• permutation_factors contains the indices of factors requiring collective-permute if they are sharded, such as the padding dimensions in a pad operation.
• All other factors are considered as pass-through factors, ie, factors that don't require any communication if sharded in the same way across all tensors that are mapped to them.

blocked_propagation_factors contains the factors along which shardings are not allowed to be propagated. It is orthogonal to the factor types. Namely, a blocked-propagation factor can be any of the factor types.

is_custom_rule describes whether this is a rule defined by a user. Users can define sharding rules for their custom calls or overwrite the pre-defined sharding rules for the standard operations. A custom rule is always preserved/never removed.

Constraints:

• Number of operand/result mappings must match the number of operands/results of the op.
• There is at least one mapping (can't have a rule for an op with no operands/results).
• Rank of each TensorMappingAttr matches the rank of the corresponding tensor type.
• For each group of factors ( reduction_factors , need_replication_factors , permutation_factors ):
  • Elements must be in range [0, $factor_sizes ].
  • No duplicate factor indices within each group and across groups.

Parameters:

PropagationEdgesAttr

Propagation edge metadata for all propagation steps.

نحو:

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

A list of per-axis propagation details for a value, grouped by step index.

Parameters:

PropagationOneStepAttr

Per-step propagation metadata.

نحو:

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

Propagation details for all axes for a single propagation step.

Parameters:

SubAxisInfoAttr

Info about how this sub-axis is derived from the full axis

نحو:

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

When splitting a full axis into n sub-axes, the axis is reshaped into [k_1,...,k_n], and the ith sub-axis can be expressed by the product of all axis sizes to its left m=prod(k_1,...,k_(i-1)) (aka pre-size) and size k_i. Therefore, the sub-axis-info attribute holds those two numbers and is denoted as follows: (m)k for pre-size m and size k.

Constraints:

• pre-size is at least 1.
• size is greater than 1.
• pre-size must divide the size of the full axis, ie, both pre-size and size divide the size of the full axis, and the sub-axis doesn't go beyond the full axis.
• The size of the sub-axis isn't equal to the size of the corresponding full axis, in which case the full axis should be used instead.

Parameters:

TensorMappingAttr

Factor mappings for each dimension of a tensor.

نحو:

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

Constraints:

• Elements in dim_mappings must satisfy the constraints in DimMappingAttr .
• No duplicate factors indices across dimensions.

Parameters:

TensorShardingAttr

Tensor sharding

نحو:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>,   # replicated_axes
  ::llvm::ArrayRef<AxisRefAttr>,   # unreduced_axes
  `sum` | `max` | `min`   # reduction_op
>

A tensor sharding is bound to a specific mesh, and can only reference axis names from that mesh. The dimension shardings tell us for each dimension of the tensor, along which axes (or sub-axes) it is sharded from major to minor. All other axes that don't shard a dimension are either implicitly or explicitly (if they appear in the list of replicated axes) replicated.

Note that no sharding attribute on a tensor is equivalent to a fully open tensor sharding.

The mesh this sharding is bound to can either be specified by a symbol name, referencing a corresponding MeshOp symbol, or an inlined MeshAttr .

A sharding can have unreduced axes (specified by unreduced_axes ), meaning the tensor is unreduced along these axes. For example, if the contracting dimension of a matmul is sharded along axis x in both the lhs and rhs, the result is unreduced along x . Applying an all-reduce on the tensor along the unreduced axes will make the tensor replicated along those axes. However, a tensor with unreduced axes doesn't have to be all-reduced immediately, it can remain unreduced when passed to linear operations like stablehlo.add (as long as both lhs and rhs are unreduced) and all-reduced afterwards. We assume the reduction type is sum, other reductions may be supported in the future.

Constraints:

• Elements in dim_shardings must satisfy the constraints listed in DimensionShardingAttr .
• Elements in replicated_axes must satisfy the constraints listed in AxisRefListAttr .
• Elements in unreduced_axes must satisfy the constraints listed in AxisRefListAttr .
• If the corresponding tensor type isn't a ShapedType , the sharding must have rank 0 and no replicated axes.
• If it is a ShapedType , then:
  • The tensor should have a rank.
  • The number of dimension shardings is equal to the rank of the tensor.
  • Dimensions of size 0 aren't sharded.
• There are no duplicate axis-refs or sub-axes that overlap with one another across dim_shardings , replicated_axes , and unreduced_axes .
• Items in replicated_axes and unreduced_axes are ordered wrt mesh_or_ref (see AxisRefAttr::getMeshComparator ).

Parameters:

TensorShardingPerValueAttr

Tensor sharding per operand/result of an op

نحو:

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

A list of TensorShardingAttr s, one for each operand/result of an op.

Constraints:

• Elements in shardings must satisfy the constraints of TensorShardingAttr .

Parameters:

انوم‌ها

EdgeNodeType

Edge node type enum

Cases:

PropagationDirection

Propagation direction enum

Cases:

ReductionOp

Reduction op enum

Cases:

The Shardy (SDY) dialect

The Shardy (SDY) dialect defines an axis-based tensor sharding representation and additional API components to attach shardings to tensors.

Version log: 0.0.1: Add unreduced axes to TensorShardingAttr.

عملیات

sdy.all_gather (sdy::AllGatherOp)

Performs an all-gather communication along axes

نحو:

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

Gathers chunks of a tensor along axes specified in gathering_axes .

The gathering_axes is a list of lists of axes. The outer list is over the dimensions of the tensor. Each inner list specifies the axes along which a separate gather should be performed on the respective dimension. It will be applied to the sharding of the operand ( tensor ) to obtain the sharding of the result ( out_sharding ).

Note that out_sharding is not used to determine the sharding of the result. Instead, the sharding of the result is determined by the sharding of the operand and the gathering_axes , and out_sharding must match this inferred sharding.

Example:

%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>

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• Elements in gathering_axes must satisfy the constraints listed in AxisRefListAttr .
• Applying gathering_axes to the operand sharding gets out_sharding .

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.all_reduce (sdy::AllReduceOp)

Perform an all-reduce comunication along axes

نحو:

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

Reduces chunks of a tensor along axes specified in reduction_axes . The order of reduction_axes is not important for the result, but can affect the order of the corresponding replica groups.

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• reduction_axes must satisfy the constraints listed in AxisRefListAttr .
• reduction_axes must be sorted wrt the mesh.
• The operand sharding and out_sharding must have equivalent dimension shardings.
• reduction_axes must not overlap with the operand dimension sharding and replicated axes (it can overlap with unreduced axes).
• reduction_axes must not overlap with the unreduced axes of out_sharding . In other words, out_sharding must be be replicated along reduction_axes (implicitly or explicitly).

Traits: SameOperandsAndResultType

Interfaces: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.all_slice (sdy::AllSliceOp)

Performs a dynamic-slice operation along axes

نحو:

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

Slices chunks of a tensor along axes specified in slicing_axes . There is an algebric duality between sdy.all_slice and sdy.all_gather .

The slicing_axes is a list of lists of axes. The outer list is over the dimensions of the tensor. Each inner list specifies the axes along which a slice should be performed on the respective dimension. It will be applied to the sharding of the operand ( tensor ) to obtain the sharding of the result ( out_sharding ).

Note that out_sharding is not used to determine the sharding of the result. Instead, the sharding of the result is determined by the sharding of the operand and the slicing_axes , and out_sharding must match this inferred sharding.

Example:

%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>

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• Elements in slicing_axes must satisfy the constraints listed in AxisRefListAttr .
• Applying slicing_axes to the operand sharding gets out_sharding .

Traits: SameOperandsAndResultType

Interfaces: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.all_to_all (sdy::AllToAllOp)

Performs an all-to-all communication along axes

نحو:

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

For each (axes, src_dim, tgt_dim) tuple in the parameter list, this operation slices chunks of a tensor along dimension tgt_dim and axes specified in axes , scatteres those chunks along the axes, and concatenates them along dimension src_dim .

This operation is essentially a combination of an all-gather along src_dim and axes , followed by an all-slice along tgt_dim and axes , ie, a suffix of the axes sharding dimension src_dim on the input tensor is appended to the axes sharding dimension tgt_dim on the output tensor.

The all-to-all will be applied to the sharding of the operand ( tensor ) to obtain the sharding of the result ( out_sharding ).

Note that out_sharding is not used to determine the sharding of the result. Instead, the sharding of the result is determined by the sharding of the operand, src_dim , tgt_dim , and axes , and out_sharding must match this inferred sharding.

Example:

%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>

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• The parameter list must not be empty.
• For each parameter in params :
  • Elements in axes must satisfy the constraints of AxisRefAttr .
  • src_dim and tgt_dim must be valid dimensions (non-negative and less than rank of tensor).
  • Any src_dim or tgt_dim must be unique across all parameters.
  • src_dim must be sorted in ascending order across all parameters.
• Moving axes from src_dim to tgt_dim in the operand sharding gets out_sharding .

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.collective_permute (sdy::CollectivePermuteOp)

Performs a collective-permute communication to replace axes

نحو:

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

Sends a chunk of the input tensor from each device to another to reorder/replace the axes that shard the tensor.

A collective permute can transform the input sharding such that each dimension must be as sharded as it was before, ie, it must be sharded along axes whose product of sizes matches that of the axes that previously sharded the tensor.

This is useful for reordering axes in a single dimension or across different dimensions, and swapping sharded axes with replicated ones.

In the below example, the sharded tensor size is tensor<1x4x2xf32> , and that is preserved by the collective permute.

Example:

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>

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• If input and output sharding have different meshes, then those meshes must have exactly the same axes and different order of device ids.
• For each dimension, the product of sharding axis sizes in out_sharding must match that of the corresponding operand dimension sharding.

Traits: SameOperandsAndResultType

Interfaces: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.constant (sdy::ConstantOp)

Constant operation

Produces an output tensor from a constant value .

See: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

Example:

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

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable , InferTypeOpInterface , NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

نتایج:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Data flow edge op.

نحو:

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

A data flow edge of some op X defines a bridge between a set of sources (each is either an operand of X or an operand of X's block terminator) and a set of targets (each is either a result of X or a block argument of X), such that all sources and targets should be sharded in the same way.

An op can have multiple data flow edges that are orthogonal to one another.

For example:

  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
                  })

This while op has n data flow edges, the i-th data flow edges is between sources x_i , return_value_i and targets y_i , pred_arg_i , body_arg_i .

An sdy.data_flow_edge takes as input the owner of an edge (can be any of the targets, but preferably an op result rather than a block argument), which shouldn't have any other uses. This op isn't pure because it can take an input that originally didn't have any uses.

The sdy.data_flow_edge also holds an optional sharding for all targets of the edge, and that sharding should be updated instead of the targets' sharding (if can be attached) during propagation. This is useful when an op has many edges, as it's much more efficient to:

• propagate through each edge separately.
• update the sharding of each edge separately instead of all targets at once (eg an op has a single immutable TensorShardingPerValueAttr for result shardings).
• add each edge to the worklist separately when the sharding of a source has changed.

Propagation will propagate shardings between all sources and targets of a sdy.data_flow_edge as if it was a regular op with the sources as operands and targets as results, and an identity sdy.op_sharding_rule . That means that forward propagation is from sources to targets and backwards propagation is from targets to sources.

We don't allow the input of a sdy.data_flow_edge to be defined by an SdyDialect op, so we can assume that it's defined by an op that has unregistered sdy.sharding attribute.

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.func_data_flow_edge (sdy::FuncDataFlowEdgeOp)

Func input/output data flow edge op.

نحو:

operation ::= `sdy.func_data_flow_edge` $operand attr-dict `:` type($result)

A data flow edge op but for func arguments or call results. When its operand is a BlockArgument; it is a bridge from the caller callOp's argument to the users of the func argument. There is one func data flow edge for each func argument. When its operand is an OpResult; it is a bridge from the called funcOp's return value to the users of the call result. There is one func data flow edge for each call result.

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , SymbolUserOpInterface

Operands:

نتایج:

sdy.manual_computation (sdy::ManualComputationOp)

Multi-device parallelism operation with manual collectives

نحو:

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)

Jump into a region written in terms of per-device local code with explicit collectives, where logical shapes match local per-device physical buffer shapes and collectives correspond exactly to physical cross-device communication.

The body is local wrt the manual_axes. Propagation will occur through the body on any free axes - those not in the manual_axes list.

Note that any unranked tensors are expected to have a sharding with rank 0, ie fully replicated.

Constraints:

• Elements in in_shardings and out_shardings must satisfy the constraints listed in TensorShardingAttr .
• The number of global and local tensor inputs/outputs of the op region must match.
• The manual axes must come before any free axes in each dim sharding.
• The manual axes cannot introduce padding. Namely, the dimension size must be divisible by the corresponding manual axes size.
• The global and local shapes of the op regions arguments/results must match.

Traits: IsolatedFromAbove , RecursiveMemoryEffects , SingleBlockImplicitTerminator<ReturnOp> , SingleBlock

Interfaces: ShardableDataFlowOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.mesh (sdy::MeshOp)

Named mesh

نحو:

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

Defines a new named mesh. All meshes in a module must have the same number of devices (except for meshes with a single device_id). The mesh is a Symbol operation that appears in the module's SymbolTable and can be referenced by its name .

Traits: HasParent<ModuleOp>

Interfaces: Symbol

Attributes:

sdy.named_computation (sdy::NamedComputationOp)

Named computation operation

نحو:

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)

Groups a computation, ie a block of operations, and gives it a name. Propagation will flow in/out of the region as if everything was inlined.

This can be used to handle propagating through call instructions to other functions. Any users of Shardy should write an import/export pass that converts their call ops to sdy.named_computation ops, duplicating/copying the body of the called function into the body of the named_computation .

The type of each block arguments and returned values in the region must be the same as the type of the operands and results type of the op.

Example:

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

Traits: IsolatedFromAbove , RecursiveMemoryEffects , RecursivelySpeculatableImplTrait , SingleBlockImplicitTerminator<ReturnOp> , SingleBlock

Interfaces: ConditionallySpeculatable , InferTypeOpInterface , ShardableDataFlowOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Propagation barrier operation

نحو:

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

This op operates like an identity op, outputting the same value it took as input. But in terms of propagation, this will only allow propagation to flow through it in a certain direction.

This prevents shardings from being propagated between the uses of the result of the barrier op and its operand.

• FORWARD means shardings can only flow from the operand to the result.
• BACKWARD means shardings can only flow from the result to the operand.
• NONE means no sharding can propagate through this op.
• Cannot specify BOTH , as this op would be redundant.

Traits: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

Interfaces: ConditionallySpeculatable , InferTypeOpInterface , NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

نتایج:

sdy.reduce_scatter (sdy::ReduceScatterOp)

Performs a reduce-scatter communication along axes

نحو:

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

Reduces chunks of a tensor along axes specified in reduce_scatter_axes , and then scatters the result along the same axes. This operation is essentially a combination of an sdy.all_reduce followed by an sdy.all_slice along the same reduce_scatter_axes .

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• Elements in reduce_scatter_axes must satisfy the constraints listed in AxisRefListAttr .
• Applying reduce_scatter_axes to the operand sharding gets out_sharding .

Traits: SameOperandsAndResultType

Interfaces: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

Move implicitly or explicitly replicated axes to unreduced axes.

نحو:

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

The axes should be implicitly or explicitly replicated in the operand. This operation makes them unreduced in the result. We have the following relationship:

all-reduce(replicated-to-unreduced(x, axes), axes) = x

Example:

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

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• axes must satisfy the constraints listed in AxisRefListAttr .
• axes must be sorted wrt the mesh.
• axes are not empty.
• The input and output sharding must have the same dimension shardings.
• axes must be implicitly or explicitly replicated in the operand sharding.
• inUnreducedAxes + axes = outUnreducedAxes.

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.reshard (sdy::ReshardOp)

Reshards a tensor to a different sharding

نحو:

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

Reshards the input tensor with the specified sharding, which is different from the input tensor's existing sharding.

Both ShardingConstraintOp and ReshardOp attach a sharding to a tensor. Their lifespan is:

1. Before sharding propagation, ShardingConstraintOp is added by users.
2. Sharding propagation consumes ShardingConstraintOp. There is no ShardingConstraintOp in the results of sharding propagation. Instead, ReshardOp may be added if needed.
3. A partitioner converts a ReshardOp into a collective op (or an identity op). There should be no ReshardOp in the results of the partitioner.

Traits: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

Interfaces: ConditionallySpeculatable , InferTypeOpInterface , NoMemoryEffect (MemoryEffectOpInterface) , SymbolUserOpInterface

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

نتایج:

sdy.return (sdy::ReturnOp)

The sdy.return operation terminates the regions attached to sdy region-based ops and any other Shardy region-based ops. It is variadic: it takes as arguments a list of values whose types can be any (but of the same kind, eg AnyTensor ) and therefore can be reused at various levels of the Shardy IR stack.

نحو:

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

Traits: AlwaysSpeculatableImplTrait , ReturnLike , Terminator

Interfaces: ConditionallySpeculatable , NoMemoryEffect (MemoryEffectOpInterface) , RegionBranchTerminatorOpInterface

Effects: MemoryEffects::Effect{}

Operands:

sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

Move some sharded axes of the operand to unreduced axes of the result.

نحو:

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

The axes should be used to shard the operand. This operation makes them unreduced in the result. We have the following relationship:

all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes), where all-gather, sharded-to-unreduced, all-reduce are applied on the same axes.

Example:

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

Constraints:

• Must satisfy the constraints listed in Sdy_CollectiveOpInterface .
• Elements in axes must satisfy the constraints listed in AxisRefListAttr .
• Applying axes to the operand sharding gets out_sharding .

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Constrains a tensor to the specified sharding

نحو:

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

Attaches a sharding to an intermediate tensor (eg the result of a matmul) to indicate that this is how that tensor, or a subset of its uses, should be sharded.

If the sharding has open dimensions and unconstraint axes, it means the tensor can be further sharded along the open dimensions.

This op can either:

• Have no uses (dangling) - which means the attached sharding is how the input tensor itself should be sharded.
• Have uses - which means the attached sharding is how the uses of the sharding constraint op should be sharded, while other uses of the input tensor might have a different sharding (if the input tensor has no other uses then the behavior is the same as the no uses case).

Traits: SameOperandsAndResultType

Interfaces: InferTypeOpInterface , SymbolUserOpInterface

Attributes:

Operands:

نتایج:

sdy.sharding_group (sdy::ShardingGroupOp)

Constrains tensors in the group to have the same sharding.

نحو:

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

This op provides an interface to assign tensors to sharding groups ( groups of tensors that will be enforced to have identical shardings). During propagation, as soon as one group element is sharded, all other members will be sharded in exactly the same way. This operation takes the argument group ID and returns no result, but instead modifies the internal sharding group representation to add the input tensor to the group with the given ID.

Interfaces: InferTypeOpInterface

Attributes:

Operands:

ویژگی‌ها

AllToAllParamAttr

All-to-all parameter

نحو:

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

A tuple containing the axes and source/target dimensions to perform all-to-all on.

Parameters:

AllToAllParamListAttr

List of all-to-all parameters

نحو:

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

Parameters:

AxisRefAttr

Reference to either a full axis or a split sub-axis

نحو:

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

Constraints:

• name must be present in the bound MeshAttr .
• If sub_axis_info is present, it must satisfy the constraints of SubAxisInfoAttr .

Parameters:

AxisRefListAttr

List of axis refs

نحو:

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

Constraints:

• Elements in value must satisfy the constraints of AxisRefAttr .
• There are no duplicate axis-refs or sub-axes that overlap with one another.
• No two adjacent axis-refs are consecutive sub-axes of that same full axis, ie, they can be merged into one sub-axis or the full axis.

Parameters:

AxisToPropagationDetailsAttr

Propagation edge flow details for a specific axis and source.

نحو:

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

Maps a source value reference to a list of target value references along a particular axis.

Parameters:

DimMappingAttr

List of factor indices for a dimension

An empty list indicates that this is a null mapping (this is parsed/printed with * ), ie the dimension isn't mapped to any factors.

Constraints:

• There is at least one factor index.
• Factor indices must be in range [0, $factor_sizes ).
• If there are multiple factors, none of them can have size 1.
• No duplicate factor indices.

Parameters:

DimensionShardingAttr

Dimension sharding

List of axis names to shard a tensor dimension on from major to minor, a boolean indicating whether the dimension can be further sharded, and an optional integer denoting the priority of this dimension sharding, which will respected during sharding propagation. Priorities originate from user sharding annotations and a lower value denotes a higher priority. The highest priority is assumed when the priority is missing in the annotation.

Constraints:

• Elements in axes must satisfy the constraints listed in AxisRefListAttr .
• If a dimension sharding has a priority:
  • The priority is greater than or equal to 0.
  • The dimension has at least one axis if it is closed.

Parameters:

EdgeValueRefAttr

Reference to a particular index of a value edge of type type .

نحو:

#sdy.edge_value_ref<
  `operand` | `result`,   # type
  int64_t   # index
>

Parameters:

ListOfAxisRefListsAttr

List of axis ref lists

نحو:

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

Parameters:

ManualAxesAttr

A list of axes that a ManualComputationOp is manual on

نحو:

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

Parameters:

MeshAttr

Mesh of axes and a list of devices

نحو:

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

A mesh is a list of axes and an optional list of device IDs specifying the device ordering.

If the list of axes is empty

• If the device_ids is not provided, it is an empty mesh.
• If the device_ids is provided, it must be a single non-negative integer, we call it a maximal-sharding mesh .

If the list of axes is provided

• If a device ID list is specified, the product of the axis sizes should match the number of devices.
• If a device ID list is not specified, the implicit device ID list is iota(product(axes)). For simplicity, we also disallow specifying a device ID list that is the same as iota(product(axes)); in this case, a device ID list shouldn't be specified.
• It is not a maximal-sharding mesh even if the total size of axes is 1.

Here are some examples of meshes:

• An empty mesh represents a placeholder mesh that can be replaced during propagation: <[]>
• A mesh without axes list and a single non-negative device ID, which is a maximal-sharding mesh: <[], device_ids=[3]>
• A mesh with two axes and implicit device IDs iota(6): <["a"=2, "b"=3]>
• A mesh with two axes and explicit device IDs specifying the device ordering: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Constraints:

• Elements in device_ids should be non-negative.
• If axes is empty, the size of device_ids can be 0 (empty mesh) or 1 (maximal-sharding mesh).
• If axes is not empty,
  • Elements in axes must not have duplicate names.
  • If device_ids is specified, the original device_ids is not iota(product(axis_sizes)) and the sorted device_ids is iota(product(axis_sizes)) .

Parameters:

MeshAxisAttr

Named axis in a mesh

نحو:

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

Parameters:

OpShardingRuleAttr

Specifies how an operation can be partitioned.

نحو:

#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
>