'sdy' lehçesi

Shardy (SDY) lehçesi, eksene dayalı bir tensör bölme temsili ve bölme işlemlerini tenörlere eklemek için ek API bileşenleri tanımlar.

İşlemler

sdy.all_gather (sdy::AllGatherOp)

Eksenler boyunca tümünü toplama iletişimi gerçekleştirir

Söz dizimi:

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

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

gathering_axes, eksen listelerinin bir listesidir. Dış liste, tensörün boyutlarının üzerindedir. Her iç liste, ilgili boyutta ayrı bir toplama işleminin yapılması gereken eksenleri belirtir. Sonuç (out_sharding) için parçalamayı elde etmek üzere operandın (tensor) parçalamasına uygulanır.

Sonuçların bölümlendirilmesini belirlemek için out_sharding değerinin kullanılmadığını unutmayın. Bunun yerine, sonucun bölümlenmesi, işlenenin ve gathering_axes'ün bölümlenmesiyle belirlenir ve out_sharding bu tahmini bölümlemeyle 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'teki öğeler, AxisRefListAttr'te listelenen kısıtlamaları karşılamalıdır.
  • Operand bölme işlemine gathering_axes uygulandığında out_sharding elde edilir.

Özellikler: SameOperandsAndResultType

Arayüzler: InferTypeOpInterface, Sdy_CollectiveOpInterface

Özellikler:

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

İşlemler:

İşlem Gören Açıklama
tensor Herhangi bir türdeki değerlerin tensörü

Sonuçlar:

Sonuç Açıklama
result Herhangi bir türdeki değerlerin tensörü

sdy.all_reduce (sdy::AllReduceOp)

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

Söz dizimi:

operation ::= `sdy.all_reduce` $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 bölümünde listelenen kısıtlamaları karşılamalıdır;
  • reduction_axes, operand bölme eksenleriyle çakışmamalıdır;

Özellikler: SameOperandsAndResultType

Arayüzler: CollectiveOpInterface, InferTypeOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
reduction_axes::mlir::sdy::AxisRefListAttrEksen referanslarının listesi
out_sharding::mlir::sdy::TensorShardingAttrTensör bölme

İşlemler:

İşlem Gören Açıklama
tensor Herhangi bir türdeki değerlerin tensörü

Sonuçlar:

Sonuç Açıklama
result Herhangi bir türdeki değerlerin tensörü

sdy.all_slice (sdy::AllSliceOp)

Dinamik dilimleme işlemini eksenler boyunca 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 dilimlere ayırır. sdy.all_slice ile sdy.all_gather arasında cebirsel bir dualite vardır.

slicing_axes, eksen listelerinin bir listesidir. Dış liste, tensörün boyutlarının üzerindedir. Her iç liste, ilgili boyutta bir dilimin hangi eksenler boyunca gerçekleştirileceğini belirtir. Sonuç (out_sharding) için parçalamayı elde etmek üzere operandın (tensor) parçalamasına uygulanır.

Sonuçların bölümlendirilmesini belirlemek için out_sharding değerinin kullanılmadığını unutmayın. Bunun yerine, sonucun bölümlenmesi, operatörün ve slicing_axes'ün bölümlenmesine göre belirlenir ve out_sharding, bu tahmini bölümlemeyle 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:

  • slicing_axes'teki öğeler, AxisRefListAttr'te listelenen kısıtlamaları karşılamalıdır.
  • Sdy_CollectiveOpInterface bölümünde listelenen kısıtlamaları karşılamalıdır.
  • Operand bölme işlemine slicing_axes uygulandığında out_sharding elde edilir.

Özellikler: SameOperandsAndResultType

Arayüzler: CollectiveOpInterface, InferTypeOpInterface

Özellikler:

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

İşlemler:

İşlem Gören Açıklama
tensor Herhangi bir türdeki değerlerin tensörü

Sonuçlar:

Sonuç Açıklama
result Herhangi bir türdeki değerlerin tensörü

sdy.all_to_all (sdy::AllToAllOp)

Eksenler boyunca her bir cihazın diğer tüm cihazlarla iletişim kurmasını sağlar

Söz dizimi:

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

Bu işlem, parametre listesindeki her (axes, src_dim, tgt_dim) tuple için bir tensörün parçalarını tgt_dim boyutu ve axes'te belirtilen eksenler boyunca dilimlendirir, bu parçaları eksenler boyunca dağıtır ve src_dim boyutu boyunca birleştirir.

Bu işlem temel olarak src_dim ve axes boyunca bir tüm toplama işleminin ardından tgt_dim ve axes boyunca bir tüm dilimleme işleminin birleşimidir. Yani, giriş tenzorundaki eksen bölme boyutu src_dim'nin son eki, çıkış tenzorundaki eksen bölme boyutu tgt_dim'ye eklenir.

Sonuç (out_sharding) için parçalara ayırma işlemini elde etmek üzere tümünü tümüne işlemi, operanın (tensor) parçalara ayırma işlemine uygulanır.

Sonuçların bölümlendirilmesini belirlemek için out_sharding değerinin kullanılmadığını unutmayın. Bunun yerine, sonucun bölümlenmesi, operandın (src_dim, tgt_dim ve axes) bölümlenmesiyle belirlenir ve out_sharding bu tahmin edilen bölümlemeyle 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ş bırakılmamalıdır.
  • params içindeki her parametre için:
    • axes içindeki öğeler AxisRefAttr'un kısıtlamalarını karşılamalıdır.
    • src_dim ve tgt_dim geçerli boyutlar (negatif olmayan ve tensörün rütbesinden az) olmalıdır.
    • Tüm src_dim veya tgt_dim değerleri tüm parametreler arasında benzersiz olmalıdır.
    • src_dim tüm parametrelerde artan düzende sıralanmış olmalıdır.
  • Operand bölme işleminde axes, src_dim'den tgt_dim'ye taşındığında out_sharding elde edilir.

Özellikler: SameOperandsAndResultType

Arayüzler: InferTypeOpInterface, Sdy_CollectiveOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
params::mlir::sdy::AlltoAllParamListAttrHerkese açık parametrelerin listesi
out_sharding::mlir::sdy::TensorShardingAttrTensör bölme

İşlemler:

İşlem Gören Açıklama
tensor Herhangi bir türdeki değerlerin tensörü

Sonuçlar:

Sonuç Açıklama
result Herhangi bir türdeki değerlerin tensörü

sdy.collective_permute (sdy::CollectivePermuteOp)

Eksenler yerine kolektif 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ü bölen eksenleri yeniden sıralamak/değiştirmek için giriş tensörünün bir bölümünü her cihazdan diğerine gönderir.

Toplu permütasyon, giriş bölme işlemini her boyutun önceki gibi bölüneceği şekilde dönüştürebilir. Yani boyutlar, boyutlarının çarpımı daha önce tenzoru bölen eksenlerin çarpımıyla eşleşen eksenler boyunca bölünmelidir.

Bu, tek bir boyuttaki veya farklı boyutlardaki eksenleri yeniden sıralamak ve parçalara ayrılmış eksenleri kopyalanan eksenlerle değiştirmek için yararlıdır.

Aşağıdaki örnekte, parçalara ayrılmış tenör boyutu tensor<1x4x2xf32>'tü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ış bölmelerinin farklı ağları varsa 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'teki bölme ekseni boyutlarının çarpımı, ilgili operand boyut bölme boyutuyla eşleşmelidir.

Özellikler: SameOperandsAndResultType

Arayüzler: CollectiveOpInterface, InferTypeOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
out_sharding::mlir::sdy::TensorShardingAttrTensör bölme

İşlemler:

İşlem Gören Açıklama
tensor Herhangi bir türdeki değerlerin tensörü

Sonuçlar:

Sonuç Açıklama
result Herhangi bir türdeki değerlerin tensörü

sdy.constant (sdy::ConstantOp)

Sabit işlem

Sabit bir value değerinden output tensörü oluşturur.

Ş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 Her türden değerin 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)

Bir X işleminin veri akışı kenarı, bir kaynak kümesi (her biri X'in bir operandı veya X'in blok sonlandırıcısı) ile bir hedef kümesi (her biri X'in bir sonucu veya X'in bir blok bağımsız değişkeni) arasında bir köprü tanımlar. Bu köprü, tüm kaynakların ve hedeflerin aynı şekilde bölümlenmesini sağlar.

Bir işlem, birbirine dik olan birden fazla veri akışı kenarı içerebilir.

Ö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 sırada işlem n tane veri akışı kenarı içerir. i. veri akışı kenarı, x_i, return_value_i kaynakları ile y_i, pred_arg_i, body_arg_i hedefleri arasındadır.

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

sdy.data_flow_edge, kenardaki tüm hedefler için isteğe bağlı bir bölme de içerir ve bu bölme, dağıtım sırasında hedeflerin bölme yerine (eklenebilirse) güncellenmelidir. Bu, bir op'nin çok fazla kenarı olduğunda kullanışlıdır. Çünkü aşağıdakiler çok daha verimlidir:

  • her kenar boyunca ayrı ayrı yayılır.
  • Tüm hedefleri tek seferde güncellemek yerine her kenar için ayrı ayrı bölümlendirmeyi güncelleyin (ör. bir işlem, sonuç bölümlendirmeleri için tek bir değişmez TensorShardingPerValueAttr değerine sahiptir).
  • Bir kaynağın bölümlenmesi değiştiğinde her kenarı ayrı ayrı iş listesine ekleyin.

Yayma işlemi, sdy.data_flow_edge'ün tüm kaynakları ve hedefleri arasında, kaynaklar operatör, hedefler sonuç ve sdy.op_sharding_rule kimliği normal bir işlemmiş gibi parçalama işlemlerini dağıtır. Yani ileri yayılma kaynaklardan hedeflere, geri yayılma ise hedeflerden kaynaklara olur.

sdy.data_flow_edge girişinin bir SdyDialect işlemi tarafından tanımlanmasına izin verilmez. Bu nedenle, kayıtlı olmayan sdy.sharding özelliğine sahip bir işlem tarafından tanımlandığını varsayabiliriz.

Özellikler: SameOperandsAndResultType

Arabirimler: InferTypeOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
sharding::mlir::sdy::TensorShardingAttrTensör bölme

İşlemler:

İşlem Gören Açıklama
input Herhangi bir türde değerden oluşur.

Sonuçlar:

Sonuç Açıklama
result Herhangi bir türde değerden oluşur.

sdy.manual_computation (sdy::ManualComputationOp)

Manuel koleksiyonlarla çok cihazlı paralellik 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)

Mantıksal şekillerin cihaz başına yerel fiziksel arabellek şekilleriyle eşleştiği ve koleksiyonların tam olarak cihazlar arası fiziksel iletişime karşılık geldiği, cihaz başına yerel kod terimleriyle yazılmış bir bölgeye açık koleksiyonlarla atlayın.

Body, manual_axes ile ilgili yereldir. Yayma işlemi, body aracılığıyla tüm serbest eksenlerde (manual_axes listesinde olmayanlar) gerçekleşir.

Kısıtlamalar:

  • in_shardings ve out_shardings'teki öğeler, TensorShardingAttr'de listelenen kısıtlamaları karşılamalıdır.
  • İşlem bölgesinin küresel ve yerel tensör giriş/çıkışlarının sayısı eşleşmelidir.
  • Manuel eksenler, her boyut bölme işleminde ücretsiz eksenlerden önce gelmelidir.
  • Manuel eksenler dolgu ekleyemez. Yani boyut boyutu, ilgili manuel eksen boyutuna bölünebilir olmalıdır.
  • op regions bağımsız değişkenlerinin/sonuçlarının genel ve yerel şekilleri eşleşmelidir.
  • Manuel eksenler bölünmez.

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

Arabirimler: ShardableDataFlowOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
in_shardings::mlir::sdy::TensorShardingPerValueAttrBir işlemin operand/sonuç başına tensör bölme
out_shardings::mlir::sdy::TensorShardingPerValueAttrBir işlemin operand/sonuç başına tensör bölme
manual_axes::mlir::sdy::ManualAxesAttrManuelComputationOp'un manuel olduğu eksenlerin listesi

İşlemler:

İşlem Gören Açıklama
tensors Her türden değere sahip sıralı tenzorun değişken sayısı

Sonuçlar:

Sonuç Açıklama
results herhangi bir türdeki değerlerin sıralı tenörünün değişken sayısı

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 içeren ağlar hariç). Örgü, modülün SymbolTable bölümünde görünen ve name ile referans verilebilen bir Symbol işlemidir.

Özellikler: HasParent<ModuleOp>

Arabirimler: Symbol

Özellikler:

ÖzellikMLIR TürüAçıklama
sym_name::mlir::StringAttrdize özelliği
mesh::mlir::sdy::MeshAttrEksen örgüsü ve cihazların 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ı (ör. işlem bloğu) gruplandırır ve ona bir ad verir. Yayma işlemi, her şey satır içiymiş gibi bölgeye girip çıkar.

Bu, diğer işlevlere çağrı talimatları aracılığıyla yayma işlemini gerçekleştirmek için kullanılabilir. Shardy kullanıcıları, çağrı işlemlerini sdy.named_computation işlemlerine dönüştüren, çağrılan işlevin gövdesini named_computation gövdesine kopyalayan/yayan bir içe aktarma/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şlenenin türü ve işlemin 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

Özellikler:

ÖzellikMLIR TürüAçıklama
name::mlir::StringAttrdize özelliği
in_shardings::mlir::sdy::TensorShardingPerValueAttrBir işlemin operand/sonuç başına tensör bölme
out_shardings::mlir::sdy::TensorShardingPerValueAttrBir işlemin operand/sonuç başına tensör bölme

İşlemler:

İşlem Gören Açıklama
operands Her türden değişken

Sonuçlar:

Sonuç Açıklama
«unnamed» Her türden değişken

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Yayınlama engeli 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 döndürür. Ancak dağıtım açısından bu, yalnızca dağıtımın belirli bir yönde akmasına izin verir.

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

  • FORWARD, bölmelerin yalnızca operatandan sonuca akabileceği anlamına gelir.
  • BACKWARD, parçalara ayırma işlemlerinin yalnızca sonuçtan operatöre akabileceği anlamına gelir.
  • NONE, bu işlemde hiçbir bölme 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

İşlemler:

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

Sonuçlar:

Sonuç Açıklama
result Herhangi bir türdeki değerlerin sıralı tensörü

sdy.reshard (sdy::ReshardOp)

Bir tenzoru farklı bir bölme işlemine yeniden böler

Söz dizimi:

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

Giriş tenzorunu, giriş tenzorunun mevcut bölme işleminden farklı olan belirtilen bölme işlemiyle yeniden böler.

Hem ShardingConstraintOp hem de ReshardOp, bir tenöre bir bölme ekler. Bu anahtarların kullanım ömrü:

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

// TODO(b/331680067). Gereksiz // yeniden bölümlendirme işlemlerini kaldırmak için bir standartlaştırma kalıbı ekleyin.

Özellikler: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

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

Etkiler: MemoryEffects::Effect{}

Özellikler:

ÖzellikMLIR TürüAçıklama
sharding::mlir::sdy::TensorShardingAttrTensör bölme

İşlemler:

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

Sonuçlar:

Sonuç Açıklama
result Herhangi bir türdeki değerlerin tensörü

sdy.return (sdy::ReturnOp)

sdy.return işlemi, sdy bölgeye dayalı işlemlere ve diğer Shardy bölgeye dayalı işlemlere bağlı bölgeleri sonlandırır. Değişkendir: Türleri herhangi bir türde (ancak aynı türde, ör. AnyTensor) olabilen ve bu nedenle Shardy IR yığınının çeşitli seviyelerinde yeniden kullanılabilen bir değer listesi bağımsız değişken olarak alır.

Söz dizimi:

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

Özellikler: AlwaysSpeculatableImplTrait, Terminator

Arayüzler: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Etkiler: MemoryEffects::Effect{}

İşlemler:

İşlem Gören Açıklama
results Her türden değişken

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Bir tenzoru belirtilen bölmeyle sınırlandırır

Söz dizimi:

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

Bir ara tenöre (ör.bir matmul sonucu) bir bölme ekler. Bu, söz konusu tenörün veya kullanımlarının bir alt kümesinin bu şekilde bölünmesi gerektiğini belirtir.

Bölme işleminde açık boyutlar ve sınırsız eksenler varsa tensör, açık boyutlar boyunca daha da bölünebilir.

Bu işlem şunları yapabilir:

  • Kullanımı yoktur (boştadır). Bu, ekli bölme işleminin, giriş tensörünün kendisinin nasıl bölünmesi gerektiği anlamına gelir.
  • Kullanım alanı var: Bu, ekteki bölme işleminin, bölme kısıtlaması işleminin kullanım alanlarının nasıl bölünmesi gerektiği anlamına gelir. Giriş tenörünün diğer kullanım alanları farklı bir bölme işlemine sahip olabilir (Giriş tenörünün başka bir kullanım alanı yoksa davranış, kullanım alanı yok durumuyla aynıdır).

Özellikler: SameOperandsAndResultType

Arabirimler: InferTypeOpInterface

Özellikler:

ÖzellikMLIR TürüAçıklama
sharding::mlir::sdy::TensorShardingAttrTensör bölme

İşlemler:

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

Sonuçlar:

Sonuç Açıklama
result Herhangi bir türdeki değerlerin tensörü

sdy.sharding_group (sdy::ShardingGroupOp)

Gruptaki tenzorları aynı bölme işlemine sahip olacak şekilde kısıtlar.

Söz dizimi:

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

Bu işlem, tenzorları bölme gruplarına (aynı bölmelere sahip olmaları zorunlu kılınacak tenzor grupları) atama arayüzü sağlar. Yayma sırasında, bir grup öğesi bölünür bölünmez diğer tüm üyeler de tam olarak aynı şekilde bölünür. Bu işlem, bağımsız değişken grup kimliğini alır ve sonuç döndürmez. Bunun yerine, giriş tenzorunu belirtilen kimliğe sahip gruba eklemek için dahili bölme grubu temsilini değiştirir.

Arabirimler: InferTypeOpInterface

Özellikler:

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

İşlemler:

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

Özellikler

AllToAllParamAttr

Herkese açık parametre

Söz dizimi:

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

Her şeyin her şeyle karşılaştırmasını gerçekleştirmek için eksenleri ve kaynak/hedef boyutları içeren bir tuple.

Parametreler:

Parametre C++ türü Açıklama
eksenler ::llvm::ArrayRef<AxisRefAttr> tümünü tümüne göre karşılaştırma yapılacak eksenler
src_dim int64_t kaynak boyut dizini
tgt_dim int64_t hedef boyut dizini

AlltoAllParamListAttr

Herkese açık 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'ün 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 Bu bir alt eksense ek bilgi

AxisRefListAttr

Eks referanslarının listesi

Söz dizimi:

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

Kısıtlamalar:

  • value içindeki öğeler AxisRefAttr'un kısıtlamalarını karşılamalıdır.
  • Birbiriyle çakışan yinelenen eksen referansları veya alt eksenler yoktur.
  • İki bitişik eksen referansı, aynı tam eksenin art arda gelen alt eksenleri değildir. Yani, tek bir alt eksene veya tam eksene birleştirilebilirler.

Parametreler:

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

DimMappingAttr

Bir boyut için faktör dizinlerinin listesi

