Định nghĩa OpInterface

CollectiveOpInterface (Sdy_CollectiveOpInterface)

Giao diện cho tất cả các thao tác tập thể. Đóng gói phương thức get/set phổ biến cho thuộc tính outSharding.

Các quy tắc ràng buộc:

  • Toán hạng phải có phân đoạn hoặc allowMissingInputSharding() trả về giá trị true.
  • out_sharding hợp lệ so với loại tương ứng.
  • Phân đoạn toán hạng và kết quả phải có cùng một lưới nếu allowDifferentMeshes() trả về giá trị false.
  • Cùng thứ hạng cho toán hạng và phân đoạn kết quả.

Phương thức:

getOutSharding

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

Trả về phân đoạn tensor đầu ra của toán tử tập thể.

setOutShardingAttr

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

Đặt tính năng phân đoạn tensor đầu ra của toán tử tập hợp.

getTensor

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

Lấy toán hạng tensor của toán tử tập hợp.

getType

::mlir::Type getType();

Lấy loại kết quả của toán tử tập hợp.

allowDifferentMeshes

bool allowDifferentMeshes();

Cho biết liệu toán tử tập hợp có cho phép phân đoạn đầu vào và đầu ra có lưới khác nhau hay không.

allowMissingInputSharding

bool allowMissingInputSharding();

Cho biết liệu toán tử tập hợp có cho phép đầu vào không phân đoạn hay không, tức là được sao chép đầy đủ một cách ngầm ẩn.

ShardableDataFlowOpInterface (Sdy_ShardableDataFlowOpInterface)

Giao diện toán tử cho phép phân đoạn truyền bá các phân đoạn thông qua các cạnh luồng dữ liệu của các toán tử mở rộng giao diện này.

Cạnh luồng dữ liệu của một số toán tử X xác định một cầu nối giữa một tập hợp các nguồn (mỗi nguồn là một toán hạng của X hoặc một toán hạng của trình kết thúc khối của X) và một tập hợp các mục tiêu (mỗi mục tiêu là kết quả của X hoặc đối số khối của X), sao cho tất cả các nguồn và mục tiêu phải được phân đoạn theo cùng một cách. Một toán tử có thể có nhiều cạnh luồng dữ liệu vuông góc với nhau.

Chủ sở hữu là một mục tiêu do người dùng chỉ định của cạnh luồng dữ liệu mà tính năng truyền tải của shardy sử dụng. Người dùng có thể chọn giá trị này tuỳ ý nhưng giá trị này cần phải tĩnh.

Ví dụ:

  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 này có hai loại cạnh luồng dữ liệu, mỗi loại có n cạnh giữa return_value_i (nguồn) và y_i (đích) và n cạnh giữa x_i(nguồn) và body_arg_i(đích). Trong trường hợp này, chủ sở hữu cạnh giống với mục tiêu.

Dưới đây là ví dụ về một thao tác có nhiều mục tiêu:

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

Toán tử while này có n cạnh luồng dữ liệu, cạnh luồng dữ liệu thứ i nằm giữa nguồn x_i, return_value_i và đích y_i, pred_arg_i, body_arg_i.

Phương thức:

getBlockArgumentEdgeOwnerShardings

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

Trả về các phân đoạn của tất cả chủ sở hữu cạnh luồng dữ liệu đối số khối.

setBlockArgumentEdgeOwnerShardings

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

Đặt shardings của tất cả các chủ sở hữu cạnh đối số khối.

getOpResultEdgeOwnerShardings

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

Trả về các phân đoạn của tất cả chủ sở hữu cạnh luồng dữ liệu kết quả hoạt động.

setOpResultEdgeOwnerShardings

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

Đặt shardings của tất cả các chủ sở hữu cạnh kết quả hoạt động.

transformTargetSharding

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

Biến đổi sharding của mục tiêu tuỳ thuộc vào transformType

Vui lòng xem DataFlowShardingTransformType để biết thêm thông tin.

getBlockArgumentEdgeOwners

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

Lấy tất cả chủ sở hữu cạnh đối số khối.

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

Lấy tất cả chủ sở hữu cạnh kết quả hoạt động.

getEdgeSources

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

Lấy nguồn cạnh luồng dữ liệu cho cạnh owner.

getEdgeOwnerFromTarget

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

Lấy target chủ sở hữu của một cạnh luồng dữ liệu, với target có thể là hoặc không phải là chủ sở hữu.

getEdgeOwnerFromSource

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

Lấy mục tiêu chủ sở hữu của một cạnh luồng dữ liệu cho một source.

getNonEdgeOwnerTargets

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

Lấy các mục tiêu không phải chủ sở hữu của một cạnh luồng dữ liệu dựa trên cạnh owner.

ShardingRuleOpInterface (Sdy_ShardingRuleOpInterface)

Giao diện toán tử cho phép toán tử xác định quy tắc phân đoạn riêng. Quy tắc phân đoạn chỉ định cách phân vùng một toán tử theo nhiều thuộc tính trên toán tử – bất kỳ thuộc tính nào, hình dạng của toán hạng, hình dạng của kết quả, v.v. Hãy xem OpShardingRuleAttr để biết thêm thông tin chi tiết.

Phương thức:

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

Trả về quy tắc phân đoạn của toán tử.