تعاریف OpInterface

ShardableDataFlowOpInterface ( ShardableDataFlowOpInterface )

یک رابط عملیاتی که به Shardy اجازه می‌دهد تا شاردینگ‌ها را از طریق لبه‌های جریان داده عملیات‌هایی که این رابط را گسترش می‌دهند، منتشر کند.

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

مالک یک هدف مشخص شده توسط کاربر از لبه جریان داده است که توسط انتشار Shardy استفاده می شود. کاربر می تواند آن را خودسرانه انتخاب کند اما باید ثابت باشد.

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

  y_1, ..., y_n = custom_op (x_1, ..., x_n)
                  ((body_arg_1,..., body_arg_n) {
                    ...
                    return return_value_1, ..., return_value_n
                  })

این custom_op دو نوع برای لبه‌های جریان داده دارد، n یال هر کدام بین return_value_i (منبع) و y_i (هدف‌ها) و n یال بین x_i (منبع) و body_arg_i (هدف‌ها). در این مورد صاحبان لبه همان اهداف هستند.

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

  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 است.

روش ها:

getBlockArgumentEdgeOwnerShardings

mlir::ArrayRef<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

اشتراک گذاری های همه صاحبان لبه جریان داده آرگومان بلوکی را برمی گرداند.

setBlockArgumentEdgeOwnerSharding

void setBlockArgumentEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);

sharding مالک لبه آرگومان بلوک را با index داده شده تنظیم می کند.

setBlockArgumentEdgeOwnerShardings

void setBlockArgumentEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

shardings همه صاحبان لبه آرگومان بلوک را تنظیم می کند.

getOpResultEdgeOwnerShardings

mlir::ArrayRef<mlir::sdy::TensorShardingAttr> getOpResultEdgeOwnerShardings();

اشتراک گذاری همه صاحبان لبه جریان داده نتیجه عملیات را برمی گرداند.

setOpResultEdgeOwnerSharding

void setOpResultEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);

sharding مالک لبه نتیجه عملیات را با index داده شده تنظیم می کند.

setOpResultEdgeOwnerShardings

void setOpResultEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

shardings همه صاحبان لبه نتیجه عملیات را تنظیم می کند.

getBlockArgumentEdgeOwners

mlir::ArrayRef<mlir::BlockArgument> getBlockArgumentEdgeOwners();

همه صاحبان لبه آرگومان بلوک را دریافت می کند.

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

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

getEdgeSources

mlir::SmallVector<mlir::Value> getEdgeSources(mlir::Value target);

منابع لبه جریان داده را با یک مقدار target دریافت می کند.

getEdgeOwnerFromTarget

mlir::Value getEdgeOwnerFromTarget(mlir::Value target);

با توجه به target که ممکن است مالک باشد یا نباشد، target مالک لبه جریان داده را دریافت می کند.

getEdgeOwnerFromSource

mlir::Value getEdgeOwnerFromSource(mlir::OpOperand&source);

هدف مالک لبه جریان داده را با یک source دریافت می کند.

getNonEdgeOwnerTargets

mlir::SmallVector<mlir::Value> getNonEdgeOwnerTargets(mlir::Value owner);

با توجه به owner لبه، اهداف غیر مالک لبه جریان داده را دریافت می کند.

ShardingRuleOpInterface ( ShardingRuleOpInterface )

یک رابط op که به op اجازه می دهد تا قانون اشتراک گذاری خود را تعریف کند. یک قانون تقسیم بندی مشخص می کند که چگونه یک عملیات را می توان بر اساس ویژگی های مختلف در عملیات تقسیم بندی کرد - هر ویژگی، شکل عملوندها، شکل نتایج و غیره. برای جزئیات بیشتر به OpShardingRuleAttr مراجعه کنید.

روش ها:

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

قانون اشتراک گذاری عملیات را برمی گرداند.