OpInterface 정의

ShardableDataFlowOpInterface (ShardableDataFlowOpInterface)

shardy가 이 인터페이스를 확장하는 작업의 데이터 흐름 가장자리를 통해 샤딩을 전파할 수 있는 작업 인터페이스입니다.

일부 작업 X의 데이터 흐름 에지는 소스 집합(각각 X의 피연산자 또는 X의 블록 종결자 피연산자)과 타겟 집합 (각각 X의 결과 또는 X의 블록 인수) 간의 브리지를 정의합니다. 그러면 모든 소스와 타겟이 동일한 방식으로 샤딩되어야 합니다. 연산에는 서로 직교하는 여러 데이터 흐름 가장자리가 있을 수 있습니다.

소유자는 Shardy의 전파에 사용되는 데이터 흐름 가장자리의 사용자 지정 타겟입니다. 사용자는 임의로 선택할 수 있지만 정적여야 합니다.

예를 들면 다음과 같습니다.

  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에는 return_value_i (소스)와 y_i (타겟) 사이에 있는 n개의 에지와 x_i(소스)와 body_arg_i(타겟) 사이의 에지가 n개 있는 데이터 흐름 에지의 두 가지 유형이 있습니다. 이 경우 가장자리 소유자는 타겟과 동일합니다.

다음은 여러 타겟이 있는 작업의 예입니다.

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

이 while 연산에는 n개의 데이터 흐름 가장자리가 있으며, i번째 데이터 흐름 가장자리가 소스 x_i, return_value_i와 타겟 y_i, pred_arg_i, body_arg_i 사이입니다.

메서드

getBlockArgumentEdgeOwnerShardings

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

모든 블록 인수 데이터 흐름 가장자리 소유자의 샤딩을 반환합니다.

setBlockArgumentEdgeOwnerSharding

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

지정된 index를 사용하여 블록 인수 에지 소유자의 sharding를 설정합니다.

setBlockArgumentEdgeOwnerShardings

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

모든 블록 인수 가장자리 소유자의 shardings를 설정합니다.

getOpResultEdgeOwnerShardings

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

모든 op 결과 데이터 흐름 에지 소유자의 샤딩을 반환합니다.

setOpResultEdgeOwnerSharding

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

지정된 index로 op 결과 가장자리 소유자의 sharding를 설정합니다.

setOpResultEdgeOwnerShardings

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

모든 op 결과 에지 소유자의 shardings를 설정합니다.

getBlockArgumentEdgeOwners

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

모든 블록 인수 에지 소유자를 가져옵니다.

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

모든 op 결과 에지 소유자를 가져옵니다.

getEdgeSources

mlir::SmallVector<mlir::Value> getEdgeSources(mlir::Value target);

target 값을 사용하여 데이터 흐름 가장자리 소스를 가져옵니다.

getEdgeOwnerFromTarget

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

소유자일 수도 있고 아닐 수도 있는 target를 기준으로 데이터 흐름 가장자리의 소유자 target를 가져옵니다.

getEdgeOwnerFromSource

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

source가 지정된 데이터 흐름 에지의 소유자 타겟을 가져옵니다.

getNonEdgeOwnerTargets

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

owner 가장자리가 주어질 때 데이터 흐름 가장자리의 소유자가 아닌 타겟을 가져옵니다.

ShardingRuleOpInterface (ShardingRuleOpInterface)

작업이 자체 샤딩 규칙을 정의할 수 있는 작업 인터페이스입니다. 샤딩 규칙은 op의 다양한 속성(속성, 피연산자의 모양, 결과의 모양 등)에 따라 작업을 파티셔닝하는 방법을 지정합니다. 자세한 내용은 OpShardingRuleAttr를 참고하세요.

메서드

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

연산의 샤딩 규칙을 반환합니다.