Boş liste, bunun null eşleme olduğunu (* ile ayrıştırılır/bası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 bunların hiçbiri 1 boyutunda olamaz.
  • Yinelenen faktör dizini 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 bölme ekseni adlarının listesi, boyutun daha fazla bölünüp bölünemeyeceğini belirten bir boole değeri ve bölme işleminin yayılması sırasında dikkate alınacak bu boyut bölme işleminin önceliğini belirten isteğe bağlı bir tam sayı. Öncelikler, kullanıcı bölme ek açıklamalarından gelir ve daha düşük bir değer daha yüksek önceliği gösterir. Ek açıklamada öncelik eksik olduğunda en yüksek öncelik varsayılır.

Kısıtlamalar:

  • axes'teki öğeler, AxisRefListAttr'te listelenen kısıtlamaları karşılamalıdır.
  • Boyut bölme işleminin önceliği varsa:
    • Öncelik 0'dan büyük veya 0'a eşit olmalıdır.
    • Kapalıysa boyut en az bir eksene sahiptir.

Parametreler:

Parametre C++ türü Açıklama
eksenler ::llvm::ArrayRef<AxisRefAttr> eksen referansları
is_closed bool Bu boyutun daha fazla bölünüp bölünemeyeceği
önceliği std::optional<int64_t> Kullanıcı önceliğine dayalı yayınlama sırasında kullanılan öncelik

ListOfAxisRefListsAttr

Eksen referans listeleri 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

ManuelComputationOp'un 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

Eksenler örgüsü ve cihazların listesi

Söz dizimi:

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

Örgü, cihaz sıralamasını belirten eksenlerin ve isteğe bağlı olarak cihaz kimliklerinin listesidir.

Eksen listesi boşsa örgünün boyutu 1 olan, adı belirtilmemiş bir gizli ekseni vardır. Bu durumda, cihaz kimliği listesi sağlanmazsa gizli cihaz kimliği listesi [0] olur; cihaz kimliği listesi sağlanırsa sıfırdan farklı bir tam sayı içermelidir. Buna maksimum bölme durumu diyoruz.

Maksimum olmayan tüm bölme işlemlerinde, bir cihaz kimliği listesi belirtilirse eksen boyutlarının çarpımı cihaz sayısıyla eşleşmelidir. Cihaz kimliği listesi belirtilmezse gizli 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, cihaz kimliği listesi belirtilmemelidir.

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

  • Boş bir ağ, dağıtım sırasında değiştirilebilecek bir yer tutucu ağı temsil eder: <[]>
  • Adsız bir eksen ve açık cihaz kimliği içeren bir ağ. Genellikle maksimum bölme işlemini temsil etmek için kullanılır: <[], device_ids=[3]>
  • İki eksen ve gizli cihaz kimlikleri içeren bir ağ iota(6): <["a"=2, "b"=3]>
  • Cihaz sıralamasını belirten iki eksen ve açık cihaz kimlikleri içeren bir ağ: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Kısıtlamalar:

  • axes içindeki öğelerin adları yinelenmemelidir.
  • device_ids belirtilmişse:
    • Eksen boyutlarının çarpımı, cihaz sayısıyla eşleşmelidir.
    • Tüm öğeleri negatif olmamalıdır.
    • device_ids, iota(product(axis_sizes)) değerine eşit olmamalıdır.
    • Sıralı device_ids, iota(product(axis_sizes)) olmalıdır.

Parametreler:

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

MeshAxisAttr

Örgüdeki 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
>

Bölme kuralında, bir işlemin işlemdeki çeşitli özelliklere (özellikler, işlenen öğelerin şekli, sonuçların şekli vb.) göre nasıl bölünebileceği belirtilir. Ö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>

Boyutu 1 olan faktörlere, bölümlenememeleri rağmen izin verildiğini unutmayın. Bu, noktasal işlemler gibi birçok işlemin, işlemler ve sonuçlar arasında karşılık gelen boyutu bir boyutlara sahip olması nedeniyle esas olarak eksiksizlik içindir.

Faktör türleri:

  • reduction_factors, nokta işleminde daraltılan boyutlar gibi azaltılması gereken faktörlerin dizinlerini içerir.
  • need_replication_factors, tam çoğaltma gerektiren faktörlerin dizinlerini (ör. bir sıralama işleminde sıralanmış boyut) içerir.
  • permutation_factors, bir dolgu işlemindeki dolgu boyutları gibi, parçalara ayrılmışsa toplu permütasyon gerektiren faktörlerin dizinlerini içerir.
  • Diğer tüm faktörler, geçiş faktörleri olarak kabul edilir. Yani, bunlarla eşlenen tüm tensörler arasında aynı şekilde bölünmüşlerse herhangi bir iletişim gerektirmeyen faktörlerdir.

blocked_propagation_factors, bölmelerin dağıtılmasına izin verilmeyen faktörleri içerir. Faktör türlerine diktir. Yani, engellenen yayılma faktörü faktör türlerinden herhangi biri olabilir.

is_custom_rule, bunun bir kullanıcı tarafından tanımlanan bir kural olup olmadığını belirtir. Kullanıcılar, özel çağrıları için bölme kuralları tanımlayabilir veya standart işlemler için önceden tanımlanmış bölme kurallarının üzerine yazabilir. Özel kurallar her zaman korunur ve hiçbir zaman kaldırılmaz.

Kısıtlamalar:

  • Operand/sonuç eşlemelerinin sayısı, işlemin operand/sonuç sayısıyla eşleşmelidir.
  • En az bir eşleme vardır (operan/sonuç içermeyen bir işlem için kural olamaz).
  • Her TensorMappingAttr'ün sıralaması, ilgili 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 bulunmamalı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> İş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> Bölmelerin dağıtılmadığı faktörler
is_custom_rule bool Kuralın stablehlo.custom_call için olup olmadığı

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] şeklinde yeniden şekillendirilir ve i. alt eksen, solundaki tüm eksen boyutlarının m=prod(k_1,...,k_(i-1)) (ön boyut olarak da bilinir) ve boyut k_i'nin çarpımıyla ifade edilebilir. Bu nedenle, alt eksen bilgisi özelliği bu iki sayıyı içerir ve aşağıdaki şekilde gösterilir: (m)k için ön boyut m ve boyut k.

Kısıtlamalar:

  • pre-size en az 1 olmalıdır.
  • size 1'den büyüktür.
  • pre-size, tam eksenin boyutunu bölmelidir. Yani hem pre-size hem de size, tam eksenin boyutunu böler ve alt eksen, tam eksenden öte geçmez.
  • Alt eksenin boyutu, ilgili tam eksenin boyutuna eşit değilse bunun yerine tam eksen kullanılmalıdır.

Parametreler:

Parametre C++ türü Açıklama
pre_size int64_t Bu alt eksenin solunda yer alan 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'teki öğeler DimMappingAttr'taki kısıtlamalara uygun olmalıdır.
  • Boyutlar arasında yinelenen faktör dizini yoktur.

Parametreler:

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

TensorShardingAttr

Tensör bölme

Söz dizimi:

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

Tensör bölme işlemi belirli bir ağa bağlıdır ve yalnızca bu ağdaki eksen adlarına referans verebilir. Boyut bölme işlemleri, tensörün her boyutu için hangi eksenler (veya alt eksenler) boyunca büyükten küçüğe bölündüğünü bize bildirir. Bir boyutu bölmeyen diğer tüm eksenler, dolaylı veya açıkça (kopyalanan eksenler listesinde görünüyorlarsa) kopyalanır.

Bu bölme işleminin bağlı olduğu ağ, ilgili bir MeshOp simgesine referans veren bir simge adı veya satır içi bir MeshAttr ile belirtilebilir.

Kısıtlamalar:

  • dim_shardings'teki öğeler, DimensionShardingAttr'te listelenen kısıtlamaları karşılamalıdır.
  • replicated_axes'teki öğeler, AxisRefListAttr'te listelenen kısıtlamaları karşılamalıdır.
  • İlgili tenör türü ShapedType değilse bölme işleminin 0. sırada olması ve kopyalanan eksen olmaması gerekir.
  • Tensör bir sıraya sahip olmalıdır.
  • Boyut bölmelerinin sayısı, tensörün rütbesine eşittir.
  • 0 boyutundaki boyutlar bölümlenmez.
  • replicated_axes öğeleri, mesh_or_ref'ye göre sıralanır (AxisRefAttr::getMeshComparator bölümüne bakın).

Parametreler:

Parametre C++ türü Açıklama
mesh_or_ref ::mlir::Attribute mesh özl veya düz örgülü simge referans özl
dim_shardings ::llvm::ArrayRef<DimensionShardingAttr> boyut bölme işlemleri
replicated_axes ::llvm::ArrayRef<AxisRefAttr> eksen referansları

TensorShardingPerValueAttr

Bir işleme ait operand/sonuç başına tensor bölme

Söz dizimi:

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

Bir işlemin her operandı/sonucu için bir tane olmak üzere TensorShardingAttr değerlerinin listesi.

Kısıtlamalar:

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

Parametreler:

Parametre C++ türü Açıklama
parçalamalar ::llvm::ArrayRef<TensorShardingAttr> Değer başına bölme

Sıralamalar

PropagationDirection

Yayınlama yönü enum

Vakalar:

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