Dialek 'sdy'

Dialek Shardy (SDY) menentukan representasi sharding tensor berbasis sumbu dan komponen API tambahan untuk melampirkan sharding ke tensor.

Operasi

sdy.constant (sdy::ConstantOp)

Operasi konstan

Menghasilkan tensor output dari value konstan.

Lihat: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

Contoh:

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

Ciri: AlwaysSpeculatableImplTrait

Antarmuka: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Efek: MemoryEffects::Effect{}

Atribut:

Hasil:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Operasi edge aliran data.

Sintaksis:

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

Tepi aliran data dari beberapa op X menentukan jembatan antara kumpulan sumber (masing-masing adalah operand X atau operand terminator blok X) dan kumpulan target (masing-masing adalah hasil X atau argumen blok X), sehingga semua sumber dan target harus di-shard dengan cara yang sama.

Operasi dapat memiliki beberapa tepi aliran data yang saling ortogonal.

Contoh:

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

Meskipun operasi memiliki n tepi aliran data, tepi aliran data ke-i berada di antara sumber x_i, return_value_i dan target y_i, pred_arg_i, body_arg_i.

sdy.data_flow_edge menggunakan target root tepi sebagai input (dapat berupa target apa pun, tetapi sebaiknya hasil op, bukan argumen blok), yang tidak boleh memiliki penggunaan lain. Operasi ini tidak murni karena dapat mengambil input yang awalnya tidak memiliki kegunaan apa pun.

sdy.data_flow_edge juga menyimpan sharding opsional untuk semua target tepi, dan sharding tersebut harus diperbarui, bukan sharding target (jika dapat dilampirkan) selama penyebaran. Hal ini berguna saat op memiliki banyak edge, karena jauh lebih efisien untuk:

• ditransmisikan melalui setiap edge secara terpisah.
• memperbarui sharding setiap edge secara terpisah, bukan semua target sekaligus (misalnya, op memiliki satu TensorShardingPerValueAttr yang tidak dapat diubah untuk sharding hasil).
• tambahkan setiap tepi ke daftar kerja secara terpisah saat sharding sumber telah berubah.

Propagasi akan menyebarkan sharding di antara semua sumber dan target sdy.data_flow_edge seolah-olah itu adalah operasi reguler dengan sumber sebagai operand dan target sebagai hasil, serta identitas sdy.op_sharding_rule. Artinya, penyebaran maju adalah dari sumber ke target dan penyebaran mundur adalah dari target ke sumber.

Kami tidak mengizinkan input sdy.data_flow_edge ditentukan oleh op SdyDialect, sehingga kita dapat mengasumsikan bahwa input tersebut ditentukan oleh op yang memiliki atribut sdy.sharding yang tidak terdaftar.

Ciri: SameOperandsAndResultType

Antarmuka: InferTypeOpInterface

Atribut:

Operand:

Hasil:

sdy.manual_computation (sdy::ManualComputationOp)

Operasi paralelisme multiperangkat dengan kolektif manual

Sintaksis:

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)

Masuk ke wilayah yang ditulis dalam kode lokal per perangkat dengan kolektif eksplisit, dengan bentuk logis yang cocok dengan bentuk buffering fisik per perangkat lokal dan kolektif yang sama persis dengan komunikasi lintas perangkat fisik.

Isi adalah lokal terkait manual_axes. Propagasi akan terjadi melalui body pada sumbu bebas apa pun - yang tidak ada dalam daftar manual_axes.

Sifat: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Atribut:

Operand:

Hasil:

sdy.mesh (sdy::MeshOp)

Mesh bernama

Sintaksis:

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

Menentukan mesh baru yang diberi nama. Semua mesh dalam modul harus memiliki jumlah perangkat yang sama (kecuali untuk mesh dengan satu device_id). Mesh adalah operasi Symbol yang muncul di SymbolTable modul dan dapat direferensikan oleh name-nya.

Ciri: HasParent<ModuleOp>

Antarmuka: Symbol

Atribut:

sdy.named_computation (sdy::NamedComputationOp)

Operasi komputasi bernama

Sintaksis:

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)

Mengelompokkan komputasi, yaitu blok operasi, dan memberinya nama. Propagasi akan mengalir masuk/keluar dari region seolah-olah semuanya telah disisipkan.

Ini dapat digunakan untuk menangani penyebaran melalui petunjuk panggilan ke fungsi lain. Setiap pengguna Shardy harus menulis kartu impor/ekspor yang mengonversi operasi panggilan mereka menjadi operasi sdy.named_computation, menduplikasi/menyalin isi fungsi yang dipanggil ke dalam isi named_computation.

Jenis setiap argumen blok dan nilai yang ditampilkan di region harus sama dengan jenis operand dan jenis hasil op.

Contoh:

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

Sifat: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Antarmuka: ConditionallySpeculatable, ShardableDataFlowOpInterface

Atribut:

Operand:

Hasil:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Operasi penghalang propagasi

Sintaksis:

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

Operasi ini beroperasi seperti operasi identitas, yang menghasilkan nilai yang sama dengan yang diperlukan sebagai input. Namun dalam hal propagasi, hal ini hanya akan memungkinkan propagasi mengalir melaluinya dalam arah tertentu.

Hal ini mencegah sharding disebarkan di antara penggunaan hasil operasi penghalang dan operandnya.

• FORWARD berarti sharding hanya dapat mengalir dari operand ke hasil.
• BACKWARD berarti sharding hanya dapat mengalir dari hasil ke operand.
• NONE berarti tidak ada sharding yang dapat di-propagasi melalui operasi ini.
• Tidak dapat menentukan BOTH, karena operasi ini akan berlebihan.

Karakteristik: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

Antarmuka: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Efek: MemoryEffects::Effect{}

Atribut:

Operand:

Hasil:

sdy.reshard (sdy::ReshardOp)

Membuat ulang sharding tensor ke sharding yang berbeda

Sintaksis:

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

Melakukan resharding pada tensor input dengan sharding yang ditentukan, yang berbeda dengan sharding tensor input yang ada.

ShardingConstraintOp dan ReshardOp melampirkan sharding ke tensor. Masa aktifnya adalah:

1. Sebelum penerapan sharding, ShardingConstraintOp ditambahkan oleh pengguna.
2. Penyebaran sharding menggunakan ShardingConstraintOp. Tidak ada ShardingConstraintOp dalam hasil propagasi sharding. Sebagai gantinya, ReshardOp dapat ditambahkan jika diperlukan.
3. Partisioner mengonversi ReshardOp menjadi operasi kolektif (atau operasi identitas). Tidak boleh ada ReshardOp dalam hasil partisi.

// TODO(b/331680067). Tambahkan pola kanonisisasi untuk menghapus operasi // reshard yang redundan.

Sifat: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

Antarmuka: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Efek: MemoryEffects::Effect{}

Atribut:

Operand:

Hasil:

sdy.return (sdy::ReturnOp)

Operasi sdy.return menghentikan region yang dilampirkan ke operasi berbasis region sdy dan operasi berbasis region Shardy lainnya. Fungsi ini variadik: fungsi ini menggunakan daftar nilai sebagai argumen yang jenisnya dapat berupa apa saja (tetapi dari jenis yang sama, misalnya AnyTensor) sehingga dapat digunakan kembali di berbagai tingkat stack IR Shardy.

Sintaksis:

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

Sifat: AlwaysSpeculatableImplTrait, Terminator

Antarmuka: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Efek: MemoryEffects::Effect{}

Operand:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Membatasi tensor ke sharding yang ditentukan

Sintaksis:

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

Melampirkan sharding ke tensor perantara (misalnya, hasil matmul) untuk menunjukkan bahwa ini adalah cara tensor tersebut, atau sebagian penggunaannya, harus di-sharding.

Jika sharding memiliki dimensi terbuka dan sumbu yang tidak dibatasi, berarti tensor dapat di-sharding lebih lanjut di sepanjang dimensi terbuka.

Operasi ini dapat:

• Tidak memiliki penggunaan (tergantung) - yang berarti sharding yang dilampirkan adalah cara tensor input itu sendiri harus di-sharding.
• Memiliki penggunaan - yang berarti sharding yang dilampirkan adalah cara penggunaan op batasan sharding harus di-sharding, sedangkan penggunaan lain dari tensor input mungkin memiliki sharding yang berbeda (jika tensor input tidak memiliki penggunaan lain, perilakunya sama dengan kasus tidak ada penggunaan).

Karakteristik: Elementwise, SameOperandsAndResultType

Antarmuka: InferTypeOpInterface

Atribut:

Operand:

Hasil:

sdy.sharding_group (sdy::ShardingGroupOp)

Operasi grup sharding

Sintaksis:

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

Operasi ini menyediakan antarmuka untuk menetapkan tensor ke grup sharding ( grup tensor yang akan diberlakukan untuk memiliki sharding yang identik). Selama penyebaran, segera setelah satu elemen grup di-shard, semua anggota lainnya akan di-shard dengan cara yang sama persis. Operasi ini mengambil ID grup argumen dan tidak menampilkan hasil, tetapi mengubah representasi grup sharding internal untuk menambahkan tensor input ke grup dengan ID yang diberikan.

Atribut:

Operand:

Atribut

AxisRefAttr

Referensi ke sumbu penuh atau sub-sumbu terpisah

Sintaksis:

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

Parameter:

DimMappingAttr

Daftar indeks faktor untuk dimensi

Semua indeks faktor harus berada dalam rentang [0, num_factors) dan daftar kosong menunjukkan bahwa ini adalah pemetaan null (ini diuraikan/dicetak dengan *), yaitu dimensi tidak dipetakan ke faktor apa pun.

Parameter:

DimensionShardingAttr

Sharding dimensi

