이 페이지는 Cloud Translation API를 통해 번역되었습니다.
Switch to English

XLA 용 MLIR CodeGen

XLA는 HloInstruction 작동 HloInstruction 표현에 대해 많은 최적화를 수행하여 대상 장치간에 많은 것을 공유합니다. 어느 시점에서 선형 스케줄이 계산되고 메모리 버퍼가 각 값에 정적으로 할당됩니다. 장치 별 codegen은이 시퀀스를 순회하고 "이미 터"를 호출하여 장치에 적합한 표현을 생성합니다 (예 : CPU에서 XLA 계산 ​​당 단일 LLVM 함수 또는 GPU 작업을 캡슐화하는 "썽크"시퀀스 및 PTX가 생성 될 때 타겟팅 GPU).

준비 단계로서, 우리는 현재 XLA가 버퍼 할당 단계를 완료하고 대신 lhlo 언어로 MLIR 모듈을 방출 한 직후에 프로세스를 가로채는 과정에 있습니다. 여기에서 장치에 따라 MLIR 구성 요소 (주로 Linalg, affine 및 GPU 방언)를 사용하여 codegen을 수행합니다.

다음은 lhlo 를 codegen 입력으로 사용하여 XLA / GPU를 점진적으로 마이그레이션하는 레코드 계획입니다.

작업

주최자 장치
입력 형식 HloInstruction * (작업 1) HloInstruction * (작업 1)
출력 형식 xla :: 썽크 (작업 2) LLVM IR (작업 3)
  • 작업 1 은 호스트 및 장치 입력 형식을 HloInstruction *에서 LHLO로 변경합니다.
  • 작업 2 는 호스트의 출력 형식을 썽크에서 "호스트 용 랜딩 패드"로 변경합니다 (아래 참조).
  • 작업 3 은 LLVM IR의 장치 출력을 일부 형태의 MLIR로 마이그레이션합니다. 이 프로젝트의 옵션이며 자세한 내용은 "장치 LLVM IR 마이그레이션"섹션을 참조하십시오.

이 프로젝트는 LHLO 이미 터가 가능한 한 엔드-투-엔드 실행 가능 모델의 우선 순위를 정합니다. 이는 우선 순위 별 목표의 다음 순서 목록을 의미합니다.

  • 기존 Thunk 및 이미 터가 수정되지 않은 LHLO 이미 터로 XLA / GPU를 실행 가능하게 만듭니다.
  • 경우에 따라 LHLO에서 HloInstruction *에 대한 참조를 제거하십시오.
    • 레거시 이미 터를 MLIR 기반 이미 터 (예 : Linalg)로 전환하거나
    • MLIR 표현을 취하도록 기존 이미 터를 기계적으로 변환합니다 (GPU Dialect를 사용하여 표준으로 마이그레이션).

썽크 마이그레이션 (작업 2)

xla :: gpu :: Thunk는 다음과 같은 데이터 구조입니다.

  • 호스트 (xla :: gpu :: Thunk :: ExecuteOnStream ())에서 호출 할 수 있습니다.
  • 서브 클래스에 다양한 데이터를 전달합니다.
  • BufferAllocation :: Slice 및 StreamExecutor와 상호 작용합니다.
  • 커널을 시작합니다
  • 모든 런타임 라이브러리를 호출합니다.

그 비용은 다음과 같습니다.

  • op 별 구성 데이터 (예 : 컨볼 루션 구성)를 나타냅니다.
  • op 셰이프 및 피연산자 셰이프 마이그레이션
  • 썽크 트리를 나타내는 동안 (조건, 조건 등).

마이그레이션 작업은 LHLO / 이미 터 마이그레이션과 독립적입니다. 제한된 리소스에서는 LHLO / 이미 터 마이그레이션보다 우선 순위가 높습니다.

LHLO에서 호스트 측 부분을 낮추는 방법에 대한 몇 가지 선택 사항이 있습니다.

  • TFRT
    • (Pro) 훌륭한 CUDA 및 HIP 래퍼 사용.
    • (Pro) TFRT op가 C ++ 코드로 해석되므로 라이브러리 호출 (cuDNN, cuBLAS, cuFFT 등)을 쉽게 구현할 수 있습니다.
    • (Con) 호스트 측은 개발 중이며 테스트되지 않았습니다.
  • 지트 CPU 코드
    • (프로) 뛰어난 저 능력. 몇 가지 루프와 조건을 작성하면 완료됩니다.
    • (Con) GPUDialect는 아직 체인 / 스트림 / 비동기 성 / 장치 할당을 모델링하지 않습니다.
    • (Con) CUDA / HIP 런타임 지원은 최소한입니다 (툴킷 경로, 버전, 동적 로딩 등).
  • 기존 (해석) XLA 런타임

의사 결정 : TFRT를 채택하고 TFRT에서 CPU 코드 지팅도 지원합니다.

장치 LLVM IR 마이그레이션 (작업 3)

원소 이미 터는 요소별로 요소를 채워서 대상 op를 생성합니다. 각 출력 요소는 피연산자의 요소 세트에 따라 다릅니다. 버퍼를 동적 인덱스와 결합하여 모든 요소를 ​​설명합니다. 거의 모든 "math"연산을 설명하는 것으로 충분하지만 성능상의 이유로 "math"연산의 큰 부분 집합 만 (Cpu | Gpu) ElementalIrEmitter에서 직접 구현됩니다.

ElementalIrEmitter는 다음에서 고유합니다.

  • 코드의 많은 부분이 XLA / GPU와 CPU간에 공유됩니다.
  • 모든 요소 별 연산을 포함하여 모델에서 볼 수있는 대부분의 연산을 나타냅니다.
  • 대부분의 융합은 ElementalIrEmitter에만 의존합니다.
  • op 요소와 피연산자 요소 사이의 데이터 종속성 DAG를 설명하므로 구조적으로 간단합니다.
  • GPU kReduce 및 GPU kCopy와 달리 대부분 휴대용이며 높은 수준입니다.
  • 동적 요소 지원은 최소한 요소 별 연산에 용이합니다.

이제 원소 방출 여부에 관계없이 모든 연산에 대해 각 XLA 연산의 종료 상태에는 몇 가지 맛이 있습니다.

  1. 장치 코드는 LLVM IR로 유지됩니다.
  2. 이전 이미 터를 LHLO-> MLIR LLVM Dialect와 같이 리팩터링하십시오.
    • (비용) 궁극적으로 표준으로 마이그레이션하려는 경우 버리기 작업이됩니다.
    • (장점) 쉽고 기계적입니다. 단기간에 수행 할 수 있습니다.
    • (이점) (1)에 비해 더 많은 이점이 없습니다.
  3. 기존 이미 터를 LHLO-> MLIR GPU + 표준 + 루프처럼 리팩토링하십시오.
    • (비용) 기존 이미 터를 표준으로 올리면 몇 가지 문제가 발생합니다. 포인터와 GEP는 MemRefs와 SubView로 변환해야합니다. Amdgpu 완전성을 보장하는 것은 또 다른 것입니다.
    • (비용) XLA / GPU는 LLVM 메타 데이터에 크게 의존합니다.
      • 블록 / 스레드 인덱스의 range .
      • 로드 / 저장에 대한 align , dereferenceable , invariant.load , alias.scope , noalias
      • 순차 루프의 경우 llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable .
    • (장점) 장기적 일 수 있습니다. 더 휴대하기 편리합니다.
  4. 기존 이미 터를 LHLO-> Linalg로 리팩터링하고 새로운 Linalg 이미 터 작성
    • (비용) 이것은 경우에 따라 다릅니다. 이전 옵션과 비교하여 XLA의 성능과 일치하는 새로운 구현은 벤치 마크 <-> 최적화 워크 플로우를 거쳐야하며, 이는 일부 운영에 상당한 비용이 될 수 있습니다.
    • (이점) 통합 스택; 지역 사회 지원; 휴대 성; 더 많은 최적화 잠재력.

결론 :

  • 가지 마세요 (2). (1) 또는 (3)이 (2)보다 낫습니다. (2) 많은 기계적 리팩토링이 필요하기 때문에 (1) 이상의 비용이 듭니다. (1)을 통해 XLA가 MLIR 이미 터를 픽업 할 수 있도록하는 목표를 달성 할 수 있습니다. 이것은 LHLO-> LLVM IR-> 레거시 장치 이미 터 실행을 수행하는 것입니다.
  • ElementalIrEmitter ops는 (4)가되지만 점진적으로 증가하지는 않습니다. 모든 원소 방출 op가 동일한 그래프에 연결되어 있기 때문에 op로 op를 수행 할 수있는 방법이 없습니다. 이 작업은 여러 진행중인 힘 (xla / service / mlir_gpu, 커널 생성기, Linalg)의 통합 지점으로도 사용될 수 있습니다.
  • 다른 모든 작전은 (1)로 진행됩니다. 확장 목표로, 이들은 (3) 또는 (4)로 마이그레이션 될 수 있습니다.

우선 순위

위에서 언급 한 세 가지 작업은 모두 병렬화 할 수 있지만 제한된 리소스에서는 직렬화해야합니다. 우선 순위는 각 작업을 완료하기위한 가시적 인 결과에 중점을 둡니다.

우선 순위는 태스크 1 (레거시 이미 터의 경우 LHLO)> 태스크 2 (썽크)> 태스크 3 (MLIR 이미 터)입니다.

작업 1이 끝나면 XLA 사용자는 LHLO (예 : 커널 생성기)를 생성하고 실행할 수 있습니다. 컴파일 형식은 직렬화 가능한 MLIR이 아닙니다.

작업 2가 끝나면 LHLO는 적절한 직렬화 가능한 MLIR로 낮아집니다. 오프라인 컴파일이 가능합니다.

작업 3이 끝날 때 모든 XLA 이미 터는 구현시 MLIR 기반입니다.

상세 설계

1 단계 : (작업 1) LHLO 완료 및 레거시 이미 터가 LHLO 사용

이 단계는 모든 기존 XLA / GPU 이미 터가 MLIR op와 상호 작용하도록합니다. 이 단계는 순수한 리팩토링 및 NFC입니다.

이 단계는 대부분 기계적인 것이지만, 중첩되지 않은 HloComputation과 LHLO 간의 다음 불일치를 주목할 가치가 있습니다.

  • 각 HloInstruction은 피연산자 (데이터 흐름 DAG)에 직접 액세스 할 수 있습니다. 반대로 각 LHLO op는 피연산자 버퍼 (ops와 버퍼 사이의 이분)에만 액세스 할 수 있습니다. LHLO op는 피연산자 op에 액세스하기 위해 use-def 체인을 거쳐야합니다.
  • 중첩되지 않은 레거시 이미 터는 경험적으로 피연산자에 거의 액세스하지 않습니다. 유일한 예외는 kReduce입니다.
  • 중첩되지 않은 레거시 이미 터는 슬라이스를 가져 오기 위해서만 BufferAssignment에 액세스하며 dataflow_analysis () 또는 alias_analysis ()와 같은 보조 데이터 구조에는 액세스하지 않습니다. llvm_ir은 슬라이스 정보를 기반으로 고유 한 alias_analysis ()를 빌드합니다.

결론은 LHLO가 큰 번거 로움없이 바로 맞아야한다는 것입니다.

2 단계 : (선택 사항) 프로파일 링 지원

이 단계는 XLA Thunk 논리 중 일부를 폐기하기 시작하는 경우에만 필요합니다 (다음 단계 참조).

MLIR 기반 이미 터를 실제로 켜기 전에 MLIR 기반 이미 터에 대한 프로파일 링이 필요합니다.

현재 XLA는 StreamExecutor의 타이머를 호출하여 자체 프로파일 링을 수행합니다. 후드 아래의 타이머는 커널 시작 전후에 두 개의 이벤트를 삽입하고이 두 이벤트 사이의 동기화 시간을 측정합니다.

MLIR에서 프로파일 링을 지원하는 대략 3 가지 접근 방식이 있습니다.

  • 종단 간 프로파일 러 실행
  • 주입 된 프로파일 러를 사용하여 LHLO의 각 op에 대한 프로파일 op를 추가하십시오.

"end-to-end"접근 방식은 MLIR에 투명하지만 XLA를 우선 사용하지 않는 것과 같은 문제가 있습니다. 프로파일 러 (nvprof / ...)에 의해 수집 된 라이브러리 호출은 HLO와 쉽게 관련 될 수 없습니다 작전. 예를 들어, cuDNN은 각 HLO에 대해 여러 개의 커널을 시작하며 어떤 커널이 어떤 HLO에 해당하는지 알기가 어렵습니다.

"주입 된 프로파일 러"접근 방식에는 다음이 필요합니다.

  • 프로파일 러를 매개 변수로 사용하는 LHLO.
  • 각 op 전후에 profile.start / profile.end 삽입
  • 그것에서 profile. {start, end}를 C ++ 구현으로 전달합니다.

MLIR 생성 op에 대해서는 정확한 프로파일 링을 쉽게 수행 할 수 없습니다.

  • MLIR에는 타이머가 없으며 TFRT / StreamExecutor에 의존합니다.
  • MLIR은 복잡한 매개 변수를 사용하여 C 함수를 쉽게 호출하지 않습니다.

3 단계 : (작업 2) 썽크 마이그레이션

참고로 대략 3 가지 종류의 썽크가 있습니다.

  • KernelThunk는 커널을 시작합니다.
  • 호스트 제어 흐름 논리 (조건부, 동안, 순서) 및 시작 본문 커널이있는 제어 흐름 썽크.
  • 라이브러리 썽크 : cuDNN, cuBLAS, cuFFT, NCCL 등

계획은 다음과 같습니다

  • 썽크 직렬화 가능하게 만듭니다.
  • 이러한 의미를 지원할 수있는 상태로 TFRT를 개선하십시오.
  • 상태가 개선되면 개별 썽크를 점진적으로 마이그레이션하십시오.

이러한 조치 항목은 부분적으로 만 주문됩니다. 실제 실행 순서 / 엔지니어링 병렬 처리는 그대로 평가됩니다.

4 단계 : (작업 3) 마이그레이션 된 ElementalIrEmitter

프로파일 링이 준비되면 MLIR에서 모든 ElementalIrEmitter 기반 이미 터를 완료하고 조정할 수 있습니다. 그런 다음 모든 MLIR 기반 이미 터가 단일 스트림을 사용한다고 가정하여 기본적으로 설정합니다.

XLA / CPU의 ElementalIrEmitter도 코드의 많은 부분을 공유하기 때문에 마이그레이션하는 것이 좋습니다.

모든 벤치마킹 및 성능 헌팅 (TODO : 성능 패리티 정의)이 완료되면 새로운 MLIR 기반 Elemental Emitter를 켜고 레거시 ElementalIrEmitter를 삭제합니다.

이 단계는 또한 나중에 마이그레이션 할 수있는 간편한 융합 전환 (중첩 된 운영 체제)을 제공합니다.

5 단계 : 멀티 스트림 지원 또는 삭제

MLIR에서 이미 터 를 지원하거나 피처를 드롭 할 때까지 일부 이미 터를 삭제할 수 없습니다. MLIR에서 비교적 많은 양의 작업과 XLA의 적은 양입니다. 멀티 스트림 XLA / GPU 사용자의 현재 사용자를 조사하고 합리적인 경우이 기능을 삭제해야합니다.

6 단계 : (작업 3) 마이그레이션 된 기기 운영

이 단계에서는 중첩되지 않은 모든 op를 마이그레이션 한 다음 중첩되지 않은 이미 터를 모두 삭제할 수 있습니다.

이것은 kCopy 및 kReduce에 대한 다시 쓰기 / 리 팩터를 호출합니다. kReduce는 이미 많은 작업을 수행 했으므로 수행해야하는 실제 작업량은 여전히 ​​남아 있습니다.