'sdy' گویش

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

عملیات

sdy.constant (sdy::ConstantOp)

عملیات ثابت

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

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

مثال:

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

صفات: AlwaysSpeculatableImplTrait

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

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

ویژگی ها:

نتایج:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

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

نحو:

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

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

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

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

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

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

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

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

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

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

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

صفات: SameOperandsAndResultType

رابط ها: InferTypeOpInterface

ویژگی ها:

عملگرها:

نتایج:

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 نیستند - رخ می دهد.

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

ویژگی ها:

عملگرها:

نتایج:

sdy.mesh (sdy::MeshOp)

مش به نام

نحو:

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

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

صفات: HasParent<ModuleOp>

رابط ها: Symbol

ویژگی ها:

sdy.named_computation (sdy::NamedComputationOp)

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

نحو:

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

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

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

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

مثال:

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

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

رابط ها: ConditionallySpeculatable ، ShardableDataFlowOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

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

نحو:

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

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

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

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

صفات: AlwaysSpeculatableImplTrait ، Elementwise ، SameOperandsAndResultType

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

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

ویژگی ها:

عملگرها:

نتایج:

sdy.reshard (sdy::ReshardOp)

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

نحو:

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

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

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

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

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

صفات: AlwaysSpeculatableImplTrait ، Elementwise ، SameOperandsAndResultType

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

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

ویژگی ها:

عملگرها:

نتایج:

sdy.return (sdy::ReturnOp)

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

نحو:

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

ویژگی ها: AlwaysSpeculatableImplTrait , Terminator

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

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

عملگرها:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

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

نحو:

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

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

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

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

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

صفات: Elementwise ، SameOperandsAndResultType

رابط ها: InferTypeOpInterface

ویژگی ها:

عملگرها:

نتایج:

sdy.sharding_group (sdy::ShardingGroupOp)

عملیات گروهی شاردینگ

نحو:

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

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

ویژگی ها:

عملگرها:

صفات

AxisRefAttr

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

نحو:

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

پارامترها:

DimMappingAttr

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

همه شاخص‌های فاکتور باید در محدوده [0، num_factors) باشند و یک لیست خالی نشان می‌دهد که این یک نگاشت تهی است (این با * تجزیه/چاپ می‌شود)، یعنی بعد به هیچ عاملی نگاشت نشده است.

پارامترها:

DimensionShardingAttr

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

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

پارامترها:

ManualAxesAttr

نحو:

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

پارامترها:

MeshAttr

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

نحو:

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

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

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

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

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

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

پارامترها:

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

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

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

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

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

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

پارامترها:

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)) بیان کرد. m=prod(k_1,...,k_(i-1)) (معروف به پیش اندازه) و اندازه k_i. بنابراین، ویژگی sub-axis-info آن دو عدد را نگه می‌دارد و به صورت زیر نشان داده می‌شود: (m)k برای پیش اندازه m و اندازه k.

پارامترها:

TensorMappingAttr

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

نحو:

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

پارامترها:

TensorShardingAttr

تقسیم تانسور

نحو:

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

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

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

پارامترها:

TensorShardingPerValueAttr

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

نحو:

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

پارامترها:

Enums

جهت انتشار

جهت انتشار enum

موارد: