'mpmd' 언어

작업

mpmd.assign (mpmd::AssignOp)

할당 작업

구문:

operation ::= `mpmd.assign` attr-dict $tensor `:` functional-type(operands, results)

메시 내에서 완전히 복제된 것으로 메시에 로컬 텐서를 할당합니다.

이는 jax 작업을 낮출 때 로컬 유형에서 메시 유형으로 이동하기 위해 도입된 임시 작업입니다. 이러한 작업은 가져오기 중에 제거됩니다. 이때 func 작업의 입력과 결과가 메시 텐서가 됩니다.

결과 유형의 메시 이름은 토폴로지의 메시와 일치해야 하며 전역 유형은 피연산자 유형과 동일해야 합니다.

할당 작업의 출처는 메쉬의 출처입니다(예: named_computation, 메쉬 추론 등).

특성: AlwaysSpeculatableImplTrait, HasParent<::mlir::func::FuncOp, ForOp>

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

효과: MemoryEffects::Effect{}

속성:

피연산자:

결과:

mpmd.broadcast (mpmd::BroadcastOp)

브로드캐스트 작업

구문:

operation ::= `mpmd.broadcast` attr-dict $tensor `:` type($tensor)

텐서가 사용되는 모든 메시에서 텐서를 전송 (또는 복제)할 수 있습니다. 전송될 때마다 전송의 출처는 피연산자의 현재 위치입니다.

특성: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

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

효과: MemoryEffects::Effect{}

피연산자:

결과:

mpmd.call (mpmd::CallOp)

MPMD 특정 호출 함수

구문:

operation ::= `mpmd.call` $callee `(` $tensors `)` attr-dict `:` functional-type(operands, results)

함수 호출 작업입니다. 예를 들어 코드 크기를 줄이기 위해 함수 선언에서 루프 본문을 래핑하는 데 유용합니다.

인터페이스: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

mpmd.for (mpmd::ForOp)

운영자

고정된 횟수의 반복에 대해 본문 함수를 실행한 결과를 반환하며, 반복 색인은 본문에서 사용할 수 있습니다.

반복 횟수를 나눌 수 있는 선택적 언롤링 요소는 해당 요소로 작업의 본문을 언롤링하도록 지정할 수 있습니다. 즉, 언롤링 요소 N의 경우 본문이 복제되어 N개의 사본이 생성되고 반복 횟수가 1/N만큼 감소합니다. 첫 번째 복사본을 제외한 각 복사본은 블록 인수 대신 이전 복사본의 결과를 사용하며 반복 색인은 언롤링 요소가 곱해지고 각 복사본 후에 증가합니다.

for 연산자는 모든 유형을 허용하고 반환할 수 있지만 이러한 유형의 TypeID는 동일해야 합니다(예: 모든 텐서 유형 또는 모든 MPMD 메시 유형 등). 이를 통해 다양한 수준에서 구현과 변환을 공유하는 작업을 사용할 수 있습니다.

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

인터페이스: ConditionallySpeculatable, LoopLikeOpInterface, OpAsmOpInterface, ShardableDataFlowOpInterface

속성:

피연산자:

결과:

mpmd.fragment (mpmd::FragmentOp)

프래그먼트 작업

MPMD 토폴로지의 특정 메시에 계산(작업 블록)을 할당합니다. 이는 개별 SPMD 프로그램 프래그먼트로 실행하기 위한 것입니다.

프래그먼트는 프래그먼트와 동일한 메시에 할당된 메시 텐서만 가져오고 반환합니다.

프래그먼트의 메시 이름은 토폴로지의 메시와 일치해야 합니다.

이 프래그먼트에는 출처 목록, 즉 이 프래그먼트를 형성한 원래 named_computations에 관한 정보가 포함된 메타데이터와 사용자 정의 프래그먼트인 경우, 즉 출처 목록이 비어 있지 않은 경우 iff로 정의된 staged_id가 포함됩니다. 선택적 in_shardings는 피연산자에 해당하는 프래그먼트의 블록 인수의 샤딩을 지정합니다. 선택사항인 out_shardings는 결과의 샤딩을 지정합니다.

프래그먼트의 리전에는 자유 변수가 없어야 하며 리전의 각 블록 인수와 반환 값의 유형은 해당 메시 텐서의 전역 텐서 유형입니다.

특성: HasParent<::mlir::func::FuncOp, ForOp>, IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

인터페이스: ConditionallySpeculatable, ShardableDataFlowOpInterface

속성:

피연산자:

결과:

mpmd.fragment_call (mpmd::FragmentCallOp)

프래그먼트 호출 작업

MPMD 프레그먼트 본문을 보유하는 함수 호출을 나타냅니다. 즉, MPMD 토폴로지의 특정 메시지에 할당된 계산으로, 개별 SPMD 프로그램 프레그먼트로 실행되도록 설계되었습니다.

프래그먼트의 메시 이름은 둘러싸는 함수의 토폴로지에 있는 메시와 일치해야 하며 메시 모양은 호출자의 메시 모양과 일치해야 합니다.

출처는 병합 등을 통해 이 프래그먼트 호출에 기여한 사용자 명명 계산을 지정합니다.

호출된 함수의 입력 및 결과 유형은 이 작업의 피연산자 및 결과의 해당 메시 텐서의 로컬 텐서 유형이어야 합니다.

예:

%2 = mpmd.fragment_call<mesh="m1",> origin=[] @my_fragment(%0, %1) :
  (m<esh>_tensor..., m<esh>_te>nsor...) - m<esh>_tensor...

특성: HasParent<::mlir::func::FuncOp>, MemRefsNormalizable

인터페이스: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

속성:

피연산자:

결과:

mpmd.named_computation (mpmd::NamedComputationOp)

이름이 지정된 범위 작업

계산(작업 블록)을 그룹화하고 UserOrigin 속성을 통해 이름과 전치 수를 지정합니다. 이 NamedComputation은 MPMD의 계산에 메시를 할당하거나 최적화하는 데 사용할 수 있습니다.

전치 개수 (기본값=0)는 명명된 계산이 특정 수의 JAX AD 전치 변환에 의해 생성되었는지 여부를 나타냅니다.

작업의 리전에는 자유 변수가 없어야 하며 리전의 각 블록 인수와 반환 값의 유형은 입력 유형 및 작업의 반환 유형과 동일해야 합니다.

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

인터페이스: ConditionallySpeculatable

속성:

피연산자:

결과:

mpmd.named_tensor (mpmd::NamedTensorOp)

메시에 텐서 할당

구문:

operation ::= `mpmd.named_tensor` $tensor `name````=```$name attr-dict `:` type($result)

텐서의 결과를 지정된 이름과 연결하는 ID 작업입니다. 이 NamedTensor는 MPMD에서 텐서에 메시를 할당하는 데 사용할 수 있습니다.

특성: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

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

효과: MemoryEffects::Effect{}

속성:

피연산자:

결과:

mpmd.reduce (mpmd::ReduceOp)

크로스 메시 축소 작업

구문:

operation ::= `mpmd.reduce` `` $reduction attr-dict $tensors `:` functional-type(operands, results)

텐서를 여러 메시에서 축소한 다음 사용해야 하는 곳에 브로드캐스트할 수 있습니다.

특성: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

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

효과: MemoryEffects::Effect{}

속성:

피연산자:

결과:

mpmd.return (mpmd::ReturnOp)

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

구문:

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

특성: AlwaysSpeculatableImplTrait, Terminator

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

효과: MemoryEffects::Effect{}

피연산자:

mpmd.transfer (mpmd::TransferOp)

전송 작업

구문:

operation ::= `mpmd.transfer` attr-dict $tensor `:` functional-type(operands, results)

분산 텐서를 한 메시에서 다른 메시로 전송합니다.

피연산자와 결과 유형의 메시 이름은 토폴로지의 메시와 일치해야 하며 전역 유형은 동일해야 합니다.

특성: AlwaysSpeculatableImplTrait, HasParent<::mlir::func::FuncOp>

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

효과: MemoryEffects::Effect{}

피연산자:

결과:

mpmd.unassign (mpmd::UnassignOp)

할당 해제 작업

구문:

operation ::= `mpmd.unassign` attr-dict $tensor `:` functional-type(operands, results)

메시에서 완전히 복제된 텐서를 할당 해제합니다.

이는 jax 작업을 낮출 때 로컬 유형에서 메시 유형으로 이동하기 위해 도입된 임시 작업입니다. 이러한 작업은 가져오기 중에 제거됩니다. 이때 func 작업의 입력과 결과가 메시 텐서가 됩니다.

피연산자 유형의 메시 이름은 토폴로지의 메시와 일치해야 하며 전역 유형은 결과 유형과 동일해야 합니다.

특성: AlwaysSpeculatableImplTrait, HasParent<::mlir::func::FuncOp, ForOp>, InferTensorType

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

효과: MemoryEffects::Effect{}

속성:

피연산자:

결과:

속성

MeshWithOriginsAttr

기원과 조화롭게 어울립니다.

구문:

#mpmd.mesh_with_origins<
  ::llvm::StringRef,   # mesh_name
  ::llvm::ArrayRef<OriginAttr>   # origins
>

매개변수:

MeshesWithOriginsAttr

원점이 포함된 메시 목록

구문:

#mpmd.meshes_with_origins<
  ::llvm::ArrayRef<MeshWithOriginsAttr>   # value
>

매개변수:

NamedMeshAttr

이름과 메시가 있는 쌍입니다.

구문:

#mpmd.named_mesh<
  ::llvm::StringRef,   # name
  sdy::MeshAttr   # mesh
>

매개변수:

OriginAttr

메시 할당의 출처입니다.

구문:

#mpmd.origin<
  ::llvm::StringRef   # origin_label
>

메시 할당의 출처입니다.

origin_label은 출처에 대한 사람이 읽을 수 있는 라벨입니다. 디버깅 목적으로 사용하기 위한 것입니다.

매개변수:

ReductionAttr

감소를 나타냅니다.

구문:

#mpmd.reduction<
  `none` | `add` | `max` | `min` | `mul` | `or` | `and`   # reduction_type
>

매개변수:

TopologyAttr

이름이 지정된 메시의 토폴로지입니다.

구문:

#mpmd.topology<
  ::llvm::ArrayRef<NamedMeshAttr>   # meshes
>

매개변수:

UserOriginAttr

사용자 지정 계산의 출처입니다.

구문:

#mpmd.user_origin<
  ::mlir::StringAttr,   # userName
  int64_t   # transposeCount
>

매개변수:

유형

MeshTensorType

메시 텐서 유형

프로그램의 MPMD 메시 토폴로지에서 특정 SPMD 메시에 RankedTensorType을 할당합니다. 이 유형은 텐서가 SPMD 메시지와 관련하여 샤딩되는 방식을 지정하는 선택적 샤딩을 보유합니다. 샤딩이 없으면 텐서가 완전히 복제됩니다.

매개변수:

열거형

ReductionType

감소 유형 속성

케이스:

EdgeNodeType

에지 노드 유형 열거형

케이스:

PropagationDirection

전파 방향 열거형

케이스:

ReductionOp

감소 작업 enum

케이스: