Definisi OpInterface

CollectiveOpInterface (Sdy_CollectiveOpInterface)

Antarmuka untuk semua operasi kolektif. Mengenkapsulasi get/set umum untuk atribut outSharding.

Batasan:

  • Operand harus memiliki sharding atau allowMissingInputSharding() menampilkan true.
  • out_sharding valid sehubungan dengan jenis yang sesuai.
  • Sharding operand dan hasil harus memiliki mesh yang sama jika allowDifferentMeshes() menampilkan nilai salah.
  • Peringkat yang sama untuk sharding operand dan hasil.

Metode:

getOutSharding

::mlir::sdy::TensorShardingAttr getOutSharding();

Menampilkan sharding tensor output dari operasi kolektif.

setOutShardingAttr

void setOutShardingAttr(::mlir::sdy::TensorShardingAttr sharding);

Menetapkan sharding tensor output dari operasi kolektif.

getTensor

::mlir::TypedValue<::mlir::TensorType> getTensor();

Dapatkan operand tensor dari operasi kolektif.

getType

::mlir::Type getType();

Mendapatkan jenis hasil operasi kolektif.

allowDifferentMeshes

bool allowDifferentMeshes();

Menunjukkan apakah operasi kolektif memungkinkan sharding input dan output memiliki mesh yang berbeda.

allowMissingInputSharding

bool allowMissingInputSharding();

Menunjukkan apakah operasi kolektif memungkinkan input tidak memiliki sharding, yaitu direplikasi sepenuhnya secara implisit.

ShardableDataFlowOpInterface (Sdy_ShardableDataFlowOpInterface)

Antarmuka op yang memungkinkan shardy menyebarkan sharding melalui tepi aliran data op yang memperluas antarmuka ini.

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.

Pemilik adalah target tepi aliran data yang ditentukan pengguna yang digunakan oleh penyebaran shardy. Pengguna dapat memilihnya secara sewenang-wenang, tetapi harus statis.

Contoh:

  y_1, ..., y_n = custom_op (x_1, ..., x_n)
                  ((body_arg_1,..., body_arg_n) {
                    ...
                    return return_value_1, ..., return_value_n
                  })

custom_op ini memiliki dua jenis untuk tepi aliran data, masing-masing tepi n antara return_value_i (sumber) dan y_i (target) dan tepi n antara x_i(sumber) dan body_arg_i(target). Dalam hal ini, pemilik edge sama dengan target.

Berikut adalah contoh operasi dengan beberapa target:

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

Operasi while ini 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.

Metode:

getBlockArgumentEdgeOwnerShardings

mlir::SmallVector<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

Menampilkan sharding dari semua pemilik tepi aliran data argumen blok.

setBlockArgumentEdgeOwnerShardings

void setBlockArgumentEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

Menetapkan shardings dari semua pemilik tepi argumen blok.

getOpResultEdgeOwnerShardings

mlir::SmallVector<mlir::sdy::TensorShardingAttr> getOpResultEdgeOwnerShardings();

Menampilkan sharding dari semua pemilik tepi aliran data hasil op.

setOpResultEdgeOwnerShardings

void setOpResultEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

Menetapkan shardings dari semua pemilik edge hasil op.

transformTargetSharding

mlir::sdy::TensorShardingAttr transformTargetSharding(mlir::Value target, mlir::sdy::TensorShardingAttr sharding, mlir::sdy::DataFlowShardingTransformType transformType);

Mengubah sharding target bergantung pada transformType

Lihat DataFlowShardingTransformType untuk mengetahui informasi selengkapnya.

getBlockArgumentEdgeOwners

mlir::ArrayRef<mlir::BlockArgument> getBlockArgumentEdgeOwners();

Mendapatkan semua pemilik tepi argumen blok.

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

Mendapatkan semua pemilik edge hasil op.

getEdgeSources

mlir::SmallVector<mlir::OpOperand*> getEdgeSources(mlir::Value owner);

Mendapatkan sumber tepi aliran data dengan owner tepi.

getEdgeOwnerFromTarget

mlir::Value getEdgeOwnerFromTarget(mlir::Value target);

Mendapatkan target pemilik tepi alur data dengan target yang mungkin atau mungkin bukan pemilik.

getEdgeOwnerFromSource

mlir::Value getEdgeOwnerFromSource(mlir::OpOperand&source);

Mendapatkan target pemilik tepi aliran data dengan source.

getNonEdgeOwnerTargets

mlir::SmallVector<mlir::Value> getNonEdgeOwnerTargets(mlir::Value owner);

Mendapatkan target non-pemilik tepi alur data dengan owner tepi.

ShardingRuleOpInterface (Sdy_ShardingRuleOpInterface)

Antarmuka op yang memungkinkan op menentukan aturan sharding-nya sendiri. Aturan sharding menentukan cara operasi dapat dipartisi sesuai dengan berbagai properti pada op - atribut apa pun, bentuk operand, bentuk hasil, dll. Lihat OpShardingRuleAttr untuk mengetahui detail selengkapnya.

Metode:

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

Menampilkan aturan sharding op.