'sdy' 방언

Shardy (SDY) 언어

Shardy (SDY) 다이얼렉트는 축 기반 텐서 샤딩 표현식과 텐서에 샤딩을 연결하는 추가 API 구성요소를 정의합니다.

버전 로그: 0.0.1: TensorShardingAttr에 축소되지 않은 축 추가

작업

sdy.all_gather (sdy::AllGatherOp)

축을 따라 all-gather 통신을 실행합니다.

구문:

operation ::= `sdy.all_gather` $gathering_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

gathering_axes에 지정된 축을 따라 텐서의 청크를 수집합니다.

gathering_axes는 축 목록의 목록입니다. 외부 목록이 텐서의 차원을 초과합니다. 각 내부 목록은 해당 차원에서 별도의 수집을 실행해야 하는 축을 지정합니다. 피연산자 (tensor)의 샤딩에 적용되어 결과 (out_sharding)의 샤딩을 얻습니다.

out_sharding은 결과의 샤딩을 결정하는 데 사용되지 않습니다. 대신 결과의 샤딩은 피연산자와 gathering_axes의 샤딩에 따라 결정되며 out_sharding은 이 추론된 샤딩과 일치해야 합니다.

예:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_gather [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\]> : tensor<8x8x8xf32>

제약 조건:

• Sdy_CollectiveOpInterface에 나열된 제약 조건을 충족해야 합니다.
• gathering_axes의 요소는 AxisRefListAttr에 나열된 제약 조건을 충족해야 합니다.
• 피연산자 샤딩에 gathering_axes를 적용하면 out_sharding이 됩니다.

특성: SameOperandsAndResultType

인터페이스: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.all_reduce (sdy::AllReduceOp)

축을 따라 올리듀스 통신 실행

구문:

operation ::= `sdy.all_reduce` ($reduction_op^)? $reduction_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

reduction_axes에 지정된 축을 따라 텐서의 청크를 줄입니다. 결과에서 reduction_axes의 순서는 중요하지 않지만 해당 복제본 그룹의 순서에 영향을 줄 수 있습니다.

제약 조건:

• Sdy_CollectiveOpInterface에 나열된 제약 조건을 충족해야 합니다.
• reduction_axes는 AxisRefListAttr에 나열된 제약 조건을 충족해야 합니다.
• reduction_axes은(는) 메시와 관련하여 정렬되어야 합니다.
• 피연산자 샤딩과 out_sharding의 차원 샤딩이 동일해야 합니다.
• reduction_axes은 피연산자 차원 샤딩 및 복제된 축과 겹쳐서는 안 됩니다 (축소되지 않은 축과는 겹칠 수 있음).
• reduction_axes은 out_sharding의 축소되지 않은 축과 겹쳐서는 안 됩니다. 즉, out_sharding는 reduction_axes를 따라 복제되어야 합니다 (암시적 또는 명시적으로).

특성: SameOperandsAndResultType

인터페이스: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.all_slice (sdy::AllSliceOp)

축을 따라 동적 슬라이스 작업을 실행합니다.

구문:

operation ::= `sdy.all_slice` $slicing_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

slicing_axes에 지정된 축을 따라 텐서의 청크를 슬라이스합니다. sdy.all_slice과 sdy.all_gather 사이에는 대수적 이중성이 있습니다.

slicing_axes는 축 목록의 목록입니다. 외부 목록이 텐서의 차원을 초과합니다. 각 내부 목록은 해당 차원에서 슬라이스를 실행해야 하는 축을 지정합니다. 피연산자 (tensor)의 샤딩에 적용되어 결과(out_sharding)의 샤딩을 얻습니다.

out_sharding은 결과의 샤딩을 결정하는 데 사용되지 않습니다. 대신 결과의 샤딩은 피연산자와 slicing_axes의 샤딩에 따라 결정되며 out_sharding은 이 추론된 샤딩과 일치해야 합니다.

예:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}, {}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_slice [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a", "b", "c"}, {}, {"d"}\]> : tensor<8x8x8xf32>

제약 조건:

• Sdy_CollectiveOpInterface에 나열된 제약 조건을 충족해야 합니다.
• slicing_axes의 요소는 AxisRefListAttr에 나열된 제약 조건을 충족해야 합니다.
• 피연산자 샤딩에 slicing_axes를 적용하면 out_sharding이 됩니다.

특성: SameOperandsAndResultType

인터페이스: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.all_to_all (sdy::AllToAllOp)

축을 따라 all-to-all 통신을 실행합니다.

구문:

operation ::= `sdy.all_to_all` $params $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

매개변수 목록의 각 (axes, src_dim, tgt_dim) 튜플에 대해 이 작업은 tgt_dim 차원과 axes에 지정된 축을 따라 텐서의 청크를 슬라이스하고, 축을 따라 이러한 청크를 분산하고, src_dim 차원을 따라 연결합니다.

이 작업은 기본적으로 src_dim 및 axes를 따라 all-gather를 수행한 후 tgt_dim 및 axes를 따라 all-slice를 수행하는 것입니다. 즉, 입력 텐서의 축 샤딩 차원 src_dim의 접미사가 출력 텐서의 축 샤딩 차원 tgt_dim에 추가됩니다.

all-to-all은 피연산자 (tensor)의 샤딩에 적용되어 결과 (out_sharding)의 샤딩을 얻습니다.

out_sharding은 결과의 샤딩을 결정하는 데 사용되지 않습니다. 대신 결과의 샤딩은 피연산자, src_dim, tgt_dim, axes의 샤딩에 의해 결정되며 out_sharding는 이 추론된 샤딩과 일치해야 합니다.

예:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b"}, {"c"}, {}, {}\]>]>} : tensor<8x8x4x4x32>
%2 = sdy.all_to_all [{"b"}: 0->2, {"c"}: 1->3] %1 out_sharding=<@mesh, [{"a"}, {}, {"b"}, {"c"}\]> : tensor<8x8x4x4x32>

제약 조건:

• Sdy_CollectiveOpInterface에 나열된 제약 조건을 충족해야 합니다.
• 매개변수 목록은 비워 둘 수 없습니다.
• params의 각 매개변수에 대해 다음을 실행합니다.
  • axes의 요소는 AxisRefAttr의 제약 조건을 충족해야 합니다.
  • src_dim 및 tgt_dim은 유효한 차원 (음수가 아니고 텐서의 순위보다 작음)이어야 합니다.
  • src_dim 또는 tgt_dim는 모든 매개변수에서 고유해야 합니다.
  • src_dim은 모든 매개변수에서 오름차순으로 정렬해야 합니다.
• 피연산자 샤딩에서 axes을 src_dim에서 tgt_dim로 이동하면 out_sharding이 발생합니다.

특성: SameOperandsAndResultType

인터페이스: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.collective_permute (sdy::CollectivePermuteOp)

축을 대체하기 위해 집단 순열 통신을 실행합니다.

구문:

operation ::= `sdy.collective_permute` $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

텐서를 샤딩하는 축을 재정렬/바꾸기 위해 각 기기에서 다른 기기로 입력 텐서의 청크를 전송합니다.

집단 순열은 각 차원이 이전과 같이 샤딩되어야 하도록 입력 샤딩을 변환할 수 있습니다. 즉, 크기의 곱이 이전에 텐서를 샤딩한 축의 곱과 일치하는 축을 따라 샤딩되어야 합니다.

이는 단일 차원 또는 여러 차원에서 축의 순서를 바꾸고 샤딩된 축을 복제된 축으로 바꾸는 데 유용합니다.

아래 예시에서 샤딩된 텐서 크기는 tensor<1x4x2xf32>이며, 이는 집단 순열에 의해 유지됩니다.

예:

sdy.mesh @mesh = <["a"=2, "b"=2, "c"=4, "d"=2, "e"=2, "f"=2]>
%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "c"}, {"f"}, {"d", "e"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.collective_permute %1 out_sharding=<@mesh, [{"c":(1)2, "b", "f"}, {"a"}, {"e", "d"}\]> : tensor<8x8x8xf32>

제약 조건:

• Sdy_CollectiveOpInterface에 나열된 제약 조건을 충족해야 합니다.
• 입력 및 출력 샤딩의 메시가 서로 다른 경우 이러한 메시에는 정확히 동일한 축과 서로 다른 기기 ID 순서가 있어야 합니다.
• 각 차원에서 out_sharding의 샤딩 축 크기의 곱은 해당 피연산자 차원 샤딩의 곱과 일치해야 합니다.

특성: SameOperandsAndResultType

인터페이스: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.constant (sdy::ConstantOp)

상수 작업

상수 value에서 output 텐서를 생성합니다.

참고: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

예:

%output = sdy.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>

특성: AlwaysSpeculatableImplTrait

인터페이스: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

효과: MemoryEffects::Effect{}

속성:

결과:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

데이터 흐름 에지 작업

구문:

operation ::= `sdy.data_flow_edge` $input (`sharding````=``` $sharding^)? attr-dict `:` type($result)

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

작업에는 서로 직교하는 여러 데이터 흐름 가장자리가 있을 수 있습니다.

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

  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 사이에 있습니다.

sdy.data_flow_edge는 에지의 소유자 (타겟일 수 있지만 블록 인수보다는 작업 결과가 바람직함)를 입력으로 사용하며 다른 용도로 사용해서는 안 됩니다. 이 작업은 원래 용도가 없던 입력을 사용할 수 있으므로 순수하지 않습니다.

sdy.data_flow_edge는 에지의 모든 타겟에 대한 선택적 샤딩도 보유하며, 전파 중에 타겟의 샤딩 대신 이 샤딩을 업데이트해야 합니다 (연결할 수 있는 경우). 이는 작업에 여러 개의 가장자리가 있는 경우에 유용합니다.

• 각 에지를 통해 별도로 전파됩니다.
• 모든 타겟을 한 번에 업데이트하는 대신 각 에지의 샤딩을 별도로 업데이트합니다(예: 작업에 결과 샤딩을 위한 단일 불변 TensorShardingPerValueAttr가 있음).
• 소스의 샤딩이 변경된 경우 각 에지를 작업 목록에 별도로 추가합니다.

전파는 소스를 피연산자로, 타겟을 결과로, 항등 sdy.op_sharding_rule을 사용하여 일반 작업인 것처럼 sdy.data_flow_edge의 모든 소스와 타겟 간에 샤딩을 전파합니다. 즉, 순방향 전파는 소스에서 타겟으로, 역방향 전파는 타겟에서 소스로 이루어집니다.

SdyDialect 작업으로 정의된 sdy.data_flow_edge의 입력은 허용되지 않으므로 등록되지 않은 sdy.sharding 속성이 있는 작업으로 정의된다고 가정할 수 있습니다.

특성: SameOperandsAndResultType

인터페이스: InferTypeOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.func_data_flow_edge (sdy::FuncDataFlowEdgeOp)

함수 입력/출력 데이터 흐름 가장자리 작업

구문:

operation ::= `sdy.func_data_flow_edge` $operand attr-dict `:` type($result)

함수 인수 또는 호출 결과를 위한 데이터 흐름 가장자리 작업입니다. 피연산자가 BlockArgument인 경우 호출자 callOp의 인수에서 func 인수의 사용자로 연결하는 브리지입니다. 각 함수 인수에는 하나의 함수 데이터 흐름 에지가 있습니다. 피연산자가 OpResult인 경우 호출된 funcOp의 반환 값에서 호출 결과 사용자로 연결되는 브리지입니다. 각 호출 결과에 하나의 func 데이터 흐름 에지가 있습니다.

특성: SameOperandsAndResultType

인터페이스: InferTypeOpInterface, SymbolUserOpInterface

피연산자:

결과:

sdy.manual_computation (sdy::ManualComputationOp)

수동 집합을 사용한 멀티 디바이스 병렬 처리 작업

구문:

operation ::= `sdy.manual_computation` `(`operands`)`
              `in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)
              `out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)
              `manual_axes````=```$manual_axes
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:`
              functional-type(operands, results)

논리적 모양이 기기별 로컬 실제 버퍼 모양과 일치하고 집합이 실제 교차 기기 통신과 정확히 일치하는 명시적 집합이 있는 기기별 로컬 코드로 작성된 영역으로 이동합니다.

본체는 manual_axes와 관련하여 로컬입니다. 전파는 manual_axes 목록에 없는 자유 축의 본체를 통해 발생합니다.

순위가 지정되지 않은 텐서는 순위가 0인 샤딩, 즉 완전히 복제된 샤딩을 갖는 것으로 예상됩니다.

제약 조건:

• in_shardings 및 out_shardings의 요소는 TensorShardingAttr에 나열된 제약 조건을 충족해야 합니다.
• 작업 영역의 전역 및 로컬 텐서 입력/출력 수가 일치해야 합니다.
• 수동 축은 각 dim 샤딩에서 자유 축 앞에 와야 합니다.
• 수동 축은 패딩을 도입할 수 없습니다. 즉, 측정기준 크기는 해당 수동 축 크기로 나눌 수 있어야 합니다.
• 작업 영역 인수/결과의 전역 및 로컬 모양이 일치해야 합니다.

특성: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

인터페이스: ShardableDataFlowOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.mesh (sdy::MeshOp)

이름이 지정된 메시

구문:

operation ::= `sdy.mesh` $sym_name `=` $mesh attr-dict

새 명명된 메시를 정의합니다. 모듈의 모든 메시에는 동일한 수의 기기가 있어야 합니다 (single device_id가 있는 메시 제외). 메시는 모듈의 SymbolTable에 표시되는 Symbol 작업이며 name에서 참조할 수 있습니다.

특성: HasParent<ModuleOp>, SymbolName

인터페이스: Symbol

속성:

sdy.named_computation (sdy::NamedComputationOp)

이름이 지정된 계산 작업

구문:

operation ::= `sdy.named_computation` `<`$name`>` `` `(` $operands `)`
              (`in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)^)?
              (`out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)^)?
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:` functional-type($operands, results)

계산(작업 블록)을 그룹화하고 이름을 지정합니다. 전파는 모든 항목이 인라인된 것처럼 리전 안팎으로 흐릅니다.

이는 호출 명령어를 통해 다른 함수로 전파하는 데 사용할 수 있습니다. Shardy 사용자는 호출 작업을 sdy.named_computation 작업으로 변환하여 호출된 함수의 본문을 named_computation의 본문에 복제/복사하는 가져오기/내보내기 패스를 작성해야 합니다.

영역에 있는 각 블록 인수와 반환된 값의 유형은 op의 피연산자 유형 및 결과 유형과 동일해야 합니다.

예:

%1 = sdy.named_computation<"foo">(%0) (%arg1: tensor<16x32xf32>) {
  sdy.return %arg1 : tensor<16x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

특성: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

인터페이스: ConditionallySpeculatable, InferTypeOpInterface, ShardableDataFlowOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

전파 장벽 작업

구문:

operation ::= `sdy.propagation_barrier` $input `allowed_direction````=```$allowed_direction attr-dict `:` type($input)

이 작업은 입력으로 가져온 것과 동일한 값을 출력하는 ID 작업과 유사하게 작동합니다. 하지만 전파 측면에서는 특정 방향으로만 전파가 흐르도록 허용합니다.

이렇게 하면 장벽 작업의 결과와 피연산자 간에 샤딩이 전파되지 않습니다.

• FORWARD는 샤딩이 피연산자에서 결과로만 흐를 수 있음을 의미합니다.
• BACKWARD는 샤딩이 결과에서 피연산자로만 흐를 수 있음을 의미합니다.
• NONE은 이 작업을 통해 샤딩이 전파될 수 없음을 의미합니다.
• 이 작업은 중복되므로 BOTH를 지정할 수 없습니다.

특성: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

인터페이스: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

효과: MemoryEffects::Effect{}

속성:

피연산자:

결과:

sdy.reduce_scatter (sdy::ReduceScatterOp)

축을 따라 reduce-scatter 통신을 실행합니다.

구문:

operation ::= `sdy.reduce_scatter` ($reduction_op^)? $reduce_scatter_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

reduce_scatter_axes에 지정된 축을 따라 텐서의 청크를 축소한 다음 동일한 축을 따라 결과를 분산합니다. 이 작업은 기본적으로 동일한 reduce_scatter_axes를 따라 sdy.all_reduce 다음에 sdy.all_slice가 오는 조합입니다.

제약 조건:

• Sdy_CollectiveOpInterface에 나열된 제약 조건을 충족해야 합니다.
• reduce_scatter_axes의 요소는 AxisRefListAttr에 나열된 제약 조건을 충족해야 합니다.
• 피연산자 샤딩에 reduce_scatter_axes를 적용하면 out_sharding이 됩니다.

특성: SameOperandsAndResultType

인터페이스: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

암시적으로 또는 명시적으로 복제된 축을 축소되지 않은 축으로 이동합니다.

구문:

operation ::= `sdy.replicated_to_unreduced` $axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

axes는 피연산자에서 암시적 또는 명시적으로 복제되어야 합니다. 이 작업으로 인해 결과에서 축소되지 않습니다. 다음과 같은 관계가 있습니다.

all-reduce(replicated-to-unreduced(x, axes), axes) = x

예:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}, {}\], replicated={"c", "d"}, unreduced={"e"}>]>} : tensor<8x8x8xf32>
%2 = sdy.replicated_to_unreduced {"a", "c", "f"} %1 out_sharding=<@mesh, [{"b"}, {}, {}\], replicated={"d"}, unreduced={"a", "c", "e", "f"}> : tensor<8x8x8xf32>

제약 조건:

• Sdy_CollectiveOpInterface에 나열된 제약 조건을 충족해야 합니다.
• axes는 AxisRefListAttr에 나열된 제약 조건을 충족해야 합니다.
• axes은(는) 메시와 관련하여 정렬되어야 합니다.
• axes이 비어 있지 않습니다.
• 입력과 출력 샤딩의 차원 샤딩이 동일해야 합니다.
• axes는 피연산자 샤딩에서 암시적 또는 명시적으로 복제되어야 합니다.
• inUnreducedAxes + axes = outUnreducedAxes

특성: SameOperandsAndResultType

인터페이스: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.reshard (sdy::ReshardOp)

텐서를 다른 샤딩으로 리샤드합니다.

구문:

operation ::= `sdy.reshard` $input $sharding attr-dict `:` type($result)

입력 텐서의 기존 샤딩과 다른 지정된 샤딩으로 입력 텐서를 리샤딩합니다.

ShardingConstraintOp와 ReshardOp는 모두 텐서에 샤딩을 연결합니다. 수명은 다음과 같습니다.

1. 샤딩 전파 전에 사용자가 ShardingConstraintOp를 추가합니다.
2. 샤딩 전파는 ShardingConstraintOp를 사용합니다. 샤딩 전파 결과에 ShardingConstraintOp이 없습니다. 대신 필요한 경우 ReshardOp이 추가될 수 있습니다.
3. 파티셔너는 ReshardOp을 집단 연산 (또는 ID 연산)으로 변환합니다. 파티셔너 결과에 ReshardOp이 없어야 합니다.

특성: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

인터페이스: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface), SymbolUserOpInterface

효과: MemoryEffects::Effect{}

속성:

피연산자:

결과:

sdy.return (sdy::ReturnOp)

sdy.return 작업은 sdy 리전 기반 작업 및 기타 모든 Shardy 리전 기반 작업에 연결된 리전을 종료합니다. 가변적입니다. 유형이 모든 유형 (하지만 동일한 종류, 예를 들어 AnyTensor)일 수 있는 값 목록을 인수로 취하므로 Shardy IR 스택의 다양한 수준에서 재사용할 수 있습니다.

구문:

operation ::= `sdy.return` attr-dict ($results^ `:` type($results))?

특성: AlwaysSpeculatableImplTrait, ReturnLike, Terminator

인터페이스: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

효과: MemoryEffects::Effect{}

피연산자:

sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

피연산자의 일부 샤딩된 축을 결과의 축소되지 않은 축으로 이동합니다.

구문:

operation ::= `sdy.sharded_to_unreduced` $axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

axes는 피연산자를 샤딩하는 데 사용해야 합니다. 이 작업은 결과에서 축소되지 않도록 합니다. 다음과 같은 관계가 있습니다.

all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes)입니다. 여기서 all-gather, sharded-to-unreduced, all-reduce는 동일한 축에 적용됩니다.

예:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\], unreduced={"e"}>]>} : tensor<8x8x8xf32>
%2 = sdy.sharded_to_unreduced [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\], unreduced={"b", "c", "d", "e"}> : tensor<8x8x8xf32>

제약 조건:

• Sdy_CollectiveOpInterface에 나열된 제약 조건을 충족해야 합니다.
• axes의 요소는 AxisRefListAttr에 나열된 제약 조건을 충족해야 합니다.
• 피연산자 샤딩에 axes를 적용하면 out_sharding이 됩니다.

특성: SameOperandsAndResultType

인터페이스: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

텐서를 지정된 샤딩으로 제한합니다.

구문:

operation ::= `sdy.sharding_constraint` $input $sharding attr-dict `:` type($result)

텐서 또는 텐서 사용의 하위 집합이 샤딩되는 방식을 나타내기 위해 샤딩을 중간 텐서 (예: matmul의 결과)에 연결합니다.

샤딩에 열린 차원과 제한되지 않은 축이 있으면 텐서를 열린 차원을 따라 추가로 샤딩할 수 있습니다.

이 작업은 다음 중 하나일 수 있습니다.

• 사용되지 않음 (매달림) - 즉, 연결된 샤딩이 입력 텐서 자체가 샤딩되는 방식입니다.
• 사용이 있습니다. 즉, 연결된 샤딩은 샤딩 제약 조건 작업의 사용이 샤딩되는 방식인 반면 입력 텐서의 다른 사용에는 다른 샤딩이 있을 수 있습니다 (입력 텐서에 다른 사용이 없는 경우 동작은 사용이 없는 경우와 동일함).

특성: SameOperandsAndResultType

인터페이스: InferTypeOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

sdy.sharding_group (sdy::ShardingGroupOp)

그룹의 텐서가 동일한 샤딩을 갖도록 제한합니다.

구문:

operation ::= `sdy.sharding_group` $input `group_id````=```$group_id attr-dict `:` type($input)

이 작업은 텐서를 샤딩 그룹(동일한 샤딩이 적용되는 텐서 그룹)에 할당하는 인터페이스를 제공합니다. 전파 중에 한 그룹 요소가 샤딩되면 다른 모든 구성원이 정확히 동일한 방식으로 샤딩됩니다. 이 작업은 그룹 ID 인수를 가져오고 결과를 반환하지 않지만 대신 내부 샤딩 그룹 표현식을 수정하여 지정된 ID가 있는 그룹에 입력 텐서를 추가합니다.

인터페이스: InferTypeOpInterface

속성:

피연산자:

속성

AllToAllParamAttr

All-to-all 매개변수

구문:

#sdy.all_to_all_param<
  ::llvm::ArrayRef<AxisRefAttr>,   # axes
  int64_t,   # src_dim
  int64_t   # tgt_dim
>

모든-모두를 실행할 축과 소스/타겟 측정기준을 포함하는 튜플입니다.

매개변수:

AllToAllParamListAttr

all-to-all 매개변수 목록

구문:

#sdy.all_to_all_param_list<
  ::llvm::ArrayRef<AllToAllParamAttr>   # value
>

매개변수:

AxisRefAttr

전체 축 또는 분할된 하위 축에 대한 참조

구문:

#sdy.axis_ref<
  ::llvm::StringRef,   # name
  SubAxisInfoAttr   # sub_axis_info
>

제약 조건:

• name는 바인드된 MeshAttr에 있어야 합니다.
• sub_axis_info가 있는 경우 SubAxisInfoAttr의 제약 조건을 충족해야 합니다.

매개변수:

AxisRefListAttr

축 참조 목록

구문:

#sdy.axis_ref_list<
  ::llvm::ArrayRef<AxisRefAttr>   # value
>

제약 조건:

• value의 요소는 AxisRefAttr의 제약 조건을 충족해야 합니다.
• 중복된 축 참조 또는 서로 겹치는 하위 축이 없습니다.
• 인접한 두 축 참조가 동일한 전체 축의 연속된 하위 축이 아닙니다. 즉, 하나의 하위 축 또는 전체 축으로 병합할 수 있습니다.

매개변수:

AxisToPropagationDetailsAttr

특정 축과 소스의 전파 가장자리 흐름 세부정보입니다.

구문:

#sdy.axis_to_propagation_details<
  ::mlir::sdy::AxisRefAttr,   # axis_name
  ::mlir::sdy::EdgeValueRefAttr,   # source
  ::llvm::ArrayRef<EdgeValueRefAttr>   # targets
>

특정 축을 따라 소스 값 참조를 타겟 값 참조 목록에 매핑합니다.

매개변수:

DimMappingAttr

측정기준의 요소 색인 목록

빈 목록은 null 매핑 (*로 파싱/인쇄됨)을 나타냅니다. 즉, 측정기준이 어떤 요소에도 매핑되지 않습니다.

제약 조건:

• 요인 지수가 하나 이상 있습니다.
• 요인 색인은 [0, $factor_sizes) 범위에 있어야 합니다.
• 요인이 여러 개인 경우 크기가 1인 요인은 있을 수 없습니다.
• 중복된 요소 색인이 없습니다.

매개변수:

DimensionShardingAttr

측정기준 샤딩

텐서 차원을 샤딩할 축 이름 목록(주에서 부로), 차원을 추가로 샤딩할 수 있는지를 나타내는 불리언, 이 차원 샤딩의 우선순위를 나타내는 선택적 정수(샤딩 전파 중에 적용됨)입니다. 우선순위는 사용자 샤딩 주석에서 비롯되며 값이 낮을수록 우선순위가 높습니다. 주석에 우선순위가 누락된 경우 가장 높은 우선순위가 가정됩니다.

제약 조건:

• axes의 요소는 AxisRefListAttr에 나열된 제약 조건을 충족해야 합니다.
• 측정기준 샤딩에 우선순위가 있는 경우:
  • 우선순위가 0 이상입니다.
  • 닫힌 경우 측정기준에 축이 하나 이상 있습니다.

매개변수:

EdgeValueRefAttr

type 유형의 값 가장자리의 특정 색인에 대한 참조

구문:

#sdy.edge_value_ref<
  `operand` | `result`,   # type
  int64_t   # index
>

매개변수:

ListOfAxisRefListsAttr

축 참조 목록

구문:

#sdy.list_of_axis_ref_lists<
  ::llvm::ArrayRef<AxisRefListAttr>   # value
>

매개변수:

ManualAxesAttr

ManualComputationOp가 수동인 축 목록

구문:

#sdy.manual_axes<
  ::llvm::ArrayRef<StringAttr>   # value
>

매개변수:

MeshAttr

축 메시 및 기기 목록

구문:

#sdy.mesh<
  ::llvm::ArrayRef<MeshAxisAttr>,   # axes
  ::llvm::ArrayRef<int64_t>   # device_ids
>

메시는 축 목록과 기기 순서를 지정하는 선택적 기기 ID 목록입니다.

축 목록이 비어 있는 경우

• device_ids가 제공되지 않으면 빈 메시입니다.
• device_ids가 제공되는 경우 단일 음이 아닌 정수여야 하며 이를 최대 샤딩 메시라고 합니다.

축 목록이 제공된 경우

• 기기 ID 목록이 지정된 경우 축 크기의 곱이 기기 수와 일치해야 합니다.
• 기기 ID 목록이 지정되지 않은 경우 암시적 기기 ID 목록은 iota(product(axes))입니다. 간단하게 하기 위해 iota(product(axes))와 동일한 기기 ID 목록을 지정하는 것도 허용되지 않습니다. 이 경우 기기 ID 목록을 지정하면 안 됩니다.
• 축의 총 크기가 1이더라도 최대 샤딩 메시가 아닙니다.

다음은 메시의 몇 가지 예입니다.

• 빈 메시는 전파 중에 대체할 수 있는 자리표시자 메시를 나타냅니다. <[]>
• 축 목록이 없고 단일 음수가 아닌 기기 ID가 있는 메시(최대 샤딩 메시): <[], device_ids=[3]>
• 두 축과 암시적 기기 ID iota(6)가 있는 메시: <["a"=2, "b"=3]>
• 두 축과 기기 순서를 지정하는 명시적 기기 ID가 있는 메시: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

제약 조건:

• device_ids의 요소는 음수가 아니어야 합니다.
• axes가 비어 있으면 device_ids의 크기는 0 (빈 메시) 또는 1(최대 샤딩 메시)일 수 있습니다.
• axes가 비어 있지 않으면
  • axes의 요소에는 중복 이름이 없어야 합니다.
  • device_ids이 지정되면 원래 device_ids은 iota(product(axis_sizes))이 아니고 정렬된 device_ids은 iota(product(axis_sizes))입니다.

매개변수:

MeshAxisAttr

메시의 이름이 지정된 축

구문:

#sdy.mesh_axis<
  ::llvm::StringRef,   # name
  int64_t   # size
>

매개변수:

OpShardingRuleAttr

작업을 파티셔닝할 수 있는 방법을 지정합니다.

구문:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  ::llvm::ArrayRef<int64_t>,   # reduction_factors
  ::llvm::ArrayRef<int64_t>,   # need_replication_factors
  ::llvm::ArrayRef<int64_t>,   # permutation_factors
  ::llvm::ArrayRef<int64_t>,   # blocked_propagation_factors
  bool   # is_custom_rule
>

샤딩 규칙은 작업의 다양한 속성(속성, 피연산자의 모양, 결과의 모양 등)에 따라 작업을 파티셔닝하는 방법을 지정합니다. 예를 들면 다음과 같습니다.

%0 = stablehlo.add %arg0, %arg1 {
    sdy.sharding_rule = #sdy.op_sharding_rule<
        ([i, j],[i, j])->([i, j])
        {i=8, j=8}>
} : tensor<8x8xf32>

%1 = stablehlo.dot_general %arg2, %arg3, contracting_dims = [1] x [0] {
  sdy.sharding_rule = #sdy.op_sharding_rule<
      ([i, k],[k, j])->([i, j])
      {i=8, j=16, k=8}>
}: (tensor<8x8xf32>, tensor<8x16xf32>) -> tensor<8x16xf32>

샤딩할 수 없는 크기가 1인 요소는 허용됩니다. 이는 주로 포인트별 작업과 같은 많은 작업에 피연산자와 결과에 걸쳐 대응하는 크기가 1인 차원이 있기 때문에 완전성을 위한 것입니다.

요인 유형:

• reduction_factors에는 점 연산에서 축소되는 차원과 같이 축소가 필요한 요소의 색인이 포함됩니다. 이러한 요소는 피연산자에는 있을 수 있지만 결과에는 있을 수 없습니다.
• need_replication_factors에는 정렬 작업에서 정렬된 측정기준과 같이 전체 복제가 필요한 요소의 색인이 포함됩니다.
• permutation_factors에는 패드 작업의 패딩 차원과 같이 샤딩된 경우 collective-permute가 필요한 요소의 색인이 포함됩니다.
• 다른 모든 요인은 패스스루 요인, 즉 매핑된 모든 텐서에서 동일한 방식으로 샤딩되는 경우 통신이 필요하지 않은 요인으로 간주됩니다.

blocked_propagation_factors에는 샤딩이 전파될 수 없는 요소가 포함되어 있습니다. 요소 유형과 직교합니다. 즉, 차단된 전파 요소는 요소 유형 중 하나일 수 있습니다.

is_custom_rule는 사용자가 정의한 규칙인지 여부를 설명합니다. 사용자는 맞춤 호출의 샤딩 규칙을 정의하거나 표준 작업의 사전 정의된 샤딩 규칙을 덮어쓸 수 있습니다. 맞춤 규칙은 항상 보존되며 삭제되지 않습니다.

제약 조건:

• 피연산자/결과 매핑 수는 작업의 피연산자/결과 수와 일치해야 합니다.
• 매핑이 하나 이상 있습니다 (피연산자/결과가 없는 작업에는 규칙이 있을 수 없음).
• 각 TensorMappingAttr의 순위는 해당 텐서 유형의 순위와 일치합니다.
• 각 요인 그룹 (reduction_factors, need_replication_factors, permutation_factors)의 경우:
  • 요소는 [0, $factor_sizes] 범위에 있어야 합니다.
  • 각 그룹 내 및 그룹 간에 중복된 요소 색인이 없습니다.

매개변수:

PropagationEdgesAttr

모든 전파 단계의 전파 에지 메타데이터입니다.

구문:

#sdy.propagation_edges<
  ::llvm::ArrayRef<PropagationOneStepAttr>   # value
>

값의 축별 전파 세부정보 목록으로, 단계 색인별로 그룹화됩니다.

매개변수:

PropagationOneStepAttr

단계별 전파 메타데이터입니다.

구문:

#sdy.propagation_one_step<
  int64_t,   # step_index
  ::llvm::ArrayRef<AxisToPropagationDetailsAttr>   # axis_entries
>

단일 전파 단계의 모든 축에 대한 전파 세부정보입니다.

매개변수:

SubAxisInfoAttr

이 하위 축이 전체 축에서 파생되는 방식에 관한 정보

구문:

#sdy.sub_axis_info<
  int64_t,   # pre_size
  int64_t   # size
>

전체 축을 n개의 하위 축으로 분할하면 축이 [k_1,…,k_n]으로 재구성되고 i번째 하위 축은 왼쪽의 모든 축 크기 m=prod(k_1,...,k_(i-1)) (사전 크기)와 크기 k_i의 곱으로 표현할 수 있습니다. 따라서 sub-axis-info 속성은 이 두 숫자를 보유하며 사전 크기 m 및 크기 k의 경우 (m)k로 표시됩니다.

제약 조건:

• pre-size이 1 이상입니다.
• size이 1보다 큽니다.
• pre-size는 전체 축의 크기를 나눠야 합니다. 즉, pre-size와 size 모두 전체 축의 크기를 나누고 하위 축은 전체 축을 벗어나지 않습니다.
• 하위 축의 크기가 해당 전체 축의 크기와 같지 않은 경우 전체 축을 대신 사용해야 합니다.

매개변수:

TensorMappingAttr

텐서의 각 차원에 대한 팩터 매핑입니다.

구문:

#sdy.tensor_mapping<
  ::llvm::ArrayRef<DimMappingAttr>   # dim_mappings
>

제약 조건:

• dim_mappings의 요소는 DimMappingAttr의 제약 조건을 충족해야 합니다.
• 차원 간에 중복된 요인 색인이 없습니다.

매개변수:

TensorShardingAttr

텐서 샤딩

구문:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>,   # replicated_axes
  ::llvm::ArrayRef<AxisRefAttr>,   # unreduced_axes
  `sum` | `max` | `min`   # reduction_op
>

텐서 샤딩은 특정 메시와 바인드되며 해당 메시의 축 이름만 참조할 수 있습니다. 차원 샤딩은 텐서의 각 차원에 대해 주요 축에서 부축으로 샤딩되는 축 (또는 하위 축)을 알려줍니다. 측정기준을 샤딩하지 않는 다른 모든 축은 암시적으로 또는 명시적으로 (복제된 축 목록에 표시되는 경우) 복제됩니다.

텐서에 샤딩 속성이 없으면 완전히 열린 텐서 샤딩과 같습니다.

이 샤딩이 바인드된 메시지는 해당 MeshOp 기호를 참조하는 기호 이름이나 인라인 MeshAttr로 지정할 수 있습니다.

샤딩에는 축소되지 않은 축 (unreduced_axes로 지정)이 있을 수 있습니다. 즉, 텐서가 이러한 축을 따라 축소되지 않습니다. 예를 들어 matmul의 축소되는 차원이 lhs와 rhs 모두에서 x 축을 따라 샤딩되면 결과는 x를 따라 축소되지 않습니다. 축소되지 않은 축을 따라 텐서에 all-reduce를 적용하면 해당 축을 따라 텐서가 복제됩니다. 하지만 축이 축소되지 않은 텐서가 즉시 완전히 축소될 필요는 없습니다. stablehlo.add와 같은 선형 작업에 전달될 때 축소되지 않은 상태로 유지될 수 있으며 (lhs와 rhs가 모두 축소되지 않은 경우) 이후에 완전히 축소될 수 있습니다. 감소 유형은 합계라고 가정합니다. 향후 다른 감소가 지원될 수 있습니다.

제약 조건:

• dim_shardings의 요소는 DimensionShardingAttr에 나열된 제약 조건을 충족해야 합니다.
• replicated_axes의 요소는 AxisRefListAttr에 나열된 제약 조건을 충족해야 합니다.
• unreduced_axes의 요소는 AxisRefListAttr에 나열된 제약 조건을 충족해야 합니다.
• 해당 텐서 유형이 ShapedType이 아닌 경우 샤딩의 순위는 0이어야 하며 복제된 축이 없어야 합니다.
• ShapedType인 경우 다음을 수행합니다.
  • 텐서에 순위가 있어야 합니다.
  • 차원 샤딩 수는 텐서의 순위와 같습니다.
  • 크기가 0인 차원은 샤딩되지 않습니다.
• dim_shardings, replicated_axes, unreduced_axes에 중복된 축 참조나 서로 겹치는 하위 축이 없습니다.
• replicated_axes 및 unreduced_axes의 항목은 mesh_or_ref와 관련하여 정렬됩니다 (AxisRefAttr::getMeshComparator 참고).

매개변수:

TensorShardingPerValueAttr

작업의 피연산자/결과별 텐서 샤딩

구문:

#sdy.sharding_per_value<
  ::llvm::ArrayRef<TensorShardingAttr>   # shardings
>

작업의 각 피연산자/결과에 해당하는 TensorShardingAttr 목록입니다.

제약 조건:

• shardings의 요소는 TensorShardingAttr의 제약 조건을 충족해야 합니다.

매개변수:

열거형

EdgeNodeType

에지 노드 유형 enum

케이스:

PropagationDirection

전파 방향 열거형

케이스:

ReductionOp

감소 작업 enum

케이스: