'sdy' lehçesi

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

İşlemler

sdy.constant (sdy::ConstantOp)

Sabit işlem

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

Şu sayfaya göz atı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:

Sonuçlar:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Veri akışı uç operasyonu

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 dikey olan birden çok 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 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 kök hedefini (hedeflerden herhangi biri olabilir ancak tercihen blok bağımsız değişkeni yerine bir op sonucu olabilir) giriş olarak alır. Bu değerin 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 işlem, bir opsiyon çok sayıda kenara sahip olduğunda kullanışlıdır ve şunları yapmak ç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:

İşlemler:

Sonuçlar:

sdy.manual_computation (sdy::ManualComputationOp)

Manuel koleksiyonlarla çoklu cihaz 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 atlayın.

Body, manual_axes ile ilgili yereldir. Yayılım, herhangi bir serbest eksende (manual_axes listesinde olmayanlar) gövde boyunca gerçekleşir.

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

Özellikler:

İşlem görenler:

Sonuçlar:

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 değerine sahip 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:

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, ShardableDataFlowOpInterface

Özellikler:

İşlemler:

Sonuçlar:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Yayılım bariyeri işlemi

Söz dizimi:

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

Bu işlem, bir kimlik operasyonu gibi çalışır ve girişle aynı değeri verir. Ancak yayılım açısından bu, yayılmanın üzerinden yalnızca belirli bir yönde akışına izin verir.

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

• FORWARD, parçalamaların yalnızca işlenenden sonuca akabileceği anlamına gelir.
• BACKWARD, bölmelerin 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 belirtilemiyor.

Özellikler: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

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

Etkiler: MemoryEffects::Effect{}

Özellikler:

İşlem görenler:

Sonuçlar:

sdy.reshard (sdy::ReshardOp)

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

Söz dizimi:

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

Giriş tenörünü, giriş tenörünün mevcut bölme işleminden farklı olan belirtilen bölme işlemiyle yeniden böler.

Hem ShardingConstraintOp hem de ReshardOp, tensöre kırma işlemi ekler. Kullanım ömrü:

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

// TODO(b/331680067). Gereksiz işlemleri kaldırmak için // bir standartlaştırma kalıbı ekleyin.

Özellikler: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

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

Etkiler: MemoryEffects::Effect{}

Özellikler:

İşlemler:

Sonuçlar:

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 şey olabilen (ancak aynı türde, ör. AnyTensor) bir değer listesi alır ve bu nedenle Shardy IR yığınının çeşitli seviyelerinde yeniden kullanılabilir.

Söz dizimi:

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

Özellikler: AlwaysSpeculatableImplTrait, Terminator

Arayüzler: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Etkiler: MemoryEffects::Effect{}

İşlemler:

sdy.sharding_constraint (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 tensöre (ör. bir matmul'un sonucu) tensörün veya kullanımlarının bir alt kümesinin bu şekilde kırılması gerektiğini belirtmek için kırma ekler.

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: Elementwise, SameOperandsAndResultType

Arayüzler: InferTypeOpInterface

Özellikler:

İşlemler:

Sonuçlar:

sdy.sharding_group (sdy::ShardingGroupOp)

Grup parçalama işlemi

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.

Özellikler:

İşlemler:

Özellikler

AxisRefAttr

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

Söz dizimi:

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

Parametreler:

DimMappingAttr

Bir boyuta ilişkin faktör dizinlerinin listesi

Tüm faktör indeksleri [0, num_factors) aralığında olmalıdır ve boş bir liste, bunun boş bir eşleme olduğunu (* ile ayrıştırılır/yazılır). Diğer bir deyişle, boyut herhangi bir faktörle eşlenmez.

Parametreler:

DimensionShardingAttr

Boyut kırma

Bir tensör boyutunun büyükten küçüğe bölüneceği eksen adlarının listesi, boyutun daha fazla bölünüp bölünemeyeceğini belirten bir boole değeri ve bu boyut bölme işleminin önceliğini belirten isteğe bağlı bir tam sayı (bölme işleminin yayılması sırasında dikkate alınır). Önceliklerin kaynağı, kullanıcının parçalama ek açıklamalarıdır. Düşük değer, daha yüksek öncelik anlamına gelir. Ek açıklamada öncelik eksik olduğunda en yüksek öncelik varsayılır.

Parametreler:

ManualAxesAttr

Söz dizimi:

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

Parametreler:

MeshAttr

Eksen ö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 bir eksen listesi ve isteğe bağlı bir cihaz kimlikleri 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 varsayılan cihaz kimliği listesi [0] olur; cihaz kimliği listesi sağlanırsa sıfırdan büyük herhangi bir tam sayı içermelidir. Buna maksimum bölme durumu diyoruz.

Maksimum olmayan tüm bölme durumlarında, bir cihaz kimliği listesi belirtilirse eksen boyutlarının çarpımı cihaz sayısıyla eşleşmelidir. Cihaz kimliği listesi belirtilmezse varsayılan 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 örgü, yayılma 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]>

Parametreler:

MeshAxisAttr

Örgüdeki adlandırılmış eksen

Söz dizimi:

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

Parametreler:

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
  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ümlenebileceğ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 eksiksizliği sağlamak içindir.

is_custom_rule, bunun bir stablehlo.custom_call işlemi için kullanıcı tarafından tanımlanan bir kural olup olmadığını belirtir. Bölmeleyici bu işlemleri nasıl böleceğini bilmediği için kullanıcının bunu belirtmesi gerekir. Özel bir kural olduğunda kural her zaman korunur/hiçbir zaman kaldırılmaz. is_custom_rule yalnızca stablehlo.custom_call işlemleri için doğru olabilir.

Parametreler:

SubAxisInfoAttr

Bu alt eksenin tam eksenden nasıl elde edildiğ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.

Parametreler:

TensorMappingAttr

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

Söz dizimi:

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

Parametreler:

TensorShardingAttr

Tensor parçalama

Söz dizimi:

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

Tensör kırma, 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 parçalamanın bağlı olduğu ağ, ilgili bir MeshOp simgesine referans veren bir simge adı veya satır içi bir MeshAttr ile belirtilebilir.

Parametreler:

TensorShardingPerValueAttr

Bir işlemin işleneni/sonucu başına tensör bölme

Söz dizimi:

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

Parametreler:

Sıralamalar

PropagationDirection

yayılım yönü sıralaması

Vakalar: