Definisi OpInterface

ShardableDataFlowOpInterface (ShardableDataFlowOpInterface)

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

Tepi aliran data pada beberapa operasi X menentukan jembatan antara sekumpulan sumber (masing-masing merupakan operand dari X atau operand dari terminator blok X) dan kumpulan target (masing-masing adalah hasil dari X atau argumen blok X), sehingga semua sumber dan target harus di-sharding 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 penerapan 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::ArrayRef<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

Menampilkan sharding dari semua pemilik tepi aliran data argumen blok.

setBlockArgumentEdgeOwnerSharding

void setBlockArgumentEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);

Menetapkan sharding pemilik tepi argumen blok dengan index yang diberikan.

setBlockArgumentEdgeOwnerShardings

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

Menetapkan shardings dari semua pemilik tepi argumen blok.

getOpResultEdgeOwnerShardings

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

Menampilkan sharding dari semua pemilik edge aliran data hasil op.

setOpResultEdgeOwnerSharding

void setOpResultEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);

Menetapkan sharding pemilik tepi hasil op dengan index yang diberikan.

setOpResultEdgeOwnerShardings

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

Menetapkan shardings dari semua pemilik tepi hasil op.

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::Value> getEdgeSources(mlir::Value target);

Mendapatkan sumber edge aliran data dengan nilai target.

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 (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.