-sdy-close-shardings

텐서 샤딩을 닫고 복제된 축을 삭제합니다.

-sdy-drop-sharding-rules

등록된 모든 작업에서 OpShardingRuleAttr를 삭제합니다.

-sdy-insert-explicit-reshards

모든 작업에 호환되는 샤딩이 적용되도록 명시적 리샤딩을 삽입합니다.

호환되는 샤딩은 기본적으로 작업이 샤딩된 피연산자를 수락하고 샤딩된 결과를 생성할 수 있으며, 샤딩된 통신이 필요하지 않다는 것을 의미합니다 (작업에 여전히 all-reduce 또는 halo-swap과 같은 통신이 필요할 수 있음).

전파 후에도 일부 작업에는 호환되지 않는 샤딩이 있을 수 있습니다.

축 (또는 하위 축)이 여러 텐서에서 상응하지 않는 크기 (예: matmul의 비수축 크기)를 샤딩하는 데 사용되거나 축이 한 텐서의 크기를 샤딩하지만 다른 텐서의 상응하는 크기는 샤딩하지 않는 경우 작업에 샤딩 충돌이 있다고 합니다. 따라서 이 패스 후에는 작업이 충돌이 없게 됩니다.

이 패스는 각 연산에 대해 상응하는 측정기준이 모든 피연산자와 결과에서 동일한 방식으로 샤딩되고 모든 축 (또는 하위 축)이 단일 측정기준 유형을 샤딩하는 데만 사용될 수 있도록 재샤딩 작업을 명시적으로 삽입합니다.

명확한 예시:

입력:

mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"y"},{"x"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
stablehlo.dot %lhs, %rhs {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

출력:

sdy.mesh = <"x"=4, "y"=2>
%lhs : tensor<8x32xf32> {sdy.sharding=<@mesh, \[{"x"}, {"y"}\]>}
%rhs : tensor<32x16xf32> {sdy.sharding=<@mesh, \[{"y"}, {"x"}\]>}
%0 = sdy.reshard %rhs <@mesh, \[{"y"}, {}\]> : tensor<32x16xf32>
stablehlo.dot %lhs, %0 {sdy.sharding_per_value=<[<@mesh, \[{"x"}, {}\]>]>}
  : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>

위의 예에서는 lhs 및 rhs 텐서가 모두 비축소 차원의 'x' 축에서 샤딩되므로 충돌이 발생합니다. 여기서 rhs 텐서는 첫 번째 크기와 'x' 축에서만 샤딩되도록 명시적으로 점 연산 전에 샤딩됩니다. 이렇게 하면 점 연산이 호환됩니다.

-sdy-remove-sharding-groups

전파 후 ShardingGroupOps를 삭제합니다.

-sdy-sharding-constraint-to-reshard

ShardingConstraintOp를 ReshardOp로 변환합니다.

-sdy-sink-data-flow-edges

모든 DataFlowEdgeOp를 입력으로 싱크합니다.

각 DataFlowEdgeOp의 샤딩을 입력 (에지의 루트 타겟)으로 이동하고 연산자를 입력으로 대체합니다.

TODO(tomnatan): 샤딩을 연결할 수 있는 모든 타겟으로 샤딩을 이동하는 것이 좋습니다.

-sdy-update-non-divisible-input-output-shardings

FuncOp 입력/출력을 균등하게 샤딩하여, 나눌 수 없는 샤딩으로 인한 패딩이 필요하지 않도록 합니다.

Shardy 사용자는 텐서 패딩이 필요하지 않도록 함수 입력/출력을 균등하게 나눌 수 있고 샤딩할 수 있기를 기대합니다. 전파로 인해 입력/출력에 나눌 수 없는 샤딩이 생길 수 있으므로 이 패스는 이를 균등하게 샤딩된 원래 샤딩의 가장 큰 측정기준 샤딩 접두사로 업데이트합니다.