'sdy' 방언

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

운영

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

연산에는 n개의 데이터 흐름 에지가 있지만, i번째 데이터 흐름 에지는 소스 x_i, return_value_i와 타겟 y_i, pred_arg_i, body_arg_i 사이에 있습니다.

sdy.data_flow_edge는 다른 용도가 없어야 하는 에지의 루트 타겟 (어떤 타겟이든 될 수 있지만 블록 인수가 아닌 op 결과가 바람직함)을 입력으로 사용합니다. 이 작업은 원래 아무런 용도가 없었던 입력을 취할 수 있으므로 순수하지 않습니다

sdy.data_flow_edge는 또한 에지의 모든 타겟에 대한 선택적 샤딩을 보유하며, 전파 중에 타겟의 샤딩 대신 이 샤딩을 업데이트해야 합니다 (연결할 수 있는 경우). 이는 다음과 같은 작업을 훨씬 더 효율적으로 수행할 수 있으므로 연산에 많은 가장자리가 있는 경우에 유용합니다.

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

전파는 소스가 피연산자이고 대상이 결과인 일반 연산자와 ID sdy.op_sharding_rule인 것처럼 sdy.data_flow_edge의 모든 소스와 대상 간에 샤딩을 전파합니다. 즉, 정방향 전파는 소스에서 대상으로 이루어지고 역전파는 대상에서 소스로 진행됩니다.

sdy.data_flow_edge의 입력이 SdyDialect 연산자에 의해 정의되는 것은 허용되지 않으므로 등록되지 않은 sdy.sharding 속성이 있는 연산자에 의해 정의되었다고 가정할 수 있습니다.

트레잇: SameOperandsAndResultType

인터페이스: InferTypeOpInterface

속성:

피연산자:

결과:

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를 기준으로 로컬입니다. 전파는 수동_axes 목록에 없는 자유 축의 본문을 통해 발생합니다.

트레잇: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

속성:

피연산자:

결과:

sdy.mesh (sdy::MeshOp)

이름이 지정된 메시

구문:

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

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

트레잇: HasParent<ModuleOp>

인터페이스: 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의 본문으로 복제/복사하는 가져오기/내보내기 패스를 작성해야 합니다.

리전의 각 블록 인수 및 반환 값의 유형은 피연산자의 유형 및 연산자의 결과 유형과 동일해야 합니다.

예:

%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, ShardableDataFlowOpInterface

속성:

피연산자:

결과:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

전파 장벽 작업

구문:

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

이 연산자는 ID 연산자처럼 작동하여 입력으로 사용한 것과 동일한 값을 출력합니다. 하지만 전파 측면에서는 전파가 특정 방향으로만 흐르도록 허용합니다.

이렇게 하면 배리어 연산의 결과와 피연산자의 사용 간에 샤딩이 전파되는 것을 방지할 수 있습니다.

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

트레잇: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

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

효과: MemoryEffects::Effect{}

속성:

피연산자:

결과:

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가 없어야 합니다.

// TODO(b/331680067). 표준화 패턴을 추가하여 중복 // 재샤딩 작업을 삭제하세요.

트레잇: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

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

효과: MemoryEffects::Effect{}

속성:

피연산자:

결과:

sdy.return (sdy::ReturnOp)

sdy.return 작업은 sdy 리전 기반 작업 및 기타 Shardy 리전 기반 작업에 연결된 리전을 종료합니다. 이는 가변 인수입니다. 유형이 무엇이든 될 수 있지만 동일한 종류 (예: AnyTensor)여야 하는 값 목록을 인수로 사용하며, 따라서 Shardy IR 스택의 다양한 수준에서 재사용할 수 있습니다.

구문:

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

트레잇: AlwaysSpeculatableImplTrait, Terminator

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

효과: MemoryEffects::Effect{}

피연산자:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

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

구문:

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

중간 텐서 (예: matmul의 결과)에 샤딩을 연결하여 이 텐서 또는 사용의 하위 집합을 샤딩해야 함을 나타냅니다.

샤딩에 개방형 측정기준과 제약 조건이 없는 축이 있는 경우 텐서를 개방형 측정기준을 따라 추가로 샤딩할 수 있음을 의미합니다.

이 작업은 다음 중 하나를 수행할 수 있습니다.

• 무용지물 (댕글링)이 없습니다. 즉, 연결된 샤딩이 입력 텐서 자체를 샤딩해야 합니다.
• 사용이 있음 - 즉, 연결된 샤딩은 샤딩 제약 조건 연산자의 사용이 샤딩되는 방식이지만 입력 텐서의 다른 사용에는 다른 샤딩이 있을 수 있습니다. 입력 텐서에 다른 사용이 없는 경우 동작은 사용이 없는 케이스와 동일합니다.

트레잇: Elementwise, SameOperandsAndResultType

인터페이스: InferTypeOpInterface

속성:

피연산자:

결과:

sdy.sharding_group (sdy::ShardingGroupOp)

샤딩 그룹 작업

구문:

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

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

속성:

피연산자:

속성

AxisRefAttr

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

구문:

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

매개변수:

DimMappingAttr

측정기준의 요인 색인 목록

모든 계수 색인은 [0, num_factors) 범위 내에 있어야 하며 빈 목록은 null 매핑임을 나타냅니다 (*로 파싱/인쇄됨). 즉, 측정기준이 어떤 요소에도 매핑되지 않습니다.

매개변수:

DimensionShardingAttr

측정기준 샤딩

텐서 측정기준을 주요 측정기준에서 하위 측정기준으로 샤딩할 축 이름 목록, 측정기준을 더 샤딩할 수 있는지 나타내는 불리언, 샤딩 전파 중에 준수되는 이 측정기준 샤딩의 우선순위를 나타내는 선택적 정수입니다. 우선순위는 사용자 샤딩 주석에서 비롯되며 값이 낮을수록 우선순위가 높습니다. 주석에 우선순위가 없으면 가장 높은 우선순위가 가정됩니다.

매개변수:

ManualAxesAttr

구문:

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

매개변수:

MeshAttr

축 메시 및 기기 목록

구문:

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

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

축 목록이 비어 있으면 메시에 크기가 1인 이름이 지정되지 않은 암시적 축이 있습니다. 이 경우 기기 ID 목록이 제공되지 않으면 암시적 기기 ID 목록은 [0]입니다. 기기 ID 목록이 제공된 경우 목록에는 0이 아닌 값의 단일 정수가 포함되어야 합니다. 이를 최대 샤딩 사례라고 합니다.

최대 샤딩이 아닌 모든 케이스의 경우 기기 ID 목록이 지정된 경우 축 크기의 곱은 기기 수와 일치해야 합니다. 기기 ID 목록이 지정되지 않으면 암시적 기기 ID 목록은 iota(product(axes))입니다. 편의상 iota(product(axes)와 동일한 기기 ID 목록 지정은 허용되지 않음) 이 경우 기기 ID 목록을 지정해서는 안 됩니다.

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

• 빈 메시는 전파 중에 대체될 수 있는 자리표시자 메시를 나타냅니다. <[]>
• 이름이 지정되지 않은 축과 명시적 기기 ID가 있는 메시로, 일반적으로 최대 샤딩을 나타내는 데 사용됩니다. <[], device_ids=[3]>
• 축이 2개이고 기기 ID가 암시된 메시 iota(6): <["a"=2, "b"=3]>
• 두 축과 기기 순서를 지정하는 명시적 기기 ID가 있는 메시: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

매개변수:

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
  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의 측정기준을 갖기 때문입니다.

is_custom_rule은 이 규칙이 stablehlo.custom_call 작업에 대해 사용자가 정의한 규칙인지 여부를 설명합니다. 파티션 나누기는 이러한 작업을 분할하는 방법을 알지 못하므로 사용자가 방법을 알려야 합니다. 맞춤 규칙인 경우 규칙은 항상 보존되며 삭제되지 않습니다. is_custom_rule는 stablehlo.custom_call 연산에 대해서만 true일 수 있습니다.

매개변수:

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의 곱으로 표현할 수 있습니다. 따라서 하위 축 정보 속성에는 이 두 숫자가 저장되며 사전 크기 m과 크기 k의 경우 (m)k과 같이 표시됩니다.

매개변수:

TensorMappingAttr

텐서의 각 차원에 대한 인수 매핑

구문:

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

매개변수:

TensorShardingAttr

텐서 샤딩

구문:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>   # replicated_axes
>

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

이 샤딩이 바인딩된 메시는 기호 이름, 상응하는 MeshOp 기호 참조 또는 인라인 MeshAttr로 지정할 수 있습니다.

매개변수:

TensorShardingPerValueAttr

연산의 피연산자/결과당 텐서 샤딩

구문:

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

매개변수:

열거형

PropagationDirection

전파 방향 enum

케이스: