'sdy' lehçesi

Shardy (SDY) lehçesi

Shardy (SDY) lehçesi, eksen tabanlı tensör parçalama gösterimini ve parçalamaları tensörlere eklemek için ek API bileşenlerini tanımlar.

Sürüm günlüğü: 0.0.1: TensorShardingAttr'ye azaltılmamış eksenler ekleyin.

İşlemler

sdy.all_gather (sdy::AllGatherOp)

Eksenler boyunca tüm birleştirme iletişimi gerçekleştirir

Söz dizimi:

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

Bir tensörün parçalarını gathering_axes içinde belirtilen eksenler boyunca toplar.

gathering_axes, eksen listelerinin listesidir. Dış liste, tensörün boyutlarını aşıyor. Her iç liste, ilgili boyutta ayrı bir toplama işleminin gerçekleştirilmesi gereken eksenleri belirtir. Sonucun (out_sharding) parçalanmasını elde etmek için işlenenin (tensor) parçalanmasına uygulanır.

out_sharding, sonucun parçalanmasını belirlemek için kullanılmaz. Bunun yerine, sonucun parçalama işlemi işlenenin ve gathering_axes parçalama işlemine göre belirlenir ve out_sharding bu çıkarılan parçalama işlemiyle eşleşmelidir.

Örnek:

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

Kısıtlamalar:

  • Sdy_CollectiveOpInterface bölümünde listelenen kısıtlamaları karşılamalıdır.
  • gathering_axes içindeki öğeler, AxisRefListAttr bölümünde listelenen kısıtlamaları karşılamalıdır.
  • gathering_axes, işlenen parçalama işlemine uygulandığında out_sharding elde edilir.

Özellikler: SameOperandsAndResultType

Arayüzler: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
gathering_axes::mlir::sdy::ListOfAxisRefListsAttrEksen referans listelerinin listesi
out_sharding::mlir::sdy::TensorShardingAttrTensör parçalama

İşlenenler:

İşlem Gören Açıklama
tensor şekillendirilmiş herhangi bir jeton dışı türde değer

Sonuçlar:

Sonuç Açıklama
result şekillendirilmiş herhangi bir jeton dışı türde değer

sdy.all_reduce (sdy::AllReduceOp)

Eksenler boyunca tüm azaltma iletişimi gerçekleştirme

Söz dizimi:

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

Bir tensörün parçalarını reduction_axes içinde belirtilen eksenler boyunca azaltır. reduction_axes sırası sonuç için önemli değildir ancak ilgili kopya gruplarının sırasını etkileyebilir.

Kısıtlamalar:

  • Sdy_CollectiveOpInterface bölümünde listelenen kısıtlamaları karşılamalıdır.
  • reduction_axes, AxisRefListAttr içinde listelenen kısıtlamaları karşılamalıdır.
  • reduction_axes, ağa göre sıralanmalıdır.
  • İşlenen parçalama ve out_sharding, eşdeğer boyut parçalamalarına sahip olmalıdır.
  • reduction_axes, işlenen boyut parçalama ve çoğaltılmış eksenlerle çakışmamalıdır (azaltılmamış eksenlerle çakışabilir).
  • reduction_axes, out_sharding'nin azaltılmamış eksenleriyle çakışmamalıdır. Başka bir deyişle, out_sharding, reduction_axes boyunca (örtülü veya açıkça) çoğaltılmalıdır.

Özellikler: SameOperandsAndResultType

Arayüzler: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
reduction_axes::mlir::sdy::AxisRefListAttrEksen referanslarının listesi
reduction_op::mlir::sdy::ReductionOpAttrreduction op enum
out_sharding::mlir::sdy::TensorShardingAttrTensör parçalama

İşlenenler:

İşlem Gören Açıklama
tensor şekillendirilmiş herhangi bir jeton dışı türde değer

Sonuçlar:

Sonuç Açıklama
result şekillendirilmiş herhangi bir jeton dışı türde değer

sdy.all_slice (sdy::AllSliceOp)

Eksenler boyunca dinamik dilimleme işlemi gerçekleştirir

Söz dizimi:

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

Bir tensörün parçalarını slicing_axes içinde belirtilen eksenler boyunca dilimler. sdy.all_slice ve sdy.all_gather arasında cebirsel bir ikilik vardır.

slicing_axes, eksen listelerinin listesidir. Dış liste, tensörün boyutlarını aşıyor. Her iç liste, ilgili boyutta bir dilimleme işleminin gerçekleştirileceği eksenleri belirtir. Sonucun parçalanmasını (out_sharding) elde etmek için işlenenin (tensor) parçalanmasına uygulanır.

out_sharding, sonucun parçalanmasını belirlemek için kullanılmaz. Bunun yerine, sonucun parçalama işlemi işlenenin ve slicing_axes parçalama işlemine göre belirlenir ve out_sharding bu çıkarılan parçalama işlemiyle eşleşmelidir.

Örnek:

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

Kısıtlamalar:

  • Sdy_CollectiveOpInterface bölümünde listelenen kısıtlamaları karşılamalıdır.
  • slicing_axes içindeki öğeler, AxisRefListAttr bölümünde listelenen kısıtlamaları karşılamalıdır.
  • slicing_axes, işlenen parçalama işlemine uygulandığında out_sharding elde edilir.

Özellikler: SameOperandsAndResultType

Arayüzler: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
slicing_axes::mlir::sdy::ListOfAxisRefListsAttrEksen referans listelerinin listesi
out_sharding::mlir::sdy::TensorShardingAttrTensör parçalama

İşlenenler:

İşlem Gören Açıklama
tensor şekillendirilmiş herhangi bir jeton dışı türde değer

Sonuçlar:

Sonuç Açıklama
result şekillendirilmiş herhangi bir jeton dışı türde değer

sdy.all_to_all (sdy::AllToAllOp)

Eksenler boyunca tümden tümüne iletişim gerçekleştirir

Söz dizimi:

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

Parametre listesindeki her (axes, src_dim, tgt_dim) demeti için bu işlem, bir tensörün parçalarını tgt_dim boyutu ve axes içinde belirtilen eksenler boyunca dilimler, bu parçaları eksenler boyunca dağıtır ve src_dim boyutu boyunca birleştirir.

Bu işlem, temelde src_dim ve axes boyunca bir all-gather'ın ardından tgt_dim ve axes boyunca bir all-slice'ın birleşimidir. Yani, giriş tensöründeki eksen parçalama boyutu src_dim'nın soneki, çıkış tensöründeki eksen parçalama boyutu tgt_dim'ye eklenir.

Tümüyle tümü, sonuç parçalama (out_sharding) elde etmek için işlenenin (tensor) parçalanmasına uygulanır.

out_sharding, sonucun parçalanmasını belirlemek için kullanılmaz. Bunun yerine, sonucun parçalama işlemi işlenenlerin (src_dim, tgt_dim ve axes) parçalama işlemine göre belirlenir ve out_sharding, bu çıkarılan parçalama işlemiyle eşleşmelidir.

Örnek:

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

Kısıtlamalar:

  • Sdy_CollectiveOpInterface bölümünde listelenen kısıtlamaları karşılamalıdır.
  • Parametre listesi boş olmamalıdır.
  • params içindeki her parametre için:
    • axes içindeki öğeler, AxisRefAttr kısıtlamalarını karşılamalıdır.
    • src_dim ve tgt_dim geçerli boyutlar olmalıdır (negatif olmayan ve tensörün sıralamasından küçük).
    • Tüm parametrelerdeki src_dim veya tgt_dim benzersiz olmalıdır.
    • src_dim tüm parametrelerde artan düzende sıralanmalıdır.
  • İşlenen parçalama işleminde axes öğesinin src_dim konumundan tgt_dim konumuna taşınması out_sharding sonucunu verir.

Özellikler: SameOperandsAndResultType

Arayüzler: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
params::mlir::sdy::AllToAllParamListAttrTümden tüm parametrelerin listesi
out_sharding::mlir::sdy::TensorShardingAttrTensör parçalama

İşlenenler:

İşlem Gören Açıklama
tensor şekillendirilmiş herhangi bir jeton dışı türde değer

Sonuçlar:

Sonuç Açıklama
result şekillendirilmiş herhangi bir jeton dışı türde değer

sdy.collective_permute (sdy::CollectivePermuteOp)

Eksenleri değiştirmek için toplu permütasyon iletişimi gerçekleştirir

Söz dizimi:

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

Tensörü parçalayan eksenleri yeniden sıralamak/değiştirmek için her cihazdaki giriş tensörünün bir parçasını başka bir cihaza gönderir.

Toplu permütasyon, giriş parçalamayı her boyutun daha önce olduğu gibi parçalanması (yani, boyut çarpımı daha önce tensörü parçalayan eksenlerin boyut çarpımına eşit olan eksenler boyunca parçalanması) şeklinde dönüştürebilir.

Bu özellik, tek bir boyuttaki veya farklı boyutlardaki eksenleri yeniden sıralamak ve parçalanmış eksenleri çoğaltılmış eksenlerle değiştirmek için kullanışlıdır.

Aşağıdaki örnekte, parçalanmış tensör boyutu tensor<1x4x2xf32>'dır ve bu boyut, toplu permütasyon tarafından korunur.

Örnek:

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>

Kısıtlamalar:

  • Sdy_CollectiveOpInterface bölümünde listelenen kısıtlamaları karşılamalıdır.
  • Giriş ve çıkış parçalama farklı ağlara sahipse bu ağlar tam olarak aynı eksenlere ve farklı cihaz kimliği sırasına sahip olmalıdır.
  • Her boyut için, out_sharding içindeki parçalama ekseni boyutlarının çarpımı, ilgili işlenen boyut parçalamanın çarpımıyla eşleşmelidir.

Özellikler: SameOperandsAndResultType

Arayüzler: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
out_sharding::mlir::sdy::TensorShardingAttrTensör parçalama

İşlenenler:

İşlem Gören Açıklama
tensor şekillendirilmiş herhangi bir jeton dışı türde değer

Sonuçlar:

Sonuç Açıklama
result şekillendirilmiş herhangi bir jeton dışı türde değer

sdy.constant (sdy::ConstantOp)

Sabit işlem

Sabit bir value değerinden output tensörü üretir.

Şu sayfaya bakın: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

Örnek:

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

Özellikler: AlwaysSpeculatableImplTrait

Arayüzler: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Etkiler: MemoryEffects::Effect{}

Özellikler:

ÖzellikMLIR TürüAçıklama
value::mlir::ElementsAttrsabit vektör/tensör özelliği

Sonuçlar:

Sonuç Açıklama
output Herhangi bir jeton dışı türde değerlerin statik olarak şekillendirilmiş tensörü

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Veri akışı kenar işlemi

Söz dizimi:

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

Bazı X işlemlerinin veri akışı kenarı, bir dizi kaynak (her biri X'in işleneni veya X'in blok sonlandırıcısının işleneni) ile bir dizi hedef (her biri X'in sonucu veya X'in blok bağımsız değişkeni) arasında bir köprü tanımlar. Bu durumda, tüm kaynaklar ve hedefler aynı şekilde parçalanmalıdır.

Bir işlem, birbirine dik olan birden fazla veri akışı kenarına sahip olabilir.

Örneğin:

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

Bu işlem n veri akışı kenarına sahipken i. veri akışı kenarı kaynaklar x_i, return_value_i ve hedefler y_i, pred_arg_i, body_arg_i arasındadır.

Bir sdy.data_flow_edge, bir kenarın sahibini giriş olarak alır (hedeflerden herhangi biri olabilir ancak tercihen bir blok bağımsız değişkeni yerine bir işlem sonucu olmalıdır) ve başka bir kullanımı olmamalıdır. Bu işlem, başlangıçta herhangi bir kullanımı olmayan bir girişi alabileceğinden saf değildir.

sdy.data_flow_edge, tüm uç hedefleri için isteğe bağlı bir parçalama da içerir ve bu parçalama, yayılma sırasında hedeflerin parçalaması yerine güncellenmelidir (eklenebiliyorsa). Bu, bir işlem çok sayıda kenara sahip olduğunda kullanışlıdır. Çünkü bu durumda:

  • her kenarda ayrı ayrı yayılır.
  • Tüm hedefleri aynı anda güncellemek yerine her kenarın parçalama işlemini ayrı ayrı güncelleme (ör. bir işlem, sonuç parçalama işlemleri için tek bir değişmez TensorShardingPerValueAttr içerir).
  • Bir kaynağın parçalama işlemi değiştiğinde her kenarı ayrı ayrı iş listesine ekleyin.

Yayma, kaynaklar işlenenler ve hedefler sonuçlar, kimlik ise sdy.op_sharding_rule olarak kullanılıyormuş gibi davranarak sdy.data_flow_edge'nın tüm kaynakları ve hedefleri arasında parçalama işlemini yayar. Bu nedenle, ileri yayılım kaynaklardan hedeflere, geri yayılım ise hedeflerden kaynaklara doğru gerçekleşir.

sdy.data_flow_edge girişinin bir SdyDialect işlemiyle tanımlanmasına izin vermediğimiz için, girişin, kayıtlı olmayan sdy.sharding özelliğine sahip bir işlemle tanımlandığını varsayabiliriz.

Özellikler: SameOperandsAndResultType

Arayüzler: InferTypeOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
sharding::mlir::sdy::TensorShardingAttrTensör parçalama

İşlenenler:

İşlem Gören Açıklama
input şekillendirilmiş herhangi bir jeton dışı türde değer

Sonuçlar:

Sonuç Açıklama
result şekillendirilmiş herhangi bir jeton dışı türde değer

sdy.func_data_flow_edge (sdy::FuncDataFlowEdgeOp)

Func input/output data flow edge op.

Söz dizimi:

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

Bir veri akışı kenarı işlemi ancak işlev bağımsız değişkenleri veya çağrı sonuçları için. İşlemi BlockArgument olduğunda, arayan callOp'un argümanından func argümanının kullanıcılarına bir köprüdür. Her işlev bağımsız değişkeni için bir işlev veri akışı kenarı vardır. İşleneni bir OpResult olduğunda; çağrılan funcOp'un dönüş değerinden çağrı sonucunun kullanıcılarına bir köprüdür. Her bir çağrı sonucu için bir işlev veri akışı kenarı vardır.

Özellikler: SameOperandsAndResultType

Arayüzler: InferTypeOpInterface, SymbolUserOpInterface

İşlenenler:

İşlem Gören Açıklama
operand şekillendirilmiş herhangi bir jeton dışı türde değer

Sonuçlar:

Sonuç Açıklama
result şekillendirilmiş herhangi bir jeton dışı türde değer

sdy.manual_computation (sdy::ManualComputationOp)

Manuel topluluklarla çeşitli cihaz paralel işlemi

Söz dizimi:

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)

Açık kolektiflerle cihaz başına yerel kod açısından yazılmış bir bölgeye geçin. Burada mantıksal şekiller, cihaz başına yerel fiziksel arabellek şekilleriyle eşleşir ve kolektifler, cihazlar arası fiziksel iletişimle tam olarak uyuşur.

Gövde, manual_axes'e göre yereldir. Yayılma, manuel_axes listesinde olmayan serbest eksenlerdeki gövde aracılığıyla gerçekleşir.

Derecelendirilmemiş tüm tensörlerin, tam olarak kopyalanmış şekilde 0 dereceli bir parçalama içermesi beklenir.

Kısıtlamalar:

  • in_shardings ve out_shardings öğeleri, TensorShardingAttr içinde listelenen kısıtlamaları karşılamalıdır.
  • İşlem bölgesinin genel ve yerel tensör giriş/çıkış sayısı eşleşmelidir.
  • Manuel eksenler, her bir boyut parçalama işleminde serbest eksenlerden önce gelmelidir.
  • Manuel eksenler dolgu oluşturamaz. Yani boyut, ilgili manuel eksen boyutuna bölünebilmelidir.
  • Op bölgeleri bağımsız değişkenlerinin/sonuçlarının küresel ve yerel şekilleri eşleşmelidir.

Özellikler: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Arayüzler: ShardableDataFlowOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
in_shardings::mlir::sdy::TensorShardingPerValueAttrBir işlemin işleneni/sonucu başına tensör parçalama
out_shardings::mlir::sdy::TensorShardingPerValueAttrBir işlemin işleneni/sonucu başına tensör parçalama
manual_axes::mlir::sdy::ManualAxesAttrBir ManualComputationOp'nin manuel olduğu eksenlerin listesi

İşlenenler:

İşlem Gören Açıklama
tensors Herhangi bir jeton dışı türün değişken sayıda bağımsız değişkeni

Sonuçlar:

Sonuç Açıklama
results Herhangi bir jeton dışı türün değişken sayıda bağımsız değişkeni

sdy.mesh (sdy::MeshOp)

Adlandırılmış ağ

Söz dizimi:

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

Yeni bir adlandırılmış ağ tanımlar. Bir modüldeki tüm ağlar aynı sayıda cihaza sahip olmalıdır (tek bir device_id'ye sahip ağlar hariç). Ağ, modülün Symbol bölümünde görünen ve SymbolTable ile referans verilebilen bir name işlemidir.

Özellikler: HasParent<ModuleOp>, SymbolName

Arayüzler: Symbol

Özellikler:

ÖzellikMLIR TürüAçıklama
sym_name::mlir::StringAttrdize özelliği
mesh::mlir::sdy::MeshAttrEksenlerin ağı ve cihaz listesi

sdy.named_computation (sdy::NamedComputationOp)

Adlandırılmış hesaplama işlemi

Söz dizimi:

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)

Bir hesaplamayı (yani bir işlem bloğunu) gruplandırır ve ona bir ad verir. Her şey satır içi olarak yerleştirilmiş gibi yayılım bölgeye girer/bölgeden çıkar.

Bu, arama talimatlarının diğer işlevlere yayılmasını sağlamak için kullanılabilir. Shardy'yi kullanan herkes, çağrı işlemlerini sdy.named_computation işlemlerine dönüştüren, çağrılan işlevin gövdesini named_computation'nin gövdesine kopyalayan/çoğaltan bir içe/dışa aktarma geçişi yazmalıdır.

Bölgedeki her blok bağımsız değişkeninin ve döndürülen değerlerin türü, işlemin işlenenlerinin türü ve sonuç türüyle aynı olmalıdır.

Örnek:

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

Özellikler: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Arayüzler: ConditionallySpeculatable, InferTypeOpInterface, ShardableDataFlowOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
name::mlir::StringAttrdize özelliği
in_shardings::mlir::sdy::TensorShardingPerValueAttrBir işlemin işleneni/sonucu başına tensör parçalama
out_shardings::mlir::sdy::TensorShardingPerValueAttrBir işlemin işleneni/sonucu başına tensör parçalama

İşlenenler:

İşlem Gören Açıklama
operands Herhangi bir jeton dışı türün değişken sayıda bağımsız değişkeni

Sonuçlar:

Sonuç Açıklama
"unnamed" Herhangi bir jeton dışı türün değişken sayıda bağımsız değişkeni

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Yayılma bariyeri işlemi

Söz dizimi:

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

Bu işlem, kimlik işlemi gibi çalışır ve giriş olarak aldığı değeri çıkış olarak verir. Ancak yayılma açısından bu, yayılmanın yalnızca belirli bir yönde akmasına izin verir.

Bu, bariyer işleminin sonucunun ve işleneninin kullanımları arasında parçalama işlemlerinin yayılmasını önler.

  • FORWARD, parçalama işlemlerinin yalnızca işlenenden sonuca doğru yapılabileceği anlamına gelir.
  • BACKWARD, parçalamaların yalnızca sonuçtan işlenene akabileceği anlamına gelir.
  • NONE, bu işlem aracılığıyla parçalama işleminin yayılamayacağı anlamına gelir.
  • Bu işlem gereksiz olacağından BOTH belirtilemez.

Özellikler: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Arayüzler: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Etkiler: MemoryEffects::Effect{}

Özellikler:

ÖzellikMLIR TürüAçıklama
allowed_direction::mlir::sdy::PropagationDirectionAttryayılım yönü enum

İşlenenler:

İşlem Gören Açıklama
input Herhangi bir jeton dışı türdeki değerlerin sıralanmış tensörü

Sonuçlar:

Sonuç Açıklama
result Herhangi bir jeton dışı türdeki değerlerin sıralanmış tensörü

sdy.reduce_scatter (sdy::ReduceScatterOp)

Eksenler boyunca dağılımı azaltan bir iletişim gerçekleştirir

Söz dizimi:

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

Bir tensörün parçalarını reduce_scatter_axes içinde belirtilen eksenler boyunca azaltır ve ardından sonucu aynı eksenler boyunca dağıtır. Bu işlem, temelde aynı reduce_scatter_axes üzerinde sdy.all_reduce ve sdy.all_slice hareketlerinin birleşimidir.

Kısıtlamalar:

  • Sdy_CollectiveOpInterface bölümünde listelenen kısıtlamaları karşılamalıdır.
  • reduce_scatter_axes içindeki öğeler, AxisRefListAttr bölümünde listelenen kısıtlamaları karşılamalıdır.
  • reduce_scatter_axes, işlenen parçalama işlemine uygulandığında out_sharding elde edilir.

Özellikler: SameOperandsAndResultType

Arayüzler: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
reduce_scatter_axes::mlir::sdy::ListOfAxisRefListsAttrEksen referans listelerinin listesi
reduction_op::mlir::sdy::ReductionOpAttrreduction op enum
out_sharding::mlir::sdy::TensorShardingAttrTensör parçalama

İşlenenler:

İşlem Gören Açıklama
tensor şekillendirilmiş herhangi bir jeton dışı türde değer

Sonuçlar:

Sonuç Açıklama
result şekillendirilmiş herhangi bir jeton dışı türde değer

sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

Örtülü veya açıkça çoğaltılmış eksenleri azaltılmamış eksenlere taşıyın.

Söz dizimi:

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

axes, işlenende örtülü veya açık bir şekilde çoğaltılmalıdır. Bu işlem, sonuçta boyutları küçültülmemiş olarak gösterilmelerini sağlar. Aşağıdaki ilişkiyi kurduk:

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

Örnek:

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

Kısıtlamalar:

  • Sdy_CollectiveOpInterface bölümünde listelenen kısıtlamaları karşılamalıdır.
  • axes, AxisRefListAttr içinde listelenen kısıtlamaları karşılamalıdır.
  • axes, ağa göre sıralanmalıdır.
  • axes boş değil.
  • Giriş ve çıkış parçalama işlemleri aynı boyut parçalama işlemlerine sahip olmalıdır.
  • axes, işlenen parçalama işleminde örtülü veya açık bir şekilde kopyalanmalıdır.
  • inUnreducedAxes + axes = outUnreducedAxes.

Özellikler: SameOperandsAndResultType

Arayüzler: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
axes::mlir::sdy::AxisRefListAttrEksen referanslarının listesi
out_sharding::mlir::sdy::TensorShardingAttrTensör parçalama

İşlenenler:

İşlem Gören Açıklama
tensor şekillendirilmiş herhangi bir jeton dışı türde değer

Sonuçlar:

Sonuç Açıklama
result şekillendirilmiş herhangi bir jeton dışı türde değer

sdy.reshard (sdy::ReshardOp)

Bir tensörü farklı bir parçalamaya yeniden parçalar.

Söz dizimi:

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

Giriş tensörünü, giriş tensörünün mevcut parçalamasından farklı olan belirtilen parçalama ile yeniden parçalar.

Hem ShardingConstraintOp hem de ReshardOp, bir tensöre parçalama ekler. Bu cihazların kullanım ömrü:

  1. Parçalama yayılımından önce, kullanıcılar tarafından ShardingConstraintOp eklenir.
  2. Parçalama yayılımı, ShardingConstraintOp'u kullanır. Parçalama yayılımı sonuçlarında ShardingConstraintOp yok. Bunun yerine, gerekirse ReshardOp eklenebilir.
  3. Bir bölümleyici, ReshardOp'u toplu bir işleme (veya kimlik işlemine) dönüştürür. Bölümleyici sonuçlarında ReshardOp olmamalıdır.

Özellikler: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Arayüzler: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface), SymbolUserOpInterface

Etkiler: MemoryEffects::Effect{}

Özellikler:

ÖzellikMLIR TürüAçıklama
sharding::mlir::sdy::TensorShardingAttrTensör parçalama

İşlenenler:

İşlem Gören Açıklama
input jeton olmayan herhangi bir tür

Sonuçlar:

Sonuç Açıklama
result jeton olmayan herhangi bir tür

sdy.return (sdy::ReturnOp)

sdy.return işlemi, sdy bölge tabanlı işlemlere ve diğer tüm Shardy bölge tabanlı işlemlere bağlı bölgeleri sonlandırır. Bu işlev, değişken sayıda bağımsız değişken alır: Türleri herhangi bir şey olabilen (ancak aynı türden, örneğin AnyTensor) bir değer listesini bağımsız değişken olarak alır ve bu nedenle Shardy IR yığınının çeşitli düzeylerinde yeniden kullanılabilir.

Söz dizimi:

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

Özellikler: AlwaysSpeculatableImplTrait, ReturnLike, Terminator

Arayüzler: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Etkiler: MemoryEffects::Effect{}

İşlenenler:

İşlem Gören Açıklama
results Herhangi bir jeton dışı türün değişken sayıda bağımsız değişkeni

sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

İşlenenin bazı parçalanmış eksenlerini, sonucun azaltılmamış eksenlerine taşıyın.

Söz dizimi:

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

İşlenen, parçalara ayırmak için axes kullanılmalıdır. Bu işlem, sonuçta bu öğelerin küçültülmemesini sağlar. Aşağıdaki ilişkimiz var:

all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes), burada all-gather, sharded-to-unreduced, all-reduce aynı eksenlerde uygulanır.

Örnek:

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

Kısıtlamalar:

  • Sdy_CollectiveOpInterface bölümünde listelenen kısıtlamaları karşılamalıdır.
  • axes içindeki öğeler, AxisRefListAttr içinde listelenen kısıtlamaları karşılamalıdır.
  • axes, işlenen parçalama işlemine uygulandığında out_sharding elde edilir.

Özellikler: SameOperandsAndResultType

Arayüzler: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
axes::mlir::sdy::ListOfAxisRefListsAttrEksen referans listelerinin listesi
out_sharding::mlir::sdy::TensorShardingAttrTensör parçalama

İşlenenler:

İşlem Gören Açıklama
tensor şekillendirilmiş herhangi bir jeton dışı türde değer

Sonuçlar:

Sonuç Açıklama
result şekillendirilmiş herhangi bir jeton dışı türde değer

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Bir tensörü belirtilen parçalama ile kısıtlar

Söz dizimi:

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

Bu tensörün veya kullanım alt kümesinin nasıl parçalanması gerektiğini belirtmek için ara tensöre (ör. matmul sonucu) parçalama ekler.

Parçalama işleminde açık boyutlar ve kısıtlanmamış eksenler varsa tensörün açık boyutlar boyunca daha fazla parçalanabileceği anlamına gelir.

Bu işlem şunlardan biri olabilir:

  • Kullanımı olmamalıdır (bağlantısız). Bu, eklenen parçalamanın giriş tensörünün kendisinin nasıl parçalanması gerektiği anlamına gelir.
  • Kullanımları vardır. Bu, eklenen parçalama işleminin, parçalama kısıtlaması işleminin kullanımlarının nasıl parçalanması gerektiği anlamına gelir. Giriş tensörünün diğer kullanımları farklı bir parçalamaya sahip olabilir (giriş tensörünün başka kullanımı yoksa davranış, kullanım yok durumuyla aynıdır).

Özellikler: SameOperandsAndResultType

Arayüzler: InferTypeOpInterface, SymbolUserOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
sharding::mlir::sdy::TensorShardingAttrTensör parçalama

İşlenenler:

İşlem Gören Açıklama
input jeton olmayan herhangi bir tür

Sonuçlar:

Sonuç Açıklama
result jeton olmayan herhangi bir tür

sdy.sharding_group (sdy::ShardingGroupOp)

Gruptaki tensörlerin aynı parçalamaya sahip olmasını zorlar.

Söz dizimi:

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

Bu işlem, tensörleri parçalama gruplarına atamak için bir arayüz sağlar ( aynı parçalamaya sahip olması zorunlu kılınacak tensör grupları). Yayma sırasında, bir grup öğesi parçalandığı anda diğer tüm üyeler de aynı şekilde parçalanır. Bu işlem, bağımsız değişken grup kimliğini alır ve sonuç döndürmez. Bunun yerine, giriş tensörünü belirtilen kimliğe sahip gruba eklemek için dahili parçalama grubu gösterimini değiştirir.

Arayüzler: InferTypeOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
group_id::mlir::IntegerAttr64 bit işaretsiz tam sayı özelliği

İşlenenler:

İşlem Gören Açıklama
input Herhangi bir jeton dışı türdeki değerlerin sıralanmış tensörü

Özellikler

AllToAllParamAttr

Tümden tüm parametresi

Söz dizimi:

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

Tümden tüme işlemi gerçekleştirmek için eksenleri ve kaynak/hedef boyutlarını içeren bir demet.

Parametreler:

Parametre C++ türü Açıklama
baltalar ::llvm::ArrayRef<AxisRefAttr> tümden tümeye işleminin gerçekleştirileceği eksenler
src_dim int64_t kaynak boyutu dizini
tgt_dim int64_t hedef boyut dizini

AllToAllParamListAttr

Tümden tüm parametrelerin listesi

Söz dizimi:

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

Parametreler:

Parametre C++ türü Açıklama
value ::llvm::ArrayRef<AllToAllParamAttr>

AxisRefAttr

Tam eksene veya bölünmüş alt eksene referans

Söz dizimi:

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

Kısıtlamalar:

  • name, bağlı MeshAttr içinde bulunmalıdır.
  • sub_axis_info varsa SubAxisInfoAttr kısıtlamalarını karşılamalıdır.

Parametreler:

Parametre C++ türü Açıklama
ad ::llvm::StringRef bu eksenin adı
sub_axis_info SubAxisInfoAttr Alt eksen ise ek bilgiler

AxisRefListAttr

Eksen referanslarının listesi

Söz dizimi:

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

Kısıtlamalar:

  • value içindeki öğeler, AxisRefAttr kısıtlamalarını karşılamalıdır.
  • Yinelenen eksen referansları veya birbiriyle çakışan alt eksenler yoktur.
  • Bitişik iki eksen referansı, aynı tam eksenin ardışık alt eksenleri değildir. Yani, tek bir alt eksende veya tam eksende birleştirilebilirler.

Parametreler:

Parametre C++ türü Açıklama
value ::llvm::ArrayRef<AxisRefAttr>

AxisToPropagationDetailsAttr

Belirli bir eksen ve kaynak için yayılma kenarı akışı ayrıntıları.

Söz dizimi:

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

Bir kaynak değeri referansını belirli bir eksen boyunca hedef değeri referanslarının listesiyle eşler.

Parametreler:

Parametre C++ türü Açıklama
axis_name ::mlir::sdy::AxisRefAttr Tam eksene veya bölünmüş alt eksene referans
kaynak ::mlir::sdy::EdgeValueRefAttr type türündeki bir değer kenarının belirli bir dizinine referans.
hedef ::llvm::ArrayRef<EdgeValueRefAttr> uç hedef değerlerinin listesi

DimMappingAttr

Bir boyut için faktör dizinlerinin listesi

Boş liste, bunun boş bir eşleme olduğunu (* ile ayrıştırılır/yazdırılır), yani boyutun herhangi bir faktörle eşlenmediğini gösterir.

Kısıtlamalar:

  • En az bir faktör dizini vardır.
  • Faktör dizinleri [0, $factor_sizes) aralığında olmalıdır.
  • Birden fazla faktör varsa hiçbirinin boyutu 1 olamaz.
  • Yinelenen faktör dizinleri yok.

Parametreler:

Parametre C++ türü Açıklama
factor_indices ::llvm::ArrayRef<int64_t> Bu boyutun eşlendiği faktörler

DimensionShardingAttr

Boyut parçalama

Bir tensör boyutunu büyükten küçüğe doğru parçalamak için kullanılan eksen adlarının listesi, boyutun daha fazla parçalanıp parçalanamayacağını belirten bir boole değeri ve bu boyut parçalama işleminin önceliğini gösteren isteğe bağlı bir tam sayı. Bu öncelik, parçalama yayılımı sırasında dikkate alınır. Öncelikler, kullanıcı parçalama ek açıklamalarından kaynaklanır ve daha düşük bir değer, daha yüksek bir önceliği ifade eder. Öncelik, ek açıklamada eksik olduğunda en yüksek öncelik varsayılır.

Kısıtlamalar:

  • axes içindeki öğeler, AxisRefListAttr'de listelenen kısıtlamaları karşılamalıdır.
  • Boyut parçalama işleminin önceliği varsa:
    • Öncelik, 0'dan büyük veya 0'a eşit olmalıdır.
    • Boyut kapalıysa en az bir ekseni vardır.

Parametreler:

Parametre C++ türü Açıklama
baltalar ::llvm::ArrayRef<AxisRefAttr> axis refs
is_closed bool Bu boyutun daha fazla parçalanamayıp parçalanamayacağı
önceliği std::optional<int64_t> Kullanıcı önceliğine dayalı yayma sırasında kullanılan öncelik

EdgeValueRefAttr

type türünde bir değer kenarının belirli bir dizinine referans.

Söz dizimi:

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

Parametreler:

Parametre C++ türü Açıklama
tür ::mlir::sdy::EdgeNodeType EdgeNodeType türünde bir enum
indeks int64_t Tam sayı dizini (0, 1, 2 vb.)

ListOfAxisRefListsAttr

Eksen referans listelerinin listesi

Söz dizimi:

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

Parametreler:

Parametre C++ türü Açıklama
value ::llvm::ArrayRef<AxisRefListAttr>

ManualAxesAttr

Bir ManualComputationOp'nin manuel olduğu eksenlerin listesi

Söz dizimi:

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

Parametreler:

Parametre C++ türü Açıklama
value ::llvm::ArrayRef<StringAttr>

MeshAttr

Eksen ağı ve cihaz listesi

Söz dizimi:

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

Ağ, eksenlerin listesi ve cihaz sıralamasını belirten isteğe bağlı cihaz kimlikleri listesidir.

Eksen listesi boşsa

  • device_ids sağlanmazsa boş bir ağdır.
  • device_ids sağlanırsa negatif olmayan tek bir tam sayı olmalıdır. Buna maksimum parçalama ağı diyoruz.

Eksen listesi sağlanırsa

  • Bir cihaz kimliği listesi belirtilmişse eksen boyutlarının çarpımı cihaz sayısıyla eşleşmelidir.
  • Cihaz kimliği listesi belirtilmemişse örtülü cihaz kimliği listesi iota(product(axes)) olur. Basitlik açısından, iota(product(axes)) ile aynı olan bir cihaz kimliği listesinin belirtilmesine de izin verilmez. Bu durumda, bir cihaz kimliği listesi belirtilmemelidir.
  • Eksenlerin toplam boyutu 1 olsa bile bu, maksimum parçalama ağı değildir.

Aşağıda bazı ağ örnekleri verilmiştir:

  • Boş bir ağ, yayma sırasında değiştirilebilecek bir yer tutucu ağı temsil eder: <[]>
  • Eksen listesi olmayan ve tek bir negatif olmayan cihaz kimliği içeren bir ağ. Bu, maksimum parçalama ağıdır: <[], device_ids=[3]>
  • İki eksenli ve örtülü cihaz kimlikli bir ağ iota(6): <["a"=2, "b"=3]>
  • Cihaz sıralamasını belirten iki eksenli ve açık cihaz kimliklerine sahip bir ağ: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Kısıtlamalar:

  • device_ids içindeki öğeler negatif olmamalıdır.
  • axes boşsa device_ids boyutu 0 (boş ağ) veya 1 (maksimum parçalama ağı) olabilir.
  • axes boş değilse:
    • axes içindeki öğeler yinelenen adlara sahip olmamalıdır.
    • device_ids belirtilmişse orijinal device_ids iota(product(axis_sizes)) değildir ve sıralanmış device_ids iota(product(axis_sizes)) olur.

Parametreler:

Parametre C++ türü Açıklama
baltalar ::llvm::ArrayRef<MeshAxisAttr> örgü eksenleri
device_ids ::llvm::ArrayRef<int64_t> açık cihaz sıralaması veya maksimum cihaz kimliği

MeshAxisAttr

Ağdaki adlandırılmış eksen

Söz dizimi:

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

Parametreler:

Parametre C++ türü Açıklama
ad ::llvm::StringRef ad
beden int64_t bu eksenin boyutu

OpShardingRuleAttr

Bir işlemin nasıl bölümlenebileceğini belirtir.

Söz dizimi:

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

Parçalama kuralı, bir işlemin, işlemdeki çeşitli özelliklere (ör. özellikler, işlenenlerin şekli, sonuçların şekli) göre nasıl bölümlenebileceğini belirtir. Örneğin:

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

Parçalanamasalar bile boyutu 1 olan faktörlere izin verdiğimizi unutmayın. Bunun temel nedeni, nokta nokta işlemleri gibi birçok işlemin, işlenenler ve sonuçlar arasında karşılık gelen boyutu 1 olan boyutlara sahip olmasıdır.

Faktör türleri:

  • reduction_factors, nokta işleminde daraltma boyutları gibi azaltma gerektiren faktörlerin dizinlerini içerir. Bu faktörler işlenenlerde olabilir ancak sonuçlarda olamaz.
  • need_replication_factors, sıralama işleminde sıralanan boyut gibi tam kopyalama gerektiren faktörlerin dizinlerini içerir.
  • permutation_factors, parçalanmışlarsa toplu permütasyon gerektiren faktörlerin dizinlerini içerir. Örneğin, bir doldurma işlemindeki doldurma boyutları.
  • Diğer tüm faktörler, geçiş faktörleri olarak kabul edilir. Yani, kendileriyle eşlenen tüm tensörlerde aynı şekilde parçalanırsa herhangi bir iletişim gerektirmeyen faktörlerdir.

blocked_propagation_factors, parçalamaların yayılmasına izin verilmeyen faktörleri içerir. Faktör türlerine diktir. Yani, engellenen yayılma faktörü herhangi bir faktör türü olabilir.

is_custom_rule, bunun kullanıcı tarafından tanımlanan bir kural olup olmadığını açıklar. Kullanıcılar, özel çağrıları için parçalama kuralları tanımlayabilir veya standart işlemler için önceden tanımlanmış parçalama kurallarını geçersiz kılabilir. Özel kurallar her zaman korunur/hiçbir zaman kaldırılmaz.

Kısıtlamalar:

  • İşlenen/sonuç eşlemelerinin sayısı, işlemin işlenen/sonuç sayısıyla eşleşmelidir.
  • En az bir eşleme vardır (işlenenleri/sonuçları olmayan bir işlem için kural olamaz).
  • Her TensorMappingAttr öğesinin sıralaması, karşılık gelen tensör türünün sıralamasıyla eşleşir.
  • Her faktör grubu için (reduction_factors, need_replication_factors, permutation_factors):
    • Öğeler [0, $factor_sizes] aralığında olmalıdır.
    • Her grupta ve gruplar arasında yinelenen faktör dizini olmamalıdır.

Parametreler:

Parametre C++ türü Açıklama
factor_sizes ::llvm::ArrayRef<int64_t> Bu kuraldaki tüm faktörlerin boyutları
operand_mappings ::llvm::ArrayRef<TensorMappingAttr> işlenen eşlemeleri
result_mappings ::llvm::ArrayRef<TensorMappingAttr> sonuç eşlemeleri
reduction_factors ::llvm::ArrayRef<int64_t> azaltılması gereken faktörler
need_replication_factors ::llvm::ArrayRef<int64_t> tam kopyalama gerektiren faktörler
permutation_factors ::llvm::ArrayRef<int64_t> collective-permute gerektiren faktörler
blocked_propagation_factors ::llvm::ArrayRef<int64_t> Parçalama işlemlerinin yayılmadığı faktörler
is_custom_rule bool Kuralın stablehlo.custom_call için olup olmadığını

PropagationEdgesAttr

Tüm yayma adımları için yayma kenarı meta verileri.

Söz dizimi:

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

Bir değerin eksen başına yayılma ayrıntılarının adım dizinine göre gruplandırılmış listesi.

Parametreler:

Parametre C++ türü Açıklama
value ::llvm::ArrayRef<PropagationOneStepAttr>

PropagationOneStepAttr

Adım adım yayma meta verileri

Söz dizimi:

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

Tek bir yayma adımı için tüm eksenlerin yayma ayrıntıları.

Parametreler:

Parametre C++ türü Açıklama
step_index int64_t adım dizini
axis_entries ::llvm::ArrayRef<AxisToPropagationDetailsAttr> Yayma kararı başına eksen yayma ayrıntıları

SubAxisInfoAttr

Bu alt eksenin tam eksenden nasıl türetildiği hakkında bilgi

Söz dizimi:

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

Tam bir eksen n alt eksene bölündüğünde eksen [k_1,...,k_n] olarak yeniden şekillendirilir ve i'inci alt eksen, solundaki tüm eksen boyutlarının m=prod(k_1,...,k_(i-1)) (ön boyut olarak da bilinir) ve k_i boyutunun çarpımıyla ifade edilebilir. Bu nedenle, sub-axis-info özelliği bu iki sayıyı içerir ve şu şekilde gösterilir: (m)k (ön boyut m ve boyut k için).

Kısıtlamalar:

  • pre-size en az 1 olmalıdır.
  • size, 1'den büyük olmalıdır.
  • pre-size, tam eksenin boyutunu bölmelidir. Yani hem pre-size hem de size, tam eksenin boyutunu böler ve alt eksen, tam eksenin ötesine geçmez.
  • Alt eksenin boyutu, karşılık gelen tam eksenin boyutuna eşit değildir. Bu durumda bunun yerine tam eksen kullanılmalıdır.

Parametreler:

Parametre C++ türü Açıklama
pre_size int64_t bu alt eksenin solundaki alt eksen boyutlarının çarpımı
beden int64_t bu alt eksenin boyutu

TensorMappingAttr

Bir tensörün her boyutu için faktör eşlemeleri.

Söz dizimi:

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

Kısıtlamalar:

  • dim_mappings içindeki öğeler, DimMappingAttr içindeki kısıtlamaları karşılamalıdır.
  • Boyutlar arasında yinelenen faktör dizinleri olmamalıdır.

Parametreler:

Parametre C++ türü Açıklama
dim_mappings ::llvm::ArrayRef<DimMappingAttr> boyut eşlemeleri

TensorShardingAttr

Tensor parçalama

Söz dizimi:

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

Tensör parçalama işlemi belirli bir ağa bağlıdır ve yalnızca bu ağdaki eksen adlarına referans verebilir. Boyut parçalama, tensörün her boyutu için hangi eksenlerde (veya alt eksenlerde) büyükten küçüğe doğru parçalandığını gösterir. Bir boyutu parçalamayan diğer tüm eksenler, dolaylı olarak veya açıkça (çoğaltılmış eksenler listesinde görünürlerse) çoğaltılır.

Bir tensörde parçalama özelliğinin olmaması, tamamen açık tensör parçalamaya eşdeğerdir.

Bu parçalamanın bağlı olduğu ağ, karşılık gelen bir MeshOp sembolüne referans veren bir sembol adıyla veya satır içi bir MeshAttr ile belirtilebilir.

Parçalama işleminde, azaltılmamış eksenler (unreduced_axes ile belirtilir) olabilir. Bu, tensörün bu eksenler boyunca azaltılmadığı anlamına gelir. Örneğin, bir matmul'un sözleşme boyutu hem sol hem de sağ tarafta x ekseni boyunca parçalanmışsa sonuç x boyunca azaltılmaz. Tensor üzerinde, indirgenmemiş eksenler boyunca bir all-reduce işlemi uygulamak, tensorun bu eksenler boyunca çoğaltılmasını sağlar. Ancak, indirgenmemiş eksenlere sahip bir tensörün hemen tamamen indirgenmesi gerekmez. stablehlo.add gibi doğrusal işlemlere geçirildiğinde indirgenmemiş olarak kalabilir (hem lhs hem de rhs indirgenmemiş olduğu sürece) ve daha sonra tamamen indirgenmiş olabilir. İndirim türünün toplam olduğunu varsayıyoruz. Gelecekte başka indirim türleri de desteklenebilir.

Kısıtlamalar:

  • dim_shardings içindeki öğeler, DimensionShardingAttr bölümünde listelenen kısıtlamaları karşılamalıdır.
  • replicated_axes içindeki öğeler, AxisRefListAttr bölümünde listelenen kısıtlamaları karşılamalıdır.
  • unreduced_axes içindeki öğeler, AxisRefListAttr bölümünde listelenen kısıtlamaları karşılamalıdır.
  • İlgili tensör türü ShapedType değilse parçalama, 0 sıralı olmalı ve tekrarlanan eksen içermemelidir.
  • ShapedType ise:
    • Tensörün bir sıralaması olmalıdır.
    • Boyut parçalama sayısı, tensörün derecesine eşittir.
    • 0 boyutundaki boyutlar parçalanmaz.
  • dim_shardings, replicated_axes ve unreduced_axes genelinde yinelenen eksen referansları veya birbiriyle çakışan alt eksenler yoktur.
  • replicated_axes ve unreduced_axes içindeki öğeler mesh_or_ref'ye göre sıralanır (bkz. AxisRefAttr::getMeshComparator).

Parametreler:

Parametre C++ türü Açıklama
mesh_or_ref ::mlir::Attribute mesh attr veya flat mesh symbol reference attr
dim_shardings ::llvm::ArrayRef<DimensionShardingAttr> boyut parçalama
replicated_axes ::llvm::ArrayRef<AxisRefAttr> axis refs
unreduced_axes ::llvm::ArrayRef<AxisRefAttr> axis refs
reduction_op ::mlir::sdy::ReductionOp ReductionOp türünde bir enum

TensorShardingPerValueAttr

Bir işlemin işleneni/sonucu başına tensör parçalama

Söz dizimi:

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

Her bir işlem operandı/sonucu için bir tane olmak üzere TensorShardingAttr listesi.

Kısıtlamalar:

  • shardings içindeki öğeler, TensorShardingAttr kısıtlamalarını karşılamalıdır.

Parametreler:

Parametre C++ türü Açıklama
parçalama ::llvm::ArrayRef<TensorShardingAttr> değer başına parçalama

Sıralamalar

EdgeNodeType

Edge node type enum

Vakalar:

Sembol Değer Dize
OPERAND 0 işlenen
SONUÇ 1 sonuç

PropagationDirection

Yayılma yönü numaralandırması

Vakalar:

Sembol Değer Dize
YOK 0 YOK
FORWARD 1 FORWARD
GERİ 2 GERİ
BOTH 3 BOTH

ReductionOp

Reduction op enum

Vakalar:

Sembol Değer Dize
TOPLA 0 toplam
Maksimum 1 maks.
MIN 2 dk