Daftar nama sumbu untuk membuat shard dimensi tensor dari utama ke minor, Boolean yang menunjukkan apakah dimensi dapat di-shard lebih lanjut, dan bilangan bulat opsional yang menunjukkan prioritas sharding dimensi ini, yang akan diikuti selama penyebaran sharding. Prioritas berasal dari anotasi sharding pengguna dan nilai yang lebih rendah menunjukkan prioritas yang lebih tinggi. Prioritas tertinggi diasumsikan jika prioritas tidak ada dalam anotasi.

Parameter:

ManualAxesAttr

Sintaksis:

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

Parameter:

MeshAttr

Mesh sumbu dan daftar perangkat

Sintaksis:

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

Mesh adalah daftar sumbu dan daftar ID perangkat opsional yang menentukan pengurutan perangkat.

Jika daftar sumbu kosong, mesh memiliki sumbu tanpa nama yang implisit berukuran 1. Dalam hal ini, jika daftar ID perangkat tidak diberikan, daftar ID perangkat implisit adalah [0]; jika daftar ID perangkat diberikan, daftar tersebut harus berisi satu bilangan bulat dari nilai non-negatif. Kami menyebutnya kasus sharding maksimum.

Untuk semua kasus sharding non-maksimum, jika daftar ID perangkat ditentukan, produk ukuran sumbu harus cocok dengan jumlah perangkat. Jika daftar ID perangkat tidak ditentukan, daftar ID perangkat implisit adalah iota(product(axes)). Untuk memudahkan, kami juga tidak mengizinkan penentuan daftar ID perangkat yang sama dengan iota(product(axes)); dalam hal ini, daftar ID perangkat tidak boleh ditentukan.

Berikut beberapa contoh mesh:

• Mesh kosong mewakili mesh placeholder yang dapat diganti selama penyebaran: <[]>
• Mesh dengan sumbu tanpa nama dan ID perangkat eksplisit, yang biasanya digunakan untuk merepresentasikan sharding maksimum: <[], device_ids=[3]>
• Mesh dengan dua sumbu dan ID perangkat implisit iota(6): <["a"=2, "b"=3]>
• Mesh dengan dua sumbu dan ID perangkat eksplisit yang menentukan urutan perangkat: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Parameter:

MeshAxisAttr

Sumbu bernama dalam mesh

Sintaksis:

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

Parameter:

OpShardingRuleAttr

Menentukan cara operasi dapat dipartisi.

Sintaksis:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  bool   # is_custom_rule
>

Aturan sharding menentukan cara operasi dapat dipartisi sesuai dengan berbagai properti pada op - atribut apa pun, bentuk operand, bentuk hasil, dll. Misalnya:

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

Perhatikan bahwa kami mengizinkan faktor dengan ukuran 1 meskipun tidak dapat di-shard, hal ini terutama untuk kelengkapan karena banyak operasi seperti operasi pointwise memiliki dimensi ukuran satu yang sesuai di seluruh operand dan hasil.

is_custom_rule menjelaskan apakah ini adalah aturan yang ditentukan oleh pengguna untuk operasi stablehlo.custom_call. Partisioner tidak tahu cara mempartisi operasi ini, sehingga pengguna harus memberi tahu caranya. Jika aturan tersebut adalah aturan kustom, aturan tersebut akan selalu dipertahankan/tidak pernah dihapus. is_custom_rule hanya dapat bernilai benar untuk operasi stablehlo.custom_call.

Parameter:

SubAxisInfoAttr

Info tentang cara sub-sumbu ini berasal dari sumbu penuh

Sintaksis:

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

Saat membagi sumbu penuh menjadi n sub-sumbu, sumbu akan dibentuk ulang menjadi [k_1,...,k_n], dan sub-sumbu ke-i dapat dinyatakan dengan produk semua ukuran sumbu di sebelah kirinya m=prod(k_1,...,k_(i-1)) (alias pra-ukuran) dan ukuran k_i. Oleh karena itu, atribut sub-axis-info menyimpan dua angka tersebut dan ditunjukkan sebagai berikut: (m)k untuk ukuran awal m dan ukuran k.

Parameter:

TensorMappingAttr

Pemetaan faktor untuk setiap dimensi tensor.

Sintaksis:

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

Parameter:

TensorShardingAttr

Sharding tensor

Sintaksis:

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

Sharding tensor terikat ke mesh tertentu, dan hanya dapat mereferensikan nama sumbu dari mesh tersebut. Sharding dimensi akan memberi tahu kita untuk setiap dimensi tensor, beserta sumbu (atau sub-sumbu) yang di-sharding dari besar ke kecil. Semua sumbu lain yang tidak melakukan sharding dimensi dapat direplikasi secara implisit atau eksplisit (jika muncul dalam daftar sumbu replika).

Mesh yang terikat dengan sharding ini dapat ditentukan oleh nama simbol, yang mereferensikan simbol MeshOp yang sesuai, atau MeshAttr yang disisipkan.

Parameter:

TensorShardingPerValueAttr

Sharding tensor per operand/hasil operasi

Sintaksis:

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

Parameter:

Enum

PropagationDirection

enum arah penyebaran

Kasus: